diff --git a/PWGHF/TableProducer/trackIndexSkimCreator.cxx b/PWGHF/TableProducer/trackIndexSkimCreator.cxx index 591533cc5aa..c8456183f31 100644 --- a/PWGHF/TableProducer/trackIndexSkimCreator.cxx +++ b/PWGHF/TableProducer/trackIndexSkimCreator.cxx @@ -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 + 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 pVec{}; ///< momentum at the PCA to the primary vertex + std::array 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 + void fillPropagatedProngCache(TCollision const& collision, TTrackIndices const& trackIndices, std::vector& prongs) + { + prongs.clear(); + const auto thisCollId = collision.globalIndex(); + for (auto trackIndex = trackIndices.begin(); trackIndex != trackIndices.end(); ++trackIndex) { + const auto track = trackIndex.template track_as(); + 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 @@ -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; + std::vector prongsPos{}; + std::vector prongsNeg{}; + for (const auto& collision : collisions) { /// retrieve PV contributors for the current collision @@ -2398,39 +2436,36 @@ struct HfTrackIndexSkimCreator { std::optionalsliceByCached(aod::track::collisionId, 0, cache))> groupedTrackIndicesSoftPionsPos; std::optionalsliceByCached(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(); + + // propagate each track to this collision's PV once, instead of once per pair or triplet + fillPropagatedProngCache(collision, groupedTrackIndicesPos1, prongsPos); + fillPropagatedProngCache(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(); + 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) { @@ -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 @@ -2655,18 +2691,13 @@ struct HfTrackIndexSkimCreator { isSelected3ProngCand = 0; } - const auto trackPos2 = trackIndexPos2.template track_as(); - - 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++) { @@ -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 @@ -2930,17 +2962,13 @@ struct HfTrackIndexSkimCreator { isSelected3ProngCand = 0; } - auto trackNeg2 = trackIndexNeg2.template track_as(); - 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++) {