algorithms/fardetectors/FarDetectorMLReconstruction.h

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