=====================================
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 :class:`dynabench.dataset.EquationMovingWindowIterator` to iterate over the data generated for 
the :class:`dynabench.equation.FitzhughNagumoEquation`. The Fitzhugh-Nagumo equation is a simplified model of the electrical activity in a
neuron. The equation is given by:

.. math::
    \frac{\partial v}{\partial t} = \nabla^2 v + v - \frac{v^3}{3} - w + s

.. math::
    \frac{\partial w}{\partial t} = \frac{v + a - b * w}{\tau}

where :math:`v` is the membrane potential and :math:`w` is the recovery variable, and :math:`s` (stimulus), :math:`a`, :math:`b`, and :math:`\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 :doc:`Cahn-Hilliard Equation </examples/cahnhiliard>`.

This example consists of the following steps:

1. :ref:`Generate data. <generate_data>`
2. :ref:`Iterate over simulation. <iterate_simulation>`
3. :ref:`Load whole simulation. <get_whole_simulation>`
4. :ref:`Summary. <summary_customequation>`

.. _generate_data:

***************
Generate Data
***************
Similar to :doc:`Cahn-Hilliard Equation </examples/cahnhiliard>` using the :class:`dynabench.solver.PyPDESolver` solver. 
As the Fitzhugh-Nagumo equation is a two-component equation, we will use the :class:`dynabench.initial.Composite` class to generate the
initial condition for both components of the equation.

.. code-block::

    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 :class:`dynabench.equation.FitzhughNagumoEquation` class.

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

- :math:`a = 0.0`
- :math:`b = 0.0`
- :math:`\tau = 10.0`
- :math:`stimulus = 0.5`

.. _iterate_simulation:

**************************************
Iterate over simulation
**************************************

Next, we will iterate over the simulation generated by the PyPDESolver using the :class:`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 :class:`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.

.. code-block::

    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 :class:`dynabench.dataset.DynaBenchIterator`:

.. code-block::

    import tqdm

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

.. _get_whole_simulation:

**************************************
Load the whole simulation
**************************************

For some applications (e.g. sparse discovery of goverining equations using `SINDy <https://pysindy.readthedocs.io>`_) it might be useful to load the whole simulation length.
This can be done using the :meth:`dynabench.dataset.EquationMovingWindowIterator.get_full_simulation_data` method:

.. code-block::

    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_customequation:

**************************************
Summary
**************************************

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

.. code-block::

    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)