LikeModule.cxx

Go to the documentation of this file.

// C++
#include <cmath>
#include <iomanip>
#include <iostream>
#include <stdio.h>
#include <stdlib.h>
#include <sstream>

// ROOT
#include "TFile.h"
#include "TTree.h"

#include "LikeModule.h"
 
using namespace std;

//-------------------------------------------------------------------------------------------
Lit::LikeModule::LikeModule()
   :fDimn(0),
    fTree(0)
{
}

//-------------------------------------------------------------------------------------------
Lit::LikeModule::~LikeModule()
{
   if(fTree)
   {
      delete fTree;
   }
}

//-------------------------------------------------------------------------------------------
bool Lit::LikeModule::Write(const string &filepath, const string &treename)
{ 
   TFile *file = new TFile(filepath.c_str(), "UPDATE");
   if(!file || !file -> IsOpen())
   {
      cout << "LikeModule::Write() - failed to recreate ROOT file:\n" << filepath << endl;      
      return false;  
   }

   // write out all events in TTree
   Write(file, treename);

   file -> Write();
   file -> Close();

   return true;
}

//-------------------------------------------------------------------------------------------
void Lit::LikeModule::Write(TDirectory *dir, const string &treename)
{ 
   if(!dir)
   {
      cerr << "LikeModule::Write() - null TDirectory pointer" << endl;
      return;
   }

   if(dir -> Get(treename.c_str()))
   {
      const string treename_tmp = treename + ";*";
      cout << "LikeModule::Write - delete object with name " << treename_tmp << endl;
      dir -> Delete(treename_tmp.c_str());
   }

   Event *event = new Event();
   TTree *tree = new TTree(treename.c_str(), "event tree");
   tree -> SetDirectory(dir);
   tree -> Branch("event", "Lit::Event", &event);

   double size = 0.0;
   for(EventVec::const_iterator it = fEvent.begin(); it != fEvent.end(); ++it)
   {
      (*event) = (*it);
      size += tree -> Fill();
   }

   cout << "LikeModule::Write() - filled tree with " << size/1048576.0 << "MB and " 
        << fEvent.size () << " events" << endl;

   delete event; 
}

//-------------------------------------------------------------------------------------------
bool Lit::LikeModule::Read(const string &filepath, const int nevent, const string &treename)
{ 
   TFile *file = new TFile(filepath.c_str(), "READ");
   if(!file || !file -> IsOpen())
   {
      cout << "LikeModule::Read() - failed to read ROOT file:\n" << filepath << endl;      
      return false;  
   }

   const bool result = Read(file, nevent, treename);

   file -> Close();

   return result;
}

//-------------------------------------------------------------------------------------------
bool Lit::LikeModule::Read(TDirectory *dir, int nevent, const string &treename)
{ 
   TTree *tree = dynamic_cast<TTree *>(dir -> Get(treename.c_str()));
   if(!tree)
   {
      cout << "LikeModule::Read() - failed to find " << treename << " tree" << endl;
      return false;
   }

   Event *event = new Event();
   tree -> SetBranchAddress("event", &event);

   if(nevent < 1 || nevent > tree -> GetEntries())
   {
      nevent = tree -> GetEntries();
   }

   double size = 0.0;
   for(int i = 0; i < nevent; ++i)
   {
      size += tree -> GetEntry(i);
      Add(*event);
   }

   cout << "LikeModule::Read() - read " << size/1048576.0 << "MB and " 
        << fEvent.size () << " events" << endl;

   delete event;

   return true;
}

//-------------------------------------------------------------------------------------------
void Lit::LikeModule::Add(const Event &event)
{   
   if(fTree)
   {
      cerr << "LikeModule::Add() - can not add event: tree is already built" << endl;
      return;      
   }

   if(fDimn < 1)
   {
      fDimn = event.GetNVar();
   }
   else if(fDimn != event.GetNVar())
   {
      cerr << "LikeModule::Add() - number of dimension does not match previous events" << endl;
      return;
   }

   fEvent.push_back(event);

   for(unsigned int ivar = 0; ivar < fDimn; ++ivar)
   {
      fVar[ivar].push_back(event.GetVar(ivar));
   }

   std::map<short, unsigned int>::iterator cit = fCount.find(event.GetType());
   if(cit == fCount.end())
   {
      fCount[event.GetType()] = 1;
   }
   else
   {
      ++(cit -> second);
   }
}

//-------------------------------------------------------------------------------------------
bool Lit::LikeModule::Fill(const unsigned short optimize_depth, const std::string &option)
{
   if(fTree)
   {
      cerr << "LikeModule::Fill - tree already have been created" << endl;
      return false;
   }

   // If trim option is set then find class with lowest number of events
   // and set that as maximum number of events for all other classes.
   unsigned int min = 0;
   if(option.find("trim") != string::npos)
   {  
      for(map<short, unsigned int>::const_iterator it = fCount.begin(); it != fCount.end(); ++it)
      {
         if(min == 0 || min > it -> second)
         {
            min = it -> second;
         }
      }
      
      cout << "LikeModule::Fill - trimming all classes to " << min << " events" << endl;
      for(map<short, unsigned int>::const_iterator it = fCount.begin(); it != fCount.end(); ++it)
      {
         cout << "class " << it -> first << " has " 
              << setfill(' ') << setw(8) << it -> second << " events" << endl;
      }

      fCount.clear();
      fVar.clear();

      EventVec evec;

      for(EventVec::const_iterator event = fEvent.begin(); event != fEvent.end(); ++event)
      {
         std::map<short, unsigned int>::iterator cit = fCount.find(event -> GetType());
         if(cit == fCount.end())
         {
            fCount[event -> GetType()] = 1;
         }
         else if(cit -> second < min)
         {
            ++(cit -> second);
         }
         else
         {
            continue;
         }
         
         for(unsigned int d = 0; d < fDimn; ++d)
         {
            fVar[d].push_back(event -> GetVar(d));
         }       

         evec.push_back(*event);
      }

      cout << "LikeModule::Fill - erased " << fEvent.size() - evec.size() << " events" << endl;

      fEvent = evec;
   }

   //clear event count
   fCount.clear();

   // sort each variable for all events - needs this before Balance()
   for(VarMap::iterator it = fVar.begin(); it != fVar.end(); ++it)
   {
      std::sort((it -> second).begin(), (it -> second).end());
   }

   if(option.find("metric") != string::npos)
   {
      ComputeMetric();

      // sort again each variable for all events - needs this before Balance()
      // rescaling has changed variable values
      for(VarMap::iterator it = fVar.begin(); it != fVar.end(); ++it)
      {
         std::sort((it -> second).begin(), (it -> second).end());
      }
   }

   // If optimize_depth > 0 then fill first optimize_depth levels
   // with empty nodes that split separating variable in half for
   // all child nodes. If optimize_depth = 0 then split variable 0
   // at the median (in half) and return it as root node
   fTree = Balance(optimize_depth);

   if(!fTree)
   {
      cout << "LikeModule::Fill - failed to create tree" << endl;
      return false;      
   }      
   
   for(EventVec::const_iterator event = fEvent.begin(); event != fEvent.end(); ++event)
   {
      fTree -> Add(*event, 0);

      std::map<short, unsigned int>::iterator cit = fCount.find(event -> GetType());
      if(cit == fCount.end())
      {
         fCount[event -> GetType()] = 1;
      }
      else
      {
         ++(cit -> second);
      }
   }
   
   for(map<short, unsigned int>::const_iterator it = fCount.begin(); it != fCount.end(); ++it)
   {
      cout << "class " << it -> first << " has " 
           << setfill(' ') << setw(8) << it -> second << " events" << endl;
   }
   
   return true;
}

//-------------------------------------------------------------------------------------------
bool Lit::LikeModule::Find(Event event, const unsigned int nfind) const
{  
   // if tree has been filled then search for nfind closest events 
   // if metic (fVarScale map) is computed then rescale event variables
   // using previsouly computed width of variable distribution

   if(!fTree)
   {
      cerr << "LikeModule::Find() - tree has not been filled" << endl;
      return false;
   }
   if(fDimn != event.GetNVar())
   {
      cerr << "LikeModule::Find() - number of dimension does not match training events" << endl;
      return false;
   }
   if(nfind < 1)
   {
      cerr << "LikeModule::Find() - requested 0 nearest neighbors" << endl;
      return false;
   }

   // if variable widths are computed then rescale variable in this event
   // to same widths as events in stored kd-tree
   if(!fVarScale.empty())
   {
      event = Scale(event);
   }

   fResult.Clear();

   fResult.fEvent = event;

   // recursive kd-tree search for nfind-nearest neighbors
   Lit::Find(fResult.fList, fTree, event, nfind);

   for(List::const_iterator lit = fResult.GetList().begin(); lit != fResult.GetList().end(); ++lit)
   {
      const Node *node = lit -> first;

      map<short, double>::iterator pit = fResult.fProbNN.find(node -> GetType());
      if(pit == fResult.fProbNN.end())
      {
         fResult.fProbNN[node -> GetType()] = node -> GetWeight()/double(nfind);
      }
      else
      {
         pit -> second += node -> GetWeight()/double(nfind);
      }
   }

   return true;
}

