PointCountFM-multi-geometry

PointCountFM, multi-geometry pre-trained models

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: it samples the number of points in each calorimeter layer, which then conditions the shower point cloud model 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

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 code, point the model loader at the downloaded folder:

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.

Citation

If you use these weights, please cite:

@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"
}