Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
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
80 changes: 36 additions & 44 deletions PWGJE/Core/JetTaggingUtilities.h
Original file line number Diff line number Diff line change
Expand Up @@ -643,53 +643,49 @@ int vertexClustering(AnyCollision const& collision, AnalysisJet const& jet, AnyT

std::vector<int> tempTrkVtxIndex;

int i=0;
int i = 0;
for (const auto& constituent : tracks) {
if (!constituent.has_mcParticle() || !constituent.template mcParticle_as<AnyParticles>().isPhysicalPrimary() || constituent.pt() < trackPtMin)
tempTrkVtxIndex.push_back(-1);
else
tempTrkVtxIndex.push_back(i++);
}
tempTrkVtxIndex.push_back(i); // temporary index for PV
if (n_trks < 1) { // the process should be done for n_trks == 1 as well
if (n_trks < 1) { // the process should be done for n_trks == 1 as well
trkLabels["trkVtxIndex"] = tempTrkVtxIndex;
return n_trks;
}

int n_pos = n_trks + 1;
std::vector<float> dists(n_pos * (n_pos - 1) / 2);
auto trk_pair_idx = [n_pos](int ti, int tj) {
if (ti==tj || ti>=n_pos || tj>=n_pos || ti<0 || tj<0) {
LOGF(info,"Track pair index out of range");
if (ti == tj || ti >= n_pos || tj >= n_pos || ti < 0 || tj < 0) {
LOGF(info, "Track pair index out of range");
return -1;
}
else
} else
return (ti < tj) ? (ti * n_pos - (ti * (ti + 1)) / 2 + tj - ti - 1) : (tj * n_pos - (tj * (tj + 1)) / 2 + ti - tj - 1);
}; // index n_trks is for PV

for (int ti=0; ti<n_pos-1; ti++)
for (int tj=ti+1; tj<n_pos; tj++) {
for (int ti = 0; ti < n_pos - 1; ti++)
for (int tj = ti + 1; tj < n_pos; tj++) {
std::array<float, 3> posi, posj;

if (tj < n_trks) {
if (tracks[tj].has_mcParticle()) {
const auto& pj = tracks[tj].template mcParticle_as<AnyParticles>().template mcParticle_as<AnyOriginalParticles>();
posj = std::array<float, 3>{pj.vx(), pj.vy(), pj.vz()};
}
else {
} else {
dists[trk_pair_idx(ti, tj)] = std::numeric_limits<float>::max();
continue;
}
}
else {
} else {
posj = std::array<float, 3>{collision.posX(), collision.posY(), collision.posZ()};
}

if (tracks[ti].has_mcParticle()) {
const auto& pi = tracks[ti].template mcParticle_as<AnyParticles>().template mcParticle_as<AnyOriginalParticles>();
posi = std::array<float, 3>{pi.vx(), pi.vy(), pi.vz()};
}
else {
} else {
dists[trk_pair_idx(ti, tj)] = std::numeric_limits<float>::max();
continue;
}
Expand All @@ -698,13 +694,13 @@ int vertexClustering(AnyCollision const& collision, AnalysisJet const& jet, AnyT
}

int clusteri = -1, clusterj = -1;
float min_min_dist = -1.f; // If there is an not-merge-able min_dist pair, check the 2nd-min_dist pair.
float min_min_dist = -1.f; // If there is an not-merge-able min_dist pair, check the 2nd-min_dist pair.
while (true) {

float min_dist = -1.f; // Get min_dist pair
for (int ti=0; ti<n_pos-1; ti++)
for (int tj=ti+1; tj<n_pos; tj++)
if (tempTrkVtxIndex[ti] != tempTrkVtxIndex[tj] && tempTrkVtxIndex[ti]>=0 && tempTrkVtxIndex[tj]>=0) {
for (int ti = 0; ti < n_pos - 1; ti++)
for (int tj = ti + 1; tj < n_pos; tj++)
if (tempTrkVtxIndex[ti] != tempTrkVtxIndex[tj] && tempTrkVtxIndex[ti] >= 0 && tempTrkVtxIndex[tj] >= 0) {
float dist = dists[trk_pair_idx(ti, tj)];
if ((dist < min_dist || min_dist < 0.f) && dist > min_min_dist) {
min_dist = dist;
Expand All @@ -714,22 +710,22 @@ int vertexClustering(AnyCollision const& collision, AnalysisJet const& jet, AnyT
}
if (clusteri < 0 || clusterj < 0)
break;

bool mrg = true; // Merge-ability check
for (int ti=0; ti<n_pos && mrg; ti++)
if (tempTrkVtxIndex[ti] == tempTrkVtxIndex[clusteri] && tempTrkVtxIndex[ti]>=0)
for (int tj=0; tj<n_pos && mrg; tj++)
if (tj != ti && tempTrkVtxIndex[tj] == tempTrkVtxIndex[clusterj] && tempTrkVtxIndex[tj]>=0)
for (int ti = 0; ti < n_pos && mrg; ti++)
if (tempTrkVtxIndex[ti] == tempTrkVtxIndex[clusteri] && tempTrkVtxIndex[ti] >= 0)
for (int tj = 0; tj < n_pos && mrg; tj++)
if (tj != ti && tempTrkVtxIndex[tj] == tempTrkVtxIndex[clusterj] && tempTrkVtxIndex[tj] >= 0)
if (dists[trk_pair_idx(ti, tj)] > vtxResParam) { // If there is more distant pair compared to vtx_res between two clusters, they cannot be merged.
mrg = false;
min_min_dist = min_dist;
}
if (min_dist > vtxResParam || min_dist < 0.f)
break;

if (mrg) { // Merge two clusters
int old_index = tempTrkVtxIndex[clusterj];
for (int t=0; t<n_pos; t++)
for (int t = 0; t < n_pos; t++)
if (tempTrkVtxIndex[t] == old_index)
tempTrkVtxIndex[t] = tempTrkVtxIndex[clusteri];
}
Expand All @@ -739,61 +735,57 @@ int vertexClustering(AnyCollision const& collision, AnalysisJet const& jet, AnyT

// Sort the indices from PV (as 0) to the most distant SV (as 1~).
int idxPV = tempTrkVtxIndex[n_trks];
for (int t=0; t<n_trks; t++)
for (int t = 0; t < n_trks; t++)
if (tempTrkVtxIndex[t] == idxPV) {
tempTrkVtxIndex[t] = -2;
n_vertices = 1; // There is a track originating from PV
}

std::unordered_map<int, float> avgDistances;
std::unordered_map<int, int> count;
for (int t=0; t<n_trks; t++) {
for (int t = 0; t < n_trks; t++) {
if (tempTrkVtxIndex[t] >= 0) {
avgDistances[tempTrkVtxIndex[t]] += dists[trk_pair_idx(t, n_trks)];
count[tempTrkVtxIndex[t]]++;
}
}

trkLabels["trkVtxIndex"] = std::vector<int>(n_trks, -1);
if (count.size() != 0) { // If there is any SV cluster not only PV cluster
for (auto& [idx, avgDistance] : avgDistances)
avgDistance /= count[idx];

n_vertices += avgDistances.size();

std::vector<std::pair<int, float>> sortedIndices(avgDistances.begin(), avgDistances.end());
std::sort(sortedIndices.begin(), sortedIndices.end(), [](const auto& a, const auto& b) { return a.second < b.second; });
int rank = 1;
for (const auto& [idx, avgDistance] : sortedIndices) {
bool found = false;
for (int t=0; t<n_trks; t++)
for (int t = 0; t < n_trks; t++)
if (tempTrkVtxIndex[t] == idx) {
trkLabels["trkVtxIndex"][t] = rank;
found = true;
}
rank += found;
}
}
for (int t=0; t<n_trks; t++)

for (int t = 0; t < n_trks; t++)
if (tempTrkVtxIndex[t] == -2)
trkLabels["trkVtxIndex"][t] = 0;

// trkOrigin

int trkIdx = 0;
for (auto &constituent : jet.template tracks_as<AnyTracks>())
{
if (!constituent.has_mcParticle() || !constituent.template mcParticle_as<AnyParticles>().isPhysicalPrimary() || constituent.pt() < trackPtMin)
{
for (auto& constituent : jet.template tracks_as<AnyTracks>()) {
if (!constituent.has_mcParticle() || !constituent.template mcParticle_as<AnyParticles>().isPhysicalPrimary() || constituent.pt() < trackPtMin) {
trkLabels["trkOrigin"].push_back(0);
}
else
{
const auto &particle = constituent.template mcParticle_as<AnyParticles>();
} else {
const auto& particle = constituent.template mcParticle_as<AnyParticles>();
int orig = RecoDecay::getParticleOrigin(particles, particle, searchUpToQuark);
trkLabels["trkOrigin"].push_back((orig > 0) ? orig :
(trkLabels["trkVtxIndex"][trkIdx] == 0) ? 3 : 4);
trkLabels["trkOrigin"].push_back((orig > 0) ? orig : (trkLabels["trkVtxIndex"][trkIdx] == 0) ? 3
: 4);
}

trkIdx++;
Expand Down
9 changes: 4 additions & 5 deletions PWGJE/Tasks/bjetTreeCreator.cxx
Original file line number Diff line number Diff line change
Expand Up @@ -286,8 +286,7 @@ struct BJetTreeCreator {
registry.add("h2_Response_DetjetpT_PartjetpT_lfjet", "Response matrix lf-jet;#it{p}_{T,jet}^{det} (GeV/#it{c});#it{p}_{T,jet}^{part} (GeV/#it{c})", {HistType::kTH2F, {{200, 0., 200.}, {200, 0., 200.}}});
}

if (doprocessMCJetsForGNN)
{
if (doprocessMCJetsForGNN) {
//+jet
registry.add("h_jet_pt", "jet_pt;#it{p}_{T}^{ch jet} (GeV/#it{c});Entries", {HistType::kTH1F, {{100, 0., 200.}}});
registry.add("h_jet_eta", "jet_eta;#it{#eta}_{ch jet};Entries", {HistType::kTH1F, {{100, -2., 2.}}});
Expand Down Expand Up @@ -464,7 +463,7 @@ struct BJetTreeCreator {
}

template <typename AnyCollision, typename AnalysisJet, typename AnyTracks, typename AnyOriginalTracks>
void analyzeJetTrackInfoForGNN(AnyCollision const& /*collision*/, AnalysisJet const& analysisJet, AnyTracks const& /*allTracks*/, AnyOriginalTracks const&, std::vector<int>& trackIndices, int jetFlavor = 0, double eventweight = 1.0, TrackLabelMap *trkLabels = nullptr)
void analyzeJetTrackInfoForGNN(AnyCollision const& /*collision*/, AnalysisJet const& analysisJet, AnyTracks const& /*allTracks*/, AnyOriginalTracks const&, std::vector<int>& trackIndices, int jetFlavor = 0, double eventweight = 1.0, TrackLabelMap* trkLabels = nullptr)
{
int trkIdx = -1;
for (auto& constituent : analysisJet.template tracks_as<AnyTracks>()) {
Expand Down Expand Up @@ -523,7 +522,7 @@ struct BJetTreeCreator {
registry.fill(HIST("h_trk_origin"), trkOrigin);

if (produceTree) {
bjetTracksExtraTable(/*bjetParamsTable.lastIndex() + 1, */origConstit.phi(), constituent.sign(), origConstit.itsChi2NCl(), origConstit.tpcChi2NCl(), origConstit.itsNCls(), origConstit.tpcNClsFound(), origConstit.tpcNClsCrossedRows(), trkOrigin, trkVtxIndex); //+
bjetTracksExtraTable(/*bjetParamsTable.lastIndex() + 1, */ origConstit.phi(), constituent.sign(), origConstit.itsChi2NCl(), origConstit.tpcChi2NCl(), origConstit.itsNCls(), origConstit.tpcNClsFound(), origConstit.tpcNClsCrossedRows(), trkOrigin, trkVtxIndex); //+
}
}

Expand Down Expand Up @@ -693,7 +692,7 @@ struct BJetTreeCreator {
using MCDJetTableNoSV = soa::Filtered<soa::Join<aod::ChargedMCDetectorLevelJets, aod::ChargedMCDetectorLevelJetConstituents, aod::ChargedMCDetectorLevelJetsMatchedToChargedMCParticleLevelJets, aod::ChargedMCDetectorLevelJetEventWeights>>;
using JetParticleswID = soa::Join<aod::JetParticles, aod::JMcParticlePIs>;

void processMCJetsForGNN(FilteredCollisionMCD::iterator const &collision, aod::JMcCollisions const &, MCDJetTableNoSV const &MCDjets, MCPJetTable const &MCPjets, JetTracksMCDwID const &allTracks, JetParticleswID const &MCParticles, OriginalTracks const& origTracks, aod::McParticles const& origParticles)
void processMCJetsForGNN(FilteredCollisionMCD::iterator const& collision, aod::JMcCollisions const&, MCDJetTableNoSV const& MCDjets, MCPJetTable const& MCPjets, JetTracksMCDwID const& allTracks, JetParticleswID const& MCParticles, OriginalTracks const& origTracks, aod::McParticles const& origParticles)
{
if (!jetderiveddatautilities::selectCollision(collision, eventSelection) || (static_cast<double>(std::rand()) / RAND_MAX < eventReductionFactor)) {
return;
Expand Down
Loading