DeepONet Model Documentation

Overview

The DeepONet model is a neural network architecture specifically designed for learning operators, making it ideal for applications involving Partial Differential Equations (PDEs). Unlike traditional neural networks, which learn mappings between finite-dimensional inputs and outputs, DeepONet learns mappings between functions (i.e., infinite-dimensional spaces). This enables it to approximate solutions of PDEs with high efficiency and flexibility.

The model operates using a two-branch architecture:

  • Branch Network: Encodes the input functions (e.g., source terms, coefficients).

  • Trunk Network: Encodes spatial coordinates of the solution domain.

DeepONet seamlessly combines these two components to predict scalar values corresponding to the solution fields.

Usage Example

Here’s an example of how to set up and train the DeepONet model:

from ml.deeponet_trainer import DeepONetTrainer
from data_loaders.local_loader import LocalLoader

# Initialize the data loader
file_path = "simulations/poisson_equation.h5"
data_loader = LocalLoader(file_path)

# Initialize the trainer
trainer = DeepONetTrainer(
    branch_hidden_layers=[128, 128, 128],
    trunk_hidden_layers=[128, 128, 128],
    data_loader=data_loader
)

# Train the model
trainer.train(epochs=10000, batch_size=32, learning_rate=0.001, metrics_file="metrics.csv")

Model Architecture

General Design

The DeepONet model uses a dual-branch architecture:

  • Branch Network: - Dynamically adjusts the number of input neurons based on the dimensions of the input data (e.g., number of fields in the PDE). - Example for Poisson Equation: 1089 neurons (corresponding to grid points for field_input_f).

  • Trunk Network: - Dynamically adjusts the input size based on the spatial dimensions (e.g., 2 for ((x, y))).

Layers

  • Branch Network: [Dynamic Input Size, 128, 128, 128]

  • Trunk Network: [Dynamic Input Size, 128, 128, 128]

  • Output Layer: Single scalar output per grid point.

Activation Function

  • Branch Network: ReLU

  • Trunk Network: ReLU

Optimizer

  • Adam optimizer with learning rate tuning.

Metrics

The following metrics are calculated during training and evaluation:

  • Mean Squared Error (MSE): Measures the average squared difference between predicted and actual values.

  • L2 Relative Error: Provides a normalized error measure.

Key Features

  • Dynamic Input Handling: - The first layer of both the branch and trunk networks dynamically adjusts to the dimensions of the input data. This ensures flexibility across different PDEs and grid resolutions.

  • Support for Multiple Field Inputs: - The branch network can encode multiple input fields, such as source terms or coefficients.

  • Efficient Data Preprocessing: - Includes a Loader for local HDF5 file.

Example Training Workflow

  1. Load Data: - Use LocalLoader or a custom data loader to load simulation data in HDF5 format. - Preprocess data to extract branch inputs, trunk inputs, and solution values.

  2. Model Configuration: - Specify the hidden layers for both the branch and trunk networks. - Example: [128, 128, 128] for both networks.

  3. Training: - Train the model using the DeepONetTrainer class. - Specify hyperparameters such as batch size, learning rate, and number of epochs.

  4. Evaluation: - Evaluate the trained model using metrics like MSE and L2 relative error.

Dependencies

Ensure the following dependencies are installed:

  • DeepXDE: For implementing the DeepONet model.

  • H5py: For reading HDF5 files.

  • Scikit-learn: For data preprocessing and splitting.

  • NumPy: For efficient numerical operations.

Install dependencies using the following command:

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

Advanced Usage

Configurable Branch and Trunk Layers

The number of hidden layers and neurons per layer can be customized for specific applications by modifying the branch_hidden_layers and trunk_hidden_layers parameters in the DeepONetTrainer class.

Custom Data Loader

Users can implement their own data loaders by extending the DataLoader abstract base class. This allows seamless integration of new datasets and file formats.

Extending Metrics

To add custom evaluation metrics, modify the evaluate method in the DeepONetTrainer class.

Contributing

Contributions are welcome! If you’d like to add new features, improve documentation, or fix bugs, please submit a pull request.

Contact

For questions or feedback, please contact the project maintainers: