WCTrackAlg.cxx

Go to the documentation of this file.

/////////////////////////////////////////////////////////////////////////////////
/// \file    WCTrackAlg.cxx
/// \brief   Algorithm implementation for wire chamber reconstruction.
///          Part of the beamline reconstruction for NOvA test beam.
/// \author  Mike Wallbank (University of Cincinnati) <wallbank@fnal.gov>
///          Fan Gao(University of Pittsburgh) <fag16@pitt.edu>
///          Based on LArIAT code lariatsoft/LArIATRecoAlg/WCTrackBuilderAlg.cxx
///          by Greg Pulliam gkpullia@syr.edu
/// \date    December 2018
/////////////////////////////////////////////////////////////////////////////////

// nova
#include "BeamlineReco/WCTrackAlg.h"

using beamlinegeo::BeamlineComponent;
using beamlinegeo::BeamlineCoordSystem;

namespace beamlinereco {

  WCTrackAlg::WCTrackAlg(const fhicl::ParameterSet& pset) {
    this->reconfigure(pset);
  }

  WCTrackAlg::~WCTrackAlg() { }

  // ------------------------------------------------------------------------------
  void WCTrackAlg::reconfigure(fhicl::ParameterSet const& pset) {

    // Just a number to use to initialize things before they get filled correctly.
    fInitialDefault = -999;

    // These four are parameters from histos I produced from picky-good tracks
    //fCentralYKink         = pset.get<float >("CentralYKink",        -0.01    );
    //fSigmaYKink           = pset.get<float >("SigmaYKink",          0.03      );
    //fCentralYDist         = pset.get<float >("CentralYDist",        0.69      );
    //fSigmaYDist           = pset.get<float >("SigmaYDist",          18.0      );

    // Reconstruction options
    fVerbose     = pset.get<bool>("Verbose", false);
    fPickyTracks = pset.get<bool>("PickyTracks", false);

    // Get coordinate system
    unsigned int coordSystem  = pset.get<unsigned int>("CoordSystem", 0);
    switch (coordSystem) {
    case 0:
      fBeamlineCoordSystem = BeamlineCoordSystem::Beamline;
      break;
    case 1:
      fBeamlineCoordSystem = BeamlineCoordSystem::NOvADetector;
      break;
    default:
      std::cout << "Unknown coordinate system (" << coordSystem << "). "
                << "Options are 0: Beamline; 1: NOvA Detector." << std::endl;
      abort();
    }

    return;

  }

  // ------------------------------------------------------------------------------
  /// Main function called for each event
  void WCTrackAlg::reconstructTracks(std::vector<double> & reco_p_list,
                                     std::vector<TVector3>& mag_int_list,
                                     std::vector<double>& dist_to_mag_axis_list,
                                     std::vector<TVector2>& face_list,
                                     std::vector<std::vector<TVector3> > & wc_hit_pos_list,
                                     std::vector<TVector3> & incoming_dir_list,
                                     std::vector<WCHitList> & event_final_tracks,
                                     std::vector<std::vector<WCHitList> > & good_hits,
                                     std::vector<double> & y_kink_list,
                                     std::vector<TVector3>& dist_list,
                                     int & WCMissed,
                                     //std::vector<TH2F*> & Recodiff,
                                     TH1F* & WCdistribution,
                                     float & residual) {

    WCMissed = fInitialDefault;

    // Determine if one should skip this trigger based on whether there is exactly one good hit in each wire chamber and axis
    // If there isn't, continue after adding a new empty vector to the reco_p_array contianer.
    // If there is, then move on with the function.
    // This can be modified to permit more than one good hit in each wire chamber axis - see comments in function
    bool skip = shouldSkipTrigger(good_hits, WCMissed, WCdistribution);
    if (skip == true)
      return;

    // Assign value to class member variable
    fWCMissed = WCMissed;

    // Depending on if an event has a hit in all 4 WC or whether it missed WC2 or WC3 (but not both),
    // we reconstruct the momentum differently.
    // This code doesn't change from before we allowed 3 point tracks.
    if (fNHits == 4) {

      // At this point, we should have a list of good hits with at least one good hit in X,Y for each WC.
      // Now find all possible combinations of these hits that could form a track, sans consideration
      // of the kinks or end displacements (for now). For the "exactly one" condition set in the above
      // step, this won't matter, but if you want to set the condition to "at least one hit in each WC axis,
      // this will give many combinations.
      residual = buildFourPointTracks(good_hits,
                                      reco_p_list,
                                      mag_int_list,
                                      dist_to_mag_axis_list,
                                      face_list,
                                      wc_hit_pos_list,
                                      incoming_dir_list,
                                      event_final_tracks,
                                      y_kink_list,
                                      dist_list,
                                      WCMissed);

    }

    return;

  }

