README.md

---
language:
  - en
library_name: pytorch
tags:
  - physics
  - physical-simulation
  - trajectory-prediction
  - hamiltonian
  - diffusion
datasets:
  - HamiFormer/Hamiballs
---

# 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](https://github.com/HamiFormer/HamiFormer) · [HamiBalls datasets](https://huggingface.co/datasets/HamiFormer/Hamiballs)

## Available weights

| File under `weights/` | Dataset | Method | Evaluation argument |
|---|---|---|---|
| `hami1_ours.pt` | HamiBalls-1 | HamiFormer | `--dataset h1 --method ours` |
| `hami1_physiformer.pt` | HamiBalls-1 | PhysiFormer | `--dataset h1 --method physiformer` |
| `hami1_hgdpf.pt` | HamiBalls-1 | Hamiltonian-Guided Diffusion Fields | `--dataset h1 --method hgdpf` |
| `hami2_ours.pt` | HamiBalls-2 | HamiFormer | `--dataset h2 --method ours` |
| `hami2_physiformer.pt` | HamiBalls-2 | PhysiFormer | `--dataset h2 --method physiformer` |
| `hami2_dit.pt` | HamiBalls-2 | Diffusion Transformer | `--dataset h2 --method dit` |
| `hami2_transformer_ar1.pt` | HamiBalls-2 | Transformer-AR, context 1 | `--dataset h2 --method transformer_ar1` |
| `hami2_transformer_ar4.pt` | HamiBalls-2 | Transformer-AR, context 4 | `--dataset h2 --method transformer_ar4` |

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.

```sh
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:

```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:

```sh
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:

```python
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](https://huggingface.co/datasets/HamiFormer/Hamiballs).

The project code is distributed under the [MIT license](https://github.com/HamiFormer/HamiFormer/blob/main/LICENSE). See [third-party notices](https://github.com/HamiFormer/HamiFormer/blob/main/THIRD_PARTY_NOTICES.md) for baseline references and attribution.