diff --git a/PWGDQ/Core/MuonMatchingMlResponse.h b/PWGDQ/Core/MuonMatchingMlResponse.h index 1c36f26ab5d..a77deb53397 100644 --- a/PWGDQ/Core/MuonMatchingMlResponse.h +++ b/PWGDQ/Core/MuonMatchingMlResponse.h @@ -40,6 +40,16 @@ break; \ } +// Check if the index of mCachedIndices (index associated to a FEATURE) +// matches the entry in EnumInputFeatures associated to this FEATURE +// if so, the inputFeatures vector is filled with the FEATURE's value +// by calling the corresponding GETTER=FEATURE from track +#define CHECK_AND_FILL_MUONGLOB_TRACK(FEATURE, GETTER) \ + case static_cast(InputFeaturesMFTMuonMatch::FEATURE): { \ + inputFeature = muonglob.GETTER(); \ + break; \ + } + // Check if the index of mCachedIndices (index associated to a FEATURE) // matches the entry in EnumInputFeatures associated to this FEATURE // if so, the inputFeatures vector is filled with the FEATURE's value @@ -112,6 +122,7 @@ enum class InputFeaturesMFTMuonMatch : uint8_t { nClustersMCH, chi2MCH, pdca, + Rabs, cXXMFT, cXYMFT, cYYMFT, @@ -164,7 +175,8 @@ enum class InputFeaturesMFTMuonMatch : uint8_t { centFT0M, centFT0A, centFT0C, - chi2MCHMFT + chi2MCHMFT, + chi2GlobMUON }; template @@ -251,6 +263,33 @@ class MlResponseMFTMuonMatch : public MlResponse return inputFeature; } + template + float returnFeatureGlob(uint8_t idx, T1 const& muonglob, T2 const& muon, T3 const& mft, C1 const& mftcov, U const& collision) + { + float inputFeature = 0.; + switch (idx) { + CHECK_AND_FILL_MFT_TRACK(xMFT, getX); + CHECK_AND_FILL_MFT_TRACK(yMFT, getY); + CHECK_AND_FILL_MFT_TRACK(qOverptMFT, getInvQPt); + CHECK_AND_FILL_MFT_TRACK(tglMFT, getTanl); + CHECK_AND_FILL_MFT_TRACK(phiMFT, getPhi); + CHECK_AND_FILL_MFT_TRACK(chi2MFT, getTrackChi2); + CHECK_AND_FILL_MUON_TRACK(xMCH, getX); + CHECK_AND_FILL_MUON_TRACK(yMCH, getY); + CHECK_AND_FILL_MUON_TRACK(qOverptMCH, getInvQPt); + CHECK_AND_FILL_MUON_TRACK(tglMCH, getTanl); + CHECK_AND_FILL_MUON_TRACK(phiMCH, getPhi); + CHECK_AND_FILL_MUON_TRACK(chi2MCH, getTrackChi2); + CHECK_AND_FILL_MUONGLOB_TRACK(chi2MCHMFT, chi2MatchMCHMFT); + CHECK_AND_FILL_MUONGLOB_TRACK(chi2GlobMUON, chi2); + CHECK_AND_FILL_MUONGLOB_TRACK(Rabs, rAtAbsorberEnd); + // Below are dummy files to remove warning of unused parameters + CHECK_AND_FILL_MFTMUON_COLLISION(posZ); + CHECK_AND_FILL_MFT_COV(cXXMFT, cXX); + } + return inputFeature; + } + template float returnFeatureTest(uint8_t idx, T1 const& muon) { @@ -286,6 +325,17 @@ class MlResponseMFTMuonMatch : public MlResponse return inputFeatures; } + template + std::vector getInputFeaturesGlob(T1 const& muonglob, T2 const& muon, T3 const& mft, C1 const& mftcov, U const& collision) + { + std::vector inputFeatures; + for (const auto& idx : MlResponse::mCachedIndices) { + float inputFeature = returnFeatureGlob(idx, muonglob, muon, mft, mftcov, collision); + inputFeatures.emplace_back(inputFeature); + } + return inputFeatures; + } + /// Method to get the value of variable chosen for binning /// \param track is the single track, \param collision is the collision /// \return binning variable @@ -328,6 +378,7 @@ class MlResponseMFTMuonMatch : public MlResponse FILL_MAP_MFTMUON_MATCH(nClustersMCH), FILL_MAP_MFTMUON_MATCH(chi2MCH), FILL_MAP_MFTMUON_MATCH(pdca), + FILL_MAP_MFTMUON_MATCH(Rabs), FILL_MAP_MFTMUON_MATCH(cXXMFT), FILL_MAP_MFTMUON_MATCH(cXYMFT), FILL_MAP_MFTMUON_MATCH(cYYMFT), @@ -358,7 +409,8 @@ class MlResponseMFTMuonMatch : public MlResponse FILL_MAP_MFTMUON_MATCH(c1PtPhiMCH), FILL_MAP_MFTMUON_MATCH(c1PtTglMCH), FILL_MAP_MFTMUON_MATCH(c1Pt21Pt2MCH), - FILL_MAP_MFTMUON_MATCH(chi2MCHMFT)}; + FILL_MAP_MFTMUON_MATCH(chi2MCHMFT), + FILL_MAP_MFTMUON_MATCH(chi2GlobMUON)}; } uint8_t mCachedIndexBinning; // index correspondance between configurable and available input features diff --git a/PWGDQ/Core/VarManager.cxx b/PWGDQ/Core/VarManager.cxx index eb263b73373..4a5769e0560 100644 --- a/PWGDQ/Core/VarManager.cxx +++ b/PWGDQ/Core/VarManager.cxx @@ -27,6 +27,7 @@ std::map VarManager::fgVarNamesMap; bool VarManager::fgUsedVars[VarManager::kNVars] = {false}; bool VarManager::fgUsedKF = false; float VarManager::fgMagField = 0.5; +float VarManager::fgzMatching = -77.5; float VarManager::fgValues[VarManager::kNVars] = {0.0f}; float VarManager::fgTPCInterSectorBoundary = 1.0; // cm int VarManager::fgITSROFbias = 0; @@ -241,7 +242,7 @@ double VarManager::ComputePIDcalibration(int species, double nSigmaValue) default: LOG(fatal) << "Invalid species for PID calibration: " << species; return -999.0; // Return zero if species is invalid - }; + } TH3F* calibMeanHist = reinterpret_cast(fgCalibs[calibMean]); TH3F* calibSigmaHist = reinterpret_cast(fgCalibs[calibSigma]); @@ -291,7 +292,7 @@ double VarManager::ComputePIDcalibration(int species, double nSigmaValue) default: LOG(fatal) << "Invalid species for PID calibration: " << species; return -999.0; // Return zero if species is invalid - }; + } THnF* calibMeanHist = reinterpret_cast(fgCalibs[calibMean]); THnF* calibSigmaHist = reinterpret_cast(fgCalibs[calibSigma]); @@ -345,7 +346,7 @@ double VarManager::ComputePIDcalibration(int species, double nSigmaValue) default: return nSigmaValue; // unknown status, return the original nSigma value break; - }; + } } else { // unknown calibration type, return the original nSigma value LOG(fatal) << "Unknown calibration type: " << fgCalibrationType; diff --git a/PWGDQ/Core/VarManager.h b/PWGDQ/Core/VarManager.h index f5b7b67f16c..1f5f7d0c59c 100644 --- a/PWGDQ/Core/VarManager.h +++ b/PWGDQ/Core/VarManager.h @@ -891,7 +891,8 @@ class VarManager : public TObject // Index used to set different options for Muon propagation kToVertex = 0, // propagtion to vertex by default kToDCA, - kToRabs + kToRabs, + kToMatching }; static TString fgVariableNames[kNVars]; // variable names @@ -938,6 +939,17 @@ class VarManager : public TObject fgMagField = magField; } + // Setup plane position for MFT-MCH matching + static void SetMatchingPlane(float z) + { + fgzMatching = z; + } + + static float GetMatchingPlane() + { + return fgzMatching; + } + // Setup the 2 prong KFParticle static void SetupTwoProngKFParticle(float magField) { @@ -1047,8 +1059,12 @@ class VarManager : public TObject return deltaPsi; } + template + static o2::track::TrackParCovFwd FwdToTrackPar(const T& track, const C& cov); template static o2::dataformats::GlobalFwdTrack PropagateMuon(const T& muon, const C& collision, int endPoint = kToVertex); + template + static o2::track::TrackParCovFwd PropagateFwd(const T& track, const C& cov, float z); template static void FillMuonPDca(const T& muon, const C& collision, float* values = nullptr); template @@ -1208,6 +1224,7 @@ class VarManager : public TObject static void SetVariableDependencies(); // toggle those variables on which other used variables might depend static float fgMagField; + static float fgzMatching; static float fgCenterOfMassEnergy; // collision energy static float fgMassofCollidingParticle; // mass of the colliding particle static float fgTPCInterSectorBoundary; // TPC inter-sector border size at the TPC outer radius, in cm @@ -1252,6 +1269,19 @@ class VarManager : public TObject ClassDef(VarManager, 4); }; +template +o2::track::TrackParCovFwd VarManager::FwdToTrackPar(const T& track, const C& cov) +{ + double chi2 = track.chi2(); + SMatrix5 tpars(track.x(), track.y(), track.phi(), track.tgl(), track.signed1Pt()); + std::vector v1{cov.cXX(), cov.cXY(), cov.cYY(), cov.cPhiX(), cov.cPhiY(), + cov.cPhiPhi(), cov.cTglX(), cov.cTglY(), cov.cTglPhi(), cov.cTglTgl(), + cov.c1PtX(), cov.c1PtY(), cov.c1PtPhi(), cov.c1PtTgl(), cov.c1Pt21Pt2()}; + SMatrix55 tcovs(v1.begin(), v1.end()); + o2::track::TrackParCovFwd trackparCov{track.z(), tpars, tcovs, chi2}; + return trackparCov; +} + template auto VarManager::getRotatedCovMatrixXX(const T& matrix, U phi, V theta) { @@ -1300,13 +1330,7 @@ KFPTrack VarManager::createKFPTrackFromTrack(const T& track) template KFPTrack VarManager::createKFPFwdTrackFromFwdTrack(const T& muon) { - double chi2 = muon.chi2(); - SMatrix5 tpars(muon.x(), muon.y(), muon.phi(), muon.tgl(), muon.signed1Pt()); - std::vector v1{muon.cXX(), muon.cXY(), muon.cYY(), muon.cPhiX(), muon.cPhiY(), - muon.cPhiPhi(), muon.cTglX(), muon.cTglY(), muon.cTglPhi(), muon.cTglTgl(), - muon.c1PtX(), muon.c1PtY(), muon.c1PtPhi(), muon.c1PtTgl(), muon.c1Pt21Pt2()}; - SMatrix55 tcovs(v1.begin(), v1.end()); - o2::track::TrackParCovFwd trackparCov{muon.z(), tpars, tcovs, chi2}; + o2::track::TrackParCovFwd trackparCov = FwdToTrackPar(muon, muon); std::array trk_cov; trackparCov.getCovXYZPxPyPzGlo(trk_cov); @@ -1321,7 +1345,7 @@ KFPTrack VarManager::createKFPFwdTrackFromFwdTrack(const T& muon) kfpTrack.SetCovarianceMatrix(trkcov_KF); kfpTrack.SetCharge(muon.sign()); kfpTrack.SetNDF(muon.nClusters() - 5); - kfpTrack.SetChi2(chi2); + kfpTrack.SetChi2(muon.chi2()); return kfpTrack; } @@ -1340,19 +1364,13 @@ KFPVertex VarManager::createKFPVertexFromCollision(const T& collision) template o2::dataformats::GlobalFwdTrack VarManager::PropagateMuon(const T& muon, const C& collision, const int endPoint) { - double chi2 = muon.chi2(); - SMatrix5 tpars(muon.x(), muon.y(), muon.phi(), muon.tgl(), muon.signed1Pt()); - std::vector v1{muon.cXX(), muon.cXY(), muon.cYY(), muon.cPhiX(), muon.cPhiY(), - muon.cPhiPhi(), muon.cTglX(), muon.cTglY(), muon.cTglPhi(), muon.cTglTgl(), - muon.c1PtX(), muon.c1PtY(), muon.c1PtPhi(), muon.c1PtTgl(), muon.c1Pt21Pt2()}; - SMatrix55 tcovs(v1.begin(), v1.end()); - o2::track::TrackParCovFwd fwdtrack{muon.z(), tpars, tcovs, chi2}; + o2::track::TrackParCovFwd fwdtrack = FwdToTrackPar(muon, muon); o2::dataformats::GlobalFwdTrack propmuon; if (static_cast(muon.trackType()) > 2) { o2::dataformats::GlobalFwdTrack track; - track.setParameters(tpars); + track.setParameters(fwdtrack.getParameters()); track.setZ(fwdtrack.getZ()); - track.setCovariances(tcovs); + track.setCovariances(fwdtrack.getCovariances()); auto mchTrack = mMatching.FwdtoMCH(track); if (endPoint == kToVertex) { @@ -1364,6 +1382,9 @@ o2::dataformats::GlobalFwdTrack VarManager::PropagateMuon(const T& muon, const C if (endPoint == kToRabs) { o2::mch::TrackExtrap::extrapToZ(mchTrack, -505.); } + if (endPoint == kToMatching) { + o2::mch::TrackExtrap::extrapToVertexWithoutBranson(mchTrack, fgzMatching); + } auto proptrack = mMatching.MCHtoFwd(mchTrack); propmuon.setParameters(proptrack.getParameters()); @@ -1384,6 +1405,14 @@ o2::dataformats::GlobalFwdTrack VarManager::PropagateMuon(const T& muon, const C return propmuon; } +template +o2::track::TrackParCovFwd VarManager::PropagateFwd(const T& track, const C& cov, float z) +{ + o2::track::TrackParCovFwd fwdtrack = FwdToTrackPar(track, cov); + fwdtrack.propagateToZhelix(z, fgMagField); + return fwdtrack; +} + template void VarManager::FillMuonPDca(const T& muon, const C& collision, float* values) { @@ -1491,12 +1520,7 @@ void VarManager::FillGlobalMuonRefitCov(T1 const& muontrack, T2 const& mfttrack, if constexpr ((MuonfillMap & MuonCov) > 0) { if constexpr ((MFTfillMap & MFTCov) > 0) { o2::dataformats::GlobalFwdTrack propmuon = PropagateMuon(muontrack, collision); - SMatrix5 tpars(mfttrack.x(), mfttrack.y(), mfttrack.phi(), mfttrack.tgl(), mfttrack.signed1Pt()); - std::vector v1{mftcov.cXX(), mftcov.cXY(), mftcov.cYY(), mftcov.cPhiX(), mftcov.cPhiY(), - mftcov.cPhiPhi(), mftcov.cTglX(), mftcov.cTglY(), mftcov.cTglPhi(), mftcov.cTglTgl(), - mftcov.c1PtX(), mftcov.c1PtY(), mftcov.c1PtPhi(), mftcov.c1PtTgl(), mftcov.c1Pt21Pt2()}; - SMatrix55 tcovs(v1.begin(), v1.end()); - o2::track::TrackParCovFwd mft{mfttrack.z(), tpars, tcovs, mfttrack.chi2()}; + o2::track::TrackParCovFwd mft = FwdToTrackPar(mfttrack, mftcov); o2::dataformats::GlobalFwdTrack globalRefit = o2::aod::fwdtrackutils::refitGlobalMuonCov(propmuon, mft); values[kX] = globalRefit.getX(); @@ -3685,20 +3709,8 @@ void VarManager::FillPairVertexing(C const& collision, T const& t1, T const& t2, procCode = fgFitterTwoProngBarrel.process(pars1, pars2); } else if constexpr ((pairType == kDecayToMuMu) && muonHasCov) { // Initialize track parameters for forward - double chi21 = t1.chi2(); - double chi22 = t2.chi2(); - SMatrix5 t1pars(t1.x(), t1.y(), t1.phi(), t1.tgl(), t1.signed1Pt()); - std::vector v1{t1.cXX(), t1.cXY(), t1.cYY(), t1.cPhiX(), t1.cPhiY(), - t1.cPhiPhi(), t1.cTglX(), t1.cTglY(), t1.cTglPhi(), t1.cTglTgl(), - t1.c1PtX(), t1.c1PtY(), t1.c1PtPhi(), t1.c1PtTgl(), t1.c1Pt21Pt2()}; - SMatrix55 t1covs(v1.begin(), v1.end()); - o2::track::TrackParCovFwd pars1{t1.z(), t1pars, t1covs, chi21}; - SMatrix5 t2pars(t2.x(), t2.y(), t2.phi(), t2.tgl(), t2.signed1Pt()); - std::vector v2{t2.cXX(), t2.cXY(), t2.cYY(), t2.cPhiX(), t2.cPhiY(), - t2.cPhiPhi(), t2.cTglX(), t2.cTglY(), t2.cTglPhi(), t2.cTglTgl(), - t2.c1PtX(), t2.c1PtY(), t2.c1PtPhi(), t2.c1PtTgl(), t2.c1Pt21Pt2()}; - SMatrix55 t2covs(v2.begin(), v2.end()); - o2::track::TrackParCovFwd pars2{t2.z(), t2pars, t2covs, chi22}; + o2::track::TrackParCovFwd pars1 = FwdToTrackPar(t1, t1); + o2::track::TrackParCovFwd pars2 = FwdToTrackPar(t2, t2); procCode = fgFitterTwoProngFwd.process(pars1, pars2); } else { return; @@ -3960,20 +3972,8 @@ void VarManager::FillPairVertexing(C const& collision, T const& t1, T const& t2, } if (propToSV) { if constexpr ((pairType == kDecayToMuMu) && muonHasCov) { - double chi21 = t1.chi2(); - double chi22 = t2.chi2(); - SMatrix5 t1pars(t1.x(), t1.y(), t1.phi(), t1.tgl(), t1.signed1Pt()); - std::vector c1{t1.cXX(), t1.cXY(), t1.cYY(), t1.cPhiX(), t1.cPhiY(), - t1.cPhiPhi(), t1.cTglX(), t1.cTglY(), t1.cTglPhi(), t1.cTglTgl(), - t1.c1PtX(), t1.c1PtY(), t1.c1PtPhi(), t1.c1PtTgl(), t1.c1Pt21Pt2()}; - SMatrix55 t1covs(c1.begin(), c1.end()); - o2::track::TrackParCovFwd pars1{t1.z(), t1pars, t1covs, chi21}; - SMatrix5 t2pars(t2.x(), t2.y(), t2.phi(), t2.tgl(), t2.signed1Pt()); - std::vector c2{t2.cXX(), t2.cXY(), t2.cYY(), t2.cPhiX(), t2.cPhiY(), - t2.cPhiPhi(), t2.cTglX(), t2.cTglY(), t2.cTglPhi(), t2.cTglTgl(), - t2.c1PtX(), t2.c1PtY(), t2.c1PtPhi(), t2.c1PtTgl(), t2.c1Pt21Pt2()}; - SMatrix55 t2covs(c2.begin(), c2.end()); - o2::track::TrackParCovFwd pars2{t2.z(), t2pars, t2covs, chi22}; + o2::track::TrackParCovFwd pars1 = FwdToTrackPar(t1, t1); + o2::track::TrackParCovFwd pars2 = FwdToTrackPar(t2, t2); auto geoMan1 = o2::base::GeometryManager::meanMaterialBudget(t1.x(), t1.y(), t1.z(), KFGeoTwoProng.GetX(), KFGeoTwoProng.GetY(), KFGeoTwoProng.GetZ()); auto geoMan2 = o2::base::GeometryManager::meanMaterialBudget(t2.x(), t2.y(), t2.z(), KFGeoTwoProng.GetX(), KFGeoTwoProng.GetY(), KFGeoTwoProng.GetZ()); @@ -4282,29 +4282,10 @@ void VarManager::FillDileptonTrackVertexing(C const& collision, T1 const& lepton mlepton2 = o2::constants::physics::MassMuon; mtrack = o2::constants::physics::MassMuon; - double chi21 = lepton1.chi2(); - double chi22 = lepton2.chi2(); - double chi23 = track.chi2(); - SMatrix5 t1pars(lepton1.x(), lepton1.y(), lepton1.phi(), lepton1.tgl(), lepton1.signed1Pt()); - std::vector v1{lepton1.cXX(), lepton1.cXY(), lepton1.cYY(), lepton1.cPhiX(), lepton1.cPhiY(), - lepton1.cPhiPhi(), lepton1.cTglX(), lepton1.cTglY(), lepton1.cTglPhi(), lepton1.cTglTgl(), - lepton1.c1PtX(), lepton1.c1PtY(), lepton1.c1PtPhi(), lepton1.c1PtTgl(), lepton1.c1Pt21Pt2()}; - SMatrix55 t1covs(v1.begin(), v1.end()); - o2::track::TrackParCovFwd pars1{lepton1.z(), t1pars, t1covs, chi21}; - - SMatrix5 t2pars(lepton2.x(), lepton2.y(), lepton2.phi(), lepton2.tgl(), lepton2.signed1Pt()); - std::vector v2{lepton2.cXX(), lepton2.cXY(), lepton2.cYY(), lepton2.cPhiX(), lepton2.cPhiY(), - lepton2.cPhiPhi(), lepton2.cTglX(), lepton2.cTglY(), lepton2.cTglPhi(), lepton2.cTglTgl(), - lepton2.c1PtX(), lepton2.c1PtY(), lepton2.c1PtPhi(), lepton2.c1PtTgl(), lepton2.c1Pt21Pt2()}; - SMatrix55 t2covs(v2.begin(), v2.end()); - o2::track::TrackParCovFwd pars2{lepton2.z(), t2pars, t2covs, chi22}; - - SMatrix5 t3pars(track.x(), track.y(), track.phi(), track.tgl(), track.signed1Pt()); - std::vector v3{track.cXX(), track.cXY(), track.cYY(), track.cPhiX(), track.cPhiY(), - track.cPhiPhi(), track.cTglX(), track.cTglY(), track.cTglPhi(), track.cTglTgl(), - track.c1PtX(), track.c1PtY(), track.c1PtPhi(), track.c1PtTgl(), track.c1Pt21Pt2()}; - SMatrix55 t3covs(v3.begin(), v3.end()); - o2::track::TrackParCovFwd pars3{track.z(), t3pars, t3covs, chi23}; + o2::track::TrackParCovFwd pars1 = FwdToTrackPar(lepton1, lepton1); + o2::track::TrackParCovFwd pars2 = FwdToTrackPar(lepton2, lepton2); + o2::track::TrackParCovFwd pars3 = FwdToTrackPar(track, track); + procCode = VarManager::fgFitterThreeProngFwd.process(pars1, pars2, pars3); procCodeJpsi = VarManager::fgFitterTwoProngFwd.process(pars1, pars2); } else if constexpr ((candidateType == kBtoJpsiEEK || candidateType == kDstarToD0KPiPi) && trackHasCov) { diff --git a/PWGDQ/TableProducer/tableMakerMC_withAssoc.cxx b/PWGDQ/TableProducer/tableMakerMC_withAssoc.cxx index cd38a312311..23cfd2d9042 100644 --- a/PWGDQ/TableProducer/tableMakerMC_withAssoc.cxx +++ b/PWGDQ/TableProducer/tableMakerMC_withAssoc.cxx @@ -217,6 +217,7 @@ struct TableMakerMC { Configurable fUseML{"cfgUseML", false, "Import ONNX model from ccdb to decide which matching candidates to keep"}; Configurable fMuonMatchEtaMin{"cfgMuonMatchEtaMin", -4.0f, "Definition of the acceptance of muon tracks to be matched with MFT"}; Configurable fMuonMatchEtaMax{"cfgMuonMatchEtaMax", -2.5f, "Definition of the acceptance of muon tracks to be matched with MFT"}; + Configurable fzMatching{"cfgzMatching", -77.5f, "Plane for MFT-MCH matching"}; Configurable> fModelPathsCCDB{"fModelPathsCCDB", std::vector{"Users/m/mcoquet/MLTest"}, "Paths of models on CCDB"}; Configurable> fInputFeatures{"cfgInputFeatures", std::vector{"chi2MCHMFT"}, "Names of ML model input features"}; Configurable> fModelNames{"cfgModelNames", std::vector{"model.onnx"}, "ONNX file names for each pT bin (if not from CCDB full path)"}; @@ -379,6 +380,7 @@ struct TableMakerMC { fCCDB->get(fConfigCCDB.fGeoPath); } fCCDBApi.init(fConfigCCDB.fConfigCcdbUrl.value); + VarManager::SetMatchingPlane(fConfigVariousOptions.fzMatching.value); if (fConfigVariousOptions.fUseML.value) { // TODO : for now we use hard coded values since the current models use 1 pT bin @@ -929,6 +931,36 @@ struct TableMakerMC { } } + template + void skimBestMuonMatchesML(TMuons const& muons, TMFTTracks const& /*mfttracks*/, TMFTCovs const& mfCovs, TEvent const& collision) + { + std::unordered_map> mCandidates; + for (const auto& muon : muons) { + if (static_cast(muon.trackType()) < 2) { + auto muonID = muon.matchMCHTrackId(); + auto muontrack = muon.template matchMCHTrack_as(); + auto muonprop = VarManager::PropagateMuon(muontrack, collision, VarManager::kToMatching); + auto mfttrack = muon.template matchMFTTrack_as(); + auto const& mfttrackcov = mfCovs.rawIteratorAt(map_mfttrackcovs[mfttrack.globalIndex()]); + o2::track::TrackParCovFwd mftprop = VarManager::PropagateFwd(mfttrack, mfttrackcov, VarManager::GetMatchingPlane()); + std::vector output; + std::vector inputML = matchingMlResponse.getInputFeaturesGlob(muon, muonprop, mftprop, mfttrackcov, collision); + matchingMlResponse.isSelectedMl(inputML, 0, output); + float score = output[0]; + if (mCandidates.find(muonID) == mCandidates.end()) { + mCandidates[muonID] = {score, muon.globalIndex()}; + } else { + if (score < mCandidates[muonID].first) { + mCandidates[muonID] = {score, muon.globalIndex()}; + } + } + } + } + for (auto& pairCand : mCandidates) { + fBestMatch[pairCand.second.second] = true; + } + } + template void skimMuons(TEvent const& collision, TMuons const& muons, FwdTrackAssoc const& muonAssocs, aod::McParticles const& mcTracks, TMFTTracks const& /*mftTracks*/, TMFTCovs const& mfCovs) { @@ -1154,6 +1186,7 @@ struct TableMakerMC { fGrpMag = fCCDB->getForTimeStamp(fConfigCCDB.fGrpMagPath, bcs.begin().timestamp()); if (fGrpMag != nullptr) { o2::base::Propagator::initFieldFromGRP(fGrpMag); + VarManager::SetMagneticField(fGrpMag->getNominalL3Field()); } if (fConfigVariousOptions.fPropMuon) { VarManager::SetupMuonMagField(); @@ -1244,7 +1277,13 @@ struct TableMakerMC { if constexpr (static_cast(TMFTFillMap)) { auto groupedMuonIndices = fwdTrackAssocs.sliceBy(fwdtrackIndicesPerCollision, origIdx); if (fConfigVariousOptions.fKeepBestMatch) { - skimBestMuonMatches(muons); + if constexpr (static_cast(TMFTFillMap & VarManager::ObjTypes::MFTCov)) { + if (fConfigVariousOptions.fUseML.value) { + skimBestMuonMatchesML(muons, mftTracks, mftCovs, collision); + } + } else { + skimBestMuonMatches(muons); + } } if constexpr (static_cast(TMFTFillMap & VarManager::ObjTypes::MFTCov)) { skimMuons(collision, muons, groupedMuonIndices, mcParticles, mftTracks, mftCovs); diff --git a/PWGDQ/TableProducer/tableMaker_withAssoc.cxx b/PWGDQ/TableProducer/tableMaker_withAssoc.cxx index 97cf93bfd87..415c237d3ad 100644 --- a/PWGDQ/TableProducer/tableMaker_withAssoc.cxx +++ b/PWGDQ/TableProducer/tableMaker_withAssoc.cxx @@ -21,6 +21,8 @@ #include #include #include +#include +#include #include // other includes #include "PWGDQ/Core/AnalysisCompositeCut.h" @@ -28,6 +30,7 @@ #include "PWGDQ/Core/CutsLibrary.h" #include "PWGDQ/Core/HistogramManager.h" #include "PWGDQ/Core/HistogramsLibrary.h" +#include "PWGDQ/Core/MuonMatchingMlResponse.h" #include "PWGDQ/Core/VarManager.h" #include "PWGDQ/DataModel/ReducedInfoTables.h" @@ -255,8 +258,13 @@ struct TableMaker { Configurable fPropMuon{"cfgPropMuon", true, "Propagate muon tracks through absorber (do not use if applying pairing)"}; Configurable fRefitGlobalMuon{"cfgRefitGlobalMuon", true, "Correct global muon parameters"}; Configurable fKeepBestMatch{"cfgKeepBestMatch", false, "Keep only the best match global muons in the skimming"}; + Configurable fUseML{"cfgUseML", false, "Import ONNX model from ccdb to decide which matching candidates to keep"}; Configurable fMuonMatchEtaMin{"cfgMuonMatchEtaMin", -4.0f, "Definition of the acceptance of muon tracks to be matched with MFT"}; Configurable fMuonMatchEtaMax{"cfgMuonMatchEtaMax", -2.5f, "Definition of the acceptance of muon tracks to be matched with MFT"}; + Configurable fzMatching{"cfgzMatching", -77.5f, "Plane for MFT-MCH matching"}; + Configurable> fModelPathsCCDB{"fModelPathsCCDB", std::vector{"Users/m/mcoquet/MLTest"}, "Paths of models on CCDB"}; + Configurable> fInputFeatures{"cfgInputFeatures", std::vector{"chi2MCHMFT"}, "Names of ML model input features"}; + Configurable> fModelNames{"cfgModelNames", std::vector{"model.onnx"}, "ONNX file names for each pT bin (if not from CCDB full path)"}; // TPC occupancy related variables Configurable fTPCShortPast{"cfgTPCShortPast", 8.0f, "Time in short past to look for occupancy (micro-seconds)"}; @@ -289,6 +297,12 @@ struct TableMaker { std::map fBestMatch; std::unordered_map map_mfttrackcovs; + + o2::analysis::MlResponseMFTMuonMatch matchingMlResponse; + std::vector binsPtMl; + std::array cutValues; + std::vector cutDirMl; + // FIXME: For now, the skimming is done using the Common track-collision association task, which does not allow to use // our own Filtered tracks. If the filter is very selective, then it may be worth to run the association in this workflow // using the Common/CollisionAssociation class @@ -423,6 +437,20 @@ struct TableMaker { } VarManager::SetUseVars(fHistMan->GetUsedVars()); // provide the list of required variables so that VarManager knows what to fill fOutputList.setObject(fHistMan->GetMainHistogramList()); + + VarManager::SetMatchingPlane(fConfigVariousOptions.fzMatching.value); + + if (fConfigVariousOptions.fUseML.value) { + // TODO : for now we use hard coded values since the current models use 1 pT bin + binsPtMl = {-1e-6, 1000.0}; + cutValues = {0.0}; + cutDirMl = {cuts_ml::CutNot}; + o2::framework::LabeledArray mycutsMl(cutValues.data(), 1, 1, std::vector{"pT bin 0"}, std::vector{"score"}); + matchingMlResponse.configure(binsPtMl, mycutsMl, cutDirMl, 1); + matchingMlResponse.setModelPathsCCDB(fConfigVariousOptions.fModelNames.value, fCCDBApi, fConfigVariousOptions.fModelPathsCCDB.value, fConfigCCDB.fConfigNoLaterThan.value); + matchingMlResponse.cacheInputFeaturesIndices(fConfigVariousOptions.fInputFeatures.value); + matchingMlResponse.init(); + } } void DefineCuts() @@ -1157,6 +1185,36 @@ struct TableMaker { } } + template + void skimBestMuonMatchesML(TMuons const& muons, TMFTTracks const& /*mfttracks*/, TMFTCovs const& mfCovs, TEvent const& collision) + { + std::unordered_map> mCandidates; + for (const auto& muon : muons) { + if (static_cast(muon.trackType()) < 2) { + auto muonID = muon.matchMCHTrackId(); + auto muontrack = muon.template matchMCHTrack_as(); + auto muonprop = VarManager::PropagateMuon(muontrack, collision, VarManager::kToMatching); + auto mfttrack = muon.template matchMFTTrack_as(); + auto const& mfttrackcov = mfCovs.rawIteratorAt(map_mfttrackcovs[mfttrack.globalIndex()]); + o2::track::TrackParCovFwd mftprop = VarManager::PropagateFwd(mfttrack, mfttrackcov, VarManager::GetMatchingPlane()); + std::vector output; + std::vector inputML = matchingMlResponse.getInputFeaturesGlob(muon, muonprop, mftprop, mfttrackcov, collision); + matchingMlResponse.isSelectedMl(inputML, 0, output); + float score = output[0]; + if (mCandidates.find(muonID) == mCandidates.end()) { + mCandidates[muonID] = {score, muon.globalIndex()}; + } else { + if (score < mCandidates[muonID].first) { + mCandidates[muonID] = {score, muon.globalIndex()}; + } + } + } + } + for (auto& pairCand : mCandidates) { + fBestMatch[pairCand.second.second] = true; + } + } + template void skimMuons(TEvent const& collision, TBCs const& /*bcs*/, TMuons const& muons, FwdTrackAssoc const& muonAssocs, TMFTTracks const& /*mftTracks*/, TMFTCovs const& mfCovs) { @@ -1368,6 +1426,7 @@ struct TableMaker { fGrpMag = fCCDB->getForTimeStamp(fConfigCCDB.fConfigGrpMagPath, bcs.begin().timestamp()); if (fGrpMag != nullptr) { o2::base::Propagator::initFieldFromGRP(fGrpMag); + VarManager::SetMagneticField(fGrpMag->getNominalL3Field()); } if (fConfigVariousOptions.fPropMuon) { VarManager::SetupMuonMagField(); @@ -1445,7 +1504,13 @@ struct TableMaker { if constexpr (static_cast(TMFTFillMap)) { auto groupedMuonIndices = fwdTrackAssocs.sliceBy(fwdtrackIndicesPerCollision, origIdx); if (fConfigVariousOptions.fKeepBestMatch) { - skimBestMuonMatches(muons); + if constexpr (static_cast(TMFTFillMap & VarManager::ObjTypes::MFTCov)) { + if (fConfigVariousOptions.fUseML.value) { + skimBestMuonMatchesML(muons, mftTracks, mftCovs, collision); + } + } else { + skimBestMuonMatches(muons); + } } if constexpr (static_cast(TMFTFillMap & VarManager::ObjTypes::MFTCov)) { skimMuons(collision, bcs, muons, groupedMuonIndices, mftTracks, mftCovs);