Skip to content

[PWGHF] Cache prongs propagation in trackIndexSkimCreator #17906

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
fgrosa merged 1 commit into from
Sep 14, 2026
+67 −39
Merged

[PWGHF] Cache prongs propagation in trackIndexSkimCreator #17906

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
106 changes: 67 additions & 39 deletions PWGHF/TableProducer/trackIndexSkimCreator.cxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1724,6 +1724,38 @@ struct HfTrackIndexSkimCreator {
}
}

/// One track of the collision under study, propagated to that collision's primary vertex.
/// \tparam TTrackIndex is the iterator type of the track-collision associations
/// \tparam TTrack is the iterator type of the track table
template <typename TTrackIndex, typename TTrack>
struct HfPropagatedProng {
TTrackIndex trackIndex; ///< track-collision association, carries isSelProng and isIdentifiedPid
TTrack track; ///< track row
o2::track::TrackParCov trackParVar; ///< track parameters at the PCA to the primary vertex
std::array<float, 3> pVec{}; ///< momentum at the PCA to the primary vertex
std::array<float, 2> dcaInfo{}; ///< DCA (xy, z) to the primary vertex
};

/// Fill the per-collision cache of tracks propagated to the collision's primary vertex, so that
/// each track is propagated once per collision instead of once per pair or triplet.
/// \param collision is the collision under study
/// \param trackIndices are the track associations of this collision
/// \param prongs is the cache to be filled, in the order of trackIndices
template <typename TTracks, typename TCollision, typename TTrackIndices, typename TProng>
void fillPropagatedProngCache(TCollision const& collision, TTrackIndices const& trackIndices, std::vector<TProng>& prongs)
{
prongs.clear();
const auto thisCollId = collision.globalIndex();
for (auto trackIndex = trackIndices.begin(); trackIndex != trackIndices.end(); ++trackIndex) {
const auto track = trackIndex.template track_as<TTracks>();
auto& prong = prongs.emplace_back(TProng{trackIndex, track, getTrackParCov(track), track.pVector(), {track.dcaXY(), track.dcaZ()}});
if (thisCollId != track.collisionId()) { // this is not the "default" collision for this track, we have to re-propagate it
o2::base::Propagator::Instance()->propagateToDCABxByBz({collision.posX(), collision.posY(), collision.posZ()}, prong.trackParVar, 2.f, noMatCorr, &prong.dcaInfo);
getPxPyPz(prong.trackParVar, prong.pVec);
}
}
}

/// Method to perform selections for 2-prong candidates before vertex reconstruction
/// \param pVecTrack0 is the momentum array of the first daughter track
/// \param pVecTrack1 is the momentum array of the second daughter track
Expand Down Expand Up @@ -2319,6 +2351,12 @@ struct HfTrackIndexSkimCreator {
}
*/

// per-collision caches of the tracks propagated to the primary vertex, declared here so that their capacity is reused
using TrackIndexIterator = decltype(positiveFor2And3Prongs->sliceByCached(aod::track::collisionId, 0, cache).begin());
using PropagatedProng = HfPropagatedProng<TrackIndexIterator, decltype(tracks.rawIteratorAt(0))>;
std::vector<PropagatedProng> prongsPos{};
std::vector<PropagatedProng> prongsNeg{};

