From a94f86fec4c71489ecdd6cc43b9be1e8a78055f3 Mon Sep 17 00:00:00 2001
From: Sylvester Joosten <sjoosten@anl.gov>
Date: Fri, 20 Nov 2020 17:26:57 +0000
Subject: [PATCH] Finished first iteration of analysis script
---
dvmp/analysis/mt.h | 11 ++++
dvmp/analysis/plot.h | 40 ++++++++++++
dvmp/analysis/util.h | 127 ++++++++++++++++++++++++++++++++++++++
dvmp/analysis/vm_mass.cxx | 124 +++++++++++++++++++++++++++++--------
dvmp/dvmp.sh | 23 ++++---
5 files changed, 291 insertions(+), 34 deletions(-)
create mode 100644 dvmp/analysis/mt.h
create mode 100644 dvmp/analysis/plot.h
create mode 100644 dvmp/analysis/util.h
diff --git a/dvmp/analysis/mt.h b/dvmp/analysis/mt.h
new file mode 100644
index 00000000..198050c6
--- /dev/null
+++ b/dvmp/analysis/mt.h
@@ -0,0 +1,11 @@
+#ifndef MT_H
+#define MT_H
+
+// Defines the number of threads to run within the ROOT analysis scripts.
+// TODO: make this a file configured by the CI scripts so we can specify
+// the number of threads (and the number of processes) at a global
+// level
+
+constexpr const int kNumThreads = 8;
+
+#endif
diff --git a/dvmp/analysis/plot.h b/dvmp/analysis/plot.h
new file mode 100644
index 00000000..69ebba34
--- /dev/null
+++ b/dvmp/analysis/plot.h
@@ -0,0 +1,40 @@
+#ifndef PLOT_H
+#define PLOT_H
+
+#include <TColor.h>
+#include <fmt/core.h>
+#include <vector>
+
+namespace plot {
+
+const int kMpBlue = TColor::GetColor(0x1f, 0x77, 0xb4);
+const int kMpOrange = TColor::GetColor(0xff, 0x7f, 0x0e);
+const int kMpGreen = TColor::GetColor(0x2c, 0xa0, 0x2c);
+const int kMpRed = TColor::GetColor(0xd6, 0x27, 0x28);
+const int kMpPurple = TColor::GetColor(0x94, 0x67, 0xbd);
+const int kMpBrown = TColor::GetColor(0x8c, 0x56, 0x4b);
+const int kMpPink = TColor::GetColor(0xe3, 0x77, 0xc2);
+const int kMpGrey = TColor::GetColor(0x7f, 0x7f, 0x7f);
+const int kMpMoss = TColor::GetColor(0xbc, 0xbd, 0x22);
+const int kMpCyan = TColor::GetColor(0x17, 0xbe, 0xcf);
+
+const std::vector<int> kPalette = {kMpBlue, kMpOrange, kMpGreen, kMpRed,
+ kMpPurple, kMpBrown, kMpPink, kMpGrey,
+ kMpMoss, kMpCyan};
+
+void draw_label(int ebeam, int pbeam, const std::string_view detector,
+ std::string_view vm, std::string_view what) {
+ auto t = new TPaveText(.15, 0.800, .7, .925, "NB NDC");
+ t->SetFillColorAlpha(kWhite, 0.4);
+ t->SetTextFont(43);
+ t->SetTextSize(25);
+ t->AddText(fmt::format("#bf{{{} }}SIMULATION", detector).c_str());
+ t->AddText(fmt::format("{} for {} DVMP.", what, vm).c_str());
+ t->AddText(fmt::format("{} GeV on {} GeV", ebeam, pbeam, what).c_str());
+ t->SetTextAlign(12);
+ t->Draw();
+}
+
+} // namespace plot
+
+#endif
diff --git a/dvmp/analysis/util.h b/dvmp/analysis/util.h
new file mode 100644
index 00000000..874b7024
--- /dev/null
+++ b/dvmp/analysis/util.h
@@ -0,0 +1,127 @@
+#ifndef UTIL_H
+#define UTIL_H
+
+#include <algorithm>
+#include <cmath>
+#include <exception>
+#include <fmt/core.h>
+#include <limits>
+#include <string>
+#include <vector>
+
+#include <Math/Vector4D.h>
+
+#include "dd4pod/Geant4ParticleCollection.h"
+#include "eicd/TrackParametersCollection.h"
+
+namespace util {
+
+// Exception definition for unknown particle errors
+// FIXME: A utility exception base class should be included in the analysis
+// utility library, so we can skip most of this boilerplate
+class unknown_particle_error : public std::exception {
+public:
+ unknown_particle_error(std::string_view particle) : m_particle{particle} {}
+ virtual const char* what() const throw() {
+ return fmt::format("Unknown particle type: {}", m_particle).c_str();
+ }
+ virtual const char* type() const throw() { return "unknown_particle_error"; }
+
+private:
+ const std::string m_particle;
+};
+
+// Simple function to return the appropriate PDG mass for the particles
+// we care about for this process.
+// FIXME: consider something more robust (maybe based on hepPDT) to the
+// analysis utility library
+inline double get_pdg_mass(std::string_view part) {
+ if (part == "electron") {
+ return 0.0005109989461;
+ } else if (part == "muon") {
+ return .1056583745;
+ } else if (part == "jpsi") {
+ return 3.0969;
+ } else if (part == "upsilon") {
+ return 9.49630;
+ } else {
+ throw unknown_particle_error{part};
+ }
+}
+
+// Get a vector of 4-momenta from raw tracking info, using an externally
+// provided particle mass assumption.
+// FIXME: should be part of utility library
+inline auto
+momenta_from_tracking(const std::vector<eic::TrackParametersData>& tracks,
+ const double mass) {
+ std::vector<ROOT::Math::PxPyPzMVector> momenta{tracks.size()};
+ // transform our raw tracker info into proper 4-momenta
+ std::transform(tracks.begin(), tracks.end(), momenta.begin(),
+ [mass](const auto& track) {
+ // make sure we don't divide by zero
+ if (fabs(track.qOverP) < 1e-9) {
+ return ROOT::Math::PtEtaPhiMVector{};
+ }
+ const double pt = 1. / track.qOverP * sin(track.theta);
+ const double eta = -log(tan(track.theta / 2));
+ const double phi = track.phi;
+ return ROOT::Math::PtEtaPhiMVector{pt, eta, phi, mass};
+ });
+ return momenta;
+}
+
+// Get a vector of 4-momenta from the simulation data.
+// FIXME: should be part of utility library
+// TODO: Add PID selector (maybe using ranges?)
+inline auto
+momenta_from_simulation(const std::vector<dd4pod::Geant4ParticleData>& parts) {
+ std::vector<ROOT::Math::PxPyPzMVector> momenta{parts.size()};
+ // transform our simulation particle data into 4-momenta
+ std::transform(parts.begin(), parts.end(), momenta.begin(),
+ [](const auto& part) {
+ return ROOT::Math::PxPyPzMVector{part.psx, part.psy,
+ part.psz, part.mass};
+ });
+ return momenta;
+}
+
+// Find the decay pair candidates from a vector of particles (parts),
+// with invariant mass closest to a desired value (pdg_mass)
+// FIXME: not sure if this belongs here, or in the utility library. Probably the
+// utility library
+inline std::pair<ROOT::Math::PxPyPzMVector, ROOT::Math::PxPyPzMVector>
+find_decay_pair(const std::vector<ROOT::Math::PxPyPzMVector>& parts,
+ const double pdg_mass) {
+ int first = -1;
+ int second = -1;
+ double best_mass = -1;
+
+ // go through all particle combinatorics, calculate the invariant mass
+ // for each combination, and remember which combination is the closest
+ // to the desired pdg_mass
+ for (int i = 0; i < parts.size(); ++i) {
+ for (int j = i + 1; j < parts.size(); ++j) {
+ const double new_mass{(parts[i] + parts[j]).mass()};
+ if (fabs(new_mass - pdg_mass) < fabs(best_mass - pdg_mass)) {
+ first = i;
+ second = j;
+ best_mass = new_mass;
+ }
+ }
+ }
+ if (first < 0) {
+ return {{}, {}};
+ }
+ return {parts[first], parts[second]};
+}
+// Calculate the invariant mass of a given pair of particles
+inline double
+get_im(const std::pair<ROOT::Math::PxPyPzMVector, ROOT::Math::PxPyPzMVector>&
+ particle_pair) {
+ return (particle_pair.first + particle_pair.second).mass();
+}
+
+} // namespace util
+
+#endif
diff --git a/dvmp/analysis/vm_mass.cxx b/dvmp/analysis/vm_mass.cxx
index 320ecd6c..19c56db8 100644
--- a/dvmp/analysis/vm_mass.cxx
+++ b/dvmp/analysis/vm_mass.cxx
@@ -1,34 +1,108 @@
+#include "mt.h"
+#include "plot.h"
+#include "util.h"
+
#include <ROOT/RDataFrame.hxx>
#include <cmath>
+#include <fmt/color.h>
+#include <fmt/core.h>
#include <iostream>
#include <string>
#include <vector>
-#include "dd4pod/Geant4ParticleCollection.h"
-#include "eicd/TrackParametersCollection.h"
-
-// 4-momenta for all tracks, using the raw reconstruction results
-auto momenta(const std::vector<eic::TrackParametersData>& in) {
- std::vector<ROOT::Math::XYZMVector> results;
- std::transform(in.begin(), in.end(), results.begin(), [](const auto& part) {
- return {
- part.p.px, part.p.py, part.p.pz,
- .000511 // hard-assume these are electrons for now, will need PID down
- // the road FIXME
- };
- });
- return results;
-}
+// Run VM invariant-mass-based benchmarks on an input reconstruction file for
+// a desired vector meson (e.g. jpsi) and a desired decay particle (e.g. muon)
+// Output figures are written to our output prefix (which includes the output
+// file prefix), and labeled with our detector name.
+// TODO: I think it would be better to pass small json configuration file to
+// the test, instead of this ever-expanding list of function arguments.
+int vm_mass(std::string_view rec_file, std::string_view vm_name,
+ std::string_view decay_name, std::string_view detector,
+ std::string output_prefix) {
+ fmt::print(fmt::emphasis::bold | fg(fmt::color::forest_green),
+ "Running VM invariant mass analysis...\n");
+ fmt::print(" - Vector meson: {}\n", vm_name);
+ fmt::print(" - Decay particle: {}\n", decay_name);
+ fmt::print(" - Detector package: {}\n", detector);
+ fmt::print(" - output prefix: {}\n", output_prefix);
+
+ // Run this in multi-threaded mode if desired
+ ROOT::EnableImplicitMT(kNumThreads);
-// masses for all combinatorics
-auto calc_masses(const std::vector<ROOT::Math::XYZMVector>& p) {
- std::vector<double> masses;
- for (int i = 0; i < p.size(); ++i) {
- for (int j = i + 1; i < p.size(); ++j) {
- masses.push_back((p[i] + p[j]).mass());
- }
+ // The particles we are looking for. E.g. J/psi decaying into e+e-
+ const double vm_mass = util::get_pdg_mass(vm_name);
+ const double decay_mass = util::get_pdg_mass(decay_name);
+
+ // Ensure our output prefix always ends on a dot, a slash or a dash
+ if (output_prefix.back() != '.' && output_prefix.back() != '/' &&
+ output_prefix.back() != '-') {
+ output_prefix += "-";
}
- return masses
-}
-auto jpsi_mass(const std::string
+ // Open our input file file as a dataframe
+ ROOT::RDataFrame d{"events", rec_file};
+
+ // utility lambda functions to bind the vector meson and decay particle types
+ auto momenta_from_tracking =
+ [decay_mass](const std::vector<eic::TrackParametersData>& tracks) {
+ return util::momenta_from_tracking(tracks, decay_mass);
+ };
+ auto find_decay_pair =
+ [vm_mass](const std::vector<ROOT::Math::PxPyPzMVector>& parts) {
+ return util::find_decay_pair(parts, vm_mass);
+ };
+
+ // Define analysis flow
+ auto d_im =
+ d.Define("p_rec", momenta_from_tracking, {"outputTrackParameters"})
+ .Define("N", "p_rec.size()")
+ .Define("p_sim", util::momenta_from_simulation, {"mcparticles2"})
+ .Define("decay_pair_rec", find_decay_pair, {"p_rec"})
+ .Define("decay_pair_sim", find_decay_pair, {"p_sim"})
+ .Define("mass_rec", util::get_im, {"decay_pair_rec"})
+ .Define("mass_sim", util::get_im, {"decay_pair_sim"});
+
+ // Define output histograms
+ auto h_im_rec = d_im.Histo1D(
+ {"h_im_rec", ";m_{ll'} (GeV);#", 100, -1.1, vm_mass + 5}, "mass_rec");
+ auto h_im_sim = d_im.Histo1D(
+ {"h_im_sim", ";m_{ll'} (GeV);#", 100, -1.1, vm_mass + 5}, "mass_sim");
+
+ // Plot our histograms.
+ // TODO: to start I'm explicitly plotting the histograms, but want to
+ // factorize out the plotting code moving forward.
+ {
+ TCanvas c{"canvas", "canvas", 800, 800};
+ gPad->SetLogx(false);
+ gPad->SetLogy(false);
+ auto& h0 = *h_im_sim;
+ auto& h1 = *h_im_rec;
+ // histogram style
+ h0.SetLineColor(plot::kMpBlue);
+ h0.SetLineWidth(2);
+ h1.SetLineColor(plot::kMpOrange);
+ h1.SetLineWidth(2);
+ // axes
+ h0.GetXaxis()->CenterTitle();
+ h0.GetYaxis()->CenterTitle();
+ // draw everything
+ h0.DrawClone("hist");
+ h1.DrawClone("hist same");
+ // FIXME hardcoded beam configuration
+ plot::draw_label(10, 100, detector, vm_name, "Invariant mass");
+ TText* tptr;
+ auto t = new TPaveText(.6, .8417, .9, .925, "NB NDC");
+ t->SetFillColorAlpha(kWhite, 0);
+ t->SetTextFont(43);
+ t->SetTextSize(25);
+ tptr = t->AddText("simulated");
+ tptr->SetTextColor(plot::kMpBlue);
+ tptr = t->AddText("reconstructed");
+ tptr->SetTextColor(plot::kMpOrange);
+ t->Draw();
+ // Print canvas to output file
+ c.Print(fmt::format("{}vm_mass.png", output_prefix).c_str());
+ }
+ // That's all!
+ return 0;
+}
diff --git a/dvmp/dvmp.sh b/dvmp/dvmp.sh
index 379dedc3..908f94ef 100644
--- a/dvmp/dvmp.sh
+++ b/dvmp/dvmp.sh
@@ -55,7 +55,7 @@ npsim --runType batch \
--compactFile ${DETECTOR_PATH}/${JUGGLER_DETECTOR}.xml \
--inputFiles ${JUGGLER_GEN_FILE} \
--outputFile ${JUGGLER_SIM_FILE}
-if [[ "$?" -ne "0" ]] ; then
+if [ "$?" -ne "0" ] ; then
echo "ERROR running npsim"
exit 1
fi
@@ -66,7 +66,7 @@ echo "Running the digitization and reconstruction"
# FIXME Need to figure out how to pass file name to juggler from the commandline
xenv -x ${JUGGLER_INSTALL_PREFIX}/Juggler.xenv \
gaudirun.py options/tracker_reconstruction.py
-if [[ "$?" -ne "0" ]] ; then
+if [ "$?" -ne "0" ] ; then
echo "ERROR running juggler"
exit 1
fi
@@ -74,12 +74,17 @@ ls -l
## =============================================================================
## Step 4: Analysis
-echo "STAND-IN FOR ANALYSIS SCRIPT"
-#root -b -q "dis/scripts/rec_dis_electrons.cxx(\"${JUGGLER_DETECTOR}/${JUGGLER_REC_FILE}\")"
-#if [[ "$?" -ne "0" ]] ; then
-# echo "ERROR running root script"
-# exit 1
-#fi
+root -b -q "dvmp/analysis/vm_mass.cxx(\
+ \"${JUGGLER_REC_FILE}\", \
+ \"jpsi\", \
+ \"electron\", \
+ \"${JUGGLER_DETECTOR}\", \
+ \"results/dvmp/plot\")"
+
+if [ "$?" -ne "0" ] ; then
+ echo "ERROR running root script"
+ exit 1
+fi
## =============================================================================
## Step 5: finalize
@@ -87,7 +92,7 @@ echo "Finalizing ${JUGGLER_FILE_NAME_TAG} benchmark"
## Copy over reconsturction artifacts as long as we don't have
## too many events
-if [[ "${JUGGLER_N_EVENTS}" -lt "500" ]] ; then
+if [ "${JUGGLER_N_EVENTS}" -lt "500" ] ; then
cp ${JUGGLER_REC_FILE} results/dvmp/.
fi
--
GitLab