make_dst_cosrejbdttrain.C

Go to the documentation of this file.

// Dan's P. heritage

#include "CAFAna/Core/SAMQuerySource.h"
#include "3FlavorAna/Vars/NumuVars.h"

#include "StandardRecord/StandardRecord.h"

#include "TFile.h"
#include "TTree.h"
#include "TVector3.h"
#include "TH1.h"

#include <vector>
#include <string>
#include <iostream>

using namespace ana;

// Macro for stripping only most interesting info out of a set of CAF files
// Intended for making a dst for training a CosRej MVA.  This macro will
// make these files light enough to run a standard TMVA training script on
void make_dst_cosrejbdttrain(std::string fname, std::string out)
{

  // Needed for deteriming cosZ
  TVector3 beamDir = NuMIBeamDirection(caf::kFARDET).Unit();

  // New test vars
  
  float inel, hpp, wshw, lidcos;
  int nshw;
  float starty, stopy, cosang, maxy, efrac, nplfr, lid;

  // Basic size variables
  float nhit,calE;
  // PID variables
  float cvne, remid, cvnsse, cvnprelec, cvnprmuon, lidmulll, lidelll, lidemullt, lidemulll;
  // Basic reco variables
  float prongl,sparseassym,costhetaZ,ptp,pxp,pyp,vtxx,vtxy,vtxz, xyhit_assym;
  // Save distances to all walls, and dist to any wall but the top
  float disttop,distbottom,distfront,distback,disteast,distwest,distnottop;

  // Needed to specifiy signal / bkg and for making plots after training
  // with osc weights so that we can reasonably gauge performance
  float woscdumb;
  int pdg,pdgorig;
  bool isGENIE, isCC;

  // ROOT complains about needing the output file opened before makeing
  // TTree
  TFile* outFile = TFile::Open(out.c_str(), "recreate");

  // Need three things - a TTree for sig/bkg, POT, and livetime
  TH1F *pot = new TH1F("pot","",1,0,1);
  TH1F *livetime = new TH1F("livetime","",1,0,1);
  TTree *sigTree = new TTree("sigTree","");
  TTree *bkgTree = new TTree("bkgTree","");

  // Set branch address of all the variables we're saving
  //general
  sigTree->Branch("nhit",        &nhit);
  sigTree->Branch("hitsperpl",   &hpp);
  sigTree->Branch("calE",        &calE);
  sigTree->Branch("inelast",     &inel);
  sigTree->Branch("costheta",    &lidcos);
  sigTree->Branch("cosang",      &cosang);
  sigTree->Branch("costhetaZ",   &costhetaZ);
  sigTree->Branch("numofshow",   &nshw);
  sigTree->Branch("widthofshow", &wshw);
  sigTree->Branch("efrac",       &efrac);  
  sigTree->Branch("xyhitassym",  &xyhit_assym);
  sigTree->Branch("prongl",      &prongl);
  sigTree->Branch("costhetaZ",   &costhetaZ);
  // pid vars
  sigTree->Branch("cvne",        &cvne);
  sigTree->Branch("remid",       &remid);
  sigTree->Branch("lid",         &lid);
  sigTree->Branch("cvnsse",      &cvnsse);    
  sigTree->Branch("cvnprelec",   &cvnprelec); 
  sigTree->Branch("cvnprmuon",   &cvnprmuon); 
  sigTree->Branch("lidmulll",    &lidmulll); 
  sigTree->Branch("lidelll",     &lidelll);
  sigTree->Branch("lidemullt",   &lidemullt);
  sigTree->Branch("lidemulll",   &lidemulll);
  //cosrej vars
  sigTree->Branch("sparseassym", &sparseassym);
  sigTree->Branch("ptp",         &ptp);
  sigTree->Branch("pxp",         &pxp);
  sigTree->Branch("pyp",         &pyp);
  //geometry 
  sigTree->Branch("starty",      &starty);
  sigTree->Branch("stopy",       &stopy);
  sigTree->Branch("maxy",        &maxy);
  sigTree->Branch("nplfr",       &nplfr);
  sigTree->Branch("vtxx",        &vtxx);
  sigTree->Branch("vtxy",        &vtxy);
  sigTree->Branch("vtxz",        &vtxz);
  sigTree->Branch("disttop",     &disttop);
  sigTree->Branch("distbottom",  &distbottom);
  sigTree->Branch("distfront",   &distfront);
  sigTree->Branch("distback",    &distback);
  sigTree->Branch("disteast",    &disteast);
  sigTree->Branch("distwest",    &distwest);
  sigTree->Branch("distnottop",  &distnottop);

  sigTree->Branch("pdg",         &pdg);
  sigTree->Branch("pdgorig",     &pdgorig);
  sigTree->Branch("woscdumb",    &woscdumb);
  sigTree->Branch("isGENIE",     &isGENIE);
  sigTree->Branch("isCC",        &isCC);

  // And for bkg
  // general 
  bkgTree->Branch("nhit",        &nhit);
  bkgTree->Branch("hitsperpl",   &hpp);
  bkgTree->Branch("calE",        &calE);
  bkgTree->Branch("inelast",     &inel);
  bkgTree->Branch("costheta",    &lidcos);
  bkgTree->Branch("cosang",      &cosang);
  bkgTree->Branch("costhetaZ",   &costhetaZ);
  bkgTree->Branch("numofshow",   &nshw);
  bkgTree->Branch("widthofshow", &wshw);
  bkgTree->Branch("efrac",       &efrac);
  bkgTree->Branch("xyhitassym",  &xyhit_assym);
  bkgTree->Branch("prongl",      &prongl);
  bkgTree->Branch("costhetaZ",   &costhetaZ);
  // pid vars
  bkgTree->Branch("cvne",        &cvne);
  bkgTree->Branch("remid",       &remid);
  bkgTree->Branch("lid",         &lid);
  bkgTree->Branch("cvnsse",      &cvnsse);    
  bkgTree->Branch("cvnprelec",   &cvnprelec);
  bkgTree->Branch("cvnprmuon",   &cvnprmuon);
  bkgTree->Branch("lidmulll",    &lidmulll);
  bkgTree->Branch("lidelll",     &lidelll); 
  bkgTree->Branch("lidemullt",   &lidemullt);
  bkgTree->Branch("lidemulll",   &lidemulll);
  //cosrej vars
  bkgTree->Branch("sparseassym", &sparseassym);
  bkgTree->Branch("ptp",         &ptp);
  bkgTree->Branch("pxp",         &pxp);
  bkgTree->Branch("pyp",         &pyp);
  //geometry
  bkgTree->Branch("starty",      &starty);
  bkgTree->Branch("stopy",       &stopy);
  bkgTree->Branch("maxy",        &maxy);
  bkgTree->Branch("nplfr",       &nplfr);
  bkgTree->Branch("vtxx",        &vtxx);
  bkgTree->Branch("vtxy",        &vtxy);
  bkgTree->Branch("vtxz",        &vtxz);
  bkgTree->Branch("disttop",     &disttop);
  bkgTree->Branch("distbottom",  &distbottom);
  bkgTree->Branch("distfront",   &distfront);
  bkgTree->Branch("distback",    &distback);
  bkgTree->Branch("disteast",    &disteast);
  bkgTree->Branch("distwest",    &distwest);
  bkgTree->Branch("distnottop",  &distnottop);

  bkgTree->Branch("pdg",         &pdg);
  bkgTree->Branch("pdgorig",     &pdgorig);
  bkgTree->Branch("woscdumb",    &woscdumb);
  bkgTree->Branch("isGENIE",     &isGENIE);
  bkgTree->Branch("isCC",        &isCC);


  // Need a way to run over a set of files on the grid through SAM.
  // Use CAFAna's way of dealing with this so that it is maintained
  int stride = getenv("CAFANA_STRIDE") ? atoi(getenv("CAFANA_STRIDE")) : -1;
  int offset = getenv("CAFANA_OFFSET") ? atoi(getenv("CAFANA_OFFSET")) : -1;
  SAMQuerySource fSource(fname, stride, offset, -1);
  std::cout << "Running over dataset with stride: " << stride << std::endl;
  std::cout << "Running over dataset with offset: " << offset << std::endl;

  // Ask for the next file
  while (TFile *inFile = fSource.GetNextFile()){

    // We have a file, let's get fill our TTree

    // Grab POT / livetime
    TH1* curpot = (TH1*)inFile->Get("TotalPOT");
    pot->Fill(0.5,curpot->Integral()); // Only nonzero for MC, where we don't
                                       // care about spill cuts
    isGENIE = curpot->Integral() > 0;
    delete curpot;
    TH1* curlivetime = (TH1*)inFile->Get("TotalLivetime");
    livetime->Fill(0.5,curlivetime->Integral()); // Only nonzero for cosmics
    delete curlivetime;

    // Grap recTree from current file
    TTree *recTree = (TTree*)inFile->Get("recTree");
    caf::StandardRecord* sr = 0;
    recTree->SetBranchAddress("rec", &sr);

    // Run over the TTree's entries
    for(int i = 0; i<recTree->GetEntries(); i++) {
    //  if (i > 10000) break;
      recTree->GetEntry(i);

      // Apply some preselection - dominated by the CVN cut
      if (sr->sel.cvnProd3Train.nueid < 0.5) // new cvn -ok
        continue; // Quickly cut lost of stuff
      if (!sr->sel.veto.keep)
        continue;      // Quickly cut lots of stuff
      if ( !sr->vtx.elastic.IsValid )
        continue;   // Need Vertex for prongs
     // if (sr->slc.nhit < 20 || sr->slc.nhit > 150)
     //   continue;      // Require event reasonably sized

      // Grab the variables you want before storing them in a TTree

      inel = 1.f;
      if (sr->vtx.elastic.IsValid ){
        if (sr->vtx.elastic.fuzzyk.nshwlid >= 1) {
          inel = (sr->slc.calE - sr->vtx.elastic.fuzzyk.png[0].shwlid.calE)/(sr->slc.calE);
        }
      }

      hpp = sr->sel.nuecosrej.hitsperplane;

      nshw = 0u;
      if (sr->vtx.elastic.IsValid) { nshw = sr->vtx.elastic.fuzzyk.nshwlid; }
      
      wshw = -1000.;
      if (sr->vtx.elastic.IsValid ) {
        if (sr->vtx.elastic.fuzzyk.nshwlid >= 1) {
          wshw = sr->vtx.elastic.fuzzyk.png[0].shwlid.width;
        }
      } 

      lidcos = -1000.;
      if (sr->vtx.elastic.IsValid) {
        if (sr->vtx.elastic.fuzzyk.nshwlid >= 1) {
          lidcos = sr->vtx.elastic.fuzzyk.png[0].shwlid.lid.costheta;
        }
      } 
     
      starty = -1000.;
      if (sr->vtx.elastic.IsValid ) {
        if (sr->vtx.elastic.fuzzyk.nshwlid >= 1) {
          starty = sr->vtx.elastic.fuzzyk.png[0].shwlid.start.y;
        }
      }

      stopy = -1000.;
      if (sr->vtx.elastic.IsValid) {
        if (sr->vtx.elastic.fuzzyk.nshwlid >= 1) {
          stopy = sr->vtx.elastic.fuzzyk.png[0].shwlid.stop.y;
        }
      }

     cosang = sr->sel.nuecosrej.cosdang;

     efrac = -1000.;
     if (sr->vtx.elastic.IsValid) {
        if (sr->vtx.elastic.fuzzyk.nshwlid >= 1) {
          efrac = sr->vtx.elastic.fuzzyk.png[0].shwlid.lid.shwEFrac;
        }
      }

     nplfr = sr->sel.contain.nplanestofront;

     maxy = -1000.0f;
     if (sr->vtx.elastic.IsValid) {
       if (sr->vtx.elastic.fuzzyk.nshwlid != 0) {
         float maxyall = -999.0;
         for (unsigned int i=0;i<sr->vtx.elastic.fuzzyk.nshwlid;i++){
            maxyall = std::max(std::max(sr->vtx.elastic.fuzzyk.png[i].shwlid.start.Y(),sr->vtx.elastic.fuzzyk.png[i].shwlid.stop.Y()),maxyall);
            maxy = maxyall;
         }
       }
     }

      // Energy estimators
      nhit        = sr->slc.nhit;
      calE        = sr->slc.calE;

      // Various PID's
      cvne        = sr->sel.cvn2017.nueid;
      cvnsse      = sr->sel.cvnProd3Train.nueid;
      remid       = (sr->trk.kalman.ntracks>0) ? float(sr->sel.remid.pid) : float(-1);
      lid         = sr->sel.lid.ann;
      cvnprelec   = -1000.;
      if (sr->vtx.elastic.IsValid ) {
        if (sr->vtx.elastic.fuzzyk.npng >= 1) {
          cvnprelec = sr->vtx.elastic.fuzzyk.png[0].cvnpart.electronid;
        }
      }
      cvnprmuon   = -1000.;
      if (sr->vtx.elastic.IsValid ) {
        if (sr->vtx.elastic.fuzzyk.npng >= 1) {
          cvnprmuon = sr->vtx.elastic.fuzzyk.png[0].cvnpart.muonid;
        }
      }
      lidmulll    = -1000.;
      if (sr->vtx.elastic.IsValid ) {
        if (sr->vtx.elastic.fuzzyk.nshwlid >= 1) {
          lidmulll = sr->vtx.elastic.fuzzyk.png[0].shwlid.lid.mulll;
        }
      }
      lidelll     = -1000.;
      if (sr->vtx.elastic.IsValid ) {
        if (sr->vtx.elastic.fuzzyk.nshwlid >= 1) {
          lidelll = sr->vtx.elastic.fuzzyk.png[0].shwlid.lid.elll;
        }
      }
      lidemullt   = -1000.;
      if (sr->vtx.elastic.IsValid ) {
        if (sr->vtx.elastic.fuzzyk.nshwlid >= 1) {
          lidemullt = sr->vtx.elastic.fuzzyk.png[0].shwlid.lid.emullt;
        }
      }
      lidemulll   = -1000.;
      if (sr->vtx.elastic.IsValid ) {
        if (sr->vtx.elastic.fuzzyk.nshwlid >= 1) {
          lidemulll = sr->vtx.elastic.fuzzyk.png[0].shwlid.lid.emulll;
        }
      }
      // Some reco vars that can help pid / reject back
      prongl      = -1;
      if (sr->vtx.elastic.IsValid ){
        if (sr->vtx.elastic.fuzzyk.npng > 0){
          prongl = sr->vtx.elastic.fuzzyk.png[0].len;
        }
      }
      sparseassym = sr->sel.nuecosrej.sparsenessasymm;
      xyhit_assym = -1000.0;
      if (sr->vtx.elastic.IsValid ) {
        if (sr->vtx.elastic.fuzzyk.nshwlid >= 1) {
          float xyhitdiff = abs(sr->vtx.elastic.fuzzyk.png[0].shwlid.nhitx - 
                                sr->vtx.elastic.fuzzyk.png[0].shwlid.nhity);
          float xyhitsum  = (sr->vtx.elastic.fuzzyk.png[0].shwlid.nhitx + 
                                 sr->vtx.elastic.fuzzyk.png[0].shwlid.nhity);
          if(xyhitsum>0) xyhit_assym =  xyhitdiff/xyhitsum;
        }
      }

      costhetaZ   = -2;
      if (sr->vtx.elastic.IsValid ){
        if (sr->vtx.elastic.fuzzyk.npng > 0){
          TVector3 dir(sr->vtx.elastic.fuzzyk.png[0].dir);
          costhetaZ = dir.Dot(beamDir);
        }
      }
      ptp         = sr->sel.nuecosrej.photptp;
      pxp         = sr->sel.nuecosrej.photpxp;
      pyp         = sr->sel.nuecosrej.photpyp;

      vtxx        = -1e3;
      vtxy        = -1e3;
      vtxz        = -1e3;
      if (sr->vtx.elastic.IsValid ){
        TVector3 vtx = TVector3(sr->vtx.elastic.vtx);
        vtxx = vtx[0];
        vtxy = vtx[1];
        vtxz = vtx[2];
      }

      // Containment variables
      disttop     = sr->sel.nuecosrej.distallpngtop;
      distbottom  = sr->sel.nuecosrej.distallpngbottom;
      distwest    = sr->sel.nuecosrej.distallpngwest;
      disteast    = sr->sel.nuecosrej.distallpngeast;
      distfront   = sr->sel.nuecosrej.distallpngfront;
      distback    = sr->sel.nuecosrej.distallpngback;
      // A way to do this
      distnottop = distbottom;
      if (distwest  < distnottop) distnottop = distwest;
      if (disteast  < distnottop) distnottop = disteast;
      if (distfront < distnottop) distnottop = distfront;
      if (distback  < distnottop) distnottop = distback;

      // Require an MCNeutrino to fill truth
      // Don't continue if no neutrino - cosmic slices have no neutrino
      if (sr->mc.nnu != 0){
        pdg      = sr->mc.nu[0].pdg;
        pdgorig  = sr->mc.nu[0].pdgorig;
        isCC     = sr->mc.nu[0].iscc;
        woscdumb = sr->mc.nu[0].woscdumb;
      }

      if(isGENIE && isCC && std::abs(pdgorig)==14 && std::abs(pdg)==12) // it will train numu->nue CC (FHC mode) vs everything else
        sigTree->Fill();
      else
        bkgTree->Fill();
    } // run over recTree

  }


  // Go back to your output file and save your TTree
  outFile->cd();
  sigTree->Write();
  bkgTree->Write();
  pot->Write();
  livetime->Write();
  outFile->Close();
}