algorithms/fardetectors/FarDetectorTransportationPostML.h

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2024 - 2025 Simon Gardner
0003
0004 #pragma once
0005
0006 #include <algorithms/algorithm.h>
0007 #include <edm4eic/ReconstructedParticleCollection.h>
0008 #include <edm4eic/MCRecoParticleAssociationCollection.h>
0009 #include <edm4eic/MCRecoTrackParticleAssociationCollection.h>
0010 #include <edm4eic/TensorCollection.h>
0011 #include <edm4hep/MCParticleCollection.h>
0012 #include <mutex>
0013 #include <optional>
0014 #include <string>
0015 #include <string_view>
0016
0017 #include "algorithms/fardetectors/FarDetectorTransportationPostMLConfig.h"
0018 #include "algorithms/interfaces/WithPodConfig.h"
0019
0020 namespace eicrecon {
0021
0022 using FarDetectorTransportationPostMLAlgorithm = algorithms::Algorithm<
0023     algorithms::Input<edm4eic::TensorCollection,
0024                       std::optional<edm4eic::MCRecoTrackParticleAssociationCollection>,
0025                       std::optional<edm4hep::MCParticleCollection>>,
0026     algorithms::Output<edm4eic::ReconstructedParticleCollection,
0027                        edm4eic::MCRecoParticleAssociationCollection>>;
0028
0029 class FarDetectorTransportationPostML
0030     : public FarDetectorTransportationPostMLAlgorithm,
0031       public WithPodConfig<FarDetectorTransportationPostMLConfig> {
0032
0033 public:
0034   FarDetectorTransportationPostML(std::string_view name)
0035       : FarDetectorTransportationPostMLAlgorithm{
0036             name,
0037             {"inputPredictionsTensor", "trackAssociations", "beamElectrons"},
0038             {"outputParticles", "outputAssociations"},
0039             "Convert ML output tensor into reconstructed electron"} {}
0040
0041   void init() final;
0042   void process(const Input&, const Output&) const final;
0043
0044 private:
0045   double m_mass         = 0.000511; // Default to electron mass in GeV
0046   float m_charge        = -1.0;     // Default to electron charge
0047   mutable float m_beamE = 10.0;
0048   mutable std::once_flag m_initBeamE;
0049 };
0050
0051 } // namespace eicrecon