.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "auto_examples/scientific_machine_learning/fourier_neural_operator/burgers.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code or to run this example in your browser via Binder .. rst-class:: sphx-glr-example-title .. _sphx_glr_auto_examples_scientific_machine_learning_fourier_neural_operator_burgers.py: Learning the Burgers' Operator ============================== In this example, we train a Fourier Neural Operator to learn the mapping :math:`y(x, 0) \mapsto y(x, 0.5)` where :math:`y(x,t)` is the solution to the Burgers' equation given by .. math:: \frac{\partial}{\partial t}u(x, t) + u(x, t) \frac{\partial}{\partial x} u(x, t) = \nu \frac{\partial^2}{\partial x^2} u(x,t) .. GENERATED FROM PYTHON SOURCE LINES 13-17 We learn the mapping from the initial condition :math:`y(x, 0)` to a point at time :math:`t=0.5`. The training data for this example was computed separately and is available on GitHub. First, import the necessary modules. .. GENERATED FROM PYTHON SOURCE LINES 19-34 .. code-block:: default import logging import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader import matplotlib.pyplot as plt import UQpy.scientific_machine_learning as sml torch.manual_seed(123) plt.style.use("ggplot") logger = logging.getLogger("UQpy") logger.setLevel(logging.INFO) .. GENERATED FROM PYTHON SOURCE LINES 35-36 Below we define the data and architecture of our neural operator as subclasses of PyTorch objects. .. GENERATED FROM PYTHON SOURCE LINES 38-109 .. code-block:: default class BurgersDataset(Dataset): def __init__(self, filename_initial_condition: str, filename_solution: str): r"""Construct a dataset for training a FNO to learn the Burgers' operator :math:`\mathcal{B}:y(x,0) \to y(x, t^*)` :param filename_initial_condition: Relative path to the `.pt` file containing a torch tensor of shape :math:`(N_\text{samples}, 1, N_x)` :param filename_solution: Relative path to the `.pt` file containing a torch tensor of shape :math:`(N_\text{samples}, 1, N_x, N_t)` """ self.filename_initial_condition = filename_initial_condition self.filename_solution = filename_solution self.initial_conditions = torch.load(self.filename_initial_condition) self.burgers_solution = torch.load(self.filename_solution) self.n_samples = self.initial_conditions.size(0) def __len__(self): return self.n_samples def __getitem__(self, item): return self.initial_conditions[item], self.burgers_solution[item] class FourierNeuralOperator(nn.Module): def __init__(self): """Construct a Bayesian FNO The lifting layers a single, deterministic, and fully connected linear layer. There are four Bayesian Fourier layers performing a 1d Fourier transform with ReLU activation functions and batch normalization in between. The projection layers are also a single, deterministic, and fully connected linear layer. """ super().__init__() modes = 16 width = 8 self.lifting_layers = nn.Sequential( sml.Permutation((0, 2, 1)), nn.Linear(1, width), sml.Permutation((0, 2, 1)), ) self.fourier_blocks = nn.Sequential( sml.Fourier1d(width, modes), nn.ReLU(), sml.Fourier1d(width, modes), nn.ReLU(), sml.Fourier1d(width, modes), nn.ReLU(), sml.Fourier1d(width, modes), ) self.projection_layers = nn.Sequential( sml.Permutation((0, 2, 1)), nn.Linear(width, 1), sml.Permutation((0, 2, 1)), ) def forward(self, x: torch.Tensor) -> torch.Tensor: """Forward computational call of the FNO. Apply the lifting layers, then theFourier blocks, and finally projection layers. :param x: Tensor of shape :math:`(N, C, L)` :return: Tensor of shape :math:`(N, C, L)` """ x = self.lifting_layers(x) x = self.fourier_blocks(x) x = self.projection_layers(x) return x .. GENERATED FROM PYTHON SOURCE LINES 110-111 We create two instances of the :code:`BurgersDataset` class, one for training data and one for testing data. .. GENERATED FROM PYTHON SOURCE LINES 113-125 .. code-block:: default train_dataset = BurgersDataset( "initial_conditions_train.pt", "burgers_solutions_train.pt", ) test_dataset = BurgersDataset( "initial_conditions_test.pt", "burgers_solutions_test.pt", ) train_data = DataLoader(train_dataset, batch_size=100, shuffle=True) test_data = DataLoader(test_dataset) .. GENERATED FROM PYTHON SOURCE LINES 126-131 Finally, we are ready to instantiate the model and train is using the Adam optimizer and UQpy's :py:class:`sml.Trainer`. The model is quite small, with only a few thousand parameters, and we only train for 100 epochs to save on computational time. In many applications much larger Fourier neural operators are needed and are trained for longer. .. GENERATED FROM PYTHON SOURCE LINES 133-143 .. code-block:: default model = FourierNeuralOperator() print(model) optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.1) trainer = sml.Trainer(model, optimizer, scheduler=scheduler) trainer.run(train_data=train_data, test_data=test_data, epochs=100) # torch.save(model.state_dict(), "fno_state_dict.pt") # optionally save the trained model .. GENERATED FROM PYTHON SOURCE LINES 144-145 The hard part is done! The rest of this example plots the loss history and a prediction from the trained model. .. GENERATED FROM PYTHON SOURCE LINES 147-184 .. code-block:: default # plot loss history fig, ax = plt.subplots() ax.semilogy(trainer.history["train_loss"], label="Train Loss") ax.semilogy(trainer.history["test_loss"], label="Test NLL") ax.set_title("Deterministic FNO Training History") ax.set(xlabel="Epoch", ylabel="MSE") ax.legend() fig.tight_layout() # plot FNO prediction # optional, load weights from save file. fno_state_dict.pt is available on our GitHub. # model.load_state_dict(torch.load("fno_state_dict.pt", weights_only=True)) x = torch.linspace(0.0, 1.0, 256) # spacial domain fig, ax = plt.subplots() initial_condition, burgers_solution = train_dataset[0:1] with torch.no_grad(): prediction = model(initial_condition) initial_condition = initial_condition.squeeze() burgers_solution = burgers_solution.squeeze() prediction = prediction.squeeze() ax.plot(x, initial_condition, label="$y(x,0)$", color="black", linestyle="dashed") ax.plot( x, burgers_solution, label="$y(x, 0.5)$", color="black", ) ax.plot(x, prediction, label="FNO $\hat{y}(x, 0.5)$", color="tab:blue") ax.set_title("Deterministic FNO on Training Data", loc="left") ax.set(ylabel="$y(x, t)$", xlabel="Time $t$") ax.legend(fontsize="small", ncols=2) plt.show() .. rst-class:: sphx-glr-timing **Total running time of the script:** ( 0 minutes 0.000 seconds) .. _sphx_glr_download_auto_examples_scientific_machine_learning_fourier_neural_operator_burgers.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: binder-badge .. image:: images/binder_badge_logo.svg :target: https://mybinder.org/v2/gh/SURGroup/UQpy/master?urlpath=lab/tree/notebooks/auto_examples/scientific_machine_learning/fourier_neural_operator/burgers.ipynb :alt: Launch binder :width: 150 px .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: burgers.py ` .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: burgers.ipynb ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_