//-------------------------------------------------------------------------------------------
bool Lit::LikeModule::Find(const unsigned int nfind, const string &option) const
{
   if(fCount.empty() || !fTree)
   {
      return false;
   }

   static std::map<short, unsigned int>::const_iterator cit = fCount.end();

   if(cit == fCount.end())
   {
      cit = fCount.begin();
   }

   const short etype = (cit++) -> first;

   if(option == "flat")
   {
      VarVec dvec;
      for(unsigned int d = 0; d < fDimn; ++d)
      {
         VarMap::const_iterator vit = fVar.find(d);
         if(vit == fVar.end())
         {
            return false;
         }

         const std::vector<double> &vvec = vit -> second;
         
         if(vvec.empty())
         {
            return false;
         }

         // assume that vector elements of fVar are sorted
         const VarType min = vvec.front();
         const VarType max = vvec.back();
         const VarType width = max - min;

         if(width < 0.0 || width > 0.0)
         {
            dvec.push_back(min + width*double(std::rand())/double(RAND_MAX));
         }
         else
         {
            return false;
         }
      }

      const Event event(dvec, 1.0, etype);
      
      Find(event, nfind);  

      return true;
   }

   return false;
}

//-------------------------------------------------------------------------------------------
Lit::Node* Lit::LikeModule::Balance(const unsigned int bdepth)
{
   // Balance() balances binary tree for first bdepth levels
   // for each depth we split sorted depth % dimension variables
   // into 2^bdepth parts

   if(fVar.empty() || fDimn != fVar.size())
   {
      cout << "LikeModule::Balance - can not build a tree" << endl;
      return 0;
   }

   const unsigned int size = (fVar.begin() -> second).size();
   if(size < 1)
   {
      cout << "LikeModule::Balance - can not build a tree without events" << endl;
      return 0;
   }

   for(VarMap::const_iterator it = fVar.begin(); it != fVar.end(); ++it)
   {
      if((it -> second).size() != size)
      {
         cout << "LikeModule::Balance - # of variables doesn't match between dimensions" << endl;
         return 0;
      }
   }

   if(double(fDimn*size) < pow(2.0, double(bdepth)))
   {
      cout << "LikeModule::Balance - balance depth exceeds number of events" << endl;
      return 0;      
   }   

   cout << "Balancing tree for " << fDimn << " variables with " << size << " values" << endl;

   vector<Node *> pvec, cvec;

   VarMap::const_iterator it = fVar.find(0);
   if(it == fVar.end() || (it -> second).size() < 2)
   {
      cout << "Missing 0 variable" << endl;
      return 0;
   }

   const Event pevent(VarVec(fDimn, (it -> second)[size/2]), -1.0, -1);
   
   Node *tree = new Node(0, pevent, 0);
   
   pvec.push_back(tree);

   for(unsigned int depth = 1; depth < bdepth; ++depth)
   {            
      const unsigned int mod = depth % fDimn;
   
      VarMap::const_iterator vit = fVar.find(mod);
      if(vit == fVar.end())
      {
         cerr << "Missing " << mod << " variable" << endl;
         return 0;
      }
      const vector<double> &dvec = vit -> second;

      if(dvec.size() < 2)
      {
         cerr << "Missing " << mod << " variable" << endl;
         return 0;       
      }

      //cout << "Mod = " << mod << endl;

      unsigned int ichild = 1;
      for(vector<Node *>::iterator pit = pvec.begin(); pit != pvec.end(); ++pit)
      {
         Node *parent = *pit;

         const VarType lmedian = dvec[size*ichild/(2*pvec.size() + 1)];
         //cout << "ChildL " << setw(2) << ichild << " lmedian " << lmedian << endl;
         ++ichild;

         const VarType rmedian = dvec[size*ichild/(2*pvec.size() + 1)];
         //cout << "ChildR " << setw(2) << ichild << " rmedian " << rmedian << endl;
         ++ichild;
      
         const Event levent(VarVec(fDimn, lmedian), -1.0, -1);
         const Event revent(VarVec(fDimn, rmedian), -1.0, -1);

         Node *lchild = new Node(parent, levent, mod);
         Node *rchild = new Node(parent, revent, mod);

         parent -> SetNodeL(lchild);
         parent -> SetNodeR(rchild);

         cvec.push_back(lchild);
         cvec.push_back(rchild);
      }

      pvec = cvec;
      cvec.clear();
   }

   return tree;
}

