Learning Lagrangian Interaction Dynamics with Sampling-Based Model Order Reduction

Transactions on Machine Learning Research (TMLR)

     

Paper   Project Page   Weights   Data

Overview

GIOROM Simulating physical systems governed by Lagrangian dynamics often entails solving partial differential equations (PDEs) over high-resolution spatial domains, leading to significant computational expense. Reduced-order modeling (ROM) mitigates this cost by evolving low-dimensional latent representations of the underlying system. While neural ROMs enable querying solutions from latent states at arbitrary spatial points, their latent states typically represent the global domain and struggle to capture localized, highly dynamic behaviors such as fluids. We propose a sampling-based reduction framework that evolves Lagrangian systems directly in physical space, over the particles themselves, reducing the number of active degrees of freedom via data-driven neural PDE operators. To enable querying at arbitrary spatial locations, we introduce a learnable kernel parameterization that uses local spatial information from time-evolved sample particles to infer the underlying solution manifold. Empirically, our approach achieves a 6.6–32x reduction in input dimensionality while maintaining high-fidelity evaluations across diverse Lagrangian regimes, including fluid flows, granular media, and elastoplastic dynamics. We refer to this framework as GIOROM (Geometry-InfOrmed Reduced-Order Modeling).

Features

• 25× Speedup: Vastly outperforms existing neural network-based physics simulators while delivering high-fidelity predictions.
• Massive Upsampling: Infers dense point clouds of ~100,000 points from highly sparse sensor graphs of ~1,000 points with negligible computational overhead.
• Discretization-Agnostic: Geometry-aware architecture that seamlessly generalizes to new initial conditions, velocities, and unseen geometries post-training.
• 5 Complex Physical Systems: Empirically validated on elastic solids, Newtonian fluids, Non-Newtonian fluids, Drucker-Prager granular flows, and von Mises elastoplasticity.
• Comprehensive ROM Benchmark Suite: Includes clean, well-tuned, modular implementations of 6 state-of-the-art baseline Reduced Order Models (PCA, GNO, LiCROM, DINo, CORAL, and CoLoRA).
• Fully Automated Pipelines: Single-command execution for model training, metrics aggregation (Chamfer, Rel L2, VRAM), and professional Blender 3D video rendering.

Repository Structure

To ensure modularity and clean benchmarking, this repository is divided into two distinct phases. Our core algorithmic contributions lie in the Online Phase.

Note

A Note on Methodology: In a practical deployment, the sparse inputs used to train the online phase are generated by a high-fidelity offline model, which requires training for about 20 million steps (GNS based models). However, to enable fast benchmarking of the online model, our codebase allows for sampling the sparse observations directly from the ground-truth offline datasets to train and evaluate the online ROMs.

1. online/ (Function space manifold parameterization & ROM Benchmarks)

Contains the JAX-based GIOROM architecture alongside PyTorch implementations of all baseline ROMs. It includes a unified evaluation engine that guarantees fair, apples-to-apples comparisons across all models regarding inference time, memory footprint, and reconstruction accuracy. See the online_phase/README.md for full execution details.

2. offline/ (Neural time-stepper & neural operator Baselines)

Contains the code for processing raw simulation data (GNS / NCLAW) and training the offline full-order and time-stepper models. While we provide our specific offline training code for reproducibility, The online phase is strictly independent to the offline model. Any suitable high-fidelity neural operator based physics simulator can be used to generate the prior trajectories. See the offline_phase/README.md for details.

Quick Start

Installation

We recommend using Conda to manage dependencies. GIOROM utilizes both JAX (for the core online method) and PyTorch (for the offline model and baseline benchmark suite).

# Clone the repository
git clone [https://github.com/HrishikeshVish/GIOROM.git](https://github.com/HrishikeshVish/GIOROM.git)
cd GIOROM

# Create the environment
conda create --name giorom_env python=3.10
conda activate giorom_env

Because GIOROM utilizes both JAX (for the core method) and PyTorch (for the baseline suite and data loaders), the environment must be built carefully to prevent CUDA library conflicts.

We highly recommend following these exact steps to ensure both frameworks correctly bind to your GPU, and that torch-geometric compiles properly.

Step 1: Create the Environment

conda create --name giorom_env python=3.10 -y
conda activate giorom_env

Step 2: Install PyTorch (CUDA 12.1) We use PyTorch 2.4.0. Ensure your system CUDA driver supports 12.1.

pip install torch==2.4.0 torchvision torchaudio --index-url [https://download.pytorch.org/whl/cu121](https://download.pytorch.org/whl/cu121)

Step 3: Install Torch Geometric & Dependencies These wheels must exactly match the PyTorch version (2.4.0) and CUDA version (cu121).

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.4.0+cu121.html](https://data.pyg.org/whl/torch-2.4.0+cu121.html)

Step 4: Install JAX (CUDA 12) We use JAX 0.4.29 built for CUDA 12.

pip install --upgrade "jax[cuda12]==0.4.29"

Step 5: Install Remaining Generic Dependencies

pip install -r requirements.txt

Data & Weights

All necessary datasets (NCLAW/GNS formats) and pre-trained offline weights for baseline ROM models are available on our Google Drive.

For the automated scripts to work out-of-the-box, extract the downloaded data into a /data/ directory. The master execution script expects the following path variables (which you can modify at the top of run_experiments.sh):

H5_DATA_BASE="online/data/CROM_dataset/CROM_Ready_Data"

PT_DATA_BASE="online/data/pt_dataset"

OFFLINE_BASE="online/data/CROM_offline_training"

Running the Online Benchmarks Navigate to the online_phase directory to run the master experimental pipeline. This script trains the baselines, evaluates GIOROM, aggregates all metrics into CSVs, and optionally stitches 3D Blender renders into .mp4 comparison videos.

cd online
chmod +x run_experiments.sh

# Run the master sweep across all models and datasets
./run_experiments.sh

To run ablation studies on GIOROM's grid resolution and sparsity tolerances:

chmod +x run_ablations.sh
./run_ablations.sh

Acknowledgments

This codebase is inspired by or partly uses code from the following repositories:

• NCLaw for the dataset generation and material simulation framework.
• CROM for the Continuous Reduced-Order Modeling framework and offline baseline implementations.
• LiCROM for the LiCROM baseline implementation.
• Learning to Simulate (GNS) for the dataset structure and graph-based physics simulation utilities.
• NeuralOperator for the core neural operator architectures.
• GNOT for the Neural Operator Transformer baseline implementations.

Citation

If you find this codebase or benchmark suite useful in your research, please consider citing our paper:

@article{
viswanath2026learning,
title={Learning Lagrangian Interaction Dynamics with Sampling-Based Model Order Reduction},
author={Hrishikesh Viswanath and Yue Chang and Aleksey Panas and Julius Berner and Peter Yichen Chen and Aniket Bera},
journal={Transactions on Machine Learning Research},
issn={2835-8856},
year={2026},
url={https://openreview.net/forum?id=vXCQA1EzaG},
note={Expert Certification}
}