Changed the part of the code that finds the vertex. Also passed the AtPatternEvent to the GetModifiedTrack method since it may be useful too.

AtReconstruction/AtPatternModification/AtBraggCurveFinder.cxx

@@ -31,7 +31,8 @@ void AtBraggCurveFinder::ModifyPatternEvent(AtPatternEvent *patternEvent, AtRawE
    AtPatternModification::ModifyPatternEvent(patternEvent, rawEvent, event);
 }
 
-AtTrack AtBraggCurveFinder::GetModifiedTrack(const AtTrack &track, AtRawEvent *rawEvent, AtEvent *event)
+AtTrack AtBraggCurveFinder::GetModifiedTrack(const AtTrack &track, AtPatternEvent *patternEvent, AtRawEvent *rawEvent,
+                                             AtEvent *event)
 {
    // Create a copy of the AtTrack as an AtTrackBragg (not yet, still AtTrack for now).
    AtTrack modifiedTrack(track);
@@ -40,17 +41,15 @@ AtTrack AtBraggCurveFinder::GetModifiedTrack(const AtTrack &track, AtRawEvent *r
    auto *pattern = modifiedTrack.GetPattern();
 
    // Find the vertex of this track.
-   std::vector<AtTrack> trackToFindVtx;
-   trackToFindVtx.push_back(modifiedTrack);
-   AtFindVertex findVtx(fLineDistThreshold);
-   findVtx.FindVertex(trackToFindVtx, 1);
-   std::vector<tracksFromVertex> tv = findVtx.GetTracksVertex();
-   if (tv.size() != 1) {
-      LOG(warning) << "Found " << tv.size()
-                   << " vertex. We need to have 1 and only 1 to find the Bragg curve! Skipping this track!";
+   bool foundVertex{false};
+   XYZPoint vertex = FindVertex(modifiedTrack, patternEvent, foundVertex);
+
+   if (!foundVertex) {
+      LOG(warning) << "Could not find the vertex for the track with ID " << modifiedTrack.GetTrackID()
+                   << "! Maybe you need to change the distance threshold or the number of tracks per vertex. Skipping "
+                      "this track!";
       return modifiedTrack;
    }
-   XYZPoint vertex = (XYZPoint)tv.at(0).vertex;
 
    // Extract the AtHits.
    std::vector<AtHit> hitArray = modifiedTrack.GetHitArrayObject();
@@ -65,6 +64,30 @@ AtTrack AtBraggCurveFinder::GetModifiedTrack(const AtTrack &track, AtRawEvent *r
    return modifiedTrack;
 }
 
+AtBraggCurveFinder::XYZPoint
+AtBraggCurveFinder::FindVertex(const AtTrack &modifiedTrack, AtPatternEvent *patternEvent, bool &foundVertex)
+{
+   std::vector<AtTrack> trackToFindVtx = patternEvent->GetTrackCand();
+   AtFindVertex findVtx(fLineDistThreshold);
+   findVtx.FindVertex(trackToFindVtx, fNumTracksPerVtx);
+   std::vector<tracksFromVertex> tv = findVtx.GetTracksVertex();
+
+   XYZPoint vertex;
+   for (auto trackVertex : tv) {
+      for (auto track : trackVertex.tracks) {
+         if (track.GetTrackID() == modifiedTrack.GetTrackID()) {
+            vertex = trackVertex.vertex;
+            foundVertex = true;
+            break;
+         }
+      }
+      if (foundVertex)
+         break;
+   }
+
+   return vertex;
+}
+
 void AtBraggCurveFinder::ProcessHit(XYZPoint vertex, AtHit hit, AtTrack &modifiedTrack, AtRawEvent *rawEvent)
 {
    if (rawEvent == nullptr) {

AtReconstruction/AtPatternModification/AtBraggCurveFinder.h

@@ -18,8 +18,9 @@ class AtBraggCurveFinder : public AtPatternModification {
    using XYZPoint = ROOT::Math::XYZPoint;
 
 protected:
-   // Line distance threshold to be used by the AtFindVertex.
-   Double_t fLineDistThreshold{30};
+   // Parameters to be used by the AtFindVertex.
+   double fLineDistThreshold{30};
+   int fNumTracksPerVtx{1};
 
    // Pointer to AtPSAHitPerTB and parameters related to it.
    std::unique_ptr<AtPSAHitPerTB> fPSA{nullptr};
@@ -45,7 +46,8 @@ class AtBraggCurveFinder : public AtPatternModification {
    virtual void
    ModifyPatternEvent(AtPatternEvent *patternEvent, AtRawEvent *rawEvent = nullptr, AtEvent *event = nullptr) override;
 
-   void SetLineDistThreshold(Double_t lineDistThreshold) { fLineDistThreshold = lineDistThreshold; }
+   void SetLineDistThreshold(double lineDistThreshold) { fLineDistThreshold = lineDistThreshold; }
+   void SetNumTracksPerVtx(int numTracksPerVtx) { fNumTracksPerVtx = numTracksPerVtx; }
    void SetTSSemiWidth(int value) { fTSSemiWidth = value; }
    void SetNeedPSA(bool value) { fNeedPSA = value; }
 
@@ -55,8 +57,9 @@ class AtBraggCurveFinder : public AtPatternModification {
    void SetNumSmoothingSteps(int value) { fNumSmoothingSteps = value; }
 
 protected:
-   virtual AtTrack
-   GetModifiedTrack(const AtTrack &track, AtRawEvent *rawEvent = nullptr, AtEvent *event = nullptr) override;
+   virtual AtTrack GetModifiedTrack(const AtTrack &track, AtPatternEvent *patternEvent, AtRawEvent *rawEvent = nullptr,
+                                    AtEvent *event = nullptr) override;
+   XYZPoint FindVertex(const AtTrack &modifiedTrack, AtPatternEvent *patternEvent, bool &foundVertex);
    void ProcessHit(XYZPoint vertex, AtHit hit, AtTrack &modifiedTrack, AtRawEvent *rawEvent);
    void ProcessHit(XYZPoint vertex, AtHit hit, AtTrack &modifiedTrack);
 

AtReconstruction/AtPatternModification/AtPatternModification.cxx

@@ -10,7 +10,7 @@ void AtPatternModification::ModifyPatternEvent(AtPatternEvent *patternEvent, AtR
 
    // Iterate over the original tracks and store the modified ones.
    for (auto track : tracks)
-      modifiedTracks.push_back(GetModifiedTrack(track, rawEvent, event));
+      modifiedTracks.push_back(GetModifiedTrack(track, patternEvent, rawEvent, event));
 
    // Replace the vector of track candidates with the new modified one.
    patternEvent->SetTrackCand(modifiedTracks);

AtReconstruction/AtPatternModification/AtPatternModification.h

@@ -23,7 +23,8 @@ class AtPatternModification {
    /**
     * Actually implements ths track modification.
     */
-   virtual AtTrack GetModifiedTrack(const AtTrack &track, AtRawEvent *rawEvent = nullptr, AtEvent *event = nullptr) = 0;
+   virtual AtTrack GetModifiedTrack(const AtTrack &track, AtPatternEvent *patternEvent, AtRawEvent *rawEvent = nullptr,
+                                    AtEvent *event = nullptr) = 0;
 };
 
 #endif