Skip to content

Commit bd9e617

Browse files
DraTeotsDraTeots
committed
FF Lambdas script
1 parent 380ce6f commit bd9e617

File tree

3 files changed

+515
-190
lines changed

3 files changed

+515
-190
lines changed

‎csv_convert/CMakeLists.txt‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,4 +54,15 @@ target_link_libraries(csv_reco_dis
5454
podio::podio podio::podioRootIO podio::podioDict
5555
EDM4HEP::edm4hep EDM4HEP::edm4hepDict
5656
EDM4EIC::edm4eic EDM4EIC::edm4eicDict
57-
fmt::fmt)
57+
fmt::fmt)
58+
59+
60+
add_executable(csv_reco_ff_lambda csv_reco_ff_lambda.cxx)
61+
62+
target_link_libraries(csv_reco_ff_lambda
63+
PRIVATE
64+
ROOT::Core ROOT::RIO ROOT::Tree ROOT::Hist ROOT::Gpad
65+
podio::podio podio::podioRootIO podio::podioDict
66+
EDM4HEP::edm4hep EDM4HEP::edm4hepDict
67+
EDM4EIC::edm4eic EDM4EIC::edm4eicDict
68+
fmt::fmt)

‎csv_convert/csv_reco_ff_lambda.cxx‎

