diff --git a/include/SimCore/TrackMap.h b/include/SimCore/TrackMap.h index e4cc46e..d63068e 100644 --- a/include/SimCore/TrackMap.h +++ b/include/SimCore/TrackMap.h @@ -41,6 +41,12 @@ class TrackMap { return ancestry_.find(track->GetTrackID()) != ancestry_.end(); } + /** + * Check if the track with the given ID is a descendant of the track with + * the given ancestor ID up to a given depth. + */ + bool isDescendant(int trackID, int ancestorID, int maximum_depth) const; + /** * Find a trajectory's nearest parent that is incident on the calorimeter * region. We assume that the primary particles have a parent ID of 0. diff --git a/src/SimCore/TrackMap.cxx b/src/SimCore/TrackMap.cxx index cd25c54..0731ea7 100644 --- a/src/SimCore/TrackMap.cxx +++ b/src/SimCore/TrackMap.cxx @@ -6,6 +6,26 @@ namespace simcore { +bool TrackMap::isDescendant(int trackID, int ancestorID, + int maximum_depth) const { + int current_depth{0}; + int current_track{trackID}; + // Walk the tree until we either no longer have a parent or we reach the + // desired depth + while (current_depth < maximum_depth && + (ancestry_.find(current_track) != ancestry_.end())) { + // See if we have encountered the parent of the current track + // + // operator[] is not const, so we need to use at() + current_track = ancestry_.at(current_track).first; + if (current_track == ancestorID) { + // If one of the parents is the track of interest, we are done! + return true; + } + current_depth++; + } + return false; +} void TrackMap::insert(const G4Track* track) { ancestry_[track->GetTrackID()] = std::make_pair(track->GetParentID(), isInCalorimeterRegion(track));