README.md

---
license: apache-2.0
tags:
- physics
- calorimeter
- fast-simulation
- generative-model
- flow-matching
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"
}
```