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