SimG4Core/Notification/src/SimTrackManager.cc - master - githubmirror/cmssw

SimG4Core/Notification/src/SimTrackManager.cc
468 строк15 KB
Felice Pantaleo
SimG4Core,PhysicsTools/TruthInfo: regression tests guarding the SimTrack/SimVertex history
20 июн 2026, 16:00
20 июн 2026, 16:005db28cf
// -*- C++ -*-
//
// Package:     Application
// Class  :     SimTrackManager
//
// Implementation:
//     <Notes on implementation>
//
// Original Author:
//         Created:  Fri Nov 25 17:44:19 EST 2005
//

// system include files
#include <iostream>
#include <unordered_set>

// user include files
#include "SimG4Core/Notification/interface/SimTrackManager.h"
#include "SimG4Core/Notification/interface/TmpSimTrack.h"
#include "SimG4Core/Notification/interface/TmpSimVertex.h"
#include "SimG4Core/Notification/interface/TmpSimEvent.h"
#include "SimG4Core/Notification/interface/TrackInformation.h"
#include "SimDataFormats/TrackingHit/interface/PSimHit.h"

#include "FWCore/MessageLogger/interface/MessageLogger.h"
#include "FWCore/Utilities/interface/Exception.h"

#include "G4VProcess.hh"
#include "G4Track.hh"
#include "G4ThreeVector.hh"
#include <CLHEP/Units/SystemOfUnits.h>

//#define DebugLog
namespace {
  const double invcm = 1.0 / CLHEP::cm;
  const double r_limit2 = 1.e-6;  // 10 micron in CMS units
}  // namespace

SimTrackManager::SimTrackManager(TmpSimEvent* ptr, int) : m_simEvent(ptr) {
  idsave.reserve(1000);
  ancestorList.reserve(1000);
  m_trackContainer.reserve(1000);
  m_endPoints.reserve(1000);
}

SimTrackManager::~SimTrackManager() { reset(); }

void SimTrackManager::reset() {
  deleteTracks();
  cleanVertexMap();
  idsave.clear();
  ancestorList.clear();
  m_parentOfAll.clear();
  m_droppedParentRedirect.clear();
  lastTrack = 0;
  lastHist = 0;
}

void SimTrackManager::deleteTracks() {
  if (!m_trackContainer.empty()) {
    for (auto const& ptr : m_trackContainer) {
      delete ptr;
    }
    m_trackContainer.clear();
    m_endPoints.clear();
  }
}

void SimTrackManager::cleanVertexMap() {
  m_vertexMap.clear();
  m_vertexMap.swap(m_vertexMap);
  m_nVertices = 0;
}

void SimTrackManager::addTrack(TrackWithHistory* iTrack, const G4Track* track, bool inHistory, bool withAncestor) {
  std::pair<int, int> thePair(iTrack->trackID(), iTrack->parentID());
  idsave.push_back(thePair);
  // Independent, persistent-for-the-event copy of the full parent map: idsave is
  // cleared per primary by fillMotherList(), but the redirect needs to walk the
  // whole chain (through dropped intermediates) at end of event.
  if (m_reconnectDroppedAncestors)
    m_parentOfAll[iTrack->trackID()] = iTrack->parentID();
  if (inHistory) {
    auto info = static_cast<const TrackInformation*>(track->GetUserInformation());
    if (info->isInTrkFromBackscattering())
      iTrack->setFromBackScattering();
    // set there for the *non-primary* tracks the G4Track ID of the last stored ancestor
    // for the primaries the genparticle id is saved in the TrackWithHistory constructor.
    // In the constructor of TrackWithHistory the info isPrimary is saved and used
    // to give -1 if the track is not a primary.
    if (not iTrack->isPrimary())
      iTrack->setGenParticleID(info->idLastStoredAncestor());
    m_trackContainer.push_back(iTrack);
    const auto& v = track->GetStep()->GetPostStepPoint()->GetPosition();
    std::pair<int, math::XYZVectorD> p(iTrack->trackID(),
                                       math::XYZVectorD(v.x() * invcm, v.y() * invcm, v.z() * invcm));
    m_endPoints.push_back(p);
  }
  if (withAncestor) {
    std::pair<int, int> thisPair(iTrack->trackID(), 0);
    ancestorList.push_back(thisPair);
  }
}