//-------------------------------------------------------------------------------------------
void Lit::LikeModule::ComputeMetric(const unsigned int ifrac)
{
   // compute scale factor for each variable (dimension) so that 
   // distance is computed uniformely along each dimension
   // compute width of interval that includes (100 - 2*ifrac)% of events
   // below, assume that in fVar each vector of values is sorted

   if(ifrac > 49)
   {
      cerr << "LikeModule::ComputeMetric() - fraction can not exceed 50%" << endl;
      return;
   }
   
   if(!fVarScale.empty())
   {
      cerr << "LikeModule::ComputeMetric() - metric is already computed" << endl;
      return;
   }

   cout << "Computing scale for 1d axes: ifrac = " << ifrac << endl;

   fVarScale.clear();
   
   for(VarMap::const_iterator vit = fVar.begin(); vit != fVar.end(); ++vit)
   {
      const vector<double> &dvec = vit -> second;
      
      vector<double>::const_iterator beg_it = dvec.end();
      vector<double>::const_iterator end_it = dvec.end();
      for(vector<double>::const_iterator dit = dvec.begin(); dit != dvec.end(); ++dit)
      {
         const int dist = distance(dvec.begin(), dit);

         if((100*dist)/dvec.size() == ifrac && beg_it == dvec.end())
         {
            beg_it = dit;
         }

         if((100*dist)/dvec.size() == 100 - ifrac && end_it == dvec.end())
         {
            end_it = dit;
         }
      }

      if(beg_it == dvec.end() || end_it == dvec.end())
      {
         cerr << "LikeModule::ComputeMetric() - failed to determine scale " << vit -> first << endl;
         continue;
      }

      const double lpos = *beg_it;
      const double rpos = *end_it;

      if(!(lpos < rpos))
      {
         cerr << "LikeModule::ComputeMetric() - min value is greater than max value" << endl;
         continue;       
      }
      
      cout << "Variable " << vit -> first 
           << " included " << distance(beg_it, end_it) + 1
           << " events: width = " << setfill(' ') << setw(5) << setprecision(3) << rpos - lpos
           << ", (min, max) = (" << setfill(' ') << setw(5) << setprecision(3) << lpos 
           << ", " << setfill(' ') << setw(5) << setprecision(3) << rpos << ")" << endl;

      fVarScale[vit -> first] = rpos - lpos;
   }

   fVar.clear();

   for(unsigned int ievent = 0; ievent < fEvent.size(); ++ievent)
   {      
      fEvent[ievent] = Scale(fEvent[ievent]);

      for(unsigned int ivar = 0; ivar < fDimn; ++ivar)
      {
         fVar[ivar].push_back(fEvent[ievent].GetVar(ivar));
      }  
   }
}

//-------------------------------------------------------------------------------------------
const Lit::Event Lit::LikeModule::Scale(const Event &event) const
{
   // scale each event variable so that rms of variables is approximately 1.0
   // this allows comparisons of variables with distinct scales and units
   
   if(fVarScale.empty())
   {
      return event;
   }

   if(event.GetNVar() != fVarScale.size())
   {
      cerr << "LikeModule::Scale() - mismatched metric and event size" << endl;
      return event;
   }

   VarVec vvec(event.GetNVar(), 0.0);

   for(unsigned int ivar = 0; ivar < event.GetNVar(); ++ivar)
   {
      map<int, double>::const_iterator fit = fVarScale.find(ivar);
      if(fit == fVarScale.end())
      {
         cerr << "LikeModule::Scale() - failed to find scale for " << ivar << endl;
         continue;
      }
      
      if(fit -> second > 0.0)
      {
         vvec[ivar] = event.GetVar(ivar)/fit -> second;
      }
      else
      {
         cerr << "Variable " << ivar << " has zero width" << endl;
      }
   }

   return Event(vvec, event.GetWeight(), event.GetType());
}

//-------------------------------------------------------------------------------------------
void Lit::LikeModule::Print() const
{
   Print(std::cout);
}

//-------------------------------------------------------------------------------------------
void Lit::LikeModule::Print(std::ostream &os) const
{
   os << "---------------------------------------------------------------------------------"<< endl;
   os << "Dumping stringstream content of LikeModule" << endl;
   os << "---------------------------------------------------------------------------------"<< endl;
}

---