Import Geant4 11.1.0 source tree

This commit is contained in:
Gabriele Cosmo
2022-12-09 14:43:28 +01:00
parent c07cea1fe0
commit 9f34590941
3810 changed files with 200490 additions and 182326 deletions
@@ -1,20 +1,17 @@
-------------------------------------------------------------------
///\file "parameterisations/Par03/.README.txt"
///\brief Example Par04 README page
=========================================================
Geant4 - an Object-Oriented Toolkit for Simulation in HEP
=========================================================
Example Par04
-------------
/*! \page ExamplePar04 Example Par04
This example demonstrates how to use the Machine Learning (ML) inference
to create energy deposits as a fast simulation model using
<a href="https://github.com/microsoft/onnxruntime">ONNX runtime</a>
and <a href="https://github.com/lwtnn/lwtnn">LWTNN</a> libraries.
<a href="https://github.com/microsoft/onnxruntime">ONNX runtime</a>,
<a href="https://github.com/lwtnn/lwtnn">LWTNN</a>, and
<a href="https://pytorch.org/cppdocs/frontend.html">LibTorch</a> libraries.
The model used in this example was trained externally (in Python) on data
from this examples' full simulation and can be applied to perform fast simulation.
The python scripts are availbale in the training folder.
The python scripts are available in the training folder.
The geometry used in the example is a cylindrical setup of layers: tungsten
absorber and silicon as the active material. 3D readout geometry (cylindrical)
@@ -32,7 +29,7 @@
Input macro can specify which layer is considered an active layer (sensitive
detector is attached to it). For fast simulation both layers should be marked
as sensitive. It is connected to the wway the deposits are created: position is
as sensitive. It is connected to the way the deposits are created: position is
centre of the layer, which may often fall within the absorber (which is thicker
than the active material). In a realistic detector setup, the positions used in
fast simulation would be calculated properly, to deposit energy within the active
@@ -73,7 +70,7 @@
## 6. ML Inference
- Par04MLFastSimModel : model used for parametrisation of źelectrons, positrons,
- Par04MLFastSimModel : model used for parametrisation of electrons, positrons,
and gammas. Energy is deposited and
distributed according to inferred values from the ML model.
This class triggers the inference setup, asks for values,
@@ -91,8 +88,8 @@
- Par04InferenceInterface : is a base class that allows to read in the ML model, configure
and execute inference.
- Par04OnnxInference and Par04LWTNNInference : inference library specific classes that inherit
from the base class Par04InferenceInterface.
- Par04OnnxInference and Par04LWTNNInference and Par04TorchInference : inference library specific
classes that inherit from the base class Par04InferenceInterface.
## 7. Output
@@ -102,41 +99,67 @@
The macro file examplePar04.mac is used to run full simulation. It will simulate 100
events, for single 10 GeV electron beams.
If CMake is able to find inference libraries (lwtnn and/or ONNX Runtime), a configuration
macro will be available for that library (examplePar04_lwtnn.mac and/or examplePar04_onnx.mac).
It will use a trained model to run inference and create showers in the detector by directly
depositing energy.
If CMake is able to find inference libraries (LWTNN and/or ONNX Runtime and/or LibTorch), a configuration
macro will be available for that library (examplePar04_lwtnn.mac and/or examplePar04_onnx.mac
and/or examplePar04_torch.mac). It will use a trained model to run inference and create showers
in the detector by directly depositing energy.
## 8. How to build and run the example
- LWTNN and ONNX Runtime are available on LCG. In order to use them, one can setup the envirnment:
% source /cvmfs/sft.cern.ch/lcg/views/LCG_100/x86_64-centos7-gcc10-opt/setup.sh
- LWTNN, ONNX Runtime, and LibTorch are available on LCG. In order to use them, you can set a `CMAKE_PREFIX_PATH`:
\verbatim
% source /cvmfs/sft.cern.ch/lcg/contrib/gcc/11.3.0/x86_64-centos7/setup.sh
% cmake -DCMAKE_PREFIX_PATH="/cvmfs/sft.cern.ch/lcg/releases/LCG_102b/lwtnn/2.11.1/x86_64-centos7-gcc11-opt/;/cvmfs/sft.cern.ch/lcg/releases/LCG_102b/onnxruntime/1.11.1/x86_64-centos7-gcc11-opt/;/cvmfs/sft.cern.ch/lcg/releases/LCG_102b/torch/1.11.0/x86_64-centos7-gcc11-opt/lib/python3.9/site-packages/torch/" <Par04_SOURCE>
\endverbatim
- Compile and link to generate the executable (in your CMAKE build directory):
% cmake <Par04_SOURCE>
% make
- Compile and link to generate the executable (in your CMake build directory):
\verbatim
% cmake <Par04_SOURCE>
% make
\endverbatim
- Execute the application (in batch mode):
% ./examplePar04 -m examplePar04.mac
\verbatim
% ./examplePar04 -m examplePar04.mac
\endverbatim
which produces two root file for full simulation.
- Execute the application (in interactive mode):
% ./examplePar04 -i -m vis.mac
\verbatim
% ./examplePar04 -i -m vis.mac
\endverbatim
which allows to visualize hits (from full simulation).
- If ONNX Runtime is available:
% ./examplePar04 -m examplePar04_onnx.mac
\verbatim
% ./examplePar04 -m examplePar04_onnx.mac
\endverbatim
For interactive mode with visualization:
% ./examplePar04 -i -m vis_onnx.mac
\verbatim
% ./examplePar04 -i -m vis_onnx.mac
\endverbatim
- If LWTNN is available:
% ./examplePar04 -m examplePar04_lwtnn.mac
\verbatim
% ./examplePar04 -m examplePar04_lwtnn.mac
\endverbatim
For interactive mode with visualization:
% ./examplePar04 -i -m vis_lwtnn.mac
\verbatim
% ./examplePar04 -i -m vis_lwtnn.mac
\endverbatim
- If LibTorch is available:
\verbatim
% ./examplePar04 -m examplePar04_torch.mac
\endverbatim
For interactive mode with visualization:
\verbatim
% ./examplePar04 -i -m vis_torch.mac
\endverbatim
By default, CMake will attempt to build fast simulation with ONNX Runtime and LWTNN. However, if none
of those libraries is found, it will proceed with full simulation only. The search can be switched
off manually switching CMake flag INFERENCE_LIB to OFF (-DINFERENCE_LIB=OFF)
off manually switching CMake flag `INFERENCE_LIB` to `OFF` (`-DINFERENCE_LIB=OFF`)
## 9. Macros
@@ -144,11 +167,15 @@
It can be used to visualize full simulation.
vis_onnx.mac - Allows to run visualization with ONNX Runtime inference. Pass it to the example in interactive mode
("-i" passed to the executable). It contains ecessary settings of the inference, and it treats full
("-i" passed to the executable). It contains necessary settings of the inference, and it treats full
calorimeter as sensitive material (due to deposition of hits regardless of the volume).
vis_lwtnn.mac - Allows to run visualization with LWTNN inference. Pass it to the example in interactive mode
("-i" passed to the executable). It contains ecessary settings of the inference, and it treats full
("-i" passed to the executable). It contains necessary settings of the inference, and it treats full
calorimeter as sensitive material (due to deposition of hits regardless of the volume).
vis_torch.mac - Allows to run visualization with LibTorch inference. Pass it to the example in interactive mode
("-i" passed to the executable). It contains necessary settings of the inference, and it treats full
calorimeter as sensitive material (due to deposition of hits regardless of the volume).
examplePar04.mac - Runs full simulation. It will run 100 events with single electrons, 10 GeV and
@@ -160,91 +187,70 @@
examplePar04_lwtnn.mac - Available only if LWTNN is found by CMake. Runs fast simulation with
a NN stored in json file.
examplePar04_torch.mac - Available only if LibTorch is found by CMake. Runs fast simulation with
a NN stored in pt file.
## 10. UI commands
UI commands useful in this example:
- activation/disactivation of the fast simulation model:
/param/ActivateModel inferenceModel
/param/InActivateModel inferenceModel
\verbatim
/param/ActivateModel inferenceModel
/param/InActivateModel inferenceModel
\endverbatim
- particle gun commands
/gun/particle e-
/gun/energy 10 GeV
/gun/direction 0 1 0
/gun/position 0 0 0
\verbatim
/gun/particle e-
/gun/energy 10 GeV
/gun/direction 0 1 0
/gun/position 0 0 0
\endverbatim
UI commands defined in this example:
- detector settings
/Par04/detector/setDetectorInnerRadius 80 cm
/Par04/detector/setDetectorLength 2 m
/Par04/detector/setNbOfLayers 90
/Par04/detector/setAbsorber 0 G4_W 1.4 mm false
/Par04/detector/setAbsorber 1 G4_Si 0.3 mm true
\verbatim
/Par04/detector/setDetectorInnerRadius 80 cm
/Par04/detector/setDetectorLength 2 m
/Par04/detector/setNbOfLayers 90
/Par04/detector/setAbsorber 0 G4_W 1.4 mm false
/Par04/detector/setAbsorber 1 G4_Si 0.3 mm true
\endverbatim
- readout mesh
/Par04/mesh/setSizeOfRhoCells 2.325 mm
/Par04/mesh/setSizeOfZCells 3.4 mm
/Par04/mesh/setNbOfRhoCells 18
/Par04/mesh/setNbOfPhiCells 50
/Par04/mesh/setNbOfZCells 45
\verbatim
/Par04/mesh/setSizeOfRhoCells 2.325 mm
/Par04/mesh/setSizeOfZCells 3.4 mm
/Par04/mesh/setNbOfRhoCells 18
/Par04/mesh/setNbOfPhiCells 50
/Par04/mesh/setNbOfZCells 45
\endverbatim
- inference setup
/Par04/inference/setSizeLatentVector 10
/Par04/inference/setSizeConditionVector 4
/Par04/inference/setModelPathName MLModels/Generator.onnx
/Par04/inference/setProfileFlag 0
/Par04/inference/setOptimizationFlag 0
/Par04/inference/setInferenceLibrary ONNX
/Par04/inference/setSizeOfRhoCells 2.325 mm
/Par04/inference/setSizeOfZCells 3.4 mm
/Par04/inference/setNbOfRhoCells 18
/Par04/inference/setNbOfPhiCells 50
/Par04/inference/setNbOfZCells 45
\verbatim
/Par04/inference/setSizeLatentVector 10
/Par04/inference/setSizeConditionVector 4
/Par04/inference/setModelPathName MLModels/Generator.onnx
/Par04/inference/setProfileFlag 0
/Par04/inference/setOptimizationFlag 0
/Par04/inference/setInferenceLibrary ONNX
/Par04/inference/setSizeOfRhoCells 2.325 mm
/Par04/inference/setSizeOfZCells 3.4 mm
/Par04/inference/setNbOfRhoCells 18
/Par04/inference/setNbOfPhiCells 50
/Par04/inference/setNbOfZCells 45
\endverbatim
## 11. Python scripts for training
The scripts available in the training folder were used to train
the VAE model of this example.
the VAE model of this example. More details can be found in
training/README.
- model: defines the VAE model as a class which contains the architecture,
the loss function and the training function.
- utils: defines the data loading and preprocessing function and returns
the preprocessed array of shower energies and the condition arrays of energy,
angle and geometry. The input is expected to be in local directory 'detector_*'.
In this example, the model is trained on 2 detector geometries and the input
directories are 'detector_SiW' and 'detector_SciPb'. Each directory contains the
HDF5 files for each primary particle energy and angle.
- train: defines data loading parameters and calls the data preprocessing function.
It defines the model parameters and instantiates the VAE model, performs the training
and then coverts the model into an ONNX format.
A history object is returned from the training function of the VAE. The history is a callback
object registered during the training which records metrics for each epoch such as the loss.
The user can list the metrics collected in the history object using:
% print(history.history.keys())
The data collected in the history object can be used to create plots such as the loss function
as function of the epochs. If the loss metric collected in the history object is called loss,
then to plot it using for example the matplotlib library:
% matplotlib.pyplot.plot(history.history['loss'])
The traing function can also generate intermediate files representing checkpoints of the model
that can be used for validation purposes. These checkpoints are generated if the early stopping flag
of the training is off, which means to train the model for the predefined number of epochs. The weights
of the model are saved every 100 epochs.
After the training, the model (only the decoder part) is saved as an HDF5 file and then is converted
into an ONNX format. The final output model called Generator.onnx can be used to perform the inference
in this example.
To perform the training run:
% python train.py
## 12. Public data
Data generated with full simulation with this example has been published on <a href="https://doi.org/10.5281/zenodo.6082201">zenodo</a>.
*/
@@ -54,10 +54,22 @@ if(INFERENCE_LIB)
endif()
endif()
# TORCH
if(INFERENCE_LIB)
find_package(Torch QUIET)
if(Torch_FOUND)
message("Torch inference library found.")
add_definitions(-DUSE_INFERENCE)
add_definitions(-DUSE_INFERENCE_TORCH)
else()
message("Torch not found!")
endif()
endif()
#----------------------------------------------------------------------------
# Locate sources and headers for this project
#
include_directories(${PROJECT_SOURCE_DIR}/include
include_directories(${PROJECT_SOURCE_DIR}/include
${Geant4_INCLUDE_DIR})
file(GLOB sources ${PROJECT_SOURCE_DIR}/src/*.cc)
file(GLOB headers ${PROJECT_SOURCE_DIR}/include/*.hh)
@@ -82,6 +94,14 @@ if(OnnxRuntime_FOUND)
add_dependencies(examplePar04 examplePar04onnxdata)
endif()
if(Torch_FOUND)
target_include_directories(examplePar04 PUBLIC ${TORCH_INCLUDE_DIRS})
target_link_libraries(examplePar04 ${TORCH_LIBRARIES})
message(STATUS "${TORCH_LIBRARIES}")
# Depend on data for runtime
add_dependencies(examplePar04 examplePar04torchdata)
endif()
#----------------------------------------------------------------------------
# Copy all scripts to the build directory, i.e. the directory in which we
# build Par04. This is so that we can run the executable directly because it
@@ -96,6 +116,9 @@ endif()
if(OnnxRuntime_FOUND)
set(Par04_SCRIPTS ${Par04_SCRIPTS} examplePar04_onnx.mac vis_onnx.mac)
endif()
if(Torch_FOUND)
set(Par04_SCRIPTS ${Par04_SCRIPTS} examplePar04_torch.mac vis_torch.mac)
endif()
foreach(_script ${Par04_SCRIPTS})
configure_file(
@@ -132,6 +155,17 @@ if(OnnxRuntime_FOUND)
DOWNLOAD_NO_EXTRACT true
)
endif()
if(Torch_FOUND)
ExternalProject_Add(examplePar04torchdata
DOWNLOAD_DIR ${PROJECT_BINARY_DIR}/MLModels
URL https://cern.ch/geant4-data/datasets/examples/extended/parameterisations/Par04/Generator.pt
URL_MD5 a43337f7f976e976f1127015f2ba61db
CONFIGURE_COMMAND ""
BUILD_COMMAND ""
INSTALL_COMMAND ""
DOWNLOAD_NO_EXTRACT true
)
endif()
#----------------------------------------------------------------------------
# Add program to the project targets
@@ -4,6 +4,16 @@ See `CONTRIBUTING.rst` for details of **required** info/format for each entry,
which **must** added in reverse chronological order (newest at the top). It must **not**
be used as a substitute for writing good git commit messages!
## 2022-11-08 D. Salamani (expar04-V11-00-05)
- Add updated version of the training code (PEP 8 style guide, formatting with YAPF,
annotation with types, code upgrade to TF2.9, logic for GPU usage management)
## 2022-11-08 A. Zaborowska (expar04-V11-00-04)
- Add support of LibTorch for inference
## 2022-10-25 I. Hrivnacova (expar04-V11-00-03)
- Fixes in Doxygen documentation (links, formatting)
## 2022-02-28 Dalila Salamani (expar04-V11-00-02)
- Add python training scripts
- Add model conversion to ONNX and LWTNN to the train script and update README
@@ -8,12 +8,12 @@
-------------
This example demonstrates how to use the Machine Learning (ML) inference
to create energy deposits as a fast simulation model using ONNX runtime [1]
and LWTNN [2] libraries.
to create energy deposits as a fast simulation model using ONNX Runtime [1],
LWTNN [2], and LibTorch [3] libraries.
The model used in this example was trained externally (in Python) on data
from this examples' full simulation and can be applied to perform fast simulation.
The python scripts are availbale in the training folder.
The python scripts are available in the training folder.
The geometry used in the example is a cylindrical setup of layers: tungsten
absorber and silicon as the active material. 3D readout geometry (cylindrical)
@@ -24,6 +24,7 @@
[1]: https://github.com/microsoft/onnxruntime
[2]: https://github.com/lwtnn/lwtnn
[3]: https://pytorch.org/cppdocs/frontend.html
1. Detector description
-----------------------
@@ -35,7 +36,7 @@
Input macro can specify which layer is considered an active layer (sensitive
detector is attached to it). For fast simulation both layers should be marked
as sensitive. It is connected to the wway the deposits are created: position is
as sensitive. It is connected to the way the deposits are created: position is
centre of the layer, which may often fall within the absorber (which is thicker
than the active material). In a realistic detector setup, the positions used in
fast simulation would be calculated properly, to deposit energy within the active
@@ -80,7 +81,7 @@
6. ML Inference
----------------------------------------------------------
- Par04MLFastSimModel : model used for parametrisation of źelectrons, positrons,
- Par04MLFastSimModel : model used for parametrisation of electrons, positrons,
and gammas. Energy is deposited and
distributed according to inferred values from the ML model.
This class triggers the inference setup, asks for values,
@@ -98,8 +99,8 @@
- Par04InferenceInterface : is a base class that allows to read in the ML model, configure
and execute inference.
- Par04OnnxInference and Par04LWTNNInference : inference library specific classes that inherit
from the base class Par04InferenceInterface.
- Par04OnnxInference and Par04LWTNNInference and Par04TorchInference : inference library specific
classes that inherit from the base class Par04InferenceInterface.
7. Output
@@ -110,17 +111,18 @@
The macro file examplePar04.mac is used to run full simulation. It will simulate 100
events, for single 10 GeV electron beams.
If CMake is able to find inference libraries (lwtnn and/or ONNX Runtime), a configuration
macro will be available for that library (examplePar04_lwtnn.mac and/or examplePar04_onnx.mac).
It will use a trained model to run inference and create showers in the detector by directly
depositing energy.
If CMake is able to find inference libraries (LWTNN and/or ONNX Runtime and/or LibTorch), a configuration
macro will be available for that library (examplePar04_lwtnn.mac and/or examplePar04_onnx.mac
and/or examplePar04_torch.mac). It will use a trained model to run inference and create showers
in the detector by directly depositing energy.
8. How to build and run the example
-----------------------------------
- LWTNN and ONNX Runtime are available on LCG. In order to use them, one can setup the envirnment:
% source /cvmfs/sft.cern.ch/lcg/views/LCG_100/x86_64-centos7-gcc10-opt/setup.sh
- LWTNN, ONNX Runtime, and LibTorch are available on LCG. In order to use them, you can set a CMAKE_PREFIX_PATH:
% source /cvmfs/sft.cern.ch/lcg/contrib/gcc/11.3.0/x86_64-centos7/setup.sh
% cmake -DCMAKE_PREFIX_PATH="/cvmfs/sft.cern.ch/lcg/releases/LCG_102b/lwtnn/2.11.1/x86_64-centos7-gcc11-opt/;/cvmfs/sft.cern.ch/lcg/releases/LCG_102b/onnxruntime/1.11.1/x86_64-centos7-gcc11-opt/;/cvmfs/sft.cern.ch/lcg/releases/LCG_102b/torch/1.11.0/x86_64-centos7-gcc11-opt/lib/python3.9/site-packages/torch/" <Par04_SOURCE>
- Compile and link to generate the executable (in your CMAKE build directory):
- Compile and link to generate the executable (in your CMake build directory):
% cmake <Par04_SOURCE>
% make
@@ -141,7 +143,11 @@
% ./examplePar04 -m examplePar04_lwtnn.mac
For interactive mode with visualization:
% ./examplePar04 -i -m vis_lwtnn.mac
- If LibTorch is available:
% ./examplePar04 -m examplePar04_torch.mac
For interactive mode with visualization:
% ./examplePar04 -i -m vis_torch.mac
By default, CMake will attempt to build fast simulation with ONNX Runtime and LWTNN. However, if none
of those libraries is found, it will proceed with full simulation only. The search can be switched
off manually switching CMake flag INFERENCE_LIB to OFF (-DINFERENCE_LIB=OFF)
@@ -153,11 +159,15 @@
It can be used to visualize full simulation.
vis_onnx.mac - Allows to run visualization with ONNX Runtime inference. Pass it to the example in interactive mode
("-i" passed to the executable). It contains ecessary settings of the inference, and it treats full
("-i" passed to the executable). It contains necessary settings of the inference, and it treats full
calorimeter as sensitive material (due to deposition of hits regardless of the volume).
vis_lwtnn.mac - Allows to run visualization with LWTNN inference. Pass it to the example in interactive mode
("-i" passed to the executable). It contains ecessary settings of the inference, and it treats full
("-i" passed to the executable). It contains necessary settings of the inference, and it treats full
calorimeter as sensitive material (due to deposition of hits regardless of the volume).
vis_torch.mac - Allows to run visualization with LibTorch inference. Pass it to the example in interactive mode
("-i" passed to the executable). It contains necessary settings of the inference, and it treats full
calorimeter as sensitive material (due to deposition of hits regardless of the volume).
examplePar04.mac - Runs full simulation. It will run 100 events with single electrons, 10 GeV and
@@ -169,6 +179,9 @@
examplePar04_lwtnn.mac - Available only if LWTNN is found by CMake. Runs fast simulation with
a NN stored in json file.
examplePar04_torch.mac - Available only if LibTorch is found by CMake. Runs fast simulation with
a NN stored in pt file.
10. UI commands
--------------
@@ -216,45 +229,8 @@
--------------
The scripts available in the training folder were used to train
the VAE model of this example.
- model: defines the VAE model as a class which contains the architecture,
the loss function and the training function.
- utils: defines the data loading and preprocessing function and returns
the preprocessed array of shower energies and the condition arrays of energy,
angle and geometry. The input is expected to be in local directory 'detector_*'.
In this example, the model is trained on 2 detector geometries and the input
directories are 'detector_SiW' and 'detector_SciPb'. Each directory contains the
HDF5 files for each primary particle energy and angle.
- train: defines data loading parameters and calls the data preprocessing function.
It defines the model parameters and instantiates the VAE model, performs the training
and then coverts the model into an ONNX format.
A history object is returned from the training function of the VAE. The history is a callback
object registered during the training which records metrics for each epoch such as the loss.
The user can list the metrics collected in the history object using:
% print(history.history.keys())
The data collected in the history object can be used to create plots such as the loss function
as function of the epochs. If the loss metric collected in the history object is called loss,
then to plot it using for example the matplotlib library:
% matplotlib.pyplot.plot(history.history['loss'])
The traing function can also generate intermediate files representing checkpoints of the model
that can be used for validation purposes. These checkpoints are generated if the early stopping flag
of the training is off, which means to train the model for the predefined number of epochs. The weights
of the model are saved every 100 epochs.
After the training, the model (only the decoder part) is saved as an HDF5 file and then is converted
into an ONNX format. The final output model called Generator.onnx can be used to perform the inference
in this example.
To perform the training run:
% python train.py
the VAE model of this example. More details can be found in
training/README.
12. Public data
@@ -11,7 +11,7 @@ Environment variable "G4FORCE_RUN_MANAGER_TYPE" enabled with value == Serial. Fo
**************************************************************
Geant4 version Name: geant4-11-01-beta-01 (30-June-2022)
Geant4 version Name: geant4-11-01-ref-00 (9-December-2022)
Copyright : Geant4 Collaboration
References : NIM A 506 (2003), 250-303
: IEEE-TNS 53 (2006), 270-278
@@ -34,6 +34,8 @@ Registered graphics systems are:
RayTracer (RayTracer)
VRML2FILE (VRML2FILE)
gMocrenFile (gMocrenFile)
TOOLSSG_OFFSCREEN (TSG_OFFSCREEN)
TOOLSSG_OFFSCREEN (TSG_OFFSCREEN, TSG_FILE)
OpenGLImmediateQt (OGLIQt, OGLI)
OpenGLStoredQt (OGLSQt, OGL, OGLS)
OpenGLImmediateXm (OGLIXm, OGLIQt_FALLBACK)
@@ -278,7 +280,6 @@ Checking overlaps for volume Layer:179 (G4Tubs) ... OK!
e+ : fastSimProcess_massGeom[geom:World]
e- : fastSimProcess_massGeom[geom:World]
gamma : fastSimProcess_massGeom[geom:World]
Set file name: 10GeV_100events_fullsim.root
Model defineMesh activated.
Model inferenceModel not found.
@@ -417,6 +418,11 @@ Model inferenceModel not found.
Process: hFritiofCaptureAtRest
---------------------------------------------------
Hadronic Processes for anti_hypertriton
Process: hFritiofCaptureAtRest
---------------------------------------------------
Hadronic Processes for anti_lambda
@@ -0,0 +1,47 @@
# examplePar04_torch.mac
#
# Detector Construction
/Par04/detector/setDetectorInnerRadius 80 cm
/Par04/detector/setDetectorLength 2 m
/Par04/detector/setNbOfLayers 90
/Par04/detector/setAbsorber 0 G4_W 1.4 mm true
/Par04/detector/setAbsorber 1 G4_Si 0.3 mm true
## 2.325 mm of tungsten =~ 0.25 * 9.327 mm = 0.25 * R_Moliere
/Par04/mesh/setSizeOfRhoCells 2.325 mm
## 2 * 1.4 mm of tungsten =~ 0.65 X_0
/Par04/mesh/setSizeOfZCells 3.4 mm
/Par04/mesh/setNbOfRhoCells 18
/Par04/mesh/setNbOfPhiCells 50
/Par04/mesh/setNbOfZCells 45
# Initialize
/run/initialize
/gun/energy 10 GeV
/gun/position 0 0 0
/gun/direction 0 1 0
# Inference Setup
## dimension of the latent vector (encoded vector in a Variational Autoencoder model)
/Par04/inference/setSizeLatentVector 10
## size of the condition vector (energy, angle and geometry)
/Par04/inference/setSizeConditionVector 4
## path to the model which is set to download by cmake
/Par04/inference/setModelPathName MLModels/Generator.pt
/Par04/inference/setInferenceLibrary TORCH
## set mesh size for inference == mesh size of a full sim that
## was used for training; it coincides with readout mesh size
/Par04/inference/setSizeOfRhoCells 2.325 mm
/Par04/inference/setSizeOfZCells 3.4 mm
/Par04/inference/setNbOfRhoCells 18
/Par04/inference/setNbOfPhiCells 50
/Par04/inference/setNbOfZCells 45
# Fast Simulation
/analysis/setFileName 10GeV_100events_fastsim_libtorch.root
## dynamically set readout mesh from particle direction
## needs to be the first fast sim model!
/param/ActivateModel defineMesh
## ML fast sim, configured with the inference setup /Par04/inference
/param/ActivateModel inferenceModel
/run/beamOn 100
@@ -0,0 +1,62 @@
//
// ********************************************************************
// * License and Disclaimer *
// * *
// * The Geant4 software is copyright of the Copyright Holders of *
// * the Geant4 Collaboration. It is provided under the terms and *
// * conditions of the Geant4 Software License, included in the file *
// * LICENSE and available at http://cern.ch/geant4/license . These *
// * include a list of copyright holders. *
// * *
// * Neither the authors of this software system, nor their employing *
// * institutes,nor the agencies providing financial support for this *
// * work make any representation or warranty, express or implied, *
// * regarding this software system or assume any liability for its *
// * use. Please see the license in the file LICENSE and URL above *
// * for the full disclaimer and the limitation of liability. *
// * *
// * This code implementation is the result of the scientific and *
// * technical work of the GEANT4 collaboration. *
// * By using, copying, modifying or distributing the software (or *
// * any work based on the software) you agree to acknowledge its *
// * use in resulting scientific publications, and indicate your *
// * acceptance of all terms of the Geant4 Software license. *
// ********************************************************************
//
#ifdef USE_INFERENCE_TORCH
#ifndef PAR04TORCHINFERENCE_HH
#define PAR04TORCHINFERENCE_HH
#include <G4String.hh> // for G4String
#include <G4Types.hh> // for G4int, G4double
#include <memory> // for unique_ptr
#include <vector> // for vector
#include "Par04InferenceInterface.hh" // for Par04InferenceInterface
#include <torch/script.h>
/**
* @brief Inference using the TORCH.
*
* Runs the inference with LibTorch using the input vector from Par04InferenceSetup.
*
**/
class Par04TorchInference : public Par04InferenceInterface
{
public:
Par04TorchInference(G4String);
Par04TorchInference();
/// Run inference
/// @param[in] aGenVector Input latent space and conditions
/// @param[out] aEnergies Model output = generated shower energies
/// @param[in] aSize Size of the output
void RunInference(std::vector<float> aGenVector, std::vector<G4double>& aEnergies, int aSize);
private:
torch::jit::script::Module fModule;
};
#endif /* PAR04TORCHINFERENCE_HH */
#endif
@@ -33,6 +33,9 @@
#ifdef USE_INFERENCE_LWTNN
#include "Par04LwtnnInference.hh" // for Par04LwtnnInference
#endif
#ifdef USE_INFERENCE_TORCH
#include "Par04TorchInference.hh" // for Par04TorchInference
#endif
#include <CLHEP/Units/SystemOfUnits.h> // for pi, GeV, deg
#include <CLHEP/Vector/Rotation.h> // for HepRotation
#include <CLHEP/Vector/ThreeVector.h> // for Hep3Vector
@@ -82,6 +85,12 @@ void Par04InferenceSetup::SetInferenceLibrary(G4String aName)
fInferenceInterface =
std::unique_ptr<Par04InferenceInterface>(new Par04LwtnnInference(fModelPathName));
#endif
#ifdef USE_INFERENCE_TORCH
if(fInferenceLibrary == "TORCH")
fInferenceInterface =
std::unique_ptr<Par04InferenceInterface>(new Par04TorchInference(fModelPathName));
#endif
CheckInferenceLibrary();
}
@@ -94,7 +103,10 @@ void Par04InferenceSetup::CheckInferenceLibrary()
msg += "ONNX,";
#endif
#ifdef USE_INFERENCE_LWTNN
msg += "LWTNN";
msg += "LWTNN,";
#endif
#ifdef USE_INFERENCE_TORCH
msg += "TORCH";
#endif
if(fInferenceInterface == nullptr)
G4Exception("Par04InferenceSetup::CheckInferenceLibrary()", "InvalidSetup", FatalException,
@@ -0,0 +1,90 @@
//
// ********************************************************************
// * License and Disclaimer *
// * *
// * The Geant4 software is copyright of the Copyright Holders of *
// * the Geant4 Collaboration. It is provided under the terms and *
// * conditions of the Geant4 Software License, included in the file *
// * LICENSE and available at http://cern.ch/geant4/license . These *
// * include a list of copyright holders. *
// * *
// * Neither the authors of this software system, nor their employing *
// * institutes,nor the agencies providing financial support for this *
// * work make any representation or warranty, express or implied, *
// * regarding this software system or assume any liability for its *
// * use. Please see the license in the file LICENSE and URL above *
// * for the full disclaimer and the limitation of liability. *
// * *
// * This code implementation is the result of the scientific and *
// * technical work of the GEANT4 collaboration. *
// * By using, copying, modifying or distributing the software (or *
// * any work based on the software) you agree to acknowledge its *
// * use in resulting scientific publications, and indicate your *
// * acceptance of all terms of the Geant4 Software license. *
// ********************************************************************
//
#ifdef USE_INFERENCE_TORCH
#include "Par04TorchInference.hh"
#include <algorithm> // for copy, max
#include <cassert> // for assert
#include <cstddef> // for size_t
#include <cstdint> // for int64_t
#include <utility> // for move
#include "Par04InferenceInterface.hh" // for Par04InferenceInterface
#include <torch/torch.h>
//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
Par04TorchInference::Par04TorchInference(G4String modelPath)
: Par04InferenceInterface()
{
fModule = torch::jit::load( modelPath );
}
//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
void Par04TorchInference::RunInference(std::vector<float> aGenVector, std::vector<G4double>& aEnergies,
int aSize)
{
// latentSize : size of the latent space
// 4 is the size of the condition vector
int latentSize = aGenVector.size() - 4;
// split into latent and condition vectors
std::vector<float> latent;
for ( int i=0;i<latentSize;i++) {
latent.push_back(aGenVector[i]);
}
std::vector<float> energy;
energy.push_back(aGenVector[latentSize+1]);
std::vector<float> angle;
energy.push_back(aGenVector[latentSize+2]);
std::vector<float> geo;
for ( int i=latentSize+2;i<latentSize+4;i++) {
geo.push_back(aGenVector[i]);
}
// convert vectors to tensors
torch::Tensor latentVector = torch::tensor(latent);
torch::Tensor eTensor = torch::tensor(energy);
torch::Tensor angleTensor = torch::tensor(angle);
torch::Tensor geoTensor = torch::tensor(geo);
std::vector<torch::jit::IValue> genInput;
genInput.push_back( latentVector );
genInput.push_back( eTensor );
genInput.push_back( angleTensor );
genInput.push_back( geoTensor );
at::Tensor outTensor = fModule.forward( genInput).toTensor().contiguous();
std::vector<G4double> output( outTensor.data_ptr<float>(), outTensor.data_ptr<float>() + outTensor.numel() );
aEnergies.assign(aSize, 0);
for(int i = 0; i < aSize; i++) {
aEnergies[i] = output[i];
}
}
#endif
@@ -0,0 +1,98 @@
This repository contains the set of scripts used to train, generate and validate the generative model used
in this example.
- core/constants.py: defines the set of common variables.
- core/model.py: defines the VAE model class and a handler to construct the model.
- utils/preprocess.py: defines the data loading and preprocessing functions.
- utils/hyperparameter_tuner.py: defines the HyperparameterTuner class.
- utils/gpu_limiter.py: defines a logic responsible for GPU memory management.
- utils/observables.py: defines a set of observable possibly calculated from a shower.
- utils/plotter.py: defines plotting classes responsible for manufacturing various plots of observables.
- train.py: performs model training.
- generate.py: generate showers using a saved VAE model.
- observables.py: defines a set of shower observables.
- validate.py: creates validation plots using shower observables.
- convert.py: defines the conversion function to an ONNX file.
- tune_model.py: performs hyperparameters optimization.
## Getting Started
`setup.py` script creates necessary folders used to save model checkpoints, generate showers and validation plots.
```
python3 setup.py
```
## Full simulation dataset
The full simulation dataset can be downloaded from/linked to [Zenodo](https://zenodo.org/record/6082201#.Ypo5UeDRaL4).
## Training
In order to launch the training:
```
python3 train.py
```
You may specify those three following flags. If you do not, then default values will be used.
```--max-gpu-memory-allocation``` specifies a maximum memory allocation on a single, logic GPU unit. Should be given as
an integer.
```--gpu-ids``` specifies IDs of physical GPUs. Should be given as a string, separated with comas, no spaces.
If you specify more than one GPU then automatically ```tf.distribute.MirroredStrategy``` will be applied to the
training.
```--study-name``` specifies a study name. This name is used as an experiment name in W&B dashboard and as a name of
directory for saving models.
## Hyperparameters tuning
If you want to tune hyperparameters, specify in `tune_model.py` parameters to be tuned. There are three types of
parameters: discrete, continuous and categorical. Discrete and continuous require range specification (low, high), while
the categorical parameter requires a list of possible values to be chosen. Then run it with:
```
python3 tune_model.py
```
If you want to parallelize tuning process you need to specify a common storage (preferable MySQL database) by
setting `--storage="URL_TO_MYSQL_DATABASE"`. Then you can run multiple processes with the same command:
```
python3 tune_model.py --storage="URL_TO_MYSQL_DATABASE"
```
Similarly to training procedure, you may specify ```--max-gpu-memory-allocation```, ```--gpu-ids``` and
```--study-name```.
## ML shower generation (MLFastSim)
In order to generate showers using the ML model, use `generate.py` script and specify information of geometry, energy
and angle of the particle and the epoch of the saved checkpoint model. The number of events to generate can also be
specified (by default is set to 10.000):
```
python3 generate.py --geometry=SiW --energy=64 --angle=90 --epoch=1000 --study-name=YOUR_STUDY_NAME
```
If you do not specify an epoch number the based model (saved as ```VAEbest```) will be used for shower generation.
## Validation
In order to validate the MLFastSim and the full simulation, use `validate.py` script and specify information of
geometry, energy and angle of the particle:
```
python3 validate.py --geometry=SiW --energye=64 --angle=90
```
## Conversion
After training and validation, the model can be converted into a format that can be used in C++, such as ONNX,
use `convert.py` script:
```
python3 convert.py --epoch 1000
```
@@ -0,0 +1,72 @@
"""
** convert **
defines the conversion function to and ONNX file
"""
import argparse
import sys
import tf2onnx
import numpy as np
from onnxruntime import InferenceSession
from core.constants import GLOBAL_CHECKPOINT_DIR, CONV_DIR, ORIGINAL_DIM
from core.model import VAEHandler
"""
epoch: epoch of the saved checkpoint model
study-name: study-name for which the model is trained for
"""
def parse_args(argv):
p = argparse.ArgumentParser()
p.add_argument("--epoch", type=int, default=None)
p.add_argument("--study-name", type=str, default="default_study_name")
args = p.parse_args()
return args
# main function
def main(argv):
# 1. Set up the model to convert
# Parse commandline arguments
args = parse_args(argv)
epoch = args.epoch
study_name = args.study_name
# Instantiate and load a saved model
vae = VAEHandler()
# Load the saved weights
weights_dir = f"VAE_epoch_{epoch:03}" if epoch is not None else "VAE_best"
vae.model.load_weights(
f"{GLOBAL_CHECKPOINT_DIR}/{study_name}/{weights_dir}/model_weights"
).expect_partial()
# 2. Convert the model to ONNX format
# Create the Keras model, convert it into an ONNX model, and save.
keras_model = vae.model.decoder
output_path = f"{CONV_DIR}/{study_name}/Generator_{weights_dir}.onnx"
onnx_model = tf2onnx.convert.from_keras(keras_model,
output_path=output_path)
# Checking the converted model
input_1 = np.random.randn(10).astype(np.float32).reshape(1, -1)
input_2 = np.random.randn(1).astype(np.float32).reshape(1, -1)
input_3 = np.random.randn(1).astype(np.float32).reshape(1, -1)
input_4 = np.random.randn(2).astype(np.float32).reshape(1, -1)
sess = InferenceSession(output_path)
# TODO: @Piyush-555 Find a way to use predefined names
result = sess.run(
None, {
'input_9': input_1,
'input_6': input_2,
'input_7': input_3,
'input_8': input_4
})
assert result[0].shape[1] == ORIGINAL_DIM
if __name__ == "__main__":
exit(main(sys.argv[1:]))
@@ -0,0 +1,86 @@
from utils.optimizer import OptimizerType
"""
Experiment constants.
"""
# Number of calorimeter layers (z-axis segmentation).
N_CELLS_Z = 45
# Segmentation in the r,phi direction.
N_CELLS_R = 18
N_CELLS_PHI = 50
# Cell size in the r and z directions
SIZE_R = 2.325
SIZE_Z = 3.4
# Minimum and maximum primary particle energy to consider for training in GeV units.
MIN_ENERGY = 1
MAX_ENERGY = 1024
# Minimum and maximum primary particle angle to consider for training in degrees units.
MIN_ANGLE = 50
MAX_ANGLE = 90
"""
Directories.
"""
# Directory to load the full simulation dataset.
INIT_DIR = "./dataset/"
# Directory to save VAE checkpoints
GLOBAL_CHECKPOINT_DIR = "./checkpoint"
# Directory to save model after conversion to a format that can be used in C++.
CONV_DIR = "./conversion"
# Directory to save validation plots.
VALID_DIR = "./validation"
# Directory to save VAE generated showers.
GEN_DIR = "./generation"
"""
Model default parameters.
"""
BATCH_SIZE_PER_REPLICA = 128
# Total number of readout cells (represents the number of nodes in the input/output layers of the model).
ORIGINAL_DIM = N_CELLS_Z * N_CELLS_R * N_CELLS_PHI
INTERMEDIATE_DIMS = [100, 50, 20, 14]
LATENT_DIM = 10
EPOCHS = 1000
LEARNING_RATE = 0.001
ACTIVATION = "leaky_relu"
OUT_ACTIVATION = "sigmoid"
VALIDATION_SPLIT = 0.10
NUMBER_OF_K_FOLD_SPLITS = 1
OPTIMIZER_TYPE = OptimizerType.ADAM
KERNEL_INITIALIZER = "RandomNormal"
BIAS_INITIALIZER = "Zeros"
EARLY_STOP = False
SAVE_BEST_MODEL = True
SAVE_MODEL_EVERY_EPOCH = True
PATIENCE = 10
MIN_DELTA = 0.01
BEST_MODEL_FILENAME = "VAE_best"
# GPU identifiers separated by comma, no spaces.
GPU_IDS = "0"
# Maximum allowed memory on one of the GPUs (in GB)
MAX_GPU_MEMORY_ALLOCATION = 32
# Buffer size used while shuffling the dataset.
BUFFER_SIZE = 1000
"""
Optimizer parameters.
"""
N_TRIALS = 50
# Maximum size of a hidden layer
MAX_HIDDEN_LAYER_DIM = 2000
"""
Validator parameter.
"""
FULL_SIM_HISTOGRAM_COLOR = "blue"
ML_SIM_HISTOGRAM_COLOR = "red"
FULL_SIM_GAUSSIAN_COLOR = "green"
ML_SIM_GAUSSIAN_COLOR = "orange"
HISTOGRAM_TYPE = "step"
"""
W&B parameters.
"""
# Change this to your entity name.
WANDB_ENTITY = "entity-name"
@@ -0,0 +1,429 @@
import gc
from dataclasses import dataclass, field
from typing import List, Tuple
import numpy as np
import tensorflow as tf
import wandb
from sklearn.model_selection import KFold
from tensorflow.keras import backend as K
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint, History, Callback
from tensorflow.keras.layers import BatchNormalization, Input, Dense, Layer, concatenate
from tensorflow.keras.losses import BinaryCrossentropy, Reduction
from tensorflow.keras.models import Model
from tensorflow.python.data import Dataset
from tensorflow.python.distribute.distribute_lib import Strategy
from tensorflow.python.distribute.mirrored_strategy import MirroredStrategy
from wandb.keras import WandbCallback
from core.constants import ORIGINAL_DIM, LATENT_DIM, BATCH_SIZE_PER_REPLICA, EPOCHS, LEARNING_RATE, ACTIVATION, \
OUT_ACTIVATION, OPTIMIZER_TYPE, KERNEL_INITIALIZER, GLOBAL_CHECKPOINT_DIR, EARLY_STOP, BIAS_INITIALIZER, \
INTERMEDIATE_DIMS, SAVE_MODEL_EVERY_EPOCH, SAVE_BEST_MODEL, PATIENCE, MIN_DELTA, BEST_MODEL_FILENAME, \
NUMBER_OF_K_FOLD_SPLITS, VALIDATION_SPLIT, WANDB_ENTITY
from utils.optimizer import OptimizerFactory, OptimizerType
class _Sampling(Layer):
""" Custom layer to do the reparameterization trick: sample random latent vectors z from the latent Gaussian
distribution.
The sampled vector z is given by sampled_z = mean + std * epsilon
"""
def __call__(self, inputs, **kwargs):
z_mean, z_log_var, epsilon = inputs
z_sigma = K.exp(0.5 * z_log_var)
return z_mean + z_sigma * epsilon
# KL divergence computation
class _KLDivergenceLayer(Layer):
def call(self, inputs, **kwargs):
mu, log_var = inputs
kl_loss = -0.5 * (1 + log_var - K.square(mu) - K.exp(log_var))
kl_loss = K.mean(K.sum(kl_loss, axis=-1))
self.add_loss(kl_loss)
return inputs
class VAE(Model):
def get_config(self):
config = super().get_config()
config["encoder"] = self.encoder
config["decoder"] = self.decoder
return config
def call(self, inputs, training=None, mask=None):
_, e_input, angle_input, geo_input, _ = inputs
z = self.encoder(inputs)
return self.decoder([z, e_input, angle_input, geo_input])
def __init__(self, encoder, decoder, **kwargs):
super(VAE, self).__init__(**kwargs)
self.encoder = encoder
self.decoder = decoder
self._set_inputs(inputs=self.encoder.inputs, outputs=self(self.encoder.inputs))
@dataclass
class VAEHandler:
"""
Class to handle building and training VAE models.
"""
_wandb_project_name: str = None
_wandb_tags: List[str] = field(default_factory=list)
_original_dim: int = ORIGINAL_DIM
latent_dim: int = LATENT_DIM
_batch_size_per_replica: int = BATCH_SIZE_PER_REPLICA
_intermediate_dims: List[int] = field(default_factory=lambda: INTERMEDIATE_DIMS)
_learning_rate: float = LEARNING_RATE
_epochs: int = EPOCHS
_activation: str = ACTIVATION
_out_activation: str = OUT_ACTIVATION
_number_of_k_fold_splits: float = NUMBER_OF_K_FOLD_SPLITS
_optimizer_type: OptimizerType = OPTIMIZER_TYPE
_kernel_initializer: str = KERNEL_INITIALIZER
_bias_initializer: str = BIAS_INITIALIZER
_checkpoint_dir: str = GLOBAL_CHECKPOINT_DIR
_early_stop: bool = EARLY_STOP
_save_model_every_epoch: bool = SAVE_MODEL_EVERY_EPOCH
_save_best_model: bool = SAVE_BEST_MODEL
_patience: int = PATIENCE
_min_delta: float = MIN_DELTA
_best_model_filename: str = BEST_MODEL_FILENAME
_validation_split: float = VALIDATION_SPLIT
_strategy: Strategy = MirroredStrategy()
def __post_init__(self) -> None:
# Calculate true batch size.
self._batch_size = self._batch_size_per_replica * self._strategy.num_replicas_in_sync
self._build_and_compile_new_model()
# Setup Wandb.
if self._wandb_project_name is not None:
self._setup_wandb()
def _setup_wandb(self) -> None:
config = {
"learning_rate": self._learning_rate,
"batch_size": self._batch_size,
"epochs": self._epochs,
"optimizer_type": self._optimizer_type,
"intermediate_dims": self._intermediate_dims,
"latent_dim": self.latent_dim
}
# Reinit flag is needed for hyperparameter tuning. Whenever new training is started, new Wandb run should be
# created.
wandb.init(project=self._wandb_project_name, entity=WANDB_ENTITY, reinit=True, config=config,
tags=self._wandb_tags)
def _build_and_compile_new_model(self) -> None:
""" Builds and compiles a new model.
VAEHandler keep a list of VAE instance. The reason is that while k-fold cross validation is performed,
each fold requires a new, clear instance of model. New model is always added at the end of the list of
existing ones.
Returns: None
"""
# Build encoder and decoder.
encoder = self._build_encoder()
decoder = self._build_decoder()
# Compile model within a distributed strategy.
with self._strategy.scope():
# Build VAE.
self.model = VAE(encoder, decoder)
# Manufacture an optimizer and compile model with.
optimizer = OptimizerFactory.create_optimizer(self._optimizer_type, self._learning_rate)
reconstruction_loss = BinaryCrossentropy(reduction=Reduction.SUM)
self.model.compile(optimizer=optimizer, loss=[reconstruction_loss], loss_weights=[ORIGINAL_DIM])
def _prepare_input_layers(self, for_encoder: bool) -> List[Input]:
"""
Create four Input layers. Each of them is responsible to take respectively: batch of showers/batch of latent
vectors, batch of energies, batch of angles, batch of geometries.
Args:
for_encoder: Boolean which decides whether an input is full dimensional shower or a latent vector.
Returns:
List of Input layers (five for encoder and four for decoder).
"""
e_input = Input(shape=(1,))
angle_input = Input(shape=(1,))
geo_input = Input(shape=(2,))
if for_encoder:
x_input = Input(shape=self._original_dim)
eps_input = Input(shape=self.latent_dim)
return [x_input, e_input, angle_input, geo_input, eps_input]
else:
x_input = Input(shape=self.latent_dim)
return [x_input, e_input, angle_input, geo_input]
def _build_encoder(self) -> Model:
""" Based on a list of intermediate dimensions, activation function and initializers for kernel and bias builds
the encoder.
Returns:
Encoder is returned as a keras.Model.
"""
with self._strategy.scope():
# Prepare input layer.
x_input, e_input, angle_input, geo_input, eps_input = self._prepare_input_layers(for_encoder=True)
x = concatenate([x_input, e_input, angle_input, geo_input])
# Construct hidden layers (Dense and Batch Normalization).
for intermediate_dim in self._intermediate_dims:
x = Dense(units=intermediate_dim, activation=self._activation,
kernel_initializer=self._kernel_initializer,
bias_initializer=self._bias_initializer)(x)
x = BatchNormalization()(x)
# Add Dense layer to get description of multidimensional Gaussian distribution in terms of mean
# and log(variance).
z_mean = Dense(self.latent_dim, name="z_mean")(x)
z_log_var = Dense(self.latent_dim, name="z_log_var")(x)
# Add KLDivergenceLayer responsible for calculation of KL loss.
z_mean, z_log_var = _KLDivergenceLayer()([z_mean, z_log_var])
# Sample a probe from the distribution.
encoder_output = _Sampling()([z_mean, z_log_var, eps_input])
# Create model.
encoder = Model(inputs=[x_input, e_input, angle_input, geo_input, eps_input], outputs=encoder_output,
name="encoder")
return encoder
def _build_decoder(self) -> Model:
""" Based on a list of intermediate dimensions, activation function and initializers for kernel and bias builds
the decoder.
Returns:
Decoder is returned as a keras.Model.
"""
with self._strategy.scope():
# Prepare input layer.
latent_input, e_input, angle_input, geo_input = self._prepare_input_layers(for_encoder=False)
x = concatenate([latent_input, e_input, angle_input, geo_input])
# Construct hidden layers (Dense and Batch Normalization).
for intermediate_dim in reversed(self._intermediate_dims):
x = Dense(units=intermediate_dim, activation=self._activation,
kernel_initializer=self._kernel_initializer,
bias_initializer=self._bias_initializer)(x)
x = BatchNormalization()(x)
# Add Dense layer to get output which shape is compatible in an input's shape.
decoder_outputs = Dense(units=self._original_dim, activation=self._out_activation)(x)
# Create model.
decoder = Model(inputs=[latent_input, e_input, angle_input, geo_input], outputs=decoder_outputs,
name="decoder")
return decoder
def _manufacture_callbacks(self) -> List[Callback]:
"""
Based on parameters set by the user, manufacture callbacks required for training.
Returns:
A list of `Callback` objects.
"""
callbacks = []
# If the early stopping flag is on then stop the training when a monitored metric (validation) has stopped
# improving after (patience) number of epochs.
if self._early_stop:
callbacks.append(
EarlyStopping(monitor="val_loss",
min_delta=self._min_delta,
patience=self._patience,
verbose=True,
restore_best_weights=True))
# Save model after every epoch.
if self._save_model_every_epoch:
callbacks.append(ModelCheckpoint(filepath=f"{self._checkpoint_dir}/VAE_epoch_{{epoch:03}}/model_weights",
monitor="val_loss",
verbose=True,
save_weights_only=True,
mode="min",
save_freq="epoch"))
# Pass metadata to wandb.
callbacks.append(WandbCallback(
monitor="val_loss", verbose=0, mode="auto", save_model=False))
return callbacks
def _get_train_and_val_data(self, dataset: np.array, e_cond: np.array, angle_cond: np.array, geo_cond: np.array,
noise: np.array, train_indexes: np.array, validation_indexes: np.array) \
-> Tuple[Dataset, Dataset]:
"""
Splits data into train and validation set based on given lists of indexes.
"""
# Prepare training data.
train_dataset = dataset[train_indexes, :]
train_e_cond = e_cond[train_indexes]
train_angle_cond = angle_cond[train_indexes]
train_geo_cond = geo_cond[train_indexes, :]
train_noise = noise[train_indexes, :]
# Prepare validation data.
val_dataset = dataset[validation_indexes, :]
val_e_cond = e_cond[validation_indexes]
val_angle_cond = angle_cond[validation_indexes]
val_geo_cond = geo_cond[validation_indexes, :]
val_noise = noise[validation_indexes, :]
# Gather them into tuples.
train_x = (train_dataset, train_e_cond, train_angle_cond, train_geo_cond, train_noise)
train_y = train_dataset
val_x = (val_dataset, val_e_cond, val_angle_cond, val_geo_cond, val_noise)
val_y = val_dataset
# Wrap data in Dataset objects.
# TODO(@mdragula): This approach requires loading the whole data set to RAM. It
# would be better to read the data partially when needed. Also one should bare in mind that using tf.Dataset
# slows down training process.
train_data = Dataset.from_tensor_slices((train_x, train_y))
val_data = Dataset.from_tensor_slices((val_x, val_y))
# The batch size must now be set on the Dataset objects.
train_data = train_data.batch(self._batch_size)
val_data = val_data.batch(self._batch_size)
# Disable AutoShard.
options = tf.data.Options()
options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.DATA
train_data = train_data.with_options(options)
val_data = val_data.with_options(options)
return train_data, val_data
def _k_fold_training(self, dataset: np.array, e_cond: np.array, angle_cond: np.array, geo_cond: np.array,
noise: np.array, callbacks: List[Callback], verbose: bool = True) -> List[History]:
"""
Performs K-fold cross validation training.
Number of fold is defined by (self._number_of_k_fold_splits). Always shuffle the dataset.
Args:
dataset: A matrix representing showers. Shape =
(number of samples, ORIGINAL_DIM = N_CELLS_Z * N_CELLS_R * N_CELLS_PHI).
e_cond: A matrix representing an energy for each sample. Shape = (number of samples, ).
angle_cond: A matrix representing an angle for each sample. Shape = (number of samples, ).
geo_cond: A matrix representing a geometry of the detector for each sample. Shape = (number of samples, 2).
noise: A matrix representing an additional noise needed to perform a reparametrization trick.
callbacks: A list of callback forwarded to the fitting function.
verbose: A boolean which says there the training should be performed in a verbose mode or not.
Returns: A list of `History` objects.`History.history` attribute is a record of training loss values and
metrics values at successive epochs, as well as validation loss values and validation metrics values (if
applicable).
"""
# TODO(@mdragula): KFold cross validation can be parallelized. Each fold is independent from each the others.
k_fold = KFold(n_splits=self._number_of_k_fold_splits, shuffle=True)
histories = []
for i, (train_indexes, validation_indexes) in enumerate(k_fold.split(dataset)):
print(f"K-fold: {i + 1}/{self._number_of_k_fold_splits}...")
train_data, val_data = self._get_train_and_val_data(dataset, e_cond, angle_cond, geo_cond, noise,
train_indexes, validation_indexes)
self._build_and_compile_new_model()
history = self.model.fit(x=train_data,
shuffle=True,
epochs=self._epochs,
verbose=verbose,
validation_data=val_data,
callbacks=callbacks
)
histories.append(history)
if self._save_best_model:
self.model.save_weights(f"{self._checkpoint_dir}/VAE_fold_{i + 1}/model_weights")
print(f"Best model from fold {i + 1} was saved.")
# Remove all unnecessary data from previous fold.
del self.model
del train_data
del val_data
tf.keras.backend.clear_session()
gc.collect()
return histories
def _single_training(self, dataset: np.array, e_cond: np.array, angle_cond: np.array, geo_cond: np.array,
noise: np.ndarray, callbacks: List[Callback], verbose: bool = True) -> List[History]:
"""
Performs a single training.
A fraction of dataset (self._validation_split) is used as a validation data.
Args:
dataset: A matrix representing showers. Shape =
(number of samples, ORIGINAL_DIM = N_CELLS_Z * N_CELLS_R * N_CELLS_PHI).
e_cond: A matrix representing an energy for each sample. Shape = (number of samples, ).
angle_cond: A matrix representing an angle for each sample. Shape = (number of samples, ).
geo_cond: A matrix representing a geometry of the detector for each sample. Shape = (number of samples, 2).
noise: A matrix representing an additional noise needed to perform a reparametrization trick.
callbacks: A list of callback forwarded to the fitting function.
verbose: A boolean which says there the training should be performed in a verbose mode or not.
Returns: A one-element list of `History` objects.`History.history` attribute is a record of training loss
values and metrics values at successive epochs, as well as validation loss values and validation metrics
values (if applicable).
"""
dataset_size, _ = dataset.shape
permutation = np.random.permutation(dataset_size)
split = int(dataset_size * self._validation_split)
train_indexes, validation_indexes = permutation[split:], permutation[:split]
train_data, val_data = self._get_train_and_val_data(dataset, e_cond, angle_cond, geo_cond, noise, train_indexes,
validation_indexes)
history = self.model.fit(x=train_data,
shuffle=True,
epochs=self._epochs,
verbose=verbose,
validation_data=val_data,
callbacks=callbacks
)
if self._save_best_model:
self.model.save_weights(f"{self._checkpoint_dir}/VAE_best/model_weights")
print("Best model was saved.")
return [history]
def train(self, dataset: np.array, e_cond: np.array, angle_cond: np.array, geo_cond: np.array,
verbose: bool = True) -> List[History]:
"""
For a given input data trains and validates the model.
If the numer of K-fold splits > 1 then it runs K-fold cross validation, otherwise it runs a single training
which uses (self._validation_split * 100) % of dataset as a validation data.
Args:
dataset: A matrix representing showers. Shape =
(number of samples, ORIGINAL_DIM = N_CELLS_Z * N_CELLS_R * N_CELLS_PHI).
e_cond: A matrix representing an energy for each sample. Shape = (number of samples, ).
angle_cond: A matrix representing an angle for each sample. Shape = (number of samples, ).
geo_cond: A matrix representing a geometry of the detector for each sample. Shape = (number of samples, 2).
verbose: A boolean which says there the training should be performed in a verbose mode or not.
Returns: A list of `History` objects.`History.history` attribute is a record of training loss values and
metrics values at successive epochs, as well as validation loss values and validation metrics values (if
applicable).
"""
callbacks = self._manufacture_callbacks()
noise = np.random.normal(0, 1, size=(dataset.shape[0], self.latent_dim))
if self._number_of_k_fold_splits > 1:
return self._k_fold_training(dataset, e_cond, angle_cond, geo_cond, noise, callbacks, verbose)
else:
return self._single_training(dataset, e_cond, angle_cond, geo_cond, noise, callbacks, verbose)
@@ -0,0 +1,87 @@
"""
** generate **
generate showers using a saved VAE model
"""
import argparse
import numpy as np
import tensorflow as tf
from tensorflow.python.data import Dataset
from core.constants import GLOBAL_CHECKPOINT_DIR, GEN_DIR, BATCH_SIZE_PER_REPLICA, MAX_GPU_MEMORY_ALLOCATION, GPU_IDS
from utils.gpu_limiter import GPULimiter
from utils.preprocess import get_condition_arrays
def parse_args():
argument_parser = argparse.ArgumentParser()
argument_parser.add_argument("--geometry", type=str, default="")
argument_parser.add_argument("--energy", type=int, default="")
argument_parser.add_argument("--angle", type=int, default="")
argument_parser.add_argument("--events", type=int, default=10000)
argument_parser.add_argument("--epoch", type=int, default=None)
argument_parser.add_argument("--study-name", type=str, default="default_study_name")
argument_parser.add_argument("--max-gpu-memory-allocation", type=int, default=MAX_GPU_MEMORY_ALLOCATION)
argument_parser.add_argument("--gpu-ids", type=str, default=GPU_IDS)
args = argument_parser.parse_args()
return args
# main function
def main():
# 0. Parse arguments.
args = parse_args()
energy = args.energy
angle = args.angle
geometry = args.geometry
events = args.events
epoch = args.epoch
study_name = args.study_name
max_gpu_memory_allocation = args.max_gpu_memory_allocation
gpu_ids = args.gpu_ids
# 1. Set GPU memory limits.
GPULimiter(_gpu_ids=gpu_ids, _max_gpu_memory_allocation=max_gpu_memory_allocation)()
# 2. Load a saved model.
# Create a handler and build model.
# This import must be local because otherwise it is impossible to call GPULimiter.
from core.model import VAEHandler
vae = VAEHandler()
# Load the saved weights
weights_dir = f"VAE_epoch_{epoch:03}" if epoch is not None else "VAE_best"
vae.model.load_weights(f"{GLOBAL_CHECKPOINT_DIR}/{study_name}/{weights_dir}/model_weights").expect_partial()
# The generator is defined as the decoder part only
generator = vae.model.decoder
# 3. Prepare data. Get condition values. Sample from the prior (normal distribution) in d dimension (d=latent_dim,
# latent space dimension). Gather them into tuples. Wrap data in Dataset objects. The batch size must now be set
# on the Dataset objects. Disable AutoShard.
e_cond, angle_cond, geo_cond = get_condition_arrays(geometry, energy, events)
z_r = np.random.normal(loc=0, scale=1, size=(events, vae.latent_dim))
data = ((z_r, e_cond, angle_cond, geo_cond),)
data = Dataset.from_tensor_slices(data)
batch_size = BATCH_SIZE_PER_REPLICA
data = data.batch(batch_size)
options = tf.data.Options()
options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF
data = data.with_options(options)
# 4. Generate showers using the VAE model.
generated_events = generator.predict(data) * (energy * 1000)
# 5. Save the generated showers.
np.save(f"{GEN_DIR}/VAE_Generated_Geo_{geometry}_E_{energy}_Angle_{angle}.npy", generated_events)
if __name__ == "__main__":
exit(main())
@@ -1,134 +0,0 @@
"""
** model **
defines the VAE model class
"""
# Setup
import keras
from tensorflow.keras.layers import Input, Dense, Lambda, Layer, Multiply, Add, concatenate
from tensorflow.keras.layers import BatchNormalization
from tensorflow.keras.models import Model
from tensorflow.keras import backend as K
from tensorflow.keras import metrics
# VAE model class
class VAE:
def __init__(self, **kwargs):
self.original_dim = kwargs.get('original_dim')
self.latent_dim = kwargs.get('latent_dim')
self.batch_size = kwargs.get('batch_size')
self.intermediate_dim1 = kwargs.get('intermediate_dim1')
self.intermediate_dim2 = kwargs.get('intermediate_dim2')
self.intermediate_dim3 = kwargs.get('intermediate_dim3')
self.intermediate_dim4 = kwargs.get('intermediate_dim4')
self.epsilon_std = kwargs.get('epsilon_std')
self.mu = kwargs.get('mu')
self.lr = kwargs.get('lr')
self.epochs = kwargs.get('epochs')
self.activ = kwargs.get('activ')
self.outActiv = kwargs.get('outActiv')
self.validation_split = kwargs.get('validation_split')
self.wReco = kwargs.get('wReco')
self.wkl = kwargs.get('wkl')
self.optimizer = kwargs.get('optimizer')
self.ki = kwargs.get('ki')
self.bi = kwargs.get('bi')
self.checkpoint_dir = kwargs.get('checkpoint_dir')
self.earlyStop = kwargs.get('earlyStop')
# KL divergence computation
class KLDivergenceLayer(Layer):
def __init__(self, *args, **kwargs):
self.is_placeholder = True
super(KLDivergenceLayer, self).__init__(*args, **kwargs)
def call(self, inputs):
mu, log_var = inputs
kl_batch = -self.wkl * K.sum(1 + log_var - K.square(mu) - K.exp(log_var), axis=-1)
self.add_loss(K.mean(kl_batch), inputs=inputs)
return inputs
# Build the encoder
xIn = Input((input_dim,))
eCond = Input(shape=(1,))
angleCond = Input(shape=(1,))
GeoCond = Input(shape=(2,))
mergedInput = concatenate([xIn, eCond, angleCond, GeoCond],)
h1 = Dense(self.intermediate_dim1, activation=self.activ,
kernel_initializer=self.ki, bias_initializer=self.bi)(mergedInput)
h1 = BatchNormalization()(h1)
h2 = Dense(self.intermediate_dim2, activation=self.activ,
kernel_initializer=self.ki, bias_initializer=self.bi)(h1)
h2 = BatchNormalization()(h2)
h3 = Dense(self.intermediate_dim3, activation=self.activ,
kernel_initializer=self.ki, bias_initializer=self.bi)(h2)
h3 = BatchNormalization()(h3)
h4 = Dense(self.intermediate_dim4, activation=self.activ,
kernel_initializer=self.ki, bias_initializer=self.bi)(h3)
h = BatchNormalization()(h4)
z_mu = Dense(self.latent_dim,)(h)
z_log_var = Dense(self.latent_dim,)(h)
# compute the KL divergence
z_mu, z_log_var = KLDivergenceLayer()([z_mu, z_log_var])
# Reparameterization trick
z_sigma = Lambda(lambda t: K.exp(.5*t))(z_log_var)
eps = Input(tensor=K.random_normal(shape=(K.shape(xIn)[0], self.latent_dim)))
z_eps = Multiply()([z_sigma, eps])
z = Add()([z_mu, z_eps])
zCond = concatenate([z,eCond,angleCond,GeoCond],)
# This defines the encoder which takes noise and input and outputs the latent variable z
self.encoder = Model(inputs=[xIn,eCond,angleCond,GeoCond,eps], outputs=zCond)
# Build the decoder / Generator
decoL4 = Dense(self.intermediate_dim4, input_dim=(self.latent_dim+4),
activation=self.activ, kernel_initializer=self.ki, bias_initializer=self.bi)
decoL4_BN = BatchNormalization()
decoL3 = Dense(self.intermediate_dim3, input_dim=self.intermediate_dim4,
activation=self.activ, kernel_initializer=self.ki, bias_initializer=self.bi)
decoL3_BN = BatchNormalization()
decoL2 = Dense(self.intermediate_dim2, input_dim=self.intermediate_dim3,
activation=self.activ, kernel_initializer=self.ki, bias_initializer=self.bi)
decoL2_BN = BatchNormalization()
decoL1 = Dense(self.intermediate_dim1, input_dim=self.intermediate_dim2,
activation=self.activ, kernel_initializer=self.ki, bias_initializer=self.bi)
decoL1_BN = BatchNormalization()
x_reco = Dense(self.original_dim, activation=self.outActiv)
zDecoInput = Input(shape=(latent_dim+4,))
x_recoDeco = x_reco((((decoL1_BN(decoL1(decoL2_BN(decoL2(decoL3_BN(decoL3(decoL4_BN(decoL4(zDecoInput))))))))))))
# This defines the decoder which takes an input of size latent dimension + condition size dimension and outputs the reconstructed input version
self.decoder = Model(inputs=[zDecoInput], outputs=[x_recoDeco])
# This defines the reconstruction loss of the VAE model
def reconstructionLoss(G4_Event, VAE_Event):
return K.mean(self.wReco*K.sum(metrics.binary_crossentropy(G4_Event, VAE_Event)))
# This defines the VAE model (encoder and decoder)
self.vae = Model(inputs=[xIn,eCond,angleCond,GeoCond,eps], outputs=[self.decoder(self.encoder([xIn, eCond,angleCond,GeoCond,eps]))])
self.vae.compile(optimizer=self.optimizer, loss=[reconstructionLoss] )
# Training function
def train(self, trainSet, eCond, angleCond, GeoCond):
# If the early stopping flag is on then stop the training when a monitored metric (validation) has stopped improving after (patience) number of epochs
if(self.earlyStop):
from tensorflow.keras.callbacks import EarlyStopping
cP = EarlyStopping(monitor='val_loss', min_delta=0.01, patience=5,verbose=1)
# If the early stopping flag is off then run the training for the number of epochs and save the model every (period) epochs
else:
cP = keras.callbacks.ModelCheckpoint('%s/VAE-{epoch:02d}.h5'%self.checkpoint_dir, monitor='val_loss',
verbose=0, save_best_only=False, save_weights_only=False, mode='auto',
period=100)
noise = np.random.normal(0,1, size = (trainSet.shape[0],latent_dim))
history = self.vae.fit([trainSet, eCond, angleCond, GeoCond,noise], [trainSet],
shuffle=True,
epochs=self.epochs,
verbose=1,
validation_split=self.validation_split,
batch_size=self.batch_size,
callbacks=[cP]
)
return history
# Encode function uses only the encoder to generate the latent representation of an input
def encode(self, dataSet):
return self.encoder.predict(dataSet, batch_size=self.batch_size)
# Generate function uses only the decoder to generate new showers using the z_sample which is a vector of 10D Gaussians in addition to
def generate(self, z_sample):
return self.decoder.predict([z_sample])
# Encode function
def predict(self, dataSet):
return self.vae.predict(dataSet, batch_size=self.batch_size)
# Encode function
def evaluate(self, dataSet):
return self.vae.evaluate(dataSet, batch_size=self.batch_size)
@@ -0,0 +1,12 @@
tensorflow==2.9.1
numpy==1.23.1
h5py==3.7.0
matplotlib==3.5.2
optuna==2.10.1
mysqlclient==2.1.1
pymysql==1.0.2
scikit-learn==1.1.1
scipy==1.8.1
wandb==0.13.1
tf2onnx==1.12.0
onnxruntime==1.12.1
@@ -0,0 +1,16 @@
"""
** setup **
creates necessary folders
"""
import os
from core.constants import INIT_DIR, GLOBAL_CHECKPOINT_DIR, CONV_DIR, VALID_DIR, GEN_DIR
for folder in [INIT_DIR, # Directory to load the full simulation dataset
GLOBAL_CHECKPOINT_DIR, # Directory to save VAE checkpoints
CONV_DIR, # Directory to save model after conversion to a format that can be used in C++
VALID_DIR, # Directory to save validation plots
GEN_DIR, # Directory to save VAE generated showers
]:
os.system(f"mkdir {folder}")
@@ -1,87 +1,52 @@
"""
** train **
- defines data loading parameters and calls the data preprocessing function
- defines the model parameters and instantiates the VAE model
- performs the training
"""
from argparse import ArgumentParser
# 1. Data loading/preprocessing
from utils import *
# Directory where the HDF5 files are saved
init_dir = './detector_'
# Number of calorimeter layers
nCells_z = 45
# Segmentation in the r,phi direction
nCells_r = 18
nCells_phi = 50
# Total number of readout cells (represents the number of nodes in the input/output layers of the model)
original_dim = nCells_z*nCells_r*nCells_phi
# Minimum and maximum primary particle energy to consider for training in GeV units
min_energy = 1
max_energy = 1024
# Minimum and maximum primary particle angle to consider for training in degrees units
min_angle = 50
max_angle = 90
# The preprocess function reads the data and performs preprocessing and encoding for the values of energy, angle and geometry
energies_Train,condE_Train,condAngle_Train,condGeo_Train = preprocess(init_dir,original_dim,min_angle,max_angle,min_energy,max_energy)
from core.constants import GPU_IDS, MAX_GPU_MEMORY_ALLOCATION, GLOBAL_CHECKPOINT_DIR
from utils.gpu_limiter import GPULimiter
from utils.preprocess import preprocess
# 2. Model architecture
import model
# Instantiate a VAE model and define all the parameters
vae = model.VAE(batch_size=100 ,
original_dim=original_dim,
intermediate_dim1=100,
intermediate_dim2=50,
intermediate_dim3=20,
intermediate_dim4=10+4,
latent_dim=10,
epsilon_std=1.,
mu=0,
epochs=10000,
lr=0.001,
activ=tf.keras.layers.LeakyReLU(),
outActiv='sigmoid',
validation_split=0.05,
wReco=original_dim,
wkl=0.5,
optimizer=optimizers.Adam(),
ki='RandomNormal',
bi='Zeros',
earlyStop=False,
checkpoint_dir = "."
)
# 3. Model training
history = vae.train(energies_Train,
condE_Train,
condAngle_Train,
condGeo_Train
)
def parse_args():
argument_parser = ArgumentParser()
argument_parser.add_argument("--max-gpu-memory-allocation", type=int, default=MAX_GPU_MEMORY_ALLOCATION)
argument_parser.add_argument("--gpu-ids", type=str, default=GPU_IDS)
argument_parser.add_argument("--study-name", type=str, default="default_study_name")
args = argument_parser.parse_args()
return args
# 4. Save the model (ony the decoder part) after traing
self.vae.decoder.save("decoder.h5")
# 5. Convert the model to ONNX format
import keras2onnx
import tensorflow
# Create the Keras model and convert itinto an ONNX model
kerasModel = tensorflow.keras.models.load_model("decoder.h5")
onnxModel = keras2onnx.convert_keras(kerasModel,"name")
# Save the ONNX model. Generator.onnx can then be used to perform the inference in the example
keras2onnx.save_model(onnxModel,"Generator.onnx")
def main():
# 0. Parse arguments.
args = parse_args()
max_gpu_memory_allocation = args.max_gpu_memory_allocation
gpu_ids = args.gpu_ids
study_name = args.study_name
checkpoint_dir = f"{GLOBAL_CHECKPOINT_DIR}/{study_name}"
"""
# In order to convert the model into a format that can be used with the LWTNN library
# 1. After training :
# serialize model to JSON
json_model = self.vae.decoder.to_json()
with open("decoder.json", "w") as json_file:
json_file.write(json_model)
# serialize weights to HDF5
self.vae.decoder.save_weights("decoder.h5")
# 2. Externally, after building the LWTNN code available at https://github.com/lwtnn/lwtnn
# 2.1 Run the kerasfunc2json python script (available in lwtnn/ converters/) to generate a template file of your functional model input variables by calling:
# $ kerasfunc2json.py decoder.json decoder.h5 > inputs.json
# 2.2 Run again kerasfunc2json script to get your output file that would be used for the inference in the example
# $ kerasfunc2json.py decoder.json decoder.h5 inputs.json > Generator.json
"""
# 1. Set GPU memory limits.
GPULimiter(_gpu_ids=gpu_ids, _max_gpu_memory_allocation=max_gpu_memory_allocation)()
# 2. Data loading/preprocessing
# The preprocess function reads the data and performs preprocessing and encoding for the values of energy,
# angle and geometry
energies_train, cond_e_train, cond_angle_train, cond_geo_train = preprocess()
# 3. Manufacture model handler.
# This import must be local because otherwise it is impossible to call GPULimiter.
from core.model import VAEHandler
vae = VAEHandler(_wandb_project_name=study_name, _wandb_tags=["single training"], _checkpoint_dir=checkpoint_dir)
# 4. Train model.
histories = vae.train(energies_train,
cond_e_train,
cond_angle_train,
cond_geo_train
)
# Note : One history object can be used to plot the loss evaluation as function of the epochs. Remember that the
# function returns a list of those objects. Each of them represents a different fold of cross validation.
if __name__ == "__main__":
exit(main())
@@ -0,0 +1,48 @@
from argparse import ArgumentParser
from core.constants import MAX_GPU_MEMORY_ALLOCATION, GPU_IDS
from utils.gpu_limiter import GPULimiter
from utils.optimizer import OptimizerType
# Hyperparemeters to be optimized.
discrete_parameters = {"nb_hidden_layers": (1, 6), "latent_dim": (15, 100)}
continuous_parameters = {"learning_rate": (0.0001, 0.005)}
categorical_parameters = {"optimizer_type": [OptimizerType.ADAM, OptimizerType.RMSPROP]}
def parse_args():
argument_parser = ArgumentParser()
argument_parser.add_argument("--study-name", type=str, default="default_study_name")
argument_parser.add_argument("--storage", type=str)
argument_parser.add_argument("--max-gpu-memory-allocation", type=int, default=MAX_GPU_MEMORY_ALLOCATION)
argument_parser.add_argument("--gpu-ids", type=str, default=GPU_IDS)
args = argument_parser.parse_args()
return args
def main():
# 0. Parse arguments.
args = parse_args()
study_name = args.study_name
storage = args.storage
max_gpu_memory_allocation = args.max_gpu_memory_allocation
gpu_ids = args.gpu_ids
# 1. Set GPU memory limits.
GPULimiter(_gpu_ids=gpu_ids, _max_gpu_memory_allocation=max_gpu_memory_allocation)()
# 2. Manufacture hyperparameter tuner.
# This import must be local because otherwise it is impossible to call GPULimiter.
from utils.hyperparameter_tuner import HyperparameterTuner
hyperparameter_tuner = HyperparameterTuner(discrete_parameters, continuous_parameters, categorical_parameters,
storage, study_name)
# 3. Run main tuning function.
hyperparameter_tuner.tune()
# Watch out! This script neither deletes the study in DB nor deletes the database itself. If you are using
# parallelized optimization, then you should care about deleting study in the database by yourself.
if __name__ == "__main__":
exit(main())
@@ -1,57 +0,0 @@
"""
** utils **
defines the data loading and preprocessing function
"""
# Setup
import h5py
import numpy as np
# preprocess function returns the array of the shower energies and the condition arrays
"""
- init_dir: the name of the directory which contains the HDF5 files
- size_1DVec: represents the size of the input and output layer of the VAE which corresponds to the total number of readout cells
- min_energy,max_energy: minimum and maximum primary particle energy to consider for training in GeV units
- min_angle and max_angle: minimum and maximum primary particle angle to consider for training in degrees units
"""
def preprocess(init_dir,size_1DVec,min_angle,max_angle,min_energy,max_energy):
energies_Train = []
condE_Train = []
condAngle_Train = []
condGeo_Train = []
# This example is trained using 2 detector geometries
for geo in [ 'SiW' , 'SciPb' ]:
dirGeo = init_dir + geo + '/'
energyParticle=min_energy
# loop over the energies in powers of 2
while(energyParticle<=max_energy):
# loop over the angles in a step of 10
for angleParticle in range(min_angle,max_angle+10,10):
fName = 'Energy_%s_Angle_%s.hdf5' %(energyParticle,angleParticle)
fName = dirGeo + fName
# read the HDF5 file
h5 = h5py.File(fName,'r')
# get the key value of the group from the HDF5 file
GroupKey = 'Grp_Angle_%s_E_%s'%(angleParticle,energyParticle)
# get all key values of one group
listKeys = list( h5[GroupKey].keys() )
# loop over the events
for ckey in listKeys:
# scale the energy of each cell to the energy of the primary particle (in MeV units)
energyArray = np.array(h5[GroupKey][ckey])/(energyParticle*1000)
energies_Train.append( energyArray.reshape(size_1DVec) )
# build the energy and angle condition vectors
condE_Train.append( [energyParticle/mamax_energyxE]*len(listKeys) )
condAngle_Train.append( [angleParticle/max_angle]*len(listKeys) )
# build the geometry condition vector (1 hot encoding vector)
if( geo == 'SiW' ):
condGeo_Train.append( [[0,1]]*len(listKeys) )
else:
condGeo_Train.append( [[1,0]]*len(listKeys) )
energyParticle*=2
# return numpy arrays
energies_Train = np.array(energies_Train)
condE_Train = np.concatenate(condE_Train)
condAngle_Train = np.concatenate(condAngle_Train)
condGeo_Train = np.concatenate(condGeo_Train)
return energies_Train,condE_Train,condAngle_Train,condGeo_Train
@@ -0,0 +1,36 @@
import os
from dataclasses import dataclass
import tensorflow as tf
@dataclass
class GPULimiter:
"""
Class responsible to set the limits of possible GPU usage by TensorFlow. Currently, the limiter creates one
instance of logical device per physical device. This can be changed in a future.
Attributes:
_gpu_ids: A string representing visible devices for the process. Identifiers of physical GPUs should
be separated by commas (no spaces).
_max_gpu_memory_allocation: An integer specifying limit of allocated memory per logical device.
"""
_gpu_ids: str
_max_gpu_memory_allocation: int
def __call__(self):
os.environ["CUDA_VISIBLE_DEVICES"] = f"{self._gpu_ids}"
gpus = tf.config.list_physical_devices('GPU')
if gpus:
# Restrict TensorFlow to only allocate max_gpu_memory_allocation*1024 MB of memory on one of the GPUs
try:
for gpu in gpus:
tf.config.set_logical_device_configuration(
gpu,
[tf.config.LogicalDeviceConfiguration(memory_limit=1024 * self._max_gpu_memory_allocation)])
logical_gpus = tf.config.list_logical_devices('GPU')
print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
except RuntimeError as e:
# Virtual devices must be set before GPUs have been initialized
print(e)
@@ -0,0 +1,216 @@
from dataclasses import dataclass
from typing import Tuple, Dict, Any, List
import numpy as np
from optuna import Trial, create_study, get_all_study_summaries, load_study
from optuna.pruners import MedianPruner
from optuna.samplers import TPESampler
from optuna.trial import TrialState
from core.constants import LEARNING_RATE, BATCH_SIZE_PER_REPLICA, ACTIVATION, OUT_ACTIVATION, \
OPTIMIZER_TYPE, KERNEL_INITIALIZER, BIAS_INITIALIZER, N_TRIALS, LATENT_DIM, \
INTERMEDIATE_DIMS, MAX_HIDDEN_LAYER_DIM, GLOBAL_CHECKPOINT_DIR
from core.model import VAEHandler
from utils.preprocess import preprocess
@dataclass
class HyperparameterTuner:
"""Tuner which looks for the best hyperparameters of a Variational Autoencoder specified in model.py.
Currently, supported hyperparameters are: dimension of latent space, number of hidden layers, learning rate,
activation function, activation function after the final layer, optimizer type, kernel initializer,
bias initializer, batch size.
Attributes:
_discrete_parameters: A dictionary of hyperparameters taking discrete values in the range [low, high].
_continuous_parameters: A dictionary of hyperparameters taking continuous values in the range [low, high].
_categorical_parameters: A dictionary of hyperparameters taking values specified by the list of them.
_storage: A string representing URL to a database required for a distributed training
_study_name: A string, a name of study.
"""
_discrete_parameters: Dict[str, Tuple[int, int]]
_continuous_parameters: Dict[str, Tuple[float, float]]
_categorical_parameters: Dict[str, List[Any]]
_storage: str = None
_study_name: str = None
def _check_hyperparameters(self):
available_hyperparameters = ["latent_dim", "nb_hidden_layers", "learning_rate", "activation", "out_activation",
"optimizer_type", "kernel_initializer", "bias_initializer",
"batch_size_per_replica"]
hyperparameters_to_be_optimized = list(self._discrete_parameters.keys()) + list(
self._continuous_parameters.keys()) + list(self._categorical_parameters.keys())
for hyperparameter_name in hyperparameters_to_be_optimized:
if hyperparameter_name not in available_hyperparameters:
raise Exception(f"Unknown hyperparameter: {hyperparameter_name}")
def __post_init__(self):
self._check_hyperparameters()
self._energies_train, self._cond_e_train, self._cond_angle_train, self._cond_geo_train = preprocess()
if self._storage is not None and self._study_name is not None:
# Parallel optimization
study_summaries = get_all_study_summaries(self._storage)
if any(self._study_name == study_summary.study_name for study_summary in study_summaries):
# The study is already created in the database. Load it.
self._study = load_study(self._study_name, self._storage)
else:
# The study does not exist in the database. Create a new one.
self._study = create_study(storage=self._storage, sampler=TPESampler(), pruner=MedianPruner(),
study_name=self._study_name, direction="minimize")
else:
# Single optimization
self._study = create_study(sampler=TPESampler(), pruner=MedianPruner(), direction="minimize")
def _create_model_handler(self, trial: Trial) -> VAEHandler:
"""For a given trail builds the model.
Optuna suggests parameters like dimensions of particular layers of the model, learning rate, optimizer, etc.
Args:
trial: Optuna's trial
Returns:
Variational Autoencoder (VAE)
"""
# Discrete parameters
if "latent_dim" in self._discrete_parameters.keys():
latent_dim = trial.suggest_int(name="latent_dim",
low=self._discrete_parameters["latent_dim"][0],
high=self._discrete_parameters["latent_dim"][1])
else:
latent_dim = LATENT_DIM
if "nb_hidden_layers" in self._discrete_parameters.keys():
nb_hidden_layers = trial.suggest_int(name="nb_hidden_layers",
low=self._discrete_parameters["nb_hidden_layers"][0],
high=self._discrete_parameters["nb_hidden_layers"][1])
all_possible = np.arange(start=latent_dim + 5, stop=MAX_HIDDEN_LAYER_DIM)
chunks = np.array_split(all_possible, nb_hidden_layers)
ranges = [(chunk[0], chunk[-1]) for chunk in chunks]
ranges = reversed(ranges)
# Cast from np.int to int allows to become JSON serializable.
intermediate_dims = [trial.suggest_int(name=f"intermediate_dim_{i}", low=int(low), high=int(high)) for
i, (low, high)
in enumerate(ranges)]
else:
intermediate_dims = INTERMEDIATE_DIMS
if "batch_size_per_replica" in self._discrete_parameters.keys():
batch_size_per_replica = trial.suggest_int(name="batch_size_per_replica",
low=self._discrete_parameters["batch_size_per_replica"][0],
high=self._discrete_parameters["batch_size_per_replica"][1])
else:
batch_size_per_replica = BATCH_SIZE_PER_REPLICA
# Continuous parameters
if "learning_rate" in self._continuous_parameters.keys():
learning_rate = trial.suggest_float(name="learning_rate",
low=self._continuous_parameters["learning_rate"][0],
high=self._continuous_parameters["learning_rate"][1])
else:
learning_rate = LEARNING_RATE
# Categorical parameters
if "activation" in self._categorical_parameters.keys():
activation = trial.suggest_categorical(name="activation",
choices=self._categorical_parameters["activation"])
else:
activation = ACTIVATION
if "out_activation" in self._categorical_parameters.keys():
out_activation = trial.suggest_categorical(name="out_activation",
choices=self._categorical_parameters["out_activation"])
else:
out_activation = OUT_ACTIVATION
if "optimizer_type" in self._categorical_parameters.keys():
optimizer_type = trial.suggest_categorical(name="optimizer_type",
choices=self._categorical_parameters["optimizer_type"])
else:
optimizer_type = OPTIMIZER_TYPE
if "kernel_initializer" in self._categorical_parameters.keys():
kernel_initializer = trial.suggest_categorical(name="kernel_initializer",
choices=self._categorical_parameters["kernel_initializer"])
else:
kernel_initializer = KERNEL_INITIALIZER
if "bias_initializer" in self._categorical_parameters.keys():
bias_initializer = trial.suggest_categorical(name="bias_initializer",
choices=self._categorical_parameters["bias_initializer"])
else:
bias_initializer = BIAS_INITIALIZER
checkpoint_dir = f"{GLOBAL_CHECKPOINT_DIR}/{self._study_name}/trial_{trial.number:03d}"
return VAEHandler(_wandb_project_name=self._study_name,
_wandb_tags=["hyperparameter tuning", f"trial {trial.number}"],
_batch_size_per_replica=batch_size_per_replica,
_intermediate_dims=intermediate_dims,
latent_dim=latent_dim,
_learning_rate=learning_rate,
_activation=activation,
_out_activation=out_activation,
_optimizer_type=optimizer_type,
_kernel_initializer=kernel_initializer,
_bias_initializer=bias_initializer,
_checkpoint_dir=checkpoint_dir,
_early_stop=True,
_save_model_every_epoch=False,
_save_best_model=True,
)
def _objective(self, trial: Trial) -> float:
"""For a given trial trains the model and returns an average validation loss.
Args:
trial: Optuna's trial
Returns: One float numer which is a validation loss. It can be either calculated as an average of k trainings
performed in cross validation mode or is one number obtained from validation on unseen before, some fraction
of the dataset.
"""
# Generate the trial model.
model_handler = self._create_model_handler(trial)
# Train the model.
verbose = True
histories = model_handler.train(self._energies_train, self._cond_e_train, self._cond_angle_train,
self._cond_geo_train, verbose)
# Return validation loss (currently it is treated as an objective goal). Notice that we take into account the
# best model according to the validation loss.
final_validation_losses = [np.min(history.history["val_loss"]) for history in histories]
avg_validation_loss = np.mean(final_validation_losses).item()
return avg_validation_loss
def tune(self) -> None:
"""Main tuning function.
Based on a given study, tunes the model and prints detailed information about the best trial (value of the
objective function and adjusted parameters).
"""
self._study.optimize(func=self._objective, n_trials=N_TRIALS, gc_after_trial=True)
pruned_trials = self._study.get_trials(deepcopy=False, states=(TrialState.PRUNED,))
complete_trials = self._study.get_trials(deepcopy=False, states=(TrialState.COMPLETE,))
print("Study statistics: ")
print(" Number of finished trials: ", len(self._study.trials))
print(" Number of pruned trials: ", len(pruned_trials))
print(" Number of complete trials: ", len(complete_trials))
print("Best trial:")
trial = self._study.best_trial
print(" Value: ", trial.value)
print(" Params: ")
for key, value in trial.params.items():
print(f" {key}: {value}")
@@ -0,0 +1,216 @@
from dataclasses import dataclass
from enum import Enum
import numpy as np
from core.constants import N_CELLS_Z, N_CELLS_R, SIZE_Z, SIZE_R
@dataclass
class Observable:
""" An abstract class defining interface of all observables.
Do not use this class directly.
Attributes:
_input: A numpy array with shape = (NE, R, PHI, Z), where NE stays for number of events.
"""
_input: np.ndarray
class ProfileType(Enum):
""" Enum class of various profile types.
"""
LONGITUDINAL = 0
LATERAL = 1
@dataclass
class Profile(Observable):
""" An abstract class describing behaviour of LongitudinalProfile and LateralProfile.
Do not use this class directly. Use LongitudinalProfile or LateralProfile instead.
"""
def calc_profile(self) -> np.ndarray:
pass
def calc_first_moment(self) -> np.ndarray:
pass
def calc_second_moment(self) -> np.ndarray:
pass
@dataclass
class LongitudinalProfile(Profile):
""" A class defining observables related to LongitudinalProfile.
Attributes:
_energies_per_event: A numpy array with shape = (NE, Z) where NE stays for a number of events. An
element [i, j] is a sum of energies detected in all cells located in a jth layer for an ith event.
_total_energy_per_event: A numpy array with shape = (NE, ). An element [i] is a sum of energies detected in all
cells for an ith event.
_w: A numpy array = [0, 1, ..., Z - 1] which represents weights used in computation of first and second moment.
"""
def __post_init__(self):
self._energies_per_event = np.sum(self._input, axis=(1, 2))
self._total_energy_per_event = np.sum(self._energies_per_event, axis=1)
self._w = np.arange(N_CELLS_Z)
def calc_profile(self) -> np.ndarray:
""" Calculates a longitudinal profile.
A longitudinal profile for a given layer l (l = 0, ..., Z - 1) is defined as:
sum_{i = 0}^{NE - 1} energy_per_event[i, l].
Returns:
A numpy array of longitudinal profiles for each layer with a shape = (Z, ).
"""
return np.sum(self._energies_per_event, axis=0)
def calc_first_moment(self) -> np.ndarray:
""" Calculates a first moment of profile.
A first moment of a longitudinal profile for a given event e (e = 0, ..., NE - 1) is defined as:
FM[e] = alpha * (sum_{i = 0}^{Z - 1} energies_per_event[e, i] * w[i]) / total_energy_per_event[e], where
w = [0, 1, 2, ..., Z - 1],
alpha = SIZE_Z defined in core/constants.py.
Returns:
A numpy array of first moments of longitudinal profiles for each event with a shape = (NE, ).
"""
return SIZE_Z * np.dot(self._energies_per_event, self._w) / self._total_energy_per_event
def calc_second_moment(self) -> np.ndarray:
""" Calculates a second moment of a longitudinal profile.
A second moment of a longitudinal profile for a given event e (e = 0, ..., NE - 1) is defined as:
SM[e] = (sum_{i = 0}^{Z - 1} (w[i] - alpha - FM[e])^2 * energies_per_event[e, i]) total_energy_per_event[e],
where
w = [0, 1, 2, ..., Z - 1],
alpha = SIZE_Z defined in ochre/constants.py
Returns:
A numpy array of second moments of longitudinal profiles for each event with a shape = (NE, ).
"""
first_moment = self.calc_first_moment()
first_moment = np.expand_dims(first_moment, axis=1)
w = np.expand_dims(self._w, axis=0)
# w has now a shape = [1, Z] and first moment has a shape = [NE, 1]. There is a broadcasting in the line
# below how that one create an array with a shape = [NE, Z]
return np.sum(np.multiply(np.power(w * SIZE_Z - first_moment, 2), self._energies_per_event),
axis=1) / self._total_energy_per_event
@dataclass
class LateralProfile(Profile):
""" A class defining observables related to LateralProfile.
Attributes:
_energies_per_event: A numpy array with shape = (NE, R) where NE stays for a number of events. An
element [i, j] is a sum of energies detected in all cells located in a jth layer for an ith event.
_total_energy_per_event: A numpy array with shape = (NE, ). An element [i] is a sum of energies detected in all
cells for an ith event.
_w: A numpy array = [0, 1, ..., R - 1] which represents weights used in computation of first and second moment.
"""
def __post_init__(self):
self._energies_per_event = np.sum(self._input, axis=(2, 3))
self._total_energy_per_event = np.sum(self._energies_per_event, axis=1)
self._w = np.arange(N_CELLS_R)
def calc_profile(self) -> np.ndarray:
""" Calculates a lateral profile.
A lateral profile for a given layer l (l = 0, ..., R - 1) is defined as:
sum_{i = 0}^{NE - 1} energy_per_event[i, l].
Returns:
A numpy array of longitudinal profiles for each layer with a shape = (R, ).
"""
return np.sum(self._energies_per_event, axis=0)
def calc_first_moment(self) -> np.ndarray:
""" Calculates a first moment of profile.
A first moment of a lateral profile for a given event e (e = 0, ..., NE - 1) is defined as:
FM[e] = alpha * (sum_{i = 0}^{R - 1} energies_per_event[e, i] * w[i]) / total_energy_per_event[e], where
w = [0, 1, 2, ..., R - 1],
alpha = SIZE_R defined in core/constants.py.
Returns:
A numpy array of first moments of lateral profiles for each event with a shape = (NE, ).
"""
return SIZE_R * np.dot(self._energies_per_event, self._w) / self._total_energy_per_event
def calc_second_moment(self) -> np.ndarray:
""" Calculates a second moment of a lateral profile.
A second moment of a lateral profile for a given event e (e = 0, ..., NE - 1) is defined as:
SM[e] = (sum_{i = 0}^{R - 1} (w[i] - alpha - FM[e])^2 * energies_per_event[e, i]) total_energy_per_event[e],
where
w = [0, 1, 2, ..., R - 1],
alpha = SIZE_R defined in ochre/constants.py
Returns:
A numpy array of second moments of lateral profiles for each event with a shape = (NE, ).
"""
first_moment = self.calc_first_moment()
first_moment = np.expand_dims(first_moment, axis=1)
w = np.expand_dims(self._w, axis=0)
# w has now a shape = [1, R] and first moment has a shape = [NE, 1]. There is a broadcasting in the line
# below how that one create an array with a shape = [NE, R]
return np.sum(np.multiply(np.power(w * SIZE_R - first_moment, 2), self._energies_per_event),
axis=1) / self._total_energy_per_event
@dataclass
class Energy(Observable):
""" A class defining observables total energy per event and cell energy.
"""
def calc_total_energy(self):
""" Calculates total energy detected in an event.
Total energy for a given event e (e = 0, ..., NE - 1) is defined as a sum of energies detected in all cells
for this event.
Returns:
A numpy array of total energy values with shape = (NE, ).
"""
return np.sum(self._input, axis=(1, 2, 3))
def calc_cell_energy(self):
""" Calculates cell energy.
Cell energy for a given event (e = 0, ..., NE - 1) is defined by an array with shape (R * PHI * Z) storing
values of energy in particular cells.
Returns:
A numpy array of cell energy values with shape = (NE * R * PHI * Z, ).
"""
return np.copy(self._input).reshape(-1)
def calc_energy_per_layer(self):
""" Calculates total energy detected in a particular layer.
Energy per layer for a given event (e = 0, ..., NE - 1) is defined by an array with shape (Z, ) storing
values of total energy detected in a particular layer
Returns:
A numpy array of cell energy values with shape = (NE, Z).
"""
return np.sum(self._input, axis=(1, 2))
@@ -0,0 +1,55 @@
from enum import IntEnum
from tensorflow.keras.optimizers import Optimizer, Adadelta, Adagrad, Adam, Adamax, Ftrl, SGD, Nadam, RMSprop
class OptimizerType(IntEnum):
""" Enum class of various optimizer types.
This class must be IntEnum to be JSON serializable. This feature is important because, when Optuna's study is
saved in a relational DB, all objects must be JSON serializable.
"""
SGD = 0
RMSPROP = 1
ADAM = 2
ADADELTA = 3
ADAGRAD = 4
ADAMAX = 5
NADAM = 6
FTRL = 7
class OptimizerFactory:
"""Factory of optimizer like Stochastic Gradient Descent, RMSProp, Adam, etc.
"""
@staticmethod
def create_optimizer(optimizer_type: OptimizerType, learning_rate: float) -> Optimizer:
"""For a given type and a learning rate creates an instance of optimizer.
Args:
optimizer_type: a type of optimizer
learning_rate: a learning rate that should be passed to an optimizer
Returns:
An instance of optimizer.
"""
if optimizer_type == OptimizerType.SGD:
return SGD(learning_rate)
elif optimizer_type == OptimizerType.RMSPROP:
return RMSprop(learning_rate)
elif optimizer_type == OptimizerType.ADAM:
return Adam(learning_rate)
elif optimizer_type == OptimizerType.ADADELTA:
return Adadelta(learning_rate)
elif optimizer_type == OptimizerType.ADAGRAD:
return Adagrad(learning_rate)
elif optimizer_type == OptimizerType.ADAMAX:
return Adamax(learning_rate)
elif optimizer_type == OptimizerType.NADAM:
return Nadam(learning_rate)
else:
# i.e. optimizer_type == OptimizerType.FTRL
return Ftrl(learning_rate)
@@ -0,0 +1,518 @@
from dataclasses import dataclass
from typing import Tuple
import numpy as np
from matplotlib import pyplot as plt
from scipy.optimize import curve_fit
from core.constants import N_CELLS_Z, N_CELLS_R, VALID_DIR, SIZE_Z, SIZE_R, HISTOGRAM_TYPE, FULL_SIM_HISTOGRAM_COLOR, \
ML_SIM_HISTOGRAM_COLOR, FULL_SIM_GAUSSIAN_COLOR, ML_SIM_GAUSSIAN_COLOR
from utils.observables import LongitudinalProfile, ProfileType, Profile, Energy
plt.rcParams.update({"font.size": 22})
@dataclass
class Plotter:
""" An abstract class defining interface of all plotters.
Do not use this class directly. Use ProfilePlotter or EnergyPlotter instead.
Attributes:
_particle_energy: An integer which is energy of the primary particle in GeV units.
_particle_angle: An integer which is an angle of the primary particle in degrees.
_geometry: A string which is a name of the calorimeter geometry (e.g. SiW, SciPb).
"""
_particle_energy: int
_particle_angle: int
_geometry: str
def plot_and_save(self):
pass
def _gaussian(x: np.ndarray, a: float, mu: float, sigma: float) -> np.ndarray:
""" Computes a value of a Gaussian.
Args:
x: An argument of a function.
a: A scaling parameter.
mu: A mean.
sigma: A variance.
Returns:
A value of a function for given arguments.
"""
return a * np.exp(-((x - mu)**2 / (2 * sigma**2)))
def _best_fit(data: np.ndarray,
bins: np.ndarray,
hist: bool = False) -> Tuple[np.ndarray, np.ndarray]:
""" Finds estimated shape of a Gaussian using Use non-linear least squares.
Args:
data: A numpy array with values of observables from multiple events.
bins: A numpy array specifying histogram bins.
hist: If histogram is calculated. Then data is the frequencies.
Returns:
A tuple of two lists. Xs and Ys of predicted curve.
"""
# Calculate histogram.
if not hist:
hist, _ = np.histogram(data, bins)
else:
hist = data
# Choose only those bins which are nonzero. Nonzero() return a tuple of arrays. In this case it has a length = 1,
# hence we are interested in its first element.
indices = hist.nonzero()[0]
# Based on previously chosen nonzero bin, calculate position of xs and ys_bar (true values) which will be used in
# fitting procedure. Len(bins) == len(hist + 1), so we choose middles of bins as xs.
bins_middles = (bins[:-1] + bins[1:]) / 2
xs = bins_middles[indices]
ys_bar = hist[indices]
# Set initial parameters for curve fitter.
a0 = np.max(ys_bar)
mu0 = np.mean(xs)
sigma0 = np.var(xs)
# Fit a Gaussian to the prepared data.
(a, mu, sigma), _ = curve_fit(f=_gaussian,
xdata=xs,
ydata=ys_bar,
p0=[a0, mu0, sigma0],
method="trf",
maxfev=1000)
# Calculate values of an approximation in given points and return values.
ys = _gaussian(xs, a, mu, sigma)
return xs, ys
@dataclass
class ProfilePlotter(Plotter):
""" Plotter responsible for preparing plots of profiles and their first and second moments.
Attributes:
_full_simulation: A numpy array representing a profile of data generated by Geant4.
_ml_simulation: A numpy array representing a profile of data generated by ML model.
_plot_gaussian: A boolean. Decides whether first and second moment should be plotted as a histogram or
a fitted gaussian.
_profile_type: An enum. A profile can be either lateral or longitudinal.
"""
_full_simulation: Profile
_ml_simulation: Profile
_plot_gaussian: bool = False
def __post_init__(self):
# Check if profiles are either both longitudinal or lateral.
full_simulation_type = type(self._full_simulation)
ml_generation_type = type(self._ml_simulation)
assert full_simulation_type == ml_generation_type, "Both profiles within a ProfilePlotter must be the same " \
"type."
# Set an attribute with profile type.
if full_simulation_type == LongitudinalProfile:
self._profile_type = ProfileType.LONGITUDINAL
else:
self._profile_type = ProfileType.LATERAL
def _plot_and_save_customizable_histogram(
self,
full_simulation: np.ndarray,
ml_simulation: np.ndarray,
bins: np.ndarray,
xlabel: str,
observable_name: str,
plot_profile: bool = False,
y_log_scale: bool = False) -> None:
""" Prepares and saves a histogram for a given pair of observables.
Args:
full_simulation: A numpy array of observables coming from full simulation.
ml_simulation: A numpy array of observables coming from ML simulation.
bins: A numpy array specifying histogram bins.
xlabel: A string. Name of x-axis on the plot.
observable_name: A string. Name of plotted observable.
plot_profile: A boolean. If set to True, full_simulation and ml_simulation are histogram weights while x is
defined by the number of layers. This means that in order to plot histogram (and gaussian), one first
need to create a data repeating each layer or R index appropriate number of times. Should be set to True
only while plotting profiles not first or second moments.
y_log_scale: A boolean. Used log scale on y-axis is set to True.
Returns:
None.
"""
fig, axes = plt.subplots(2,
1,
figsize=(15, 10),
clear=True,
sharex="all")
# Plot histograms.
if plot_profile:
# We already have the bins (layers) and freqencies (energies),
# therefore directly plotting a step plot + lines instead of a hist plot.
axes[0].step(bins[:-1],
full_simulation,
label="FullSim",
color=FULL_SIM_HISTOGRAM_COLOR)
axes[0].step(bins[:-1],
ml_simulation,
label="MLSim",
color=ML_SIM_HISTOGRAM_COLOR)
axes[0].vlines(x=bins[0],
ymin=0,
ymax=full_simulation[0],
color=FULL_SIM_HISTOGRAM_COLOR)
axes[0].vlines(x=bins[-2],
ymin=0,
ymax=full_simulation[-1],
color=FULL_SIM_HISTOGRAM_COLOR)
axes[0].vlines(x=bins[0],
ymin=0,
ymax=ml_simulation[0],
color=ML_SIM_HISTOGRAM_COLOR)
axes[0].vlines(x=bins[-2],
ymin=0,
ymax=ml_simulation[-1],
color=ML_SIM_HISTOGRAM_COLOR)
axes[0].set_ylim(0, None)
# For using it later for the ratios.
energy_full_sim, energy_ml_sim = full_simulation, ml_simulation
else:
energy_full_sim, _, _ = axes[0].hist(
x=full_simulation,
bins=bins,
label="FullSim",
histtype=HISTOGRAM_TYPE,
color=FULL_SIM_HISTOGRAM_COLOR)
energy_ml_sim, _, _ = axes[0].hist(x=ml_simulation,
bins=bins,
label="MLSim",
histtype=HISTOGRAM_TYPE,
color=ML_SIM_HISTOGRAM_COLOR)
# Plot Gaussians if needed.
if self._plot_gaussian:
if plot_profile:
(xs_full_sim, ys_full_sim) = _best_fit(full_simulation,
bins,
hist=True)
(xs_ml_sim, ys_ml_sim) = _best_fit(ml_simulation,
bins,
hist=True)
else:
(xs_full_sim, ys_full_sim) = _best_fit(full_simulation, bins)
(xs_ml_sim, ys_ml_sim) = _best_fit(ml_simulation, bins)
axes[0].plot(xs_full_sim,
ys_full_sim,
color=FULL_SIM_GAUSSIAN_COLOR,
label="FullSim")
axes[0].plot(xs_ml_sim,
ys_ml_sim,
color=ML_SIM_GAUSSIAN_COLOR,
label="MLSim")
if y_log_scale:
axes[0].set_yscale("log")
axes[0].legend(loc="best")
axes[0].set_xlabel(xlabel)
axes[0].set_ylabel("Energy [Mev]")
axes[0].set_title(
f" $e^-$, {self._particle_energy} [GeV], {self._particle_angle}$^{{\circ}}$, {self._geometry}"
)
# Calculate ratios.
ratio = np.divide(energy_ml_sim,
energy_full_sim,
out=np.ones_like(energy_ml_sim),
where=(energy_full_sim != 0))
# Since len(bins) == 1 + data, we calculate middles of bins as xs.
bins_middles = (bins[:-1] + bins[1:]) / 2
axes[1].plot(bins_middles, ratio, "-o")
axes[1].set_xlabel(xlabel)
axes[1].set_ylabel("MLSim/FullSim")
axes[1].axhline(y=1, color="black")
plt.savefig(
f"{VALID_DIR}/{observable_name}_Geo_{self._geometry}_E_{self._particle_energy}_"
+ f"Angle_{self._particle_angle}.png")
plt.clf()
def _plot_profile(self) -> None:
""" Plots profile of an observable.
Returns:
None.
"""
full_simulation_profile = self._full_simulation.calc_profile()
ml_simulation_profile = self._ml_simulation.calc_profile()
if self._profile_type == ProfileType.LONGITUDINAL:
# matplotlib will include the right-limit for the last bar,
# hence extending by 1.
bins = np.linspace(0, N_CELLS_Z, N_CELLS_Z + 1)
observable_name = "LongProf"
xlabel = "Layer index"
else:
bins = np.linspace(0, N_CELLS_R, N_CELLS_R + 1)
observable_name = "LatProf"
xlabel = "R index"
self._plot_and_save_customizable_histogram(full_simulation_profile,
ml_simulation_profile,
bins,
xlabel,
observable_name,
plot_profile=True)
def _plot_first_moment(self) -> None:
""" Plots and saves a first moment of an observable's profile.
Returns:
None.
"""
full_simulation_first_moment = self._full_simulation.calc_first_moment(
)
ml_simulation_first_moment = self._ml_simulation.calc_first_moment()
if self._profile_type == ProfileType.LONGITUDINAL:
xlabel = "$<\lambda> [mm]$"
observable_name = "LongFirstMoment"
bins = np.linspace(0, 0.4 * N_CELLS_Z * SIZE_Z, 128)
else:
xlabel = "$<r> [mm]$"
observable_name = "LatFirstMoment"
bins = np.linspace(0, 0.75 * N_CELLS_R * SIZE_R, 128)
self._plot_and_save_customizable_histogram(
full_simulation_first_moment, ml_simulation_first_moment, bins,
xlabel, observable_name)
def _plot_second_moment(self) -> None:
""" Plots and saves a second moment of an observable's profile.
Returns:
None.
"""
full_simulation_second_moment = self._full_simulation.calc_second_moment(
)
ml_simulation_second_moment = self._ml_simulation.calc_second_moment()
if self._profile_type == ProfileType.LONGITUDINAL:
xlabel = "$<\lambda^{2}> [mm^{2}]$"
observable_name = "LongSecondMoment"
bins = np.linspace(0, pow(N_CELLS_Z * SIZE_Z, 2) / 35., 128)
else:
xlabel = "$<r^{2}> [mm^{2}]$"
observable_name = "LatSecondMoment"
bins = np.linspace(0, pow(N_CELLS_R * SIZE_R, 2) / 8., 128)
self._plot_and_save_customizable_histogram(
full_simulation_second_moment, ml_simulation_second_moment, bins,
xlabel, observable_name)
def plot_and_save(self) -> None:
""" Main plotting function.
Calls private methods and prints the information about progress.
Returns:
None.
"""
if self._profile_type == ProfileType.LONGITUDINAL:
profile_type_name = "longitudinal"
else:
profile_type_name = "lateral"
print(f"Plotting the {profile_type_name} profile...")
self._plot_profile()
print(f"Plotting the first moment of {profile_type_name} profile...")
self._plot_first_moment()
print(f"Plotting the second moment of {profile_type_name} profile...")
self._plot_second_moment()
@dataclass
class EnergyPlotter(Plotter):
""" Plotter responsible for preparing plots of profiles and their first and second moments.
Attributes:
_full_simulation: A numpy array representing a profile of data generated by Geant4.
_ml_simulation: A numpy array representing a profile of data generated by ML model.
"""
_full_simulation: Energy
_ml_simulation: Energy
def _plot_total_energy(self, y_log_scale=True) -> None:
""" Plots and saves a histogram with total energy detected in an event.
Args:
y_log_scale: A boolean. Used log scale on y-axis is set to True.
Returns:
None.
"""
full_simulation_total_energy = self._full_simulation.calc_total_energy(
)
ml_simulation_total_energy = self._ml_simulation.calc_total_energy()
plt.figure(figsize=(12, 8))
bins = np.linspace(
np.min(full_simulation_total_energy) -
np.min(full_simulation_total_energy) * 0.05,
np.max(full_simulation_total_energy) +
np.max(full_simulation_total_energy) * 0.05, 50)
plt.hist(x=full_simulation_total_energy,
histtype=HISTOGRAM_TYPE,
label="FullSim",
bins=bins,
color=FULL_SIM_HISTOGRAM_COLOR)
plt.hist(x=ml_simulation_total_energy,
histtype=HISTOGRAM_TYPE,
label="MLSim",
bins=bins,
color=ML_SIM_HISTOGRAM_COLOR)
plt.legend(loc="upper left")
if y_log_scale:
plt.yscale("log")
plt.xlabel("Energy [MeV]")
plt.ylabel("# events")
plt.title(
f" $e^-$, {self._particle_energy} [GeV], {self._particle_angle}$^{{\circ}}$, {self._geometry} "
)
plt.savefig(
f"{VALID_DIR}/E_tot_Geo_{self._geometry}_E_{self._particle_energy}_Angle_{self._particle_angle}.png"
)
plt.clf()
def _plot_cell_energy(self) -> None:
""" Plots and saves a histogram with number of detector's cells across whole
calorimeter with particular energy detected.
Returns:
None.
"""
full_simulation_cell_energy = self._full_simulation.calc_cell_energy()
ml_simulation_cell_energy = self._ml_simulation.calc_cell_energy()
log_full_simulation_cell_energy = np.log10(
full_simulation_cell_energy,
out=np.zeros_like(full_simulation_cell_energy),
where=(full_simulation_cell_energy != 0))
log_ml_simulation_cell_energy = np.log10(
ml_simulation_cell_energy,
out=np.zeros_like(ml_simulation_cell_energy),
where=(ml_simulation_cell_energy != 0))
plt.figure(figsize=(12, 8))
bins = np.linspace(-4, 1, 1000)
plt.hist(x=log_full_simulation_cell_energy,
bins=bins,
histtype=HISTOGRAM_TYPE,
label="FullSim",
color=FULL_SIM_HISTOGRAM_COLOR)
plt.hist(x=log_ml_simulation_cell_energy,
bins=bins,
histtype=HISTOGRAM_TYPE,
label="MLSim",
color=ML_SIM_HISTOGRAM_COLOR)
plt.xlabel("log10(E/MeV)")
plt.ylim(bottom=1)
plt.yscale("log")
plt.ylim(bottom=1)
plt.ylabel("# entries")
plt.title(
f" $e^-$, {self._particle_energy} [GeV], {self._particle_angle}$^{{\circ}}$, {self._geometry} "
)
plt.grid(True)
plt.legend(loc="upper left")
plt.savefig(
f"{VALID_DIR}/E_cell_Geo_{self._geometry}_E_{self._particle_energy}_Angle_{self._particle_angle}.png"
)
plt.clf()
def _plot_energy_per_layer(self):
""" Plots and saves N_CELLS_Z histograms with total energy detected in particular layers.
Returns:
None.
"""
full_simulation_energy_per_layer = self._full_simulation.calc_energy_per_layer(
)
ml_simulation_energy_per_layer = self._ml_simulation.calc_energy_per_layer(
)
number_of_plots_in_row = 9
number_of_plots_in_column = 5
bins = np.linspace(np.min(full_simulation_energy_per_layer - 10),
np.max(full_simulation_energy_per_layer + 10), 25)
fig, ax = plt.subplots(number_of_plots_in_column,
number_of_plots_in_row,
figsize=(20, 15),
sharex="all",
sharey="all",
constrained_layout=True)
for layer_nb in range(N_CELLS_Z):
i = layer_nb // number_of_plots_in_row
j = layer_nb % number_of_plots_in_row
ax[i][j].hist(full_simulation_energy_per_layer[:, layer_nb],
histtype=HISTOGRAM_TYPE,
label="FullSim",
bins=bins,
color=FULL_SIM_HISTOGRAM_COLOR)
ax[i][j].hist(ml_simulation_energy_per_layer[:, layer_nb],
histtype=HISTOGRAM_TYPE,
label="MLSim",
bins=bins,
color=ML_SIM_HISTOGRAM_COLOR)
ax[i][j].set_title(f"Layer {layer_nb}", fontsize=13)
ax[i][j].set_yscale("log")
ax[i][j].tick_params(axis='both', which='major', labelsize=10)
fig.supxlabel("Energy [MeV]", fontsize=14)
fig.supylabel("# entries", fontsize=14)
fig.suptitle(
f" $e^-$, {self._particle_energy} [GeV], {self._particle_angle}$^{{\circ}}$, {self._geometry} "
)
# Take legend from one plot and make it a global legend.
handles, labels = ax[0][0].get_legend_handles_labels()
fig.legend(handles, labels, bbox_to_anchor=(1.15, 0.5))
plt.savefig(
f"{VALID_DIR}/E_layer_Geo_{self._geometry}_E_{self._particle_energy}_Angle_{self._particle_angle}.png",
bbox_inches="tight")
plt.clf()
def plot_and_save(self):
""" Main plotting function.
Calls private methods and prints the information about progress.
Returns:
None.
"""
print("Plotting total energy...")
self._plot_total_energy()
print("Plotting cell energy...")
self._plot_cell_energy()
print("Plotting energy per layer...")
self._plot_energy_per_layer()
@@ -0,0 +1,83 @@
import h5py
import numpy as np
from core.constants import INIT_DIR, ORIGINAL_DIM, MAX_ENERGY, MAX_ANGLE, MIN_ANGLE, MIN_ENERGY
# preprocess function loads the data and returns the array of the shower energies and the condition arrays
def preprocess():
energies_train = []
cond_e_train = []
cond_angle_train = []
cond_geo_train = []
# This example is trained using 2 detector geometries
for geo in ["SiW", "SciPb"]:
dir_geo = INIT_DIR + geo + "/"
# loop over the angles in a step of 10
for angle_particle in range(MIN_ANGLE, MAX_ANGLE + 10, 10):
f_name = f"{geo}_angle_{angle_particle}.h5"
f_name = dir_geo + f_name
# read the HDF5 file
h5 = h5py.File(f_name, "r")
# loop over energies from min_energy to max_energy
energy_particle = MIN_ENERGY
while energy_particle <= MAX_ENERGY:
# scale the energy of each cell to the energy of the primary particle (in MeV units)
events = np.array(h5[f"{energy_particle}"]) / (energy_particle * 1000)
energies_train.append(events.reshape(len(events), ORIGINAL_DIM))
# build the energy and angle condition vectors
cond_e_train.append([energy_particle / MAX_ENERGY] * len(events))
cond_angle_train.append([angle_particle / MAX_ANGLE] * len(events))
# build the geometry condition vector (1 hot encoding vector)
if geo == "SiW":
cond_geo_train.append([[0, 1]] * len(events))
if geo == "SciPb":
cond_geo_train.append([[1, 0]] * len(events))
energy_particle *= 2
# return numpy arrays
energies_train = np.concatenate(energies_train)
cond_e_train = np.concatenate(cond_e_train)
cond_angle_train = np.concatenate(cond_angle_train)
cond_geo_train = np.concatenate(cond_geo_train)
return energies_train, cond_e_train, cond_angle_train, cond_geo_train
# get_condition_arrays function returns condition values from a single geometry, a single energy and angle of primary
# particles
"""
- geo : name of the calorimeter geometry (eg: SiW, SciPb)
- energy_particle : energy of the primary particle in GeV units
- nb_events : number of events
"""
def get_condition_arrays(geo, energy_particle, nb_events):
cond_e = [energy_particle / MAX_ENERGY] * nb_events
cond_angle = [energy_particle / MAX_ENERGY] * nb_events
if geo == "SiW":
cond_geo = [[0, 1]] * nb_events
else: # geo == "SciPb"
cond_geo = [[1, 0]] * nb_events
cond_e = np.array(cond_e)
cond_angle = np.array(cond_angle)
cond_geo = np.array(cond_geo)
return cond_e, cond_angle, cond_geo
# load_showers function loads events from a single geometry, a single energy and angle of primary particles
"""
- init_dir: the name of the directory which contains the HDF5 files
- geo : name of the calorimeter geometry (eg: SiW, SciPb)
- energy_particle : energy of the primary particle in GeV units
- angle_particle : angle of the primary particle in degrees
"""
def load_showers(init_dir, geo, energy_particle, angle_particle):
dir_geo = init_dir + geo + "/"
f_name = f"{geo}_angle_{angle_particle}.h5"
f_name = dir_geo + f_name
# read the HDF5 file
h5 = h5py.File(f_name, "r")
energies = np.array(h5[f"{energy_particle}"])
return energies
@@ -0,0 +1,62 @@
import argparse
import numpy as np
from core.constants import INIT_DIR, GEN_DIR, N_CELLS_PHI, N_CELLS_R, N_CELLS_Z
from utils.observables import LongitudinalProfile, LateralProfile, Energy
from utils.plotters import ProfilePlotter, EnergyPlotter
from utils.preprocess import load_showers
def parse_args():
p = argparse.ArgumentParser()
p.add_argument("--geometry", type=str, default="")
p.add_argument("--energy", type=int, default="")
p.add_argument("--angle", type=int, default="")
args = p.parse_args()
return args
# main function
def main():
# Parse commandline arguments
args = parse_args()
particle_energy = args.energy
particle_angle = args.angle
geometry = args.geometry
# 1. Full simulation data loading
# Load energy of showers from a single geometry, energy and angle
e_layer_g4 = load_showers(INIT_DIR, geometry, particle_energy,
particle_angle)
# 2. Fast simulation data loading, scaling to original energy range & reshaping
vae_energies = np.load(f"{GEN_DIR}/VAE_Generated_Geo_{geometry}_E_{particle_energy}_Angle_{particle_angle}.npy")
# Reshape the events into 3D
e_layer_vae = vae_energies.reshape((len(vae_energies), N_CELLS_R, N_CELLS_PHI, N_CELLS_Z))
print("Data has been loaded.")
# 3. Create observables from raw data.
full_sim_long = LongitudinalProfile(_input=e_layer_g4)
full_sim_lat = LateralProfile(_input=e_layer_g4)
full_sim_energy = Energy(_input=e_layer_g4)
ml_sim_long = LongitudinalProfile(_input=e_layer_vae)
ml_sim_lat = LateralProfile(_input=e_layer_vae)
ml_sim_energy = Energy(_input=e_layer_vae)
print("Created observables.")
# 4. Plot observables
longitudinal_profile_plotter = ProfilePlotter(particle_energy, particle_angle, geometry, full_sim_long, ml_sim_long,
_plot_gaussian=False)
lateral_profile_plotter = ProfilePlotter(particle_energy, particle_angle,
geometry, full_sim_lat, ml_sim_lat, _plot_gaussian=False)
energy_plotter = EnergyPlotter(particle_energy, particle_angle, geometry, full_sim_energy, ml_sim_energy)
longitudinal_profile_plotter.plot_and_save()
lateral_profile_plotter.plot_and_save()
energy_plotter.plot_and_save()
print("Done.")
if __name__ == "__main__":
exit(main())
@@ -0,0 +1,118 @@
/Par04/detector/setDetectorInnerRadius 80 cm
/Par04/detector/setDetectorLength 4 m
/Par04/detector/setNbOfLayers 90
/Par04/detector/setAbsorber 0 G4_W 1.4 mm true
/Par04/detector/setAbsorber 1 G4_Si 0.3 mm true
/Par04/mesh/setSizeOfRhoCells 2.325 mm
/Par04/mesh/setSizeOfZCells 3.4 mm
/Par04/mesh/setNbOfRhoCells 18
/Par04/mesh/setNbOfPhiCells 50
/Par04/mesh/setNbOfZCells 45
/Par04/detector/print
# Use default detector dimensions and initialize
/run/initialize
# If inference model is active, de-activate it because it needs configuration
/param/InActivateModel inferenceModel
# Use this open statement to create an OpenGL view:
/vis/open OGL 600x600-0+0
#
# Use this open statement to create a .prim file suitable for
# viewing in DAWN:
#/vis/open DAWNFILE
#
# Use this open statement to create a .heprep file suitable for
# viewing in HepRApp:
#/vis/open HepRepFile
#
# Use this open statement to create a .wrl file suitable for
# viewing in a VRML viewer:
#/vis/open VRML2FILE
#
# Disable auto refresh and quieten vis messages whilst scene and
# trajectories are established:
/vis/viewer/set/autoRefresh false
/vis/verbose errors
#
# Draw geometry:
/vis/drawVolume worlds
#
# Specify view angle:
/vis/viewer/set/viewpointThetaPhi 0 90 deg
/vis/viewer/set/targetPoint 0 800 0 mm
#
# Specify zoom value:
/vis/viewer/zoom 10
#
# Specify style (surface or wireframe):
#/vis/viewer/set/style wireframe
#
# Draw coordinate axes:
#/vis/scene/add/axes 0 0 0 1 m
#
# Draw smooth trajectories at end of event, showing trajectory points
# as markers 2 pixels wide:
#/vis/scene/add/trajectories smooth
/vis/scene/add/trajectories
/vis/modeling/trajectories/create/drawByCharge
/vis/modeling/trajectories/drawByCharge-0/default/setDrawStepPts true
/vis/modeling/trajectories/drawByCharge-0/default/setStepPtsSize 2
# (if too many tracks cause core dump => /tracking/storeTrajectory 0)
#
# Draw hits at end of event:
/vis/scene/add/hits
#
# To draw only gammas:
#/vis/filtering/trajectories/create/particleFilter
#/vis/filtering/trajectories/particleFilter-0/add gamma
#
# To invert the above, drawing all particles except gammas,
# keep the above two lines but also add:
#/vis/filtering/trajectories/particleFilter-0/invert true
#
# Many other options are available with /vis/modeling and /vis/filtering.
# For example, to select colour by particle ID:
#/vis/modeling/trajectories/create/drawByParticleID
#/vis/modeling/trajectories/drawByParticleID-0/set e- blue
#
# Create an attribute filter to draw only particles with certain (high) momentum
/vis/filtering/trajectories/create/attributeFilter
# Select attribute "IMag"
/vis/filtering/trajectories/attributeFilter-0/setAttribute IMag
# Select trajectories with 25 MeV <= IMag < 1000 GeV
/vis/filtering/trajectories/attributeFilter-0/addInterval 25 MeV 1000 GeV
#
# To superimpose all of the events from a given run:
/vis/scene/endOfEventAction accumulate
#
# Re-establish auto refreshing and verbosity:
/vis/viewer/set/autoRefresh true
/vis/verbose warnings
#
# For file-based drivers, use this to create an empty detector view:
#/vis/viewer/flush
/vis/viewer/set/background 1 1 1
# Fast Simulation
# Inference Setup
## dimension of the latent vector (encoded vector in a Variational Autoencoder model)
/Par04/inference/setSizeLatentVector 10
## size of the condition vector (energy, angle and geometry)
/Par04/inference/setSizeConditionVector 4
## path to the model which is set to download by cmake
/Par04/inference/setModelPathName MLModels/Generator.pt
/Par04/inference/setInferenceLibrary TORCH
## set mesh size for inference == mesh size of a full sim that
## was used for training; it coincides with readout mesh size
/Par04/inference/setSizeOfRhoCells 2.325 mm
/Par04/inference/setSizeOfZCells 3.4 mm
/Par04/inference/setNbOfRhoCells 18
/Par04/inference/setNbOfPhiCells 50
/Par04/inference/setNbOfZCells 45
## Dynamic readout mesh from particle direction needs to be the first fast sim model!
/param/ActivateModel defineMesh
## ML fast sim, configured with the inference setup /Par04/inference
/param/ActivateModel inferenceModel