From 816a8daaf87d8323bc3dee4e0616290e56d1085f Mon Sep 17 00:00:00 2001
From: Xin Dong <XDong@lbl.gov>
Date: Mon, 26 Aug 2024 13:00:11 -0700
Subject: [PATCH] IterativeVertexFinder: fill vertex-to-particle relation
 (#1576)

### Briefly, what does this PR introduce?
- added association for ReconstructionParticles to output vertices

### What kind of change does this PR introduce?
- [ ] Bug fix (issue #__)
- [X] New feature (issue #1575)
- [ ] Documentation update
- [ ] Other: __

### Please check if this PR fulfills the following:
- [X] Tests for the changes have been added
- [X] Documentation has been added / updated
- [X] Changes have been communicated to collaborators

### Does this PR introduce breaking changes? What changes might users
need to make to their code?

### Does this PR change default behavior?
- ReconstructedParticle association added to the default output
CKFCentralVertices

---------

Co-authored-by: Xin Dong <xdong@eic0102.rcf.bnl.gov>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Dmitry Kalinkin <dmitry.kalinkin@gmail.com>
Co-authored-by: Xin Dong <xdong@eic0104.rcf.bnl.gov>
---
 .../tracking/IterativeVertexFinder.cc         | 52 +++++++++++++++++--
 .../tracking/IterativeVertexFinder.h          |  3 +-
 .../tracking/IterativeVertexFinder_factory.h  |  4 +-
 src/global/tracking/tracking.cc               |  2 +-
 4 files changed, 55 insertions(+), 6 deletions(-)

diff --git a/src/algorithms/tracking/IterativeVertexFinder.cc b/src/algorithms/tracking/IterativeVertexFinder.cc
index c933a96feb..b0fcbb4fd8 100644
--- a/src/algorithms/tracking/IterativeVertexFinder.cc
+++ b/src/algorithms/tracking/IterativeVertexFinder.cc
@@ -7,13 +7,18 @@
 #include <Acts/Definitions/Common.hpp>
 #include <Acts/Definitions/Direction.hpp>
 #include <Acts/Definitions/TrackParametrization.hpp>
+#include <Acts/Definitions/Units.hpp>
 #include <Acts/EventData/GenericBoundTrackParameters.hpp>
 #include <Acts/EventData/GenericParticleHypothesis.hpp>
 #include <Acts/EventData/ParticleHypothesis.hpp>
 #include <Acts/EventData/TrackParameters.hpp>
 #include <Acts/Propagator/EigenStepper.hpp>
 #include <Acts/Propagator/Propagator.hpp>
+#include <ActsExamples/EventData/Track.hpp>
 #include <boost/move/utility_core.hpp>
+#include <edm4eic/Track.h>
+#include <fmt/core.h>
+#include <podio/RelationRange.h>
 #if Acts_VERSION_MAJOR >= 32
 #include <Acts/Propagator/VoidNavigator.hpp>
 #else
@@ -26,17 +31,23 @@
 #include <Acts/Vertexing/HelicalTrackLinearizer.hpp>
 #include <Acts/Vertexing/ImpactPointEstimator.hpp>
 #include <Acts/Vertexing/IterativeVertexFinder.hpp>
+#include <Acts/Vertexing/TrackAtVertex.hpp>
 #include <Acts/Vertexing/Vertex.hpp>
 #include <Acts/Vertexing/VertexingOptions.hpp>
 #include <Acts/Vertexing/ZScanVertexFinder.hpp>
 #include <ActsExamples/EventData/Trajectories.hpp>
 #include <boost/container/vector.hpp>
 #include <edm4eic/Cov4f.h>
-#include <math.h>
+#include <edm4eic/ReconstructedParticleCollection.h>
+#include <edm4eic/TrackParameters.h>
+#include <edm4eic/Trajectory.h>
+#include <edm4eic/unit_system.h>
+#include <edm4hep/Vector2f.h>
 #include <Eigen/Core>
 #include <Eigen/Geometry>
 #include <Eigen/LU>
 #include <algorithm>
+#include <cmath>
 #include <optional>
 #include <tuple>
 #include <utility>
@@ -56,7 +67,8 @@ void eicrecon::IterativeVertexFinder::init(std::shared_ptr<const ActsGeometryPro
 }
 
 std::unique_ptr<edm4eic::VertexCollection> eicrecon::IterativeVertexFinder::produce(
-    std::vector<const ActsExamples::Trajectories*> trajectories) {
+    std::vector<const ActsExamples::Trajectories*> trajectories,
+    const edm4eic::ReconstructedParticleCollection* reconParticles) {
 
   auto outputVertices = std::make_unique<edm4eic::VertexCollection>();
 
@@ -172,11 +184,14 @@ std::unique_ptr<edm4eic::VertexCollection> eicrecon::IterativeVertexFinder::prod
     }
     /// CKF can provide multiple track trajectories for a single input seed
     for (auto& tip : tips) {
+      ActsExamples::TrackParameters par = trajectory->trackParameters(tip);
+
 #if Acts_VERSION_MAJOR >= 33
       inputTracks.emplace_back(&(trajectory->trackParameters(tip)));
 #else
       inputTrackPointers.push_back(&(trajectory->trackParameters(tip)));
 #endif
+      m_log->trace("Track local position at input = {} mm, {} mm", par.localPosition().x() / Acts::UnitConstants::mm, par.localPosition().y() / Acts::UnitConstants::mm);
     }
   }
 
@@ -207,7 +222,38 @@ std::unique_ptr<edm4eic::VertexCollection> eicrecon::IterativeVertexFinder::prod
          (float)vtx.time(),
     }); // vtxposition
     eicvertex.setPositionError(cov);                          // covariance
