diff --git a/PWGLF/DataModel/LFDoubleOmegaTrackTables.h b/PWGLF/DataModel/LFDoubleOmegaTrackTables.h new file mode 100644 index 00000000000..fb75f3c3f46 --- /dev/null +++ b/PWGLF/DataModel/LFDoubleOmegaTrackTables.h @@ -0,0 +1,109 @@ +// Copyright 2019-2020 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 LFDoubleOmegaTrackTables.h +/// \brief Tables of doubleOmegaTrackTreeCreator: the training sample of the double-Omega track selection +/// (DOmegaTrackTrain) and the V0 and cascade candidates built from the tracks selected by the track +/// BDT (DOmegaV0s, DOmegaCascades), analogous to aod::V0s and aod::Cascades (same index getters: +/// collision, posTrack, negTrack, v0, bachelor). + +#ifndef PWGLF_DATAMODEL_LFDOUBLEOMEGATRACKTABLES_H_ +#define PWGLF_DATAMODEL_LFDOUBLEOMEGATRACKTABLES_H_ + +#include +#include + +#include + +namespace o2::aod +{ +namespace domegatrack +{ +DECLARE_SOA_COLUMN(Pt, pt, float); +DECLARE_SOA_COLUMN(Eta, eta, float); +DECLARE_SOA_COLUMN(Phi, phi, float); +DECLARE_SOA_COLUMN(Sign, sign, int8_t); +DECLARE_SOA_COLUMN(TpcInnerParam, tpcInnerParam, float); //! Momentum at the TPC inner wall (PID) +DECLARE_SOA_COLUMN(TpcSignal, tpcSignal, float); +DECLARE_SOA_COLUMN(TpcNSigmaPr, tpcNSigmaPr, float); +DECLARE_SOA_COLUMN(TpcNSigmaPi, tpcNSigmaPi, float); +DECLARE_SOA_COLUMN(TpcNSigmaKa, tpcNSigmaKa, float); +DECLARE_SOA_COLUMN(TofNSigmaPr, tofNSigmaPr, float); //! -999 without TOF +DECLARE_SOA_COLUMN(TofNSigmaPi, tofNSigmaPi, float); //! -999 without TOF +DECLARE_SOA_COLUMN(TofNSigmaKa, tofNSigmaKa, float); //! -999 without TOF +DECLARE_SOA_COLUMN(DcaXY, dcaXY, float); //! Signed transverse DCA to the primary vertex (cm), -999 if the propagation failed +DECLARE_SOA_COLUMN(DcaZ, dcaZ, float); //! Longitudinal DCA to the primary vertex (cm) +DECLARE_SOA_COLUMN(SigmaDcaXY, sigmaDcaXY, float); //! Uncertainty of DcaXY, track and vertex (cm) +DECLARE_SOA_COLUMN(SigmaDcaZ, sigmaDcaZ, float); //! Uncertainty of DcaZ, track and vertex (cm) +DECLARE_SOA_COLUMN(ItsClusterMap, itsClusterMap, uint8_t); +DECLARE_SOA_COLUMN(ItsClusterSizes, itsClusterSizes, uint32_t); +DECLARE_SOA_COLUMN(ItsChi2NCl, itsChi2NCl, float); +DECLARE_SOA_COLUMN(TpcNClsFound, tpcNClsFound, int16_t); +DECLARE_SOA_COLUMN(TpcNClsCrossedRows, tpcNClsCrossedRows, int16_t); +DECLARE_SOA_COLUMN(TpcNClsShared, tpcNClsShared, uint8_t); +DECLARE_SOA_COLUMN(TpcChi2NCl, tpcChi2NCl, float); +DECLARE_SOA_COLUMN(DetectorMap, detectorMap, uint8_t); //! ITS, TPC, TRD, TOF bits +DECLARE_SOA_COLUMN(IsPVContributor, isPVContributor, bool); +DECLARE_SOA_COLUMN(NumContrib, numContrib, uint16_t); //! Primary-vertex contributors of the collision (multiplicity proxy) +DECLARE_SOA_COLUMN(PosZ, posZ, float); //! Primary-vertex z of the collision (cm) +DECLARE_SOA_COLUMN(Role, role, int8_t); //! See doubleOmegaTrackTreeCreator.cxx +DECLARE_SOA_COLUMN(PdgCode, pdgCode, int); //! Of the MC particle, 0 without MC particle +DECLARE_SOA_COLUMN(MotherPdgCode, motherPdgCode, int); +DECLARE_SOA_COLUMN(IsPhysicalPrimary, isPhysicalPrimary, bool); +DECLARE_SOA_COLUMN(PtMC, ptMC, float); +DECLARE_SOA_COLUMN(ProductionRadius, productionRadius, float); //! Transverse radius of the MC production vertex (cm) +DECLARE_SOA_COLUMN(McKey, mcKey, int64_t); //! MC particle, unique across dataframes (dataframe << 32 | index), -1 without MC particle +DECLARE_SOA_COLUMN(AncestorPdgCode, ancestorPdgCode, int); //! First particle of the decay chain (see doubleOmegaTrackTreeCreator.cxx), 0 without MC particle +DECLARE_SOA_COLUMN(AncestorPtMC, ancestorPtMC, float); //! Its generated pT (GeV/c), -999 without MC particle +} // namespace domegatrack + +DECLARE_SOA_TABLE(DOmegaTrackTrain, "AOD", "DOMEGATRKTRAIN", //! Tracks for the training of the double-Omega track selection + domegatrack::Pt, domegatrack::Eta, domegatrack::Phi, domegatrack::Sign, + domegatrack::TpcInnerParam, domegatrack::TpcSignal, + domegatrack::TpcNSigmaPr, domegatrack::TpcNSigmaPi, domegatrack::TpcNSigmaKa, + domegatrack::TofNSigmaPr, domegatrack::TofNSigmaPi, domegatrack::TofNSigmaKa, + domegatrack::DcaXY, domegatrack::DcaZ, domegatrack::SigmaDcaXY, domegatrack::SigmaDcaZ, + domegatrack::ItsClusterMap, domegatrack::ItsClusterSizes, domegatrack::ItsChi2NCl, + domegatrack::TpcNClsFound, domegatrack::TpcNClsCrossedRows, + domegatrack::TpcNClsShared, domegatrack::TpcChi2NCl, + domegatrack::DetectorMap, domegatrack::IsPVContributor, + domegatrack::NumContrib, domegatrack::PosZ, + domegatrack::Role, domegatrack::PdgCode, domegatrack::MotherPdgCode, + domegatrack::IsPhysicalPrimary, domegatrack::PtMC, domegatrack::ProductionRadius, domegatrack::McKey, + domegatrack::AncestorPdgCode, domegatrack::AncestorPtMC); + +// V0 and cascade candidates from the tracks selected by the track BDT. The track indices point to the Tracks +// of the AO2D the candidates were built from: like aod::V0s and aod::Cascades, the tables are used together +// with it. +namespace domegav0 +{ +DECLARE_SOA_INDEX_COLUMN(Collision, collision); //! Collision index +DECLARE_SOA_INDEX_COLUMN_FULL(PosTrack, posTrack, int, Tracks, "_Pos"); //! Positive track (proton of a Lambda, pion of an anti-Lambda) +DECLARE_SOA_INDEX_COLUMN_FULL(NegTrack, negTrack, int, Tracks, "_Neg"); //! Negative track (pion of a Lambda, antiproton of an anti-Lambda) +DECLARE_SOA_COLUMN(Sign, sign, int8_t); //! +1 Lambda, -1 anti-Lambda +} // namespace domegav0 + +DECLARE_SOA_TABLE(DOmegaV0s, "AOD", "DOMEGAV0", //! Lambda candidates, analogous to aod::V0s + o2::soa::Index<>, domegav0::CollisionId, domegav0::PosTrackId, domegav0::NegTrackId, domegav0::Sign); + +namespace domegacascade +{ +DECLARE_SOA_INDEX_COLUMN(Collision, collision); //! Collision index +DECLARE_SOA_INDEX_COLUMN_FULL(V0, v0, int, DOmegaV0s, ""); //! Lambda candidate +DECLARE_SOA_INDEX_COLUMN_FULL(Bachelor, bachelor, int, Tracks, ""); //! Bachelor kaon +DECLARE_SOA_COLUMN(Sign, sign, int8_t); //! Charge of the Omega: -1 Omega-, +1 Omega+ +} // namespace domegacascade + +DECLARE_SOA_TABLE(DOmegaCascades, "AOD", "DOMEGACASCADE", //! Omega candidates, analogous to aod::Cascades + o2::soa::Index<>, domegacascade::CollisionId, domegacascade::V0Id, domegacascade::BachelorId, domegacascade::Sign); +} // namespace o2::aod + +#endif // PWGLF_DATAMODEL_LFDOUBLEOMEGATRACKTABLES_H_ diff --git a/PWGLF/TableProducer/Strangeness/CMakeLists.txt b/PWGLF/TableProducer/Strangeness/CMakeLists.txt index a0a64e95190..6dee03abe14 100644 --- a/PWGLF/TableProducer/Strangeness/CMakeLists.txt +++ b/PWGLF/TableProducer/Strangeness/CMakeLists.txt @@ -75,6 +75,11 @@ o2physics_add_dpl_workflow(double-omega-builder PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2Physics::EventFilteringUtils COMPONENT_NAME Analysis) +o2physics_add_dpl_workflow(double-omega-track-index-skim-creator + SOURCES doubleOmegaTrackIndexSkimCreator.cxx + PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2Physics::MLCore O2::DCAFitter + COMPONENT_NAME Analysis) + o2physics_add_dpl_workflow(lambdakzerofinder SOURCES lambdakzerofinder.cxx REUSE_FROM LfStrangenessPCH diff --git a/PWGLF/TableProducer/Strangeness/doubleOmegaTrackIndexSkimCreator.cxx b/PWGLF/TableProducer/Strangeness/doubleOmegaTrackIndexSkimCreator.cxx new file mode 100644 index 00000000000..a8418d79ac3 --- /dev/null +++ b/PWGLF/TableProducer/Strangeness/doubleOmegaTrackIndexSkimCreator.cxx @@ -0,0 +1,651 @@ +// Copyright 2019-2020 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 doubleOmegaTrackIndexSkimCreator.cxx +/// \brief Table producer for V0s and Cascade tables from tracks selected by a track-level BDT. +/// Designed to improve the efficiency in the double Omega search, the BDT is supposed to +/// discriminate between primary tracks and daughters of Lambda and Omega decays, +/// regardless of the fact that those particles are from a double Omega or not. +/// A table producer for the AO2Ds required for the training of the BDT is also provided. +/// +/// \author Giorgio Alberto Lucia (giorgio.alberto.lucia@cern.ch) + +#include "PWGLF/DataModel/LFDoubleOmegaTrackTables.h" +#include "PWGLF/Utils/doubleOmegaBuilderHelper.h" + +#include "Common/CCDB/EventSelectionParams.h" +#include "Common/Core/RecoDecay.h" +#include "Common/Core/trackUtilities.h" +#include "Common/DataModel/EventSelection.h" +#include "Common/DataModel/PIDResponseTOF.h" +#include "Common/DataModel/PIDResponseTPC.h" +#include "Tools/ML/MlResponse.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#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; +using namespace o2::pwglf::doubleomega; + +using Collisions = soa::Join; +using CollisionsMC = soa::Join; +using TracksData = soa::Join; +using TracksMC = soa::Join; + +struct doubleOmegaTrackIndexSkimCreator { + Produces trackTable; + Produces v0Table; + Produces cascadeTable; + Service ccdb{}; + o2::ccdb::CcdbApi ccdbApi; + + static constexpr int kDoubleOmegaPdg = 1060020020; + + enum Role : int8_t { + kBackground = 0, + kOmegaProton, + kOmegaPion, + kOmegaKaon, + kDirectLambdaProton, + kDirectLambdaPion, + kDirectKaon, + kOtherLambdaProton, + kOtherLambdaPion, + kOtherOmegaKaon, + kNRoles + }; + static constexpr std::array RoleLabels{ + "Background", "p #leftarrow #Lambda #leftarrow #Omega #leftarrow #Omega#Omega", "#pi #leftarrow #Lambda #leftarrow #Omega #leftarrow #Omega#Omega", + "K #leftarrow #Omega #leftarrow #Omega#Omega", "p #leftarrow #Lambda #leftarrow #Omega#Omega", "#pi #leftarrow #Lambda #leftarrow #Omega#Omega", + "K #leftarrow #Omega#Omega", "p #leftarrow other #Lambda", "#pi #leftarrow other #Lambda", "K #leftarrow other #Omega"}; + + // Species of the chain starts in hGeneratedAncestors (charge conjugates included) + enum Species : int8_t { + kSpeciesDoubleOmega = 0, + kSpeciesOmega, + kSpeciesXi, + kSpeciesLambda, + kSpeciesPion, + kSpeciesKaon, + kSpeciesProton, + kNSpecies + }; + static constexpr std::array SpeciesLabels{ + "#Omega#Omega", "#Omega", "#Xi^{-}, #Xi^{0}", "#Lambda", "#pi^{#pm}", "K^{#pm}", "p"}; + static constexpr int kMaxChainDepth = 10; + + // Classes of the track BDT used for the candidates (candOpts.classIndices, classThresholds) + enum TrackClass : uint8_t { + kClassProton = 0, + kClassPion, + kClassKaon, + kNTrackClasses + }; + + Preslice tracksDataPerCollision = aod::track::collisionId; + Preslice tracksMCPerCollision = aod::track::collisionId; + + Configurable ccdburl{"ccdb-url", "http://alice-ccdb.cern.ch", "url of the ccdb repository"}; + Configurable cfgMaterialCorrection{"cfgMaterialCorrection", static_cast(o2::base::Propagator::MatCorrType::USEMatCorrLUT), "Material correction for the propagation to the primary vertex"}; + Configurable cfgApplyEventSelection{"cfgApplyEventSelection", true, "Apply the standard collision event selection"}; + Configurable zVtxMax{"zVtxMax", 10.f, "Maximum |z| of the primary vertex (cm)"}; + + // Track preselection, the baseline of the selection: it should be looser than any selection studied with it. + // No track-quality cut (e.g. TPC crossed rows): the quality variables are stored and left to the selections. + Configurable requireTPC{"requireTPC", true, "Require TPC (PID of all the daughters)"}; + Configurable etaMax{"etaMax", 1.f, "Maximum |eta|"}; + Configurable genEtaMax{"genEtaMax", 1.f, "Maximum |eta| of the generated daughters (efficiency denominator)"}; + ConfigurableAxis ptAxis{"ptAxis", {100, 0.f, 10.f}, "pT binning of the efficiency and purity histograms (GeV/c)"}; + Configurable ancestorRapidityMax{"ancestorRapidityMax", 0.5f, "Maximum |y| of the generated chain starts (hGeneratedAncestors)"}; + Configurable ptMin{"ptMin", 0.f, "Minimum pT (GeV/c)"}; + + // Downsampling (roles 1-6 are always stored) + Configurable fractionBackground{"fractionBackground", 0.01f, "Fraction of the background tracks (role 0) stored"}; + Configurable fractionOtherDaughters{"fractionOtherDaughters", 0.1f, "Fraction of the daughters of other Lambdas and Omegas (roles 7-9) stored"}; + Configurable randomSeed{"randomSeed", 0, "Seed of the downsampling (0: random)"}; + + // V0 and cascade candidates (processCandidates), from the preselected tracks above selected by the track BDT + struct : ConfigurableGroup { + std::string prefix = "candOpts"; + Configurable onnxFileName{"onnxFileName", "model.onnx", "ONNX file of the track BDT (local path, or file name on CCDB)"}; + Configurable loadModelFromCCDB{"loadModelFromCCDB", false, "Load the ONNX model from CCDB"}; + Configurable modelPathCCDB{"modelPathCCDB", "Users/g/galucia/DoubleOmega/TrackBDT", "CCDB path of the model"}; + Configurable timestampCCDB{"timestampCCDB", -1, "Timestamp of the model on CCDB"}; + Configurable> inputFeatures{"inputFeatures", std::vector{"fPt", "fTpcInnerParam", "fTpcNSigmaPr", "fTpcNSigmaPi", "fTpcNSigmaKa", "fTofNSigmaPr", "fTofNSigmaPi", "fTofNSigmaKa", "fDcaXY", "fDcaZ", "fTpcNClsCrossedRows", "fTpcNClsShared", "fTpcChi2NCl", "fIsPVContributor"}, "Model inputs, in order (model_inputs of config.json of the training; names of PWGLF/Utils/doubleOmegaBuilderHelper.h)"}; + Configurable nClasses{"nClasses", 4, "Number of model outputs (classes)"}; + Configurable> classIndices{"classIndices", std::vector{1, 2, 3}, "Output index of the proton, pion and kaon classes (class_names of config.json)"}; + Configurable> classThresholds{"classThresholds", std::vector{0.5f, 0.5f, 0.5f}, "Minimum probability of the proton, pion and kaon classes"}; + Configurable lambdaMassWindow{"lambdaMassWindow", 0.05f, "Lambda mass window of the V0s, from the momenta at the innermost update (GeV/c^2; negative: no selection)"}; + Configurable omegaMassWindow{"omegaMassWindow", 0.1f, "Omega mass window of the cascades, from the momenta at the innermost update and the Lambda mass (GeV/c^2; negative: no selection)"}; + } candOpts; + + HistogramRegistry histos{"histos", {}, OutputObjHandlingPolicy::AnalysisObject}; + + o2::analysis::MlResponse mlResponse; // track BDT + std::vector mlFeatureIndices; // model inputs in the feature array of computeTrackFeatures + std::array mlClassIndex{}; + std::array mlThreshold{}; + + int mRunNumber = 0; + std::mt19937 mRandom; + std::uniform_real_distribution mUniform{0.f, 1.f}; + std::unordered_map mRoles; // MC particle index -> role, per dataframe + std::unordered_set mAcceptedMcCollisions; + int64_t mDataframe = 0; + + enum Stage : uint8_t { + kPreselected = 0, + kStored, + kNStages + }; + + void init(InitContext const&) + { + ccdb->setURL(ccdburl); + ccdb->setCaching(true); + ccdb->setLocalObjectValidityChecking(); + ccdb->setFatalWhenNull(false); + mRandom.seed(randomSeed.value != 0 ? static_cast(randomSeed.value) : std::random_device{}()); + + histos.add("hEvents", ";;Collisions", HistType::kTH1F, {{2, -0.5f, 1.5f}}); + histos.get(HIST("hEvents"))->GetXaxis()->SetBinLabel(1, "All"); + histos.get(HIST("hEvents"))->GetXaxis()->SetBinLabel(2, "Selected"); + const AxisSpec roleAxis{kNRoles, -0.5f, static_cast(kNRoles) - 0.5f, "Role"}; + auto hTracks = histos.add("hTracks", ";Role;Stage;#it{p}_{T} (GeV/#it{c})", HistType::kTH3F, {roleAxis, {kNStages, -0.5f, static_cast(kNStages) - 0.5f}, ptAxis}); + hTracks->GetYaxis()->SetBinLabel(kPreselected + 1, "Preselected"); + hTracks->GetYaxis()->SetBinLabel(kStored + 1, "Stored"); + auto hGenerated = histos.add("hGenerated", ";Role;#it{p}_{T}^{gen} (GeV/#it{c})", HistType::kTH2F, {roleAxis, ptAxis}); + auto hGeneratedAncestorPt = histos.add("hGeneratedAncestorPt", ";Role;#it{p}_{T}^{gen} (GeV/#it{c});#it{p}_{T}^{gen} of the first particle of the decay chain (GeV/#it{c})", + HistType::kTH3F, {roleAxis, ptAxis, ptAxis}); + auto hGeneratedAncestors = histos.add("hGeneratedAncestors", ";Species;#it{p}_{T}^{gen} (GeV/#it{c})", HistType::kTH2F, + {{kNSpecies, -0.5f, static_cast(kNSpecies) - 0.5f, "Species"}, ptAxis}); + for (int species = 0; species < kNSpecies; ++species) { + hGeneratedAncestors->GetXaxis()->SetBinLabel(species + 1, SpeciesLabels[species]); + } + for (auto* axis : {hTracks->GetXaxis(), hGenerated->GetXaxis(), hGeneratedAncestorPt->GetXaxis()}) { + for (int role = 0; role < kNRoles; ++role) { + axis->SetBinLabel(role + 1, RoleLabels[role]); + } + } + + if (doprocessCandidates) { + initTrackMl(); + histos.add("Candidates/hEvents", ";;Collisions", HistType::kTH1F, {{2, -0.5f, 1.5f}}); + histos.get(HIST("Candidates/hEvents"))->GetXaxis()->SetBinLabel(1, "All"); + histos.get(HIST("Candidates/hEvents"))->GetXaxis()->SetBinLabel(2, "Selected"); + auto hPools = histos.add("Candidates/hPools", ";Class;#it{p}_{T} (GeV/#it{c})", HistType::kTH2F, {{2 * kNTrackClasses, -0.5f, static_cast(2 * kNTrackClasses) - 0.5f}, ptAxis}); + for (const auto& [bin, label] : {std::pair{1, "p"}, std::pair{2, "#bar{p}"}, std::pair{3, "#pi^{+}"}, std::pair{4, "#pi^{-}"}, std::pair{5, "K^{+}"}, std::pair{6, "K^{-}"}}) { + hPools->GetXaxis()->SetBinLabel(bin, label); + } + histos.add("Candidates/hMassLambda", ";#it{p}_{T} (GeV/#it{c});Sign;#it{M}(p#pi) (GeV/#it{c}^{2})", HistType::kTH3F, {ptAxis, {2, -2.f, 2.f}, {200, 1.05f, 1.25f}}); + histos.add("Candidates/hMassOmega", ";#it{p}_{T} (GeV/#it{c});Sign;#it{M}(#LambdaK) (GeV/#it{c}^{2})", HistType::kTH3F, {ptAxis, {2, -2.f, 2.f}, {200, 1.57f, 1.77f}}); + } + } + + void initTrackMl() + { + mlFeatureIndices = trackFeatureIndices(candOpts.inputFeatures.value); + const auto& indices = candOpts.classIndices.value; + const auto& thresholds = candOpts.classThresholds.value; + if (indices.size() != kNTrackClasses || thresholds.size() != kNTrackClasses) { + LOG(fatal) << "candOpts.classIndices and candOpts.classThresholds need one entry per class (proton, pion, kaon)"; + } + for (int iClass = 0; iClass < kNTrackClasses; ++iClass) { + if (indices[iClass] < 0 || indices[iClass] >= candOpts.nClasses) { + LOG(fatal) << "candOpts.classIndices: index " << indices[iClass] << " out of the " << candOpts.nClasses << " model outputs"; + } + mlClassIndex[iClass] = indices[iClass]; + mlThreshold[iClass] = thresholds[iClass]; + } + + // One model for all pT; the class thresholds are applied here, not by MlResponse. + const int nClasses = candOpts.nClasses; + std::vector noCuts(nClasses, 0.); + std::vector labelsClasses; + for (int iClass = 0; iClass < nClasses; ++iClass) { + labelsClasses.push_back(Form("class %d", iClass)); + } + mlResponse.configure(std::vector{0., 1.e6}, LabeledArray{noCuts.data(), 1, static_cast(nClasses), {"all pT"}, labelsClasses}, + std::vector(nClasses, o2::cuts_ml::CutNot), static_cast(nClasses)); + if (candOpts.loadModelFromCCDB) { + ccdbApi.init(ccdburl); + mlResponse.setModelPathsCCDB(std::vector{candOpts.onnxFileName.value}, ccdbApi, std::vector{candOpts.modelPathCCDB.value}, candOpts.timestampCCDB); + } else { + mlResponse.setModelPathsLocal(std::vector{candOpts.onnxFileName.value}); + } + mlResponse.init(); + } + + template + void initCCDB(Bc const& bc) + { + if (mRunNumber == bc.runNumber()) { + return; + } + auto* grpmag = ccdb->getForTimeStamp("GLO/Config/GRPMagField", bc.timestamp()); + if (!grpmag) { + LOG(fatal) << "Could not retrieve GRPMagField for timestamp " << bc.timestamp(); + } + o2::base::Propagator::initFieldFromGRP(grpmag); + if (static_cast(cfgMaterialCorrection.value) == o2::base::Propagator::MatCorrType::USEMatCorrLUT) { + auto* lut = o2::base::MatLayerCylSet::rectifyPtrFromFile(ccdb->getForTimeStamp("GLO/Param/MatLUT", bc.timestamp())); + if (!lut) { + LOG(fatal) << "Could not retrieve material LUT for timestamp " << bc.timestamp(); + } + o2::base::Propagator::Instance()->setMatLUT(lut); + } + mRunNumber = bc.runNumber(); + } + + // isCandidates: counted in Candidates/hEvents (processCandidates) instead of hEvents (processMC) + template + bool acceptCollision(C const& collision) + { + initCCDB(collision.template bc_as()); + const auto fillEvents = [&](int bin) { + if constexpr (isCandidates) { + histos.fill(HIST("Candidates/hEvents"), bin); + } else { + histos.fill(HIST("hEvents"), bin); + } + }; + fillEvents(0); + if (cfgApplyEventSelection && + (!collision.sel8() || std::abs(collision.posZ()) > zVtxMax)) { + return false; + } + fillEvents(1); + return true; + } + + template + bool preselectTrack(T const& track) const + { + return track.sign() != 0 && + (!requireTPC.value || track.hasTPC()) && + std::abs(track.eta()) <= etaMax.value && + track.pt() >= ptMin.value; + } + + bool keep(Role role) + { + switch (role) { + case kBackground: + return mUniform(mRandom) < fractionBackground; + case kOtherLambdaProton: + case kOtherLambdaPion: + case kOtherOmegaKaon: + return mUniform(mRandom) < fractionOtherDaughters; + default: + return true; + } + } + + // --------------------------------------------------------------------------------------------- + // MC roles + + // Daughter of a given PDG code, -1 if none + template + static int64_t findDaughter(McParticle const& particle, int pdgCode) + { + for (const auto& daughter : particle.template daughters_as()) { + if (daughter.pdgCode() == pdgCode) { + return daughter.globalIndex(); + } + } + return -1; + } + + // First particle of the decay chain of an MC particle: the double Omega if it is an ancestor, otherwise the + // first physical primary going up the mothers (the particle itself if it is a physical primary) + static int64_t chainStart(aod::McParticles const& mcParticles, int64_t id) + { + for (int depth = 0; depth < kMaxChainDepth; ++depth) { + const auto particle = mcParticles.rawIteratorAt(id); + if (!particle.has_mothers()) { + break; + } + const int64_t motherId = particle.mothersIds()[0]; + if (motherId < 0 || motherId >= mcParticles.size()) { + break; + } + if (std::abs(mcParticles.rawIteratorAt(motherId).pdgCode()) == kDoubleOmegaPdg) { + return motherId; + } + if (particle.isPhysicalPrimary()) { + break; + } + id = motherId; + } + return id; + } + + // Species of a chain start in hGeneratedAncestors, -1 if none + static int species(int pdgCode) + { + switch (std::abs(pdgCode)) { + case kDoubleOmegaPdg: + return kSpeciesDoubleOmega; + case PDG_t::kOmegaMinus: + return kSpeciesOmega; + case PDG_t::kXiMinus: + case 3322: // Xi0 + return kSpeciesXi; + case PDG_t::kLambda0: + return kSpeciesLambda; + case PDG_t::kPiPlus: + return kSpeciesPion; + case PDG_t::kKPlus: + return kSpeciesKaon; + case PDG_t::kProton: + return kSpeciesProton; + default: + return -1; + } + } + + // Proton and pion of a Lambda (sign: of the Lambda PDG code) + void assignLambdaDaughters(aod::McParticles const& mcParticles, int64_t lambdaId, int sign, Role protonRole, Role pionRole) + { + if (lambdaId < 0) { + return; + } + const auto lambda = mcParticles.rawIteratorAt(lambdaId); + if (const auto id = findDaughter(lambda, sign * PDG_t::kProton); id >= 0) { + mRoles.emplace(id, protonRole); + } + if (const auto id = findDaughter(lambda, -sign * PDG_t::kPiPlus); id >= 0) { + mRoles.emplace(id, pionRole); + } + } + + // Roles of the MC particles of the dataframe. The double-Omega chains are assigned first, so that their + // Lambdas and Omega are not counted again as "other" ones (emplace does not overwrite). + void assignRoles(aod::McParticles const& mcParticles) + { + mRoles.clear(); + for (const auto& particle : mcParticles) { + if (std::abs(particle.pdgCode()) != kDoubleOmegaPdg) { + continue; + } + const int sign = particle.pdgCode() > 0 ? 1 : -1; + const int64_t omegaId = findDaughter(particle, sign * PDG_t::kOmegaMinus); + const int64_t directLambdaId = findDaughter(particle, sign * PDG_t::kLambda0); + if (const auto id = findDaughter(particle, -sign * PDG_t::kKPlus); id >= 0) { + mRoles.emplace(id, kDirectKaon); + } + assignLambdaDaughters(mcParticles, directLambdaId, sign, kDirectLambdaProton, kDirectLambdaPion); + if (omegaId >= 0) { + const auto omega = mcParticles.rawIteratorAt(omegaId); + if (const auto id = findDaughter(omega, -sign * PDG_t::kKPlus); id >= 0) { + mRoles.emplace(id, kOmegaKaon); + } + assignLambdaDaughters(mcParticles, findDaughter(omega, sign * PDG_t::kLambda0), sign, kOmegaProton, kOmegaPion); + } + } + for (const auto& particle : mcParticles) { + const int pdgCode = particle.pdgCode(); + const int sign = pdgCode > 0 ? 1 : -1; + if (std::abs(pdgCode) == PDG_t::kLambda0) { + assignLambdaDaughters(mcParticles, particle.globalIndex(), sign, kOtherLambdaProton, kOtherLambdaPion); + } else if (std::abs(pdgCode) == PDG_t::kOmegaMinus) { + if (const auto id = findDaughter(particle, -sign * PDG_t::kKPlus); id >= 0) { + mRoles.emplace(id, kOtherOmegaKaon); + } + } + } + } + + // --------------------------------------------------------------------------------------------- + // Filling + + template + void fillTrack(C const& collision, T const& track, Role role, int pdgCode, int motherPdgCode, bool isPhysicalPrimary, float ptMC, float productionRadius, int64_t mcKey, + int ancestorPdgCode, float ancestorPtMC) + { + // DCA to the primary vertex with the uncertainties of the track and of the vertex + auto trackParCov = getTrackParCov(track); + const o2::dataformats::VertexBase primaryVertex{o2::math_utils::Point3D{collision.posX(), collision.posY(), collision.posZ()}, + std::array{collision.covXX(), collision.covXY(), collision.covYY(), + collision.covXZ(), collision.covYZ(), collision.covZZ()}}; + o2::dataformats::DCA dca; + float dcaXY = -999.f, dcaZ = -999.f, sigmaDcaXY = -999.f, sigmaDcaZ = -999.f; + if (o2::base::Propagator::Instance()->propagateToDCABxByBz(primaryVertex, trackParCov, 2.f, static_cast(cfgMaterialCorrection.value), &dca)) { + dcaXY = dca.getY(); + dcaZ = dca.getZ(); + sigmaDcaXY = std::sqrt(std::max(0.f, static_cast(dca.getSigmaY2()))); + sigmaDcaZ = std::sqrt(std::max(0.f, static_cast(dca.getSigmaZ2()))); + } + const uint8_t detectorMap = static_cast(track.hasITS()) | (static_cast(track.hasTPC()) << 1) | + (static_cast(track.hasTRD()) << 2) | (static_cast(track.hasTOF()) << 3); + trackTable(track.pt(), track.eta(), track.phi(), static_cast(track.sign()), + track.tpcInnerParam(), track.tpcSignal(), + track.tpcNSigmaPr(), track.tpcNSigmaPi(), track.tpcNSigmaKa(), + track.hasTOF() ? track.tofNSigmaPr() : -999.f, track.hasTOF() ? track.tofNSigmaPi() : -999.f, track.hasTOF() ? track.tofNSigmaKa() : -999.f, + dcaXY, dcaZ, sigmaDcaXY, sigmaDcaZ, + track.itsClusterMap(), track.itsClusterSizes(), track.itsChi2NCl(), + static_cast(track.tpcNClsFound()), static_cast(track.tpcNClsCrossedRows()), + track.tpcNClsShared(), track.tpcChi2NCl(), + detectorMap, track.isPVContributor(), + static_cast(collision.numContrib()), collision.posZ(), + role, pdgCode, motherPdgCode, isPhysicalPrimary, ptMC, productionRadius, mcKey, + ancestorPdgCode, ancestorPtMC); + } + + void processMC(CollisionsMC const& collisions, TracksMC const& tracks, aod::McParticles const& mcParticles, + aod::McCollisions const&, aod::BCsWithTimestamps const&) + { + ++mDataframe; + assignRoles(mcParticles); + mAcceptedMcCollisions.clear(); + for (const auto& collision : collisions) { + if (!acceptCollision(collision)) { + continue; + } + if (collision.has_mcCollision()) { + mAcceptedMcCollisions.insert(collision.mcCollisionId()); + } + for (const auto& track : tracks.sliceBy(tracksMCPerCollision, collision.globalIndex())) { + Role role = kBackground; + int pdgCode = 0, motherPdgCode = 0; + bool isPhysicalPrimary = false; + float ptMC = -999.f, productionRadius = -999.f, ancestorPtMC = -999.f; + int64_t mcKey = -1; + int ancestorPdgCode = 0; + if (track.has_mcParticle()) { + const auto mcParticle = track.template mcParticle_as(); + if (const auto found = mRoles.find(mcParticle.globalIndex()); found != mRoles.end()) { + role = found->second; + } + pdgCode = mcParticle.pdgCode(); + for (const auto& mother : mcParticle.template mothers_as()) { + motherPdgCode = mother.pdgCode(); + break; + } + isPhysicalPrimary = mcParticle.isPhysicalPrimary(); + ptMC = mcParticle.pt(); + productionRadius = std::hypot(mcParticle.vx(), mcParticle.vy()); + mcKey = (mDataframe << 32) | mcParticle.globalIndex(); + const auto ancestor = mcParticles.rawIteratorAt(chainStart(mcParticles, mcParticle.globalIndex())); + ancestorPdgCode = ancestor.pdgCode(); + ancestorPtMC = ancestor.pt(); + } + if (!preselectTrack(track)) { + continue; + } + histos.fill(HIST("hTracks"), role, kPreselected, track.pt()); + if (keep(role)) { + histos.fill(HIST("hTracks"), role, kStored, track.pt()); + fillTrack(collision, track, role, pdgCode, motherPdgCode, isPhysicalPrimary, ptMC, productionRadius, mcKey, + ancestorPdgCode, ancestorPtMC); + } + } + } + + // Generated daughters (efficiency denominator) + for (const auto& [id, role] : mRoles) { + const auto particle = mcParticles.rawIteratorAt(id); + if (!mAcceptedMcCollisions.contains(particle.mcCollisionId()) || std::abs(particle.eta()) > genEtaMax) { + continue; + } + histos.fill(HIST("hGenerated"), role, particle.pt()); + histos.fill(HIST("hGeneratedAncestorPt"), role, particle.pt(), mcParticles.rawIteratorAt(chainStart(mcParticles, id)).pt()); + } + + // Generated chain starts: the MC pT shape of the double Omegas and of the primary hadrons + for (const auto& particle : mcParticles) { + const int index = species(particle.pdgCode()); + if (index < 0 || (index != kSpeciesDoubleOmega && !particle.isPhysicalPrimary()) || + !mAcceptedMcCollisions.contains(particle.mcCollisionId()) || std::abs(particle.y()) > ancestorRapidityMax) { + continue; + } + histos.fill(HIST("hGeneratedAncestors"), index, particle.pt()); + } + } + PROCESS_SWITCH(doubleOmegaTrackIndexSkimCreator, processMC, "Store the tracks of MC with their role (signal: double-Omega daughters, background downsampled)", true); + + // --------------------------------------------------------------------------------------------- + // V0 and cascade candidates + + struct PoolTrack { + int64_t id = -1; + std::array momentum{}; + }; + // Tracks of a collision selected by the BDT, per class and charge ([0] negative, [1] positive) + using Pools = std::array, 2>, kNTrackClasses>; + + struct StoredV0 { + int64_t index = -1; // row of DOmegaV0s + int64_t pionId = -1; + int8_t sign = 0; + std::array momentum{}; + }; + + template + Pools fillPools(C const& collision, T const& collisionTracks) + { + Pools pools; + std::array features{}; + std::vector inputs(mlFeatureIndices.size()); + for (const auto& track : collisionTracks) { + if (!preselectTrack(track)) { + continue; + } + computeTrackFeatures(collision, track, static_cast(cfgMaterialCorrection.value), features); + for (std::size_t iInput = 0; iInput < mlFeatureIndices.size(); ++iInput) { + inputs[iInput] = features[mlFeatureIndices[iInput]]; + } + const auto scores = mlResponse.getModelOutput(inputs, 0); + const int iCharge = track.sign() > 0 ? 1 : 0; + for (int iClass = 0; iClass < kNTrackClasses; ++iClass) { + if (scores[mlClassIndex[iClass]] <= mlThreshold[iClass]) { + continue; + } + pools[iClass][iCharge].push_back({track.globalIndex(), {track.px(), track.py(), track.pz()}}); + histos.fill(HIST("Candidates/hPools"), 2 * iClass + 1 - iCharge, track.pt()); + } + } + return pools; + } + + template + void buildCandidates(C const& collision, Pools const& pools) + { + constexpr auto MassLambda = o2::constants::physics::MassLambda0; + constexpr auto MassOmega = o2::constants::physics::MassOmegaMinus; + std::vector v0s; + + // Lambda: p pi-, anti-Lambda: pbar pi+ + for (const int8_t sign : {int8_t{1}, int8_t{-1}}) { + const auto& protons = pools[kClassProton][sign > 0 ? 1 : 0]; + const auto& pions = pools[kClassPion][sign > 0 ? 0 : 1]; + for (const auto& proton : protons) { + for (const auto& pion : pions) { + const auto momentum = RecoDecay::pVec(proton.momentum, pion.momentum); + const float mass = RecoDecay::m(std::array{proton.momentum, pion.momentum}, std::array{o2::constants::physics::MassProton, o2::constants::physics::MassPionCharged}); + if (candOpts.lambdaMassWindow >= 0.f && std::abs(mass - MassLambda) > candOpts.lambdaMassWindow) { + continue; + } + histos.fill(HIST("Candidates/hMassLambda"), RecoDecay::pt(momentum), sign, mass); + const int64_t posId = sign > 0 ? proton.id : pion.id; + const int64_t negId = sign > 0 ? pion.id : proton.id; + v0Table(collision.globalIndex(), posId, negId, sign); + v0s.push_back({v0Table.lastIndex(), pion.id, sign, momentum}); + } + } + } + + // Omega-: Lambda K-, Omega+: anti-Lambda K+; the bachelor has the charge of the pion of the V0 + for (const auto& v0 : v0s) { + const int8_t omegaSign = -v0.sign; + for (const auto& kaon : pools[kClassKaon][omegaSign > 0 ? 1 : 0]) { + if (kaon.id == v0.pionId) { + continue; + } + const float mass = RecoDecay::m(std::array{v0.momentum, kaon.momentum}, std::array{MassLambda, o2::constants::physics::MassKaonCharged}); + if (candOpts.omegaMassWindow >= 0.f && std::abs(mass - MassOmega) > candOpts.omegaMassWindow) { + continue; + } + histos.fill(HIST("Candidates/hMassOmega"), RecoDecay::pt(RecoDecay::pVec(v0.momentum, kaon.momentum)), omegaSign, mass); + cascadeTable(collision.globalIndex(), v0.index, kaon.id, omegaSign); + } + } + } + + void processCandidates(Collisions const& collisions, TracksData const& tracks, aod::BCsWithTimestamps const&) + { + for (const auto& collision : collisions) { + if (!acceptCollision(collision)) { + continue; + } + const auto pools = fillPools(collision, tracks.sliceBy(tracksDataPerCollision, collision.globalIndex())); + buildCandidates(collision, pools); + } + } + PROCESS_SWITCH(doubleOmegaTrackIndexSkimCreator, processCandidates, "V0 and cascade candidates (DOmegaV0s, DOmegaCascades) from the tracks selected by the track BDT, data or MC", false); +}; + +WorkflowSpec defineDataProcessing(ConfigContext const& cfgc) +{ + return WorkflowSpec{adaptAnalysisTask(cfgc)}; +} diff --git a/PWGLF/TableProducer/Strangeness/doubleOmegaTreeCreator.cxx b/PWGLF/TableProducer/Strangeness/doubleOmegaTreeCreator.cxx index 8622f2e626a..d4245f801d4 100644 --- a/PWGLF/TableProducer/Strangeness/doubleOmegaTreeCreator.cxx +++ b/PWGLF/TableProducer/Strangeness/doubleOmegaTreeCreator.cxx @@ -10,11 +10,11 @@ // or submit itself to any jurisdiction. #include "PWGLF/DataModel/LFDoubleOmegaTables.h" +#include "PWGLF/DataModel/LFDoubleOmegaTrackTables.h" #include "PWGLF/DataModel/LFStrangenessTables.h" #include "PWGLF/Utils/doubleOmegaBuilderHelper.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" @@ -41,8 +41,6 @@ #include #include #include -#include -#include #include #include @@ -57,6 +55,7 @@ #include #include #include +#include #include #include #include @@ -100,6 +99,7 @@ struct DoubleOmegaMCInfo { struct MixingSeed { int64_t collisionId = -1; std::array primaryVertex{}; + int runNumber = 0; float bz = 0.f; BuiltOmegaCandidate omega; BuiltLambdaCandidate lambda; @@ -122,6 +122,7 @@ struct doubleOmegaTreeCreator { kFindableSelectedTracks, kFindableCascadeAndV0, kFindableCascadeAndV0SelectedTracks, + kFindableSkimCascadeAndV0, // processFindableTracksWithSkim: in DOmegaCascades + DOmegaV0s (tracks selected by the track BDT) kNFindabilitySteps }; enum MixingType : uint8_t { @@ -204,6 +205,33 @@ struct doubleOmegaTreeCreator { using DerivedMixingBinning = ColumnBinningPolicy; static constexpr float kBzTolerance = 1.e-3f; // kG + // Magnetic fields of two collisions: kaons can be mixed between collisions recorded with the same or + // the opposite field (see DoubleOmegaBuilderHelper::mirrorCharge), Omegas and Lambdas only between + // collisions recorded with the same field. + enum class FieldRelation : uint8_t { + kSame, + kOpposite, + kIncompatible + }; + + static FieldRelation fieldRelation(float bz0, float bz1) + { + if (std::abs(bz0 - bz1) < kBzTolerance) { + return FieldRelation::kSame; + } + if (std::abs(bz0 + bz1) < kBzTolerance) { + return FieldRelation::kOpposite; + } + return FieldRelation::kIncompatible; + } + + // Effective charge of a candidate: candidates of opposite charge recorded with opposite fields have + // the same trajectories. + static int8_t effectiveCharge(int8_t charge, float bz) + { + return bz < 0.f ? static_cast(-charge) : charge; + } + HistogramRegistry histos{"histos", {}, OutputObjHandlingPolicy::AnalysisObject}; template @@ -302,6 +330,7 @@ struct doubleOmegaTreeCreator { printValue(doprocessDataFromCascades); printValue(doprocessMC); printValue(doprocessFindableTracks); + printValue(doprocessFindableTracksWithSkim); printValue(doprocessDerivedData); printValue(doprocessDerivedMixing); printValue(doprocessDerivedMixingCandidates); @@ -310,6 +339,9 @@ struct doubleOmegaTreeCreator { void init(InitContext const&) { printConfig(); + if (doprocessFindableTracks && doprocessFindableTracksWithSkim) { + LOG(fatal) << "processFindableTracks and processFindableTracksWithSkim are alternatives: enable only one of them"; + } ccdb->setURL(ccdburl); ccdb->setCaching(true); @@ -356,6 +388,7 @@ struct doubleOmegaTreeCreator { generatedAndFindable->GetXaxis()->SetBinLabel(kFindableSelectedTracks + 1, "Findable, selected tracks"); generatedAndFindable->GetXaxis()->SetBinLabel(kFindableCascadeAndV0 + 1, "Findable cascade + V0"); generatedAndFindable->GetXaxis()->SetBinLabel(kFindableCascadeAndV0SelectedTracks + 1, "Findable cascade + V0, selected tracks"); + generatedAndFindable->GetXaxis()->SetBinLabel(kFindableSkimCascadeAndV0 + 1, "Findable skim cascade + V0"); if (doprocessDerivedMixingCandidates) { histos.add("Mixing/hSeedRadiusPhi", ";Double-#Omega decay radius (cm);Double-#Omega #varphi (rad);Same-event candidates", HistType::kTH2F, {mixingOpts.binsRadius, mixingOpts.binsPhi}); @@ -369,9 +402,9 @@ struct doubleOmegaTreeCreator { histos.add("Mixing/hMixedCollisions", ";;Entries", HistType::kTH1F, {{4, -0.5f, 3.5f}}); auto mixedCollisions = histos.get(HIST("Mixing/hMixedCollisions")); mixedCollisions->GetXaxis()->SetBinLabel(1, "Pairs"); - mixedCollisions->GetXaxis()->SetBinLabel(2, "Pairs, #Delta#it{z}_{PV}, same #it{B}"); + mixedCollisions->GetXaxis()->SetBinLabel(2, "Pairs, #Delta#it{z}_{PV}, same or opposite #it{B}"); mixedCollisions->GetXaxis()->SetBinLabel(3, "Triples"); - mixedCollisions->GetXaxis()->SetBinLabel(4, "Triples, #Delta#it{z}_{PV}, same #it{B}"); + mixedCollisions->GetXaxis()->SetBinLabel(4, "Triples, #Delta#it{z}_{PV}, same or opposite #it{B}"); histos.add("Mixing/hMixedCandidates", ";;Entries", HistType::kTH1F, {{kNMixingTypes - 1, 0.5f, static_cast(kNMixingTypes) - 0.5f}}); auto mixedCandidates = histos.get(HIST("Mixing/hMixedCandidates")); mixedCandidates->GetXaxis()->SetBinLabel(kMixedOmega, "#Omega mixed"); @@ -841,30 +874,72 @@ struct doubleOmegaTreeCreator { }); } + // A mixed combination of the candidates of up to three collisions, built with respect to the primary + // vertex and in the magnetic field of the reference collision. + struct MixingJob { + int64_t reference = -1; + int64_t omegaSource = -1; + int64_t lambdaSource = -1; + int64_t kaonSource = -1; + bool mirrorKaons = false; // Kaons from a collision recorded with the opposite field. + MixingType type = kSameEvent; + }; + // Two candidates from one collision and one from the other, in both directions. The collision - // providing two candidates is the reference one. - void mixPair(CollisionCandidates const& candidates0, CollisionCandidates const& candidates1) + // providing two candidates is the reference one. With opposite fields only the kaon is mixed. + static void addPairJobs(std::vector& jobs, int64_t collision0, int64_t collision1, FieldRelation relation) { - mix(candidates0, candidates1, candidates0, candidates0, kMixedOmega); - mix(candidates1, candidates0, candidates1, candidates1, kMixedOmega); - mix(candidates0, candidates0, candidates1, candidates0, kMixedLambda); - mix(candidates1, candidates1, candidates0, candidates1, kMixedLambda); - mix(candidates0, candidates0, candidates0, candidates1, kMixedKaon); - mix(candidates1, candidates1, candidates1, candidates0, kMixedKaon); + if (relation == FieldRelation::kSame) { + jobs.push_back({collision0, collision1, collision0, collision0, false, kMixedOmega}); + jobs.push_back({collision1, collision0, collision1, collision1, false, kMixedOmega}); + jobs.push_back({collision0, collision0, collision1, collision0, false, kMixedLambda}); + jobs.push_back({collision1, collision1, collision0, collision1, false, kMixedLambda}); + } + const bool mirrorKaons = relation == FieldRelation::kOpposite; + jobs.push_back({collision0, collision0, collision0, collision1, mirrorKaons, kMixedKaon}); + jobs.push_back({collision1, collision1, collision1, collision0, mirrorKaons, kMixedKaon}); } - // All the assignments of the Omega, Lambda and kaon to three different collisions. The collision - // providing the Omega is the reference one. - void mixTriple(CollisionCandidates const& candidates0, CollisionCandidates const& candidates1, CollisionCandidates const& candidates2) + // All the assignments of the Omega, Lambda and kaon to three different collisions, with the Omega and + // Lambda from collisions recorded with the same field. The collision providing the Omega is the + // reference one. + static void addTripleJobs(std::vector& jobs, std::array const& collisions, std::array const& bz) { - const std::array sources{&candidates0, &candidates1, &candidates2}; std::array roles{0, 1, 2}; // Sources of the Omega, Lambda and kaon. do { - const auto& omegaSource = *sources[roles[0]]; - mix(omegaSource, omegaSource, *sources[roles[1]], *sources[roles[2]], kMixedAll); + const auto [omega, lambda, kaon] = roles; + const auto kaonRelation = fieldRelation(bz[omega], bz[kaon]); + if (fieldRelation(bz[omega], bz[lambda]) != FieldRelation::kSame || kaonRelation == FieldRelation::kIncompatible) { + continue; + } + jobs.push_back({collisions[omega], collisions[omega], collisions[lambda], collisions[kaon], kaonRelation == FieldRelation::kOpposite, kMixedAll}); } while (std::next_permutation(roles.begin(), roles.end())); } + // Builds the mixed combinations grouped by the field of the reference collision, so that the field is + // switched as few times as possible. + void runMixingJobs(std::vector& jobs, std::unordered_map const& candidatesPerCollision) + { + std::stable_sort(jobs.begin(), jobs.end(), [&](MixingJob const& job0, MixingJob const& job1) { + return candidatesPerCollision.at(job0.reference).bz < candidatesPerCollision.at(job1.reference).bz; + }); + std::unordered_map mirroredKaonsPerCollision; + for (const auto& job : jobs) { + const auto& reference = candidatesPerCollision.at(job.reference); + initCCDBForDerived(reference.runNumber, reference.bz); + const auto& kaons = candidatesPerCollision.at(job.kaonSource); + if (job.mirrorKaons) { + auto mirrored = mirroredKaonsPerCollision.find(job.kaonSource); + if (mirrored == mirroredKaonsPerCollision.end()) { + mirrored = mirroredKaonsPerCollision.emplace(job.kaonSource, DoubleOmegaBuilderHelper::mirroredKaons(kaons)).first; + } + mix(reference, candidatesPerCollision.at(job.omegaSource), candidatesPerCollision.at(job.lambdaSource), mirrored->second, job.type); + } else { + mix(reference, candidatesPerCollision.at(job.omegaSource), candidatesPerCollision.at(job.lambdaSource), kaons, job.type); + } + } + } + void processData(Collisions const& collision, TracksFull const& tracks, aod::V0s const& v0s, @@ -939,13 +1014,57 @@ struct doubleOmegaTreeCreator { } PROCESS_SWITCH(doubleOmegaTreeCreator, processMC, "Reconstruct and truth-match Omega + Lambda + kaon in MC", false); - void processFindableTracks(soa::Join const& collisions, - TracksFullMC const& tracks, - aod::Cascades const& cascades, - aod::V0s const& v0s, - aod::McParticles const& mcParticles, - aod::McCollisions const&, - aod::BCsWithTimestamps const&) + // Whether the tracks of the Omega (daughterTrackIds 0-2) form a cascade of the table and those of the direct Lambda + // (daughterTrackIds 3-4) a V0, whatever the order of the tracks; the collisions of the first ones found + template + static void findCascadeAndV0(Cascades const& cascades, V0s const& v0s, std::array const& daughterTrackIds, + bool& omegaInCascade, int64_t& cascadeCollisionId, bool& directLambdaInV0, int64_t& directLambdaCollisionId) + { + omegaInCascade = false; + directLambdaInV0 = false; + for (const auto& cascade : cascades) { + const auto cascadeV0 = v0s.rawIteratorAt(cascade.v0Id()); + const std::array cascadeTrackIds{ + cascadeV0.posTrackId(), cascadeV0.negTrackId(), cascade.bachelorId()}; + bool allOmegaTracksInCascade = true; + for (std::size_t iDaughter = 0; iDaughter < 3; ++iDaughter) { + if (std::find(cascadeTrackIds.begin(), cascadeTrackIds.end(), daughterTrackIds[iDaughter]) == cascadeTrackIds.end()) { + allOmegaTracksInCascade = false; + break; + } + } + if (allOmegaTracksInCascade) { + omegaInCascade = true; + cascadeCollisionId = cascade.collisionId(); + break; + } + } + + for (const auto& v0 : v0s) { + const std::array v0TrackIds{v0.posTrackId(), v0.negTrackId()}; + bool allDirectLambdaTracksInV0 = true; + for (std::size_t iDaughter = 3; iDaughter < 5; ++iDaughter) { + if (std::find(v0TrackIds.begin(), v0TrackIds.end(), daughterTrackIds[iDaughter]) == v0TrackIds.end()) { + allDirectLambdaTracksInV0 = false; + break; + } + } + if (allDirectLambdaTracksInV0) { + directLambdaInV0 = true; + directLambdaCollisionId = v0.collisionId(); + break; + } + } + } + + // skimCascades, skimV0s: tables of doubleOmegaTrackIndexSkimCreator (processFindableTracksWithSkim), nullptr otherwise + void runFindableTracks(soa::Join const& collisions, + TracksFullMC const& tracks, + aod::Cascades const& cascades, + aod::V0s const& v0s, + aod::McParticles const& mcParticles, + aod::DOmegaCascades const* skimCascades, + aod::DOmegaV0s const* skimV0s) { for (const auto& mcParticle : mcParticles) { if (std::abs(mcParticle.pdgCode()) != kDoubleOmegaPdg) { @@ -960,11 +1079,11 @@ struct doubleOmegaTreeCreator { // First stage: double-Omega -> Omega + Lambda + kaon. for (const auto& daughter : mcParticle.daughters_as()) { - if (daughter.pdgCode() == sign * 3334 && omegaId < 0) { + if (daughter.pdgCode() == sign * PDG_t::kOmegaMinus && omegaId < 0) { omegaId = daughter.globalIndex(); - } else if (daughter.pdgCode() == sign * 3122 && directLambdaId < 0) { + } else if (daughter.pdgCode() == sign * PDG_t::kLambda0 && directLambdaId < 0) { directLambdaId = daughter.globalIndex(); - } else if (daughter.pdgCode() == -sign * 321 && directKaonId < 0) { + } else if (daughter.pdgCode() == -sign * PDG_t::kKPlus && directKaonId < 0) { directKaonId = daughter.globalIndex(); } } @@ -978,9 +1097,9 @@ struct doubleOmegaTreeCreator { int64_t omegaLambdaId = -1; int64_t omegaKaonId = -1; for (const auto& daughter : omega.daughters_as()) { - if (daughter.pdgCode() == sign * 3122 && omegaLambdaId < 0) { + if (daughter.pdgCode() == sign * PDG_t::kLambda0 && omegaLambdaId < 0) { omegaLambdaId = daughter.globalIndex(); - } else if (daughter.pdgCode() == -sign * 321 && omegaKaonId < 0) { + } else if (daughter.pdgCode() == -sign * PDG_t::kKPlus && omegaKaonId < 0) { omegaKaonId = daughter.globalIndex(); } } @@ -1065,41 +1184,18 @@ struct doubleOmegaTreeCreator { } if (allDaughtersReconstructed) { histos.fill(HIST("MC/generatedAndFindable"), kFindable); - bool omegaInCascade = false; - int64_t cascadeCollisionId = -1; - for (const auto& cascade : cascades) { - const auto cascadeV0 = cascade.v0(); - const std::array cascadeTrackIds{ - cascadeV0.posTrackId(), cascadeV0.negTrackId(), cascade.bachelorId()}; - bool allOmegaTracksInCascade = true; - for (std::size_t iDaughter = 0; iDaughter < 3; ++iDaughter) { - if (std::find(cascadeTrackIds.begin(), cascadeTrackIds.end(), daughterTrackIds[iDaughter]) == cascadeTrackIds.end()) { - allOmegaTracksInCascade = false; - break; - } - } - if (allOmegaTracksInCascade) { - omegaInCascade = true; - cascadeCollisionId = cascade.collisionId(); - break; - } - } - - bool directLambdaInV0 = false; - int64_t directLambdaCollisionId = -1; - for (const auto& v0 : v0s) { - const std::array v0TrackIds{v0.posTrackId(), v0.negTrackId()}; - bool allDirectLambdaTracksInV0 = true; - for (std::size_t iDaughter = 3; iDaughter < 5; ++iDaughter) { - if (std::find(v0TrackIds.begin(), v0TrackIds.end(), daughterTrackIds[iDaughter]) == v0TrackIds.end()) { - allDirectLambdaTracksInV0 = false; - break; - } - } - if (allDirectLambdaTracksInV0) { - directLambdaInV0 = true; - directLambdaCollisionId = v0.collisionId(); - break; + bool omegaInCascade = false, directLambdaInV0 = false; + int64_t cascadeCollisionId = -1, directLambdaCollisionId = -1; + findCascadeAndV0(cascades, v0s, daughterTrackIds, omegaInCascade, cascadeCollisionId, directLambdaInV0, directLambdaCollisionId); + + // The same with the V0s and cascades of doubleOmegaTrackIndexSkimCreator (tracks selected by the track BDT) + if (skimCascades != nullptr && skimV0s != nullptr) { + bool omegaInSkimCascade = false, directLambdaInSkimV0 = false; + int64_t skimCascadeCollisionId = -1, skimDirectLambdaCollisionId = -1; + findCascadeAndV0(*skimCascades, *skimV0s, daughterTrackIds, omegaInSkimCascade, skimCascadeCollisionId, directLambdaInSkimV0, skimDirectLambdaCollisionId); + // No track selection on top: the skim tracks are already selected by the track BDT + if (omegaInSkimCascade && directLambdaInSkimV0) { + histos.fill(HIST("MC/generatedAndFindable"), kFindableSkimCascadeAndV0); } } @@ -1132,8 +1228,36 @@ struct doubleOmegaTreeCreator { } } } + + void processFindableTracks(soa::Join const& collisions, + TracksFullMC const& tracks, + aod::Cascades const& cascades, + aod::V0s const& v0s, + aod::McParticles const& mcParticles, + aod::McCollisions const&, + aod::BCsWithTimestamps const&) + { + runFindableTracks(collisions, tracks, cascades, v0s, mcParticles, nullptr, nullptr); + } PROCESS_SWITCH(doubleOmegaTreeCreator, processFindableTracks, "Count findable double-Omega decays and reconstruct selected daughter tracks", false); + // As processFindableTracks, and the findable double Omegas are also looked for in the V0s and cascades of + // doubleOmegaTrackIndexSkimCreator (bin "Findable skim cascade + V0" of MC/generatedAndFindable; it needs + // doubleOmegaTrackIndexSkimCreator with processCandidates in the same workflow) + void processFindableTracksWithSkim(soa::Join const& collisions, + TracksFullMC const& tracks, + aod::Cascades const& cascades, + aod::V0s const& v0s, + aod::DOmegaCascades const& skimCascades, + aod::DOmegaV0s const& skimV0s, + aod::McParticles const& mcParticles, + aod::McCollisions const&, + aod::BCsWithTimestamps const&) + { + runFindableTracks(collisions, tracks, cascades, v0s, mcParticles, &skimCascades, &skimV0s); + } + PROCESS_SWITCH(doubleOmegaTreeCreator, processFindableTracksWithSkim, "processFindableTracks, with the findable double Omegas also looked for in the DOmegaCascades and DOmegaV0s tables of doubleOmegaTrackIndexSkimCreator", false); + // Bin index of value on a ConfigurableAxis (variable or fixed width), -1 outside the axis. static int findBin(std::vector const& axis, double value) { @@ -1177,10 +1301,44 @@ struct doubleOmegaTreeCreator { // The track identifiers read from the derived tables (see trackFingerprint) already identify a track // in any dataframe; the collision index is made unique across dataframes. - MixingSeed makeSeed(int64_t collisionId, std::array const& primaryVertex, BuiltOmegaCandidate const& omega, + MixingSeed makeSeed(int64_t collisionId, CollisionCandidates const& candidates, BuiltOmegaCandidate const& omega, BuiltLambdaCandidate const& lambda, BuiltKaonCandidate const& kaon) const { - return {uniqueIndex(collisionId), primaryVertex, mBz, omega, lambda, kaon}; + return {uniqueIndex(collisionId), candidates.primaryVertex, candidates.runNumber, candidates.bz, omega, lambda, kaon}; + } + + // Mixed candidates whose reference seed was recorded with a field different from the current one: + // they are built after the other mixed candidates of the collision, switching the field only once. + struct DeferredMix { + std::array primaryVertex{}; + int runNumber = 0; + float bz = 0.f; + BuiltOmegaCandidate omega; + BuiltLambdaCandidate lambda; + BuiltKaonCandidate kaon; + MixingType type = kSameEvent; + }; + std::vector mDeferredMixes; + + void buildMixedCandidate(std::array const& primaryVertex, BuiltOmegaCandidate const& omega, BuiltLambdaCandidate const& lambda, + BuiltKaonCandidate const& kaon, MixingType mixingType) + { + DoubleOmegaCandidate cand; + if (!helper.buildDoubleOmega(primaryVertex, omega, lambda, kaon, cand)) { + return; + } + histos.fill(HIST("Mixing/hMixedCandidates"), mixingType); + writeMixedCandidate(cand, mixingType); + } + + void buildDeferredMixes() + { + std::stable_sort(mDeferredMixes.begin(), mDeferredMixes.end(), [](DeferredMix const& mix0, DeferredMix const& mix1) { return mix0.bz < mix1.bz; }); + for (const auto& deferred : mDeferredMixes) { + initCCDBForDerived(deferred.runNumber, deferred.bz); + buildMixedCandidate(deferred.primaryVertex, deferred.omega, deferred.lambda, deferred.kaon, deferred.type); + } + mDeferredMixes.clear(); } // Builds and writes one mixed candidate, with the primary vertex of the reference seed. @@ -1194,44 +1352,58 @@ struct doubleOmegaTreeCreator { if (!written.insert(combination).second || DoubleOmegaBuilderHelper::sharesTracks(omega, lambda, kaon)) { return; } - DoubleOmegaCandidate cand; - if (!helper.buildDoubleOmega(reference.primaryVertex, omega, lambda, kaon, cand)) { + if (std::abs(reference.bz - mBz) >= kBzTolerance) { + mDeferredMixes.push_back({reference.primaryVertex, reference.runNumber, reference.bz, omega, lambda, kaon, mixingType}); return; } - histos.fill(HIST("Mixing/hMixedCandidates"), mixingType); - writeMixedCandidate(cand, mixingType); + buildMixedCandidate(reference.primaryVertex, omega, lambda, kaon, mixingType); + } + + // Kaon of a seed as used with a seed recorded with the given field. + static BuiltKaonCandidate kaonForField(MixingSeed const& seed, float bz) + { + return fieldRelation(seed.bz, bz) == FieldRelation::kOpposite ? DoubleOmegaBuilderHelper::mirrorCharge(seed.kaon) : seed.kaon; } // One daughter of each seed replaced by the corresponding daughter of the other seed. The seed - // keeping two of its daughters provides the primary vertex. + // keeping two of its daughters provides the primary vertex. With opposite fields only the kaon is + // mixed, charge conjugated (the seeds have opposite charges, see effectiveCharge). void mixSeedPair(MixingSeed const& seed0, MixingSeed const& seed1, std::set& written) { - mixSeedDaughters(seed0, seed1.omega, seed0.lambda, seed0.kaon, kMixedOmega, written); - mixSeedDaughters(seed1, seed0.omega, seed1.lambda, seed1.kaon, kMixedOmega, written); - mixSeedDaughters(seed0, seed0.omega, seed1.lambda, seed0.kaon, kMixedLambda, written); - mixSeedDaughters(seed1, seed1.omega, seed0.lambda, seed1.kaon, kMixedLambda, written); - mixSeedDaughters(seed0, seed0.omega, seed0.lambda, seed1.kaon, kMixedKaon, written); - mixSeedDaughters(seed1, seed1.omega, seed1.lambda, seed0.kaon, kMixedKaon, written); + if (fieldRelation(seed0.bz, seed1.bz) == FieldRelation::kSame) { + mixSeedDaughters(seed0, seed1.omega, seed0.lambda, seed0.kaon, kMixedOmega, written); + mixSeedDaughters(seed1, seed0.omega, seed1.lambda, seed1.kaon, kMixedOmega, written); + mixSeedDaughters(seed0, seed0.omega, seed1.lambda, seed0.kaon, kMixedLambda, written); + mixSeedDaughters(seed1, seed1.omega, seed0.lambda, seed1.kaon, kMixedLambda, written); + } + mixSeedDaughters(seed0, seed0.omega, seed0.lambda, kaonForField(seed1, seed0.bz), kMixedKaon, written); + mixSeedDaughters(seed1, seed1.omega, seed1.lambda, kaonForField(seed0, seed1.bz), kMixedKaon, written); } // Two daughters of a seed replaced by those of two other seeds: Omega, Lambda and kaon from three - // different collisions, in all the assignments. The seed providing the Omega is the reference. + // different collisions, in all the assignments with the Omega and Lambda recorded with the same field. + // The seed providing the Omega is the reference. void mixSeedTriple(MixingSeed const& seed0, MixingSeed const& seed1, MixingSeed const& seed2, std::set& written) { const std::array seeds{&seed0, &seed1, &seed2}; std::array roles{0, 1, 2}; // Seeds providing the Omega, Lambda and kaon. do { const auto& omegaSeed = *seeds[roles[0]]; - mixSeedDaughters(omegaSeed, omegaSeed.omega, seeds[roles[1]]->lambda, seeds[roles[2]]->kaon, kMixedAll, written); + const auto& lambdaSeed = *seeds[roles[1]]; + const auto& kaonSeed = *seeds[roles[2]]; + if (fieldRelation(omegaSeed.bz, lambdaSeed.bz) != FieldRelation::kSame) { + continue; + } + mixSeedDaughters(omegaSeed, omegaSeed.omega, lambdaSeed.lambda, kaonForField(kaonSeed, omegaSeed.bz), kMixedAll, written); } while (std::next_permutation(roles.begin(), roles.end())); } // Seeds can be mixed if they come from different collisions with close primary vertices, recorded - // with the same magnetic field. + // with the same or the opposite magnetic field. bool canMix(MixingSeed const& seed0, MixingSeed const& seed1) const { return seed0.collisionId != seed1.collisionId && - std::abs(seed0.bz - seed1.bz) < 1.e-3f && + fieldRelation(seed0.bz, seed1.bz) != FieldRelation::kIncompatible && std::abs(seed0.primaryVertex[2] - seed1.primaryVertex[2]) <= mixingOpts.maxDeltaPVz.value; } @@ -1266,8 +1438,9 @@ struct doubleOmegaTreeCreator { } // Same-event candidates of a collision become seeds, binned in decay radius, momentum azimuth and - // charge (mixingOpts.binsRadius, mixingOpts.binsPhi). The seeds are mixed with the pool first and - // added to it afterwards, so seeds of the same collision are never mixed with each other. + // effective charge (mixingOpts.binsRadius, mixingOpts.binsPhi, effectiveCharge). The seeds are mixed + // with the pool first and added to it afterwards, so seeds of the same collision are never mixed with + // each other. void mixCollisionSeeds(int64_t collisionId, CollisionCandidates const& candidates, std::set& written) { const int nPhiBins = nBins(mixingOpts.binsPhi.value); @@ -1282,12 +1455,16 @@ struct doubleOmegaTreeCreator { return; } histos.fill(HIST("Mixing/hSeedRadiusPhi"), radius, cand.phi); - const int key = (radiusBin * nPhiBins + phiBin) * 2 + DoubleOmegaBuilderHelper::chargeIndex(cand.sign); - newSeeds.emplace_back(key, makeSeed(collisionId, candidates.primaryVertex, omega, lambda, kaon)); + const int key = (radiusBin * nPhiBins + phiBin) * 2 + DoubleOmegaBuilderHelper::chargeIndex(effectiveCharge(cand.sign, candidates.bz)); + newSeeds.emplace_back(key, makeSeed(collisionId, candidates, omega, lambda, kaon)); }); for (const auto& [key, seed] : newSeeds) { mixWithPool(seed, mSeedPool[key], written); } + if (!mDeferredMixes.empty()) { + buildDeferredMixes(); + initCCDBForDerived(candidates.runNumber, candidates.bz); + } for (auto& [key, seed] : newSeeds) { auto& pool = mSeedPool[key]; pool.push_back(std::move(seed)); @@ -1331,27 +1508,42 @@ struct doubleOmegaTreeCreator { return static_cast(hash); } - // Derived collisions store the run number and the magnetic field: the propagator, the material LUT - // and the fitters are configured again only when the field changes. + // Field objects of the derived collisions, one per nominal field, so that the mixing can switch between + // the fields of the mixed collisions without accessing the CCDB again. + std::vector> mGrpMagFields; + bool mHasDerivedField = false; + bool mMatLUTLoaded = false; + + // Derived collisions store the run number and the magnetic field: the propagator and the fitters are + // configured again only when the field changes. void initCCDBForDerived(int runNumber, float bz) { - if (mRunNumber != 0 && std::abs(bz - mBz) < kBzTolerance) { + if (mHasDerivedField && std::abs(bz - mBz) < kBzTolerance) { return; } - auto* grpmag = ccdb->getForRun("GLO/Config/GRPMagField", runNumber); - if (!grpmag) { - LOG(fatal) << "Could not retrieve GRPMagField for run " << runNumber; + auto grpmag = std::find_if(mGrpMagFields.begin(), mGrpMagFields.end(), + [&](auto const& field) { return std::abs(field.first - bz) < kBzTolerance; }); + if (grpmag == mGrpMagFields.end()) { + auto* ccdbGrpmag = ccdb->getForRun("GLO/Config/GRPMagField", runNumber); + if (!ccdbGrpmag) { + LOG(fatal) << "Could not retrieve GRPMagField for run " << runNumber; + } + mGrpMagFields.emplace_back(bz, *ccdbGrpmag); + grpmag = std::prev(mGrpMagFields.end()); } - o2::base::Propagator::initFieldFromGRP(grpmag); + o2::base::Propagator::initFieldFromGRP(&grpmag->second); mBz = o2::base::Propagator::Instance()->getNominalBz(); helper.setBz(mBz); - if (static_cast(cfgMaterialCorrection.value) == o2::base::Propagator::MatCorrType::USEMatCorrLUT) { + // The material LUT does not depend on the field; it is set once the field is initialised. + if (!mMatLUTLoaded && static_cast(cfgMaterialCorrection.value) == o2::base::Propagator::MatCorrType::USEMatCorrLUT) { auto* lut = o2::base::MatLayerCylSet::rectifyPtrFromFile(ccdb->getForRun("GLO/Param/MatLUT", runNumber)); if (!lut) { LOG(fatal) << "Could not retrieve material LUT for run " << runNumber; } o2::base::Propagator::Instance()->setMatLUT(lut); + mMatLUTLoaded = true; } + mHasDerivedField = true; mRunNumber = runNumber; } @@ -1363,6 +1555,8 @@ struct doubleOmegaTreeCreator { { CollisionCandidates candidates; candidates.primaryVertex = {collision.posX(), collision.posY(), collision.posZ()}; + candidates.runNumber = collision.runNumber(); + candidates.bz = collision.bz(); for (const auto& row : omegas.sliceBy(derivedOmegasPerCollision, collision.globalIndex())) { const auto posTrack = tracks.rawIteratorAt(row.posTrackId()); @@ -1442,8 +1636,9 @@ struct doubleOmegaTreeCreator { PROCESS_SWITCH(doubleOmegaTreeCreator, processDerivedData, "Same-event Omega + Lambda + kaon from the derived tables of doubleOmegaBuilder", false); // Mixed-event background from the derived tables. Saved derived data can be merged into large - // dataframes, which provides many more mixing partners than the input AO2D dataframes. Only - // collisions recorded with the same magnetic field are mixed. + // dataframes, which provides many more mixing partners than the input AO2D dataframes. Omegas and + // Lambdas are mixed only between collisions recorded with the same magnetic field, kaons also between + // collisions recorded with opposite fields (charge conjugated, see DoubleOmegaBuilderHelper::mirrorCharge). void processDerivedMixing(aod::DOmegaColls const& collisions, aod::DOmegaTracks const& tracks, aod::DOmegaOmegas const& omegas, @@ -1455,33 +1650,33 @@ struct doubleOmegaTreeCreator { } DerivedMixingBinning binning{{mixingOpts.binsVtxZ, mixingOpts.binsMultiplicity}, true}; + std::vector jobs; for (const auto& [collision0, collision1] : selfPairCombinations(binning, mixingOpts.nMixingCollisions, -1, collisions)) { histos.fill(HIST("Mixing/hMixedCollisions"), 0); - if (std::abs(collision0.posZ() - collision1.posZ()) > mixingOpts.maxDeltaPVz || - std::abs(collision0.bz() - collision1.bz()) > kBzTolerance) { + const auto relation = fieldRelation(collision0.bz(), collision1.bz()); + if (std::abs(collision0.posZ() - collision1.posZ()) > mixingOpts.maxDeltaPVz || relation == FieldRelation::kIncompatible) { continue; } histos.fill(HIST("Mixing/hMixedCollisions"), 1); - initCCDBForDerived(collision0.runNumber(), collision0.bz()); - mixPair(candidatesPerCollision.at(collision0.globalIndex()), candidatesPerCollision.at(collision1.globalIndex())); + addPairJobs(jobs, collision0.globalIndex(), collision1.globalIndex(), relation); } - if (!(mixingOpts.mixingTypes & (1 << (kMixedAll - 1)))) { - return; - } - for (const auto& [collision0, collision1, collision2] : selfTripleCombinations(binning, mixingOpts.nMixingCollisions, -1, collisions)) { - histos.fill(HIST("Mixing/hMixedCollisions"), 2); - const auto [minPVz, maxPVz] = std::minmax({collision0.posZ(), collision1.posZ(), collision2.posZ()}); - if (maxPVz - minPVz > mixingOpts.maxDeltaPVz || - std::abs(collision0.bz() - collision1.bz()) > kBzTolerance || - std::abs(collision0.bz() - collision2.bz()) > kBzTolerance) { - continue; + if (mixingOpts.mixingTypes & (1 << (kMixedAll - 1))) { + for (const auto& [collision0, collision1, collision2] : selfTripleCombinations(binning, mixingOpts.nMixingCollisions, -1, collisions)) { + histos.fill(HIST("Mixing/hMixedCollisions"), 2); + const auto [minPVz, maxPVz] = std::minmax({collision0.posZ(), collision1.posZ(), collision2.posZ()}); + if (maxPVz - minPVz > mixingOpts.maxDeltaPVz || + fieldRelation(collision0.bz(), collision1.bz()) == FieldRelation::kIncompatible || + fieldRelation(collision0.bz(), collision2.bz()) == FieldRelation::kIncompatible || + fieldRelation(collision1.bz(), collision2.bz()) == FieldRelation::kIncompatible) { + continue; + } + histos.fill(HIST("Mixing/hMixedCollisions"), 3); + addTripleJobs(jobs, {collision0.globalIndex(), collision1.globalIndex(), collision2.globalIndex()}, + {collision0.bz(), collision1.bz(), collision2.bz()}); } - histos.fill(HIST("Mixing/hMixedCollisions"), 3); - initCCDBForDerived(collision0.runNumber(), collision0.bz()); - mixTriple(candidatesPerCollision.at(collision0.globalIndex()), candidatesPerCollision.at(collision1.globalIndex()), - candidatesPerCollision.at(collision2.globalIndex())); } + runMixingJobs(jobs, candidatesPerCollision); } PROCESS_SWITCH(doubleOmegaTreeCreator, processDerivedMixing, "Mixed-event Omega + Lambda + kaon background from the derived tables of doubleOmegaBuilder", false); @@ -1501,8 +1696,8 @@ struct doubleOmegaTreeCreator { std::set written; // Every combination has a daughter from this dataframe. for (const auto& collision : collisions) { - // Seeds are mixed only with seeds recorded with the same field (canMix), so the field of the - // current collision is the right one for all its mixed candidates. + // Mixed candidates whose reference seed was recorded with the opposite field are built in that + // field (see mixSeedDaughters), the others in the field of the current collision. initCCDBForDerived(collision.runNumber(), collision.bz()); histos.fill(HIST("QA/zVtx"), collision.posZ()); const auto candidates = readCollisionCandidates(collision, tracks, omegas, lambdas); diff --git a/PWGLF/Utils/doubleOmegaBuilderHelper.h b/PWGLF/Utils/doubleOmegaBuilderHelper.h index ac071658516..a137d4e1bdc 100644 --- a/PWGLF/Utils/doubleOmegaBuilderHelper.h +++ b/PWGLF/Utils/doubleOmegaBuilderHelper.h @@ -11,7 +11,11 @@ /// \file doubleOmegaBuilderHelper.h /// \brief Building blocks of the double-Omega (Omega + Lambda + kaon) reconstruction, shared by the -/// candidate builder (derived data producer) and the tree creator. +/// candidate builder (derived data producer) and the tree creator; and the input features of the +/// double-Omega track-selection BDT, computed from a track as in the training: the columns of the +/// doubleOmegaTrackTreeCreator tree and the derived features of DoubleOmega/track_selection/common.py +/// (add_derived_features), with the same names. Missing values (no TOF, failed DCA propagation) are NaN, +/// as in the training (XGBoost missing values). #ifndef PWGLF_UTILS_DOUBLEOMEGABUILDERHELPER_H_ #define PWGLF_UTILS_DOUBLEOMEGABUILDERHELPER_H_ @@ -24,14 +28,16 @@ #include #include #include -#include +#include #include +#include #include #include #include -#include #include +#include +#include #include namespace o2::pwglf::doubleomega @@ -122,6 +128,8 @@ struct BuiltKaonCandidate { // Lambdas and kaons are stored per charge of the Omega they can be paired with (see chargeIndex). struct CollisionCandidates { std::array primaryVertex{}; + int runNumber = 0; + float bz = 0.f; // Nominal magnetic field (kG) the collision was recorded with. std::vector omegas; std::array, 2> lambdas; std::array, 2> kaons; @@ -133,6 +141,7 @@ struct DoubleOmegaSelections { float minNTPCClus = 80.f; float minNCrossedRows = 100.f; float maxNSharedTPCClus = 5.f; + float minCrossedRowsOverFindable = 0.8f; float nSigmaTPCCut = 3.f; float dcaBachToPV = 0.05f; float dcaKaonToPV = 0.05f; @@ -201,6 +210,35 @@ class DoubleOmegaBuilderHelper detectorMap(track)}; } + // Charge-conjugated kaon with the same trajectory in the opposite magnetic field: a K+ recorded with + // field -B follows the trajectory of a K- in the field +B. Only q/pt and its correlations change sign. + static BuiltKaonCandidate mirrorCharge(BuiltKaonCandidate kaon) + { + kaon.track.setQ2Pt(-kaon.track.getQ2Pt()); + for (const int element : {o2::track::kSigQ2PtY, o2::track::kSigQ2PtZ, o2::track::kSigQ2PtSnp, o2::track::kSigQ2PtTgl}) { + kaon.track.setCov(-kaon.track.getCov()[element], element); + } + kaon.sign = -kaon.sign; + return kaon; + } + + // Kaons of a collision mirrored to the opposite magnetic field (see mirrorCharge), stored per charge + // of the Omega they can be paired with. Omegas and Lambdas are not mirrored. + static CollisionCandidates mirroredKaons(CollisionCandidates const& candidates) + { + CollisionCandidates mirrored; + mirrored.primaryVertex = candidates.primaryVertex; + mirrored.runNumber = candidates.runNumber; + mirrored.bz = -candidates.bz; + for (const auto& kaons : candidates.kaons) { + for (const auto& kaon : kaons) { + const auto mirroredKaon = mirrorCharge(kaon); + mirrored.kaons[chargeIndex(mirroredKaon.sign)].push_back(mirroredKaon); + } + } + return mirrored; + } + static bool sharesTracks(BuiltOmegaCandidate const& omega, BuiltLambdaCandidate const& lambda) { const std::array omegaTrackIds{omega.posTrackId, omega.negTrackId, omega.bachelorId}; @@ -228,7 +266,7 @@ class DoubleOmegaBuilderHelper if (track.itsNCls() < selections.minNITSClus || track.tpcNClsFound() < selections.minNTPCClus || track.tpcNClsCrossedRows() < selections.minNCrossedRows || - track.tpcNClsCrossedRows() < 0.8f * track.tpcNClsFindable() || + track.tpcNClsCrossedRows() < selections.minCrossedRowsOverFindable * track.tpcNClsFindable() || track.tpcNClsShared() > selections.maxNSharedTPCClus) { LOG(debug) << "Track failed selection: itsNCls=" << static_cast(track.itsNCls()) << ", tpcNClsFound=" << track.tpcNClsFound() << ", tpcNClsCrossedRows=" << static_cast(track.tpcNClsCrossedRows()) << ", tpcNClsFindable=" << static_cast(track.tpcNClsFindable()) << ", tpcNClsShared=" << static_cast(track.tpcNClsShared()); return false; @@ -264,7 +302,8 @@ class DoubleOmegaBuilderHelper return true; } - template + // isLoose: no PID, DCA or mass-window selection (left to the caller) + template bool buildLambda(C const& collision, Track const& posTrack, Track const& negTrack, int8_t charge, BuiltLambdaCandidate& builtLambda) { if (posTrack.sign() <= 0 || negTrack.sign() >= 0 || @@ -275,7 +314,7 @@ class DoubleOmegaBuilderHelper const bool isMatter = charge < 0; const float protonNSigma = isMatter ? posTrack.tpcNSigmaPr() : negTrack.tpcNSigmaPr(); const float pionNSigma = isMatter ? negTrack.tpcNSigmaPi() : posTrack.tpcNSigmaPi(); - if (std::abs(protonNSigma) > selections.nSigmaTPCCut || std::abs(pionNSigma) > selections.nSigmaTPCCut) { + if ((std::abs(protonNSigma) > selections.nSigmaTPCCut || std::abs(pionNSigma) > selections.nSigmaTPCCut) && !isLoose) { return false; } float dcaXYPos = -999.f; @@ -284,7 +323,7 @@ class DoubleOmegaBuilderHelper LOG(debug) << "Propagation to the PV failed for Lambda daughters " << posTrack.globalIndex() << ", " << negTrack.globalIndex(); return false; } - if (std::abs(dcaXYPos) < selections.dcaV0DauToPV || std::abs(dcaXYNeg) < selections.dcaV0DauToPV) { + if ((std::abs(dcaXYPos) < selections.dcaV0DauToPV || std::abs(dcaXYNeg) < selections.dcaV0DauToPV) && !isLoose) { LOG(debug) << "Lambda dau don't pass selections: " << "nsigmaPR: " << protonNSigma << ", pionNSigma: " << pionNSigma << ", dcaXYPos: " << dcaXYPos << ", dcaXYNeg: " << dcaXYNeg; return false; } @@ -320,8 +359,9 @@ class DoubleOmegaBuilderHelper const float dcaDaughters = std::sqrt(std::abs(fitter.getChi2AtPCACandidate())); const float eta = RecoDecay::eta(lambdaMomentum); if (std::abs(eta) > selections.etaMax || - dcaDaughters > selections.dcaLambdaDaughters || - std::abs(mass - o2::constants::physics::MassLambda0) > selections.mLambdaWindow) { + ((dcaDaughters > selections.dcaLambdaDaughters || + std::abs(mass - o2::constants::physics::MassLambda0) > selections.mLambdaWindow) && + !isLoose)) { return false; } @@ -343,21 +383,22 @@ class DoubleOmegaBuilderHelper return true; } - template + // isLoose: no PID, DCA or mass-window selection (left to the caller) + template bool buildOmega(C const& collision, Track const& posTrack, Track const& negTrack, Track const& bachelor, BuiltOmegaCandidate& builtOmega) { if (bachelor.sign() == 0 || !selectTrack(bachelor) || - std::abs(bachelor.tpcNSigmaKa()) > selections.nSigmaTPCCut) { + (std::abs(bachelor.tpcNSigmaKa()) > selections.nSigmaTPCCut && !isLoose)) { return false; } float dcaXYBachelor = -999.f; - if (!dcaXYToPV(collision, bachelor, dcaXYBachelor) || std::abs(dcaXYBachelor) < selections.dcaBachToPV) { + if (!dcaXYToPV(collision, bachelor, dcaXYBachelor) || (!isLoose && std::abs(dcaXYBachelor) < selections.dcaBachToPV)) { return false; } BuiltLambdaCandidate lambda; - if (!buildLambda(collision, posTrack, negTrack, bachelor.sign(), lambda)) { + if (!buildLambda(collision, posTrack, negTrack, bachelor.sign(), lambda)) { return false; } @@ -391,9 +432,9 @@ class DoubleOmegaBuilderHelper std::array{o2::constants::physics::MassLambda0, o2::constants::physics::MassPionCharged}); const float dcaDaughters = std::sqrt(std::abs(fitter.getChi2AtPCACandidate())); const float eta = RecoDecay::eta(omegaMomentum); - if (dcaDaughters > selections.dcaV0Bach || - std::abs(eta) > selections.etaMax || - std::abs(massOmega - o2::constants::physics::MassOmegaMinus) > selections.mOmegaWindow) { + if (std::abs(eta) > selections.etaMax || + (!isLoose && dcaDaughters > selections.dcaV0Bach) || + (!isLoose && std::abs(massOmega - o2::constants::physics::MassOmegaMinus) > selections.mOmegaWindow)) { return false; } @@ -623,6 +664,138 @@ class DoubleOmegaBuilderHelper } }; +// --------------------------------------------------------------------------------------------- +// Input features of the track-selection BDT + +enum TrackFeature : uint8_t { + kFeatPt = 0, + kFeatEta, + kFeatPhi, + kFeatSign, + kFeatTpcInnerParam, + kFeatTpcSignal, + kFeatTpcNSigmaPr, + kFeatTpcNSigmaPi, + kFeatTpcNSigmaKa, + kFeatTofNSigmaPr, + kFeatTofNSigmaPi, + kFeatTofNSigmaKa, + kFeatDcaXY, + kFeatDcaZ, + kFeatSigmaDcaXY, + kFeatSigmaDcaZ, + kFeatAbsDcaXY, + kFeatAbsDcaZ, + kFeatDcaXYSig, + kFeatDcaZSig, + kFeatItsClusterMap, + kFeatItsNCls, + kFeatItsFirstLayer, + kFeatItsChi2NCl, + kFeatTpcNClsFound, + kFeatTpcNClsCrossedRows, + kFeatTpcNClsShared, + kFeatTpcChi2NCl, + kFeatHasIts, + kFeatHasTof, + kFeatIsPVContributor, + kFeatNumContrib, + kFeatPosZ, + kNTrackFeatures +}; + +// Names of the features, as the columns of the training sample +static constexpr std::array TrackFeatureNames{ + "fPt", "fEta", "fPhi", "fSign", "fTpcInnerParam", "fTpcSignal", + "fTpcNSigmaPr", "fTpcNSigmaPi", "fTpcNSigmaKa", "fTofNSigmaPr", "fTofNSigmaPi", "fTofNSigmaKa", + "fDcaXY", "fDcaZ", "fSigmaDcaXY", "fSigmaDcaZ", "fAbsDcaXY", "fAbsDcaZ", "fDcaXYSig", "fDcaZSig", + "fItsClusterMap", "fItsNCls", "fItsFirstLayer", "fItsChi2NCl", + "fTpcNClsFound", "fTpcNClsCrossedRows", "fTpcNClsShared", "fTpcChi2NCl", + "fHasIts", "fHasTof", "fIsPVContributor", "fNumContrib", "fPosZ"}; + +// Indices of the model inputs (in the order of the model) in the feature array; fatal for unknown names +inline std::vector trackFeatureIndices(std::vector const& names) +{ + std::vector indices; + for (const auto& name : names) { + int index = -1; + for (int i = 0; i < kNTrackFeatures; ++i) { + if (name == TrackFeatureNames[i]) { + index = i; + break; + } + } + if (index < 0) { + LOG(fatal) << "Unknown track-selection feature " << name << " (see PWGLF/Utils/doubleOmegaBuilderHelper.h)"; + } + indices.push_back(index); + } + return indices; +} + +// All the features of a track. The DCA to the primary vertex is computed with the uncertainties of the track +// and of the vertex, as in doubleOmegaTrackTreeCreator. Returns false if the propagation to the PV fails (the +// DCA features are then NaN). +template +bool computeTrackFeatures(C const& collision, T const& track, o2::base::Propagator::MatCorrType matCorr, std::array& features) +{ + constexpr float NaN = std::numeric_limits::quiet_NaN(); + auto trackParCov = getTrackParCov(track); + const o2::dataformats::VertexBase primaryVertex{o2::math_utils::Point3D{collision.posX(), collision.posY(), collision.posZ()}, + std::array{collision.covXX(), collision.covXY(), collision.covYY(), + collision.covXZ(), collision.covYZ(), collision.covZZ()}}; + o2::dataformats::DCA dca; + const bool propagated = o2::base::Propagator::Instance()->propagateToDCABxByBz(primaryVertex, trackParCov, 2.f, matCorr, &dca); + const float dcaXY = propagated ? dca.getY() : NaN; + const float dcaZ = propagated ? dca.getZ() : NaN; + const float sigmaDcaXY = propagated ? std::sqrt(std::max(0.f, static_cast(dca.getSigmaY2()))) : NaN; + const float sigmaDcaZ = propagated ? std::sqrt(std::max(0.f, static_cast(dca.getSigmaZ2()))) : NaN; + + const uint8_t clusterMap = track.itsClusterMap(); + int firstLayer = 7; // no ITS cluster + for (int layer = 0; layer < 7; ++layer) { + if ((clusterMap >> layer) & 1) { + firstLayer = layer; + break; + } + } + + features[kFeatPt] = track.pt(); + features[kFeatEta] = track.eta(); + features[kFeatPhi] = track.phi(); + features[kFeatSign] = track.sign(); + features[kFeatTpcInnerParam] = track.tpcInnerParam(); + features[kFeatTpcSignal] = track.tpcSignal(); + features[kFeatTpcNSigmaPr] = track.tpcNSigmaPr(); + features[kFeatTpcNSigmaPi] = track.tpcNSigmaPi(); + features[kFeatTpcNSigmaKa] = track.tpcNSigmaKa(); + features[kFeatTofNSigmaPr] = track.hasTOF() ? track.tofNSigmaPr() : NaN; + features[kFeatTofNSigmaPi] = track.hasTOF() ? track.tofNSigmaPi() : NaN; + features[kFeatTofNSigmaKa] = track.hasTOF() ? track.tofNSigmaKa() : NaN; + features[kFeatDcaXY] = dcaXY; + features[kFeatDcaZ] = dcaZ; + features[kFeatSigmaDcaXY] = sigmaDcaXY; + features[kFeatSigmaDcaZ] = sigmaDcaZ; + features[kFeatAbsDcaXY] = std::abs(dcaXY); + features[kFeatAbsDcaZ] = std::abs(dcaZ); + features[kFeatDcaXYSig] = sigmaDcaXY > 0.f ? std::abs(dcaXY) / sigmaDcaXY : NaN; + features[kFeatDcaZSig] = sigmaDcaZ > 0.f ? std::abs(dcaZ) / sigmaDcaZ : NaN; + features[kFeatItsClusterMap] = clusterMap; + features[kFeatItsNCls] = __builtin_popcount(clusterMap); + features[kFeatItsFirstLayer] = firstLayer; + features[kFeatItsChi2NCl] = track.itsChi2NCl(); + features[kFeatTpcNClsFound] = track.tpcNClsFound(); + features[kFeatTpcNClsCrossedRows] = track.tpcNClsCrossedRows(); + features[kFeatTpcNClsShared] = track.tpcNClsShared(); + features[kFeatTpcChi2NCl] = track.tpcChi2NCl(); + features[kFeatHasIts] = track.hasITS(); + features[kFeatHasTof] = track.hasTOF(); + features[kFeatIsPVContributor] = track.isPVContributor(); + features[kFeatNumContrib] = collision.numContrib(); + features[kFeatPosZ] = collision.posZ(); + return propagated; +} + } // namespace o2::pwglf::doubleomega #endif // PWGLF_UTILS_DOUBLEOMEGABUILDERHELPER_H_