ShellFlow paper models: union + E1 ablations (arXiv:2607.16144)

README.md

@@ -0,0 +1,69 @@
+---
+license: mit
+tags:
+- particle-physics
+- lhc
+- atlas-open-data
+- flow-matching
+- riemannian-flow-matching
+- event-generation
+- arxiv:2607.16144
+---
+
+# ShellFlow — pretrained models
+
+Pretrained checkpoints for the paper **[Learning Standard Model structure from LHC data with Riemannian flow matching](https://arxiv.org/abs/2607.16144)** (arXiv:2607.16144).
+
+ShellFlow is a transformer generative model of LHC collision events, trained with Riemannian flow matching on a product manifold (R × S², plus a log-mass coordinate for massive objects) using ATLAS Open Data. Code: https://github.com/pomidori/ShellFlow
+
+## Models
+
+All models share the same backbone (d_model 128, 6 layers, 4 heads, ~6.0M parameters) and were trained on the union of the ATLAS Open Data `1LMET30` and `2to4lep` skims (events with ≤ 8 reconstructed particles, de-duplicated by run/event number).
+
+| File | Model | Chart | Aux head | Training |
+|---|---|---|---|---|
+| `union_1LMET30_2to4lep/model.ckpt` | Main paper model | on-shell (manifold) | yes | 30 epochs, 717k steps |
+| `e1_v0_manifold_aux/model.ckpt` | Ablation E1-V0 | on-shell (manifold) | yes | 2 epochs, 48k steps |
+| `e1_v1_euclidean_aux/model.ckpt` | Ablation E1-V1 | free (Euclidean) | yes | 2 epochs, 48k steps |
+| `e1_v2_manifold_noaux/model.ckpt` | Ablation E1-V2 | on-shell (manifold) | no | 2 epochs, 48k steps |
+| `e1_v3_euclidean_noaux/model.ckpt` | Ablation E1-V3 | free (Euclidean) | no | 2 epochs, 48k steps |
+
+Checkpoint format: PyTorch Lightning (weights + EMA weights + hyperparameters; optimizer state stripped). The hyperparameters carry `architecture: legacy_type_film`, which selects the paper architecture in the ShellFlow code. Distribution matching (a training-time reweighting) is disabled in the stored hyperparameters so the checkpoints load without any dataset present; it has no effect on generation.
+
+## Usage
+
+Install the [ShellFlow code](https://github.com/pomidori/ShellFlow), then:
+
+```python
+from huggingface_hub import hf_hub_download
+from shellflow.training.riemannian_DiTs_lightning import RiemannianDiTsModule
+
+path = hf_hub_download("pomidori73/ShellFlow", "union_1LMET30_2to4lep/model.ckpt")
+module = RiemannianDiTsModule.load_from_checkpoint(path, map_location="cpu").eval()
+model = module.ema_model  # paper results use the EMA weights
+```
+
+Generation is conditioned on the event composition (particle types, charges, detector extras, MET). The simplest generation harness mirrors the training-time monitoring callback:
+
+```python
+import torch
+from shellflow.training.riemannian_DiTs_callback import RiemannianDiTsSamplePlotCallback
+
+cb = RiemannianDiTsSamplePlotCallback(val_dataset=None, collator=None, n_gen_steps=200)
+with torch.no_grad():
+    kinematics = cb._generate_batch_rfm(model, batch, module)  # (B, N, 4) = (px, py, pz, E) in GeV
+```
+
+where `batch` is a dict with `particle_type` (B, N) int64 codes {0=pad, 1=jet, 10=largeRjet, 11=e, 13=mu, 15=tau, 22=photon}, `particle_charge` (B, N) int64 in {0, 1, 2}, `particle_extra` (B, N, 13) float, `particle_mask` (B, N) 0/1, and optionally `met_vector` (B, 2). For physical samples, draw the conditioning from real events (see the dataset pipeline in the code repository).
+
+## Citation
+
+```bibtex
+@article{shellflow2026,
+  title   = {Learning Standard Model structure from LHC data with Riemannian flow matching},
+  author  = {Kato, Midori and Urqu{\'i}a-Calder{\'o}n, Kevin A. and Timiryasov, Inar and Ruchayskiy, Oleg},
+  journal = {arXiv preprint arXiv:2607.16144},
+  year    = {2026},
+  url     = {https://arxiv.org/abs/2607.16144}
+}
+```

e1_v0_manifold_aux/model.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:658ab7a5e7734cd52778174da68cc32ef7e59d000e796ff435b6f5adadbf6610
+size 24240179

e1_v1_euclidean_aux/model.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1ef8d8436c3fbc551862d721a30f098ae23a89aae2050305121a862591d40289
+size 24240179

e1_v2_manifold_noaux/model.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1d8490472ce909978632fb1ae1a9fbd09863f2ac68207fab068831abce999700
+size 24216483

e1_v3_euclidean_noaux/model.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ac0585246cc14cd8e078298736a14c22eee21186376b9d27ce3ac15737aeddd5
+size 24216547

union_1LMET30_2to4lep/model.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fa7a748f39c51cae12a9b529d916247d125062568719b7878231dbcbd0c245ce
+size 24240179