MNO : A Multi-modal Neural Operator for Parametric Nonlinear BVPs

This repository contains the official implementation for "MNO : A Multi-modal Neural Operator for Parametric Nonlinear BVPs".

🎯 Overview

This work introduces novel Multi-modal Neural Operator (MNO) architecture designed to learn solution operators for multi-parameter non-linear boundary value problems (BVPs). In this repo, we provide code to reproduce the results from following experiments discussed in the paper:

• Poisson 1D/2D: Elliptic linear PDEs
• Darcy 1D: Flow in porous media with heterogeneous variable coefficients
• Nonlinear BVP: Nonlinear boundary value problem with parameterized boundary conditions and multiple variable coefficients

📁 Repository Structure

MultimodalNeuralOperator/
├── README.md                    # This file
├── requirements.txt             # Python dependencies
├── config.py                    # Configuration system for all experiments
├── train.py                     # Main training script
├── evaluate.py                  # Evaluation script with error metrics
├── generate_figs.py             # Generate plots in the paper
├── download_model_weights.py    # Download pre-trained model weights from all experiments
├── train_custom_example.py      # Examples of custom training configurations
│
├── models/                      # Neural network implementations
│   ├── __init__.py             # Model exports for easy importing
│   ├── uno.py                  # Uni-modal Neural Operator (UNO) implementation with single channels
│   ├── mno.py                  # Multi-modal Neural Operator (MNO) implementation  
│   ├── fno.py                  # Fourier Neural Operator (FNO) implementation (Li, Zongyi, et al 2020)
│   └── deeponet.py             # DeepONet implementation (Lu, Lu, Pengzhan Jin, and George Em Karniadakis 2019)
│
├── data_generator.py            # Data generators for all experiments
├── data_utils.py               # Utility functions for data processing
├── utils.py                    # General utility functions
│
├── ckpts/                      # Model checkpoints (created after downloading model weights)
└── plots/                      # Generated figures (created after running generate_figs.py)

🚀 Quick Start

1. Installation

# Clone/Download the repository
cd MultimodalNeuralOperator

# Create virtual environment (recommended)
python -m venv venv
source venv/bin/activate  # On Windows: venv\Scripts\activate

# Install dependencies
pip install -r requirements.txt

2. Download Pre-trained Models

# Download all pre-trained model weights
python download_model_weights.py

# This will populate the ckpts/ directory with trained models for all experiments

3. Quick Evaluation

# Evaluate all models on all experiments
python evaluate.py

# This will output relative errors and residual errors for each model

4. Generate Paper Figures

# Generate all experiment plots from the paper
python generate_figs.py

# Figures will be saved in the plots/ directory

🔧 Usage

Training Models

Using Standard Configurations

from config import StandardConfigs
from train import train, get_loss_fn

# Load a standard configuration
config = StandardConfigs.poisson1d_uno()

# Create network and data generator
net = config.create_network()
data_generator = config.create_data_generator()
loss_fn = get_loss_fn(config)

# Train the model
train(net, epochs=1000, optim='adam', lr=config.lr, 
      adam_beta1=config.adam_beta1, adam_beta2=config.adam_beta2,
      loss_fn=loss_fn, data_generator=data_generator,
      check_interval=1000, threshold=-1, patience=2,
      folder=f"{config.folder}/{config.experiment}",
      exp_name=config.exp_name, save_every=1000)

Command Line Training

# Train UNO on Poisson 1D
python train.py --config poisson1d_uno

# Train FNO on Poisson 1D  
python train.py --config poisson1d_fno

# Train MNO on Darcy 1D
python train.py --config darcy1d_mno

# Other available configs:
# poisson1d_deeponet, poisson2d_uno, darcy1d_uno_aq, darcy1d_uno_ln, 
# darcy1d_uno_mix, nlbvp_uno, nlbvp_mno

Custom Training Configurations

from config import create_custom_config
from models import UNO_1D
from data_generator import Poisson1DDataGenerator

