Skip to content

Commit

Permalink
PWGLF: seamless treatment of cascade cov mats in builder / data model
Browse files Browse the repository at this point in the history
  • Loading branch information
@ddobrigk
ddobrigk committed Oct 29, 2024
1 parent abcc778 commit cea709b
Show file tree
Hide file tree
Showing 2 changed files with 53 additions and 13 deletions.
28 changes: 23 additions & 5 deletions PWGLF/DataModel/LFStrangenessTables.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -1419,12 +1428,21 @@ 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_FULL(KFCascCovs, "KFCascCovs", "AOD", "KFCASCCOVS", //!
DECLARE_SOA_TABLE(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", //!
cascdataext::PxLambda, cascdataext::PyLambda, cascdataext::PzLambda, cascdataext::Pt, cascdataext::P, cascdataext::Eta, cascdataext::Phi);
Expand Down
38 changes: 30 additions & 8 deletions PWGLF/TableProducer/Strangeness/cascadebuilder.cxx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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);
}
}

Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit cea709b

Please sign in to comment.