for (const auto& collision : collisions) {

/// retrieve PV contributors for the current collision
Expand Down Expand Up @@ -2398,39 +2436,36 @@ struct HfTrackIndexSkimCreator {
std::optional<decltype(positiveSoftPions->sliceByCached(aod::track::collisionId, 0, cache))> groupedTrackIndicesSoftPionsPos;
std::optional<decltype(negativeSoftPions->sliceByCached(aod::track::collisionId, 0, cache))> groupedTrackIndicesSoftPionsNeg;
int lastFilledD0 = -1; // index to be filled in table for D* mesons
for (auto trackIndexPos1 = groupedTrackIndicesPos1.begin(); trackIndexPos1 != groupedTrackIndicesPos1.end(); ++trackIndexPos1) {
const auto trackPos1 = trackIndexPos1.template track_as<TTracks>();

// propagate each track to this collision's PV once, instead of once per pair or triplet
fillPropagatedProngCache<TTracks>(collision, groupedTrackIndicesPos1, prongsPos);
fillPropagatedProngCache<TTracks>(collision, groupedTrackIndicesNeg1, prongsNeg);

for (std::size_t iPos1 = 0; iPos1 < prongsPos.size(); ++iPos1) {
const auto& trackIndexPos1 = prongsPos[iPos1].trackIndex;
const auto& trackPos1 = prongsPos[iPos1].track;
const auto& trackParVarPos1 = prongsPos[iPos1].trackParVar;
const auto& pVecTrackPos1 = prongsPos[iPos1].pVec;
const auto& dcaInfoPos1 = prongsPos[iPos1].dcaInfo;

// retrieve the selection flag that corresponds to this collision
const auto isSelProngPos1 = trackIndexPos1.isSelProng();
const bool sel2ProngStatusPos = TESTBIT(isSelProngPos1, CandidateType::Cand2Prong);
const bool sel3ProngStatusPos1 = TESTBIT(isSelProngPos1, CandidateType::Cand3Prong);

auto trackParVarPos1 = getTrackParCov(trackPos1);
std::array pVecTrackPos1{trackPos1.pVector()};
std::array dcaInfoPos1{trackPos1.dcaXY(), trackPos1.dcaZ()};
if (thisCollId != trackPos1.collisionId()) { // this is not the "default" collision for this track, we have to re-propagate it
o2::base::Propagator::Instance()->propagateToDCABxByBz({collision.posX(), collision.posY(), collision.posZ()}, trackParVarPos1, 2.f, noMatCorr, &dcaInfoPos1);
getPxPyPz(trackParVarPos1, pVecTrackPos1);
}

// first loop over negative tracks
for (auto trackIndexNeg1 = groupedTrackIndicesNeg1.begin(); trackIndexNeg1 != groupedTrackIndicesNeg1.end(); ++trackIndexNeg1) {
const auto trackNeg1 = trackIndexNeg1.template track_as<TTracks>();
for (std::size_t iNeg1 = 0; iNeg1 < prongsNeg.size(); ++iNeg1) {
const auto& trackIndexNeg1 = prongsNeg[iNeg1].trackIndex;
const auto& trackNeg1 = prongsNeg[iNeg1].track;
const auto& trackParVarNeg1 = prongsNeg[iNeg1].trackParVar;
const auto& pVecTrackNeg1 = prongsNeg[iNeg1].pVec;
const auto& dcaInfoNeg1 = prongsNeg[iNeg1].dcaInfo;

// retrieve the selection flag that corresponds to this collision
const auto isSelProngNeg1 = trackIndexNeg1.isSelProng();
const bool sel2ProngStatusNeg = TESTBIT(isSelProngNeg1, CandidateType::Cand2Prong);
const bool sel3ProngStatusNeg1 = TESTBIT(isSelProngNeg1, CandidateType::Cand3Prong);

auto trackParVarNeg1 = getTrackParCov(trackNeg1);
std::array pVecTrackNeg1{trackNeg1.pVector()};
std::array dcaInfoNeg1{trackNeg1.dcaXY(), trackNeg1.dcaZ()};
if (thisCollId != trackNeg1.collisionId()) { // this is not the "default" collision for this track, we have to re-propagate it
o2::base::Propagator::Instance()->propagateToDCABxByBz({collision.posX(), collision.posY(), collision.posZ()}, trackParVarNeg1, 2.f, noMatCorr, &dcaInfoNeg1);
getPxPyPz(trackParVarNeg1, pVecTrackNeg1);
}

uint isSelected2ProngCand = n2ProngBit; // bitmap for checking status of two-prong candidates (1 is true, 0 is rejected)

if (config.debug) {
Expand Down Expand Up @@ -2638,7 +2673,8 @@ struct HfTrackIndexSkimCreator {

if (config.do3Prong && is2ProngCandidateGoodFor3Prong) { // if 3 prongs are enabled and the first 2 tracks are selected for the 3-prong channels
// second loop over positive tracks
for (auto trackIndexPos2 = trackIndexPos1 + 1; trackIndexPos2 != groupedTrackIndicesPos1.end(); ++trackIndexPos2) {
for (std::size_t iPos2 = iPos1 + 1; iPos2 < prongsPos.size(); ++iPos2) {
const auto& trackIndexPos2 = prongsPos[iPos2].trackIndex;

uint isSelected3ProngCand = n3ProngBit;
if (!TESTBIT(trackIndexPos2.isSelProng(), CandidateType::Cand3Prong)) { // continue immediately
Expand All @@ -2655,18 +2691,13 @@ struct HfTrackIndexSkimCreator {
isSelected3ProngCand = 0;
}

const auto trackPos2 = trackIndexPos2.template track_as<TTracks>();

auto trackParVarPos2 = getTrackParCov(trackPos2);
std::array dcaInfoPos2{trackPos2.dcaXY(), trackPos2.dcaZ()};
const auto& trackPos2 = prongsPos[iPos2].track;
const auto& trackParVarPos2 = prongsPos[iPos2].trackParVar;
const auto& dcaInfoPos2 = prongsPos[iPos2].dcaInfo;

// preselection of 3-prong candidates
if (isSelected3ProngCand) {
std::array pVecTrackPos2{trackPos2.pVector()};
if (thisCollId != trackPos2.collisionId()) { // this is not the "default" collision for this track and we still did not re-propagate it, we have to re-propagate it
o2::base::Propagator::Instance()->propagateToDCABxByBz({collision.posX(), collision.posY(), collision.posZ()}, trackParVarPos2, 2.f, noMatCorr, &dcaInfoPos2);
getPxPyPz(trackParVarPos2, pVecTrackPos2);
}
const auto& pVecTrackPos2 = prongsPos[iPos2].pVec;

if (config.debug) {
for (int iDecay3P = 0; iDecay3P < kN3ProngDecays; iDecay3P++) {
Expand Down Expand Up @@ -2913,7 +2944,8 @@ struct HfTrackIndexSkimCreator {
}

// second loop over negative tracks
for (auto trackIndexNeg2 = trackIndexNeg1 + 1; trackIndexNeg2 != groupedTrackIndicesNeg1.end(); ++trackIndexNeg2) {
for (std::size_t iNeg2 = iNeg1 + 1; iNeg2 < prongsNeg.size(); ++iNeg2) {
const auto& trackIndexNeg2 = prongsNeg[iNeg2].trackIndex;

int isSelected3ProngCand = n3ProngBit;
if (!TESTBIT(trackIndexNeg2.isSelProng(), CandidateType::Cand3Prong)) { // continue immediately
Expand All @@ -2930,17 +2962,13 @@ struct HfTrackIndexSkimCreator {
isSelected3ProngCand = 0;
}

auto trackNeg2 = trackIndexNeg2.template track_as<TTracks>();
auto trackParVarNeg2 = getTrackParCov(trackNeg2);
std::array dcaInfoNeg2{trackNeg2.dcaXY(), trackNeg2.dcaZ()};
const auto& trackNeg2 = prongsNeg[iNeg2].track;
const auto& trackParVarNeg2 = prongsNeg[iNeg2].trackParVar;
const auto& dcaInfoNeg2 = prongsNeg[iNeg2].dcaInfo;

// preselection of 3-prong candidates
if (isSelected3ProngCand) {
std::array pVecTrackNeg2{trackNeg2.pVector()};
if (thisCollId != trackNeg2.collisionId()) { // this is not the "default" collision for this track and we still did not re-propagate it, we have to re-propagate it
o2::base::Propagator::Instance()->propagateToDCABxByBz({collision.posX(), collision.posY(), collision.posZ()}, trackParVarNeg2, 2.f, noMatCorr, &dcaInfoNeg2);
getPxPyPz(trackParVarNeg2, pVecTrackNeg2);
}
const auto& pVecTrackNeg2 = prongsNeg[iNeg2].pVec;

if (config.debug) {
for (int iDecay3P = 0; iDecay3P < kN3ProngDecays; iDecay3P++) {
Expand Down
Loading