LightFlavorRatios/corrections/swimLF.C

0001 #include <TCut.h>
0002 #include <TFile.h>
0003 #include <TMatrixD.h>
0004 #include <TString.h>
0005 #include <TTree.h>
0006
0007 #include <kfparticle_sphenix/KFParticle_sPHENIX.h>
0008 #include <kfparticle_sphenix/KFParticle_Tools.h>
0009
0010 #include <iostream>
0011 #include <limits>
0012
0013 R__LOAD_LIBRARY(libKFParticle.so)
0014
0015 #include <KFParticle.h>
0016
0017 using namespace std;
0018 KFParticle_Tools kfpTupleTools;
0019
0020 bool verbose = false;
0021
0022 //Cut values we want to probe
0023 float min_DIRA = 0.99;
0024 float min_IP_xy = 0.05;
0025 float min_decayLength = 0.05; //Is this a bias or a cut...
0026 vector<float> v_tpVal_TAU;
0027 vector<float> v_tpVal_FD;
0028
0029 void reset()
0030 {
0031   v_tpVal_TAU.clear();
0032   v_tpVal_FD.clear();
0033 }
0034
0035 KFParticle makeParticle(float pos[3], float mom[3], float cov[21], char charge, float mass, float chi2, float ndof)
0036 {
0037   KFParticle kfp_particle;
0038
0039   float f_trackParameters[6] = {pos[0], pos[1], pos[2], mom[0], mom[1], mom[2]};
0040
0041   kfp_particle.Create(f_trackParameters, cov, (int) charge, mass);
0042   kfp_particle.Chi2() = chi2;
0043   kfp_particle.NDF() = ndof;
0044
0045   return kfp_particle;
0046 }
0047
0048 KFParticle makeVertex(float pv_pos[3], float pv_cov[6], float pv_chi2, unsigned int pv_ndof)
0049 {
0050   float f_vertexParameters[6] = {pv_pos[0], pv_pos[1], pv_pos[2], 0, 0, 0};
0051
0052   float f_vertexCovariance[21] = {0};
0053   for (int i = 0; i < 6; ++i) f_vertexCovariance[i] = pv_cov[i];
0054
0055   KFParticle kfp_vertex;
0056   kfp_vertex.Create(f_vertexParameters, f_vertexCovariance, 0, -1);
0057   kfp_vertex.Chi2() = pv_chi2;
0058   kfp_vertex.NDF() = pv_ndof;
0059
0060   return kfp_vertex;
0061 }
0062
0063 bool calculateValues(KFParticle mother, KFParticle track_1, KFParticle track_2, KFParticle PV, std::array<float,3>& new_pv)
0064 {
0065   //Note, this is still missing a check as to whether the track would've hit an MVTX sensor or not
0066
0067   KFParticle moved_PV = PV;
0068   moved_PV.X() = new_pv[0];
0069   moved_PV.Y() = new_pv[1];
0070   moved_PV.Z() = new_pv[2];
0071
0072   float track_1_IP_XY = track_1.GetDistanceFromVertexXY(moved_PV);
0073   float track_2_IP_XY = track_2.GetDistanceFromVertexXY(moved_PV);
0074
0075   float DIRA = kfpTupleTools.eventDIRA(mother, moved_PV, true);
0076
0077   float decayTime, decayTimeErr, decayLength, decayLengthErr;
0078
0079   KFParticle particleCopy = mother;
0080   particleCopy.SetProductionVertex(moved_PV);
0081   particleCopy.TransportToDecayVertex();
0082   particleCopy.GetLifeTime(decayTime, decayTimeErr);
0083   particleCopy.GetDecayLengthXY(decayLength, decayLengthErr);
0084   const float speed = 2.99792458e-2;
0085   decayTime /= speed;
0086
0087   if(verbose)
0088   {
0089     cout << "New DIRA is " << DIRA << endl;
0090     cout << "New track 1 IP XY is " << track_1_IP_XY << endl;
0091     cout << "New track 2 IP XY is " << track_2_IP_XY << endl;
0092     cout << "New decay length is " << decayLength << endl;
0093     cout << "New decay time is " << decayTime << endl;
0094   }
0095
0096   //Could probably do some push back into a vector then use std::all_of to check if they're all true
0097
0098   bool acceptDIRA = DIRA >= min_DIRA ? true : false;
0099   bool acceptIPXY = min(abs(track_1_IP_XY), abs(track_2_IP_XY)) >= min_IP_xy ? true : false;
0100   bool acceptDecayLength = decayLength >= min_decayLength ? true : false;
0101
0102   // Must decay before it leaves the MVTX; this is a crude check on that quality
0103   bool acceptnMVTX = (decayLength <= 5.);
0104
0105   return acceptDIRA && acceptIPXY && acceptDecayLength && acceptnMVTX;
0106 }
0107
0108 std::array<float,3> movePV(KFParticle PV, KFParticle mother, float displace_dist)
0109 {
0110   float mother_p = sqrt(pow(mother.GetPx(),2)+pow(mother.GetPy(),2)+pow(mother.GetPz(),2));
0111   std::array<float,3> mother_p_unit = {mother.GetPx()/mother_p, mother.GetPy()/mother_p, mother.GetPz()/mother_p};
0112
0113   std::array<float,3> new_pv;
0114   std::array<float,3> original_pv = {PV.GetX(), PV.GetY(), PV.GetZ()};
0115   for (int k = 0; k < 3; ++k) new_pv[k] = original_pv[k] + mother_p_unit[k]*displace_dist;
0116
0117   return new_pv;
0118 }
0119
0120 // Moves vertex from (PV position + displace_low) to (PV position + displace_high) in increments of displace_step
0121 // Returns range of displacement values between which there are turning points
0122 std::vector<std::pair<float,float>> swimPV(KFParticle mother, std::vector<KFParticle> tracks, KFParticle PV,
0123                                            float displace_low, float displace_high, float displace_step)
0124 {
0125   std::vector<std::pair<float,float>> tp_PV_pos;
0126
0127   int nSteps = ceil((displace_high-displace_low)/displace_step);
0128
0129   bool lastWasAccepted = false;
0130
0131   for(int i=0; i<=nSteps; i++)
0132   {
0133     if(verbose) cout << "\nSwimming step " << i << endl;
0134
0135     float displace_dist = displace_low + i*displace_step;
0136
0137     std::array<float,3> new_pv = movePV(PV,mother,displace_dist);
0138
0139     if(verbose) cout << "PV is now at (" << new_pv[0] << ", " << new_pv[1] << ", " << new_pv[2] << ")" << endl;
0140     bool didWeAccept = calculateValues(mother, tracks[0], tracks[1], PV, new_pv);
0141     string result = didWeAccept ? "We would accept the candidate at this PV" : "We would reject the candidate at this PV";
0142     if(i>0)
0143     {
0144       if(didWeAccept != lastWasAccepted)
0145       {
0146         if(verbose)
0147         {
0148           cout << "Turning point found:" << std::endl;
0149           cout << "step " << i-1 << " -> " << i << std::endl;
0150           cout << "displacement distance " << displace_low + (i-1)*displace_step << " -> " << displace_dist << endl;
0151         }
0152         tp_PV_pos.push_back(std::make_pair(displace_low+(i-1)*displace_step,displace_dist));
0153       }
0154     }
0155     lastWasAccepted = didWeAccept;
0156     if(verbose) cout << result << endl;
0157   }
0158
0159   return tp_PV_pos;
0160 }
0161
0162 //int main(int argc, char* argv[])
0163 int swimLF(int nCandidates, std::string fileName, std::string motherName)
0164 {
0165 /*
0166   if (argc != 3)
0167   {
0168       cerr << "Usage: " << argv[0] << " <input file> <mother name>\n";
0169       return 1;
0170   }
0171 */
0172   KFParticle::SetField(15.);// Mag field in kiloGuass (is this negative, should print out the number in KFParticle...)
0173
0174   float large_step_size = 0.5; //Make large steps first to figure out where the turning points roughly are
0175   float small_step_size = 0.01; //Our PV resolution should be better than 100 microns but maybe not by much. Probably good for a first attempt
0176
0177   //string fileName = argv[1];
0178   //string motherName = argv[2];
0179
0180   TFile* inFile = new TFile(fileName.c_str());
0181   TTree* inTree = (TTree*) inFile->Get("DecayTree");
0182
0183   string outputName = fileName.substr(0,fileName.size() - 5);
0184   outputName += "_swum.root";
0185   TFile* outFile = new TFile(outputName.c_str(),"RECREATE");
0186
0187   TCut cut = "";
0188
0189   TTree* midTree = inTree->CopyTree(cut);
0190   TTree* outTree = midTree->CloneTree(nCandidates); //copying and cloning the input tree
0191   //TTree* outTree = midTree->CloneTree(-1); //copying and cloning the input tree
0192
0193   float mother_pos[3], mother_mom[3], mother_cov[21], mother_chi2;
0194   unsigned int mother_ndof;
0195   char mother_charge;
0196   float mother_mass;
0197
0198   float track_1_pos[3], track_1_mom[3], track_1_cov[21], track_1_chi2;
0199   unsigned int track_1_ndof;
0200   char track_1_charge;
0201   float track_1_mass;
0202
0203   float track_2_pos[3], track_2_mom[3], track_2_cov[21], track_2_chi2;
0204   unsigned int track_2_ndof;
0205   char track_2_charge;
0206   float track_2_mass;
0207
0208   float pv_pos[3], pv_cov[6], pv_chi2;
0209   unsigned int pv_ndof;
0210
0211   string xyz[3] = {"x", "y", "z"};
0212   for (int i = 0; i < 3; ++i)
0213   {
0214     outTree->SetBranchAddress(TString(motherName) + "_"  + TString(xyz[i]), &mother_pos[i]);
0215     outTree->SetBranchAddress(TString(motherName) + "_p" + TString(xyz[i]), &mother_mom[i]);
0216     outTree->SetBranchAddress("track_1_"  + TString(xyz[i]), &track_1_pos[i]);
0217     outTree->SetBranchAddress("track_1_p" + TString(xyz[i]), &track_1_mom[i]);
0218     outTree->SetBranchAddress("track_2_"  + TString(xyz[i]), &track_2_pos[i]);
0219     outTree->SetBranchAddress("track_2_p" + TString(xyz[i]), &track_2_mom[i]);
0220     outTree->SetBranchAddress("primary_vertex_" + TString(xyz[i]), &pv_pos[i]);
0221   }
0222
0223   outTree->SetBranchAddress(TString(motherName) + "_mass", &mother_mass);
0224   outTree->SetBranchAddress(TString(motherName) + "_Covariance", &mother_cov);
0225   outTree->SetBranchAddress(TString(motherName) + "_charge", &mother_charge);
0226   outTree->SetBranchAddress(TString(motherName) + "_chi2", &mother_chi2);
0227   outTree->SetBranchAddress(TString(motherName) + "_nDoF", &mother_ndof);
0228
0229   outTree->SetBranchAddress("track_1_Covariance", &track_1_cov);
0230   outTree->SetBranchAddress("track_1_charge", &track_1_charge);
0231   outTree->SetBranchAddress("track_1_mass", &track_1_mass);
0232   outTree->SetBranchAddress("track_1_chi2", &track_1_chi2);
0233   outTree->SetBranchAddress("track_1_nDoF", &track_1_ndof);
0234
0235   outTree->SetBranchAddress("track_2_Covariance", &track_2_cov);
0236   outTree->SetBranchAddress("track_2_charge", &track_2_charge);
0237   outTree->SetBranchAddress("track_2_mass", &track_2_mass);
0238   outTree->SetBranchAddress("track_2_chi2", &track_2_chi2);
0239   outTree->SetBranchAddress("track_2_nDoF", &track_2_ndof);
0240
0241   outTree->SetBranchAddress("primary_vertex_Covariance", &pv_cov);
0242   outTree->SetBranchAddress("primary_vertex_chi2", &pv_chi2);
0243   outTree->SetBranchAddress("primary_vertex_nDoF", &pv_ndof);
0244
0245   //This is just for checking values
0246   float original_DIRA; outTree->SetBranchAddress(TString(motherName) + "_DIRA", &original_DIRA);
0247   float original_decayLength; outTree->SetBranchAddress(TString(motherName) + "_decayLength", &original_decayLength);
0248   float original_decayTime; outTree->SetBranchAddress(TString(motherName) + "_decayTime", &original_decayTime);
0249   float original_track_1_IP_xy; outTree->SetBranchAddress("track_1_IP_xy", &original_track_1_IP_xy);
0250   float original_track_2_IP_xy; outTree->SetBranchAddress("track_2_IP_xy", &original_track_2_IP_xy);
0251
0252   TBranch *turningPoints_TAU = outTree->Branch("turningPoints_TAU", &v_tpVal_TAU);
0253   TBranch *turningPoints_FD  = outTree->Branch("turningPoints_FD",   &v_tpVal_FD);
0254
0255   int nEntries = outTree->GetEntries();
0256   for (int i = 0; i < nEntries; ++i)
0257   {
0258     reset();
0259
0260     if(i % 1000 == 0) cout << "Candidate " << i << endl;
0261     outTree->GetEntry(i);
0262
0263     KFParticle mother = makeParticle(mother_pos, mother_mom, mother_cov, mother_charge, mother_mass, mother_chi2, mother_ndof);
0264     KFParticle PV = makeVertex(pv_pos, pv_cov, pv_chi2, pv_ndof);
0265     std::vector<KFParticle> tracks;
0266     tracks.push_back(makeParticle(track_1_pos, track_1_mom, track_1_cov, track_1_charge, track_1_mass, track_1_chi2, track_1_ndof));
0267     tracks.push_back(makeParticle(track_2_pos, track_2_mom, track_2_cov, track_2_charge, track_2_mass, track_2_chi2, track_2_ndof));
0268
0269     mother.SetProductionVertex(PV);
0270     mother.TransportToDecayVertex();
0271
0272     if(verbose)
0273     {
0274
0275       cout << "PV information" << endl; kfpTupleTools.identify(PV);
0276       cout << "Mother information" << endl; kfpTupleTools.identify(mother);
0277       cout << "Track 1 information" << endl; kfpTupleTools.identify(tracks[0]);
0278       cout << "Track 2 information" << endl; kfpTupleTools.identify(tracks[1]);
0279
0280       cout << "Beginning coarse swimming" << endl;
0281     }
0282
0283     float original_decaylengthXY, original_decaylengthXYerr;
0284     float original_pt, original_pterr;
0285     mother.GetDecayLengthXY(original_decaylengthXY,original_decaylengthXYerr);
0286
0287     float mother_pt, mother_pterr;
0288     mother.GetPt(mother_pt,mother_pterr);
0289
0290     // swim so that 2D decay length goes down to zero and up to 10cm
0291     // we swim along the mother momentum vector, so the distance is affected by pz
0292     // ds^2 = dR^2 + dz^2
0293     // dz/dR = pz/pt, so ds^2 = dR^2(1+pz^2/pt^2)
0294     if(verbose) std::cout << "original decaylengthXY: " << original_decaylengthXY << std::endl;
0295     float coarse_swimdistance_min = -(10.-original_decaylengthXY)*sqrt(1.+pow(mother.GetPz()/mother.GetPt(),2.));
0296     float coarse_swimdistance_max = original_decaylengthXY*sqrt(1.+pow(mother.GetPz()/mother.GetPt(),2.));
0297     if(verbose)
0298     {
0299       std::cout << "min_displacement: " << coarse_swimdistance_min << std::endl;
0300       std::cout << "max_displacement: " << coarse_swimdistance_max << std::endl;
0301     }
0302
0303     std::vector<std::pair<float,float>> turning_points = swimPV(mother,tracks,PV,coarse_swimdistance_min,coarse_swimdistance_max,large_step_size);
0304     if(verbose) cout << "Beginning fine swimming" << endl;
0305     std::vector<float> turningpoint_displacement;
0306     for(std::pair<float,float> tp : turning_points)
0307     {
0308       std::vector<std::pair<float,float>> fine_turning_points = swimPV(mother,tracks,PV,tp.first,tp.second,small_step_size);
0309       for(std::pair<float,float> fine_tp : fine_turning_points)
0310       {
0311         turningpoint_displacement.push_back((fine_tp.first+fine_tp.second)/2.);
0312       }
0313     }
0314
0315     for(float disp : turningpoint_displacement)
0316     {
0317       std::array<float,3> new_pv = movePV(PV,mother,disp);
0318       KFParticle moved_PV = PV;
0319       moved_PV.X() = new_pv[0];
0320       moved_PV.Y() = new_pv[1];
0321       moved_PV.Z() = new_pv[2];
0322       KFParticle particleCopy = mother;
0323       particleCopy.SetProductionVertex(moved_PV);
0324       particleCopy.TransportToDecayVertex();
0325       float decayTime, decayTimeErr, decayLength, decayLengthErr;
0326       particleCopy.GetLifeTime(decayTime, decayTimeErr);
0327       particleCopy.GetDecayLengthXY(decayLength, decayLengthErr);
0328       const float speed = 2.99792458e-2;
0329       decayTime /= speed;
0330       if(verbose) std::cout << "turning point decay time " << decayTime << " decay length " << decayLength << std::endl;
0331       v_tpVal_TAU.push_back(decayTime);
0332       v_tpVal_FD.push_back(decayLength);
0333     }
0334
0335     turningPoints_TAU->Fill();
0336     turningPoints_FD->Fill();
0337
0338
0339 /*
0340     KFParticle track_1 = makeParticle(track_1_pos, track_1_mom, track_1_cov, track_1_charge, track_1_chi2, track_1_ndof);
0341     KFParticle track_2 = makeParticle(track_2_pos, track_2_mom, track_2_cov, track_2_charge, track_2_chi2, track_2_ndof);
0342     KFParticle PV = makeVertex(pv_pos, pv_cov, pv_chi2, pv_ndof);
0343
0344     cout << "PV information" << endl; kfpTupleTools.identify(PV);
0345     cout << "Mother information" << endl; kfpTupleTools.identify(mother);
0346     cout << "Track 1 information" << endl; kfpTupleTools.identify(track_1);
0347     cout << "Track 2 information" << endl; kfpTupleTools.identify(track_2);
0348
0349     float total_mother_momentum = sqrt(pow(mother_mom[0], 2) + pow(mother_mom[1], 2) + pow(mother_mom[2], 2));
0350     float mom_unit_vector[3]; for (int j = 0; j < 3; ++j) mom_unit_vector[j] = mother_mom[j]/total_mother_momentum; //Need a unit vector to swim along
0351
0352     //Super quick swimming
0353     //lets just bring us back and forward by 5cm
0354     float jump_back = 5, jump_forward = jump_back;
0355     float start_position[3], end_position[3];
0356     for (int j = 0; j < 3; ++j)
0357     {
0358       start_position[j] = pv_pos[j] - mom_unit_vector[j]*jump_back;
0359       end_position[j] = pv_pos[j] + mom_unit_vector[j]*jump_forward;
0360     }
0361
0362     //Check some values
0363     cout << "PV (x, y, z) =  (" << PV.GetX() << ", " << PV.GetY() << ", " << PV.GetZ() << ")" << endl;
0364     cout << "Start pos (x, y, z) =  (" << start_position[0] << ", " << start_position[1] << ", " << start_position[2] << ")" << endl;
0365     cout << "End pos (x, y, z) =  (" << end_position[0] << ", " << end_position[1] << ", " << end_position[2] << ")" << endl;
0366     cout << "PV radius = " << sqrt(pow(PV.GetX(), 2) + pow(PV.GetY(), 2));
0367     cout << ", Start radius = " << sqrt(pow(start_position[0], 2) + pow(start_position[1], 2));
0368     cout << ", End radius = " << sqrt(pow(end_position[0], 2) + pow(end_position[1], 2)) << endl;
0369     cout << "Unit vector value = (" << mom_unit_vector[0] << ", " << mom_unit_vector[1] << ", " << mom_unit_vector[2] << "), size: " << sqrt(pow(mom_unit_vector[0], 2) + pow(mom_unit_vector[1], 2) + pow(mom_unit_vector[2], 2)) << endl;
0370
0371     //Check original values
0372     cout << "Original DIRA was " << original_DIRA << endl;
0373     cout << "Original track 1 IP XY was " << original_track_1_IP_xy << endl;
0374     cout << "Original track 2 IP XY was " << original_track_2_IP_xy << endl;
0375     cout << "Original decay length was " << original_decayLength << endl;
0376     cout << "Original decay time was " << original_decayTime << endl;
0377
0378     int nBackSteps = (int) sqrt(pow(start_position[0], 2) + pow(start_position[1], 2) + pow(start_position[2], 2))/large_step_size;
0379     int nForwardSteps = (int) sqrt(pow(end_position[0], 2) + pow(end_position[1], 2) + pow(end_position[2], 2))/large_step_size;
0380
0381     std::vector<std::pair<int,int>> coarse_tp_index;
0382
0383     bool lastWasAccepted = false;
0384
0385     for (int j = -1*nBackSteps; j < nForwardSteps; ++j)
0386     {
0387       cout << "\nSwimming step " << j << endl;
0388
0389       float new_pv[3], original_pv[3] = {PV.GetX(), PV.GetY(), PV.GetZ()};
0390       for (int k = 0; k < 3; ++k) new_pv[k] = original_pv[k] + j*mom_unit_vector[k]*large_step_size;
0391       cout << "PV is now at (" << new_pv[0] << ", " << new_pv[1] << ", " << new_pv[2] << ")" << endl;
0392       bool didWeAccept = calculateValues(mother, track_1, track_2, PV, new_pv);
0393       string result = didWeAccept ? "We would accept the candidate at this PV" : "We would reject the candidate at this PV";
0394       if(j == -1*nBackSteps)
0395       {
0396         lastWasAccepted = didWeAccept;
0397       }
0398       else
0399       {
0400         if(didWeAccept != lastWasAccepted)
0401         {
0402           cout << "Turning point found:" << std::endl;
0403           cout << "step " << j-1 << " -> " << j << std::endl;
0404           coarse_tp_index.push_back(std::make_pair(j-1,j));
0405         }
0406       }
0407       lastWasAccepted = didWeAccept;
0408       cout << result << endl;
0409
0410       //Once we find the transition for the acceptance, we need to jump back and do fine grained steps to get a more precise number
0411       //This value for the decay length and decay time should be stored into the vectors we'll write to the nTuple and sorted into some readable order for turn on and turn off
0412     }
0413 */
0414
0415   }
0416
0417   outFile->Write();
0418   outFile->Close();
0419
0420   return 0;
0421 }