UPT++: Latent Point Set Neural Operators for Modeling System State Transitions

Description

UPT++ implements latent point set neural operators for modeling dynamic system state transitions. The approach encodes spatial particle clouds into a fixed-size latent set, processes them with a Perceiver-style transformer, and decodes to per-particle occupancy labels and density fields.

Supported domains:

• Molecular dynamics — alanine dipeptide trajectories (.dcd/.pdb via MDTraj)
• Fluid simulations — dam breaking and water drop (particle-based, multi-material)

Installation

This project uses uv for dependency management — no conda, virtualenv, or manual environment setup needed. uv creates and manages an isolated Python environment automatically.

Install uv (if not already installed):

curl -LsSf https://astral.sh/uv/install.sh | sh

Then install all dependencies:

uv sync

Dataset Setup

Place or symlink your data directory at data/ in the repo root (configured via configs/paths/default.yaml).

How to Run

Train with the default configuration:

uv run src/train.py

Train on a specific experiment (select GPU with trainer.devices):

# Water drop
uv run src/train.py experiment=waterdrop trainer.devices=[0]

# Water drop (XL)
uv run src/train.py experiment=waterdrop_xl trainer.devices=[0]

# Dam breaking 3D
uv run src/train.py experiment=dam3d trainer.devices=[0]

# Molecular dynamics
uv run src/train.py experiment=md_occupancy_p trainer.devices=[0]

Override any Hydra parameter from the command line:

uv run src/train.py experiment=waterdrop trainer.devices=[0] trainer.max_epochs=100 data.batch_size=32

Evaluate a trained checkpoint:

uv run src/eval.py ckpt_path=<path/to/checkpoint.ckpt>

Rollout evaluation (standalone scripts):

uv run scripts/evaluate_rollout.py
uv run scripts/evaluate_rollout_gns_f.py

Experiments

All experiment configs are in configs/experiment/. Key presets:

Experiment	Domain	Model
waterdrop	Water drop (fluid)	OccupancyAutoencoder
waterdrop_xl	Water drop large	OccupancyAutoencoder
waterdrop_gns_f	Water drop (fluid)	GNS_F
dam3d	Dam breaking 3D	OccupancyAutoencoder
dam3d_gns_f	Dam breaking 3D	GNS_F
md_occupancy_p	Molecular dynamics	OccupancyAutoencoder

Configuration

Configuration uses Hydra. The top-level config is configs/train.yaml. Setting experiment=<name> composes the model, data, trainer, and hyperparameter configs for that task.

Override data and output paths:

uv run src/train.py paths.data_dir=/path/to/data paths.log_dir=/path/to/logs