====================================
Example 1: NeuralPDE
====================================

In this example we will show how to train the NeuralPDE model on the Burgers' equation. The Burgers' equation is a fundamental partial differential equation in fluid dynamics. The equation describes the evolution of a fluid flow in one dimension. 
The equation is given by:

.. math::
    \frac{\partial u}{\partial t} + u \frac{\partial u}{\partial x} = \nu \frac{\partial^2 u}{\partial x^2}

where :math:`u` is the velocity field, :math:`t` is time, :math:`x` is the spatial coordinate, and :math:`\nu` is the kinematic viscosity.


We will train the NeuralPDE model on the Burgers' equation using the following steps:

1. :ref:`Load the dataset. <load_dataset_neuralpde>`
2. :ref:`Define the neural network architecture. <define_model_neuralpde>`
3. :ref:`Train the model. <train_neuralpde>`
4. :ref:`Evaluate the model. <evaluate_neuralpde>`
5. :ref:`Summary. <summary_neuralpde>`



.. _load_dataset_neuralpde:

************************************
Load the dataset
************************************


Let's start by loading the dataset. We will use the :class:`dynabench.dataset.DynabenchIterator` to download the dataset and iterate over the data. The dataset is generated by solving the Burgers' equation using a finite difference method. 
Additionally, we will use torch.utils.data.DataLoader to create a data loader for the training dataset.

.. code-block::

    from dynabench.dataset import DynabenchIterator, download_equation
    from torch.utils.data import DataLoader

    download_equation('burgers', structure='grid', resolution='low')

    burgers_train_iterator = DynabenchIterator(split="train",
                                               equation='burgers',
                                               structure='grid',
                                               resolution='low',
                                               lookback=1,
                                               squeeze_lookback_dim=True,
                                               rollout=1)

    train_loader = DataLoader(burgers_train_iterator, batch_size=32, shuffle=True)

.. _define_model_neuralpde:

**************************************
Define the neural network architecture
**************************************

Next, we will define the neural network architecture. We will use the NeuralPDE model to solve the Burgers' equation. The NeuralPDE model is a neural network that approximates the solution to the partial differential equation using differentiable ODE solvers and the method of lines.

.. code-block::

    from dynabench.model import NeuralPDE

    model = NeuralPDE(input_dim=2, hidden_channels=64, hidden_layers=3,
                    solver={'method': 'euler', 'options': {'step_size': 0.1}},
                    use_adjoint=False)


.. _train_neuralpde:

**************************************
Train the model
**************************************

Now we will train the model on the training dataset using PyTorch. We will use the Adam optimizer and the Mean Squared Error (MSE) loss function to train the model. We will train the model for 10 epochs.

.. code-block::
    
        import torch.optim as optim
        import torch.nn as nn
    
        optimizer = optim.Adam(model.parameters(), lr=1e-3)
        criterion = nn.MSELoss()

        for epoch in range(10):
            model.train()
            for i, (x, y, p) in enumerate(train_loader):
                x, y = x.float(), y.float() # only use the first channel and convert to float32
                optimizer.zero_grad()
                y_pred = model(x)
                loss = criterion(y_pred, y)
                loss.backward()
                optimizer.step()
                print(f"Epoch: {epoch}, Batch: {i}, Loss: {loss.item()}")


.. _evaluate_neuralpde:

**************************************
Evaluate the model
**************************************

Finally, we will evaluate the model on the test dataset. We will use the test dataset to evaluate the model's performance on unseen data.
To do this we need to load the test dataset and create a data loader for the test dataset.

We want to evaluate the model's performance over a longer time horizon, so we will set the rollout parameter to 16. 
This means that the model will have to predict the next 16 time steps given the input data. 
We can specify this in the forward pass of the model by passing the t_eval parameter to the model.

.. code-block::

    burgers_test_iterator = DynabenchIterator(split="test",
                                              equation='burgers',
                                              structure='grid',
                                              resolution='low',
                                              lookback=1,
                                              squeeze_lookback_dim=True,
                                              rollout=16)

    test_loader = DataLoader(burgers_test_iterator, batch_size=32, shuffle=False)

    model.eval()

    loss_values = []
    for i, (x, y, p) in enumerate(test_loader):
        x, y = x.float(), y.float() # only use the first channel and convert to float32
        y_pred = model(x, t_eval=range(17))
        loss = criterion(y_pred, y)
        loss_values.append(loss.item())

    print(f"Mean Loss: {sum(loss_values) / len(loss_values)}")

.. _summary_neuralpde:

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

Overall the code for training the NeuralPDE model on the Burgers' equation is as follows:

.. code-block::

    from dynabench.dataset import DynabenchIterator, download_equation
    from torch.utils.data import DataLoader
    from dynabench.model import NeuralPDE

    import torch.optim as optim
    import torch.nn as nn

    download_equation('burgers', structure='grid', resolution='low')

    burgers_train_iterator = DynabenchIterator(split="train",
                                               equation='burgers',
                                               structure='grid',
                                               resolution='low',
                                               lookback=1,
                                               squeeze_lookback_dim=True,
                                               rollout=1)

    train_loader = DataLoader(burgers_train_iterator, batch_size=32, shuffle=True)

    model = NeuralPDE(input_dim=2, hidden_channels=64, hidden_layers=3,
                    solver={'method': 'euler', 'options': {'step_size': 0.1}},
                    use_adjoint=False)

    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    criterion = nn.MSELoss()

    for epoch in range(10):
        model.train()
        for i, (x, y, p) in enumerate(train_loader):
            x, y = x.float(), y.float() # only use the first channel and convert to float32
            optimizer.zero_grad()
            y_pred = model(x)
            loss = criterion(y_pred, y)
            loss.backward()
            optimizer.step()
            print(f"Epoch: {epoch}, Batch: {i}, Loss: {loss.item()}")

    burgers_test_iterator = DynabenchIterator(split="test",
                                              equation='burgers',
                                              structure='grid',
                                              resolution='low',
                                              lookback=1,
                                              squeeze_lookback_dim=True,
                                              rollout=16)

    test_loader = DataLoader(burgers_test_iterator, batch_size=32, shuffle=False)

    model.eval()

    loss_values = []
    for i, (x, y, p) in enumerate(test_loader):
        x, y = x.float(), y.float() # only use the first channel and convert to float32
        y_pred = model(x, t_eval=range(17))
        loss = criterion(y_pred, y)
        loss_values.append(loss.item())

    print(f"Mean Loss: {sum(loss_values) / len(loss_values)}")