drich/script/dRICHAna_bootstrap.cpp

0001 #include <random>
0002 #include <vector>
0003 #include <numeric>
0004 #include <cmath>
0005 #include "podio/ROOTReader.h"
0006 #include "podio/Frame.h"
0007 #include "edm4eic/CherenkovParticleIDCollection.h"
0008 #include "edm4hep/MCParticleCollection.h"
0009 #include <chrono>
0010
0011 // struct to store needed data from each event,
0012 // so that we only have to loop over tree once
0013 struct EventData {
0014   std::vector<double> theta;      // theta*1000 per photon
0015   std::vector<double> thetaErr;   // |theta-expected| per photon
0016   double nPhotons;                // number of photons
0017   bool   detected;                // nPhotons > 5
0018 };
0019
0020 // loop over all events in the file, storing needed dRICH information
0021 std::vector<EventData> readEventData(const char* infile, int radiator) {
0022   std::vector<EventData> events;
0023   podio::ROOTReader reader;
0024   reader.openFile(infile);
0025   int nev = reader.getEntries("events");
0026   events.reserve(nev);
0027
0028   // gas or aerogel radiator
0029   std::string pidCol;
0030   double n_refr;
0031   double thmin, thmax;
0032   if (radiator == 0) {
0033     pidCol = "DRICHAerogelIrtCherenkovParticleID";
0034     n_refr = 1.019;
0035     thmax = 220;
0036     thmin = 150;
0037   } else if (radiator == 1) {
0038     pidCol = "DRICHGasIrtCherenkovParticleID";
0039     n_refr = 1.00076;
0040     thmax = 50;
0041     thmin = 20;
0042   }
0043   else throw std::runtime_error("bad radiator");
0044
0045   for (int i = 0; i < nev; ++i) {
0046     auto frame = podio::Frame(reader.readNextEntry("events"));
0047     auto& chcol = frame.get<edm4eic::CherenkovParticleIDCollection>(pidCol);
0048     auto& mc    = frame.get<edm4hep::MCParticleCollection>("MCParticles");
0049
0050     // get MC info, including expected cherenkov angle
0051     auto p4 = mc[0].getMomentum();
0052     double p = std::hypot(p4.x,p4.y,p4.z);
0053     double betaTrue = p/std::sqrt(p*p + mc[0].getMass()*mc[0].getMass());
0054     double chExpected = std::acos(1./(n_refr*betaTrue))*1000;
0055
0056     // get data for event.
0057     // only keep reconstructed photons in a (wide) reasonable range
0058     for (const auto& pid : chcol) {
0059       EventData ed;
0060       ed.nPhotons = 0;
0061
0062       for (auto& tp : pid.getThetaPhiPhotons()) {
0063         double th = tp[0]*1000;
0064         if (th < thmax && th > thmin) {
0065           ed.theta.push_back(th);
0066           ed.thetaErr.push_back(th - chExpected);
0067           ed.nPhotons += 1.;
0068         }
0069       }
0070       ed.detected = (ed.nPhotons > 5); // for efficiency calculation
0071       events.push_back(std::move(ed));
0072     }
0073   }
0074
0075   return events;
0076 }
0077
0078 struct BootstrapResults {
0079   double mean_nPhot;
0080   double mean_theta;
0081   double mean_theta_mae;
0082   double frac_detected;
0083 };
0084
0085 // bootstrap (sample with replacement) sampleSize events
0086 // out of the vector of simulated dRICH events
0087 BootstrapResults computeBootstrap(
0088     const std::vector<EventData>& events,
0089     int sampleSize,
0090     std::mt19937_64& rng
0091 ) {
0092   std::uniform_int_distribution<> pick(0, events.size() - 1);
0093
0094   double sum_nphot = 0.0;
0095   double sum_theta = 0.0;
0096   double sum_theta_err = 0.0;
0097   double total_photons = 0.0;
0098   int    n_detected = 0;
0099
0100   // sample with replacement sampleSize times
0101   for (int i = 0; i < sampleSize; ++i) {
0102     const auto& ev = events[pick(rng)];
0103     sum_nphot += ev.nPhotons;
0104     if (ev.detected) ++n_detected;
0105
0106     for (double th : ev.theta) {
0107       sum_theta += th;
0108       total_photons += 1.;
0109     }
0110     for (double err : ev.thetaErr) {
0111       sum_theta_err += std::fabs(err);
0112     }
0113   }
0114
0115   // for this set of sampled events, get information needed to
0116   // calculate metrics
0117   BootstrapResults out;
0118   out.mean_nPhot       = sum_nphot / sampleSize;
0119   out.frac_detected    = double(n_detected) / sampleSize;
0120   if (total_photons > 0) {
0121     out.mean_theta     = sum_theta     / total_photons;
0122     out.mean_theta_mae = sum_theta_err / total_photons;
0123   } else {
0124     out.mean_theta     = 0.0;
0125     out.mean_theta_mae = 0.0;
0126   }
0127   return out;
0128 }
0129
0130 std::tuple<double,double,double,double> bootstrapStats(const std::vector<EventData>& events_pi,
0131                                const std::vector<EventData>& events_K,
0132                                int sampleSize,
0133                                int nBootstrap
0134                                ){
0135
0136   // compute piKsep and acceptance metrics nBootstrap times
0137   std::vector<double> final_piKsep(nBootstrap), final_acc(nBootstrap);
0138   std::mt19937_64 rng(std::chrono::steady_clock::now().time_since_epoch().count());
0139
0140   for (int b = 0; b < nBootstrap; b++) {
0141     // get MAE, nPhotons, etc from each resampling
0142     auto stats_pi = computeBootstrap(events_pi, sampleSize, rng);
0143     auto stats_K  = computeBootstrap(events_K,  sampleSize, rng);
0144
0145     double cher_diff = std::fabs(stats_pi.mean_theta  - stats_K.mean_theta);
0146
0147     double avg_mae = (stats_pi.mean_theta_mae + stats_K.mean_theta_mae)/2.;
0148     double avg_acc = (stats_pi.frac_detected + stats_K.frac_detected)/2.;
0149     double avg_nphot = (stats_pi.mean_nPhot + stats_K.mean_nPhot)/2.;
0150
0151     // calculate and store pion-kaon separation
0152     final_piKsep[b] = cher_diff*sqrt(avg_nphot)/avg_mae;
0153     final_acc[b] = avg_acc;
0154   }
0155
0156   auto compute_stats = [&](const std::vector<double>& v) {
0157     double N = v.size();
0158     double sum = std::accumulate(v.begin(), v.end(), 0.0);
0159     double mean = sum/N;
0160     double sq = 0;
0161     for (double x : v) sq += (x-mean)*(x-mean);
0162     double stdev = std::sqrt(sq/(N-1));
0163
0164     return std::tuple{mean, stdev};
0165   };
0166
0167   // get mean and standard dev for acceptance and pi-K separation
0168   auto [mean_acc, sd_acc] = compute_stats(final_acc);
0169   auto [mean_piKsep, sd_piKsep] = compute_stats(final_piKsep);
0170
0171   return std::tuple{mean_acc, sd_acc, mean_piKsep, sd_piKsep};
0172 }
0173
0174
0175 void dRICHAna_bootstrap(const char* infile_pi,
0176                const char* infile_K,
0177                const char* outname,
0178                int radiator,
0179                int sampleSize = 1000,
0180                int nBootstrap = 100
0181                )
0182 {
0183
0184   auto events_pi = readEventData(infile_pi, radiator);
0185   auto events_K  = readEventData(infile_K,  radiator);
0186
0187   std::cout << "N ev pi: " << events_pi.size() << " K: " << events_K.size() << std::endl;
0188   if (sampleSize > static_cast<int>(events_pi.size())) {
0189     sampleSize = static_cast<int>(events_pi.size());
0190   }
0191   if (sampleSize > static_cast<int>(events_K.size())) {
0192     sampleSize = static_cast<int>(events_K.size());
0193   }
0194
0195   auto [mean_acc, sd_acc, mean_piKsep, sd_piKsep]  = bootstrapStats(events_pi,
0196                                     events_K,
0197                                     sampleSize,
0198                                     nBootstrap
0199                                     );
0200   if (sampleSize < static_cast<int>(events_pi.size())) {
0201     double ratio = double(sampleSize)/events_pi.size();
0202     sd_piKsep *= ratio;
0203     sd_acc *= ratio;
0204   }
0205
0206   FILE *outfile = fopen(outname,"w");
0207   fprintf(outfile, "%lf %lf %lf %lf \n", mean_acc, sd_acc, mean_piKsep, sd_piKsep);
0208   return;
0209 }
0210
0211 int main(int argc, char* argv[]) {
0212   if (argc < 7) {
0213     std::cout << "usage: dRICHAna_bootstrap [file, pi] [file, K] [output file name] [radiator: 0 - aerogel, 1 - gas] [N samples (optional)] [N bootstraps] \n";
0214     return 1;
0215   }
0216
0217   int rad         = std::stoi(argv[4]);
0218   int nsamples    = std::stoi(argv[5]);
0219   int nbootstraps = std::stoi(argv[6]);
0220
0221   dRICHAna_bootstrap(argv[1], argv[2], argv[3], rad, nsamples, nbootstraps);
0222   return 0;
0223 }