Multi-geometry pretrained count models (paper release)

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,83 @@
+---
+license: apache-2.0
+tags:
+- physics
+- calorimeter
+- fast-simulation
+- generative-model
+- flow-matching
+library_name: pytorch
+datasets:
+- FLC-QU-hep/calorimeter-showers-multi-geometry
+---
+
+# PointCountFM, multi-geometry pre-trained models
+
+[![arXiv](https://img.shields.io/badge/arXiv-2608.18233-b31b1b?logo=arxiv&logoColor=white)](https://arxiv.org/abs/2608.18233)
+[![Python Version](https://img.shields.io/badge/Python_3.13-306998?logo=python&logoColor=white)](https://www.python.org/)
+[![PyTorch Version](https://img.shields.io/badge/PyTorch_2.8-ee4c2c?logo=pytorch&logoColor=white)](https://pytorch.org/)
+
+Flow-matching models (fully connected) for the **per-layer hit counts** of
+electromagnetic calorimeter showers. PointCountFM is the condition producer in
+the cascade of *[Transferable Fast Calorimeter Shower Generation via Multi-Geometry Pre-training](https://arxiv.org/abs/2608.18233)*: it samples the number of points in
+each calorimeter layer, which then conditions the shower point cloud model
+[FLC-QU-hep/AllShowers-multi-geometry](https://huggingface.co/FLC-QU-hep/AllShowers-multi-geometry).
+
+Conditioning inputs: incident energy, sampling fraction, number of layers, and
+the direction unit vector, in the order `[E, SF, n_layers, dir_x, dir_y, dir_z]`.
+
+## Checkpoints
+
+| Folder | Pre-training data | Output dim (layers) |
+|---|---|---|
+| `simplebox/` | 4M showers, SimpleBox parametric geometry | 45 |
+| `lemurs/` | 4M showers, 4 detectors (Par04 SciPb, Par04 SiW, ODD, CLD) | 90 |
+
+Architecture (both): fully connected flow-matching network, hidden dims
+[128, 256, 512, 256, 128], 6-dim condition, 6-dim time embedding.
+
+## Files and usage
+
+```
+<folder>/
+├── best_model.pt  # best-validation checkpoint (includes fitted norm_stats)
+└── conf.yaml      # architecture + transform pipeline definitions
+```
+
+The checkpoint stores the fitted normalization statistics (`norm_stats`), so
+these two files are all that is needed. With the
+[PointCountFM repository](https://github.com/FLC-QU-hep/PointCountFM/tree/multi-geometry)
+code, point the model loader at the downloaded folder:
+
+```python
+from huggingface_hub import snapshot_download
+model_dir = snapshot_download("FLC-QU-hep/PointCountFM-multi-geometry",
+                              allow_patterns="lemurs/*") + "/lemurs"
+# then load with load_pcfm_model() from src/pcfm_conditioning.py, which reads
+# best_model.pt + conf.yaml from this directory and restores the transforms
+# from the checkpoint's norm_stats
+```
+
+For `lemurs/`, the 90-dim output is zero-padded at the tail: for a detector with
+`n_layers` layers, take entries `0..n_layers-1`.
+
+## Training data
+
+The pre-training datasets (Geant4, LEMURS + SimpleBox) are published at
+[doi:10.25592/uhhfdm.19103](https://doi.org/10.25592/uhhfdm.19103).
+
+## Citation
+
+If you use these weights, please cite:
+
+```bibtex
+@article{Buss2026b,
+    author = {Buss, Thorsten and Day-Hall, Henry and Gaede, Frank and Kasieczka, Gregor and Kr{\"u}ger, Katja and McKeown, Peter and Valente, Lorenzo},
+    title = "{Transferable Fast Calorimeter Shower Generation via Multi-Geometry Pre-training}",
+    eprint = "2608.18233",
+    archivePrefix = "arXiv",
+    primaryClass = "physics.ins-det",
+    month = "8",
+    year = "2026"
+}
+```

lemurs/best_model.pt

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

lemurs/conf.yaml

@@ -0,0 +1,41 @@
+name: pretrain_lemurs
+result_path: .
+
+data:
+  data_file: LEMURS_pretraining_4M.h5
+  batch_size: 1024
+  batch_size_val: 4096
+  train_fraction: 0.975
+  max_samples: null
+  use_nlayers_conditioning: true
+  use_direction_conditioning: true
+  transform_num_points:
+    - [Log, {alpha: 0.5}]
+    - [StandardScaler, {shape: [1, 90]}]
+  transform_fsamp:
+    - [MinMaxScaler, {shape: [1, 1], target_min: -1.0, target_max: 1.0}]
+  transform_nlayers:
+    - [MinMaxScaler, {shape: [1, 1], target_min: -1.0, target_max: 1.0}]
+  # directions are unit vectors — used as-is (no transform needed)
+
+model:
+  name: FullyConnected
+  dim_input: 90       # max layers across LEMURS detectors (par04_siw=90, par04_scipb=45, odd=48, fccee_cld=40)
+  dim_condition: 6    # energy + sampling_fraction + n_layers + dir_x + dir_y + dir_z
+  dim_time: 6
+  hidden_dims: [128, 256, 512, 256, 128]
+
+training:
+  epochs: 400
+  test_every: 100
+
+  # Early stopping on val flow-matching loss.
+  # Trainer reads `patience` (trainer.py:105). 0 = disabled.
+  patience: 30
+
+  optimizer:
+    name: Adam
+    lr: 1.0e-4
+  scheduler:
+    name: OneCycleLR
+    pct_start: 0.5

simplebox/best_model.pt

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

simplebox/conf.yaml

@@ -0,0 +1,37 @@
+name: pretrain_SimpleBox
+result_path: .
+
+data:
+  data_file: SimpleBox_pretraining_4M.h5
+  batch_size: 1024
+  batch_size_val: 4096
+  train_fraction: 0.975
+  max_samples: null
+  use_nlayers_conditioning: true
+  use_direction_conditioning: true
+  transform_num_points:
+    - [Log, {alpha: 0.5}]
+    - [StandardScaler, {shape: [1, 45]}]
+  transform_fsamp:
+    - [MinMaxScaler, {shape: [1, 1], target_min: -1.0, target_max: 1.0}]
+  transform_nlayers:
+    - [MinMaxScaler, {shape: [1, 1], target_min: -1.0, target_max: 1.0}]
+  # directions are unit vectors — used as-is (no transform needed)
+
+model:
+  name: FullyConnected
+  dim_input: 45       # num layers
+  dim_condition: 6    # energy + sampling_fraction + n_layers + dir_x + dir_y + dir_z
+  dim_time: 6
+  hidden_dims: [128, 256, 512, 256, 128]
+
+training:
+  epochs: 5000
+  test_every: 100
+
+  optimizer:
+    name: Adam
+    lr: 1.0e-4
+  scheduler:
+    name: OneCycleLR
+    pct_start: 0.5