Skip to content

Commit

Permalink
HFS simple tree + neural network predictions using pybind (#214)
Browse files Browse the repository at this point in the history
  • Loading branch information
cpecar authored Dec 27, 2022
1 parent e7f9da4 commit 5a3d757
Show file tree
Hide file tree
Showing 12 changed files with 392 additions and 43 deletions.
1 change: 1 addition & 0 deletions .gitignore
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
# main ignores
.*.sw*
*~
*.hepmc*
*.root
*.pcm
Expand Down
10 changes: 7 additions & 3 deletions config.mk
Original file line number Diff line number Diff line change
@@ -1,9 +1,10 @@
# compiler and flags
CXX = g++
FLAGS = -g -Wno-deprecated -fPIC -fno-inline -Wno-write-strings
FLAGS = -Wno-deprecated -fPIC
FLAGS += -fmax-errors=3
# extra flags
#FLAGS += -O0
# FLAGS += -fvisibility=hidden # FIXME: required by pybind, but causes unresolved symbols in cling...
# FLAGS += -g # build with debugging symbols
# FLAGS += -O0

# extra flags for Mac OS
UNAME := $(shell uname)
Expand All @@ -19,6 +20,9 @@ LIBS = $(shell root-config --glibs)
# Data Model (PODIO + EDM4hep + EDM4eic)
LIBS += -L/usr/local/lib -lpodio -lpodioRootIO -ledm4hep -ledm4eic

# PYTHON (for pybind)
LIBS += -lpython3.10

# Miscellaneous
LIBS += -lfmt

Expand Down
48 changes: 48 additions & 0 deletions macro/analysis_testml.C
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
// SPDX-License-Identifier: LGPL-3.0-or-later
// Copyright (C) 2022 Connor Pecar

R__LOAD_LIBRARY(EpicAnalysis)
#include <pybind11/embed.h>
namespace py = pybind11;
// using ML prediction for vecQ, and writing out tree of HFS four-vectors
// requires pybind includes above
void analysis_testml(
TString configFile="datarec/in.config", /* delphes tree(s) */
TString outfilePrefix="resolutions" /* output filename prefix*/
) {
// object needed at start of script using pybind11
py::scoped_interpreter guard{};
//outfilePrefix+="_DA";
// setup analysis ========================================
AnalysisEcce *A = new AnalysisEcce(
configFile,
outfilePrefix
);

A->maxEvents = 10000; // use this to limit the number of events
A->writeSimpleTree = true; // write SimpleTree (for one bin)
A->writeHFSTree = true; // write HFSTree (for one bin)
A->SetReconMethod("ML"); // set reconstruction method
A->AddFinalState("pipTrack"); // pion final state
//A->AddFinalState("KpTrack"); // kaon final state


// define cuts ====================================
A->AddBinScheme("w"); A->BinScheme("w")->BuildBin("Min",3.0); // W > 3 GeV
//A->AddBinScheme("xF"); A->BinScheme("xF")->BuildBin("Min",0.0); // xF > 0
//A->AddBinScheme("ptLab"); A->BinScheme("ptLab")->BuildBin("Min",0.1); // pT_lab > 0.1 GeV (tracking limit)


// set binning scheme ====================================
// z ranges
//A->AddBinScheme("z");
//A->BinScheme("z")->BuildBin("Min",0.2); // needed?

// y minima
A->AddBinScheme("y");
A->BinScheme("y")->BuildBin("Full"); // a bin with no y-cut

// perform the analysis ==================================
A->Execute();

};
28 changes: 16 additions & 12 deletions src/Analysis.cxx
Original file line number Diff line number Diff line change
Expand Up @@ -92,16 +92,18 @@ Analysis::Analysis(

// common settings defaults
// - these settings can be set at the macro level
verbose = false;
writeSimpleTree = false;
verbose = false;
writeSimpleTree = false;
writeHFSTree = false;
writeParticleTree = false;
maxEvents = 0;
useBreitJets = false;
errorCntMax = 1000;
jetAlg = 0; // Default to kT Algorithm
jetRad = 0.8;
jetMin = 1.0; // Minimum Jet pT
jetMatchDR = 0.5; // Delta R between true and reco jet to be considered matched

maxEvents = 0;
useBreitJets = false;
errorCntMax = 1000;
jetAlg = 0; // Default to kT Algorithm
jetRad = 0.8;
jetMin = 1.0; // Minimum Jet pT
jetMatchDR = 0.5; // Delta R between true and reco jet to be considered matched

weightInclusive = new WeightsUniform();
weightTrack = new WeightsUniform();
Expand Down Expand Up @@ -310,8 +312,9 @@ void Analysis::Prepare() {
// instantiate shared objects
kin = new Kinematics(eleBeamEn,ionBeamEn,crossingAngle);
kinTrue = new Kinematics(eleBeamEn, ionBeamEn, crossingAngle);
ST = new SimpleTree("tree",kin,kinTrue);
PT = new ParticleTree("ptree");
ST = new SimpleTree("tree",kin,kinTrue);
HFST = new HFSTree("hfstree",kin,kinTrue);
PT = new ParticleTree("ptree");

// if including jets, define a `jet` final state
#ifndef EXCLUDE_DELPHES
Expand Down Expand Up @@ -660,7 +663,8 @@ void Analysis::Finish() {
cout << sep << endl;
cout << "writing ROOT file..." << endl;
outFile->cd();
if(writeSimpleTree) ST->WriteTree();
if(writeSimpleTree) ST->WriteTree();
if(writeHFSTree) HFST->WriteTree();
if(writeParticleTree) PT->WriteTree();
HD->Payload([this](Histos *H){ H->WriteHists(outFile); }); HD->ExecuteAndClearOps();
HD->Payload([this](Histos *H){ H->Write(); }); HD->ExecuteAndClearOps();
Expand Down
5 changes: 4 additions & 1 deletion src/Analysis.h
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@
#include "HistosDAG.h"
#include "Kinematics.h"
#include "SimpleTree.h"
#include "HFSTree.h"
#include "ParticleTree.h"
#include "Weights.h"
#include "CommonConstants.h"
Expand Down Expand Up @@ -60,7 +61,8 @@ class Analysis : public TNamed

// common settings
Bool_t verbose; // if true, print a lot more information
Bool_t writeSimpleTree; // if true, write SimpleTree (not binned)
Bool_t writeSimpleTree; // if true, write SimpleTree (not binned)
Bool_t writeHFSTree; // if true, write HFSTree (not binned)
Bool_t writeParticleTree; // if true, write ParticleTree (not binned)
Long64_t maxEvents; /* default=0, which runs all events;
* if > 0, run a maximum number of `maxEvents` events (useful for quick tests)
Expand Down Expand Up @@ -124,6 +126,7 @@ class Analysis : public TNamed

// shared objects
SimpleTree *ST;
HFSTree *HFST;
ParticleTree *PT;
Kinematics *kin, *kinTrue;
HistosDAG *HD;
Expand Down
69 changes: 53 additions & 16 deletions src/AnalysisEcce.cxx
Original file line number Diff line number Diff line change
Expand Up @@ -101,14 +101,11 @@ void AnalysisEcce::Execute()
tr.SetEntriesRange(1,maxEvents);
do {
if(tr.GetCurrentEntry()%10000==0) cout << tr.GetCurrentEntry() << " events..." << endl;

// resets
kin->ResetHFS();
kinTrue->ResetHFS();




// a few maps needed to get the associated info between tracks, true particles, etec

std::map <int,int> mcidmap;
Expand Down Expand Up @@ -140,16 +137,16 @@ void AnalysisEcce::Execute()
* - find beam particles
*/
std::vector<Particles> mcpart;
double maxP = 0;
double maxPt = 0;
int genEleID = -1;
bool foundBeamElectron = false;
bool foundBeamIon = false;
int genEleBCID = -1;

for(int imc=0; imc<hepmcp_PDG.GetSize(); imc++) {

int pid_ = hepmcp_PDG[imc];


int genStatus_ = hepmcp_status[imc]; // genStatus 4: beam particle, 1: final state

double px_ = hepmcp_psx[imc];
Expand All @@ -158,10 +155,14 @@ void AnalysisEcce::Execute()
double e_ = hepmcp_E[imc];

double p_ = sqrt(pow(hepmcp_psx[imc],2) + pow(hepmcp_psy[imc],2) + pow(hepmcp_psz[imc],2));
double pt_ = sqrt(pow(hepmcp_psx[imc],2) + pow(hepmcp_psy[imc],2));
double mass_ = (fabs(pid_)==211)?constants::pimass:(fabs(pid_)==321)?constants::kmass:(fabs(pid_)==11)?constants::emass:(fabs(pid_)==13)?constants::mumass:(fabs(pid_)==2212)?constants::pmass:0.;

// add to `mcpart`
Particles part;

//cout << genStatus_ << " " << pid_ << " " << part.mcID << " " << hepmcp_BCID[imc] << " " << hepmcp_m1[imc] << " " << hepmcp_m2[imc] << " "
// << px_ << " " << py_ << " " << pz_ << " " << e_ << endl;

if(genStatus_ == 1) { // final state

Expand All @@ -170,18 +171,22 @@ void AnalysisEcce::Execute()
if (search != mcbcidmap.end()) {
imcpart = search->second;
}



if (imcpart >-1){
px_ = mcpart_psx[imcpart];
py_ = mcpart_psy[imcpart];
pz_ = mcpart_psz[imcpart];
e_ = mcpart_E[imcpart];
p_ = sqrt(pow(mcpart_psx[imcpart],2) + pow(mcpart_psy[imcpart],2) + pow(mcpart_psz[imcpart],2));
pt_ = sqrt(pow(mcpart_psx[imcpart],2) + pow(mcpart_psy[imcpart],2));
part.mcID = mcpart_ID[imcpart];
}
else
part.mcID = -1;

//cout << genStatus_ << " " << pid_ << " " << part.mcID << " " << hepmcp_BCID[imc] << " " << hepmcp_m1[imc] << " " << hepmcp_m2[imc] << " "
// << px_ << " " << py_ << " " << pz_ << " " << e_ << endl;

part.pid = pid_;
part.vecPart.SetPxPyPzE(px_, py_, pz_, e_);

Expand All @@ -193,15 +198,16 @@ void AnalysisEcce::Execute()

// identify scattered electron by max momentum
if(pid_ == 11) {
if(p_ > maxP) {
maxP = p_;
if(pt_ > maxPt) {
maxPt = pt_;
kinTrue->vecElectron.SetPxPyPzE(px_, py_, pz_, e_);
genEleID = part.mcID; //mcpart_ID[imc];
genEleBCID = hepmcp_BCID[imc];
// cout << "\t\t\t found scattered electron " << Form(" %6.2f %6.2f %6.2f %6.2f %6.2f %5.3f %6.2f %6.2f id %3d\n",px_,py_,pz_, sqrt(p_*p_ + mass_*mass_),p_,mass_,hepmcp_E[imc],mcpart_E[imcpart],genEleID);
}
}
}

else if(genStatus_ == 4) { // beam particles
if(pid_ == 11) { // electron beam
if(!foundBeamElectron) {
Expand All @@ -221,11 +227,37 @@ void AnalysisEcce::Execute()
}
}
} // end truth loop

// looking for radiated photons (mother particle == scattered electron)
// requires that we already know scattered electron index
for(int imc=0; imc<hepmcp_PDG.GetSize(); imc++) {
int pid_ = hepmcp_PDG[imc];
int genStatus_ = hepmcp_status[imc]; // genStatus 4: beam particle, 1: final state

double px_ = hepmcp_psx[imc];
double py_ = hepmcp_psy[imc];
double pz_ = hepmcp_psz[imc];
double e_ = hepmcp_E[imc];

double p_ = sqrt(pow(hepmcp_psx[imc],2) + pow(hepmcp_psy[imc],2) + pow(hepmcp_psz[imc],2));

int mother = hepmcp_m1[imc];
if(pid_ == 22 || pid_ == 23){
if(genStatus_ == 1){
if( mother == genEleBCID ){
TLorentzVector FSRmom(px_, py_, pz_, e_);
TLorentzVector eleCorr = kinTrue->vecElectron + FSRmom;
// correcting true scattered electron with FSR
kinTrue->vecElectron.SetPxPyPzE(eleCorr.Px(), eleCorr.Py(), eleCorr.Pz(), eleCorr.E());
}
}
}


}
// check beam finding
if(!foundBeamElectron || !foundBeamIon) { numNoBeam++; continue; };


// reconstructed particles loop
/* - add reconstructed particle to `recopart`
* - find the scattered electron
Expand Down Expand Up @@ -268,14 +300,15 @@ void AnalysisEcce::Execute()
part.vecPart.SetPxPyPzE(reco_px, reco_py, reco_pz, sqrt(reco_p*reco_p + reco_mass*reco_mass));


// cout << "\t\t\t track " << Form(" %4.2f %4.2f %4.2f true id %4d imc %3d mcid %3d \n",reco_px,reco_py,reco_pz,tracks_trueID[ireco],imc,part.mcID);
//cout << "\t\t\t track " << Form(" %4.2f %4.2f %4.2f true id %4d imc %3d mcid %3d \n",reco_px,reco_py,reco_pz,tracks_trueID[ireco],imc,part.mcID);

// add to `recopart` and hadronic final state sums only if there is a matching truth particle
if(part.mcID > 0) {
if(part.mcID > 0 && part.mcID != genEleID) {
if(imc>-1) {
// cout << "\t\t\t add to hadfs \n" ;
recopart.push_back(part);
kin->AddToHFS(part.vecPart);
kin->hfspid[kin->nHFS - 1] = pid_;
}
}

Expand All @@ -288,6 +321,7 @@ void AnalysisEcce::Execute()

} // end reco loop


//add all clusters for hadronic final state
int jentryt = tr.GetCurrentEntry();
Long64_t localEntry = chain->LoadTree(jentryt);
Expand Down Expand Up @@ -386,10 +420,13 @@ void AnalysisEcce::Execute()
// calculate DIS kinematics
if(!(kin->CalculateDIS(reconMethod))) continue; // reconstructed
if(!(kinTrue->CalculateDIS(reconMethod))) continue; // generated (truth)

// Get the weight for this event's Q2
auto Q2weightFactor = GetEventQ2Weight(kinTrue->Q2, inLookup[chain->GetTreeNumber()]);

//fill HFS tree for event
if( writeHFSTree && kin->nHFS > 0) HFST->FillTree(Q2weightFactor);

// fill inclusive histograms, if only `inclusive` is included in output
// (otherwise they will be filled in track and jet loops)
if(includeOutputSet["inclusive_only"]) {
Expand Down
31 changes: 31 additions & 0 deletions src/HFSTree.cxx
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
// SPDX-License-Identifier: LGPL-3.0-or-later
// Copyright (C) 2022 Connor Pecar

#include "HFSTree.h"

ClassImp(HFSTree)

// constructor
HFSTree::HFSTree(TString treeName_, Kinematics *K_, Kinematics *Ktrue_)
: treeName(treeName_)
, K(K_)
, Ktrue(Ktrue_)
{
T = new TTree(treeName,treeName);
T->Branch("vecElectron", &(K->vecElectron));
T->Branch("vecElectronTrue", &(Ktrue->vecElectron));
T->Branch("vecEleBeamTrue", &(Ktrue->vecEleBeam));
T->Branch("vecIonBeamTrue", &(Ktrue->vecIonBeam));
T->Branch("nHFS", &(K->nHFS), "nHFS/I");
T->Branch("hfsp4", &(K->hfsp4));
//T->Branch("hfspy", &(K->hfspy), "hfspy[nHFS]/F");
//T->Branch("hfspz", &(K->hfspz), "hfspz[nHFS]/F");
//T->Branch("hfsE", &(K->hfsE), "hfsE[nHFS]/F");
T->Branch("hfspid", &(K->hfspid), "hfspid[nHFS]/F");
T->Branch("weight", &(weight), "weight/D");
};

// destructor
HFSTree::~HFSTree() {
};

Loading

0 comments on commit 5a3d757

Please sign in to comment.