Skip to content

Commit fd8302a

Browse files
fgrosaalibuild
fgrosa
and
alibuild
authored
[DPG] Improve CPU in D0 data calibration task (#18469)
Co-authored-by: ALICE Action Bot <alibuild@cern.ch>
1 parent 175b110 commit fd8302a

1 file changed

Lines changed: 61 additions & 87 deletions

File tree

‎DPG/Tasks/AOTTrack/derivedDataCreatorD0Calibration.cxx‎

Lines changed: 61 additions & 87 deletions
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,14 @@ struct DerivedDataCreatorD0Calibration {
117117
std::string prefix = "ml";
118118
} cfgMl;
119119

120+
struct SelectedTrack {
121+
int globalIndex;
122+
o2::dataformats::DCA dca;
123+
o2::track::TrackParametrizationWithError<float> trackParCov;
124+
int pidPion;
125+
int pidKaon;
126+
};
127+
120128
using TracksWCovExtraPid = soa::Join<aod::Tracks, aod::TracksDCA, aod::TracksCov, aod::TracksExtra, aod::TrackSelection, aod::pidTPCFullPi, aod::pidTOFFullPi, aod::pidTPCFullKa, aod::pidTOFFullKa>;
121129
using TracksWCovExtraTmoPid = soa::Join<aod::Tracks, aod::TracksDCA, aod::TrackToTmo, aod::TracksCov, aod::TracksExtra, aod::TrackSelection, aod::pidTPCFullPi, aod::pidTOFFullPi, aod::pidTPCFullKa, aod::pidTOFFullKa>;
122130
using TracksWCovExtraTmoPidAndQa = soa::Join<aod::Tracks, aod::TracksDCA, aod::TrackToTmo, aod::TrackToTracksQA, aod::TracksCov, aod::TracksExtra, aod::TrackSelection, aod::pidTPCFullPi, aod::pidTOFFullPi, aod::pidTPCFullKa, aod::pidTOFFullKa>;
@@ -187,7 +195,7 @@ struct DerivedDataCreatorD0Calibration {
187195
template <bool withTrackQa, bool withTrackOcc, typename TTrackQa, typename TTracks, typename TTrackMeanOccs>
188196
void runDataCreation(CollisionsWEvSel const& collisions,
189197
aod::TrackAssoc const& trackIndices,
190-
TTracks const&,
198+
TTracks const& tracks,
191199
aod::BCsWithTimestamps const&,
192200
TTrackMeanOccs const&,
193201
TTrackQa const&)
@@ -227,87 +235,65 @@ struct DerivedDataCreatorD0Calibration {
227235
eventIR.setFromLong(bc.globalBC());
228236

229237
auto groupedTrackIndices = trackIndices.sliceBy(trackIndicesPerCollision, collision.globalIndex());
230-
std::vector<o2::dataformats::DCA> dcaNeg; // we cache the negative track, to avoid the same computation for each positive track
231-
dcaNeg.resize(groupedTrackIndices.size());
232-
std::vector<bool> isSelectedNeg(groupedTrackIndices.size(), true);
233-
std::vector<bool> isPropagatedNeg(groupedTrackIndices.size(), false);
234-
std::vector<int> pidTrackNegKaon(groupedTrackIndices.size(), -1);
235-
std::vector<int> pidTrackNegPion(groupedTrackIndices.size(), -1);
236-
237-
for (auto const& trackIndexPos : groupedTrackIndices) {
238-
auto trackPos = trackIndexPos.template track_as<TTracks>();
238+
// we cache tracks
239+
std::vector<SelectedTrack> tracksSelectedPos;
240+
std::vector<SelectedTrack> tracksSelectedNeg;
241+
tracksSelectedPos.reserve(groupedTrackIndices.size() / 2);
242+
tracksSelectedNeg.reserve(groupedTrackIndices.size() / 2);
243+
244+
// first loop on tracks to find the selected ones
245+
for (auto const& trackIndex : groupedTrackIndices) {
246+
auto track = trackIndex.template track_as<TTracks>();
239247
// track selections
240-
if (trackPos.sign() < 0) { // first positive track
241-
continue;
242-
}
243-
if (!trackPos.isGlobalTrackWoDCA()) {
248+
if (!track.isGlobalTrackWoDCA()) {
244249
continue;
245250
}
246-
if (trackPos.pt() < cfgTrackCuts.ptMin) {
251+
if (track.pt() < cfgTrackCuts.ptMin) {
247252
continue;
248253
}
249-
if (std::abs(trackPos.eta()) > cfgTrackCuts.absEtaMax) {
254+
if (std::abs(track.eta()) > cfgTrackCuts.absEtaMax) {
250255
continue;
251256
}
252-
auto trackParCovPos = getTrackParCov(trackPos);
253-
o2::dataformats::DCA dcaPos(trackPos.dcaXY(), trackPos.dcaZ(), trackPos.cYY(), trackPos.cZY(), trackPos.cZZ());
254-
if (trackPos.collisionId() != collision.globalIndex()) {
255-
trackParCovPos.propagateToDCA(primaryVertex, bz, &dcaPos);
257+
auto trackParCov = getTrackParCov(track);
258+
auto dcaTmp = o2::dataformats::DCA(track.dcaXY(), track.dcaZ(), track.cYY(), track.cZY(), track.cZZ());
259+
if (track.collisionId() != collision.globalIndex()) {
260+
trackParCov.propagateToDCA(primaryVertex, bz, &dcaTmp);
256261
}
257-
if (!isSelectedTrackDca(cfgTrackCuts.binsPt, cfgTrackCuts.limitsDca, trackParCovPos.getPt(), dcaPos.getY(), dcaPos.getZ())) {
262+
if (!isSelectedTrackDca(cfgTrackCuts.binsPt, cfgTrackCuts.limitsDca, trackParCov.getPt(), dcaTmp.getY(), dcaTmp.getZ())) {
258263
continue;
259264
}
260-
261-
int pidTrackPosKaon{-1};
262-
int pidTrackPosPion{-1};
265+
int pidTrackKaon{-1};
266+
int pidTrackPion{-1};
263267
if (cfgTrackCuts.usePidTpcOnly) {
264268
/// kaon TPC PID positive daughter
265-
pidTrackPosKaon = selectorKaon.statusTpc(trackPos);
269+
pidTrackKaon = selectorKaon.statusTpc(track);
266270
/// pion TPC PID positive daughter
267-
pidTrackPosPion = selectorPion.statusTpc(trackPos);
271+
pidTrackPion = selectorPion.statusTpc(track);
268272
} else {
269273
/// kaon TPC, TOF PID positive daughter
270-
pidTrackPosKaon = selectorKaon.statusTpcAndTof(trackPos);
274+
pidTrackKaon = selectorKaon.statusTpcAndTof(track);
271275
/// pion TPC, TOF PID positive daughter
272-
pidTrackPosPion = selectorPion.statusTpcAndTof(trackPos);
276+
pidTrackPion = selectorPion.statusTpcAndTof(track);
277+
}
278+
if (track.sign() > 0) { // first positive track
279+
tracksSelectedPos.push_back({trackIndex.trackId(),
280+
dcaTmp,
281+
trackParCov,
282+
pidTrackPion,
283+
pidTrackKaon});
284+
} else if (track.sign() < 0) {
285+
tracksSelectedNeg.push_back({trackIndex.trackId(),
286+
dcaTmp,
287+
trackParCov,
288+
pidTrackPion,
289+
pidTrackKaon});
273290
}
291+
}
274292

275-
int iTrack{-1};
276-
for (auto const& trackIndexNeg : groupedTrackIndices) {
277-
iTrack++;
278-
if (!isSelectedNeg[iTrack]) {
279-
continue;
280-
}
281-
auto trackNeg = trackIndexNeg.template track_as<TTracks>();
282-
// track selections
283-
if (trackNeg.sign() > 0) { // second negative track
284-
isSelectedNeg[iTrack] = false;
285-
continue;
286-
}
287-
if (!trackNeg.isGlobalTrackWoDCA()) {
288-
isSelectedNeg[iTrack] = false;
289-
continue;
290-
}
291-
if (trackNeg.pt() < cfgTrackCuts.ptMin) {
292-
isSelectedNeg[iTrack] = false;
293-
continue;
294-
}
295-
if (std::abs(trackNeg.eta()) > cfgTrackCuts.absEtaMax) {
296-
isSelectedNeg[iTrack] = false;
297-
continue;
298-
}
299-
auto trackParCovNeg = getTrackParCov(trackNeg);
300-
if (!isPropagatedNeg[iTrack]) {
301-
dcaNeg[iTrack] = o2::dataformats::DCA(trackNeg.dcaXY(), trackNeg.dcaZ(), trackNeg.cYY(), trackNeg.cZY(), trackNeg.cZZ());
302-
if (trackNeg.collisionId() != collision.globalIndex()) {
303-
trackParCovNeg.propagateToDCA(primaryVertex, bz, &dcaNeg[iTrack]);
304-
}
305-
isPropagatedNeg[iTrack] = true;
306-
}
307-
if (!isSelectedTrackDca(cfgTrackCuts.binsPt, cfgTrackCuts.limitsDca, trackParCovNeg.getPt(), dcaNeg[iTrack].getY(), dcaNeg[iTrack].getZ())) {
308-
isSelectedNeg[iTrack] = false;
309-
continue;
310-
}
293+
for (auto const& trackSelPos : tracksSelectedPos) {
294+
auto trackPos = tracks.rawIteratorAt(trackSelPos.globalIndex);
295+
for (auto const& trackSelNeg : tracksSelectedNeg) {
296+
auto trackNeg = tracks.rawIteratorAt(trackSelNeg.globalIndex);
311297

312298
// preselections
313299
// pt
@@ -337,32 +323,20 @@ struct DerivedDataCreatorD0Calibration {
337323
}
338324
}
339325

340-
if (cfgTrackCuts.usePidTpcOnly) {
341-
/// kaon TPC PID negative daughter
342-
pidTrackNegKaon[iTrack] = selectorKaon.statusTpc(trackNeg);
343-
/// pion TPC PID negative daughter
344-
pidTrackNegPion[iTrack] = selectorPion.statusTpc(trackNeg);
345-
} else {
346-
/// kaon TPC, TOF PID negative daughter
347-
pidTrackNegKaon[iTrack] = selectorKaon.statusTpcAndTof(trackNeg);
348-
/// pion TPC, TOF PID negative daughter
349-
pidTrackNegPion[iTrack] = selectorPion.statusTpcAndTof(trackNeg);
350-
}
351-
352326
// PID
353327
uint8_t massHypo{D0MassHypo::D0AndD0Bar}; // both mass hypotheses a priori
354-
if (pidTrackPosPion == TrackSelectorPID::Rejected || pidTrackNegKaon[iTrack] == TrackSelectorPID::Rejected) {
328+
if (trackSelPos.pidPion == TrackSelectorPID::Rejected || trackSelNeg.pidKaon == TrackSelectorPID::Rejected) {
355329
massHypo -= D0MassHypo::D0; // exclude D0
356330
}
357-
if (pidTrackNegPion[iTrack] == TrackSelectorPID::Rejected || pidTrackPosKaon == TrackSelectorPID::Rejected) {
331+
if (trackSelNeg.pidPion == TrackSelectorPID::Rejected || trackSelPos.pidKaon == TrackSelectorPID::Rejected) {
358332
massHypo -= D0MassHypo::D0Bar; // exclude D0Bar
359333
}
360334
if (massHypo == 0) {
361335
continue;
362336
}
363337

364338
// d0xd0
365-
if (dcaPos.getY() * dcaNeg[iTrack].getY() > cuts->get(ptBinNoVtxD0, CutsIds::MaxD0D0)) {
339+
if (trackSelPos.dca.getY() * trackSelNeg.dca.getY() > cuts->get(ptBinNoVtxD0, CutsIds::MaxD0D0)) {
366340
continue;
367341
}
368342

@@ -384,7 +358,7 @@ struct DerivedDataCreatorD0Calibration {
384358
}
385359

386360
// reconstruct vertex
387-
if (df.process(trackParCovPos, trackParCovNeg) == 0) {
361+
if (df.process(trackSelPos.trackParCov, trackSelNeg.trackParCov) == 0) {
388362
continue;
389363
}
390364
const auto& secondaryVertex = df.getPCACandidate();
@@ -411,7 +385,7 @@ struct DerivedDataCreatorD0Calibration {
411385
}
412386

413387
// d0xd0
414-
if (dcaPos.getY() * dcaNeg[iTrack].getY() > cuts->get(ptBinD0, CutsIds::MaxD0D0)) {
388+
if (trackSelPos.dca.getY() * trackSelNeg.dca.getY() > cuts->get(ptBinD0, CutsIds::MaxD0D0)) {
415389
continue;
416390
}
417391
// cospa
@@ -476,7 +450,7 @@ struct DerivedDataCreatorD0Calibration {
476450
} else {
477451
// apply BDT models
478452
if (cfgMl.apply) {
479-
featuresCandD0 = {dcaPos.getY(), dcaNeg[iTrack].getY(), chi2PCA, cosPaD0, cosPaXYD0, decLenXYD0, decLenD0, dcaPos.getY() * dcaNeg[iTrack].getY(), aod::pid_tpc_tof_utils::combineNSigma<false>(trackPos.tpcNSigmaPi(), trackPos.tofNSigmaPi()), aod::pid_tpc_tof_utils::combineNSigma<false>(trackNeg.tpcNSigmaKa(), trackNeg.tofNSigmaKa()), trackPos.tpcNSigmaPi(), trackPos.tpcNSigmaKa(), aod::pid_tpc_tof_utils::combineNSigma<false>(trackPos.tpcNSigmaKa(), trackPos.tofNSigmaKa()), trackNeg.tpcNSigmaPi(), trackNeg.tpcNSigmaKa(), aod::pid_tpc_tof_utils::combineNSigma<false>(trackNeg.tpcNSigmaPi(), trackNeg.tofNSigmaPi())};
453+
featuresCandD0 = {trackSelPos.dca.getY(), trackSelNeg.dca.getY(), chi2PCA, cosPaD0, cosPaXYD0, decLenXYD0, decLenD0, trackSelPos.dca.getY() * trackSelNeg.dca.getY(), aod::pid_tpc_tof_utils::combineNSigma<false>(trackPos.tpcNSigmaPi(), trackPos.tofNSigmaPi()), aod::pid_tpc_tof_utils::combineNSigma<false>(trackNeg.tpcNSigmaKa(), trackNeg.tofNSigmaKa()), trackPos.tpcNSigmaPi(), trackPos.tpcNSigmaKa(), aod::pid_tpc_tof_utils::combineNSigma<false>(trackPos.tpcNSigmaKa(), trackPos.tofNSigmaKa()), trackNeg.tpcNSigmaPi(), trackNeg.tpcNSigmaKa(), aod::pid_tpc_tof_utils::combineNSigma<false>(trackNeg.tpcNSigmaPi(), trackNeg.tofNSigmaPi())};
480454
if (!mlResponse.isSelectedMl(featuresCandD0, ptD0, bdtScoresD0)) {
481455
massHypo -= D0MassHypo::D0;
482456
} else { // selected, we compute cost*
@@ -501,7 +475,7 @@ struct DerivedDataCreatorD0Calibration {
501475
} else {
502476
// apply BDT models
503477
if (cfgMl.apply) {
504-
featuresCandD0bar = {dcaPos.getY(), dcaNeg[iTrack].getY(), chi2PCA, cosPaD0, cosPaXYD0, decLenXYD0, decLenD0, dcaPos.getY() * dcaNeg[iTrack].getY(), aod::pid_tpc_tof_utils::combineNSigma<false>(trackPos.tpcNSigmaPi(), trackPos.tofNSigmaPi()), aod::pid_tpc_tof_utils::combineNSigma<false>(trackNeg.tpcNSigmaKa(), trackNeg.tofNSigmaKa()), trackPos.tpcNSigmaPi(), trackPos.tpcNSigmaKa(), aod::pid_tpc_tof_utils::combineNSigma<false>(trackPos.tpcNSigmaKa(), trackPos.tofNSigmaKa()), trackNeg.tpcNSigmaPi(), trackNeg.tpcNSigmaKa(), aod::pid_tpc_tof_utils::combineNSigma<false>(trackNeg.tpcNSigmaPi(), trackNeg.tofNSigmaPi())};
478+
featuresCandD0bar = {trackSelPos.dca.getY(), trackSelNeg.dca.getY(), chi2PCA, cosPaD0, cosPaXYD0, decLenXYD0, decLenD0, trackSelPos.dca.getY() * trackSelNeg.dca.getY(), aod::pid_tpc_tof_utils::combineNSigma<false>(trackPos.tpcNSigmaPi(), trackPos.tofNSigmaPi()), aod::pid_tpc_tof_utils::combineNSigma<false>(trackNeg.tpcNSigmaKa(), trackNeg.tofNSigmaKa()), trackPos.tpcNSigmaPi(), trackPos.tpcNSigmaKa(), aod::pid_tpc_tof_utils::combineNSigma<false>(trackPos.tpcNSigmaKa(), trackPos.tofNSigmaKa()), trackNeg.tpcNSigmaPi(), trackNeg.tpcNSigmaKa(), aod::pid_tpc_tof_utils::combineNSigma<false>(trackNeg.tpcNSigmaPi(), trackNeg.tofNSigmaPi())};
505479
if (!mlResponse.isSelectedMl(featuresCandD0bar, ptD0, bdtScoresD0bar)) {
506480
massHypo -= D0MassHypo::D0Bar;
507481
} else { // selected, we compute cost*
@@ -694,8 +668,8 @@ struct DerivedDataCreatorD0Calibration {
694668
deltaRefGloParamQ2Pt,
695669
deltaTOFdX,
696670
deltaTOFdZ,
697-
o2::math_utils::detail::truncateFloatFraction(dcaPos.getY(), precisionDcaMask),
698-
o2::math_utils::detail::truncateFloatFraction(dcaPos.getZ(), precisionDcaMask),
671+
o2::math_utils::detail::truncateFloatFraction(trackSelPos.dca.getY(), precisionDcaMask),
672+
o2::math_utils::detail::truncateFloatFraction(trackSelPos.dca.getZ(), precisionDcaMask),
699673
getCompressedNumSigmaPid(trackPos.tpcNSigmaPi()),
700674
getCompressedNumSigmaPid(trackPos.tpcNSigmaKa()),
701675
getCompressedNumSigmaPid(trackPos.tofNSigmaPi()),
@@ -861,8 +835,8 @@ struct DerivedDataCreatorD0Calibration {
861835
deltaRefGloParamQ2Pt,
862836
deltaTOFdX,
863837
deltaTOFdZ,
864-
o2::math_utils::detail::truncateFloatFraction(dcaNeg[iTrack].getY(), precisionDcaMask),
865-
o2::math_utils::detail::truncateFloatFraction(dcaNeg[iTrack].getZ(), precisionDcaMask),
838+
o2::math_utils::detail::truncateFloatFraction(trackSelNeg.dca.getY(), precisionDcaMask),
839+
o2::math_utils::detail::truncateFloatFraction(trackSelNeg.dca.getZ(), precisionDcaMask),
866840
getCompressedNumSigmaPid(trackNeg.tpcNSigmaPi()),
867841
getCompressedNumSigmaPid(trackNeg.tpcNSigmaKa()),
868842
getCompressedNumSigmaPid(trackNeg.tofNSigmaPi()),

0 commit comments

Comments
 (0)