SciML - Scientific Machine Learning

📝 Description

This project implements modern neural network architectures for operator modeling and solving partial differential equations (PDEs). It is developed in collaboration with Professors Hadrien Montanelli and Samuel Kokh.

🎯 Objectives

• Implement and compare two deep learning architectures for operator modeling:

  • DeepONet (Deep Operator Network) - To learn operators between function spaces
  • FNO (Fourier Neural Operator) - To leverage spectral analysis in operator modeling

• Apply these architectures to solving partial differential equations (PDEs), with a particular focus on the heat equation

🧠 Key Concepts

• Operator Modeling: Learning to map functional input spaces to functional output spaces
• Spectral Analysis: Using Fourier transforms to efficiently capture spatial dynamics
• Supervised Learning: Training on input-output pairs generated by classical numerical solvers

📊 Data Format

Training data is structured in triplets (mu, x, sol) where:

• mu represents input functions (e.g., boundary conditions, initial conditions)
• x represents points in the spatial domain
• sol represents expected solutions at points x

�� Project Structure

sciml/
├── data/              # Training data and generation scripts
├── model/             # Architecture implementations
│   ├── deeponet/      # DeepONet implementation
│   └── fno/           # Fourier Neural Operator implementation
├── notebooks/         # Jupyter notebooks for experiments and visualizations
│   ├── deeponet/      # DeepONet notebooks
│   └── fno/          # FNO notebooks
├── utils/             # Utility functions
├── tests/             # Unit and integration tests
└── logs/              # Training logs and results

🚀 Installation and Setup

Prerequisites

• Python 3.9+
• TensorFlow 2.8+
• Virtual environment (recommended)

Installation Steps

1. Clone the repository

   git clone https://github.com/janisaiad/sciml.git
   cd sciml

2. Set up the environment and install dependencies

   chmod +x launch.sh
   ./launch.sh

3. Activate the virtual environment

   source .venv/bin/activate  # For Linux/MacOS
   # or
   .\.venv\Scripts\activate   # For Windows

📚 Usage

Data Generation

Data generation scripts are located in the data/generation/ directory. Example usage:

python data/generation/generate_big_heat_fno.py

Training Models

from sciml.model.fno import FNO
from sciml.model.deeponet import DeepONet

# Configure and train FNO
fno_model = FNO(hyper_params, regular_params, fourier_params)
fno_model.fit()

# Configure and train DeepONet
deeponet_model = DeepONet(hyper_params, regular_params)
deeponet_model.fit()

Example Notebooks

Several notebooks are available to demonstrate model usage:

• For DeepONet: notebooks/deeponet/TORUNDEEPONETCOMPARISON.ipynb
• For FNO: notebooks/fno/TORUNFNO.ipynb

Converting Notebooks with Jupytext

To convert all notebooks to Python files:

find . -name "*.ipynb" -type f -exec jupytext --to py {} \;

🔍 Troubleshooting

If you encounter issues with JAX/JAXlib during FNO model training, try the following commands:

# Activate virtual environment if not already done
source .venv/bin/activate

# Uninstall current versions
uv pip uninstall jax jaxlib

# Install compatible versions
uv add jax
uv add jaxlib==0.4.17

📝 Research Notes and Observations

FNO Performance

• Too many Fourier layers destabilize training
• Consider sequential temporal inference
• Test on other data (naca)

DeepONet Improvements

• Compare with RB and POD
• Test alternative branch/trunk architectures

Spectral Analysis

• Impact of Fourier layer count on generalization
• Compare Fourier weights between layers
• Verify cast/coeff and phase multiplication

📄 License

This project is distributed under the MIT license. See the LICENSE file for more details.