/// this saves a track and all its parents looping over the non ordered vector
void SimTrackManager::saveTrackAndItsBranch(TrackWithHistory* trkWHist) {
  TrackWithHistory* trkH = trkWHist;
  if (trkH == nullptr) {
    edm::LogError("SimTrackManager") << " SimTrackManager::saveTrackAndItsBranch got 0 pointer ";
    throw cms::Exception("SimTrackManager::saveTrackAndItsBranch") << " cannot handle hits for tracking";
  }
  trkH->setToBeSaved();
  int parent = trkH->parentID();

  auto tk_itr =
      std::lower_bound(m_trackContainer.begin(), m_trackContainer.end(), parent, SimTrackManager::StrictWeakOrdering());

  if (tk_itr != m_trackContainer.end() && (*tk_itr)->trackID() == parent) {
    saveTrackAndItsBranch(*tk_itr);
  }
}

void SimTrackManager::storeTracks() {
  cleanTracksWithHistory();

  // Precompute the nearest-stored-ancestor redirect while the full idsave parent
  // map is still intact (it is consumed by the swap below). Only when enabled.
  if (m_reconnectDroppedAncestors)
    buildDroppedAncestorRedirect();

  // fill the map with the final mother-daughter relationship
  idsave.swap(ancestorList);
  std::stable_sort(idsave.begin(), idsave.end());

  std::vector<std::pair<int, int> >().swap(ancestorList);

  // to get a backward compatible order
  std::stable_sort(m_trackContainer.begin(), m_trackContainer.end(), trkIDLess());

  // to reset the GenParticle ID of a SimTrack to its pre-LHCTransport value
  resetGenID();

  reallyStoreTracks();
}

void SimTrackManager::reallyStoreTracks() {
  // loop over the (now ordered) vector and really save the tracks
#ifdef DebugLog
  edm::LogVerbatim("SimTrackManager") << "reallyStoreTracks() NtracksWithHistory= " << m_trackContainer.size();
#endif

  int nn = m_endPoints.size();
  for (auto& trkH : m_trackContainer) {
    // at this stage there is one vertex per track,
    // so the vertex id of track N is also N
    int iParentID = trkH->parentID();
    int ig = trkH->genParticleID();  // filled only for primary tracks
    bool isBackScatter = trkH->isFromBackScattering();
    bool isPrimary = trkH->isPrimary();
    int primaryGenPartId = trkH->getPrimaryID();  // filled if the G4Track had this info
    int ivertex = getOrCreateVertex(trkH, iParentID);

    auto ptr = trkH;
    if (0 < iParentID) {
      for (auto& trk : m_trackContainer) {
        if (trk->trackID() == iParentID) {
          ptr = trk;
          break;
        }
      }
    }
    // Track at surface is the track at intersection point between tracker and calo
    // envelops if exist. If not exist the momentum is zero, position is the end of
    // the track
    const math::XYZVectorD& pm = ptr->momentum();
    math::XYZVectorD spos(0., 0., 0.);
    math::XYZTLorentzVectorD smom(0., 0., 0., 0.);
    int id = trkH->trackID();
    if (trkH->crossedBoundary()) {
      spos = trkH->getPositionAtBoundary();
      smom = trkH->getMomentumAtBoundary();
    } else {
      for (int i = 0; i < nn; ++i) {
        if (id == m_endPoints[i].first) {
          spos = m_endPoints[i].second;
          break;
        }
      }
    }

    TmpSimTrack* g4simtrack =
        new TmpSimTrack(id, trkH->particleID(), trkH->momentum(), trkH->totalEnergy(), ivertex, ig, pm, spos, smom);
    g4simtrack->copyCrossedBoundaryVars(trkH);
    if (isBackScatter)
      g4simtrack->setFromBackScattering();
    if (isPrimary)
      g4simtrack->setIsPrimary();
    g4simtrack->setGenParticleID(primaryGenPartId);
    m_simEvent->addTrack(g4simtrack);
  }
}

