software/JReconstruction/JEnergyCorrection.cc

Go to the documentation of this file.

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

#include "TROOT.h"
#include "TFile.h"
#include "TH1D.h"
#include "TH2D.h"
#include "TF1.h"
#include "TMath.h"
#include "TFitResult.h"

#include "JROOT/JMinimizer.hh"

#include "km3net-dataformat/offline/Head.hh"
#include "km3net-dataformat/offline/Evt.hh"
#include "km3net-dataformat/definitions/reconstruction.hh"
#include "km3net-dataformat/definitions/fitparameters.hh"
#include "JAAnet/JHead.hh"
#include "JAAnet/JHeadToolkit.hh"
#include "JAAnet/JAAnetToolkit.hh"
#include "JDAQ/JDAQEventIO.hh"
#include "JSupport/JTriggeredFileScanner.hh"
#include "JSupport/JMonteCarloFileSupportkit.hh"
#include "JSupport/JSupport.hh"
#include "JSupport/JMeta.hh"
#include "JReconstruction/JEvt.hh"
#include "JReconstruction/JEvtToolkit.hh"
#include "JReconstruction/JEnergyCorrection.hh"
#include "JGeometry3D/JTrack3E.hh"
#include "JPhysics/JGeane.hh"
#include "JTools/JQuantile.hh"
#include "JROOT/JRootToolkit.hh"
#include "JROOT/JManager.hh"
#include "JROOT/JGraph.hh"
#include "JROOT/JGraphErrors.hh"
#include "JGizmo/JGizmoToolkit.hh"
#include "JLang/JToken.hh"

#include "Jeep/JParser.hh"
#include "Jeep/JMessage.hh"

namespace {

  const char* const muon_t      =  "muon";
  const char* const neutrino_t  =  "neutrino";
}

/**
 * \file
 *
 * Utility program to determine energy correction.
 * \author mdejong
 */
int main(int argc, char **argv)
{
  using namespace std;
  using namespace JPP;
  using namespace KM3NETDAQ;

  typedef JTriggeredFileScanner<JEvt>                       JTriggeredFileScanner_t;
  typedef JTriggeredFileScanner_t::multi_pointer_type       multi_pointer_type;
  typedef JToken<';'>      JToken_t;
  typedef JRange<double>   JRange_t;

  JTriggeredFileScanner_t  inputFile;
  JLimit_t&                numberOfEvents = inputFile.getLimit();
  string                   outputFile;  
  JQualitySorter           quality_sorter;
  string                   formula;
  vector<JToken_t>         startValues;
  vector<JToken_t>         fixedValues;
  JRange_t                 X;
  string                   option;
  string                   primary;
  size_t                   numberOfBins;
  int                      debug;

  try {

    JParser<> zap("Utility program to determine energy correction.");

    zap['f'] = make_field(inputFile,   "event file(s) or histogram file");
    zap['o'] = make_field(outputFile)          = "energy_correction.root";
    zap['n'] = make_field(numberOfEvents)      = JLimit::max();
    zap['L'] = make_field(quality_sorter)      = JPARSER::initialised();
    zap['F'] = make_field(formula,     "fit formula, e.g: \"[0]+[1]*x\"")         = "";
    zap['@'] = make_field(startValues, "start values, e.g: \"p0 = GetMaximum;\"") = JPARSER::initialised();
    zap['='] = make_field(fixedValues, "fixed values, e.g: \"p0 = GetMaximum;\"") = JPARSER::initialised();
    zap['x'] = make_field(X,           "abscissa range")                          = JRange_t();
    zap['O'] = make_field(option,      "Fit option")                              = "";
    zap['p'] = make_field(primary)             =  muon_t, neutrino_t;
    zap['N'] = make_field(numberOfBins,"number of bins")                          = 80;
    zap['d'] = make_field(debug)               =  2;

    zap(argc, argv);
  }
  catch(const exception& error) {
    FATAL(error.what() << endl);
  }


  TFile* in = NULL;
  TH2D*  h2 = NULL;

  if (inputFile.size() == 1) {

    in = TFile::Open(inputFile[0].c_str(), "READ");

    if (in != NULL && in->IsOpen()) {
      h2 = (TH2D*) in->Get("h2");
    }
  }

  if (h2 == NULL) {

    vector<JHead> enigma;

    for (JMultipleFileScanner_t::const_iterator i = inputFile.begin(); i != inputFile.end(); ++i) {

      try {

        const JHead head = getHeader(*i);

        vector<JHead>::iterator p = find(enigma.begin(), enigma.end(), head);

        if (p == enigma.end())
          enigma.push_back(head);
        else
          p->add(head);
      }
      catch(const exception& error) {
        FATAL(error.what());
      }
    }

    double Emin =  numeric_limits<double>::max();
    double Emax =  0.0;

    for (vector<JHead>::const_iterator i = enigma.begin(); i!= enigma.end(); ++i) {

      if (!i->is_valid(&JHead::cut_nu)) {
        FATAL("Missing neutrino parameters." << endl);
      }

      if (i->cut_nu.E.getLowerLimit() < Emin) { Emin = i->cut_nu.E.getLowerLimit(); }
      if (i->cut_nu.E.getUpperLimit() > Emax) { Emax = i->cut_nu.E.getUpperLimit(); }

      if (i->is_valid(&JHead::genvol)) {
        cout << "Neutrino energy range:      " << i->cut_nu.E              << endl;
        cout << "Number of generated events: " << i->genvol.numberOfEvents << endl;
      }
    }

    const Double_t xmin  = log10(Emin);
    const Double_t xmax  = log10(Emax);
    const Double_t dx    = (xmax - xmin) / numberOfBins;

    h2 = new TH2D("h2", NULL, numberOfBins + 4, xmin - 2*dx, xmax + 2*dx, numberOfBins, xmin, xmax);

    double      W = 1.0;
    Vec         offset(0.0, 0.0, 0.0);

    for (string filename; inputFile.hasNext(); ) {

      STATUS("event: " << setw(10) << inputFile.getCounter() << '\r'); DEBUG(endl);

      if (filename != inputFile.getFilename()) {

        filename = inputFile.getFilename();

        try {
        
          const JHead head = getHeader(filename);

          offset = getOffset(head);

          vector<JHead>::const_iterator p = find(enigma.begin(), enigma.end(), head);

          if (p->is_valid(&JHead::genvol)) {
            W = 1.0 / (double) p->genvol.numberOfEvents;
          }

          cout << "File: " << filename << " weight " << W << endl;
        }
        catch(const exception& error) {
          FATAL(error.what());
        }
      }

      multi_pointer_type ps = inputFile.next();

      JEvt*  evt   = ps;
      Evt*   event = ps;

      Trk    t1;

      if        (primary == muon_t) {

        for (vector<Trk>::const_iterator i = event->mc_trks.begin(); i != event->mc_trks.end(); ++i) {
          if (is_muon(*i)) {
            if (i->E > t1.E) {
              t1 = *i;
            }
          }
        }

        if (!is_muon(t1)) {
          continue;
        }

      } else if (primary == neutrino_t) {

        if (has_neutrino(*event)) {
          t1 = get_neutrino(*event);
        }

        if (!is_neutrino(t1)) {
          continue;
        }
      }

      JEvt::iterator __end = partition(evt->begin(), evt->end(), JHistory::is_event(JMUONENERGY));

      if (evt->begin() == __end) {
        continue;
      }

      JEvt::iterator p = min_element(evt->begin(), __end, quality_sorter);

      JTrack3E ta = getTrack(t1);
      JTrack3E tb = getTrack(*p);

      ta.add(getPosition(offset));

      tb.setE(p->getW(JENERGY_ENERGY));
      tb.move(tb.getIntersection(ta.getPosition()), getSpeedOfLight(), gWater); 

      if (ta.getE() > 0.0 && tb.getE() > 0.0) {
        h2->Fill(log10(tb.getE()), log10(ta.getE()), (enigma.size() > (size_t) 1 ? W * event->w[2] : 1.0));
      }
    }
    STATUS(endl);
  }

  // short hands

  const Int_t    nx   = h2->GetXaxis()->GetNbins();
  const Double_t xmin = h2->GetXaxis()->GetXmin();
  const Double_t xmax = h2->GetXaxis()->GetXmax();

  const Int_t    ny   = h2->GetYaxis()->GetNbins();

  TH1D h1("h1", NULL, nx, xmin, xmax);

  JGraph_t       g1;
  JGraph_t       g2;
  JGraphErrors_t g3;

  JManager<int, TH1D> H0(new TH1D("py[%]", NULL, nx, xmin, xmax));

  for (Int_t ix = 1; ix <= nx; ++ix) {

    TH1D* h0 = H0[ix];

    JQuantile Q("", true);                          // start values

    for (int iy = 1; iy <= ny; ++iy) {

      const Double_t x = h2->GetYaxis()->GetBinCenter(iy);
      const Double_t y = h2->GetBinContent(ix,iy);
      const Double_t z = h2->GetBinError  (ix,iy);
      
      if (y > 0.0) {

        h0->SetBinContent(iy, y);
        h0->SetBinError  (iy, z);

        Q.put(x,y);
      }
    }

    if (Q.getCount() >= 3) {

      if (debug >= debug_t) {
        Q.print(cout);
      }

      TF1 f1("f1", "[0]*TMath::Gaus(x,[1],[2]) + [3]");
 
      const double x0    =        Q.getQuantile(0.50);
      const double sigma = 0.5 * (Q.getQuantile(0.66) - Q.getQuantile(0.33));

      f1.SetParameter(0, Q.getWmax());
      f1.SetParameter(1, x0);
      f1.SetParameter(2, sigma);
      f1.SetParameter(3, 0.0);

      TFitResultPtr result = h0->Fit(&f1, "SQ", "same");

      if (result.Get() == NULL) {
        FATAL("Invalid TFitResultPtr " << h0->GetName() << endl);
      }

      if (debug >= debug_t || !result->IsValid()) {
        cout << "Bin: "
             << setw(4)    << ix                               << ' '
             << FIXED(7,3) << h2->GetXaxis()->GetBinCenter(ix) << ' '
             << FIXED(7,3) << result->Chi2()                   << '/'
             << result->Ndf()                                  << ' '
             << (result->IsValid() ? "" : "failed")            << endl;
      }

      if (result->IsValid()) {

        const Double_t x  = h2->GetXaxis()->GetBinCenter(ix);

        h1.SetBinContent(ix, f1.GetParameter(1));
        h1.SetBinError  (ix, f1.GetParError (1));

        g1.put(x, pow(10.0, x) - pow(10.0, f1.GetParameter(1)));
        g2.put(x, pow(10.0, abs(f1.GetParameter(2)) - f1.GetParameter(1)));
        g3.put(x, f1.GetParameter(1), 0.0, f1.GetParameter(2));
      }
    }
  }

  TF1* f1 = NULL;

  if (formula != "") {                  // fit

    f1 = new TF1("f1", formula.c_str());

    if (f1->IsZombie()) {
      FATAL("Function: " << formula << " is zombie." << endl);
    }

    try {

      for (vector<JToken_t>::const_iterator i = startValues.begin(); i != startValues.end(); ++i) {
        f1->SetParameter(getParameter(*i), getValue(*i,&h1));
      }

      for (vector<JToken_t>::const_iterator i = fixedValues.begin(); i != fixedValues.end(); ++i) {
        f1->FixParameter(getParameter(*i), getValue(*i,&h1));
      }
    }
    catch(JLANG::JParseError& error) {
      FATAL(error << endl);
    }

    if (option.find('S') == string::npos) { option += "S"; }

    TFitResultPtr result = h1.Fit(f1, option.c_str(), "same", X.getLowerLimit(), X.getUpperLimit());

  } else {                              // interpolate

    Double_t x[2] = { xmin, xmax };
    Double_t y[2] = { xmin, xmax };
    Double_t z[2] = { 0.00, 0.00 };

    ostringstream os;

    os << "   "                         // extrapolate
       << "("
       << "x >= " << FIXED(5,3) << xmax
       << ")"
       << " * "
       << "(x)";
    os << "\n";

    for (Int_t ix = nx; ix >= 1; --ix) {

      x[0] = h1.GetXaxis()->GetBinCenter(ix);
      y[0] = h1.GetBinContent(ix);
      z[0] = h1.GetBinError  (ix);

      if (z[0] == 0.0 || !X(x[0])) {
        continue;
      }

      os << " + "
         << "("
         << "x >= " << FIXED(5,3) << x[0]
         << " && "
         << "x <  " << FIXED(5,3) << x[1]
         << ")"
         << " * "
         << "(" << FIXED(5,3) << y[0] << " + " << FIXED(5,3) << (y[1] - y[0]) << " * (x - " << FIXED(5,3) << x[0] << ") / " << FIXED(5,3) << (x[1] - x[0]) << ")";
      os << "\n";

      x[1] = x[0];
      y[1] = y[0];
      z[1] = z[0];
    }

    x[0] = -1;
    y[0] = -1;

    os << " + "                         // extrapolate
       << "("
       << "x <  " << FIXED(5,3) << x[1]
       << ")"
       << " * "
       << "(" << FIXED(5,3) << y[0] << " + " << FIXED(5,3) << (y[1] - y[0]) << " * (x - " << FIXED(5,3) << x[0] << ") / " << FIXED(5,3) << (x[1] - x[0]) << ")";
    os << "\n";

    string buffer = os.str();

    NOTICE(buffer << endl);

    for (string::iterator i = buffer.begin(); i != buffer.end(); ++i) {
      if (*i == '\n') {
        *i = ' ';
      }
    }

    f1 = new TF1("f1", buffer.c_str(), xmin, xmax);

    h1.GetListOfFunctions()->Add(f1);
  }


  TFile out(outputFile.c_str(), "recreate");

  out << JMeta(argc, argv);

  out << *h2 << H0 << h1;

  out << JGraph      (g1, "g1")
      << JGraph      (g2, "g2")
      << JGraphErrors(g3, "g3");

  if (f1 != NULL) {                     // separately store energy correction
    
    f1->GetFormula()->SetName(JEnergyCorrection::getName());

    out << *(f1->GetFormula());
  }

  out.Write();
  out.Close();

  if (in != NULL) {
    in->Close();
  }
}