  // ------------------------------------------------------------------------------
  bool WCTrackAlg::shouldSkipTrigger(std::vector<std::vector<WCHitList> > & good_hits,
                                     int & WCMissed,
                                     TH1F* & WCDist) {

    //Now determine if we want to skip
    bool skip = false;
    fNHits = 0;
    for (size_t iWC = 0; iWC < 4; ++iWC) {
      if (good_hits[iWC][0].hits.size() > 0 && good_hits[iWC][1].hits.size() > 0)
        ++fNHits;
      else
        WCMissed = iWC+1;
    }

    // With WCDist histogram, we can simply estimate the efficiency of each WC
    WCDist->Fill(0);
    if (fNHits > 2)
      WCDist->Fill(1);
    if (fNHits > 2 && WCMissed == 1)
      WCDist->Fill(2);
    if (fNHits > 2 && WCMissed == 2)
      WCDist->Fill(3);
    if (fNHits > 2 && WCMissed == 3)
      WCDist->Fill(4);
    if (fNHits > 2 && WCMissed == 4)
      WCDist->Fill(5);
    if (fNHits == 4)
      WCDist->Fill(6);

    // If we don't have 3 or 4 hits, skip.
    // For commissioning, only require 4 hits first
    // if (fNHits < 3)
    if (fNHits < 4)
      skip = true;

    if (fPickyTracks) {
      for (size_t iWC=0; iWC<4; ++iWC)
        for (size_t iAX=0; iAX<2; ++iAX)
          if (good_hits[iWC][iAX].hits.size() != 1) {
            // Only allow events with 1 and only 1 hit on each axis
            skip = true;
            break;
          }
    }
    else // Skip events with less than 3 X/Y hits or is missing the first or last WC
      if (WCMissed==1 || WCMissed==4)
        skip = true;

    if (skip) {
      if (fVerbose)
        std::cout << "Skipping this event." << std::endl;

      // Clear the hit lists for each WC/axis
      for (size_t iWC = 0; iWC < good_hits.size(); ++iWC)
        for (size_t iAx = 0; iAx < good_hits.at(iWC).size(); ++iAx)
          good_hits.at(iWC).at(iAx).hits.clear();
    }

    return skip;

  }

  // ------------------------------------------------------------------------------
  float WCTrackAlg::calculateRecoP(float theta_x_us, float theta_x_ds, float bestTrackSlope) {
    float magnet_angle = fGeo->MagnetAngle();
    float magnet_eff_length = fGeo->MagnetEffectiveLength();
    float alpha_1 = theta_x_us + magnet_angle*TMath::DegToRad();
    float alpha_2 = theta_x_ds + magnet_angle*TMath::DegToRad();
    float num   = (fabs(fBField) * magnet_eff_length * 0.01 /*cm_to_m*/ * 1000. /*GeV_to_MeV*/ );
    float denom = (3.33564*(sin(alpha_2) - sin(alpha_1)))*cos(atan(bestTrackSlope));
    return num / denom;
  }

  // ------------------------------------------------------------------------------
  float WCTrackAlg::buildFourPointTracks(std::vector<std::vector<WCHitList> > & good_hits,
                                         std::vector<double> & reco_p_list,
                                         std::vector<TVector3>& mag_int_list,
                                         std::vector<double>& dist_to_mag_axis_list,
                                         std::vector<TVector2>& face_list,
                                         std::vector<std::vector<TVector3> > & wc_hit_pos_list,
                                         std::vector<TVector3> & dir_list,
                                         std::vector<WCHitList> & event_final_tracks,
                                         std::vector<double> & y_kink_list,
                                         std::vector<TVector3>& dist_list,
                                         int & WCMissed) {

    std::vector<TVector3> xyz(4, TVector3(0, 0, 0));
    float reco_p=0;
    WCHitList best_track;
    std::vector<float> track_stats;
    std::vector<float> bestRegressionStats;
    float bestResSq = fInitialDefault;
    for (int i = 0; i < 3; i++)
      bestRegressionStats.push_back(fInitialDefault);

    // Loop over all combinations of hits, and find the positions
    for (size_t iHit0 = 0; iHit0 < good_hits[0][0].hits.size(); ++iHit0) {
      for (size_t iHit1 = 0; iHit1 < good_hits[0][1].hits.size(); ++iHit1) {
        for (size_t iHit2 = 0; iHit2 < good_hits[1][0].hits.size(); ++iHit2) {
          for (size_t iHit3 = 0; iHit3 < good_hits[1][1].hits.size(); ++iHit3) {
            for (size_t iHit4 = 0; iHit4 < good_hits[2][0].hits.size(); ++iHit4) {
              for (size_t iHit5 = 0; iHit5 < good_hits[2][1].hits.size(); ++iHit5) {
                for (size_t iHit6 = 0; iHit6 < good_hits[3][0].hits.size(); ++iHit6) {
                  for (size_t iHit7 = 0; iHit7 < good_hits[3][1].hits.size(); ++iHit7) {
                    WCHitList track;
                    track.hits.push_back(good_hits[0][0].hits[iHit0]);
                    track.hits.push_back(good_hits[0][1].hits[iHit1]);
                    track.hits.push_back(good_hits[1][0].hits[iHit2]);
                    track.hits.push_back(good_hits[1][1].hits[iHit3]);
                    track.hits.push_back(good_hits[2][0].hits[iHit4]);
                    track.hits.push_back(good_hits[2][1].hits[iHit5]);
                    track.hits.push_back(good_hits[3][0].hits[iHit6]);
                    track.hits.push_back(good_hits[3][1].hits[iHit7]);
                    findTheHitPositions(track,xyz,WCMissed);
                    // Do regression on the four points to find the straightest track in Y
                    track_stats = Regression(xyz,WCMissed);
                    if (track_stats[2]<fabs(bestResSq)) {
                      best_track=track;
                      bestRegressionStats=track_stats;
                      bestResSq=track_stats[2];
                    }
                  }
                }
              }
            }
          }
        }
      }
    }

    if (bestResSq<12) {

      // Now we should have the straightest track in Y, which will be the track that goes to the event.
      // Now we get the momentum and projections onto the detector
      calculateTheMomentum(best_track,xyz,reco_p, bestRegressionStats, mag_int_list, dist_to_mag_axis_list);

      std::vector<TVector3> wc_hit_pos;
      for (size_t i = 0; i < 4; ++i) {
        wc_hit_pos.push_back(TVector3(xyz[i].X()/10.,
                                      xyz[i].Y()/10.,
                                      xyz[i].Z()/10.));
      }
      wc_hit_pos_list.push_back(wc_hit_pos);
      reco_p_list.push_back(reco_p);

      // We should also have the x,y,z points of the best_track, so now find where it hits the NOvA detector
      projectToDetector(best_track,xyz,face_list,dir_list);
      calculateTrackKink_Dists(xyz,bestRegressionStats,y_kink_list,dist_list);
      event_final_tracks.push_back(best_track);

    }

    return bestResSq;

  }

  // ------------------------------------------------------------------------------
  void WCTrackAlg::findTheHitPositions(WCHitList& track,
                                       std::vector<TVector3>& xyz,
                                       int & WCMissed) {

    std::vector<float> x_wires;
    std::vector<float> y_wires;
    for (int i=0; i<4; ++i) {
      x_wires.push_back(fInitialDefault);
      y_wires.push_back(fInitialDefault);
    }

    for (size_t iHit = 0; iHit < track.hits.size(); ++iHit) {
      int var = iHit % 2;
      int jHit= iHit;
      if (var == 0 and int(jHit != 2*(WCMissed-1)+var)) // Skip iHit for the missed WC
        x_wires[(iHit-var)/2]=(track.hits.at(iHit).wire);
      if (var == 1 and int(jHit != 2*(WCMissed-1)+var))
        y_wires[(iHit-var)/2]=(track.hits.at(iHit).wire);
    }

    std::vector<TVector3> pos = {fGeo->BeamlineComponentPos(BeamlineComponent::WC1, fBeamlineCoordSystem),
                                 fGeo->BeamlineComponentPos(BeamlineComponent::WC2, fBeamlineCoordSystem),
                                 fGeo->BeamlineComponentPos(BeamlineComponent::WC3, fBeamlineCoordSystem),
                                 fGeo->BeamlineComponentPos(BeamlineComponent::WC4, fBeamlineCoordSystem)};

    for (int iWC = 0; iWC < 4; ++iWC) {
      if (iWC != WCMissed-1) {
        // Find the position of the hit, correcting for the rotation of the WC
        // Unit mm
        xyz[iWC] = TVector3(pos[iWC].X()*10. + x_wires[iWC] * TMath::Cos(TMath::DegToRad()*fGeo->WCAngle(iWC)),
                            pos[iWC].Y()*10. + y_wires[iWC],
                            pos[iWC].Z()*10. + x_wires[iWC] * TMath::Sin(TMath::DegToRad()*fGeo->WCAngle(iWC)));
      }
      if (iWC == WCMissed-1)
        // Just making sure a value is set, should never be used
        xyz[iWC] = TVector3(fInitialDefault, fInitialDefault, fInitialDefault);
    }

    return;

  }

  // ------------------------------------------------------------------------------
  std::vector<float> WCTrackAlg::Regression(const std::vector<TVector3>& xyz,
                                            int & WCMissed) {

    // Calculate the linear regression on yz plane, since the track deflects in x-axis
    std::vector<float> RegressionValues;
    int Npoints = 0;
    float sum_z = 0;
    float sum_zz = 0;
    float sum_y = 0;
    float sum_yz = 0;
    float intercept = 0;
    float slope = 0;
    float residualsquare = 0;
    float residual = 0;
    for (int i = 0; i < 4; ++i) {
      if (i != WCMissed-1) { //turning WC# to array index
        sum_y  += xyz[i].Y();
        sum_zz += xyz[i].Z()*xyz[i].Z();
        sum_z  += xyz[i].Z();
        sum_yz += xyz[i].Y()*xyz[i].Z();
        ++Npoints; //Residual is a function of number of points used.
      }
    }

    slope = (Npoints*sum_yz-sum_y*sum_z)/(Npoints*sum_zz-sum_z*sum_z);
    intercept = (sum_y-slope*sum_z)/Npoints;

    for (int i = 0; i < 4; ++i) {
      if (i != WCMissed-1) {
        residual = (xyz[i].Y()-slope*xyz[i].Z()-intercept)/std::sqrt(1+slope*slope);
        residualsquare += (residual)*(residual);
      }
    }
    float avgresidual = std::sqrt(residualsquare)/(Npoints-2);

    RegressionValues.push_back(slope);
    RegressionValues.push_back(intercept);
    RegressionValues.push_back(avgresidual);

    return RegressionValues;

  }

  // ------------------------------------------------------------------------------
  void WCTrackAlg::calculateTheMomentum(WCHitList& best_track,
                                        std::vector<TVector3>& xyz,
                                        float & reco_p,
                                        std::vector<float> & BestTrackStats,
                                        std::vector<TVector3>& mag_int_list,
                                        std::vector<double> & dist_to_mag_axis_list) {

    // We need the x,y,z again, which would be for the last track combination tried.
    // So we need to find the positions again
    findTheHitPositions(best_track,xyz,fWCMissed);

    // Calculate the angle of the track, in the x,z and y,z planes, in upstream(us) and downstream(ds) ends of the WC
    float dx_us = xyz[1].X() - xyz[0].X();
    float dx_ds = xyz[3].X() - xyz[2].X();
    float dz_us = xyz[1].Z() - xyz[0].Z();
    float dz_ds = xyz[3].Z() - xyz[2].Z();
    float theta_x_us = atan(dx_us/dz_us);
    float theta_x_ds = atan(dx_ds/dz_ds);

    float unscaled_reco_p = calculateRecoP(theta_x_us,theta_x_ds,BestTrackStats[0]);

    // project track to magnet face to find where it enters relative to the central axis
    float dist_to_axis = projectToMagnetFace(best_track,xyz,mag_int_list,dist_to_mag_axis_list);

    // scale momentum based on where particle enters magnet
    float a =  1.00068819;
    float b = -0.04635663;
    float c =  0.11454936;
    float d =  0.00971984;
    float e = -0.11526033;
    float f =  0.01645662;    
    float x = dist_to_axis/100; // convert cm to m
    float scale_factor = (a+b*x)*(1/(1+exp((x-c)/d)) - 1/(1+exp((x-e)/f)));
    reco_p = scale_factor * unscaled_reco_p;

    std::cout << "B: " << fBField << " momentum: " << reco_p << std::endl;

  }

  // ------------------------------------------------------------------------------
  float WCTrackAlg::projectToMagnetFace(WCHitList & best_track,
                                       const std::vector<TVector3>& xyz,
                                       std::vector<TVector3>& mag_int_list,
                                       std::vector<double>& dist_to_mag_axis_list
                                       ) {

    // Get magnet geometry - need center of magnet, length, and angle to find equation of front face
    float mag_pos_x = fGeo->BeamlineComponentPos(BeamlineComponent::Magnet, fBeamlineCoordSystem).X();
    float mag_pos_z = fGeo->BeamlineComponentPos(BeamlineComponent::Magnet, fBeamlineCoordSystem).Z();
    float magnet_angle = fGeo->MagnetAngle();
    float mag_length = fGeo->MagnetEffectiveLength();

    // center point of front face of magnet:
    float mag_front_x = mag_pos_x+mag_length/2*sin(magnet_angle*TMath::DegToRad());
    float mag_front_z = mag_pos_z-mag_length/2*cos(magnet_angle*TMath::DegToRad());

    // WC track intercept with front face of magnet
    float dxdz_us = (xyz[1].X() - xyz[0].X())/(xyz[1].Z() - xyz[0].Z());
    float dydz_us = (xyz[1].Y() - xyz[0].Y())/(xyz[1].Z() - xyz[0].Z());
    float tanang = tan((90-magnet_angle)*TMath::DegToRad());
    float zint_front = (xyz[1].X()/10 - dxdz_us*xyz[1].Z()/10 + tanang*mag_front_z-mag_front_x)/(tanang-dxdz_us);
    float xint_front = tanang*zint_front - tanang*mag_front_z+mag_front_x;
    float yint_front = dydz_us*(zint_front-xyz[1].Z()/10)+xyz[1].Y()/10;

    // Transverse distance from point of intercept to central axis of magnet
    int sign_of_dist = 1;
    if (xint_front < mag_front_x)
      sign_of_dist = -1;
    float dist_to_mag_axis = sign_of_dist*sqrt(pow(xint_front-mag_front_x,2)+pow(zint_front-mag_front_z,2));

    mag_int_list.push_back(TVector3(xint_front,yint_front,zint_front));
    dist_to_mag_axis_list.push_back(dist_to_mag_axis);
    
    return dist_to_mag_axis;

  }

