This repository contains the official implementation of our Latent Mamba Operator for Partial Differential Equations.
Latent Mamba Operator for Partial Differential Equations
Karn Tiwari, Niladri Dutta, N M Anoop Krishnan, Prathosh A P
Figure 1. Overview of LaMO.
- Install Python 3.12.3. For convenience, please go ahead and execute the following command.
pip install -r requirements.txt- Prepare Data. You can obtain experimental datasets from the following links (Download).
| Dataset | Task | Geometry | Link |
|---|---|---|---|
| Elasticity | Estimate material inner stress | Point Cloud | [Google Cloud] |
| Plasticity | Estimate material deformation over time | Structured Mesh | [Google Cloud] |
| Navier-Stokes | Predict future fluid velocity | Regular Grid | [Google Cloud] |
| Darcy | Estimate fluid pressure through medium | Regular Grid | [Google Cloud] |
| AirFoil | Estimate airflow velocity around airfoil | Structured Mesh | [Google Cloud] |
| Pipe | Estimate fluid velocity in a pipe | Structured Mesh | [Google Cloud] |
- Train and evaluate the model. We provide the experiment scripts of all benchmarks under the folder
./scripts/.You can reproduce the experiment results as the following examples:
bash scripts/LaMO_Elas.sh # for Elasticity
bash scripts/LaMO_Plas.sh # for Plasticity
bash scripts/LaMO_NS.sh # for Navier-Stokes
bash scripts/LaMO_Darcy.sh # for Darcy
bash scripts/LaMO_Airfoil.sh # for Airfoil
bash scripts/LaMO_Pipe.sh # for PipeNote: You must change the argument --data-path in the above script files to your dataset path.
-
Develop your own model. Here are the instructions:
- Add the model file under the folder
./models/. - Add the model name into
./model_dict.py. - Add a script file under folder
./scripts/and change the argument--model.
- Add the model file under the folder
Figure 2. Main Results Compared with SOTA Operator. L2 Error is reported.
LaMO achieves 32.3% averaged relative increment over the previous second-best model i.e, Transolver.
Figure 3. Showcases. Error heatmap is plotted for Transolver and LaMO for Darcy, Plasticity, and Navier-Stokes.
Please consider citing our paper if you find it helpful. Thank you!
@misc{tiwari2025latentmambaoperatorpartial,
title={Latent Mamba Operator for Partial Differential Equations},
author={Karn Tiwari and Niladri Dutta and N M Anoop Krishnan and Prathosh A P},
year={2025},
eprint={2505.19105},
archivePrefix={arXiv},
primaryClass={cs.LG},
url={https://arxiv.org/abs/2505.19105},
}
If you have any questions, please feel free to contact me at:
Karn Tiwari: karntiwari@iisc.ac.in
We appreciate the following GitHub repos a lot for their valuable code base or datasets on which we have built our code: