LightFlavorRatios/swimming_correction/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_ACCEPT;
0028 vector<float> v_tpVal_FD;
0029
0030 void reset()
0031 {
0032   v_tpVal_TAU.clear();
0033   v_tpVal_ACCEPT.clear();
0034   v_tpVal_FD.clear();
0035 }
0036
0037 KFParticle makeParticle(float pos[3], float mom[3], float cov[21], char charge, float mass, float chi2, float ndof)
0038 {
0039   KFParticle kfp_particle;
0040
0041   float f_trackParameters[6] = {pos[0], pos[1], pos[2], mom[0], mom[1], mom[2]};
0042
0043   kfp_particle.Create(f_trackParameters, cov, (int) charge, mass);
0044   kfp_particle.Chi2() = chi2;
0045   kfp_particle.NDF() = ndof;
0046
0047   return kfp_particle;
0048 }
0049
0050 KFParticle makeVertex(float pv_pos[3], float pv_cov[6], float pv_chi2, unsigned int pv_ndof)
0051 {
0052   float f_vertexParameters[6] = {pv_pos[0], pv_pos[1], pv_pos[2], 0, 0, 0};
0053
0054   float f_vertexCovariance[21] = {0};
0055   for (int i = 0; i < 6; ++i) f_vertexCovariance[i] = pv_cov[i];
0056
0057   KFParticle kfp_vertex;
0058   kfp_vertex.Create(f_vertexParameters, f_vertexCovariance, 0, -1);
0059   kfp_vertex.Chi2() = pv_chi2;
0060   kfp_vertex.NDF() = pv_ndof;
0061
0062   return kfp_vertex;
0063 }
0064
0065 bool calculateValues(KFParticle mother, KFParticle track_1, KFParticle track_2, KFParticle PV, std::array<float,3>& new_pv)
0066 {
0067   //Note, this is still missing a check as to whether the track would've hit an MVTX sensor or not
0068
0069   KFParticle moved_PV = PV;
0070   moved_PV.X() = new_pv[0];
0071   moved_PV.Y() = new_pv[1];
0072   moved_PV.Z() = new_pv[2];
0073
0074   float track_1_IP_XY = track_1.GetDistanceFromVertexXY(moved_PV);
0075   float track_2_IP_XY = track_2.GetDistanceFromVertexXY(moved_PV);
0076
0077   float DIRA = kfpTupleTools.eventDIRA(mother, moved_PV, true);
0078
0079   float decayTime, decayTimeErr, decayLength, decayLengthErr;
0080
0081   KFParticle particleCopy = mother;
0082   particleCopy.SetProductionVertex(moved_PV);
0083   particleCopy.TransportToDecayVertex();
0084   particleCopy.GetLifeTime(decayTime, decayTimeErr);
0085   particleCopy.GetDecayLengthXY(decayLength, decayLengthErr);
0086   const float speed = 2.99792458e-2;
0087   decayTime /= speed;
0088
0089   if(verbose)
0090   {
0091     cout << "New DIRA is " << DIRA << endl;
0092     cout << "New track 1 IP XY is " << track_1_IP_XY << endl;
0093     cout << "New track 2 IP XY is " << track_2_IP_XY << endl;
0094     cout << "New decay length is " << decayLength << endl;
0095     cout << "New decay time is " << decayTime << endl;
0096   }
0097
0098   //Could probably do some push back into a vector then use std::all_of to check if they're all true
0099
0100   bool acceptDIRA = DIRA >= min_DIRA ? true : false;
0101   bool acceptIPXY = min(abs(track_1_IP_XY), abs(track_2_IP_XY)) >= min_IP_xy ? true : false;
0102   bool acceptDecayLength = decayLength >= min_decayLength ? true : false;
0103
0104   return acceptDIRA && acceptIPXY && acceptDecayLength;
0105 }
0106
0107 std::array<float,3> movePV(KFParticle PV, KFParticle mother, float displace_dist)
0108 {
0109   float mother_p = sqrt(pow(mother.GetPx(),2)+pow(mother.GetPy(),2)+pow(mother.GetPz(),2));
0110   std::array<float,3> mother_p_unit = {mother.GetPx()/mother_p, mother.GetPy()/mother_p, mother.GetPz()/mother_p};
0111
0112   std::array<float,3> new_pv;
0113   std::array<float,3> original_pv = {PV.GetX(), PV.GetY(), PV.GetZ()};
0114   for (int k = 0; k < 3; ++k) new_pv[k] = original_pv[k] + mother_p_unit[k]*displace_dist;
0115
0116   return new_pv;
0117 }
0118
0119 // Moves vertex from (PV position + displace_low) to (PV position + displace_high) in increments of displace_step
0120 // Returns range of displacement values between which there are turning points
0121 std::vector<std::pair<float,float>> swimPV(KFParticle mother, std::vector<KFParticle> tracks, KFParticle PV,
0122                                            float displace_low, float displace_high, float displace_step)
0123 {
0124   std::vector<std::pair<float,float>> tp_PV_pos;
0125
0126   int nSteps = ceil((displace_high-displace_low)/displace_step);
0127
0128   bool lastWasAccepted = false;
0129
0130   for(int i=0; i<=nSteps; i++)
0131   {
0132     if(verbose) cout << "\nSwimming step " << i << endl;
0133
0134     float displace_dist = displace_low + i*displace_step;
0135
0136     std::array<float,3> new_pv = movePV(PV,mother,displace_dist);
0137
0138     if(verbose) cout << "PV is now at (" << new_pv[0] << ", " << new_pv[1] << ", " << new_pv[2] << ")" << endl;
0139     bool didWeAccept = calculateValues(mother, tracks[0], tracks[1], PV, new_pv);
0140     string result = didWeAccept ? "We would accept the candidate at this PV" : "We would reject the candidate at this PV";
0141     if(i>0)
0142     {
0143       if(didWeAccept != lastWasAccepted)
0144       {
0145         if(verbose)
0146         {
0147           cout << "Turning point found:" << std::endl;
0148           cout << "step " << i-1 << " -> " << i << std::endl;
0149           cout << "displacement distance " << displace_low + (i-1)*displace_step << " -> " << displace_dist << endl;
0150         }
0151         tp_PV_pos.push_back(std::make_pair(displace_low+(i-1)*displace_step,displace_dist));
0152       }
0153     }
0154     lastWasAccepted = didWeAccept;
0155     if(verbose) cout << result << endl;
0156   }
0157
0158   return tp_PV_pos;
0159 }
0160
0161 //int main(int argc, char* argv[])
0162 int swimLF(int nCandidates, std::string fileName, std::string motherName)
0163 {
0164 /*
0165   if (argc != 3)
0166   {
0167       cerr << "Usage: " << argv[0] << " <input file> <mother name>\n";
0168       return 1;
0169   }
0170 */
0171   KFParticle::SetField(14.);// Mag field in kiloGuass (is this negative, should print out the number in KFParticle...)
0172
0173   float large_step_size = 0.5; //Make large steps first to figure out where the turning points roughly are
0174   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
0175
0176   //string fileName = argv[1];
0177   //string motherName = argv[2];
0178
0179   TFile* inFile = new TFile(fileName.c_str());
0180   TTree* inTree = (TTree*) inFile->Get("DecayTree");
0181
0182   string outputName = fileName.substr(0,fileName.size() - 5);
0183   outputName += "_swum.root";
0184   TFile* outFile = new TFile(outputName.c_str(),"RECREATE");
0185
0186   TCut cut = "";
0187
0188   TTree* midTree = inTree->CopyTree(cut);
0189   TTree* outTree = midTree->CloneTree(nCandidates); //copying and cloning the input tree
0190   //TTree* outTree = midTree->CloneTree(-1); //copying and cloning the input tree
0191
0192   float mother_pos[3], mother_mom[3], mother_cov[21], mother_chi2;
0193   unsigned int mother_ndof;
0194   char mother_charge;
0195   float mother_mass;
0196
0197   float track_1_pos[3], track_1_mom[3], track_1_cov[21], track_1_chi2;
0198   unsigned int track_1_ndof;
0199   char track_1_charge;
0200   float track_1_mass;
0201
0202   float track_2_pos[3], track_2_mom[3], track_2_cov[21], track_2_chi2;
0203   unsigned int track_2_ndof;
0204   char track_2_charge;
0205   float track_2_mass;
0206
0207   float pv_pos[3], pv_cov[6], pv_chi2;
0208   unsigned int pv_ndof;
0209
0210   string xyz[3] = {"x", "y", "z"};
0211   for (int i = 0; i < 3; ++i)
0212   {
0213     outTree->SetBranchAddress(TString(motherName) + "_"  + TString(xyz[i]), &mother_pos[i]);
0214     outTree->SetBranchAddress(TString(motherName) + "_p" + TString(xyz[i]), &mother_mom[i]);
0215     outTree->SetBranchAddress("track_1_"  + TString(xyz[i]), &track_1_pos[i]);
0216     outTree->SetBranchAddress("track_1_p" + TString(xyz[i]), &track_1_mom[i]);
0217     outTree->SetBranchAddress("track_2_"  + TString(xyz[i]), &track_2_pos[i]);
0218     outTree->SetBranchAddress("track_2_p" + TString(xyz[i]), &track_2_mom[i]);
0219     outTree->SetBranchAddress("primary_vertex_" + TString(xyz[i]), &pv_pos[i]);
0220   }
0221
0222   outTree->SetBranchAddress(TString(motherName) + "_mass", &mother_mass);
0223   outTree->SetBranchAddress(TString(motherName) + "_Covariance", &mother_cov);
0224   outTree->SetBranchAddress(TString(motherName) + "_charge", &mother_charge);
0225   outTree->SetBranchAddress(TString(motherName) + "_chi2", &mother_chi2);
0226   outTree->SetBranchAddress(TString(motherName) + "_nDoF", &mother_ndof);
0227
0228   outTree->SetBranchAddress("track_1_Covariance", &track_1_cov);
0229   outTree->SetBranchAddress("track_1_charge", &track_1_charge);
0230   outTree->SetBranchAddress("track_1_mass", &track_1_mass);
0231   outTree->SetBranchAddress("track_1_chi2", &track_1_chi2);
0232   outTree->SetBranchAddress("track_1_nDoF", &track_1_ndof);
0233
0234   outTree->SetBranchAddress("track_2_Covariance", &track_2_cov);
0235   outTree->SetBranchAddress("track_2_charge", &track_2_charge);
0236   outTree->SetBranchAddress("track_2_mass", &track_2_mass);
0237   outTree->SetBranchAddress("track_2_chi2", &track_2_chi2);
0238   outTree->SetBranchAddress("track_2_nDoF", &track_2_ndof);
0239
0240   outTree->SetBranchAddress("primary_vertex_Covariance", &pv_cov);
0241   outTree->SetBranchAddress("primary_vertex_chi2", &pv_chi2);
0242   outTree->SetBranchAddress("primary_vertex_nDoF", &pv_ndof);
0243
0244   //This is just for checking values
0245   float original_DIRA; outTree->SetBranchAddress(TString(motherName) + "_DIRA", &original_DIRA);
0246   float original_decayLength; outTree->SetBranchAddress(TString(motherName) + "_decayLength", &original_decayLength);
0247   float original_decayTime; outTree->SetBranchAddress(TString(motherName) + "_decayTime", &original_decayTime);
0248   float original_track_1_IP_xy; outTree->SetBranchAddress("track_1_IP_xy", &original_track_1_IP_xy);
0249   float original_track_2_IP_xy; outTree->SetBranchAddress("track_2_IP_xy", &original_track_2_IP_xy);
0250
0251   TBranch *turningPoints_TAU = outTree->Branch("turningPoints_TAU", &v_tpVal_TAU);
0252   TBranch *turningPoints_ACCEPT = outTree->Branch("turningPoints_ACCEPT",&v_tpVal_ACCEPT);
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       float DIRA = kfpTupleTools.eventDIRA(mother, moved_PV, true);
0329       float track_1_IP_XY = tracks[0].GetDistanceFromVertexXY(moved_PV);
0330       float track_2_IP_XY = tracks[1].GetDistanceFromVertexXY(moved_PV);
0331       const float speed = 2.99792458e-2;
0332       decayTime /= speed;
0333       if(verbose) std::cout << "turning point decay time " << decayTime << " decay length " << decayLength << std::endl;
0334
0335       bool acceptDIRA = DIRA >= min_DIRA;
0336       bool acceptIPXY = min(abs(track_1_IP_XY), abs(track_2_IP_XY)) >= min_IP_xy;
0337       bool acceptDecayLength = decayLength >= min_decayLength;
0338
0339       v_tpVal_TAU.push_back(decayTime);
0340       v_tpVal_ACCEPT.push_back((acceptDIRA && acceptIPXY && acceptDecayLength)? 1:0);
0341       v_tpVal_FD.push_back(decayLength);
0342     }
0343
0344     turningPoints_TAU->Fill();
0345     turningPoints_ACCEPT->Fill();
0346     turningPoints_FD->Fill();
0347
0348
0349 /*
0350     KFParticle track_1 = makeParticle(track_1_pos, track_1_mom, track_1_cov, track_1_charge, track_1_chi2, track_1_ndof);
0351     KFParticle track_2 = makeParticle(track_2_pos, track_2_mom, track_2_cov, track_2_charge, track_2_chi2, track_2_ndof);
0352     KFParticle PV = makeVertex(pv_pos, pv_cov, pv_chi2, pv_ndof);
0353
0354     cout << "PV information" << endl; kfpTupleTools.identify(PV);
0355     cout << "Mother information" << endl; kfpTupleTools.identify(mother);
0356     cout << "Track 1 information" << endl; kfpTupleTools.identify(track_1);
0357     cout << "Track 2 information" << endl; kfpTupleTools.identify(track_2);
0358
0359     float total_mother_momentum = sqrt(pow(mother_mom[0], 2) + pow(mother_mom[1], 2) + pow(mother_mom[2], 2));
0360     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
0361
0362     //Super quick swimming
0363     //lets just bring us back and forward by 5cm
0364     float jump_back = 5, jump_forward = jump_back;
0365     float start_position[3], end_position[3];
0366     for (int j = 0; j < 3; ++j)
0367     {
0368       start_position[j] = pv_pos[j] - mom_unit_vector[j]*jump_back;
0369       end_position[j] = pv_pos[j] + mom_unit_vector[j]*jump_forward;
0370     }
0371
0372     //Check some values
0373     cout << "PV (x, y, z) =  (" << PV.GetX() << ", " << PV.GetY() << ", " << PV.GetZ() << ")" << endl;
0374     cout << "Start pos (x, y, z) =  (" << start_position[0] << ", " << start_position[1] << ", " << start_position[2] << ")" << endl;
0375     cout << "End pos (x, y, z) =  (" << end_position[0] << ", " << end_position[1] << ", " << end_position[2] << ")" << endl;
0376     cout << "PV radius = " << sqrt(pow(PV.GetX(), 2) + pow(PV.GetY(), 2));
0377     cout << ", Start radius = " << sqrt(pow(start_position[0], 2) + pow(start_position[1], 2));
0378     cout << ", End radius = " << sqrt(pow(end_position[0], 2) + pow(end_position[1], 2)) << endl;
0379     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;
0380
0381     //Check original values
0382     cout << "Original DIRA was " << original_DIRA << endl;
0383     cout << "Original track 1 IP XY was " << original_track_1_IP_xy << endl;
0384     cout << "Original track 2 IP XY was " << original_track_2_IP_xy << endl;
0385     cout << "Original decay length was " << original_decayLength << endl;
0386     cout << "Original decay time was " << original_decayTime << endl;
0387
0388     int nBackSteps = (int) sqrt(pow(start_position[0], 2) + pow(start_position[1], 2) + pow(start_position[2], 2))/large_step_size;
0389     int nForwardSteps = (int) sqrt(pow(end_position[0], 2) + pow(end_position[1], 2) + pow(end_position[2], 2))/large_step_size;
0390
0391     std::vector<std::pair<int,int>> coarse_tp_index;
0392
0393     bool lastWasAccepted = false;
0394
0395     for (int j = -1*nBackSteps; j < nForwardSteps; ++j)
0396     {
0397       cout << "\nSwimming step " << j << endl;
0398
0399       float new_pv[3], original_pv[3] = {PV.GetX(), PV.GetY(), PV.GetZ()};
0400       for (int k = 0; k < 3; ++k) new_pv[k] = original_pv[k] + j*mom_unit_vector[k]*large_step_size;
0401       cout << "PV is now at (" << new_pv[0] << ", " << new_pv[1] << ", " << new_pv[2] << ")" << endl;
0402       bool didWeAccept = calculateValues(mother, track_1, track_2, PV, new_pv);
0403       string result = didWeAccept ? "We would accept the candidate at this PV" : "We would reject the candidate at this PV";
0404       if(j == -1*nBackSteps)
0405       {
0406         lastWasAccepted = didWeAccept;
0407       }
0408       else
0409       {
0410         if(didWeAccept != lastWasAccepted)
0411         {
0412           cout << "Turning point found:" << std::endl;
0413           cout << "step " << j-1 << " -> " << j << std::endl;
0414           coarse_tp_index.push_back(std::make_pair(j-1,j));
0415         }
0416       }
0417       lastWasAccepted = didWeAccept;
0418       cout << result << endl;
0419
0420       //Once we find the transition for the acceptance, we need to jump back and do fine grained steps to get a more precise number
0421       //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
0422     }
0423 */
0424
0425   }
0426
0427   outFile->Write();
0428   outFile->Close();
0429
0430   return 0;
0431 }