INN for Calorimeter Shower Generation
Code used for "Normalizing Flows for High-Dimensional Detector Simulations" (arxiv:2312:09290) by Ernst F., Favaro L., Krause C., Plehn T., and Shih D.
The samples used in the paper are publicly available on Zenodo.
Fast calorimeter generation for CaloGAN dataset and Fast Calorimeter Challenge.
This is the main repository for the full-space CaloINN network. For the CaloGAN data results check the branch "calogan_data", while for our latent model see "VAE+INN".
Start a training:
python src/main.py params/<param_card>.yaml -c
This creates a subfolder in the results
folder named yyyymmdd_hhmmss_run_name
where the
prefix is the date and time and run_name
is specified in the param card.
Example param card used for pions in params/pions.yaml
Assuming the model is located in <dir>
and saved model with name model_<name>.pt
, to generate a new sample in this directory run:
python3 src/main.py <dir>/params.yaml -d <dir> -its <name> --generate -c
Additionaly --nsamples
provides a way to specify the number of samples generated for dataset 2.
To only run the plotting from the CaloChallenge pipeline using the previously generated samples switch
--generate
with --plot
.
This is a list of the parameters that can be used in yaml parameter files. Many have default values, such that not all parameters have to be specified.
Parameter | Explanation |
---|---|
run_name | Name for the output folder |
Parameter | Explanation |
---|---|
data_path | Name of the hdf5 file containing the data set |
val_data_path | Name of the hdf5 file containing the validation data set |
xml_path | Name of the XML file containing the calorimeter binning |
val_frac | Fraction of the data used for validation |
width_noise | Higher end of the uniform noise to be added to the data |
dtype | float16, float32 or float64; Higher precision makes training and generating slower |
single_energy | Train on a single incident energy, only for dataset 1 |
xml_ptype | Specifics for the XML file: "photon", "pion", or "electron" |
eval_dataset | Needed for the validation used in the CaloChallenge pipeline: "1-photons", "1-pions", "2", or "3" |
Parameter | Explanation |
---|---|
lr | Learning rate |
lr_sched_mode | Type of LR scheduling: "reduce_on_plateau", "step" or "one_cycle" |
lr_decay_epochs | Only step scheduler: decay interval in epochs |
lr_decay_factor | Only step scheduler: decay factor |
batch_size | Batch size |
weight_decay | L2 weight decay |
betas | List of the two Adam beta parameters |
eps | Adam eps parameter for numerical stability |
n_epochs | Number of training epochs |
save_interval | Interval in epochs for saving |
grad_clip | If given, a gradient clip with the given value is applied |
Parameter | Explanation |
---|---|
n_blocks | Number of coupling blocks |
internal_size | Internal size of the coupling block subnetworks |
layers_per_block | Number of layers in each coupling block subnetwork |
dropout | Dropout fraction for the subnetworks |
permute_soft | If True, uses random rotation matrices instead of permutations |
coupling_type | Type of coupling block: "affine", "cubic", "rational_quadratic" or "MADE" |
clamping | Only affine blocks: clamping parameter |
num_bins | Only spline blocks: number of bins |
bounds_init | Only spline blocks: bounds of the splines |
bayesian | True to enable Bayesian training |
sub_layers | A list for partial Bayesian networks, e.g. [linear, linear, VBLinear] |
prior_prec | Only Bayesian: Inverse of the prior standard deviation for the Bayesian layers |
std_init | Only Bayesian: ln of the initial standard deviation of the weight distributions |
layer_act | Activation function in the subnetwork |
norm | Apply ActNorm after preprocessing |
Parameter | Explanation |
---|---|
use_extra_dim | If true an extra dimension is added to the data containing the ratio between parton and detector level energy. This value is used to renormalize generated data. |
use_extra_dims | Adds as extra dimensions the energy variables u_i |
use_norm | If true samples are normalized to the incident energy. Do not use in combination with use_extra_dim or use_extra_dims |
log_cond | If true use the logarithm of the incident energy as condition |
alpha | Constant value to add on the data before taking the logarithm |