Skip to content

[PWGJE] improve Lund plane MC collision handling and pthard weighting #16367

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zovarga wants to merge 1 commit into
base: master
Choose a base branch
from
+74 −11
Open

[PWGJE] improve Lund plane MC collision handling and pthard weighting #16367

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
85 changes: 74 additions & 11 deletions PWGJE/Tasks/jetLundPlane.cxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,8 +55,14 @@ namespace o2::aod
{
// Parent table: one row per collision stored in the MiniAOD
DECLARE_SOA_COLUMN(MiniCollTag, miniCollTag, uint8_t);
DECLARE_SOA_COLUMN(MiniCollMcCollisionId, miniCollMcCollisionId, int64_t);
DECLARE_SOA_COLUMN(MiniCollWeight, miniCollWeight, float);
DECLARE_SOA_COLUMN(MiniCollPtHard, miniCollPtHard, float);
DECLARE_SOA_TABLE(MiniCollisions, "AOD", "MINICOLL",
MiniCollTag);
MiniCollTag,
MiniCollMcCollisionId,
MiniCollWeight,
MiniCollPtHard);

// MiniJets -> MiniCollisions
DECLARE_SOA_INDEX_COLUMN_CUSTOM(MiniCollision, miniCollision, "MINICOLLS");
Expand Down Expand Up @@ -148,6 +154,12 @@ struct JetMatchInfo {
float otherPt{};
};

struct McEventInfo {
int64_t mcCollisionId{-1};
float weight{1.f};
float ptHard{-1.f};
};

