FNO 2D Model Documentation

Overview

The Fourier Neural Operator (FNO) is a neural network architecture designed for solving partial differential equations (PDEs) by learning operators in the Fourier domain. The FNO leverages Fourier transforms to efficiently model spatial relationships across different resolutions, making it ideal for applications with complex dynamics and high-resolution data.

This implementation focuses on the 2D case and is integrated into the fair-sciml project for learning solutions of PDEs from simulation data.

Usage Example

Below is an example of how to use the implemented FNO 2D model for training:

from ml.fno_trainer import FNOTrainer

# Initialize the trainer
trainer = FNOTrainer(
    file_path="simulations/poisson_equation.h5",
    batch_size=32,
    lr=1e-3,
    epochs=50
)

# Run the training pipeline
trainer.run()

Model Architecture

The FNO 2D model operates on grid-based data with the following architecture:

  • Input Channels: - Number of input channels corresponds to the number of input fields (e.g., field_input_f). - Example for Poisson Equation: 2 channels (source field, boundary field).

  • Hidden Channels: - Default: 64 hidden channels in the intermediate Fourier layers.

  • Output Channels: - 1 output channel, representing the solution field.

Layers

  • Fourier Layers: - Captures global spatial dependencies using Fourier transforms. - Operates on the input grid with specified modes in the height and width dimensions.

  • Modes: - n_modes_height=16 - n_modes_width=16

  • Final Projection Layer: - Maps the output of Fourier layers to the solution domain.

Activation Function

  • ReLU activation is applied in intermediate layers.

Optimizer

  • Adam optimizer with configurable learning rate.

Key Features

  • Discretization Invariance: - The FNO 2D model can adapt to various grid resolutions without retraining.

  • Fourier Transform-Based Layers: - Efficiently captures both global and local spatial dynamics.

  • Dynamic Input Handling: - Supports flexible input sizes and multiple input fields.

  • Efficient Data Handling: - Automatically normalizes and prepares data for training.

Training Workflow

  1. Load Data: - HDF5 files are loaded using the load_data method. - Input fields and solution fields are extracted and normalized. - Example fields: field_input_f, solution values.

  2. Prepare Data: - Data is reshaped into tensors for training and testing. - Grid size is calculated dynamically based on the input data.

  3. Model Initialization: - The number of input channels is set dynamically based on the number of fields in the HDF5 file.

  4. Training: - The model is trained using MSE loss and the Adam optimizer. - Metrics such as test loss and L2 relative error are logged during training.

Metrics

The following metrics are logged during training and evaluation:

  • Mean Squared Error (MSE): Measures the average squared difference between predictions and ground truth.

  • L2 Relative Error: Normalized error that accounts for the magnitude of the true solution.

Dependencies

Ensure the following dependencies are installed:

  • NeuralOperator: For implementing the FNO model.

  • PyTorch: For deep learning operations.

  • H5py: For reading HDF5 files.

  • NumPy: For efficient numerical operations.

Install dependencies using the following command:

pip install -r ../../requirements.txt

Advanced Usage

Configurable Fourier Modes

You can adjust the number of Fourier modes in the height and width dimensions by modifying the n_modes_height and n_modes_width parameters in the FNOTrainer class.

Multi-Field Inputs

The FNO model supports multiple input fields, such as source terms and coefficients. Simply add these fields to the HDF5 file and ensure the data loader correctly extracts them.

Extended Metrics

To include custom evaluation metrics, modify the train_model method in the FNOTrainer class.

Contact

For questions or feedback, please contact the project maintainers: