[PWGHF] Implement ML models in D0-h correlation (AliceO2Group#8742)

Co-authored-by: ALICE Action Bot <alibuild@cern.ch>
Archita-Dash · Dec 11, 2024 · 65c50b2 · 65c50b2
1 parent 642e810
commit 65c50b2
Show file tree

Hide file tree

Showing 3 changed files with 196 additions and 126 deletions.
diff --git a/PWGHF/HFC/DataModel/CorrelationTables.h b/PWGHF/HFC/DataModel/CorrelationTables.h
@@ -57,6 +57,12 @@ DECLARE_SOA_COLUMN(PtD, ptD, float);                          //! Transverse mom
 DECLARE_SOA_COLUMN(PtHadron, ptHadron, float);                //! Transverse momentum of Hadron
 DECLARE_SOA_COLUMN(MD, mD, float);                            //! Invariant mass of D0
 DECLARE_SOA_COLUMN(MDbar, mDbar, float);                      //! Invariant mass of D0bar
+DECLARE_SOA_COLUMN(MlScoreBkgD0, mlScoreBkgD0, float);        //! ML background score for D0 selection
+DECLARE_SOA_COLUMN(MlScoreNonPromptD0, mlScoreNonPromptD0, float);       //! ML prompt score for D0 selection
+DECLARE_SOA_COLUMN(MlScorePromptD0, mlScorePromptD0, float);             //! ML prompt score for D0 selection
+DECLARE_SOA_COLUMN(MlScoreBkgD0bar, mlScoreBkgD0bar, float);             //! ML background score for D0 selection
+DECLARE_SOA_COLUMN(MlScoreNonPromptD0bar, mlScoreNonPromptD0bar, float); //! ML prompt score for D0 selection
+DECLARE_SOA_COLUMN(MlScorePromptD0bar, mlScorePromptD0bar, float);       //! ML prompt score for D0 selection
 DECLARE_SOA_COLUMN(SignalStatus, signalStatus, int);          //! Tag for D0,D0bar
 DECLARE_SOA_COLUMN(PoolBin, poolBin, int);                    //! Pool Bin for the MixedEvent
 DECLARE_SOA_COLUMN(IsAutoCorrelated, isAutoCorrelated, bool); //! Correlation Status
@@ -81,19 +87,36 @@ enum ParticleTypeMcRec {
 };
 } // namespace hf_correlation_d0_hadron
 
-DECLARE_SOA_TABLE(DHadronPair, "AOD", "DHADRONPAIR", //! D0-Hadrons pairs Informations
+DECLARE_SOA_TABLE(D0HadronPair, "AOD", "D0HPAIR", //! D0-Hadrons pairs Informations
                   aod::hf_correlation_d0_hadron::DeltaPhi,
                   aod::hf_correlation_d0_hadron::DeltaEta,
                   aod::hf_correlation_d0_hadron::PtD,
                   aod::hf_correlation_d0_hadron::PtHadron,
                   aod::hf_correlation_d0_hadron::PoolBin,
                   aod::hf_correlation_d0_hadron::IsAutoCorrelated);
 
-DECLARE_SOA_TABLE(DHadronRecoInfo, "AOD", "DHADRONRECOINFO", //! D0-Hadrons pairs Reconstructed Informations
+DECLARE_SOA_TABLE(D0HadronRecoInfo, "AOD", "D0HRECOINFO", //! D0-Hadrons pairs Reconstructed Informations
                   aod::hf_correlation_d0_hadron::MD,
                   aod::hf_correlation_d0_hadron::MDbar,
                   aod::hf_correlation_d0_hadron::SignalStatus);
 
+DECLARE_SOA_TABLE(D0HadronMlInfo, "AOD", "D0HMLINFO", //! D0-Hadrons pairs Machine Learning Information
+                  aod::hf_correlation_d0_hadron::MlScoreBkgD0,
+                  aod::hf_correlation_d0_hadron::MlScoreNonPromptD0,
+                  aod::hf_correlation_d0_hadron::MlScorePromptD0,
+                  aod::hf_correlation_d0_hadron::MlScoreBkgD0bar,
+                  aod::hf_correlation_d0_hadron::MlScoreNonPromptD0bar,
+                  aod::hf_correlation_d0_hadron::MlScorePromptD0bar);
+
+DECLARE_SOA_TABLE(D0CandRecoInfo, "AOD", "D0CANDRECOINFO", //! Ds candidates Reconstructed Information
+                  aod::hf_correlation_d0_hadron::MD,
+                  aod::hf_correlation_d0_hadron::MDbar,
+                  aod::hf_correlation_d0_hadron::PtD,
+                  aod::hf_correlation_d0_hadron::MlScoreBkgD0,
+                  aod::hf_correlation_d0_hadron::MlScorePromptD0,
+                  aod::hf_correlation_d0_hadron::MlScoreBkgD0bar,
+                  aod::hf_correlation_d0_hadron::MlScorePromptD0bar);
+
 // Note: definition of columns and tables for Lc-Hadron correlation pairs
 namespace hf_correlation_lc_hadron
 {