==================================== Example 1: NeuralPDE ==================================== In this example we will show how to train the GrINd 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 GrINd model on the Burgers' equation using the following steps: 1. :ref:`Load the dataset. ` 2. :ref:`Define the neural network architecture. ` 3. :ref:`Train the model. ` 4. :ref:`Evaluate the model. ` 5. :ref:`Summary. ` .. _load_dataset_grind: ************************************ 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 and sampled on a scattered set of points. 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='cloud', resolution='low') burgers_train_iterator = DynabenchIterator(split="train", equation='burgers', structure='cloud', resolution='low', lookback=1, squeeze_lookback_dim=True, rollout=1) train_loader = DataLoader(burgers_train_iterator, batch_size=32, shuffle=True) .. _define_model_grind: ************************************** Define the neural network architecture ************************************** Next, we will define the neural network architecture. We need the GrINd model to handle the interpolation into and back from the grid structure. We will use the NeuralPDE model to solve the Burgers' equation in the high-resolution grid space. 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 from dynabench.model.grind import GrIND prediction_net = NeuralPDE(input_dim=2, hidden_channels=64, hidden_layers=3, solver={'method': 'euler', 'options': {'step_size': 0.1}}, use_adjoint=False) model = GrIND(prediction_net, num_ks=9, grid_resolution=64, spatial_dim=2) The GrIND model requires the following parameters: - prediction_net: The neural network that will be used to predict the solution to the PDE in the high-resolution grid space. - num_ks: The number of frequencies in each spatial dimension to be used in the interpolation. - grid_resolution: The resolution of the grid in each spatial dimension on which the data will be interpolated. .. _train_grind: ************************************** 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, p = x.float(), y.float(), p.float() # only use the first channel and convert to float32 optimizer.zero_grad() y_pred = model(x, p) loss = criterion(y_pred, y) loss.backward() optimizer.step() print(f"Epoch: {epoch}, Batch: {i}, Loss: {loss.item()}") .. _evaluate_grind: ************************************** 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='cloud', 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, p = x.float(), y.float(), p.float() # only use the first channel and convert to float32 y_pred = model(x, p, 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_grind: ************************************** 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 from dynabench.model.grind import GrIND import torch.optim as optim import torch.nn as nn download_equation('burgers', structure='cloud', resolution='low') burgers_train_iterator = DynabenchIterator(split="train", equation='burgers', structure='cloud', resolution='low', lookback=1, squeeze_lookback_dim=True, rollout=1) train_loader = DataLoader(burgers_train_iterator, batch_size=32, shuffle=True) prediction_net = NeuralPDE(input_dim=2, hidden_channels=64, hidden_layers=3, solver={'method': 'euler', 'options': {'step_size': 0.1}}, use_adjoint=False) model = GrIND(prediction_net, num_ks=9, grid_resolution=64, spatial_dim=2) 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, p = x.float(), y.float(), p.float() # only use the first channel and convert to float32 optimizer.zero_grad() y_pred = model(x, p) loss = criterion(y_pred, y) loss.backward() optimizer.step() print(f"Epoch: {epoch}, Batch: {i}, Loss: {loss.item()}") break break burgers_test_iterator = DynabenchIterator(split="test", equation='burgers', structure='cloud', 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, p = x.float(), y.float(), p.float() # only use the first channel and convert to float32 y_pred = model(x, p, t_eval=range(17)) loss = criterion(y_pred, y) loss_values.append(loss.item()) print(f"Mean Loss: {sum(loss_values) / len(loss_values)}")