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