# Create custom configuration
config = create_custom_config(
    experiment='poisson1d',
    network='uno',
    network_class=UNO_1D,
    network_kwargs={
        'num_blocks': 3,  # More blocks than standard
        'diag_type': 'tri',
        'M': 32,  # Different decomposition
        'P': 8,
        'periodic': False,
        'bias': False,
        'activation': None,
        'init_type': 'uniform'
    },
    data_generator_class=Poisson1DDataGenerator,
    grid_size=512,  # Higher resolution
    lr=5e-4,        # Different learning rate
    exp_name='uno_custom'
)

# Train with custom config
# (same training code as above)

See train_custom_example.py for more examples.

Evaluation

Evaluate Specific Experiments

from evaluate import evaluate_poisson1d_models, evaluate_darcy1d_models

# Evaluate Poisson 1D models
results = evaluate_poisson1d_models(num_samples=1000, lin_comb=True)

# Evaluate Darcy 1D models  
results = evaluate_darcy1d_models(num_samples=1000, log_normal=True)

Evaluate Custom Models

from evaluate import evaluate_custom_config

# Evaluate your custom trained model
result = evaluate_custom_config(config, num_samples=500)

Configuration System

The configuration system provides a interface for all experiments:

from config import StandardConfigs

# All experiment configurations
configs = [
    StandardConfigs.poisson1d_uno(),
    StandardConfigs.poisson1d_fno(), 
    StandardConfigs.poisson1d_deeponet(),
    StandardConfigs.poisson2d_uno(),
    StandardConfigs.darcy1d_uno_aq(),
    StandardConfigs.darcy1d_uno_ln(),
    StandardConfigs.darcy1d_uno_mix(),
    StandardConfigs.darcy1d_mno(),
    StandardConfigs.nlbvp_uno(),
    StandardConfigs.nlbvp_mno(),
]

# Each config contains:
# - Network class and parameters
# - Data generator class  
# - Training hyperparameters
# - Experiment metadata

🔬 Reproducing Paper Results

1. Download Pre-trained Models

python download_model_weights.py

2. Run Evaluation

# Evaluate all models and print results
python evaluate.py

3. Generate Figures

# Generate paper figures with different configurations
python generate_figs.py

4. Train from Scratch (Optional)

# Train all models from scratch
python train.py --config poisson1d_uno
python train.py --config poisson1d_fno
python train.py --config poisson1d_deeponet
python train.py --config poisson2d_uno
python train.py --config darcy1d_uno_aq
python train.py --config darcy1d_uno_ln
python train.py --config darcy1d_uno_mix
python train.py --config darcy1d_mno
python train.py --config nlbvp_uno
python train.py --config nlbvp_mno

🛠️ Advanced Usage

Custom Data Generation

from data_generator import Poisson1DDataGenerator

# Create custom data generator
data_gen = Poisson1DDataGenerator(grid_size=512, batch_size=500)

# Generate batch with specific parameters
params, solution = data_gen.generate_batch(mix=True, rhs_sampling=False)

Model Components

The framework provides modular components that can be used independently:

from models import FMM_1D, FMM_2D, Edge, BlockXDiag

# Use individual components
fmm_1d = FMM_1D(diag_type='tri', M=16, P=16, periodic=False)
edge = Edge(P=16, bias=False, activation='relu')
block_diag = BlockXDiag('tri', M=8, P=16, periodic=True)

JIT Optimization

The 2D models use conditional JIT compilation for optimal performance:

• Training: Uses JIT for parallel execution
• Evaluation: Uses sequential execution for consistent timing

This is automatically handled by the model.train() and model.eval() modes.

🐛 Troubleshooting

Common Issues

1. CUDA Out of Memory: Reduce batch_size in configurations
2. Import Errors: Ensure all dependencies are installed with pip install -r requirements.txt
3. Missing Models: Run python download_model_weights.py to download pre-trained weights

Performance Tips

1. Use GPU when available (automatically detected)
2. For large-scale experiments, consider reducing grid_size during development
3. Use torch.compile for additional speedup (PyTorch 2.0+)

📄 License

This project is licensed under the MIT License - see the LICENSE file for details.