HamiFormer-Assets

HamiFormer models

Inference weights for HamiFormer and baseline models on HamiBalls-1 and HamiBalls-2. HamiFormer combines whole-window diffusion prediction with Hamiltonian propagation through affine symplectic maps and regime-conditioned corrections.

Code and instructions · HamiBalls datasets

Available weights

Files contain tensor dictionaries loaded by the corresponding implementations in the code repository.

Download and evaluate

Use Python 3.10–3.12, PyTorch 2.5.1, and a compatible CUDA environment.

git clone https://github.com/HamiFormer/HamiFormer.git
cd HamiFormer
python -m pip install -e ".[test,data,analysis]"
python -m pip install huggingface_hub

Download the weights and datasets from Python:

from huggingface_hub import snapshot_download

snapshot_download(
    repo_id="HamiFormer/HamiFormer-Assets",
    allow_patterns="weights/*.pt",
    local_dir="assets",
)
snapshot_download(
    repo_id="HamiFormer/Hamiballs",
    repo_type="dataset",
    allow_patterns="*.h5",
    local_dir="assets/data",
)

Run a small evaluation on each dataset:

python run.py evaluate --dataset h1 --method ours --data assets/data/hamiballs1.h5 --weights assets/weights --limit 2 --output results/h1_ours
python run.py evaluate --dataset h2 --method ours --data assets/data/hamiballs2.h5 --weights assets/weights --limit 2 --output results/h2_ours

Omit --limit to evaluate all 512 validation trajectories. The default protocol uses two noise seeds and 192 predicted transitions. HamiFormer predicts four consecutive 48-transition windows. The evaluator reports normalized state, position, and momentum MSE, contact-conditioned metrics, and disjoint-interval metrics, and writes metrics.json and trajectories.npz to the requested output directory. Use a new output directory for each run.

Inputs and outputs

Inputs are an initial phase-space state, object attributes, and physical times. HamiBalls-2 additionally supplies object and spring-graph information. Each object's state is ordered as [q, p], where p is momentum. The evaluation entry point handles normalization using the model's stored scales and returns trajectories in the dataset's physical coordinates.

For direct component loading:

from hamiformer.models.hamiformer import HamiFormer

model = HamiFormer.from_pretrained("assets/weights", dataset="h2", device="cuda")
components = model.components()

Intended use

These models support research on learned physical simulation and trajectory prediction in the HamiBalls spring-and-contact environments. The matching dataset conventions, physical time step, and object representation are described in the dataset card.

The project code is distributed under the MIT license. See third-party notices for baseline references and attribution.