epic-analysis/src/Kinematics.h

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2023 Christopher Dilks, Connor Pecar, Duane Byer, Sanghwa Park, Brian Page
0003
0004 /* NOTE:
0005  * if you make changes, MAINTAIN DOCUMENTATION IN ../doc/kinematics.md
0006  */
0007
0008 #pragma once
0009
0010 #include <stdio.h>
0011 #include <stdlib.h>
0012 #include <iostream>
0013 #include <sstream>
0014 #include <stdexcept>
0015
0016 // ROOT
0017 #include <TSystem.h>
0018 #include <TObject.h>
0019 #include <TFile.h>
0020 #include <TString.h>
0021 #include <TH1.h>
0022 #include <TMath.h>
0023 #include <TLorentzVector.h>
0024 #include <TRandom.h>
0025 #include <TRandomGen.h>
0026 #include <TClonesArray.h>
0027
0028 // Delphes
0029 #ifndef EXCLUDE_DELPHES
0030 #include <classes/DelphesClasses.h>
0031 #endif
0032 // pybind (for ML models using python packages)
0033 #ifdef SIDIS_MLPRED
0034 #include <pybind11/pybind11.h>
0035 #include <pybind11/numpy.h>
0036 #include <pybind11/embed.h>
0037 #include <pybind11/stl.h>
0038 namespace py = pybind11;
0039 #endif
0040
0041 using std::map;
0042 using std::cout;
0043 using std::cerr;
0044 using std::endl;
0045
0046 class Kinematics
0047 {
0048   public:
0049     Kinematics(Double_t enEleBeam, Double_t enIonBeam, Double_t crossAng);
0050     ~Kinematics();
0051
0052     // SIDIS calculators
0053     Bool_t CalculateDIS(TString recmethod); // return true if succeeded
0054     void CalculateHadronKinematics();
0055
0056     // hadronic final state (HFS)
0057     void AddToHFS(TLorentzVector p4_);
0058     void AddToHFSTree(TLorentzVector p4, int pid);
0059     void AddTrackToHFSTree(TLorentzVector p4, int pid);
0060     void SubtractElectronFromHFS();
0061     void ResetHFS();
0062
0063
0064     // DELPHES-specific methods //////////////////////////
0065 #ifndef EXCLUDE_DELPHES
0066
0067     // hadronic final state (HFS)
0068     void GetHFS(
0069         TObjArrayIter itTrack,
0070         TObjArrayIter itEFlowTrack,
0071         TObjArrayIter itEFlowPhoton,
0072         TObjArrayIter itEFlowNeutralHadron,
0073         TObjArrayIter itpfRICHTrack,
0074         TObjArrayIter itDIRCepidTrack,   TObjArrayIter itDIRChpidTrack,
0075         TObjArrayIter itBTOFepidTrack,   TObjArrayIter itBTOFhpidTrack,
0076         TObjArrayIter itdualRICHagTrack, TObjArrayIter itdualRICHcfTrack
0077         );
0078     void GetTrueHFS(TObjArrayIter itParticle);
0079
0080     // PID
0081     int GetTrackPID(
0082         Track *track,
0083         TObjArrayIter itpfRICHTrack,
0084         TObjArrayIter itDIRCepidTrack, TObjArrayIter itDIRChpidTrack,
0085         TObjArrayIter itBTOFepidTrack, TObjArrayIter itBTOFhpidTrack,
0086         TObjArrayIter itdualRICHagTrack, TObjArrayIter itdualRICHcfTrack
0087         );
0088
0089 #endif // ifndef EXCLUDE_DELPHES
0090     // end DELPHES-specific methods //////////////////////////
0091
0092
0093     // kinematics (should be Double_t, if going in SidisTree)
0094     Double_t W,Q2,Nu,x,y,s; // DIS
0095     Double_t pLab,pTlab,phiLab,etaLab,z,pT,qT,mX,xF,phiH,phiS; // hadron
0096     Double_t sigmah, Pxh, Pyh; // hadronic final state variables
0097     TLorentzVector hadronSumVec;
0098
0099     // nucleon transverse spin; if you set this externally,
0100     // it must be done before calculating `phiS` (before
0101     // `CalculateHadronKinematics`)
0102     Int_t tSpin; // should be +1 or -1
0103     Int_t lSpin; // should be +1 or -1
0104     Int_t hadPID;
0105
0106     // polarization
0107     Double_t polT;
0108     Double_t polL;
0109     Double_t polBeam;
0110
0111     // depolarization
0112     Double_t gamma,epsilon;
0113     // - factors A,B,C,V,W from [hep-ph/0611265] using notation from [1408.5721]
0114     Double_t depolA, depolB, depolC, depolV, depolW;
0115     // - ratios of factors, following notation of [1807.10606] eq. 2.3 (cf. eqs. 2.2a,b)
0116     Double_t depolP1; // for A_UT*sin(phiH+phiS) (collins), A_UT*sin(3phiH-phiS) (pretzelosity)
0117     Double_t depolP2; // for A_LL*const
0118     Double_t depolP3; // for twist-3 A_UT
0119     Double_t depolP4; // for A_LL*cos(phiH)
0120
0121     // 4-vectors
0122     // - lab frame
0123     TLorentzVector vecEleBeam, vecIonBeam;
0124     TLorentzVector vecElectron, vecW, vecQ;
0125     TLorentzVector vecHadron;
0126
0127     // HFS tree objects
0128     Int_t nHFS;
0129     std::vector<double> hfspx;
0130     std::vector<double> hfspy;
0131     std::vector<double> hfspz;
0132     std::vector<double> hfsE;
0133     std::vector<int> hfspid;
0134
0135     // HFS tree select FS tracks and matched true p4
0136     // (NOT full true HFS)
0137     std::vector<double> selectedHadPx;
0138     std::vector<double> selectedHadPy;
0139     std::vector<double> selectedHadPz;
0140     std::vector<double> selectedHadE;
0141     std::vector<int> selectedHadPID;
0142
0143     // TMVA for ML sidis reconstruction
0144     std::vector<std::vector<float>> hfsinfo;
0145     std::vector<float> globalinfo;
0146
0147     // particle masses
0148     static Double_t ElectronMass() { return 0.000511; };
0149     static Double_t ProtonMass()   { return 0.938272; };
0150     static Double_t KaonMass()     { return 0.493677; };
0151     static Double_t PionMass()     { return 0.139570; };
0152     Double_t IonMass;
0153
0154
0155     // lorentz transformations
0156     // - boost from Lab frame `Lvec` to photon+ion C.o.m. frame `Cvec`
0157     void BoostToComFrame(TLorentzVector Lvec, TLorentzVector &Cvec);
0158     // - boost from Lab frame `Lvec` to Ion rest frame `Ivec`
0159     void BoostToIonFrame(TLorentzVector Lvec, TLorentzVector &Ivec);
0160     // - boost from Lab frame `Lvec` to ion+electron Beam c.o.m. frame `Bvec`
0161     void BoostToBeamComFrame(TLorentzVector Lvec, TLorentzVector &Bvec);
0162     // - tranform from Lab frame `Lvec` to Head-on frame `Hvec`
0163     void TransformToHeadOnFrame(TLorentzVector Lvec, TLorentzVector &Hvec);
0164     // transform from Head-on frame `Hvec` back to Lab frame `Lvec`
0165     void TransformBackToLabFrame(TLorentzVector Hvec, TLorentzVector &Lvec);
0166
0167
0168     // misc calculations
0169     // - convert energy,mass to momentum
0170     static Double_t EMtoP(Double_t energy, Double_t mass) {
0171       return TMath::Sqrt( TMath::Power(energy,2) - TMath::Power(mass,2) );
0172     };
0173     // - vector projection: returns vA projected onto vB
0174     static TVector3 Project(TVector3 vA, TVector3 vB) {
0175       if(fabs(vB.Dot(vB))<0.0001) {
0176         //cerr << "WARNING: Kinematics::Project to null vector" << endl;
0177         return TVector3(0,0,0);
0178       };
0179       return vB * ( vA.Dot(vB) / ( vB.Dot(vB) ) );
0180     };
0181     // - vector rejection: returns vC projected onto plane transverse to vD
0182     static TVector3 Reject(TVector3 vC, TVector3 vD) {
0183       if(fabs(vD.Dot(vD))<0.0001) {
0184         //cerr << "WARNING: Kinematics::Reject to null vector" << endl;
0185         return TVector3(0,0,0);
0186       };
0187       return vC - Project(vC,vD);
0188     };
0189     // - calculate angle between two planes, spanned by vectors
0190     static Double_t PlaneAngle(TVector3 vA, TVector3 vB, TVector3 vC, TVector3 vD) {
0191       TVector3 crossAB = vA.Cross(vB); // AxB
0192       TVector3 crossCD = vC.Cross(vD); // CxD
0193       Double_t sgn = crossAB.Dot(vD); // (AxB).D
0194       if(fabs(sgn)<0.00001) {
0195         //cerr << "WARNING: Kinematics:PlaneAngle (AxB).D=0" << endl;
0196         return -10000;
0197       };
0198       sgn /= fabs(sgn); // sign of (AxB).D
0199       Double_t numer = crossAB.Dot(crossCD); // (AxB).(CxD)
0200       Double_t denom = crossAB.Mag() * crossCD.Mag(); // |AxB|*|CxD|
0201       if(fabs(denom)<0.00001) {
0202         //cerr << "WARNING: Kinematics:PlaneAngle denominator=0" << endl;
0203         return -10000;
0204       };
0205       return sgn * TMath::ACos(numer/denom);
0206     };
0207     // - shift angle to the range [-PI,+PI]
0208     static Double_t AdjAngle(Double_t ang) {
0209       while(ang>TMath::Pi()) ang-=2*TMath::Pi();
0210       while(ang<-TMath::Pi()) ang+=2*TMath::Pi();
0211       return ang;
0212     };
0213
0214     // misc. functions for hadronic final state
0215     float correctMass(int pid){
0216       float massOut = 0;
0217       switch(std::abs(pid)){
0218         case 11:
0219           massOut = ElectronMass();
0220           break;
0221         case 2212:
0222           massOut = ProtonMass();
0223           break;
0224         case 321:
0225           massOut = KaonMass();
0226           break;
0227         case 211:
0228           massOut = PionMass();
0229           break;
0230         default:
0231           cerr << "ERROR: unknown pid in Kinematics::correctMass" << endl;
0232           massOut = -1;
0233       };
0234       return massOut;
0235     };
0236
0237     // asymmetry injection
0238     void InjectFakeAsymmetry(); // test your own asymmetry, for fit code validation
0239
0240     // tests and validation
0241     void ValidateHeadOnFrame();
0242
0243     Long64_t countPIDsmeared,countPIDtrue,countHadrons;
0244
0245
0246     // settings
0247     Int_t mainFrame;
0248     enum mainFrame_enum {fLab, fHeadOn};
0249     Int_t qComponentsMethod;
0250     enum qComponentsMethod_enum {qQuadratic, qHadronic, qElectronic};
0251
0252   protected:
0253
0254     // reconstruction methods
0255     void CalculateDISbyElectron();
0256     void CalculateDISbyJB();
0257     void CalculateDISbyDA();
0258     void CalculateDISbyMixed();
0259     void CalculateDISbySigma();
0260     void CalculateDISbyeSigma();
0261     void CalculateDISbyML();
0262     // calculate 4-momenta components of q and W (`vecQ` and `vecW`) as well as
0263     // derived invariants `W` and `nu`
0264     void GetQWNu_electronic();
0265     void GetQWNu_hadronic();
0266     void GetQWNu_quadratic();
0267     void GetQWNu_ML();
0268   private:
0269     static const Int_t asymInjectN = 2;
0270     Double_t moduVal[asymInjectN];
0271     Double_t ampVal[asymInjectN];
0272     Double_t asymInject;
0273     std::unique_ptr<TRandom> RNG;
0274     Float_t RN;
0275     Bool_t reconOK;
0276
0277     // - c.o.m. frame of virtual photon and ion
0278     TLorentzVector CvecBoost;
0279     TVector3 Cboost;
0280     TLorentzVector CvecEleBeam, CvecIonBeam;
0281     TLorentzVector CvecElectron, CvecW, CvecQ;
0282     TLorentzVector CvecHadron;
0283     // - ion rest frame
0284     TLorentzVector IvecBoost;
0285     TVector3 Iboost;
0286     TLorentzVector IvecEleBeam, IvecIonBeam;
0287     TLorentzVector IvecElectron, IvecW, IvecQ;
0288     TLorentzVector IvecHadron;
0289     // - head-on frame
0290     TLorentzVector HvecEleBeam, HvecIonBeam;
0291     TLorentzVector HvecElectron, HvecW, HvecQ;
0292     TLorentzVector HvecHadron;
0293     // - other intermediate frames (for head-on frame transformation)
0294     TLorentzVector BvecBoost, OvecBoost;
0295     TVector3 Bboost, Oboost;
0296     TLorentzVector BvecEleBeam, BvecIonBeam;
0297     Double_t rotAboutX, rotAboutY;
0298     // other
0299     TLorentzVector vecSpin, IvecSpin;
0300 #ifdef SIDIS_MLPRED
0301     py::object keras, tensorflow;
0302     py::object efnpackage;
0303     py::function pfnimport;
0304     py::object model;
0305     py::object modelload;
0306     std::string modelname = "pfn_testEpic_000-2_vecQele_nHFS2_500_bs10k_bestValLoss";
0307 #endif
0308
0309   ClassDef(Kinematics,1);
0310 };