detectors/LOWQ2/LOWQ2.cc

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