Skip to content

Commit

Permalink
Please consider the following formatting changes
Browse files Browse the repository at this point in the history
  • Loading branch information
@alibuild
alibuild committed Nov 20, 2024
1 parent 8f6bdfa commit 374d5cf
Show file tree
Hide file tree
Showing 2 changed files with 40 additions and 49 deletions.
80 changes: 36 additions & 44 deletions PWGJE/Core/JetTaggingUtilities.h
View file
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
View file
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

0 comments on commit 374d5cf

Please sign in to comment.