Import Geant4 11.1.0 source tree
This commit is contained in:
@@ -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
|
||||
Reference in New Issue
Block a user