inline float deltaPhi(float phi1, float phi2)
{
return RecoDecay::constrainAngle(phi1 - phi2, -o2::constants::math::PI);
Expand Down Expand Up @@ -405,6 +417,9 @@ struct JetLundPlaneUnfolding {
(aod::jet::eta < jetEtaMax.node()) &&
(aod::jet::r == nround(jetR.node() * 100.f));

// Reconstructed collisions grouped by their associated MC collision.
PresliceUnsorted<soa::Filtered<aod::JetCollisionsMCD>> CollisionsPerMCPCollision = aod::jmccollisionlb::mcCollisionId;

// Type aliases
using RecoJets = soa::Join<aod::ChargedJets, aod::ChargedJetConstituents>;
using DetJetsMatched = soa::Join<aod::ChargedMCDetectorLevelJets,
Expand Down Expand Up @@ -530,7 +545,7 @@ struct JetLundPlaneUnfolding {

int miniCollIdx = -1;
if (writeMiniAOD.value) {
outMiniCollisions(static_cast<uint8_t>(0));
outMiniCollisions(static_cast<uint8_t>(0), int64_t{-1}, 1.f, -1.f);
miniCollIdx = outMiniCollisions.lastIndex();
}
for (auto const& jet : jets) {
Expand Down Expand Up @@ -559,9 +574,10 @@ struct JetLundPlaneUnfolding {

// MC PROCESSING (det + part + response)

void processMC(DetJetsMatched const& detJets,
PartJetsMatched const& partJets,
aod::JetCollisions const&,
void processMC(soa::Filtered<DetJetsMatched> const& detJets,
soa::Filtered<PartJetsMatched> const& partJets,
soa::Filtered<aod::JetCollisionsMCD> const& collisions,
aod::JetMcCollisions const& mcCollisions,
aod::JetTracks const& tracks,
aod::JetParticles const& particles)
{
Expand All @@ -571,6 +587,7 @@ struct JetLundPlaneUnfolding {
std::unordered_map<uint64_t, bool> truthMatchedById;
std::unordered_map<uint64_t, std::vector<SplittingObs>> truthSplittingsById;
std::unordered_map<uint64_t, JetMatchInfo> truthBestDet;
std::unordered_set<uint64_t> acceptedTruthJetKeys;
std::unordered_set<uint64_t> detSplittingsWritten;
std::unordered_set<uint64_t> truthSplittingsWritten;
auto h6 = registry.get<THnSparse>(HIST("hLundResponse6D"));
Expand All @@ -583,6 +600,21 @@ struct JetLundPlaneUnfolding {
std::unordered_map<uint64_t, int> detJetToMiniJetIdx;
std::unordered_map<uint64_t, int> partJetToMiniJetIdx;

std::unordered_map<int64_t, McEventInfo> mcEventInfoById;
for (auto const& mcCollision : mcCollisions) {
const int64_t mcCollisionId = mcCollision.globalIndex();
mcEventInfoById.emplace(mcCollisionId,
McEventInfo{mcCollisionId, mcCollision.weight(), mcCollision.ptHard()});
}

std::unordered_map<int64_t, McEventInfo> mcEventInfoByDetCollisionId;
for (auto const& collision : collisions) {
auto mcInfoIt = mcEventInfoById.find(collision.mcCollisionId());
if (mcInfoIt != mcEventInfoById.end()) {
mcEventInfoByDetCollisionId.emplace(collision.globalIndex(), mcInfoIt->second);
}
}

// --- Truth pass ---
// Fill inclusive truth histograms, cache truth splittings, write all accepted
// truth jets to MiniJets, and determine the best detector candidate for each truth jet.
Expand All @@ -591,10 +623,22 @@ struct JetLundPlaneUnfolding {
continue;
}

auto collisionsPerMCPJet = collisions.sliceBy(CollisionsPerMCPCollision, partJet.mcCollisionId());
if (collisionsPerMCPJet.size() < 1) {
continue;
}

auto partMcInfoIt = mcEventInfoById.find(partJet.mcCollisionId());
if (partMcInfoIt == mcEventInfoById.end()) {
continue;
}
const auto partMcInfo = partMcInfoIt->second;

registry.fill(HIST("hJetPtTruthAll"), partJet.pt());
fillJetCountSummary(4);

const uint64_t truthJetKey = partJet.globalIndex();
acceptedTruthJetKeys.insert(truthJetKey);
auto spl = getPrimarySplittings(partJet, particles);
truthSplittingsById[truthJetKey] = spl;

Expand All @@ -607,7 +651,7 @@ struct JetLundPlaneUnfolding {
int partMiniCollIdx = -1;
auto collIt = partMiniCollByKey.find(partCollKey);
if (collIt == partMiniCollByKey.end()) {
outMiniCollisions(static_cast<uint8_t>(0));
outMiniCollisions(static_cast<uint8_t>(0), partMcInfo.mcCollisionId, partMcInfo.weight, partMcInfo.ptHard);
partMiniCollIdx = outMiniCollisions.lastIndex();
partMiniCollByKey.emplace(partCollKey, partMiniCollIdx);
} else {
Expand All @@ -629,7 +673,7 @@ struct JetLundPlaneUnfolding {

JetMatchInfo bestDet{};
bool foundDet = false;
for (auto const& candDetJet : partJet.template matchedJetGeo_as<DetJetsMatched>()) {
for (auto const& candDetJet : partJet.template matchedJetGeo_as<soa::Filtered<DetJetsMatched>>()) {
if (!passJetFiducial(candDetJet, rWanted)) {
continue;
}
Expand Down Expand Up @@ -659,13 +703,19 @@ struct JetLundPlaneUnfolding {
}

// --- Detector loop ---
// Write all accepted detector jets to MiniJets. A final matched pair is accepted only
// Write accepted detector jets to MiniJets. A final matched pair is accepted only
// if the detector jet and truth jet are mutual best matches under the same cuts.
for (auto const& detJet : detJets) {
if (!passJetFiducial(detJet, rWanted)) {
continue;
}

auto detMcInfoIt = mcEventInfoByDetCollisionId.find(detJet.collisionId());
if (detMcInfoIt == mcEventInfoByDetCollisionId.end()) {
continue;
}
const auto detMcInfo = detMcInfoIt->second;

registry.fill(HIST("hJetPtRecoAll"), detJet.pt());
fillJetCountSummary(1);

Expand All @@ -680,7 +730,7 @@ struct JetLundPlaneUnfolding {
int detMiniCollIdx = -1;
auto collIt = detMiniCollByKey.find(detCollKey);
if (collIt == detMiniCollByKey.end()) {
outMiniCollisions(static_cast<uint8_t>(0));
outMiniCollisions(static_cast<uint8_t>(0), detMcInfo.mcCollisionId, detMcInfo.weight, detMcInfo.ptHard);
detMiniCollIdx = outMiniCollisions.lastIndex();
detMiniCollByKey.emplace(detCollKey, detMiniCollIdx);
} else {
Expand All @@ -706,7 +756,12 @@ struct JetLundPlaneUnfolding {
bool foundMatch = false;
std::vector<SplittingObs> bestPartSpl;

for (auto const& candPartJet : detJet.template matchedJetGeo_as<PartJetsMatched>()) {
for (auto const& candPartJet : detJet.template matchedJetGeo_as<soa::Filtered<PartJetsMatched>>()) {
const uint64_t candTruthKey = candPartJet.globalIndex();
if (acceptedTruthJetKeys.find(candTruthKey) == acceptedTruthJetKeys.end()) {
continue;
}

if (!passJetFiducial(candPartJet, rWanted)) {
continue;
}
Expand All @@ -723,7 +778,6 @@ struct JetLundPlaneUnfolding {
}

if (!foundMatch || dR < bestTruth.dR) {
const uint64_t candTruthKey = candPartJet.globalIndex();
bestTruth.otherKey = candTruthKey;
bestTruth.dR = dR;
bestTruth.relPt = rel;
Expand Down Expand Up @@ -813,6 +867,15 @@ struct JetLundPlaneUnfolding {
continue;
}

auto collisionsPerMCPJet = collisions.sliceBy(CollisionsPerMCPCollision, partJet.mcCollisionId());
if (collisionsPerMCPJet.size() < 1) {
continue;
}

if (mcEventInfoById.find(partJet.mcCollisionId()) == mcEventInfoById.end()) {
continue;
}

const uint64_t truthJetKey = partJet.globalIndex();
if (!truthMatchedById[truthJetKey]) {
registry.fill(HIST("hJetPtTruthMiss"), partJet.pt());
Expand Down
Loading