Par04/src/Par04InferenceSetup.cc

0001 //
0002 // ********************************************************************
0003 // * License and Disclaimer                                           *
0004 // *                                                                  *
0005 // * The  Geant4 software  is  copyright of the Copyright Holders  of *
0006 // * the Geant4 Collaboration.  It is provided  under  the terms  and *
0007 // * conditions of the Geant4 Software License,  included in the file *
0008 // * LICENSE and available at  http://cern.ch/geant4/license .  These *
0009 // * include a list of copyright holders.                             *
0010 // *                                                                  *
0011 // * Neither the authors of this software system, nor their employing *
0012 // * institutes,nor the agencies providing financial support for this *
0013 // * work  make  any representation or  warranty, express or implied, *
0014 // * regarding  this  software system or assume any liability for its *
0015 // * use.  Please see the license in the file  LICENSE  and URL above *
0016 // * for the full disclaimer and the limitation of liability.         *
0017 // *                                                                  *
0018 // * This  code  implementation is the result of  the  scientific and *
0019 // * technical work of the GEANT4 collaboration.                      *
0020 // * By using,  copying,  modifying or  distributing the software (or *
0021 // * any work based  on the software)  you  agree  to acknowledge its *
0022 // * use  in  resulting  scientific  publications,  and indicate your *
0023 // * acceptance of all terms of the Geant4 Software license.          *
0024 // ********************************************************************
0025 //
0026 /// \file Par04InferenceSetup.cc
0027 /// \brief Implementation of the Par04InferenceSetup class
0028
0029 #ifdef USE_INFERENCE
0030 #  include "Par04InferenceSetup.hh"
0031
0032 #  include "Par04InferenceInterface.hh"  // for Par04InferenceInterface
0033 #  include "Par04InferenceMessenger.hh"  // for Par04InferenceMessenger
0034 #  ifdef USE_INFERENCE_ONNX
0035 #    include "Par04OnnxInference.hh"  // for Par04OnnxInference
0036 #  endif
0037 #  ifdef USE_INFERENCE_LWTNN
0038 #    include "Par04LwtnnInference.hh"  // for Par04LwtnnInference
0039 #  endif
0040 #  ifdef USE_INFERENCE_TORCH
0041 #    include "Par04TorchInference.hh"  // for Par04TorchInference
0042 #  endif
0043 #  include "CLHEP/Random/RandGauss.h"  // for RandGauss
0044
0045 #  include "G4RotationMatrix.hh"  // for G4RotationMatrix
0046
0047 #  include <CLHEP/Units/SystemOfUnits.h>  // for pi, GeV, deg
0048 #  include <CLHEP/Vector/Rotation.h>  // for HepRotation
0049 #  include <CLHEP/Vector/ThreeVector.h>  // for Hep3Vector
0050 #  include <G4Exception.hh>  // for G4Exception
0051 #  include <G4ExceptionSeverity.hh>  // for FatalException
0052 #  include <G4ThreeVector.hh>  // for G4ThreeVector
0053 #  include <algorithm>  // for max, copy
0054 #  include <cmath>  // for cos, sin
0055 #  include <string>  // for char_traits, basic_string
0056
0057 #  include <ext/alloc_traits.h>  // for __alloc_traits<>::value_type
0058
0059 //....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
0060
0061 Par04InferenceSetup::Par04InferenceSetup() : fInferenceMessenger(new Par04InferenceMessenger(this))
0062 {}
0063
0064 //....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
0065
0066 Par04InferenceSetup::~Par04InferenceSetup() {}
0067
0068 //....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
0069
0070 G4bool Par04InferenceSetup::IfTrigger(G4double aEnergy)
0071 {
0072   /// Energy of electrons used in training dataset
0073   if (aEnergy > 1 * CLHEP::GeV || aEnergy < 1024 * CLHEP::GeV) return true;
0074   return false;
0075 }
0076
0077 //....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
0078
0079 void Par04InferenceSetup::SetInferenceLibrary(G4String aName)
0080 {
0081   fInferenceLibrary = aName;
0082
0083 #  ifdef USE_INFERENCE_ONNX
0084   if (fInferenceLibrary == "ONNX")
0085     fInferenceInterface = std::unique_ptr<Par04InferenceInterface>(new Par04OnnxInference(
0086       fModelPathName, fProfileFlag, fOptimizationFlag, fIntraOpNumThreads, fCudaFlag, cuda_keys,
0087       cuda_values, fModelSavePath, fProfilingOutputSavePath));
0088 #  endif
0089 #  ifdef USE_INFERENCE_LWTNN
0090   if (fInferenceLibrary == "LWTNN")
0091     fInferenceInterface =
0092       std::unique_ptr<Par04InferenceInterface>(new Par04LwtnnInference(fModelPathName));
0093 #  endif
0094 #  ifdef USE_INFERENCE_TORCH
0095   if (fInferenceLibrary == "TORCH")
0096     fInferenceInterface =
0097       std::unique_ptr<Par04InferenceInterface>(new Par04TorchInference(fModelPathName));
0098 #  endif
0099
0100   CheckInferenceLibrary();
0101 }
0102
0103 //....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
0104
0105 void Par04InferenceSetup::CheckInferenceLibrary()
0106 {
0107   G4String msg = "Please choose inference library from available libraries (";
0108 #  ifdef USE_INFERENCE_ONNX
0109   msg += "ONNX,";
0110 #  endif
0111 #  ifdef USE_INFERENCE_LWTNN
0112   msg += "LWTNN,";
0113 #  endif
0114 #  ifdef USE_INFERENCE_TORCH
0115   msg += "TORCH";
0116 #  endif
0117   if (fInferenceInterface == nullptr)
0118     G4Exception("Par04InferenceSetup::CheckInferenceLibrary()", "InvalidSetup", FatalException,
0119                 (msg + "). Current name: " + fInferenceLibrary).c_str());
0120 }
0121
0122 //....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
0123
0124 void Par04InferenceSetup::GetEnergies(std::vector<G4double>& aEnergies, G4double aInitialEnergy,
0125                                       G4float aTheta, G4float aPhi)
0126 {
0127   // First check if inference library was set correctly
0128   CheckInferenceLibrary();
0129   // size represents the size of the output vector
0130   int size = fMeshNumber.x() * fMeshNumber.y() * fMeshNumber.z();
0131   std::vector<G4float> genVector;
0132
0133   if (fModelType == "VAE")
0134   {
0135     genVector.assign(fSizeLatentVector + fSizeConditionVector, 0);
0136
0137   // randomly sample from a gaussian distribution in the latent space
0138   for (int i = 0; i < fSizeLatentVector; ++i) {
0139     genVector[i] = CLHEP::RandGauss::shoot(0., 1.);
0140   }
0141
0142   // Vector of condition
0143   // this is application specific it depdens on what the model was condition on
0144   // and it depends on how the condition values were encoded at the training
0145   // time in this example the energy of each particle is normlaized to the
0146   // highest energy in the considered range (1GeV-500GeV) the angle is also is
0147   // normlaized to the highest angle in the considered range (0-90 in dergrees)
0148   // the model in this example was trained on two detector geometries PBW04
0149   // and SiW  a one hot encoding vector is used to represent the geometry with
0150   // [0,1] for PBW04 and [1,0] for SiW
0151   // 1. energy
0152   genVector[fSizeLatentVector] = aInitialEnergy / fMaxEnergy;
0153   // 2. angle
0154   genVector[fSizeLatentVector + 1] = (aTheta / (CLHEP::deg)) / fMaxAngle;
0155   // 3. geometry
0156   genVector[fSizeLatentVector + 2] = 0;
0157   genVector[fSizeLatentVector + 3] = 1;
0158   } else if (fModelType == "CaloDiT-2")
0159   {
0160     // fSizeLatentVector & fSizeConditionVector are ignored for CaloDiT-2
0161     // Conditions (dim) are energy (1), phi (1), theta (1) and geo (5)
0162     // The energy range here is 1 GeV - 1TeV, phi goes from 0 to 2pi,
0163     // and theta goes from 0.87 to 2.27.
0164     // And, geo is one-hot encoding describing the 4 geometries the model
0165     // is trained on.
0166     // Order of the geo condition is Par04SiW (this one), Par04SciPb, ODD, FCCeeCLD
0167     // As CaloDiT-2 is trained on these 4 detectors, it can be quickly adapted to
0168     // any new detector (see CaloDiT-2 readme for adaptation) of your choice. Thus
0169     // reusing the knowledge from these previous detectors.
0170     // To use the adapted model, make the following changes for inference:
0171     // genVector[3] = 0.0; (turning OFF Par04SiW)
0172     // genVector[7] = 1.0; (turning ON a new detector)
0173     genVector.assign(8, 0);
0174
0175     genVector[0] = aInitialEnergy / 1000;  // convert to GeV
0176     genVector[1] = aPhi;
0177     genVector[2] = aTheta;
0178     genVector[3] = 1.0;  //Par04SiW
0179   }
0180   // Run the inference
0181   fInferenceInterface->RunInference(genVector, aEnergies, size);
0182
0183   // After the inference rescale back to the initial energy
0184
0185   if (fModelType == "VAE")
0186   // For VAE, energies of cells were normalized to the energy of the particle
0187   {
0188   for (int i = 0; i < size; ++i) {
0189     aEnergies[i] = aEnergies[i] * aInitialEnergy;
0190     }
0191   } else if (fModelType == "CaloDiT-2")
0192   // For CaloDiT-2, energies were scaled by a factor of 1000
0193   {
0194     for (int i = 0; i < size; ++i){
0195       aEnergies[i] = aEnergies[i] * 1000;
0196     }
0197   }
0198 }
0199
0200 //....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
0201
0202 void Par04InferenceSetup::GetPositions(std::vector<G4ThreeVector>& aPositions, G4ThreeVector pos0,
0203                                        G4ThreeVector direction)
0204 {
0205   aPositions.resize(fMeshNumber.x() * fMeshNumber.y() * fMeshNumber.z());
0206
0207   // Calculate rotation matrix along the particle momentum direction
0208   // It will rotate the shower axes to match the incoming particle direction
0209   G4RotationMatrix rotMatrix = G4RotationMatrix();
0210   double particleTheta = direction.theta();
0211   double particlePhi = direction.phi();
0212   rotMatrix.rotateZ(-particlePhi);
0213   rotMatrix.rotateY(-particleTheta);
0214   G4RotationMatrix rotMatrixInv = CLHEP::inverseOf(rotMatrix);
0215
0216   int cpt = 0;
0217   for (G4int iCellR = 0; iCellR < fMeshNumber.x(); iCellR++) {
0218     for (G4int iCellPhi = 0; iCellPhi < fMeshNumber.y(); iCellPhi++) {
0219       for (G4int iCellZ = 0; iCellZ < fMeshNumber.z(); iCellZ++) {
0220         aPositions[cpt] =
0221           pos0
0222           + rotMatrixInv
0223               * G4ThreeVector(
0224                 (iCellR + 0.5) * fMeshSize.x()
0225                   * std::cos((iCellPhi + 0.5) * 2 * CLHEP::pi / fMeshNumber.y() - CLHEP::pi),
0226                 (iCellR + 0.5) * fMeshSize.x()
0227                   * std::sin((iCellPhi + 0.5) * 2 * CLHEP::pi / fMeshNumber.y() - CLHEP::pi),
0228                 (iCellZ + 0.5) * fMeshSize.z());
0229         cpt++;
0230       }
0231     }
0232   }
0233 }
0234
0235 #endif