algorithms/fardetectors/FarDetectorMLReconstruction.h

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2023, Simon Gardner
0003
0004 #pragma once
0005
0006 #include <TMVA/MethodBase.h>
0007 #include <TMVA/Reader.h>
0008 #include <algorithms/algorithm.h>
0009 #include <edm4eic/TrackCollection.h>
0010 #include <edm4eic/TrackParametersCollection.h>
0011 #include <edm4eic/TrajectoryCollection.h>
0012 // Event Model related classes
0013 #include <edm4hep/MCParticleCollection.h>
0014 #include <mutex>
0015 #include <string>
0016 #include <string_view>
0017
0018 #include "FarDetectorMLReconstructionConfig.h"
0019 #include "algorithms/interfaces/WithPodConfig.h"
0020
0021 namespace eicrecon {
0022
0023   enum FarDetectorMLNNIndexIn{PosY,PosZ,DirX,DirY};
0024   enum FarDetectorMLNNIndexOut{MomX,MomY,MomZ};
0025
0026   using FarDetectorMLReconstructionAlgorithm = algorithms::Algorithm<
0027     algorithms::Input<
0028       edm4eic::TrackParametersCollection,
0029       edm4hep::MCParticleCollection
0030     >,
0031     algorithms::Output<
0032       edm4eic::TrajectoryCollection,
0033       edm4eic::TrackParametersCollection,
0034       edm4eic::TrackCollection
0035     >
0036   >;
0037
0038   class FarDetectorMLReconstruction
0039   : public FarDetectorMLReconstructionAlgorithm,
0040     public WithPodConfig<FarDetectorMLReconstructionConfig> {
0041
0042   public:
0043       FarDetectorMLReconstruction(std::string_view name)
0044         : FarDetectorMLReconstructionAlgorithm{name,
0045                               {"TrackParameters","BeamElectrons"},
0046                               {"Trajectory","TrackParameters","Track"},
0047                               "Reconstruct track parameters using ML method."} {}
0048
0049
0050       /** One time initialization **/
0051       void init();
0052
0053       /** Event by event processing **/
0054       void process(const Input&, const Output&);
0055
0056       //----- Define constants here ------
0057
0058   private:
0059       TMVA::Reader*     m_reader{nullptr};
0060       TMVA::MethodBase* m_method{nullptr};
0061       float m_beamE{10.0};
0062       std::once_flag m_initBeamE;
0063       float nnInput[4]  = {0.0,0.0,0.0,0.0};
0064
0065   };
0066
0067 } // eicrecon