Modular PINN & Neural Operator Framework

A robust, configurable research framework for Physics-Informed Neural Networks (PINNs) and Classical FEM Simulations, built on SciML, Gridap.jl, and DrWatson.jl.

This project allows you to:

1. Simulate Coupled PDEs using classical FEM (Gridap).
2. Train Neural Networks (MLP, SIREN, DeepONet) to solve the same PDEs.
3. Analyize and compare results using automated pipelines.

---

🚀 Quick Start

1. Setup

using Pkg
Pkg.activate(".")
Pkg.instantiate()

2. Run a Classical FEM Simulation

Generate ground truth data using the Finite Element Method.

julia --project=. scripts/solve_fem.jl

Output: data/sims/simulation_....jld2

3. Train a Neural Network (PINN)

Train a physics-informed model using the generated data (Hybrid Training).

# 1. Create a config (or use the demo)
cp configs/templates/default.toml configs/experiments/my_run.toml

# 2. Run Training
julia --project=. scripts/train_pinn.jl --config configs/experiments/demo_mlp.toml

Output: data/models/demo/...

4. Analyze Results

Generate Kymographs (heatmaps) and GIF animations.

julia --project=. scripts/analyze_sim.jl --input data/sims/YOUR_FILE.jld2

Output: data/analysis/...

5. Evaluate Model (Compare to Ground Truth)

Calculate error metrics (MSE, L2) and generate Difference Heatmaps ($|u_{pred} - u_{true}|$).

julia --project=. scripts/evaluate_model.jl --model data/models/my_run/checkpoint.jld2

Output: data/analysis/checkpoint/... (Error Maps)

---

📂 Project Structure

fem/
├── configs/              # ⚙️ Experiment Configuration (TOML)
│   ├── templates/        #    Default templates
│   └── experiments/      #    User experiments
├── src/                  # 🧠 Core Modules
│   ├── Physics.jl        #    Symbolic PDE Definitions & Data Loading
│   ├── Architectures.jl  #    Neural Net Factory (MLP, SIREN, DeepONet)
│   ├── Training.jl       #    Training Loop (Optimization, Loss)
│   ├── Solver.jl         #    Classical FEM Logic (Gridap)
│   └── Analysis.jl       #    Visualization Tools (Plots)
├── scripts/              # 🏃 Execution Scripts
│   ├── solve_fem.jl      #    Run FEM Simulation
│   ├── train_pinn.jl     #    Train Neural Network
│   └── analyze_sim.jl    #    Generate Plots/Movies
└── data/                 # 💾 Artifacts (gitignored)
    ├── sims/             #    FEM outputs (.jld2)
    ├── models/           #    Trained Models
    └── analysis/         #    Plots & GIFs

---

🧠 Core Modules

1. Physics (src/Physics.jl)

Defines the Coupled Liquid-Fiber system using ModelingToolkit.jl.

• Symbolic API: get_coupled_system(; params...)
• Data Assimilation: get_system_from_file(path) loads parameters using JLD2 and constructs continuous Interpolations of the solution surface.

2. Architectures (src/Architectures.jl)

A factory for creating Lux.jl chains.

• Supported Types: "MLP", "SIREN", "DeepONet".
• Hot-Swapping: Change type = "SIREN" in your TOML config to instantly switch architectures.

3. Training (src/Training.jl)

Orchestrates the optimization using NeuralPDE.jl.

• Strategies: QuadratureTraining, GridTraining, StochasticTraining.
• Hybrid Loss: Supports mixing Physics Loss (Residuals) with Data Loss (Interpolation Error) via the data_loss_weight parameter.

4. Solver (src/Solver.jl)

Encapsulates the classical Gridap.jl Finite Element solver.

• Method: Crank-Nicolson time stepping with Newton-Raphson nonlinear solver.
• Elements: Linear Lagrangian elements.

---

⚙️ Configuration (TOML)

All experiments are defined in configs/. Example:

[experiment]
name = "siren_test"
output_dir = "data/models/siren"

[physics]
source_type = "file" # Use data from simulation
data_path = "data/sims/simulation_....jld2"

[architecture]
type = "SIREN"
hidden_dim = 32
layers = 4

[training]
strategy = "Quadrature"
max_iters = 1000
data_loss_weight = 1.0 # Enable Hybrid Training

---

📊 Analysis

The analysis pipeline (scripts/analyze_sim.jl) automatically generates:

1. Kymographs: Space-Time heatmaps for $\rho_l$ and $\rho_f$.
2. Animations: evolution.gif showing the wave propagation over time.

---

🛠 Dependencies

• Modeling: ModelingToolkit.jl, NeuralPDE.jl, Gridap.jl
• ML: Lux.jl, Optimization.jl
• Data/IO: DrWatson.jl, JLD2, TOML, Interpolations.jl
• Vis: Plots.jl