Skip to content

Commit a0fdb2e

Browse files
fgrosafgrosa
committed
Speed-up D0 calibration derived data producer
1 parent a8eefb5 commit a0fdb2e

2 files changed

Lines changed: 87 additions & 35 deletions

File tree

‎DPG/Tasks/AOTTrack/D0CalibTables.h‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -242,6 +242,20 @@ static const std::vector<std::string> labelsPtCand = {
242242
// column labels
243243
static const std::vector<std::string> labelsCutVarCand = {"delta inv. mass", "max d0d0", "max pointing angle", "max pointing angle XY", "min cos pointing angle", "min cos pointing angle XY", "min norm decay length", "min norm decay length XY", "min decay length", "min decay length XY"};
244244

245+
// var indices for the cuts
246+
enum CutsIds : uint8_t {
247+
DeltaMass = 0,
248+
MaxD0D0 = 1,
249+
MaxPointingAngle = 2,
250+
MaxPointingAngleXY = 3,
251+
MinCosPointingAngle = 4,
252+
MinCosPointingAngleXY = 5,
253+
MinNormDecayLength = 6,
254+
MinNormDecayLengthXY = 7,
255+
MinDecayLength = 8,
256+
MinDecayLengthXY = 9
257+
};
258+
245259
static constexpr int NBinsPtMl = 10;
246260
// default values for the pT bin edges (can be used to configure histogram axis)
247261
// offset by 1 from the bin numbers in cuts array

‎DPG/Tasks/AOTTrack/derivedDataCreatorD0Calibration.cxx‎

Lines changed: 73 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -192,8 +192,18 @@ struct DerivedDataCreatorD0Calibration {
192192
TTrackMeanOccs const&,
193193
TTrackQa const&)
194194
{
195-
std::map<int, int> selectedCollisions; // map with indices of selected collisions (key: original AOD Collision table index, value: D0 collision index)
196-
std::map<int, int> selectedTracks; // map with indices of selected tracks (key: original AOD Track table index, value: D0 daughter track index)
195+
std::unordered_map<int, int> selectedCollisions; // map with indices of selected collisions (key: original AOD Collision table index, value: D0 collision index)
196+
std::unordered_map<int, int> selectedTracks; // map with indices of selected tracks (key: original AOD Track table index, value: D0 daughter track index)
197+
198+
std::vector<float> bdtScoresD0;
199+
std::vector<float> bdtScoresD0bar;
200+
std::vector<float> featuresCandD0;
201+
std::vector<float> featuresCandD0bar;
202+
bdtScoresD0.resize(3);
203+
bdtScoresD0bar.resize(3);
204+
featuresCandD0.resize(16);
205+
featuresCandD0bar.resize(16);
206+
auto const& cuts = cfgCandCuts.topologicalCuts;
197207

198208
for (auto const& collision : collisions) {
199209

@@ -217,6 +227,13 @@ struct DerivedDataCreatorD0Calibration {
217227
eventIR.setFromLong(bc.globalBC());
218228

219229
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+
220237
for (auto const& trackIndexPos : groupedTrackIndices) {
221238
auto trackPos = trackIndexPos.template track_as<TTracks>();
222239
// track selections
@@ -255,27 +272,40 @@ struct DerivedDataCreatorD0Calibration {
255272
pidTrackPosPion = selectorPion.statusTpcAndTof(trackPos);
256273
}
257274

275+
int iTrack{-1};
258276
for (auto const& trackIndexNeg : groupedTrackIndices) {
277+
iTrack++;
278+
if (!isSelectedNeg[iTrack]) {
279+
continue;
280+
}
259281
auto trackNeg = trackIndexNeg.template track_as<TTracks>();
260282
// track selections
261283
if (trackNeg.sign() > 0) { // second negative track
284+
isSelectedNeg[iTrack] = false;
262285
continue;
263286
}
264287
if (!trackNeg.isGlobalTrackWoDCA()) {
288+
isSelectedNeg[iTrack] = false;
265289
continue;
266290
}
267291
if (trackNeg.pt() < cfgTrackCuts.ptMin) {
292+
isSelectedNeg[iTrack] = false;
268293
continue;
269294
}
270295
if (std::abs(trackNeg.eta()) > cfgTrackCuts.absEtaMax) {
296+
isSelectedNeg[iTrack] = false;
271297
continue;
272298
}
273299
auto trackParCovNeg = getTrackParCov(trackNeg);
274-
o2::dataformats::DCA dcaNeg(trackNeg.dcaXY(), trackNeg.dcaZ(), trackNeg.cYY(), trackNeg.cZY(), trackNeg.cZZ());
275-
if (trackNeg.collisionId() != collision.globalIndex()) {
276-
trackParCovNeg.propagateToDCA(primaryVertex, bz, &dcaNeg);
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;
277306
}
278-
if (!isSelectedTrackDca(cfgTrackCuts.binsPt, cfgTrackCuts.limitsDca, trackParCovNeg.getPt(), dcaNeg.getY(), dcaNeg.getZ())) {
307+
if (!isSelectedTrackDca(cfgTrackCuts.binsPt, cfgTrackCuts.limitsDca, trackParCovNeg.getPt(), dcaNeg[iTrack].getY(), dcaNeg[iTrack].getZ())) {
308+
isSelectedNeg[iTrack] = false;
279309
continue;
280310
}
281311

@@ -307,47 +337,45 @@ struct DerivedDataCreatorD0Calibration {
307337
}
308338
}
309339

310-
int pidTrackNegKaon{-1};
311-
int pidTrackNegPion{-1};
312340
if (cfgTrackCuts.usePidTpcOnly) {
313341
/// kaon TPC PID negative daughter
314-
pidTrackNegKaon = selectorKaon.statusTpc(trackNeg);
342+
pidTrackNegKaon[iTrack] = selectorKaon.statusTpc(trackNeg);
315343
/// pion TPC PID negative daughter
316-
pidTrackNegPion = selectorPion.statusTpc(trackNeg);
344+
pidTrackNegPion[iTrack] = selectorPion.statusTpc(trackNeg);
317345
} else {
318346
/// kaon TPC, TOF PID negative daughter
319-
pidTrackNegKaon = selectorKaon.statusTpcAndTof(trackNeg);
347+
pidTrackNegKaon[iTrack] = selectorKaon.statusTpcAndTof(trackNeg);
320348
/// pion TPC, TOF PID negative daughter
321-
pidTrackNegPion = selectorPion.statusTpcAndTof(trackNeg);
349+
pidTrackNegPion[iTrack] = selectorPion.statusTpcAndTof(trackNeg);
322350
}
323351

324352
// PID
325353
uint8_t massHypo{D0MassHypo::D0AndD0Bar}; // both mass hypotheses a priori
326-
if (pidTrackPosPion == TrackSelectorPID::Rejected || pidTrackNegKaon == TrackSelectorPID::Rejected) {
354+
if (pidTrackPosPion == TrackSelectorPID::Rejected || pidTrackNegKaon[iTrack] == TrackSelectorPID::Rejected) {
327355
massHypo -= D0MassHypo::D0; // exclude D0
328356
}
329-
if (pidTrackNegPion == TrackSelectorPID::Rejected || pidTrackPosKaon == TrackSelectorPID::Rejected) {
357+
if (pidTrackNegPion[iTrack] == TrackSelectorPID::Rejected || pidTrackPosKaon == TrackSelectorPID::Rejected) {
330358
massHypo -= D0MassHypo::D0Bar; // exclude D0Bar
331359
}
332360
if (massHypo == 0) {
333361
continue;
334362
}
335363

336364
// d0xd0
337-
if (dcaPos.getY() * dcaNeg.getY() > cfgCandCuts.topologicalCuts->get(ptBinNoVtxD0, "max d0d0")) {
365+
if (dcaPos.getY() * dcaNeg[iTrack].getY() > cuts->get(ptBinNoVtxD0, CutsIds::MaxD0D0)) {
338366
continue;
339367
}
340368

341369
// invariant mass
342370
if (massHypo == D0MassHypo::D0 || massHypo == D0MassHypo::D0AndD0Bar) {
343371
float invMassNoVtxD0 = RecoDecay::m(std::array{trackPos.pVector(), trackNeg.pVector()}, std::array{o2::constants::physics::MassPiPlus, o2::constants::physics::MassKPlus});
344-
if (std::abs(invMassNoVtxD0 - o2::constants::physics::MassD0) > cfgCandCuts.topologicalCuts->get(ptBinNoVtxD0, "delta inv. mass") + invMassTolerance) {
372+
if (std::abs(invMassNoVtxD0 - o2::constants::physics::MassD0) > cuts->get(ptBinNoVtxD0, CutsIds::DeltaMass) + invMassTolerance) {
345373
massHypo -= D0MassHypo::D0;
346374
}
347375
}
348376
if (massHypo >= D0MassHypo::D0Bar) {
349377
float invMassNoVtxD0bar = RecoDecay::m(std::array{trackNeg.pVector(), trackPos.pVector()}, std::array{o2::constants::physics::MassPiPlus, o2::constants::physics::MassKPlus});
350-
if (std::abs(invMassNoVtxD0bar - o2::constants::physics::MassD0) > cfgCandCuts.topologicalCuts->get(ptBinNoVtxD0, "delta inv. mass") + invMassTolerance) {
378+
if (std::abs(invMassNoVtxD0bar - o2::constants::physics::MassD0) > cuts->get(ptBinNoVtxD0, CutsIds::DeltaMass) + invMassTolerance) {
351379
massHypo -= D0MassHypo::D0Bar;
352380
}
353381
}
@@ -383,64 +411,72 @@ struct DerivedDataCreatorD0Calibration {
383411
}
384412

385413
// d0xd0
386-
if (dcaPos.getY() * dcaNeg.getY() > cfgCandCuts.topologicalCuts->get(ptBinD0, "max d0d0")) {
414+
if (dcaPos.getY() * dcaNeg[iTrack].getY() > cuts->get(ptBinD0, CutsIds::MaxD0D0)) {
387415
continue;
388416
}
389417
// cospa
390418
float cosPaD0 = RecoDecay::cpa(std::array{primaryVertex.getX(), primaryVertex.getY(), primaryVertex.getZ()}, secondaryVertex, pVecD0);
391-
if (cosPaD0 < cfgCandCuts.topologicalCuts->get(ptBinD0, "min cos pointing angle")) {
419+
if (cosPaD0 < cuts->get(ptBinD0, CutsIds::MinCosPointingAngle)) {
392420
continue;
393421
}
394422
// pointing angle
395423
float paD0 = std::acos(cosPaD0);
396-
if (paD0 > cfgCandCuts.topologicalCuts->get(ptBinD0, "max pointing angle")) {
424+
if (paD0 > cuts->get(ptBinD0, CutsIds::MaxPointingAngle)) {
397425
continue;
398426
}
399427
// cospa XY
400428
float cosPaXYD0 = RecoDecay::cpaXY(std::array{primaryVertex.getX(), primaryVertex.getY(), primaryVertex.getZ()}, secondaryVertex, pVecD0);
401-
if (cosPaXYD0 < cfgCandCuts.topologicalCuts->get(ptBinD0, "min cos pointing angle XY")) {
429+
if (cosPaXYD0 < cuts->get(ptBinD0, CutsIds::MinCosPointingAngleXY)) {
402430
continue;
403431
}
404432
// pointing angle XY
405433
float paXYD0 = std::acos(cosPaXYD0);
406-
if (paXYD0 > cfgCandCuts.topologicalCuts->get(ptBinD0, "max pointing angle XY")) {
434+
if (paXYD0 > cuts->get(ptBinD0, CutsIds::MaxPointingAngleXY)) {
407435
continue;
408436
}
409437
// decay length
410438
float decLenD0 = RecoDecay::distance(std::array{primaryVertex.getX(), primaryVertex.getY(), primaryVertex.getZ()}, secondaryVertex);
411-
if (decLenD0 < cfgCandCuts.topologicalCuts->get(ptBinD0, "min decay length")) {
439+
if (decLenD0 < cuts->get(ptBinD0, CutsIds::MinDecayLength)) {
412440
continue;
413441
}
414442
// decay length XY
415443
float decLenXYD0 = RecoDecay::distanceXY(std::array{primaryVertex.getX(), primaryVertex.getY(), primaryVertex.getZ()}, secondaryVertex);
416-
if (decLenXYD0 < cfgCandCuts.topologicalCuts->get(ptBinD0, "min decay length XY")) {
444+
if (decLenXYD0 < cuts->get(ptBinD0, CutsIds::MinDecayLengthXY)) {
417445
continue;
418446
}
419447
// normalised decay length
420448
float phi{0.f}, theta{0.f};
421449
getPointDirection(std::array{primaryVertex.getX(), primaryVertex.getY(), primaryVertex.getZ()}, secondaryVertex, phi, theta);
422450
float errorDecayLengthD0 = std::sqrt(getRotatedCovMatrixXX(covMatrixPV, phi, theta) + getRotatedCovMatrixXX(covMatrixPCA, phi, theta));
423-
if (decLenD0 / errorDecayLengthD0 < cfgCandCuts.topologicalCuts->get(ptBinD0, "min norm decay length")) {
451+
if (decLenD0 / errorDecayLengthD0 < cuts->get(ptBinD0, CutsIds::MinNormDecayLength)) {
424452
continue;
425453
}
426454
// normalised decay length XY
427455
float errorDecayLengthXYD0 = std::sqrt(getRotatedCovMatrixXX(covMatrixPV, phi, 0.f) + getRotatedCovMatrixXX(covMatrixPCA, phi, 0.f));
428-
if (decLenXYD0 / errorDecayLengthXYD0 < cfgCandCuts.topologicalCuts->get(ptBinD0, "min norm decay length XY")) {
456+
if (decLenXYD0 / errorDecayLengthXYD0 < cuts->get(ptBinD0, CutsIds::MinNormDecayLengthXY)) {
429457
continue;
430458
}
431459

432460
float invMassD0{0.f}, invMassD0bar{0.f};
433461
float cosThetaStarD0{0.f}, cosThetaStarD0bar{0.f};
434-
std::vector<float> bdtScoresD0{0.f, 1.f, 1.f}, bdtScoresD0bar{0.f, 1.f, 1.f}; // always selected a priori
462+
// always selected a priori
463+
bdtScoresD0[0] = 0.f;
464+
bdtScoresD0[1] = 1.f;
465+
bdtScoresD0[2] = 1.f;
466+
bdtScoresD0bar[0] = 0.f;
467+
bdtScoresD0bar[1] = 1.f;
468+
bdtScoresD0bar[2] = 1.f;
435469
if (massHypo == D0MassHypo::D0 || massHypo == D0MassHypo::D0AndD0Bar) {
436470
invMassD0 = RecoDecay::m(std::array{pVecPos, pVecNeg}, std::array{o2::constants::physics::MassPiPlus, o2::constants::physics::MassKPlus});
437-
if (std::abs(invMassD0 - o2::constants::physics::MassD0) > cfgCandCuts.topologicalCuts->get(ptBinD0, "delta inv. mass")) {
471+
if (std::abs(invMassD0 - o2::constants::physics::MassD0) > cuts->get(ptBinD0, CutsIds::DeltaMass)) {
438472
massHypo -= D0MassHypo::D0;
439-
bdtScoresD0 = std::vector<float>{1.f, 0.f, 0.f};
473+
bdtScoresD0[0] = 1.f;
474+
bdtScoresD0[1] = 0.f;
475+
bdtScoresD0[2] = 0.f;
440476
} else {
441477
// apply BDT models
442478
if (cfgMl.apply) {
443-
std::vector<float> featuresCandD0 = {dcaPos.getY(), dcaNeg.getY(), chi2PCA, cosPaD0, cosPaXYD0, decLenXYD0, decLenD0, dcaPos.getY() * dcaNeg.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())};
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())};
444480
if (!mlResponse.isSelectedMl(featuresCandD0, ptD0, bdtScoresD0)) {
445481
massHypo -= D0MassHypo::D0;
446482
} else { // selected, we compute cost*
@@ -457,13 +493,15 @@ struct DerivedDataCreatorD0Calibration {
457493
}
458494
if (massHypo >= D0MassHypo::D0Bar) {
459495
invMassD0bar = RecoDecay::m(std::array{pVecNeg, pVecPos}, std::array{o2::constants::physics::MassPiPlus, o2::constants::physics::MassKPlus});
460-
if (std::abs(invMassD0bar - o2::constants::physics::MassD0) > cfgCandCuts.topologicalCuts->get(ptBinD0, "delta inv. mass")) {
496+
if (std::abs(invMassD0bar - o2::constants::physics::MassD0) > cuts->get(ptBinD0, CutsIds::DeltaMass)) {
461497
massHypo -= D0MassHypo::D0Bar;
462-
bdtScoresD0bar = std::vector<float>{1.f, 0.f, 0.f};
498+
bdtScoresD0bar[0] = 1.f;
499+
bdtScoresD0bar[1] = 0.f;
500+
bdtScoresD0bar[2] = 0.f;
463501
} else {
464502
// apply BDT models
465503
if (cfgMl.apply) {
466-
std::vector<float> featuresCandD0bar = {dcaPos.getY(), dcaNeg.getY(), chi2PCA, cosPaD0, cosPaXYD0, decLenXYD0, decLenD0, dcaPos.getY() * dcaNeg.getY(), aod::pid_tpc_tof_utils::combineNSigma<false>(trackNeg.tpcNSigmaPi(), trackNeg.tofNSigmaPi()), aod::pid_tpc_tof_utils::combineNSigma<false>(trackPos.tpcNSigmaKa(), trackPos.tofNSigmaKa()), trackNeg.tpcNSigmaPi(), trackNeg.tpcNSigmaKa(), aod::pid_tpc_tof_utils::combineNSigma<false>(trackNeg.tpcNSigmaKa(), trackNeg.tofNSigmaKa()), trackPos.tpcNSigmaPi(), trackPos.tpcNSigmaKa(), aod::pid_tpc_tof_utils::combineNSigma<false>(trackPos.tpcNSigmaPi(), trackPos.tofNSigmaPi())};
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())};
467505
if (!mlResponse.isSelectedMl(featuresCandD0bar, ptD0, bdtScoresD0bar)) {
468506
massHypo -= D0MassHypo::D0Bar;
469507
} else { // selected, we compute cost*
@@ -823,8 +861,8 @@ struct DerivedDataCreatorD0Calibration {
823861
deltaRefGloParamQ2Pt,
824862
deltaTOFdX,
825863
deltaTOFdZ,
826-
o2::math_utils::detail::truncateFloatFraction(dcaNeg.getY(), precisionDcaMask),
827-
o2::math_utils::detail::truncateFloatFraction(dcaNeg.getZ(), precisionDcaMask),
864+
o2::math_utils::detail::truncateFloatFraction(dcaNeg[iTrack].getY(), precisionDcaMask),
865+
o2::math_utils::detail::truncateFloatFraction(dcaNeg[iTrack].getZ(), precisionDcaMask),
828866
getCompressedNumSigmaPid(trackNeg.tpcNSigmaPi()),
829867
getCompressedNumSigmaPid(trackNeg.tpcNSigmaKa()),
830868
getCompressedNumSigmaPid(trackNeg.tofNSigmaPi()),

0 commit comments

Comments
 (0)