PWGLF: seamless treatment of cascade cov mats in builder / data model (#8225)

Co-authored-by: ALICE Builder <alibuild@users.noreply.github.com>

Author: ddobrigk

PWGLF/DataModel/LFStrangenessTables.h

@@ -1043,13 +1043,22 @@ DECLARE_SOA_COLUMN(BachX, bachX, float); //! bachelor track X at min
 //______________________________________________________
 // REGULAR COLUMNS FOR CASCCOVS
 // Saved from finding: covariance matrix of parent track (on request)
-DECLARE_SOA_COLUMN(PositionCovMat, positionCovMat, float[6]);                //! covariance matrix elements
-DECLARE_SOA_COLUMN(MomentumCovMat, momentumCovMat, float[6]);                //! covariance matrix elements
+DECLARE_SOA_DYNAMIC_COLUMN(PositionCovMat, positionCovMat, //! for transparent handling
+                           [](const float covMat[21]) -> std::vector<float> {
+                            std::vector<float> posCovMat { covMat[0], covMat[1], covMat[2], covMat[3], covMat[4], covMat[5] };
+                            return posCovMat; });
+DECLARE_SOA_DYNAMIC_COLUMN(MomentumCovMat, momentumCovMat, //! for transparent handling
+                           [](const float covMat[21]) -> std::vector<float> {
+                            std::vector<float> momCovMat { covMat[9], covMat[13], covMat[14], covMat[18], covMat[19], covMat[20] };
+                            return momCovMat; });
 DECLARE_SOA_COLUMN(KFTrackCovMat, kfTrackCovMat, float[21]);                 //! covariance matrix elements for KF method (Cascade)
 DECLARE_SOA_COLUMN(KFTrackCovMatV0, kfTrackCovMatV0, float[21]);             //! covariance matrix elements for KF method (V0)
 DECLARE_SOA_COLUMN(KFTrackCovMatV0DauPos, kfTrackCovMatV0DauPos, float[21]); //! covariance matrix elements for KF method (V0 pos daughter)
 DECLARE_SOA_COLUMN(KFTrackCovMatV0DauNeg, kfTrackCovMatV0DauNeg, float[21]); //! covariance matrix elements for KF method (V0 neg daughter)
 
+// for CascCovs / TraCascCovs, meant to provide consistent interface everywhere
+DECLARE_SOA_COLUMN(CovMat, covMat, float[21]); //! covariance matrix elements
+
 //______________________________________________________
 // REGULAR COLUMNS FOR CASCBBS
 // General cascade properties: position, momentum
@@ -1419,11 +1428,20 @@ DECLARE_SOA_TABLE(CascMCMothers, "AOD", "CASCMCMOTHER", //! optional table for M
 DECLARE_SOA_TABLE(CascBBs, "AOD", "CASCBB", //! bachelor-baryon correlation variables
                   cascdata::BachBaryonCosPA, cascdata::BachBaryonDCAxyToPV)
 
-DECLARE_SOA_TABLE_FULL(CascCovs, "CascCovs", "AOD", "CASCCOVS", //!
-                       cascdata::PositionCovMat, cascdata::MomentumCovMat);
+DECLARE_SOA_TABLE(CascCovs, "AOD", "CASCCOVS", //!
+                  cascdata::CovMat,
+                  cascdata::PositionCovMat<cascdata::CovMat>,
+                  cascdata::MomentumCovMat<cascdata::CovMat>,
+                  o2::soa::Marker<1>);
+
+DECLARE_SOA_TABLE(KFCascCovs, "AOD", "KFCASCCOVS", //!
+                  cascdata::KFTrackCovMat, cascdata::KFTrackCovMatV0, cascdata::KFTrackCovMatV0DauPos, cascdata::KFTrackCovMatV0DauNeg);
 
-DECLARE_SOA_TABLE_FULL(KFCascCovs, "KFCascCovs", "AOD", "KFCASCCOVS", //!
-                       cascdata::KFTrackCovMat, cascdata::KFTrackCovMatV0, cascdata::KFTrackCovMatV0DauPos, cascdata::KFTrackCovMatV0DauNeg);
+DECLARE_SOA_TABLE(TraCascCovs, "AOD", "TRACASCCOVS", //!
+                  cascdata::CovMat,
+                  cascdata::PositionCovMat<cascdata::CovMat>,
+                  cascdata::MomentumCovMat<cascdata::CovMat>,
+                  o2::soa::Marker<2>);
 
 // extended table with expression columns that can be used as arguments of dynamic columns
 DECLARE_SOA_EXTENDED_TABLE_USER(CascCores, StoredCascCores, "CascDATAEXT", //!
@@ -1499,6 +1517,8 @@ using CascadesLinked = soa::Join<Cascades, CascDataLink>;
 using CascadeLinked = CascadesLinked::iterator;
 using KFCascadesLinked = soa::Join<Cascades, KFCascDataLink>;
 using KFCascadeLinked = KFCascadesLinked::iterator;
+using TraCascadesLinked = soa::Join<Cascades, TraCascDataLink>;
+using TraCascadeLinked = TraCascadesLinked::iterator;
 
 namespace cascdata
 {

PWGLF/TableProducer/Strangeness/cascadebuilder.cxx

@@ -124,6 +124,7 @@ struct cascadeBuilder {
   Produces<aod::CascTrackXs> cascTrackXs; // if desired for replaying of position information
   Produces<aod::CascBBs> cascbb;          // if enabled
   Produces<aod::CascCovs> casccovs;       // if requested by someone
+  Produces<aod::TraCascCovs> tracasccovs; // if requested by someone
   Produces<aod::KFCascCovs> kfcasccovs;   // if requested by someone
 
   // produces calls for machine-learning selections
@@ -136,7 +137,7 @@ struct cascadeBuilder {
   Configurable<bool> d_UseAutodetectMode{"d_UseAutodetectMode", false, "Autodetect requested topo sels"};
 
   // Configurables related to table creation
-  Configurable<int> createCascCovMats{"createCascCovMats", -1, {"Produces V0 cov matrices. -1: auto, 0: don't, 1: yes. Default: auto (-1)"}};
+  Configurable<int> createCascCovMats{"createCascCovMats", -1, {"Produces casc cov matrices. -1: auto, 0: don't, 1: yes. Default: auto (-1)"}};
   Configurable<int> createCascTrackXs{"createCascTrackXs", -1, {"Produces track X at minima table. -1: auto, 0: don't, 1: yes. Default: auto (-1)"}};
 
   // Topological selection criteria
@@ -1601,15 +1602,19 @@ struct cascadeBuilder {
       // store momentum covariance matrix
       std::array<float, 21> covTv0 = {0.};
       std::array<float, 21> covTbachelor = {0.};
+      float covCascade[21];
       // std::array<float, 6> momentumCovariance;
-      float momentumCovariance[6];
       lV0Track.getCovXYZPxPyPzGlo(covTv0);
       lBachelorTrack.getCovXYZPxPyPzGlo(covTbachelor);
       constexpr int MomInd[6] = {9, 13, 14, 18, 19, 20}; // cov matrix elements for momentum component
+      for (int i = 0; i < 21; i++) {
+        covCascade[i] = 0.0f;
+      }
       for (int i = 0; i < 6; i++) {
-        momentumCovariance[i] = covTv0[MomInd[i]] + covTbachelor[MomInd[i]];
+        covCascade[i] = positionCovariance[i];
+        covCascade[MomInd[i]] = covTv0[MomInd[i]] + covTbachelor[MomInd[i]];
       }
-      casccovs(positionCovariance, momentumCovariance);
+      casccovs(covCascade);
     }
   }
 
@@ -1778,15 +1783,19 @@ struct cascadeBuilder {
         // store momentum covariance matrix
         std::array<float, 21> covTv0 = {0.};
         std::array<float, 21> covTbachelor = {0.};
+        float covCascade[21];
         // std::array<float, 6> momentumCovariance;
-        float momentumCovariance[6];
         lV0Track.getCovXYZPxPyPzGlo(covTv0);
         lBachelorTrack.getCovXYZPxPyPzGlo(covTbachelor);
         constexpr int MomInd[6] = {9, 13, 14, 18, 19, 20}; // cov matrix elements for momentum component
+        for (int i = 0; i < 21; i++) {
+          covCascade[i] = 0.0f;
+        }
         for (int i = 0; i < 6; i++) {
-          momentumCovariance[i] = covTv0[MomInd[i]] + covTbachelor[MomInd[i]];
+          covCascade[i] = positionCovariance[i];
+          covCascade[MomInd[i]] = covTv0[MomInd[i]] + covTbachelor[MomInd[i]];
         }
-        casccovs(positionCovariance, momentumCovariance);
+        casccovs(covCascade);
       }
 
       float lPt = 0.0f;
@@ -1830,7 +1839,7 @@ struct cascadeBuilder {
           continue; // safety (should be fine but depends on future stratrack dev)
         // Track casting to <TTracksTo>
         auto cascadeTrack = trackedCascade.template track_as<TTrackTo>();
-        auto cascadeTrackPar = getTrackPar(cascadeTrack);
+        auto cascadeTrackPar = getTrackParCov(cascadeTrack);
         auto const& collision = cascade.collision();
         gpu::gpustd::array<float, 2> dcaInfo;
         lCascadeTrack.setPID(o2::track::PID::XiMinus); // FIXME: not OK for omegas
@@ -1910,6 +1919,19 @@ struct cascadeBuilder {
                         cascadecandidate.v0dcapostopv, cascadecandidate.v0dcanegtopv,
                         cascadecandidate.bachDCAxy, cascadecandidate.cascDCAxy, cascadecandidate.cascDCAz,          // <--- stratrack (cascDCAxy/z)
                         trackedCascade.matchingChi2(), trackedCascade.topologyChi2(), trackedCascade.itsClsSize()); // <--- stratrack fit info
+
+        if (createCascCovMats) {
+          // create tracked cascade covariance in exactly the same way as non-tracked
+          // ensures getter consistency and full compatibility in template functions
+          // (easy switching between tracked and non-tracked)
+          std::array<float, 21> traCovMat = {0.};
+          cascadeTrackPar.getCovXYZPxPyPzGlo(traCovMat);
+          float traCovMatArray[21];
+          for (int ii = 0; ii < 21; ii++) {
+            traCovMatArray[ii] = traCovMat[ii];
+          }
+          tracasccovs(traCovMatArray);
+        }
       }
     }
     // En masse filling at end of process call