This repository provides the source code and training scripts for "Geometry Matters: Benchmarking Scientific ML Approaches for Flow Prediction around Complex Geometries." The project evaluates state-of-the-art neural operators and foundation models for predicting fluid dynamics over complex geometries.
Our study introduces a benchmark for scientific machine learning (SciML) models in predicting steady-state flow across intricate geometries using high-fidelity simulation data. The full paper can be accessed here:
"Geometry Matters: Benchmarking Scientific ML Approaches for Flow Prediction around Complex Geometries"
- Authors: Ali Rabeh, Ethan Herron, Aditya Balu, Soumik Sarkar, Chinmay Hegde, Adarsh Krishnamurthy, Baskar Ganapathysubramanian
- Preprint: ArXiv Paper
- Benchmarking of 11 SciML models, including CNO, FNO, DeepONet, WNO, and Poseidon-based architectures.
- Evaluation on FlowBench 2D Lid-Driven Cavity dataset, a publicly available dataset on Hugging Face.
- Comparison of two geometric representations: Signed Distance Fields (SDF) and Binary Masks.
- Hyperparameter tuning with WandB Sweeps.
- Residual and gradient calculations using FEM-based scripts.
This study utilizes the FlowBench 2D Lid-Driven Cavity (LDC) dataset, which is publicly accessible on Hugging Face: FlowBench LDC Dataset
The dataset is licensed under CC-BY-NC-4.0 and serves as a benchmark for the development and evaluation of scientific machine learning (SciML) models.
- Geometry representation: SDF and Binary Mask
- Resolution: 512×512
- Fields: Velocity (u, v) and Pressure (p)
- Stored as: Numpy tensors (
.npz
format)
- Neural Operators:
FNO, CNO, WNO, DeepONet, Geometric-DeepONet
- Vision Transformers:
scOT-T, Poseidon-T, scOT-B, Poseidon-B, scOT-L, Poseidon-L
To set up the environment and install dependencies for Neural Operators:
python3 -m venv sciml
source sciml/bin/activate
pip install --upgrade pip setuptools wheel Cython
pip install -r venv_requirements.txt
To set up the environment and install dependencies for Vision Transformers:
python3 -m venv scot
source scot/bin/activate
pip install --upgrade pip setuptools wheel Cython
pip install -r scot_requirements.txt
To train a Neural Operator, run the following command:
python3 main_sweep.py --model "model_name" --sweep
To train a Vision Transformer, run the following command:
python3 scot_sweep.py --model "model_name" --sweep
Before training, you need to specify the dataset paths in the configurations (YAML files):
data:
file_path_train_x: ./data/train_x.npz
file_path_train_y: ./data/train_y.npz
file_path_test_x: ./data/test_x.npz
file_path_test_y: ./data/test_y.npz
For model inference, use the scripts in the plotting_scripts
folder:
python3 process_NO_deriv.py --model "$model" --config "$config_path" --checkpoint "$checkpoint_file"
Use process_NO.py
for neural operators and process_scot.py
for vision transformers.
The plotting_scripts
folder contains Python scripts for:
- Evaluating field predictions and errors.
- Calculating residuals using finite element methods (FEM).
- Computing solution gradients.
Example usage:
python plotting_scripts/plot_predictions.py --data_path ./data/test.npz --model_name fno
If you use this code, please cite our work:
@misc{rabeh2024flowprediction,
title={Geometry Matters: Benchmarking Scientific ML Approaches for Flow Prediction around Complex Geometries},
author={Ali Rabeh, Ethan Herron, Aditya Balu, Soumik Sarkar, Chinmay Hegde, Adarsh Krishnamurthy, Baskar Ganapathysubramanian},
year={2024},
eprint={2405.19101},
archivePrefix={arXiv},
primaryClass={cs.LG}
}
We welcome contributions! If you’d like to improve this project, please fork the repository and submit a pull request.
This repository is licensed under the MIT License.