JetEvaluator.cc

Go to the documentation of this file. Or view the newest version in sPHENIX GitHub for file JetEvaluator.cc

#include "JetEvaluator.h"

#include "JetEvalStack.h"
#include "JetRecoEval.h"

#include <jetbase/Jet.h>
#include <jetbase/JetContainer.h>

#include <fun4all/Fun4AllReturnCodes.h>
#include <fun4all/SubsysReco.h>

#include <phool/getClass.h>
#include <phool/phool.h>

#include <TFile.h>
#include <TNtuple.h>

#include <cmath>
#include <cstdlib>
#include <iostream>
#include <map>
#include <utility>

JetEvaluator::JetEvaluator(const std::string &name,
                           const std::string &recojetname,
                           const std::string &truthjetname,
                           const std::string &filename)
  : SubsysReco(name)
  , _recojetname(recojetname)
  , _truthjetname(truthjetname)
  , _filename(filename)
{ }

int JetEvaluator::Init(PHCompositeNode * /*topNode*/)
{
  _ievent = 0;

  _tfile = new TFile(_filename.c_str(), "RECREATE");

  if (_do_recojet_eval)
  {
    _ntp_recojet = new TNtuple("ntp_recojet", "reco jet => max truth jet",
                               "event:id:ncomp:eta:phi:e:pt:"
                               "gid:gncomp:geta:gphi:ge:gpt:"
                               "efromtruth");
  }

  if (_do_truthjet_eval)
  {
    _ntp_truthjet = new TNtuple("ntp_truthjet", "truth jet => best reco jet",
                                "event:gid:gncomp:geta:gphi:ge:gpt:"
                                "id:ncomp:eta:phi:e:pt:"
                                "efromtruth");
  }

  return Fun4AllReturnCodes::EVENT_OK;
}

int JetEvaluator::process_event(PHCompositeNode *topNode)
{
  if (!_jetevalstack)
  {
    _jetevalstack = new JetEvalStack(topNode, _recojetname, _truthjetname);
    _jetevalstack->set_strict(_strict);
    _jetevalstack->set_verbosity(Verbosity() + 1);
  }
  else
  {
    _jetevalstack->next_event(topNode);
  }

  //-----------------------------------
  // print what is coming into the code
  //-----------------------------------

  printInputInfo(topNode);

  //---------------------------
  // fill the Evaluator NTuples
  //---------------------------

  fillOutputNtuples(topNode);

  //--------------------------------------------------
  // Print out the ancestry information for this event
  //--------------------------------------------------

  printOutputInfo(topNode);

  ++_ievent;

  return Fun4AllReturnCodes::EVENT_OK;
}

int JetEvaluator::End(PHCompositeNode * /*topNode*/)
{
  _tfile->cd();

  if (_do_recojet_eval)
  {
    _ntp_recojet->Write();
  }
  if (_do_truthjet_eval)
  {
    _ntp_truthjet->Write();
  }

  _tfile->Close();

  delete _tfile;

  if (Verbosity() > 0)
  {
    std::cout << "========================== JetEvaluator::End() ============================" << std::endl;
    std::cout << " " << _ievent << " events of output written to: " << _filename << std::endl;
    std::cout << "===========================================================================" << std::endl;
  }

  delete _jetevalstack;

  return Fun4AllReturnCodes::EVENT_OK;
}

void JetEvaluator::printInputInfo(PHCompositeNode * /*topNode*/)
{
  // to be implemented later if needed
  return;
}

void JetEvaluator::printOutputInfo(PHCompositeNode * /*topNode*/)
{
  // to be implemented later if needed
  return;
}

void JetEvaluator::fillOutputNtuples(PHCompositeNode *topNode)
{
  if (Verbosity() > 2)
  {
    std::cout << "JetEvaluator::fillOutputNtuples() entered" << std::endl;
  }


  JetRecoEval *recoeval = _jetevalstack->get_reco_eval();
  // JetTruthEval* trutheval = _jetevalstack->get_truth_eval();

  //-------------------------
  // fill the reco jet ntuple
  //-------------------------

  if (_do_recojet_eval)
  {
    if (Verbosity() > 1)
    {
      std::cout << "JetEvaluator::filling recojet ntuple..." << std::endl;
    }

    JetContainer *recojets = findNode::getClass<JetContainer>(topNode, _recojetname);
    if (!recojets)
    {
      std::cout << PHWHERE << " ERROR: Can't find " << _recojetname << std::endl;
      exit(-1);
    }

    // for every recojet
    for (auto recojet : *recojets)
    {
      /* Jet *recojet = iter.second; */
      Jet *truthjet = recoeval->max_truth_jet_by_energy(recojet);

      float id = recojet->get_id();
      float ncomp = recojet->size_comp();
      float eta = recojet->get_eta();
      float phi = recojet->get_phi();
      float e = recojet->get_e();
      float pt = recojet->get_pt();

      float gid = NAN;
      float gncomp = NAN;
      float geta = NAN;
      float gphi = NAN;
      float ge = NAN;
      float gpt = NAN;
      float efromtruth = NAN;

      if (truthjet)
      {
        gid = truthjet->get_id();
        gncomp = truthjet->size_comp();
        geta = truthjet->get_eta();
        gphi = truthjet->get_phi();
        ge = truthjet->get_e();
        gpt = truthjet->get_pt();
        efromtruth = recoeval->get_energy_contribution(recojet, truthjet);
      }

      float recojet_data[14] = {(float) _ievent,
                                id,
                                ncomp,
                                eta,
                                phi,
                                e,
                                pt,
                                gid,
                                gncomp,
                                geta,
                                gphi,
                                ge,
                                gpt,
                                efromtruth};

      _ntp_recojet->Fill(recojet_data);
    }
  }

  //-------------------------
  // fill the truth jet ntuple
  //-------------------------

  if (_do_truthjet_eval)
  {
    if (Verbosity() > 1)
    {
      std::cout << "JetEvaluator::filling truthjet ntuple..." << std::endl;
    }

    JetContainer *truthjets = findNode::getClass<JetContainer>(topNode, _truthjetname);
    if (!truthjets)
    {
      std::cout << PHWHERE << " ERROR: Can't find " << _truthjetname << std::endl;
      exit(-1);
    }

    // for every truthjet
    for (auto truthjet : *truthjets)
    {
      /* Jet *truthjet = iter.second; */
      Jet *recojet = recoeval->best_jet_from(truthjet);

      float gid = truthjet->get_id();
      float gncomp = truthjet->size_comp();
      float geta = truthjet->get_eta();
      float gphi = truthjet->get_phi();
      float ge = truthjet->get_e();
      float gpt = truthjet->get_pt();

      float id = NAN;
      float ncomp = NAN;
      float eta = NAN;
      float phi = NAN;
      float e = NAN;
      float pt = NAN;
      float efromtruth = NAN;

      if (recojet)
      {
        id = recojet->get_id();
        ncomp = recojet->size_comp();
        eta = recojet->get_eta();
        phi = recojet->get_phi();
        e = recojet->get_e();
        pt = recojet->get_pt();
        efromtruth = recoeval->get_energy_contribution(recojet, truthjet);
      }

      float truthjet_data[14] = {(float) _ievent,
                                 gid,
                                 gncomp,
                                 geta,
                                 gphi,
                                 ge,
                                 gpt,
                                 id,
                                 ncomp,
                                 eta,
                                 phi,
                                 e,
                                 pt,
                                 efromtruth};

      _ntp_truthjet->Fill(truthjet_data);
    }
  }

  return;
}