std::unordered_map<int, int> SimTrackManager::computeDroppedAncestorRedirect(
    const std::vector<std::pair<int, int> >& storedTracks, const std::unordered_map<int, int>& parentOfAll) {
  std::unordered_map<int, int> redirect;
  if (storedTracks.empty())
    return redirect;

  // The set of tracks that will be persisted as SimTracks this event.
  std::unordered_set<int> savedIds;
  savedIds.reserve(2 * storedTracks.size());
  for (auto const& [trackId, parentId] : storedTracks)
    savedIds.insert(trackId);

  for (auto const& [trackId, parentId] : storedTracks) {
    // Primary, or immediate parent already stored: getOrCreateVertex resolves it.
    if (parentId <= 0 || savedIds.count(parentId) != 0)
      continue;

    // Walk up through the dropped intermediates (parentOfAll holds the complete
    // topology) to the nearest stored ancestor. The chain terminates at the
    // primary (always stored), so it resolves; the step cap is a defensive guard
    // against a malformed parent loop.
    int p = parentId;
    std::size_t guard = 0;
    const std::size_t maxSteps = parentOfAll.size() + 1;
    while (p > 0 && savedIds.count(p) == 0 && guard++ < maxSteps) {
      auto it = parentOfAll.find(p);
      if (it == parentOfAll.end()) {
        p = 0;
        break;
      }
      p = it->second;
    }
    if (p > 0 && savedIds.count(p) != 0)
      redirect.emplace(trackId, p);
  }
  return redirect;
}

void SimTrackManager::buildDroppedAncestorRedirect() {
  m_droppedParentRedirect.clear();
  if (m_trackContainer.empty())
    return;

  std::vector<std::pair<int, int> > storedTracks;
  storedTracks.reserve(m_trackContainer.size());
  for (auto const& trk : m_trackContainer)
    storedTracks.emplace_back(trk->trackID(), trk->parentID());

  m_droppedParentRedirect = computeDroppedAncestorRedirect(storedTracks, m_parentOfAll);
}

int SimTrackManager::getOrCreateVertex(TrackWithHistory* trkH, int iParentID) {
  int parent = -1;
  for (auto const& trk : m_trackContainer) {
    int id = trk->trackID();
    if (id == iParentID) {
      parent = id;
      break;
    }
  }
  // The immediate parent track was dropped (e.g. a sub-PersistencyEmin
  // intermediate). Reattach this production vertex to the nearest stored ancestor
  // so it does not become an orphan; the target is precomputed in
  // buildDroppedAncestorRedirect(). No-op unless the redirect is enabled.
  if (parent < 0 && m_reconnectDroppedAncestors) {
    auto it = m_droppedParentRedirect.find(trkH->trackID());
    if (it != m_droppedParentRedirect.end())
      parent = it->second;
  }

  VertexMap::const_iterator iterator = m_vertexMap.find(parent);
  if (iterator != m_vertexMap.end()) {
    // loop over saved vertices
    for (auto const& xx : m_vertexMap[parent]) {
      if ((trkH->vertexPosition() - xx.second).Mag2() < r_limit2) {
        return xx.first;
      }
    }
  }

  m_simEvent->addVertex(new TmpSimVertex(trkH->vertexPosition(), trkH->time(), parent, trkH->processType()));
  m_vertexMap[parent].push_back(VertexPosition(m_nVertices, trkH->vertexPosition()));
  ++m_nVertices;
  return (m_nVertices - 1);
}

int SimTrackManager::idSavedTrack(int id) const {
  int idMother = id;
  if (id > 0) {
    unsigned int n = idsave.size();
    if (0 < n) {
      int jmax = n - 1;
      int j, id1;

      // first loop forward
      bool notFound = true;
      for (j = 0; j <= jmax; ++j) {
        if ((idsave[j]).first == idMother) {
          id1 = (idsave[j]).second;
          if (0 == id1 || id1 == idMother) {
            return id1;
          }
          jmax = j - 1;
          idMother = id1;
          notFound = false;
          break;
        }
      }
      if (notFound) {
        return 0;
      }

      // recursive loop
      do {
        notFound = true;
        // search ID scan backward
        for (j = jmax; j >= 0; --j) {
          if ((idsave[j]).first == idMother) {
            id1 = (idsave[j]).second;
            if (0 == id1 || id1 == idMother) {
              return id1;
            }
            jmax = j - 1;
            idMother = id1;
            notFound = false;
            break;
          }
        }
        if (notFound) {
          // ID not in the list of saved track - look into ancestors
          jmax = ancestorList.size() - 1;
          for (j = jmax; j >= 0; --j) {
            if ((ancestorList[j]).first == idMother) {
              idMother = (ancestorList[j]).second;
              return idMother;
            }
          }
          return 0;
        }
      } while (!notFound);
    }
  }
  return idMother;
}

