Transformer to create tracks from gen particles

Tracking/CMakeLists.txt

@@ -3,11 +3,13 @@ set(PackageName Tracking)
 project(${PackageName})
 
 #find_package(GenFit)
+FIND_PACKAGE(MarlinUtil)
 
 #if (GenFit_FOUND)
 
 file(GLOB sources
     ${PROJECT_SOURCE_DIR}/src/*.cpp
+    ${PROJECT_SOURCE_DIR}/components/*.cpp
 )
 
 file(GLOB headers
@@ -27,6 +29,7 @@ gaudi_add_module(${PackageName}
 target_include_directories(${PackageName} PUBLIC
   $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>
   $<INSTALL_INTERFACE:include>
+  ${MarlinUtil_INCLUDE_DIRS}
 )
 
 set_target_properties(${PackageName} PROPERTIES PUBLIC_HEADER "${headers}")
@@ -45,4 +48,9 @@ install(TARGETS ${PackageName}
 )
 
 install(FILES ${scripts} DESTINATION test)
+
+SET(test_name "test_TracksFromGenParticles")
+ADD_TEST(NAME ${test_name} COMMAND k4run test/runTracksFromGenParticles.py)
+set_test_env(${test_name})
+
 #endif()

Tracking/components/PlotTrackHitResiduals.cpp

@@ -0,0 +1,97 @@
+// Gaudi
+#include "Gaudi/Property.h"
+#include "GaudiAlg/Consumer.h"
+#include "Gaudi/Accumulators/RootHistogram.h"
+#include "Gaudi/Histograming/Sink/Utils.h"
+
+// edm4hep
+#include "edm4hep/MCParticleCollection.h"
+#include "edm4hep/TrackCollection.h"
+#include "edm4hep/MCRecoTrackParticleAssociationCollection.h"
+#include "edm4hep/SimTrackerHitCollection.h"
+
+// marlin
+#include <marlinutil/HelixClass_double.h>
+
+// ROOT
+#include "TH1D.h"
+
+// Define BaseClass_t
+#include "k4FWCore/BaseClass.h"
+
+#include <string>
+
+/** @class PlotTrackHitDistances
+ *
+ *  Gaudi consumer that generates a residual distribution (mm) by comparing the helix from Track AtIP and simHit position.
+ *  This is intended to be used on tracks produced from gen particles i.e. which do not have real hits attached to them.
+ *
+ *  @author Brieuc Francois
+ */
+
+struct PlotTrackHitDistances final
+  : Gaudi::Functional::Consumer<void(const edm4hep::SimTrackerHitCollection&, const edm4hep::MCRecoTrackParticleAssociationCollection&), BaseClass_t> {
+  PlotTrackHitDistances(const std::string& name, ISvcLocator* svcLoc)
+      : Consumer(
+            name, svcLoc,
+            {
+            KeyValue("InputSimTrackerHits", "DCHCollection"),
+            KeyValue("InputTracksFromGenParticlesAssociation", "TracksFromGenParticlesAssociation"),
+            }) {}
+
+  void operator()(const edm4hep::SimTrackerHitCollection& simTrackerHits, const edm4hep::MCRecoTrackParticleAssociationCollection& trackParticleAssociations) const override {
+
+    for (const auto& trackParticleAssociation : trackParticleAssociations) {
+      auto genParticle = trackParticleAssociation.getSim();
+      auto track = trackParticleAssociation.getRec();
+      edm4hep::TrackState trackStateAtIP;
+      bool found_trackStateAtIP = false;
+      for (const auto& trackState : track.getTrackStates()) {
+        if (trackState.location == edm4hep::TrackState::AtIP) {
+          trackStateAtIP = trackState;
+          found_trackStateAtIP = true;
+        }
+      }
+      if (!found_trackStateAtIP)
+        throw std::runtime_error("No track state defined AtIP, exiting!");
+
+      // Build an helix out of the trackState
+      auto helixFromTrack = HelixClass_double();
+      helixFromTrack.Initialize_Canonical(trackStateAtIP.phi, trackStateAtIP.D0, trackStateAtIP.Z0, trackStateAtIP.omega, trackStateAtIP.tanLambda, m_Bz);
+
+      // Fill the histogram with residuals for hits attached to the same gen particle
+      for (const auto& simTrackerHit : simTrackerHits) {
+        auto simTrackerHitgenParticle = simTrackerHit.getParticle();
+        if (simTrackerHitgenParticle.getObjectID() == genParticle.getObjectID()) {
+          double simTrackerHitPosition[] = {simTrackerHit.x(), simTrackerHit.y(), simTrackerHit.z()};
+          double distances[3];
+          helixFromTrack.getDistanceToPoint(simTrackerHitPosition, distances);
+          // Distance[0] - distance in R-Phi plane, Distance[1] - distance along Z axis, Distance[2] - 3D distance
+          ++m_residualHist[distances[2]];
+        }
+      }
+    }
+    return;
+  }
+  Gaudi::Property<float> m_Bz{this, "Bz", 2., "Z component of the (assumed constant) magnetic field in Tesla."};
+  mutable Gaudi::Accumulators::Histogram<1> m_residualHist{this, "", "Track-hit Distances", {100, 0, 1, "Distance [mm];Entries"}};
+
+	StatusCode finalize() override {
+			TFile file("track_hit_distances.root", "RECREATE");
+			std::string name = "";
+			// Name that will appear in the stats table
+			std::string histName = "Distances";
+			nlohmann::json jH {m_residualHist};
+			auto [histo, dir] = Gaudi::Histograming::Sink::jsonToRootHistogram<Gaudi::Histograming::Sink::Traits<false, TH1D, 1>>(name, histName, jH[0]);
+      // Add overflow bin to the last bin
+      int lastBinIndex = histo.GetNbinsX();
+      histo.SetBinContent(lastBinIndex, histo.GetBinContent(lastBinIndex) + histo.GetBinContent(lastBinIndex + 1));
+			// Name of the histogram in the ROOT file
+			histo.Write("Distances");
+			file.Close();
+			return StatusCode::SUCCESS;
+		}
+
+};
+
+DECLARE_COMPONENT(PlotTrackHitDistances)

Tracking/components/TracksFromGenParticles.cpp

@@ -0,0 +1,91 @@
+#include "Gaudi/Property.h"
+#include "GaudiAlg/Transformer.h"
+
+// edm4hep
+#include "edm4hep/MCParticleCollection.h"
+#include "edm4hep/TrackCollection.h"
+#include "edm4hep/MCRecoTrackParticleAssociationCollection.h"
+
+// marlin
+#include <marlinutil/HelixClass_double.h>
+
+// Define BaseClass_t
+#include "k4FWCore/BaseClass.h"
+
+#include <string>
+
+/** @class TracksFromGenParticles
+ *
+ *  Gaudi transformer that builds an edm4hep::TrackCollection out of an edm4hep::MCParticleCollection.
+ *  It just builds an helix out of the genParticle position, momentum, charge and user defined z component of the (constant) magnetic field.
+ *  From this helix, different edm4hep::TrackStates (AtIP, AtFirstHit, AtLastHit and AtCalorimeter) are defined. #FIXME for now these trackstates are dummy (copies of the same helix parameters)
+ *  This is meant to enable technical development needing edm4hep::Track and performance studies where having generator based trackis is a reasonable approximation.
+ *  Possible inprovement:
+ *    - Retrieve magnetic field from geometry: const DD4hep::Field::MagneticField& magneticField = detector.field(); DD4hep::DDRec::Vector3D field = magneticField.magneticField(point);
+ *    - Properly define different trackStates
+ *
+ *  @author Brieuc Francois
+ */
+
+struct TracksFromGenParticles final
+  : Gaudi::Functional::MultiTransformer<std::tuple<edm4hep::TrackCollection, edm4hep::MCRecoTrackParticleAssociationCollection>(const edm4hep::MCParticleCollection&), BaseClass_t> {
+  TracksFromGenParticles(const std::string& name, ISvcLocator* svcLoc)
+      : MultiTransformer(
+            name, svcLoc,
+            {KeyValue("InputGenParticles", "MCParticles")},
+            {KeyValue("OutputTracks", "TracksFromGenParticles"),
+            KeyValue("OutputMCRecoTrackParticleAssociation", "TracksFromGenParticlesAssociation")}) {
+  }
+
+std::tuple<edm4hep::TrackCollection, edm4hep::MCRecoTrackParticleAssociationCollection> operator()(const edm4hep::MCParticleCollection& genParticleColl) const override {
+
+    auto outputTrackCollection = edm4hep::TrackCollection();
+    auto MCRecoTrackParticleAssociationCollection = edm4hep::MCRecoTrackParticleAssociationCollection();
+
+    for (const auto& genParticle : genParticleColl) {
+      debug() << "Particle decayed in tracker: " << genParticle.isDecayedInTracker() << endmsg;
+      debug() << genParticle << endmsg;
+
+      // Building an helix out of MCParticle properties and B field
+      auto helixFromGenParticle = HelixClass_double();
+      double genParticleVertex[] = {genParticle.getVertex().x, genParticle.getVertex().y, genParticle.getVertex().z};
+      double genParticleMomentum[] = {genParticle.getMomentum().x, genParticle.getMomentum().y, genParticle.getMomentum().z};
+      helixFromGenParticle.Initialize_VP(genParticleVertex, genParticleMomentum, genParticle.getCharge(), m_Bz);
+
+      // Setting the track and trackStates properties
+      // #FIXME for now, the different trackStates are dummy
+      auto trackFromGen = edm4hep::MutableTrack();
+      auto trackState_IP = edm4hep::TrackState {};
+      trackState_IP.location = edm4hep::TrackState::AtIP;
+      trackState_IP.D0 = helixFromGenParticle.getD0();
+      trackState_IP.phi = helixFromGenParticle.getPhi0();
+      trackState_IP.omega = helixFromGenParticle.getOmega();
+      trackState_IP.Z0 = helixFromGenParticle.getZ0();
+      trackState_IP.tanLambda = helixFromGenParticle.getTanLambda();
+      trackFromGen.addToTrackStates(trackState_IP);
+      auto trackState_AtFirstHit = edm4hep::TrackState(trackState_IP);
+      trackState_AtFirstHit.location = edm4hep::TrackState::AtFirstHit;
+      trackFromGen.addToTrackStates(trackState_AtFirstHit);
+      auto trackState_AtLastHit = edm4hep::TrackState(trackState_IP);
+      trackState_AtFirstHit.location = edm4hep::TrackState::AtLastHit;
+      trackFromGen.addToTrackStates(trackState_AtLastHit);
+      auto trackState_AtCalorimeter = edm4hep::TrackState(trackState_IP);
+      trackState_AtFirstHit.location = edm4hep::TrackState::AtCalorimeter;
+      trackFromGen.addToTrackStates(trackState_AtCalorimeter);
+
+      debug() << trackFromGen << endmsg;
+      outputTrackCollection.push_back(trackFromGen);
+
+      // Building the association between tracks and genParticles
+      auto MCRecoTrackParticleAssociation = edm4hep::MutableMCRecoTrackParticleAssociation();
+      MCRecoTrackParticleAssociation.setRec(trackFromGen);
+      MCRecoTrackParticleAssociation.setSim(genParticle);
+      MCRecoTrackParticleAssociationCollection.push_back(MCRecoTrackParticleAssociation);
+    }
+    return std::make_tuple(std::move(outputTrackCollection), std::move(MCRecoTrackParticleAssociationCollection));
+  }
+
+  Gaudi::Property<float> m_Bz{this, "Bz", 2., "Z component of the (assumed constant) magnetic field in Tesla."};
+};
+
+DECLARE_COMPONENT(TracksFromGenParticles)

Tracking/test/runTracksFromGenParticles.py

@@ -0,0 +1,52 @@
+from Configurables import ApplicationMgr
+from Configurables import EventCounter
+from Gaudi.Configuration import INFO, WARNING, DEBUG
+import os
+
+if not os.path.isfile("ddsim_output_edm4hep.root"):
+    os.system("ddsim --enableGun --gun.distribution uniform --gun.energy '10*GeV' --gun.particle e- --numberOfEvents 100 --outputFile ddsim_output_edm4hep.root --random.enableEventSeed --random.seed 42 --compactFile $K4GEO/FCCee/IDEA/compact/IDEA_o1_v02/IDEA_o1_v02.xml")
+
+# Loading the output of the SIM step
+from Configurables import k4DataSvc, PodioInput, PodioOutput
+evtsvc = k4DataSvc('EventDataSvc')
+evtsvc.input = "ddsim_output_edm4hep.root"
+#evtsvc.input = "/afs/cern.ch/user/b/brfranco/work/public/MIT_tutorial/CLDConfig/CLDConfig/wzp6_ee_mumuH_ecm240_CLD_RECO_edm4hep.root"
+Nevts = -1
+input_reader = PodioInput('InputReader')
+
+# Calling TracksFromGenParticles
+from Configurables import TracksFromGenParticles
+tracksFromGenParticles = TracksFromGenParticles("TracksFromGenParticles",
+                                               InputGenParticles = "MCParticles",
+                                               OutputTracks = "TracksFromGenParticles",
+                                               OutputMCRecoTrackParticleAssociation = "TracksFromGenParticlesAssociation",
+                                               Bz = 2.0,
+                                               OutputLevel = INFO)
+
+# produce a TH1 with distances between tracks and simTrackerHits
+from Configurables import PlotTrackHitDistances
+plotTrackHitDistances = PlotTrackHitDistances("PlotTrackHitDistances",
+                                             InputSimTrackerHits = "CDCHHits",
+                                             InputTracksFromGenParticlesAssociation = tracksFromGenParticles.OutputMCRecoTrackParticleAssociation, 
+                                             Bz = 2.0)
+
+# Output
+from Configurables import PodioOutput
+out = PodioOutput("out",
+                  OutputLevel=INFO)
+out.outputCommands = ["keep *"]
+out.filename = "tracks_from_genParticle_output.root"
+
+# Set auditor service
+from Configurables import AuditorSvc, ChronoAuditor
+chra = ChronoAuditor()
+audsvc = AuditorSvc()
+audsvc.Auditors = [chra]
+
+ApplicationMgr(
+    TopAlg= [input_reader, tracksFromGenParticles, plotTrackHitDistances, out],
+    EvtSel='NONE',
+    EvtMax=Nevts,
+    ExtSvc=[evtsvc, audsvc],
+    StopOnSignal=True,
+)