Lines changed: 225 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,225 @@
1+
#ifdef __CLING__
2+
R__LOAD_LIBRARY(podioDict)
3+
R__LOAD_LIBRARY(podioRootIO)
4+
R__LOAD_LIBRARY(libedm4hepDict)
5+
R__LOAD_LIBRARY(libedm4eicDict)
6+
#endif
7+
8+
#include "podio/Frame.h"
9+
#include "podio/ROOTReader.h"
10+
#include <edm4eic/ReconstructedParticleCollection.h>
11+
12+
#include <fmt/core.h>
13+
#include <fmt/ostream.h>
14+
15+
#include <TFile.h>
16+
17+
#include <fstream>
18+
#include <string>
19+
#include <vector>
20+
#include <cstdlib>
21+
22+
using namespace edm4eic;
23+
24+
//------------------------------------------------------------------------------
25+
// globals & helpers
26+
//------------------------------------------------------------------------------
27+
int events_limit = -1; // -n <N>
28+
long total_evt_seen = 0;
29+
long total_lambdas_written = 0;
30+
std::ofstream csv;
31+
bool header_written = false;
32+
33+
//------------------------------------------------------------------------------
34+
// event processing
35+
//------------------------------------------------------------------------------
36+
void process_event(const podio::Frame& event, int evt_id) {
37+
38+
// Write header if first time
39+
if (!header_written) {
40+
csv << "evt,"
41+
<< "idx,"
42+
<< "pdg,"
43+
<< "charge,"
44+
<< "energy,"
45+
<< "mass,"
46+
<< "px,"
47+
<< "py,"
48+
<< "pz,"
49+
<< "ref_x,"
50+
<< "ref_y,"
51+
<< "ref_z,"
52+
<< "pid_goodness,"
53+
<< "type,"
54+
<< "n_clusters,"
55+
<< "n_tracks,"
56+
<< "n_particles,"
57+
<< "n_particle_ids,"
58+
<< "cov_xx,"
59+
<< "cov_xy,"
60+
<< "cov_xz,"
61+
<< "cov_yy,"
62+
<< "cov_yz,"
63+
<< "cov_zz,"
64+
<< "cov_xt,"
65+
<< "cov_yt,"
66+
<< "cov_zt,"
67+
<< "cov_tt\n";
68+
header_written = true;
69+
}
70+
71+
// Get the reconstructed far-forward lambdas collection
72+
const auto& ffLambdas = event.get<ReconstructedParticleCollection>("ReconstructedFarForwardZDCLambdas");
73+
74+
// Process each lambda in the collection
75+
int lam_idx = 0;
76+
for (const auto& lam : ffLambdas) {
77+
const auto mom = lam.getMomentum();
78+
const auto ref = lam.getReferencePoint();
79+
const auto cov = lam.getCovMatrix();
80+
81+
// Calculate the number of associated objects
82+
const size_t n_clusters = lam.getClusters().size();
83+
const size_t n_tracks = lam.getTracks().size();
84+
const size_t n_particles = lam.getParticles().size();
85+
const size_t n_particleIDs = lam.getParticleIDs().size();
86+
87+
csv << evt_id << ","
88+
<< lam_idx << ","
89+
<< lam.getPDG() << ","
90+
<< lam.getCharge() << ","
91+
<< lam.getEnergy() << ","
92+
<< lam.getMass() << ","
93+
<< mom.x << ","
94+
<< mom.y << ","
95+
<< mom.z << ","
96+
<< ref.x << ","
97+
<< ref.y << ","
98+
<< ref.z << ","
99+
<< lam.getGoodnessOfPID() << ","
100+
<< lam.getType() << ","
101+
<< n_clusters << ","
102+
<< n_tracks << ","
103+
<< n_particles << ","
104+
<< n_particleIDs << ","
105+
<< cov.xx << "," // xx
106+
<< cov.xy << "," // xy
107+
<< cov.xz << "," // xz
108+
<< cov.xt << "," // xt
109+
<< cov.yy << "," // yy
110+
<< cov.yz << "," // yz
111+
<< cov.yt << "," // yt
112+
<< cov.zz << "," // zz
113+
<< cov.zt << "," // zt
114+
<< cov.tt << "\n"; // tt
115+
116+
if (n_particles > 0) {
117+
auto part1 = lam.getParticles().at(0);
118+
}
119+
120+
total_lambdas_written++;
121+
}
122+
}
123+
124+
//------------------------------------------------------------------------------
125+
// file loop
126+
//------------------------------------------------------------------------------
127+
void process_file(const std::string& fname) {
128+
podio::ROOTReader rdr;
129+
try {
130+
rdr.openFile(fname);
131+
}
132+
catch (const std::runtime_error& e) {
133+
fmt::print(stderr, "Error opening file {}: {}\n", fname, e.what());
134+
return;
135+
}
136+
137+
const auto nEv = rdr.getEntries(podio::Category::Event);
138+
fmt::print("Processing {} events from {}\n", nEv, fname);
139+
140+
for (unsigned ie = 0; ie < nEv; ++ie) {
141+
if (events_limit > 0 && total_evt_seen >= events_limit) return;
142+
143+
podio::Frame evt(rdr.readNextEntry(podio::Category::Event));
144+
process_event(evt, total_evt_seen);
145+
++total_evt_seen;
146+
}
147+
}
148+
149+
//------------------------------------------------------------------------------
150+
// main
151+
//------------------------------------------------------------------------------
152+
int main(int argc, char* argv[]) {
153+
std::vector<std::string> infiles;
154+
std::string out_name = "reco_ff_lambdas.csv";
155+
156+
for (int i = 1; i < argc; ++i) {
157+
std::string a = argv[i];
158+
if (a == "-n" && i + 1 < argc) events_limit = std::atoi(argv[++i]);
159+
else if (a == "-o" && i + 1 < argc) out_name = argv[++i];
160+
else if (a == "-h" || a == "--help") {
161+
fmt::print("usage: {} [-n N] [-o file] input1.root [...]\n", argv[0]);
162+
fmt::print(" -n N Process only N events (default: all)\n");
163+
fmt::print(" -o file Output CSV file (default: reco_ff_lambdas.csv)\n");
164+
return 0;
165+
}
166+
else if (!a.empty() && a[0] != '-') infiles.emplace_back(a);
167+
else {
168+
fmt::print(stderr, "unknown option {}\n", a);
169+
return 1;
170+
}
171+
}
172+
173+
if (infiles.empty()) {
174+
fmt::print(stderr, "error: no input files\n");
175+
return 1;
176+
}
177+
178+
csv.open(out_name);
179+
if (!csv) {
180+
fmt::print(stderr, "error: cannot open output file {}\n", out_name);
181+
return 1;
182+
}
183+
184+
fmt::print("Processing {} file(s)\n", infiles.size());
185+
for (auto& f : infiles) {
186+
fmt::print("\n=== Processing file: {} ===\n", f);
187+
process_file(f);
188+
if (events_limit > 0 && total_evt_seen >= events_limit) break;
189+
}
190+
191+
csv.close();
192+
fmt::print("\n\nTotal events processed: {}\n", total_evt_seen);
193+
fmt::print("Total reconstructed far-forward lambdas written: {}\n", total_lambdas_written);
194+
fmt::print("Output written to: {}\n", out_name);
195+
return 0;
196+
}
197+
198+
// ---------------------------------------------------------------------------
199+
// ROOT-macro entry point.
200+
// Call it from the prompt: root -x -l -b -q 'csv_reco_ff_lambda.cxx("file.root", "output.csv", 100)'
201+
// ---------------------------------------------------------------------------
202+
void csv_reco_ff_lambda(const char* infile, const char* outfile = "reco_ff_lambdas.csv", int events = -1) {
203+
fmt::print("'csv_reco_ff_lambda' entry point is used. Arguments:\n");
204+
fmt::print(" infile: {}\n", infile);
205+
fmt::print(" outfile: {}\n", outfile);
206+
fmt::print(" events: {} {}\n", events, (events == -1 ? "(process all)" : ""));
207+
208+
csv.open(outfile);
209+
if (!csv) {
210+
fmt::print(stderr, "error: cannot open output file {}\n", outfile);
211+
return;
212+
}
213+
214+
events_limit = events;
215+
total_evt_seen = 0;
216+
total_lambdas_written = 0;
217+
header_written = false;
218+
219+
process_file(infile);
220+
221+
csv.close();
222+
fmt::print("\nTotal events processed: {}\n", total_evt_seen);
223+
fmt::print("Total reconstructed far-forward lambdas written: {}\n", total_lambdas_written);
224+
fmt::print("Output written to: {}\n", outfile);
225+
}

0 commit comments

Comments
 (0)