void SimTrackManager::fillMotherList() {
  if (!ancestorList.empty() && lastHist > ancestorList.size()) {
    lastHist = ancestorList.size();
    edm::LogError("SimTrackManager") << " SimTrackManager::fillMotherList track index corrupted";
  }
#ifdef DebugLog
  edm::LogVerbatim("SimTrackManager") << "### SimTrackManager::fillMotherList: " << idsave.size()
                                      << " saved; ancestor: " << lastHist << "  " << ancestorList.size();
  for (unsigned int i = 0; i < idsave.size(); ++i) {
    edm::LogVerbatim("SimTrackManager") << " ISV: Track ID = " << (idsave[i]).first
                                        << " Mother ID = " << (idsave[i]).second;
  }
#endif
  for (unsigned int n = lastHist; n < ancestorList.size(); ++n) {
    int theMotherId = idSavedTrack((ancestorList[n]).first);
    ancestorList[n].second = theMotherId;
#ifdef DebugLog
    LogDebug("SimTrackManager") << "Track ID = " << (ancestorList[n]).first
                                << " Mother ID = " << (ancestorList[n]).second;
#endif
  }

  lastHist = ancestorList.size();
  idsave.clear();
}

void SimTrackManager::cleanTracksWithHistory() {
  if (m_trackContainer.empty() && idsave.empty()) {
    return;
  }

#ifdef DebugLog
  LogDebug("SimTrackManager") << "SimTrackManager::cleanTracksWithHistory has " << idsave.size()
                              << " mother-daughter relationships stored with lastTrack = " << lastTrack;
#endif

  if (lastTrack > 0 && lastTrack >= m_trackContainer.size()) {
    lastTrack = 0;
    edm::LogError("SimTrackManager") << " SimTrackManager::cleanTracksWithHistory track index corrupted";
  }

  std::stable_sort(m_trackContainer.begin() + lastTrack, m_trackContainer.end(), trkIDLess());
  std::stable_sort(idsave.begin(), idsave.end());

#ifdef DebugLog
  LogDebug("SimTrackManager") << " SimTrackManager::cleanTracksWithHistory knows " << m_trackContainer.size()
                              << " tracks with history before branching";
  for (unsigned int it = 0; it < m_trackContainer.size(); it++) {
    LogDebug("SimTrackManager") << " 1 - Track in position " << it << " G4 track number "
                                << m_trackContainer[it]->trackID() << " mother " << m_trackContainer[it]->parentID()
                                << " status " << m_trackContainer[it]->saved();
  }
#endif

  for (auto const& t : m_trackContainer) {
    if (t->saved()) {
      saveTrackAndItsBranch(t);
    }
  }
  unsigned int num = lastTrack;
  for (unsigned int it = lastTrack; it < m_trackContainer.size(); ++it) {
    auto t = m_trackContainer[it];
    int g4ID = m_trackContainer[it]->trackID();
    if (t->saved()) {
      if (it > num) {
        m_trackContainer[num] = t;
      }
      ++num;
      for (auto& xx : idsave) {
        if (xx.first == g4ID) {
          xx.second = g4ID;
          break;
        }
      }
    } else {
      delete t;
    }
  }

  m_trackContainer.resize(num);

#ifdef DebugLog
  LogDebug("SimTrackManager") << " AFTER CLEANING, I GET " << m_trackContainer.size()
                              << " tracks to be saved persistently";
  for (unsigned int it = 0; it < m_trackContainer.size(); ++it) {
    LogDebug("SimTrackManager") << " Track in position " << it << " G4 track number " << m_trackContainer[it]->trackID()
                                << " mother " << m_trackContainer[it]->parentID() << " Status "
                                << m_trackContainer[it]->saved() << " id " << m_trackContainer[it]->particleID()
                                << " E(MeV)= " << m_trackContainer[it]->totalEnergy();
  }
#endif

  fillMotherList();
  lastTrack = m_trackContainer.size();
}

void SimTrackManager::resetGenID() {
  if (theLHCTlink == nullptr)
    return;

  for (auto const& trkH : m_trackContainer) {
    int genParticleID = trkH->genParticleID();
    if (genParticleID == -1) {
      continue;
    } else {
      for (auto const& xx : *theLHCTlink) {
        if (xx.afterHector() == genParticleID) {
          trkH->setGenParticleID(xx.beforeHector());
          continue;
        }
      }
    }
  }

  theLHCTlink = nullptr;
}

void SimTrackManager::ReportException(unsigned int id) const {
  throw cms::Exception("Unknown", "SimTrackManager::getTrackByID")
      << "Fail to get track " << id << " from SimTrackManager, container size= " << m_trackContainer.size();
}