From e1a6b2446745298e15ba804bba7f90ca15ccfade Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?BiaoZhang=20=28=E5=BC=A0=E5=BD=AA=29?=
 <52267892+zhangbiao-phy@users.noreply.github.com>
Date: Sun, 21 Dec 2025 23:39:28 +0100
Subject: [PATCH] Update producerCharmHadronsTrackFemtoDream.cxx

---
 .../producerCharmHadronsTrackFemtoDream.cxx   | 35 ++++++++++++++++---
 1 file changed, 30 insertions(+), 5 deletions(-)

diff --git a/PWGHF/HFC/TableProducer/producerCharmHadronsTrackFemtoDream.cxx b/PWGHF/HFC/TableProducer/producerCharmHadronsTrackFemtoDream.cxx
index f40df1c5ddf..013568b778c 100644
--- a/PWGHF/HFC/TableProducer/producerCharmHadronsTrackFemtoDream.cxx
+++ b/PWGHF/HFC/TableProducer/producerCharmHadronsTrackFemtoDream.cxx
@@ -269,6 +269,8 @@ struct HfProducerCharmHadronsTrackFemtoDream {
 
     bool useLcMl = doprocessDataLcToPKPiWithML || doprocessMcLcToPKPiWithML;
     bool useDplusMl = doprocessDataDplusToPiKPiWithML || doprocessMcDplusToPiKPiWithML;
+    bool useD0Ml = doprocessDataD0ToPiKWithML || doprocessMcD0ToPiKWithML;
+    bool useDstarMl = doprocessDataDstarToD0PiWithML || doprocessMcDstarToD0PiWithML;
 
     if (applyMlMode == FillMlFromNewBDT) {
       if (useLcMl) {
@@ -281,6 +283,16 @@ struct HfProducerCharmHadronsTrackFemtoDream {
         hfMlResponseDplus.cacheInputFeaturesIndices(namesInputFeatures);
         hfMlResponseDplus.init();
       }
+      if (useD0Ml) {
+        hfMlResponseD0.configure(binsPtMl, cutsMl, cutDirMl, nClassesMl);
+        hfMlResponseD0.cacheInputFeaturesIndices(namesInputFeatures);
+        hfMlResponseD0.init();
+      }
+      if (useDstarMl) {
+        hfMlResponseDstar.configure(binsPtMl, cutsMl, cutDirMl, nClassesMl);
+        hfMlResponseDstar.cacheInputFeaturesIndices(namesInputFeatures);
+        hfMlResponseDstar.init();
+      }
 
       if (loadModelsFromCCDB) {
         ccdbApi.init(ccdbUrl);
@@ -290,6 +302,13 @@ struct HfProducerCharmHadronsTrackFemtoDream {
         if (useDplusMl) {
           hfMlResponseDplus.setModelPathsCCDB(onnxFileNames, ccdbApi, modelPathsCCDB, timestampCCDB);
         }
+        if (useD0Ml) {
+          hfMlResponseD0.setModelPathsCCDB(onnxFileNames, ccdbApi, modelPathsCCDB, timestampCCDB);
+        }
+        if (useDstarMl) {
+          hfMlResponseDstar.setModelPathsCCDB(onnxFileNames, ccdbApi, modelPathsCCDB, timestampCCDB);
+        }
+
       } else {
         if (useLcMl) {
           hfMlResponseLc.setModelPathsLocal(onnxFileNames);
@@ -297,6 +316,12 @@ struct HfProducerCharmHadronsTrackFemtoDream {
         if (useDplusMl) {
           hfMlResponseDplus.setModelPathsLocal(onnxFileNames);
         }
+        if (useD0Ml) {
+          hfMlResponseD0.setModelPathsLocal(onnxFileNames);
+        }
+        if (useDstarMl) {
+          hfMlResponseDstar.setModelPathsLocal(onnxFileNames);
+        }
       }
     }
   }
@@ -679,9 +704,9 @@ struct HfProducerCharmHadronsTrackFemtoDream {
               outputMlD0.at(2) = candidate.mlProbD0()[2]; /// non-prompt score
             }
             if (candidate.mlProbD0bar().size() > 0) {
-              outputMlD0bar.at(0) = candidate.mlProbD0()[0]; /// bkg score
-              outputMlD0bar.at(1) = candidate.mlProbD0()[1]; /// prompt score
-              outputMlD0bar.at(2) = candidate.mlProbD0()[2]; /// non-prompt score
+              outputMlD0bar.at(0) = candidate.mlProbD0bar()[0]; /// bkg score
+              outputMlD0bar.at(1) = candidate.mlProbD0bar()[1]; /// prompt score
+              outputMlD0bar.at(2) = candidate.mlProbD0bar()[2]; /// non-prompt score
             }
 
           } else if (applyMlMode == FillMlFromNewBDT) {
@@ -703,8 +728,8 @@ struct HfProducerCharmHadronsTrackFemtoDream {
             LOGF(fatal, "Please check your Ml configuration!!");
           }
         }
-        fillTable(1, candidate.isSelD0(), outputMlD0.at(0), outputMlD0.at(1), outputMlD0.at(2));
-        fillTable(2, candidate.isSelD0bar(), outputMlD0bar.at(0), outputMlD0bar.at(1), outputMlD0bar.at(2));
+        fillTable(0, candidate.isSelD0(), outputMlD0.at(0), outputMlD0.at(1), outputMlD0.at(2));
+        fillTable(1, candidate.isSelD0bar(), outputMlD0bar.at(0), outputMlD0bar.at(1), outputMlD0bar.at(2));
 
       } else if constexpr (Channel == DecayChannel::DstarToD0Pi) {
         if constexpr (UseCharmMl) {