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