Example 3: Custom Equation DataLoader

In this example we will show how to iterate over the dataset generated by the DynaBench library for a custom equation. We will use the dynabench.dataset.EquationMovingWindowIterator to iterate over the data generated for the dynabench.equation.FitzhughNagumoEquation. The Fitzhugh-Nagumo equation is a simplified model of the electrical activity in a neuron. The equation is given by:

\[\frac{\partial v}{\partial t} = \nabla^2 v + v - \frac{v^3}{3} - w + s\]

\[\frac{\partial w}{\partial t} = \frac{v + a - b * w}{\tau}\]

where \(v\) is the membrane potential and \(w\) is the recovery variable, and \(s\) (stimulus), \(a\), \(b\), and \(\tau\) are parameters of the equation.

We will generate the data for the Fitzhugh-Nagumo equation in a similar way to the previous example for the Cahn-Hilliard Equation.

This example consists of the following steps:

Generate data.
Iterate over simulation.
Load whole simulation.
Summary.

Generate Data

Similar to Cahn-Hilliard Equation using the dynabench.solver.PyPDESolver solver. As the Fitzhugh-Nagumo equation is a two-component equation, we will use the dynabench.initial.Composite class to generate the initial condition for both components of the equation.

from dynabench.equation import FitzhughNagumoEquation
from dynabench.initial import RandomUniform, Composite
from dynabench.grid import Grid
from dynabench.solver import PyPDESolver

# Generate data for the Fitzhugh-Nagumo equation
pde_equation = FitzhughNagumoEquation()
grid = Grid(grid_limits=((0, 64), (0, 64)), grid_size=(64, 64))
intitial = Composite(RandomUniform(), RandomUniform())
solver = PyPDESolver(equation=pde_equation, grid=grid, initial_generator=intitial, parameters={'method': "RK23"})
solver.solve(t_span=[0, 100], dt_eval=1)

In this case the default parameters of the equation are used. The parameters of the Fitzhugh-Nagumo equation can be modified by passing the desired parameters to the constructor of the dynabench.equation.FitzhughNagumoEquation class.

The default parameters of the Fitzhugh-Nagumo equation are as follows:

\(a = 0.0\)
\(b = 0.0\)
\(\tau = 10.0\)
\(stimulus = 0.5\)

Iterate over simulation

Next, we will iterate over the simulation generated by the PyPDESolver using the dynabench.dataset.EquationMovingWindowIterator class. To load the data, we will need to specify the path to the saved equation data, as well as the lookback and rollout parameters.

The dynabench.solver.PyPDESolver saves the data in h5 format to the output directory specified by the out_dir parameter with a default value of “data/raw”. The name of the file is generated based on the equation name and the parameters of the equation. In our case the file name is fitzhughnagumo_357b852b_dt_1_trange_0_100_seed_42.h5, but it might be different in your case. For the lookback and rollout parameters, we will use the default values of 1.

from dynabench.dataset import EquationMovingWindowIterator

# initialize the equation iterator
eq_iterator = EquationMovingWindowIterator(
    data_path = "data/raw/fitzhughnagumo_357b852b_dt_1_trange_0_100_seed_42.h5", # path to the data file, generated by the solver. Might be different for you.
    lookback = 1,
    rollout = 1,
)

The equation iterator can be used in a similar way to the dynabench.dataset.DynaBenchIterator:

import tqdm

# iterate over the data
for sample in tqdm.tqdm(eq_iterator):
    input_data, target_data, points = sample

Load the whole simulation

For some applications (e.g. sparse discovery of goverining equations using SINDy) it might be useful to load the whole simulation length. This can be done using the dynabench.dataset.EquationMovingWindowIterator.get_full_simulation_data() method:

data, points = eq_iterator.get_full_simulation_data()
print(data.shape, points.shape)

The data array will contain the whole simulation data of shape (T, F, H, W), while the points array will contain the spatial points of the grid with shape (H, W, 2). In this case T is the number of time steps in the whole simulation, F is the number of fields in the equation, and H, W are the spatial dimensions of the grid.

Summary

Overall the code for using the dynabench.dataset.EquationMovingWindowIterator is as follows:

from dynabench.equation import FitzhughNagumoEquation
from dynabench.initial import RandomUniform, Composite
from dynabench.grid import Grid
from dynabench.solver import PyPDESolver

from dynabench.dataset import EquationMovingWindowIterator

import tqdm


# Generate data for the Fitzhugh-Nagumo equation
pde_equation = FitzhughNagumoEquation()
grid = Grid(grid_limits=((0, 64), (0, 64)), grid_size=(64, 64))
intitial = Composite(RandomUniform(), RandomUniform())
solver = PyPDESolver(equation=pde_equation, grid=grid, initial_generator=intitial, parameters={'method': "RK23"})
solver.solve(t_span=[0, 100], dt_eval=1)


# initialize the equation iterator
eq_iterator = EquationMovingWindowIterator(
    data_path = "data/raw/fitzhughnagumo_357b852b_dt_1_trange_0_100_seed_42.h5", # path to the data file, generated by the solver. Might be different for you.
    lookback = 4,
    rollout = 16,
)


# iterate over the data
for sample in tqdm.tqdm(eq_iterator):
    input_data, target_data, points = sample

# load the whole simulation length. Warning: for large simulations this might be memory intensive
data, points = eq_iterator.get_full_simulation_data()
print(data.shape, points.shape)