  // ------------------------------------------------------------------------------
  void WCTrackAlg::projectToDetector(WCHitList & best_track,
                                     const std::vector<TVector3>& xyz,
                                     std::vector<TVector2>& face_list,
                                     std::vector<TVector3>& incoming_dir_list) {

    // Direction of track
    float dx_ds = xyz[3].X() - xyz[2].X();
    float dy_ds = xyz[3].Y() - xyz[2].Y();
    float dz_ds = xyz[3].Z() - xyz[2].Z();
    TVector3 vec(dx_ds, dy_ds, dz_ds);
    incoming_dir_list.push_back(vec.Unit());

    // Get interseption with the front face of the detector
    TVector3 nova_face = fGeo->BeamlineComponentPos(BeamlineComponent::NOvA,
                                                    fBeamlineCoordSystem);
    TVector3 normal(0, 0, 1), point(0, 0, nova_face.Z()*10.);
    TVector3 start = xyz[3];
    float p = (normal*(point-start))/(normal*vec);
    TVector3 end_int = start + (p*vec);
    face_list.push_back(TVector2(end_int.X()/10., end_int.Y()/10.));

    return;

  }

  // ------------------------------------------------------------------------------
  void WCTrackAlg::calculateTrackKink_Dists(const std::vector<TVector3>& xyz,
                                            std::vector<float>& track_stats,
                                            std::vector<double>& y_kink_list,
                                            std::vector<TVector3>& dist_list) {

    // Determine the equations of the lines, x=mz+b and y=mz+b, for the US and DS legs
    float dx_us = xyz[1].X() - xyz[0].X();
    float dy_us = xyz[1].Y() - xyz[0].Y();
    float dz_us = xyz[1].Z() - xyz[0].Z();
    float dx_ds = xyz[3].X() - xyz[2].X();
    float dy_ds = xyz[3].Y() - xyz[2].Y();
    float dz_ds = xyz[3].Z() - xyz[2].Z();
    float x_us_slope = dx_us/dz_us;
    float y_us_slope = dy_us/dz_us;
    float x_ds_slope = dx_ds/dz_ds;
    float y_ds_slope = dy_ds/dz_ds;
    float x_us_int = xyz[1].X() - x_us_slope*xyz[1].Z();
    float y_us_int = xyz[1].Y() - y_us_slope*xyz[1].Z();
    float x_ds_int = xyz[3].X() - x_ds_slope*xyz[3].Z();
    float y_ds_int = xyz[3].Y() - y_ds_slope*xyz[3].Z();

    // First Y_kink, the angle difference between the us and ds legs.
    y_kink_list.push_back(atan(y_ds_slope)-atan(y_us_slope));

    // Because we need a way to compare tracks, regardless of current setting, we have to have a standard for the midplane.
    // We will use the normal midplane (fMP_X=fMP_M=1)

    float magnet_angle = fGeo->MagnetAngle()*TMath::DegToRad();
    float magnet_midplane = fGeo->MagnetMidplaneIntercept()*10.;

    // Get US intersection with midplane
    float z_mp_us = (magnet_midplane-x_us_int)/(x_us_slope-1/(tan(magnet_angle)));
    float x_mp_us = x_us_slope*z_mp_us+x_us_int;
    float y_mp_us = y_us_slope*z_mp_us+y_us_int;

    // Get DS intersection with midplane
    float z_mp_ds = (magnet_midplane-x_ds_int)/(x_ds_slope-1/(tan(magnet_angle)));
    float x_mp_ds = x_ds_slope*z_mp_ds+x_ds_int;
    float y_mp_ds = y_ds_slope*z_mp_ds+y_ds_int;

    dist_list.push_back(TVector3((x_mp_ds-x_mp_us)/10.,
                                 (y_mp_ds-y_mp_us)/10.,
                                 (z_mp_ds-z_mp_us)/10.));

    return;

  }

  // ------------------------------------------------------------------------------
  void WCTrackAlg::setMagneticField(float b_field) {
    fBField = b_field;
    return;
  }

}