This repo contains the official PyTorch code for GridMix (ICLR 2025, Oral).
GridMix is a novel approach that improves spatial modulation in neural fields by leveraging a mixture of grid-based representations. Existing methods rely on global modulation, which limits local detail reconstruction, while vanilla grid-based approaches struggle with overfitting and inadequate global information modeling. GridMix addresses these challenges by balancing global structure exploration with fine-grained locality preservation.
To further enhance robustness, we introduce spatial domain augmentation, improving generalization across varying spatial domains. GridMix serves as a key component of MARBLE, a comprehensive framework advancing PDE modeling with neural fields. Our method is extensively validated on benchmarks in dynamics modeling and geometric prediction, demonstrating significant improvements.
Clone the repo and install the dependencies.
git clone https://github.com/LeapLabTHU/GridMix.git
cd GridMix
conda create -n marble python=3.9.0
pip install -e .
pip install torch_geometric
pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.0.0+cu117.html
Setup wandb config in ~/.bashrc
export WANDB_API_TOKEN=your_key
export WANDB_DIR=your_dir
export WANDB_CACHE_DIR=your_cache_dir
export WANDB_CONFIG_DIR="${WANDB_DIR}config/wandb"
export MINICONDA_PATH=your_anaconda_path
- Dynamics Modeling: NS and SW data can be generated using DINo repo. We provide the generated data in google drive.
- Geometric-aware inference: NACA-Euler, Elasticity and pipe data can be downloaded from google drive, which is provided by Geo-FNO repo.
For dynamics modeling, we provide the training scripts in repro_dynamics. Taking the navier-stokes data as an example, the two-step training process is as follows:
# train the inr model
bash ./repro_dynamics/navier-stokes/marble_ns_s1_bash.sh
# train the dynamics model
bash ./repro_dynamics/navier-stokes/marble_ns_s1_ode_bash.sh
For geometric-aware inference, we provide the training scripts in repro_static. Taking the elasticity data as an example, the training process is as follows:
# train the inr model with global modulation for input functions
bash ./repro_static/elasticity/coral_inr_bash.sh
# train the inr model with gridmix modulation for output functions
bash ./repro_static/elasticity/marble_inr_bash.sh
# train the regression model
bash ./repro_static/elasticity/marble_regression_bash.sh
This code is developed on the top of CORAL, we thank to their efficient and neat codebase.
If you find our work is useful in your research, please consider citing:
@inproceedings{wanggridmix,
title={GridMix: Exploring Spatial Modulation for Neural Fields in PDE Modeling},
author={Wang, Honghui and Song, Shiji and Huang, Gao},
booktitle={The Thirteenth International Conference on Learning Representations},
year={2025},
}
If you have any questions, please feel free to contact us at [email protected].