A repository for fast detector simulation using Conditional Flow Matching. This is the reference code for CaloDREAM arXiv:XXXX and the corresponding entry in the Calochallenge.
Training:
python3 src/main.py path/to/yaml --use_cuda
The documenter will create a folder in results with the date as
prefix and the specified run_name.
| Parameter | Usage |
|---|---|
| run_name | Name of the output folder |
| hdf5_file | Path to the .hdf5 file used for training |
| xml_filename | Path to the .xml file used to extract the binning information |
| p_type | "photon", "pion", or "electron" |
| dtype | specify default dtype |
| eval_dataset | "1-photons", "1-pions", "2", or "3" used in the CaloChallenge evaluation |
| model_type | Model to be trained: "energy" or "shape" |
| network | Type of network (see Networks for more details) |
| Parameter | Usage |
|---|---|
| dim | Dimensionality of the input |
| n_con | Number of conditions |
| width_noise | Noise width used for the noise injection |
| val_frac | Fraction of events used for validation |
| transforms | Pre-processing steps defined as an ordered dictionary (see transforms.py for more details) |
| lr | learning rate |
| max_lr | Maximum learning rate for OneCycleLR scheduling |
| batch_size | batch size |
| validate_every | Interval between validations in epochs |
| use_scheduler | True or False |
| lr_scheduler | string that defines the learning rate scheduler |
| cycle_epochs | defines the length of the cycle for the OneCycleLR, default to # of epochs |
| save_interval | Interval between each model saving in epochs |
| n_epochs | Number of epochs |
| Parameter | Usage |
|---|---|
| intermediate_dim | Dimension of the intermediate layer |
| layers_per_block | Number of layers per block |
| n_blocks | Number of blocks |
| conditional | True/False, it should be always True |
An example yaml file is provided in ./configs/cfm_base.yaml.
| Parameter | Usage |
|---|---|
| patch_shape | Shape of a single patch |
| hidden_dim | Hidden/Embedding dimension |
| depth | Number of ViT blocks |
| num_heads | Number of transformer heads |
| mlp_ratio | Multiplicative factor for the MLP hidden layers |
| learn_pos_embed | (True or False) learnable position embedding |
An example yaml file is provided in ./configs/d2_shape_model_vit.yaml
Plotting:
To run the sampling and the evaluation of a trained model.
python3 src/main.py --use_cuda --plot --model_dir path/to/model --epoch model_name