diff --git a/PWGLF/DataModel/Vtx3BodyTables.h b/PWGLF/DataModel/Vtx3BodyTables.h index efbbb037519..034983095a7 100644 --- a/PWGLF/DataModel/Vtx3BodyTables.h +++ b/PWGLF/DataModel/Vtx3BodyTables.h @@ -326,6 +326,112 @@ DECLARE_SOA_TABLE(McVtx3BodyDatas, "AOD", "MC3BODYDATA", //! using Vtx3BodyDatasCovs = soa::Join; using Vtx3BodyDatasCovsIndexed = soa::Join; +namespace nuclei3body +{ +DECLARE_SOA_INDEX_COLUMN_FULL(Helium3Track, helium3Track, int, Tracks, "_He3"); //! +DECLARE_SOA_INDEX_COLUMN_FULL(Daughter1Track, daughter1Track, int, Tracks, "_D1"); //! +DECLARE_SOA_INDEX_COLUMN_FULL(Daughter2Track, daughter2Track, int, Tracks, "_D2"); //! +DECLARE_SOA_INDEX_COLUMN(Collision, collision); //! + +DECLARE_SOA_COLUMN(SignHe3, signHe3, int8_t); //! He3 charge sign +DECLARE_SOA_COLUMN(SignDaughter1, signDaughter1, int8_t); //! daughter 1 charge sign +DECLARE_SOA_COLUMN(SignDaughter2, signDaughter2, int8_t); //! daughter 2 charge sign +DECLARE_SOA_COLUMN(Pt, pt, float); //! candidate transverse momentum +DECLARE_SOA_COLUMN(Eta, eta, float); //! candidate pseudorapidity +DECLARE_SOA_COLUMN(Phi, phi, float); //! candidate azimuth +DECLARE_SOA_COLUMN(PHe3, pHe3, float); //! magnitude of fitted, charge-corrected He3 momentum +DECLARE_SOA_COLUMN(PDaughter1, pDaughter1, float); //! magnitude of fitted daughter 1 momentum +DECLARE_SOA_COLUMN(PDaughter2, pDaughter2, float); //! magnitude of fitted daughter 2 momentum +DECLARE_SOA_COLUMN(Chi2, chi2, float); //! DCA-fitter chi2 at PCA +DECLARE_SOA_COLUMN(DcaDaughters, dcaDaughters, float); //! square root of DCA-fitter chi2 at PCA +DECLARE_SOA_COLUMN(CosPA, cosPA, float); //! cosine of pointing angle +DECLARE_SOA_COLUMN(DecayLength, decayLength, float); //! PV-to-SV distance + +DECLARE_SOA_COLUMN(DcaXYHe3ToPv, dcaXYHe3ToPv, float); //! He3 DCAxy to PV +DECLARE_SOA_COLUMN(DcaZHe3ToPv, dcaZHe3ToPv, float); //! He3 DCAz to PV +DECLARE_SOA_COLUMN(DcaXYDaughter1ToPv, dcaXYDaughter1ToPv, float); //! daughter 1 DCAxy to PV +DECLARE_SOA_COLUMN(DcaZDaughter1ToPv, dcaZDaughter1ToPv, float); //! daughter 1 DCAz to PV +DECLARE_SOA_COLUMN(DcaXYDaughter2ToPv, dcaXYDaughter2ToPv, float); //! daughter 2 DCAxy to PV +DECLARE_SOA_COLUMN(DcaZDaughter2ToPv, dcaZDaughter2ToPv, float); //! daughter 2 DCAz to PV + +DECLARE_SOA_COLUMN(TpcNSigmaHe3, tpcNSigmaHe3, float); //! custom He3 TPC n-sigma +DECLARE_SOA_COLUMN(TofNSigmaHe3, tofNSigmaHe3, float); //! He3 TOF n-sigma; sentinel without TOF +DECLARE_SOA_COLUMN(TpcNSigmaPiDaughter1, tpcNSigmaPiDaughter1, float); +DECLARE_SOA_COLUMN(TpcNSigmaKaDaughter1, tpcNSigmaKaDaughter1, float); +DECLARE_SOA_COLUMN(TpcNSigmaPrDaughter1, tpcNSigmaPrDaughter1, float); +DECLARE_SOA_COLUMN(TofNSigmaPiDaughter1, tofNSigmaPiDaughter1, float); +DECLARE_SOA_COLUMN(TofNSigmaKaDaughter1, tofNSigmaKaDaughter1, float); +DECLARE_SOA_COLUMN(TofNSigmaPrDaughter1, tofNSigmaPrDaughter1, float); +DECLARE_SOA_COLUMN(TpcNSigmaPiDaughter2, tpcNSigmaPiDaughter2, float); +DECLARE_SOA_COLUMN(TpcNSigmaKaDaughter2, tpcNSigmaKaDaughter2, float); +DECLARE_SOA_COLUMN(TpcNSigmaPrDaughter2, tpcNSigmaPrDaughter2, float); +DECLARE_SOA_COLUMN(TofNSigmaPiDaughter2, tofNSigmaPiDaughter2, float); +DECLARE_SOA_COLUMN(TofNSigmaKaDaughter2, tofNSigmaKaDaughter2, float); +DECLARE_SOA_COLUMN(TofNSigmaPrDaughter2, tofNSigmaPrDaughter2, float); + +DECLARE_SOA_COLUMN(TpcNClsHe3, tpcNClsHe3, uint16_t); +DECLARE_SOA_COLUMN(TpcNClsDaughter1, tpcNClsDaughter1, uint16_t); +DECLARE_SOA_COLUMN(TpcNClsDaughter2, tpcNClsDaughter2, uint16_t); +DECLARE_SOA_COLUMN(TpcCrossedRowsHe3, tpcCrossedRowsHe3, uint16_t); +DECLARE_SOA_COLUMN(TpcCrossedRowsDaughter1, tpcCrossedRowsDaughter1, uint16_t); +DECLARE_SOA_COLUMN(TpcCrossedRowsDaughter2, tpcCrossedRowsDaughter2, uint16_t); +DECLARE_SOA_COLUMN(ItsNClsHe3, itsNClsHe3, uint8_t); +DECLARE_SOA_COLUMN(ItsNClsDaughter1, itsNClsDaughter1, uint8_t); +DECLARE_SOA_COLUMN(ItsNClsDaughter2, itsNClsDaughter2, uint8_t); + +DECLARE_SOA_COLUMN(McLabelHe3, mcLabelHe3, int); +DECLARE_SOA_COLUMN(McLabelDaughter1, mcLabelDaughter1, int); +DECLARE_SOA_COLUMN(McLabelDaughter2, mcLabelDaughter2, int); +DECLARE_SOA_COLUMN(McPdgHe3, mcPdgHe3, int); +DECLARE_SOA_COLUMN(McPdgDaughter1, mcPdgDaughter1, int); +DECLARE_SOA_COLUMN(McPdgDaughter2, mcPdgDaughter2, int); +DECLARE_SOA_COLUMN(McMatchStatus, mcMatchStatus, uint8_t); //! 0 missing label, 1 no common mother, 2 common mother, 3 configured decay +DECLARE_SOA_COLUMN(McMotherLabel, mcMotherLabel, int); +DECLARE_SOA_COLUMN(McMotherPdg, mcMotherPdg, int); +DECLARE_SOA_COLUMN(GenMotherPt, genMotherPt, float); +DECLARE_SOA_COLUMN(GenMotherEta, genMotherEta, float); +DECLARE_SOA_COLUMN(GenMotherPhi, genMotherPhi, float); +DECLARE_SOA_COLUMN(GenDecayLength, genDecayLength, float); +} // namespace nuclei3body + +#define NUCLEI_THREE_BODY_RECO_COLUMNS \ + nuclei3body::Helium3TrackId, nuclei3body::Daughter1TrackId, \ + nuclei3body::Daughter2TrackId, nuclei3body::CollisionId, \ + nuclei3body::SignHe3, nuclei3body::SignDaughter1, nuclei3body::SignDaughter2, \ + nuclei3body::Pt, nuclei3body::Eta, nuclei3body::Phi, \ + nuclei3body::PHe3, nuclei3body::PDaughter1, nuclei3body::PDaughter2, \ + nuclei3body::Chi2, nuclei3body::DcaDaughters, nuclei3body::CosPA, \ + nuclei3body::DecayLength, \ + nuclei3body::DcaXYHe3ToPv, nuclei3body::DcaZHe3ToPv, \ + nuclei3body::DcaXYDaughter1ToPv, nuclei3body::DcaZDaughter1ToPv, \ + nuclei3body::DcaXYDaughter2ToPv, nuclei3body::DcaZDaughter2ToPv, \ + nuclei3body::TpcNSigmaHe3, nuclei3body::TofNSigmaHe3, \ + nuclei3body::TpcNSigmaPiDaughter1, nuclei3body::TpcNSigmaKaDaughter1, \ + nuclei3body::TpcNSigmaPrDaughter1, nuclei3body::TofNSigmaPiDaughter1, \ + nuclei3body::TofNSigmaKaDaughter1, nuclei3body::TofNSigmaPrDaughter1, \ + nuclei3body::TpcNSigmaPiDaughter2, nuclei3body::TpcNSigmaKaDaughter2, \ + nuclei3body::TpcNSigmaPrDaughter2, nuclei3body::TofNSigmaPiDaughter2, \ + nuclei3body::TofNSigmaKaDaughter2, nuclei3body::TofNSigmaPrDaughter2, \ + nuclei3body::TpcNClsHe3, nuclei3body::TpcNClsDaughter1, nuclei3body::TpcNClsDaughter2, \ + nuclei3body::TpcCrossedRowsHe3, nuclei3body::TpcCrossedRowsDaughter1, \ + nuclei3body::TpcCrossedRowsDaughter2, nuclei3body::ItsNClsHe3, \ + nuclei3body::ItsNClsDaughter1, nuclei3body::ItsNClsDaughter2 + +DECLARE_SOA_TABLE(NucleiThreeBodyDatas, "AOD", "NUC3BODYDATA", + o2::soa::Index<>, + NUCLEI_THREE_BODY_RECO_COLUMNS); + +DECLARE_SOA_TABLE(McNucleiThreeBodyDatas, "AOD", "MCNUC3BODY", + o2::soa::Index<>, + NUCLEI_THREE_BODY_RECO_COLUMNS, + nuclei3body::McLabelHe3, nuclei3body::McLabelDaughter1, nuclei3body::McLabelDaughter2, + nuclei3body::McPdgHe3, nuclei3body::McPdgDaughter1, nuclei3body::McPdgDaughter2, + nuclei3body::McMatchStatus, nuclei3body::McMotherLabel, nuclei3body::McMotherPdg, + nuclei3body::GenMotherPt, nuclei3body::GenMotherEta, + nuclei3body::GenMotherPhi, nuclei3body::GenDecayLength); + +#undef NUCLEI_THREE_BODY_RECO_COLUMNS + } // namespace o2::aod #endif // PWGLF_DATAMODEL_VTX3BODYTABLES_H_ diff --git a/PWGLF/TableProducer/Nuspex/CMakeLists.txt b/PWGLF/TableProducer/Nuspex/CMakeLists.txt index be0b3c3fc2e..eeadf523aac 100644 --- a/PWGLF/TableProducer/Nuspex/CMakeLists.txt +++ b/PWGLF/TableProducer/Nuspex/CMakeLists.txt @@ -29,6 +29,11 @@ o2physics_add_dpl_workflow(tracked-hypertriton-reco-task PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2::DCAFitter KFParticle::KFParticle O2::TOFBase O2::DetectorsVertexing O2Physics::EventFilteringUtils COMPONENT_NAME Analysis) +o2physics_add_dpl_workflow(nuclei-three-body-builder + SOURCES nucleiThreeBodyBuilder.cxx + PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2::DCAFitter O2Physics::EventFilteringUtils + COMPONENT_NAME Analysis) + o2physics_add_dpl_workflow(lnn-reco-task SOURCES lnnRecoTask.cxx PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2::DCAFitter diff --git a/PWGLF/TableProducer/Nuspex/nucleiThreeBodyBuilder.cxx b/PWGLF/TableProducer/Nuspex/nucleiThreeBodyBuilder.cxx new file mode 100644 index 00000000000..a3ce4fb2d1a --- /dev/null +++ b/PWGLF/TableProducer/Nuspex/nucleiThreeBodyBuilder.cxx @@ -0,0 +1,760 @@ +// Copyright 2019-2026 CERN and copyright holders of ALICE O2. +// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders. +// All rights not expressly granted are reserved. +// +// This software is distributed under the terms of the GNU General Public +// License v3 (GPL Version 3), copied verbatim in the file "COPYING". +// +// In applying this license CERN does not waive the privileges and immunities +// granted to it by virtue of its status as an Intergovernmental Organization +// or submit itself to any jurisdiction. + +/// \file nucleiThreeBodyBuilder.cxx +/// \brief Generic He3-hadron-hadron secondary-vertex table producer +/// \author ALICE Collaboration + +#include "PWGLF/DataModel/Vtx3BodyTables.h" +#include "PWGLF/Utils/svPoolCreator.h" + +#include "Common/CCDB/EventSelectionParams.h" +#include "Common/Core/RecoDecay.h" +#include "Common/Core/Zorro.h" +#include "Common/Core/ZorroSummary.h" +#include "Common/Core/trackUtilities.h" +#include "Common/DataModel/Centrality.h" +#include "Common/DataModel/EventSelection.h" +#include "Common/DataModel/PIDResponseTOF.h" +#include "Common/DataModel/PIDResponseTPC.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +using namespace o2; +using namespace o2::framework; + +namespace +{ +using Collisions = soa::Join; +using CollisionsPbPb = soa::Join; +using CollisionsMC = soa::Join; +using Tracks = soa::Join; +using TracksMC = soa::Join; + +constexpr int He3Pdg = o2::constants::physics::Pdg::kHelium3; +constexpr float InvalidFloat = -999.f; +constexpr int InvalidLabel = -1; +constexpr int NProngs = 3; +constexpr float He3Charge = 2.f; +constexpr float UseCCDBBzThreshold = -990.f; +constexpr int He3PoolId = 1; +constexpr int Daughter1PoolId = 2; +constexpr int Daughter2PoolId = 3; +constexpr std::array DefaultHe3BetheBloch{-321.34, 0.6539, 1.591, 0.8225, 2.363, 0.09}; +const std::vector BetheBlochNames{"p0", "p1", "p2", "p3", "p4", "resolution"}; +const std::vector He3Name{"He3"}; + +enum MCMatchStatus : uint8_t { + MissingDaughterLabel = 0, + NoCommonImmediateMother, + CommonImmediateMother, + ConfiguredDecay +}; + +bool isSupportedHadronPdg(int pdg) +{ + const int absPdg = std::abs(pdg); + return absPdg == PDG_t::kPiPlus || + absPdg == PDG_t::kKPlus || absPdg == PDG_t::kProton; +} +} // namespace + +struct NucleiThreeBodyBuilder { + Produces outputData; + Produces outputMC; + + Service ccdb; + HistogramRegistry registry{"registry", {}, OutputObjHandlingPolicy::AnalysisObject}; + Zorro zorro; + OutputObj zorroSummary{"zorroSummary"}; + + struct : ConfigurableGroup { + std::string prefix{"event"}; + Configurable requireSel8{"requireSel8", true, "Require sel8 (main pp default)"}; + Configurable maxAbsZvtx{"maxAbsZvtx", 10.f, "Maximum absolute PV z; negative disables"}; + Configurable requireNoITSROFrameBorder{"requireNoITSROFrameBorder", false, "Require kNoITSROFrameBorder"}; + Configurable requireNoTimeFrameBorder{"requireNoTimeFrameBorder", false, "Require kNoTimeFrameBorder"}; + Configurable requireNoSameBunchPileup{"requireNoSameBunchPileup", false, "Require kNoSameBunchPileup"}; + Configurable requireGoodZvtxFT0vsPV{"requireGoodZvtxFT0vsPV", false, "Require kIsGoodZvtxFT0vsPV"}; + } event; + + struct : ConfigurableGroup { + std::string prefix{"pbpb"}; + Configurable minCentFT0C{"minCentFT0C", 0.f, "Minimum FT0C centrality"}; + Configurable maxCentFT0C{"maxCentFT0C", 100.f, "Maximum FT0C centrality"}; + } pbpb; + + struct : ConfigurableGroup { + std::string prefix{"he3"}; + Configurable minTPCInnerParam{"minTPCInnerParam", 0.8f, "Minimum He3 TPC rigidity"}; + Configurable maxAbsTPCNSigma{"maxAbsTPCNSigma", 5.f, "Maximum absolute custom He3 TPC n-sigma"}; + Configurable requirePIDForTracking{"requirePIDForTracking", false, "Require Helium3 or Alpha PID during tracking"}; + Configurable compensatePIDinTracking{"compensatePIDinTracking", true, "Divide TPC inner parameter by two for charge-two tracking PID"}; + Configurable> betheBlochParams{"betheBlochParams", {DefaultHe3BetheBloch.data(), 1, 6, He3Name, BetheBlochNames}, "TPC Bethe-Bloch parameters for He3"}; + } he3; + + struct : ConfigurableGroup { + std::string prefix{"track"}; + Configurable maxAbsEta{"maxAbsEta", 0.9f, "Maximum absolute eta for every daughter"}; + Configurable minTPCCrossedRows{"minTPCCrossedRows", 70, "Minimum TPC crossed rows for every daughter"}; + Configurable minAbsDCAxyToPV{"minAbsDCAxyToPV", 0.f, "Minimum absolute DCAxy to the assigned/candidate PV for every daughter; zero disables"}; + Configurable minAbsDCAzToPV{"minAbsDCAzToPV", 0.f, "Minimum absolute DCAz to the assigned/candidate PV for every daughter; zero disables"}; + } trackSelections; + + struct : ConfigurableGroup { + std::string prefix{"decay"}; + Configurable daughter1Pdg{"daughter1Pdg", PDG_t::kProton, "Signed pi/K/p PDG relative to a positive He3; used for charge pairing, optional PID selection, and MC matching"}; + Configurable daughter2Pdg{"daughter2Pdg", -PDG_t::kPiPlus, "Signed pi/K/p PDG relative to a positive He3; used for charge pairing, optional PID selection, and MC matching"}; + Configurable motherPdg{"motherPdg", 1010020040, "Mother PDG for MC matching; zero accepts either common mother"}; + } decay; + + struct : ConfigurableGroup { + std::string prefix{"vertex"}; + Configurable minCosPA{"minCosPA", 0.9f, "Minimum cosine of pointing angle; below -1 disables"}; + } vertexSelections; + + struct : ConfigurableGroup { + std::string prefix{"candidate"}; + Configurable applySelections{"applySelections", false, "Apply the additional candidate selections"}; + Configurable maxChi2{"maxChi2", 100.f, "Maximum DCA-fitter chi2"}; + Configurable maxDCADaughters{"maxDCADaughters", 1.f, "Maximum square root of fitter chi2"}; + Configurable maxAbsTPCNSigmaDaughter1{"maxAbsTPCNSigmaDaughter1", 5.f, "Maximum daughter 1 TPC n-sigma for its configured hypothesis"}; + Configurable maxAbsTOFNSigmaDaughter1{"maxAbsTOFNSigmaDaughter1", 5.f, "Maximum daughter 1 TOF n-sigma; applied only with TOF"}; + Configurable maxAbsTPCNSigmaDaughter2{"maxAbsTPCNSigmaDaughter2", 5.f, "Maximum daughter 2 TPC n-sigma for its configured hypothesis"}; + Configurable maxAbsTOFNSigmaDaughter2{"maxAbsTOFNSigmaDaughter2", 5.f, "Maximum daughter 2 TOF n-sigma; applied only with TOF"}; + Configurable minDecayLength{"minDecayLength", 0.f, "Minimum decay length"}; + Configurable maxDecayLength{"maxDecayLength", 100.f, "Maximum decay length"}; + } candidateSelections; + + struct : ConfigurableGroup { + std::string prefix{"pool"}; + Configurable timeMargin{"timeMargin", 800.f, "Track-collision compatibility margin in ns"}; + Configurable skipAmbiguousTracks{"skipAmbiguousTracks", false, "Reject unassigned ambiguous tracks"}; + } pool; + + struct : ConfigurableGroup { + std::string prefix{"zorro"}; + Configurable enabled{"enabled", false, "Enable fHe trigger selection and accounting for skimmed pp data"}; + Configurable ccdbPath{"ccdbPath", "EventFiltering/Zorro/", "Base path of Zorro CCDB objects"}; + Configurable bcTolerance{"bcTolerance", 100, "Zorro BC matching tolerance"}; + } zorroOptions; + + struct : ConfigurableGroup { + std::string prefix{"mc"}; + Configurable requireConfiguredDecayForGen{"requireConfiguredDecayForGen", true, "Fill generated quantities only for the configured daughters and optional mother PDG"}; + } mc; + + Configurable ccdbUrl{"ccdb-url", "http://alice-ccdb.cern.ch", "CCDB URL"}; // o2-linter: disable=name/configurable (Established CCDB option name.) + Configurable grpmagPath{"grpmagPath", "GLO/Config/GRPMagField", "CCDB path of magnetic-field object"}; + Configurable bzInput{"bz", -999.f, "Magnetic field in kG; below -990 uses CCDB"}; // o2-linter: disable=name/configurable (Established magnetic-field option name.) + + o2::vertexing::DCAFitterN<3> fitter; + svPoolCreator3Body tripletPool{He3PoolId, Daughter1PoolId, Daughter2PoolId}; + int runNumber = -1; + float bz = 0.f; + std::array he3BetheBloch{}; + std::vector selectedCollisions; + std::vector collisionsWithHe3; + std::vector identifiedHe3Tracks; + std::vector preselectedTracks; + + struct RecoCandidate { + std::array trackIds{InvalidLabel, InvalidLabel, InvalidLabel}; + int collisionId = InvalidLabel; + std::array signs{}; + float pt = 0.f; + float eta = 0.f; + float phi = 0.f; + std::array daughterMomentum{}; + float chi2 = 0.f; + float dcaDaughters = 0.f; + float cosPA = 0.f; + float decayLength = 0.f; + std::array dcaXYToPV{}; + std::array dcaZToPV{}; + float tpcNSigmaHe3 = InvalidFloat; + float tofNSigmaHe3 = InvalidFloat; + std::array, 2> tpcNSigmaHadron{}; + std::array, 2> tofNSigmaHadron{}; + std::array tpcNCls{}; + std::array tpcCrossedRows{}; + std::array itsNCls{}; + }; + + struct MCInfo { + std::array labels{InvalidLabel, InvalidLabel, InvalidLabel}; + std::array pdgs{}; + uint8_t matchStatus = MissingDaughterLabel; + int motherLabel = InvalidLabel; + int motherPdg = 0; + float motherPt = InvalidFloat; + float motherEta = InvalidFloat; + float motherPhi = InvalidFloat; + float decayLength = InvalidFloat; + }; + + void init(InitContext const&) + { + const int enabledProcesses = static_cast(doprocessData) + static_cast(doprocessPbPb) + static_cast(doprocessMC); + if (enabledProcesses != 1) { + LOG(fatal) << "Enable exactly one of processData, processPbPb, and processMC"; + } + if (zorroOptions.enabled && !doprocessData) { + LOG(fatal) << "Zorro is supported only by the pp data process path"; + } + if (!isSupportedHadronPdg(decay.daughter1Pdg) || !isSupportedHadronPdg(decay.daughter2Pdg)) { + LOG(fatal) << "Signed daughter PDG codes must identify a pion, kaon, or proton"; + } + + ccdb->setURL(ccdbUrl); + ccdb->setCaching(true); + ccdb->setLocalObjectValidityChecking(); + ccdb->setFatalWhenNull(false); + + zorroSummary.setObject(zorro.getZorroSummary()); + zorro.setBaseCCDBPath(zorroOptions.ccdbPath); + zorro.setBCtolerance(zorroOptions.bcTolerance); + + fitter.setPropagateToPCA(true); + fitter.setMaxR(200.f); + fitter.setMinParamChange(1.e-3f); + fitter.setMinRelChi2Change(0.9f); + fitter.setMaxDZIni(1.e9f); + fitter.setMaxChi2(1.e9f); + fitter.setUseAbsDCA(true); + + tripletPool.setTimeMargin(pool.timeMargin); + if (pool.skipAmbiguousTracks) { + tripletPool.setSkipAmbiTracks(); + } + for (size_t i = 0; i < he3BetheBloch.size(); ++i) { + he3BetheBloch[i] = he3.betheBlochParams->get("He3", BetheBlochNames[i].c_str()); + } + + registry.add("events", "Event and He3 preselection;step;events", HistType::kTH1D, {{5, -0.5, 4.5}}); + auto events = registry.get(HIST("events")); + events->GetXaxis()->SetBinLabel(1, "all"); + events->GetXaxis()->SetBinLabel(2, "Zorro fHe"); + events->GetXaxis()->SetBinLabel(3, "event selected"); + events->GetXaxis()->SetBinLabel(4, "identified He3"); + events->GetXaxis()->SetBinLabel(5, "candidate filled"); + registry.add("zorroEvents", "Zorro accounting;selection;events", HistType::kTH1D, {{2, -0.5, 1.5}}); + auto zorroEvents = registry.get(HIST("zorroEvents")); + zorroEvents->GetXaxis()->SetBinLabel(1, "fHe before event selection"); + zorroEvents->GetXaxis()->SetBinLabel(2, "fHe after event selection"); + registry.add("triplets", "Triplet building;step;count", HistType::kTH1D, {{3, -0.5, 2.5}}); + auto triplets = registry.get(HIST("triplets")); + triplets->GetXaxis()->SetBinLabel(1, "pool combinations"); + triplets->GetXaxis()->SetBinLabel(2, "fit succeeded"); + triplets->GetXaxis()->SetBinLabel(3, "stored"); + } + + void initCCDB(aod::BCsWithTimestamps::iterator const& bc) + { + if (runNumber == bc.runNumber()) { + return; + } + if (zorroOptions.enabled) { + zorro.initCCDB(ccdb.service, bc.runNumber(), bc.timestamp(), "fHe"); + zorro.populateHistRegistry(registry, bc.runNumber()); + } + auto* grpmag = ccdb->getForTimeStamp(grpmagPath, bc.timestamp()); + if (!grpmag) { + LOG(fatal) << "Missing magnetic-field object " << grpmagPath << " for timestamp " << bc.timestamp(); + } + o2::base::Propagator::initFieldFromGRP(grpmag); + bz = bzInput < UseCCDBBzThreshold ? o2::base::Propagator::Instance()->getNominalBz() : bzInput; + o2::base::Propagator::Instance()->setNominalBz(bz); + fitter.setBz(bz); + runNumber = bc.runNumber(); + } + + template + float correctedHe3Rigidity(TTrack const& track) + { + const bool chargeTwoPID = track.pidForTracking() == o2::track::PID::Helium3 || track.pidForTracking() == o2::track::PID::Alpha; + return chargeTwoPID && he3.compensatePIDinTracking ? track.tpcInnerParam() / He3Charge : track.tpcInnerParam(); + } + + template + float tpcNSigmaHe3(TTrack const& track) + { + const float rigidity = correctedHe3Rigidity(track); + const float expected = o2::common::BetheBlochAleph(static_cast(rigidity * He3Charge / constants::physics::MassHelium3), + he3BetheBloch[0], he3BetheBloch[1], he3BetheBloch[2], + he3BetheBloch[3], he3BetheBloch[4]); + return (track.tpcSignal() - expected) / (expected * he3BetheBloch[5]); + } + + template + bool isIdentifiedHe3(TTrack const& track) + { + if (!track.hasTPC() || correctedHe3Rigidity(track) < he3.minTPCInnerParam || + std::abs(tpcNSigmaHe3(track)) > he3.maxAbsTPCNSigma) { + return false; + } + if (he3.requirePIDForTracking && + track.pidForTracking() != o2::track::PID::Helium3 && + track.pidForTracking() != o2::track::PID::Alpha) { + return false; + } + return true; + } + + template + bool passesEventSelection(TCollision const& collision) + { + if (event.requireSel8 && !collision.sel8()) { + return false; + } + if (event.maxAbsZvtx >= 0.f && std::abs(collision.posZ()) > event.maxAbsZvtx) { + return false; + } + if (event.requireNoITSROFrameBorder && !collision.selection_bit(aod::evsel::kNoITSROFrameBorder)) { + return false; + } + if (event.requireNoTimeFrameBorder && !collision.selection_bit(aod::evsel::kNoTimeFrameBorder)) { + return false; + } + if (event.requireNoSameBunchPileup && !collision.selection_bit(aod::evsel::kNoSameBunchPileup)) { + return false; + } + if (event.requireGoodZvtxFT0vsPV && !collision.selection_bit(aod::evsel::kIsGoodZvtxFT0vsPV)) { + return false; + } + return true; + } + + template + void selectCollisions(TCollisions const& collisions, bool applyZorro) + { + selectedCollisions.assign(collisions.size(), false); + collisionsWithHe3.assign(collisions.size(), false); + for (const auto& collision : collisions) { + registry.fill(HIST("events"), 0.); + const auto bc = collision.template bc_as(); + initCCDB(bc); + + if (applyZorro) { + if (!zorro.isSelected(bc.globalBC(), zorroOptions.bcTolerance)) { + continue; + } + registry.fill(HIST("events"), 1.); + registry.fill(HIST("zorroEvents"), 0.); + } + if (!passesEventSelection(collision)) { + continue; + } + if constexpr (applyCentrality) { + if (collision.centFT0C() < pbpb.minCentFT0C || collision.centFT0C() > pbpb.maxCentFT0C) { + continue; + } + } + selectedCollisions[collision.globalIndex()] = true; + registry.fill(HIST("events"), 2.); + if (applyZorro) { + registry.fill(HIST("zorroEvents"), 1.); + } + } + } + + template + std::array dcaToPV(TCollision const& collision, TTrack const& track) + { + auto trackPar = getTrackParCov(track); + std::array dca{}; + o2::base::Propagator::Instance()->propagateToDCABxByBz( + {collision.posX(), collision.posY(), collision.posZ()}, trackPar, 2.f, fitter.getMatCorrType(), &dca); + return dca; + } + + template + bool passesTrackSelections(TTrack const& track, std::array const& dca) + { + return std::abs(track.eta()) <= trackSelections.maxAbsEta && + track.tpcNClsCrossedRows() >= trackSelections.minTPCCrossedRows && + std::abs(dca[0]) >= trackSelections.minAbsDCAxyToPV && + std::abs(dca[1]) >= trackSelections.minAbsDCAzToPV; + } + + template + static std::array tpcHadronNSigma(TTrack const& track) + { + return {track.tpcNSigmaPi(), track.tpcNSigmaKa(), track.tpcNSigmaPr()}; + } + + template + static std::array tofHadronNSigma(TTrack const& track) + { + if (!track.hasTOF()) { + return {InvalidFloat, InvalidFloat, InvalidFloat}; + } + return {track.tofNSigmaPi(), track.tofNSigmaKa(), track.tofNSigmaPr()}; + } + + static float selectedNSigma(std::array const& values, int pdg) + { + switch (std::abs(pdg)) { + case PDG_t::kPiPlus: + return values[0]; + case PDG_t::kKPlus: + return values[1]; + case PDG_t::kProton: + return values[2]; + default: + return 0.f; + } + } + + template + bool buildCandidate(TCollision const& collision, TTrack const& he3Track, + TTrack const& daughter1Track, TTrack const& daughter2Track, + RecoCandidate& out) + { + const std::array tracks{he3Track, daughter1Track, daughter2Track}; + for (int i = 0; i < NProngs; ++i) { + const auto dca = dcaToPV(collision, tracks[i]); + if (!passesTrackSelections(tracks[i], dca)) { + return false; + } + out.dcaXYToPV[i] = dca[0]; + out.dcaZToPV[i] = dca[1]; + } + + auto he3TrackFit = getTrackParCov(he3Track); + auto daughter1TrackFit = getTrackParCov(daughter1Track); + auto daughter2TrackFit = getTrackParCov(daughter2Track); + int nCandidates = 0; + try { + nCandidates = fitter.process(he3TrackFit, daughter1TrackFit, daughter2TrackFit); + } catch (...) { + LOG(error) << "Exception while fitting a nuclei three-body candidate"; + return false; + } + if (nCandidates == 0) { + return false; + } + registry.fill(HIST("triplets"), 1.); + + out.trackIds = {static_cast(he3Track.globalIndex()), + static_cast(daughter1Track.globalIndex()), + static_cast(daughter2Track.globalIndex())}; + out.collisionId = collision.globalIndex(); + out.signs = {static_cast(he3Track.sign() > 0 ? 1 : -1), + static_cast(daughter1Track.sign() > 0 ? 1 : -1), + static_cast(daughter2Track.sign() > 0 ? 1 : -1)}; + + std::array, 3> daughterMomenta{}; + std::array momentum{}; + for (int i = 0; i < NProngs; ++i) { + fitter.getTrack(i).getPxPyPzGlo(daughterMomenta[i]); + } + for (int component = 0; component < NProngs; ++component) { + daughterMomenta[0][component] *= He3Charge; + momentum[component] = daughterMomenta[0][component] + + daughterMomenta[1][component] + + daughterMomenta[2][component]; + } + for (int i = 0; i < NProngs; ++i) { + out.daughterMomentum[i] = RecoDecay::sqrtSumOfSquares( + daughterMomenta[i][0], daughterMomenta[i][1], daughterMomenta[i][2]); + } + out.pt = RecoDecay::sqrtSumOfSquares(momentum[0], momentum[1]); + out.eta = RecoDecay::eta(momentum); + out.phi = RecoDecay::phi(momentum); + const float momentumMagnitude = RecoDecay::sqrtSumOfSquares(momentum[0], momentum[1], momentum[2]); + + out.chi2 = fitter.getChi2AtPCACandidate(); + out.dcaDaughters = std::sqrt(std::max(0.f, out.chi2)); + const auto& secondaryVertex = fitter.getPCACandidate(); + const std::array flight{ + static_cast(secondaryVertex[0] - collision.posX()), + static_cast(secondaryVertex[1] - collision.posY()), + static_cast(secondaryVertex[2] - collision.posZ())}; + out.decayLength = RecoDecay::sqrtSumOfSquares(flight[0], flight[1], flight[2]); + out.cosPA = (flight[0] * momentum[0] + flight[1] * momentum[1] + flight[2] * momentum[2]) / + (out.decayLength * momentumMagnitude + 1.e-10f); + if (out.cosPA < vertexSelections.minCosPA) { + return false; + } + + out.tpcNSigmaHe3 = tpcNSigmaHe3(he3Track); + out.tofNSigmaHe3 = he3Track.hasTOF() ? he3Track.tofNSigmaHe() : InvalidFloat; + out.tpcNSigmaHadron[0] = tpcHadronNSigma(daughter1Track); + out.tpcNSigmaHadron[1] = tpcHadronNSigma(daughter2Track); + out.tofNSigmaHadron[0] = tofHadronNSigma(daughter1Track); + out.tofNSigmaHadron[1] = tofHadronNSigma(daughter2Track); + for (int i = 0; i < NProngs; ++i) { + out.tpcNCls[i] = tracks[i].tpcNClsFound(); + out.tpcCrossedRows[i] = tracks[i].tpcNClsCrossedRows(); + out.itsNCls[i] = tracks[i].itsNCls(); + } + + if (!candidateSelections.applySelections) { + return true; + } + if (out.chi2 > candidateSelections.maxChi2 || + out.dcaDaughters > candidateSelections.maxDCADaughters || + out.decayLength < candidateSelections.minDecayLength || + out.decayLength > candidateSelections.maxDecayLength || + std::abs(selectedNSigma(out.tpcNSigmaHadron[0], decay.daughter1Pdg)) > candidateSelections.maxAbsTPCNSigmaDaughter1 || + std::abs(selectedNSigma(out.tpcNSigmaHadron[1], decay.daughter2Pdg)) > candidateSelections.maxAbsTPCNSigmaDaughter2 || + (daughter1Track.hasTOF() && + std::abs(selectedNSigma(out.tofNSigmaHadron[0], decay.daughter1Pdg)) > candidateSelections.maxAbsTOFNSigmaDaughter1) || + (daughter2Track.hasTOF() && + std::abs(selectedNSigma(out.tofNSigmaHadron[1], decay.daughter2Pdg)) > candidateSelections.maxAbsTOFNSigmaDaughter2)) { + return false; + } + return true; + } + + void fillDataTable(RecoCandidate const& c) + { + outputData(c.trackIds[0], c.trackIds[1], c.trackIds[2], c.collisionId, + c.signs[0], c.signs[1], c.signs[2], + c.pt, c.eta, c.phi, + c.daughterMomentum[0], c.daughterMomentum[1], c.daughterMomentum[2], + c.chi2, c.dcaDaughters, c.cosPA, c.decayLength, + c.dcaXYToPV[0], c.dcaZToPV[0], + c.dcaXYToPV[1], c.dcaZToPV[1], + c.dcaXYToPV[2], c.dcaZToPV[2], + c.tpcNSigmaHe3, c.tofNSigmaHe3, + c.tpcNSigmaHadron[0][0], c.tpcNSigmaHadron[0][1], c.tpcNSigmaHadron[0][2], + c.tofNSigmaHadron[0][0], c.tofNSigmaHadron[0][1], c.tofNSigmaHadron[0][2], + c.tpcNSigmaHadron[1][0], c.tpcNSigmaHadron[1][1], c.tpcNSigmaHadron[1][2], + c.tofNSigmaHadron[1][0], c.tofNSigmaHadron[1][1], c.tofNSigmaHadron[1][2], + c.tpcNCls[0], c.tpcNCls[1], c.tpcNCls[2], + c.tpcCrossedRows[0], c.tpcCrossedRows[1], c.tpcCrossedRows[2], + c.itsNCls[0], c.itsNCls[1], c.itsNCls[2]); + } + + template + MCInfo getMCInfo(TTrack const& he3Track, TTrack const& daughter1Track, + TTrack const& daughter2Track, aod::McParticles const& mcParticles) + { + MCInfo info; + if (!he3Track.has_mcParticle() || !daughter1Track.has_mcParticle() || !daughter2Track.has_mcParticle()) { + return info; + } + const auto he3Particle = he3Track.template mcParticle_as(); + const auto daughter1Particle = daughter1Track.template mcParticle_as(); + const auto daughter2Particle = daughter2Track.template mcParticle_as(); + info.labels = {static_cast(he3Particle.globalIndex()), + static_cast(daughter1Particle.globalIndex()), + static_cast(daughter2Particle.globalIndex())}; + info.pdgs = {he3Particle.pdgCode(), daughter1Particle.pdgCode(), daughter2Particle.pdgCode()}; + info.matchStatus = NoCommonImmediateMother; + + int commonMotherLabel = InvalidLabel; + for (const auto& he3Mother : he3Particle.template mothers_as()) { + for (const auto& daughter1Mother : daughter1Particle.template mothers_as()) { + if (he3Mother.globalIndex() != daughter1Mother.globalIndex()) { + continue; + } + for (const auto& daughter2Mother : daughter2Particle.template mothers_as()) { + if (he3Mother.globalIndex() == daughter2Mother.globalIndex()) { + commonMotherLabel = he3Mother.globalIndex(); + break; + } + } + } + } + if (commonMotherLabel < 0) { + return info; + } + + const auto selectedMother = mcParticles.rawIteratorAt(commonMotherLabel); + info.matchStatus = CommonImmediateMother; + info.motherLabel = commonMotherLabel; + info.motherPdg = selectedMother.pdgCode(); + const int he3Sign = he3Track.sign() > 0 ? 1 : -1; + const bool expectedDaughters = info.pdgs[0] == he3Sign * He3Pdg && + info.pdgs[1] == he3Sign * decay.daughter1Pdg && + info.pdgs[2] == he3Sign * decay.daughter2Pdg; + const bool motherAccepted = decay.motherPdg == 0 || + info.motherPdg == he3Sign * std::abs(decay.motherPdg.value); + const bool configuredDecay = expectedDaughters && motherAccepted; + if (configuredDecay) { + info.matchStatus = ConfiguredDecay; + } + if ((!mc.requireConfiguredDecayForGen || configuredDecay) && motherAccepted) { + const std::array motherMomentum{selectedMother.px(), selectedMother.py(), selectedMother.pz()}; + info.motherPt = RecoDecay::sqrtSumOfSquares(motherMomentum[0], motherMomentum[1]); + info.motherEta = RecoDecay::eta(motherMomentum); + info.motherPhi = RecoDecay::phi(motherMomentum); + info.decayLength = RecoDecay::sqrtSumOfSquares(he3Particle.vx() - selectedMother.vx(), + he3Particle.vy() - selectedMother.vy(), + he3Particle.vz() - selectedMother.vz()); + } + return info; + } + + void fillMCTable(RecoCandidate const& c, MCInfo const& m) + { + outputMC(c.trackIds[0], c.trackIds[1], c.trackIds[2], c.collisionId, + c.signs[0], c.signs[1], c.signs[2], + c.pt, c.eta, c.phi, + c.daughterMomentum[0], c.daughterMomentum[1], c.daughterMomentum[2], + c.chi2, c.dcaDaughters, c.cosPA, c.decayLength, + c.dcaXYToPV[0], c.dcaZToPV[0], + c.dcaXYToPV[1], c.dcaZToPV[1], + c.dcaXYToPV[2], c.dcaZToPV[2], + c.tpcNSigmaHe3, c.tofNSigmaHe3, + c.tpcNSigmaHadron[0][0], c.tpcNSigmaHadron[0][1], c.tpcNSigmaHadron[0][2], + c.tofNSigmaHadron[0][0], c.tofNSigmaHadron[0][1], c.tofNSigmaHadron[0][2], + c.tpcNSigmaHadron[1][0], c.tpcNSigmaHadron[1][1], c.tpcNSigmaHadron[1][2], + c.tofNSigmaHadron[1][0], c.tofNSigmaHadron[1][1], c.tofNSigmaHadron[1][2], + c.tpcNCls[0], c.tpcNCls[1], c.tpcNCls[2], + c.tpcCrossedRows[0], c.tpcCrossedRows[1], c.tpcCrossedRows[2], + c.itsNCls[0], c.itsNCls[1], c.itsNCls[2], + m.labels[0], m.labels[1], m.labels[2], + m.pdgs[0], m.pdgs[1], m.pdgs[2], + m.matchStatus, m.motherLabel, m.motherPdg, + m.motherPt, m.motherEta, m.motherPhi, m.decayLength); + } + + template + void build(TCollisions const& collisions, TTracks const& tracks, + aod::AmbiguousTracks const& ambiguousTracks, + aod::BCsWithTimestamps const& bcs, + aod::McParticles const* mcParticles = nullptr) + { + selectCollisions(collisions, zorroOptions.enabled && !isMC); + identifiedHe3Tracks.assign(tracks.size(), false); + preselectedTracks.assign(tracks.size(), false); + + // Apply the common track preselection before identifying He3 candidates. + // Consequently, the triplet pools are never built unless a selected event + // contains at least one He3 that also passes the minimal track selections. + for (const auto& track : tracks) { + if (!track.has_collision() || track.collisionId() < 0 || + !selectedCollisions[track.collisionId()]) { + continue; + } + const auto collision = collisions.rawIteratorAt(track.collisionId()); + const auto dca = dcaToPV(collision, track); + if (!passesTrackSelections(track, dca)) { + continue; + } + preselectedTracks[track.globalIndex()] = true; + if (isIdentifiedHe3(track)) { + identifiedHe3Tracks[track.globalIndex()] = true; + collisionsWithHe3[track.collisionId()] = true; + } + } + for (size_t i = 0; i < collisionsWithHe3.size(); ++i) { + if (collisionsWithHe3[i]) { + registry.fill(HIST("events"), 3.); + } + } + + tripletPool.clearPools(); + tripletPool.fillBC2Coll(collisions, bcs); + for (const auto& track : tracks) { + if (!track.has_collision() || track.collisionId() < 0 || + !preselectedTracks[track.globalIndex()] || + !collisionsWithHe3[track.collisionId()]) { + continue; + } + if (identifiedHe3Tracks[track.globalIndex()]) { + tripletPool.appendTrackCand(track, collisions, He3PoolId, ambiguousTracks, bcs); + } else { + tripletPool.appendTrackCand(track, collisions, Daughter1PoolId, ambiguousTracks, bcs); + tripletPool.appendTrackCand(track, collisions, Daughter2PoolId, ambiguousTracks, bcs); + } + } + + const bool daughter1LikeSign = decay.daughter1Pdg > 0; + const bool daughter2LikeSign = decay.daughter2Pdg > 0; + auto& triplets = tripletPool.getSVCandPool(collisions, daughter1LikeSign, daughter2LikeSign); + for (const auto& triplet : triplets) { + registry.fill(HIST("triplets"), 0.); + const auto he3Track = tracks.rawIteratorAt(triplet.tr0Idx); + const auto daughter1Track = tracks.rawIteratorAt(triplet.tr1Idx); + const auto daughter2Track = tracks.rawIteratorAt(triplet.tr2Idx); + for (int collisionIndex = triplet.collBracket.getMin(); collisionIndex <= triplet.collBracket.getMax(); ++collisionIndex) { + if (!selectedCollisions[collisionIndex] || !collisionsWithHe3[collisionIndex]) { + continue; + } + const auto collision = collisions.rawIteratorAt(collisionIndex); + RecoCandidate reco; + if (!buildCandidate(collision, he3Track, daughter1Track, daughter2Track, reco)) { + continue; + } + if constexpr (isMC) { + fillMCTable(reco, getMCInfo(he3Track, daughter1Track, daughter2Track, *mcParticles)); + } else { + fillDataTable(reco); + } + registry.fill(HIST("events"), 4.); + registry.fill(HIST("triplets"), 2.); + } + } + } + + void processData(Collisions const& collisions, Tracks const& tracks, + aod::AmbiguousTracks const& ambiguousTracks, + aod::BCsWithTimestamps const& bcs) + { + build(collisions, tracks, ambiguousTracks, bcs); + } + PROCESS_SWITCH(NucleiThreeBodyBuilder, processData, "Process pp data", true); + + void processPbPb(CollisionsPbPb const& collisions, Tracks const& tracks, + aod::AmbiguousTracks const& ambiguousTracks, + aod::BCsWithTimestamps const& bcs) + { + build(collisions, tracks, ambiguousTracks, bcs); + } + PROCESS_SWITCH(NucleiThreeBodyBuilder, processPbPb, "Process Pb-Pb data", false); + + void processMC(CollisionsMC const& collisions, TracksMC const& tracks, + aod::AmbiguousTracks const& ambiguousTracks, + aod::BCsWithTimestamps const& bcs, + aod::McParticles const& mcParticles) + { + build(collisions, tracks, ambiguousTracks, bcs, &mcParticles); + } + PROCESS_SWITCH(NucleiThreeBodyBuilder, processMC, "Process reconstructed MC", false); +}; + +WorkflowSpec defineDataProcessing(ConfigContext const& cfgc) +{ + return WorkflowSpec{adaptAnalysisTask(cfgc)}; +} diff --git a/PWGLF/Utils/svPoolCreator.h b/PWGLF/Utils/svPoolCreator.h index bcb62b95f61..fedcaa95ce7 100644 --- a/PWGLF/Utils/svPoolCreator.h +++ b/PWGLF/Utils/svPoolCreator.h @@ -9,6 +9,10 @@ // granted to it by virtue of its status as an Intergovernmental Organization // or submit itself to any jurisdiction. +/// \file svPoolCreator.h +/// \brief Time-compatible secondary-vertex track-pair and triplet pools +/// \author ALICE Collaboration + #ifndef PWGLF_UTILS_SVPOOLCREATOR_H_ #define PWGLF_UTILS_SVPOOLCREATOR_H_ @@ -32,6 +36,7 @@ using CollBracket = o2::math_utils::Bracket; constexpr uint64_t bOffsetMax = 241; // track compatibility can never go beyond 6 mus (ITS) +constexpr int NChargeSigns = 2; struct TrackCand { int Idxtr; @@ -44,6 +49,13 @@ struct SVCand { CollBracket collBracket{}; }; +struct SVCand3 { + int tr0Idx; + int tr1Idx; + int tr2Idx; + CollBracket collBracket{}; +}; + class svPoolCreator { public: @@ -61,7 +73,7 @@ class svPoolCreator void clearPools() { - for (auto& pool : trackCandPool) { + for (auto& pool : trackCandPool) { // o2-linter: disable=const-ref-in-for-loop (The pool is cleared in place.) pool.clear(); } tmap.clear(); @@ -200,15 +212,15 @@ class svPoolCreator template std::vector& getSVCandPool(const C& collisions, bool combineLikeSign = false) { - gsl::span> track0Pool{trackCandPool.data(), 2}; - gsl::span> track1Pool{trackCandPool.data() + 2, 2}; - std::array, 2> mVtxTrack0{}; // 1st pos. and neg. track of the kink pool for each vertex + gsl::span> track0Pool{trackCandPool.data(), NChargeSigns}; + gsl::span> track1Pool{trackCandPool.data() + NChargeSigns, NChargeSigns}; + std::array, NChargeSigns> mVtxTrack0{}; // 1st pos. and neg. track of the kink pool for each vertex - for (int i = 0; i < 2; i++) { + for (int i = 0; i < NChargeSigns; i++) { mVtxTrack0[i].clear(); mVtxTrack0[i].resize(collisions.size(), -1); } - for (int iCharge = 0; iCharge < 2; iCharge++) { + for (int iCharge = 0; iCharge < NChargeSigns; iCharge++) { auto& vtxFirstT = mVtxTrack0[iCharge]; const auto& signTrack0Pool = track0Pool[iCharge]; for (unsigned i = 0; i < signTrack0Pool.size(); i++) { @@ -271,4 +283,108 @@ class svPoolCreator TrackCand trForpool; }; +/// Build time-compatible three-track combinations by joining two svPoolCreator +/// pools on a common first prong. The class deliberately only handles +/// combinatorics; the caller remains responsible for fitting and physics +/// selections. +class svPoolCreator3Body +{ + public: + svPoolCreator3Body() = default; + svPoolCreator3Body(int track0Pdg, int track1Pdg, int track2Pdg) + : track0Pdg(track0Pdg), + track1Pdg(track1Pdg), + track2Pdg(track2Pdg), + pool01(track0Pdg, track1Pdg), + pool02(track0Pdg, track2Pdg) + { + } + + void setPDGs(int pdg0, int pdg1, int pdg2) + { + track0Pdg = pdg0; + track1Pdg = pdg1; + track2Pdg = pdg2; + pool01.setPDGs(track0Pdg, track1Pdg); + pool02.setPDGs(track0Pdg, track2Pdg); + } + + void clearPools() + { + pool01.clearPools(); + pool02.clearPools(); + svCandPool.clear(); + } + + void setTimeMargin(float timeMargin) + { + pool01.setTimeMargin(timeMargin); + pool02.setTimeMargin(timeMargin); + } + + void setSkipAmbiTracks() + { + pool01.setSkipAmbiTracks(); + pool02.setSkipAmbiTracks(); + } + + template + void fillBC2Coll(const C& collisions, BC const& bcs) + { + pool01.fillBC2Coll(collisions, bcs); + pool02.fillBC2Coll(collisions, bcs); + } + + template + void appendTrackCand(const T& trackCand, const C& collisions, int pdgHypo, const AT& ambiTracks, BC const& bcs) + { + if (pdgHypo == track0Pdg) { + pool01.appendTrackCand(trackCand, collisions, pdgHypo, ambiTracks, bcs); + pool02.appendTrackCand(trackCand, collisions, pdgHypo, ambiTracks, bcs); + } else if (pdgHypo == track1Pdg) { + pool01.appendTrackCand(trackCand, collisions, pdgHypo, ambiTracks, bcs); + } else if (pdgHypo == track2Pdg) { + pool02.appendTrackCand(trackCand, collisions, pdgHypo, ambiTracks, bcs); + } else { + LOGP(debug, "Wrong PDG hypothesis for three-body pool"); + } + } + + template + std::vector& getSVCandPool(const C& collisions, bool combineLikeSign01 = false, bool combineLikeSign02 = false) + { + auto& candidates01 = pool01.getSVCandPool(collisions, combineLikeSign01); + auto& candidates02 = pool02.getSVCandPool(collisions, combineLikeSign02); + std::unordered_multimap candidates02ByTrack0; + candidates02ByTrack0.reserve(candidates02.size()); + for (const auto& candidate02 : candidates02) { + candidates02ByTrack0.emplace(candidate02.tr0Idx, &candidate02); + } + + for (const auto& candidate01 : candidates01) { + const auto [begin, end] = candidates02ByTrack0.equal_range(candidate01.tr0Idx); + for (auto candidate02It = begin; candidate02It != end; ++candidate02It) { + const auto& candidate02 = *candidate02It->second; + if (candidate01.tr1Idx == candidate02.tr1Idx || + candidate01.collBracket.isOutside(candidate02.collBracket)) { + continue; + } + svCandPool.emplace_back(SVCand3{candidate01.tr0Idx, + candidate01.tr1Idx, + candidate02.tr1Idx, + candidate01.collBracket.getOverlap(candidate02.collBracket)}); + } + } + return svCandPool; + } + + private: + int track0Pdg = 0; + int track1Pdg = 0; + int track2Pdg = 0; + svPoolCreator pool01; + svPoolCreator pool02; + std::vector svCandPool; +}; + #endif // PWGLF_UTILS_SVPOOLCREATOR_H_