Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2026-05-23 07:39:22

0001 // Copyright 2023-2025, Simon Gardner
0002 // Subject to the terms in the LICENSE file found in the top-level directory.
0003 //
0004 //
0005 
0006 #include <Evaluator/DD4hepUnits.h>
0007 #include <JANA/JApplication.h>
0008 #include <JANA/JApplicationFwd.h>
0009 #include <JANA/Utils/JTypeInfo.h>
0010 #include <edm4eic/EDM4eicVersion.h>
0011 #include <edm4eic/MCRecoTrackParticleAssociation.h>
0012 #if EDM4EIC_BUILD_VERSION >= EDM4EIC_VERSION(8, 7, 0)
0013 #include <edm4eic/MCRecoTrackParticleLinkCollection.h>
0014 #endif
0015 #include <edm4eic/Track.h>
0016 #include <edm4eic/TrackerHit.h>
0017 #include <edm4eic/unit_system.h>
0018 #include <edm4hep/SimTrackerHit.h>
0019 #include <fmt/format.h> // IWYU pragma: keep
0020 #include <podio/detail/Link.h>
0021 #include <cmath>
0022 #include <cstddef>
0023 #include <deque>
0024 #include <functional>
0025 #include <map>
0026 #include <memory>
0027 #include <string>
0028 #include <vector>
0029 
0030 #include "algorithms/meta/SubDivideFunctors.h"
0031 #include "extensions/jana/JOmniFactoryGeneratorT.h"
0032 #include "factories/digi/PulseCombiner_factory.h"
0033 #include "factories/digi/PulseGeneration_factory.h"
0034 #include "factories/digi/PulseNoise_factory.h"
0035 #include "factories/digi/SiliconChargeSharing_factory.h"
0036 #include "factories/digi/SiliconTrackerDigi_factory.h"
0037 #include "factories/fardetectors/FarDetectorLinearTracking_factory.h"
0038 #include "factories/fardetectors/FarDetectorTrackerCluster_factory.h"
0039 #include "factories/fardetectors/FarDetectorTransportationPostML_factory.h"
0040 #include "factories/fardetectors/FarDetectorTransportationPreML_factory.h"
0041 #include "factories/meta/CollectionCollector_factory.h"
0042 #include "factories/meta/ONNXInference_factory.h"
0043 #include "factories/meta/SubDivideCollection_factory.h"
0044 #include "factories/tracking/TrackerHitReconstruction_factory.h"
0045 
0046 extern "C" {
0047 void InitPlugin(JApplication* app) {
0048   InitJANAPlugin(app);
0049 
0050   using namespace eicrecon;
0051 
0052   std::string readout = "TaggerTrackerHits";
0053 
0054   app->Add(new JOmniFactoryGeneratorT<SiliconChargeSharing_factory>(
0055       "TaggerTrackerChargeSharing", {"TaggerTrackerHits"}, {"TaggerTrackerSharedHits"},
0056       {
0057           .sigma_sharingx = 15 * dd4hep::um,
0058           .sigma_sharingy = 15 * dd4hep::um,
0059           .min_edep       = 0.1 * edm4eic::unit::keV,
0060           .readout        = readout,
0061       },
0062       app));
0063   //  Generate signal pulse from hits
0064   app->Add(new JOmniFactoryGeneratorT<PulseGeneration_factory<edm4hep::SimTrackerHit>>(
0065       "TaggerTrackerPulseGeneration", {"TaggerTrackerSharedHits"}, {"TaggerTrackerHitPulses"},
0066       {
0067           .pulse_shape_function = "LandauPulse",
0068           .pulse_shape_params   = {1.0, 2 * edm4eic::unit::ns},
0069           .ignore_thres         = 15.0e-8,
0070           .timestep             = 0.2 * edm4eic::unit::ns,
0071       },
0072       app));
0073 
0074   // Combine pulses into larger pulses
0075   app->Add(new JOmniFactoryGeneratorT<PulseCombiner_factory>(
0076       "TaggerTrackerPulseCombiner", {"TaggerTrackerHitPulses"}, {"TaggerTrackerCombinedPulses"},
0077       {
0078           .minimum_separation = 25 * edm4eic::unit::ns,
0079       },
0080       app));
0081 
0082   // Add noise to pulses
0083   app->Add(new JOmniFactoryGeneratorT<PulseNoise_factory>(
0084       "TaggerTrackerPulseNoise", {"EventHeader", "TaggerTrackerCombinedPulses"},
0085       {"TaggerTrackerCombinedPulsesWithNoise"},
0086       {
0087           .poles    = 5,
0088           .variance = 1.0,
0089           .alpha    = 0.5,
0090           .scale    = 0.000002,
0091       },
0092       app));
0093 
0094   // Digitization of silicon hits
0095   app->Add(new JOmniFactoryGeneratorT<SiliconTrackerDigi_factory>(
0096       "TaggerTrackerRawHits", {"EventHeader", "TaggerTrackerHits"},
0097       {"TaggerTrackerRawHits",
0098 #if EDM4EIC_BUILD_VERSION >= EDM4EIC_VERSION(8, 7, 0)
0099        "TaggerTrackerRawHitLinks",
0100 #endif
0101        "TaggerTrackerRawHitAssociations"},
0102       {
0103           .threshold      = 1.5 * edm4eic::unit::keV,
0104           .timeResolution = 2 * edm4eic::unit::ns,
0105       },
0106       app));
0107 
0108   // Convert raw digitized hits into hits with geometry info (ready for tracking)
0109   app->Add(new JOmniFactoryGeneratorT<TrackerHitReconstruction_factory>(
0110       "TaggerTrackerRecHits", {"TaggerTrackerRawHits"}, {"TaggerTrackerRecHits"},
0111       {
0112           .timeResolution = 2,
0113       },
0114       app));
0115 
0116   // Divide collection based on geometry segmentation labels
0117   // This should really be done before digitization as summing hits in the same cell couldn't even be mixed between layers. At the moment just prep for clustering.
0118   std::vector<std::string> geometryLabels{"module", "layer"};
0119   std::vector<int> moduleIDs{1, 2};
0120   std::vector<int> layerIDs{0, 1, 2, 3};
0121   std::vector<std::vector<long int>> geometryDivisions{};
0122   std::vector<std::string> geometryDivisionCollectionNames;
0123   std::vector<std::string> outputClusterCollectionNames;
0124   std::vector<std::string> outputTrackTags;
0125 #if EDM4EIC_BUILD_VERSION >= EDM4EIC_VERSION(8, 7, 0)
0126   std::vector<std::string> outputTrackLinkTags;
0127 #endif
0128   std::vector<std::string> outputTrackAssociationTags;
0129   std::vector<std::vector<std::string>> moduleClusterTags;
0130 
0131   for (int mod_id : moduleIDs) {
0132     outputTrackTags.push_back(fmt::format("TaggerTrackerM{}LocalTracks", mod_id));
0133 #if EDM4EIC_BUILD_VERSION >= EDM4EIC_VERSION(8, 7, 0)
0134     outputTrackLinkTags.push_back(fmt::format("TaggerTrackerM{}LocalTrackLinks", mod_id));
0135 #endif
0136     outputTrackAssociationTags.push_back(
0137         fmt::format("TaggerTrackerM{}LocalTrackAssociations", mod_id));
0138     moduleClusterTags.emplace_back();
0139     for (int lay_id : layerIDs) {
0140       geometryDivisions.push_back({mod_id, lay_id});
0141       geometryDivisionCollectionNames.push_back(
0142           fmt::format("TaggerTrackerM{}L{}RecHits", mod_id, lay_id));
0143       outputClusterCollectionNames.push_back(
0144           fmt::format("TaggerTrackerM{}L{}ClusterPositions", mod_id, lay_id));
0145       moduleClusterTags.back().push_back(outputClusterCollectionNames.back());
0146     }
0147   }
0148 
0149   app->Add(new JOmniFactoryGeneratorT<SubDivideCollection_factory<edm4eic::TrackerHit>>(
0150       "TaggerTrackerSplitHits", {"TaggerTrackerRecHits"}, geometryDivisionCollectionNames,
0151       {
0152           .function = GeometrySplit{geometryDivisions, readout, geometryLabels},
0153       },
0154       app));
0155 
0156   app->Add(new JOmniFactoryGeneratorT<FarDetectorTrackerCluster_factory>(
0157       "TaggerTrackerClustering", geometryDivisionCollectionNames, outputClusterCollectionNames,
0158       {
0159           .readout        = "TaggerTrackerHits",
0160           .x_field        = "x",
0161           .y_field        = "y",
0162           .hit_time_limit = 10 * edm4eic::unit::ns,
0163       },
0164       app));
0165 
0166   // Linear tracking for each module, loop over modules
0167   for (std::size_t i = 0; i < moduleIDs.size(); i++) {
0168     std::string outputTrackTag = outputTrackTags[i];
0169 #if EDM4EIC_BUILD_VERSION >= EDM4EIC_VERSION(8, 7, 0)
0170     std::string outputTrackLinkTag = outputTrackLinkTags[i];
0171 #endif
0172     std::string outputTrackAssociationTag     = outputTrackAssociationTags[i];
0173     std::vector<std::string> inputClusterTags = moduleClusterTags[i];
0174 
0175 #if EDM4EIC_BUILD_VERSION >= EDM4EIC_VERSION(8, 7, 0)
0176     inputClusterTags.emplace_back("TaggerTrackerRawHitLinks");
0177 #endif
0178     inputClusterTags.emplace_back("TaggerTrackerRawHitAssociations");
0179 
0180     app->Add(new JOmniFactoryGeneratorT<FarDetectorLinearTracking_factory>(
0181         outputTrackTag, {inputClusterTags},
0182         {outputTrackTag,
0183 #if EDM4EIC_BUILD_VERSION >= EDM4EIC_VERSION(8, 7, 0)
0184          outputTrackLinkTag,
0185 #endif
0186          outputTrackAssociationTag},
0187         {
0188             .layer_hits_max       = 200,
0189             .chi2_max             = 0.001,
0190             .n_layer              = 4,
0191             .layer_weights        = {1.0, 1.0, 1.0, 1.0},
0192             .restrict_direction   = true,
0193             .optimum_theta        = -M_PI + 0.026,
0194             .optimum_phi          = 0,
0195             .step_angle_tolerance = 0.05,
0196         },
0197         app));
0198   }
0199 
0200   // Combine the tracks from each module into one collection
0201   app->Add(new JOmniFactoryGeneratorT<CollectionCollector_factory<edm4eic::Track, true>>(
0202       "TaggerTrackerLocalTracks", outputTrackTags, {"TaggerTrackerLocalTracks"}, app));
0203 
0204 #if EDM4EIC_BUILD_VERSION >= EDM4EIC_VERSION(8, 7, 0)
0205   // Combine the track links from each module into one collection
0206   app->Add(new JOmniFactoryGeneratorT<
0207            CollectionCollector_factory<edm4eic::MCRecoTrackParticleLink, true>>(
0208       "TaggerTrackerLocalTrackLinks", outputTrackLinkTags, {"TaggerTrackerLocalTrackLinks"}, app));
0209 #endif
0210 
0211   // Combine the associations from each module into one collection
0212   app->Add(new JOmniFactoryGeneratorT<
0213            CollectionCollector_factory<edm4eic::MCRecoTrackParticleAssociation, true>>(
0214       "TaggerTrackerLocalTrackAssociations", outputTrackAssociationTags,
0215       {"TaggerTrackerLocalTrackAssociations"}, app));
0216 
0217   app->Add(new JOmniFactoryGeneratorT<FarDetectorTransportationPreML_factory>(
0218       "TaggerTrackerTransportationPreML",
0219       {"TaggerTrackerLocalTracks", "TaggerTrackerLocalTrackAssociations", "MCBeamElectrons"},
0220       {"TaggerTrackerFeatureTensor", "TaggerTrackerTargetTensor"},
0221       {
0222           .beamE = 10.0,
0223       },
0224       app));
0225   app->Add(new JOmniFactoryGeneratorT<ONNXInference_factory>(
0226       "TaggerTrackerTransportationInference", {"TaggerTrackerFeatureTensor"},
0227       {"TaggerTrackerPredictionTensor"},
0228       {
0229           .modelPath = "calibrations/onnx/Low-Q2_Steering_Reconstruction.onnx",
0230       },
0231       app));
0232   app->Add(new JOmniFactoryGeneratorT<FarDetectorTransportationPostML_factory>(
0233       "TaggerTrackerTransportationPostML",
0234       {"TaggerTrackerPredictionTensor", "TaggerTrackerLocalTrackAssociations", "MCBeamElectrons"},
0235       {"TaggerTrackerReconstructedParticles",
0236 #if EDM4EIC_BUILD_VERSION >= EDM4EIC_VERSION(8, 7, 0)
0237        "TaggerTrackerReconstructedParticleLinks",
0238 #endif
0239        "TaggerTrackerReconstructedParticleAssociations"},
0240       {
0241           .beamE = 10.0,
0242       },
0243       app));
0244 }
0245 }