-  }
+
+    for (const auto& t : vtx.tracks()) {
+#if Acts_VERSION_MAJOR >= 33
+      const auto& trk = &t.originalParams;
+      const auto& par = finderCfg.extractParameters(trk);
+#else
+      const auto& par = *t.originalParams;
+#endif
+      m_log->trace("Track local position from vertex = {} mm, {} mm", par.localPosition().x() / Acts::UnitConstants::mm, par.localPosition().y() / Acts::UnitConstants::mm);
+      float loc_a = par.localPosition().x();
+      float loc_b = par.localPosition().y();
+
+      for (const auto& part : *reconParticles) {
+        const auto& tracks = part.getTracks();
+        for (const auto trk : tracks) {
+          const auto& traj = trk.getTrajectory();
+          const auto& trkPars = traj.getTrackParameters();
+          for (const auto par : trkPars) {
+            const double EPSILON = 1.0e-4; // mm
+            if (fabs((par.getLoc().a / edm4eic::unit::mm) - (loc_a / Acts::UnitConstants::mm)) < EPSILON
+              && fabs((par.getLoc().b / edm4eic::unit::mm) - (loc_b / Acts::UnitConstants::mm)) < EPSILON) {
+              m_log->trace("From ReconParticles, track local position [Loc a, Loc b] = {} mm, {} mm", par.getLoc().a / edm4eic::unit::mm, par.getLoc().b / edm4eic::unit::mm);
+              eicvertex.addToAssociatedParticles(part);
+            } // endif
+          } // end for par
+        } // end for trk
+      } // end for part
+    } // end for t
+    m_log->debug("One vertex found at (x,y,z) = ({}, {}, {}) mm.", vtx.position().x() / Acts::UnitConstants::mm, vtx.position().y() / Acts::UnitConstants::mm, vtx.position().z() / Acts::UnitConstants::mm);
+
+  } // end for vtx
+
 
   return std::move(outputVertices);
 }
diff --git a/src/algorithms/tracking/IterativeVertexFinder.h b/src/algorithms/tracking/IterativeVertexFinder.h
index b666f99861..1bdb81ad93 100644
--- a/src/algorithms/tracking/IterativeVertexFinder.h
+++ b/src/algorithms/tracking/IterativeVertexFinder.h
@@ -7,6 +7,7 @@
 #include <Acts/Geometry/GeometryContext.hpp>
 #include <Acts/MagneticField/MagneticFieldContext.hpp>
 #include <edm4eic/VertexCollection.h>
+#include <edm4eic/ReconstructedParticle.h>
 #include <spdlog/logger.h>
 #include <memory>
 #include <vector>
@@ -24,7 +25,7 @@ class IterativeVertexFinder
   void init(std::shared_ptr<const ActsGeometryProvider> geo_svc,
             std::shared_ptr<spdlog::logger> log);
   std::unique_ptr<edm4eic::VertexCollection>
-  produce(std::vector<const ActsExamples::Trajectories*> trajectories);
+  produce(std::vector<const ActsExamples::Trajectories*> trajectories, const edm4eic::ReconstructedParticleCollection* reconParticles);
 
 private:
   std::shared_ptr<spdlog::logger> m_log;
diff --git a/src/global/tracking/IterativeVertexFinder_factory.h b/src/global/tracking/IterativeVertexFinder_factory.h
index 7ab8450556..6f981db1ab 100644
--- a/src/global/tracking/IterativeVertexFinder_factory.h
+++ b/src/global/tracking/IterativeVertexFinder_factory.h
@@ -7,6 +7,7 @@
 #include <ActsExamples/EventData/Trajectories.hpp>
 #include <JANA/JEvent.h>
 #include <edm4eic/VertexCollection.h>
+#include <edm4eic/ReconstructedParticle.h>
 #include <memory>
 #include <string>
 #include <utility>
@@ -26,6 +27,7 @@ class IterativeVertexFinder_factory :
     std::unique_ptr<AlgoT> m_algo;
 
     Input<ActsExamples::Trajectories> m_acts_trajectories_input {this};
+    PodioInput<edm4eic::ReconstructedParticle> m_edm4eic_reconParticles_input {this};
     PodioOutput<edm4eic::Vertex> m_vertices_output {this};
 
     ParameterRef<int> m_maxVertices {this, "maxVertices", config().maxVertices,
@@ -47,7 +49,7 @@ class IterativeVertexFinder_factory :
     }
 
     void Process(int64_t run_number, uint64_t event_number) {
-        m_vertices_output() = m_algo->produce(m_acts_trajectories_input());
+        m_vertices_output() = m_algo->produce(m_acts_trajectories_input(), m_edm4eic_reconParticles_input());
     }
 };
 
diff --git a/src/global/tracking/tracking.cc b/src/global/tracking/tracking.cc
index 7a916effd8..fc1b48e971 100644
--- a/src/global/tracking/tracking.cc
+++ b/src/global/tracking/tracking.cc
@@ -225,7 +225,7 @@ void InitPlugin(JApplication *app) {
 
     app->Add(new JOmniFactoryGeneratorT<IterativeVertexFinder_factory>(
             "CentralTrackVertices",
-            {"CentralCKFActsTrajectories"},
+            {"CentralCKFActsTrajectories","ReconstructedChargedParticles"},
             {"CentralTrackVertices"},
             {},
             app