JTreeScanner.hh

Go to the documentation of this file.

#ifndef __JSUPPORT__JTREESCANNER__
#define __JSUPPORT__JTREESCANNER__

#include <vector>
#include <string>
#include <ostream>
#include <iomanip>
#include <algorithm>

#include "JLang/JNullType.hh"
#include "JLang/JPointer.hh"
#include "JLang/JConversion.hh"
#include "JROOT/JChainReader.hh"
#include "JROOT/JCounter.hh"
#include "JROOT/JRootSupportkit.hh"
#include "Jeep/JMessage.hh"
#include "JSupport/JSingleFileScanner.hh"
#include "JSupport/JMultipleFileScanner.hh"
#include "JSupport/JTreeScannerInterface.hh"


/**
 * \author mdejong
 */

namespace JSUPPORT {}
namespace JPP { using namespace JSUPPORT; }

namespace JSUPPORT {

  using JLANG::JPointer;
  using JLANG::JNullType;
  using JLANG::JAssertConversion;
  using JROOT::JChainReader;
  using JROOT::counter_type;
  using JEEP::JMessage;

  
  /**
   * Pattern match for names of sub-branches that will not be read when ordering elements in a TTree.
   */
  static const char* const BRANCH_PATTERN = "vector";


  /**
   * Base class for JTreeScanner.
   *
   * Following ROOT memory management, a JChainReader object is created 
   * at construction and not deleted at destruction of this object.
   */
  template<class JClass_t>
  class JTreeScanner_t :
    public JPointer< JChainReader<JClass_t> >
  {
  public:    
    /**
     * Type conversion.
     *
     * \return                   file list
     */
    operator JMultipleFileScanner_t() const
    {
      if (this->is_valid())
        return JMultipleFileScanner_t(*this->get());
      else
        return JMultipleFileScanner_t();
    }

  protected:
    /**
     * Default constructor.
     */
    JTreeScanner_t() :
      JPointer< JChainReader<JClass_t> >(new JChainReader<JClass_t>())
    {
      gErrorIgnoreLevel = kError;
    }
  };
  
    
  /**
   * Template definition for direct access of elements in ROOT TChain.
   *
   * The optional second template argument is used to the determine the value of an element
   * which defines the apparent order the elements in the TChain.
   * It can also provide for a method to find a corresponding entry in the TTree at O(1) look-up time.
   *
   * This class implements the JSUPPORT::JTreeScannerInterface interface.
   */
  template<class JClass_t = JNullType, class JEvaluator_t = JNullType>
  class JTreeScanner;


  /**
   * Auxiliary base class for reporting messages.
   */
  template<>
  struct JTreeScanner<JNullType, JNullType> :
    public JMessage< JTreeScanner<> >
  {
    using JMessage< JTreeScanner<> >::debug;
  };
  

  /**
   * Specialiation of class JTreeScanner for unordered direct access of elements in ROOT TChain.
   *
   * The method JROOT::actionAtFileOpen is called at opening of the first file.
   *
   * The iteration of elements in the TChain will be in order of appearance. 
   */
  template<class JDerived_t, class JBase_t>
  class JTreeScanner<JAssertConversion<JDerived_t, JBase_t>, JNullType> :
    public virtual JTreeScannerInterface<JBase_t>,         // interface
    public JTreeScanner_t<JDerived_t>,                     // data source
    public JTreeScanner<>                                  // messaging
  {
  public:

    typedef JBase_t                                                data_type;
    typedef typename JTreeScannerInterface<JBase_t>::pointer_type  pointer_type;
    
    using JTreeScannerInterface<JBase_t>::configure;
  

    /**
     * Default constructor.
     */
    JTreeScanner() :
      counter(0)
    {}

    
    /**
     * Constructor.
     *
     * \param  file_list         file list
     * \param  limit             limit
     */
    JTreeScanner(const JMultipleFileScanner_t& file_list, const JLimit& limit = JLimit()) :
      counter(0)
    {
      this->configure(file_list, limit);
    }


    /**
     * Copy constructor.
     *
     * \param  input             TTree scanner
     */
    JTreeScanner(const JTreeScanner& input)
    {
      this->configure(input, JLimit());
    }


    /**
     * Rewind.
     */
    virtual void rewind()
    {
      counter = 0;
    }


    /**
     * Check availability of next element.
     *
     * \return            true if the iteration has more elements.
     */
    virtual bool hasNext()
    {
      if (counter < this->getLowerLimit()) {
        skip(this->getLowerLimit() - counter);
      }

      return (counter < this->getEntries()  &&
              counter < this->getUpperLimit());
    }


    /**
     * Get next element.
     *
     * \return            pointer to element.
     */
    virtual const pointer_type& next() 
    {
      if (this->hasNext()) {

        this->get()->GetEvent(counter++);

        ps.reset(this->get()->getAddress());

      } else {

        ps.reset(NULL);
      }

      return ps;
    }


    /**
     * Skip items.
     *
     * \param  ns         number of items to skip
     * \return            number of items skipped
     */
    virtual skip_type skip(const skip_type ns)
    {
      return JROOT::advance(counter, ns, this->getEntries());
    }

    
    /**
     * Configure.
     *
     * \param  file_list         file list
     * \param  limit             limit
     */
    virtual void configure(const JMultipleFileScanner_t& file_list, const JLimit& limit)
    {
      using namespace JROOT;

      this->setLimit(limit);

      this->rewind();

      this->get()->Reset();
      
      for (JMultipleFileScanner<>::const_iterator i = file_list.begin(); i != file_list.end(); ++i) {

        TFile* file = TFile::Open(i->c_str());

        if (file != NULL && file->GetListOfKeys()->Contains(this->get()->getTreeName())) { 
          this->get()->Add(i->c_str());
        }

        delete file;
      }
      
      this->getEntries();     // load TTree to set internal file

      actionAtFileOpen<JDerived_t>(this->get()->GetCurrentFile());
      actionAtFileOpen<JBase_t>   (this->get()->GetCurrentFile());
    }


    /**
     * Get number of entries.
     *
     * \return              number of entries
     */
    virtual Long64_t getEntries() const
    {
      return this->get()->GetEntries();
    }


    /**
     * Get entry at given index.
     *
     * \param  index             index
     * \return                   pointer to object (may be NULL)
     */
    virtual data_type* getEntry(Long64_t index) 
    {
      if (index >= 0 && index < this->getEntries()) {

        this->get()->GetEvent(index);

        return this->get()->getAddress();
      }

      return NULL;
    }


    /**
     * Get internal counter.
     *
     * \return            counter
     */
    virtual counter_type getCounter() const
    {
      return counter;
    }


    /**
     * Get actual class name.
     *
     * \return                   class name
     */
    virtual const char* getClassName() const
    {
      return JDerived_t::Class_Name();
    }
    
  protected:
    counter_type counter;
  private:
    pointer_type ps;
  };

  
  /**
   * Specialiation of class JTreeScanner for unordered direct access of elements in ROOT TChain.
   *
   * The iteration of elements in the TChain will be in order of appearance. 
   */
  template<class JClass_t>
  struct JTreeScanner<JClass_t, JNullType> :
    public JTreeScanner<JAssertConversion<JClass_t, JClass_t>, JNullType>
  {
    /**
     * Default constructor.
     */
    JTreeScanner() 
    {}


    /**
     * Constructor.
     *
     * \param  file_list         file list
     * \param  limit             limit
     */
    JTreeScanner(const JMultipleFileScanner_t& file_list, const JLimit& limit = JLimit()) :
      JTreeScanner<JAssertConversion<JClass_t, JClass_t>, JNullType>(file_list, limit)
    {}


    /**
     * Copy constructor.
     *
     * \param  input             TTree scanner
     */
    JTreeScanner(const JTreeScanner& input) :
      JTreeScanner<JAssertConversion<JClass_t, JClass_t>, JNullType>(input)
    {}
  };

    
  /**
   * Specialiation of class JTreeScanner for ordered direct access of elements in ROOT TChain.
   *
   * The optional second template argument is used to the determine the value of an element.
   * The elements in the TChain are then ordered according these values (from low to high). 
   * To this end, the evaluator class should provide for the following operator:
   * <pre>
   *             double operator()(const JBase_t& element) const;
   * </pre>
   * As a result, the iteration of elements in the TChain will then be in the specified order.
   * It also provides for a method to find a corresponding entry at O(1) look up time.
   */
  template<class JDerived_t, class JBase_t, class JEvaluator_t>
  class JTreeScanner<JAssertConversion<JDerived_t, JBase_t>, JEvaluator_t> :
    public JTreeScannerInterface<JBase_t, JEvaluator_t>,
    public JTreeScanner<JAssertConversion<JDerived_t, JBase_t>, JNullType>
  {
    /**
     * Auxiliary data structure for sorting of objects in TChain.
     */
    struct JEntry_t {
      /**
       * Default constructor.
       */
      JEntry_t() :
        index(-1),
        value(0.0)
      {}


      /**
       * Constructor.
       *
       * \param  __index           index
       * \param  __value           value
       */
      JEntry_t(const Long64_t __index,
               const double   __value) :
        index(__index),
        value(__value)
      {}


      /**
       * Comparison between two TChain entries.
       *
       * \param  first             first  entry
       * \param  second            second entry
       * \return                   true if first value less than second; else false
       */
      friend inline bool operator<(const JEntry_t& first, const JEntry_t& second)
      {
        return first.value < second.value;
      }


      /**
       * Comparison between TChain entry and value.
       *
       * \param  entry             entry
       * \param  value             value
       * \return                   true if given entry has less value; else false
       */
      friend inline bool operator<(const JEntry_t& entry, const double value)
      {
        return entry.value < value;
      }

 
      Long64_t index;             //!< index in TChain
      double   value;             //!< corresponding value
    };


    /**
     * Type definition of internal queue for ordering the elements in the TChain.
     */
    typedef std::vector<JEntry_t> queue_type;


  public:

    typedef JBase_t                                                              data_type;
    typedef typename JTreeScannerInterface<JBase_t, JEvaluator_t>::pointer_type  pointer_type;
    
    using JTreeScannerInterface<JBase_t>::configure;    
    using JTreeScannerInterface<JBase_t, JEvaluator_t>::find;
    using JTreeScannerInterface<JBase_t, JEvaluator_t>::getValue;
    using JTreeScanner<>::debug;


    /**
     * Default constructor.
     */
    JTreeScanner() :
      JTreeScannerInterface<JBase_t, JEvaluator_t>()
    {}

    
    /**
     * Constructor.
     *
     * \param  file_list         file list
     * \param  evaluator         evaluator
     */
    JTreeScanner(const JMultipleFileScanner_t& file_list,
                 const JEvaluator_t&           evaluator = JEvaluator_t()) :
      JTreeScannerInterface<JBase_t, JEvaluator_t>(evaluator)
    {
      this->configure(file_list);
    }
    
    
    /**
     * Constructor.
     *
     * \param  file_list         file list
     * \param  limit             limit
     * \param  evaluator         evaluator
     */
    JTreeScanner(const JMultipleFileScanner_t& file_list,
                 const JLimit&                 limit,
                 const JEvaluator_t&           evaluator = JEvaluator_t()) :
      JTreeScannerInterface<JBase_t, JEvaluator_t>(evaluator)
    {
      this->configure(file_list, limit);
    }
    
    
    /**
     * Get next element.
     *
     * \return            pointer to element.
     */
    virtual const pointer_type& next()
    {
      if (this->hasNext()) {
        
        this->get()->GetEvent(queue[this->counter++].index);

        ps.reset(this->get()->getAddress());

      } else {

        ps.reset(NULL);
      }

      return ps;
    }
      
    
    /**
     * Configure.
     *
     * \param  file_list         file list
     * \param  limit             limit
     */
    virtual void configure(const JMultipleFileScanner_t& file_list, const JLimit& limit)
    {
      using namespace std;
      
      JTreeScanner<JAssertConversion<JDerived_t, JBase_t>, JNullType>::configure(file_list, limit);

      setBranchStatus(this->get()->GetBranch(this->get()->getBranchName()), BRANCH_PATTERN, false);

      queue.resize(this->getEntries());

      typename queue_type::iterator out = queue.begin();

      for (Long64_t i = 0; i != this->getEntries(); ++i, ++out) {

        STATUS(left << setw(24) << this->get()->GetName() << right << ' ' << setw(3) << (100 * (i + 1)) / this->getEntries() << "%\r"); DEBUG(endl);
        
        this->get()->GetEvent(i);

        const data_type* p = this->get()->getAddress();
        
        *out = JEntry_t(i, this->JTreeScannerInterface<JBase_t, JEvaluator_t>::getValue(*p));
      }
      STATUS(endl);

      this->get()->SetBranchStatus("*", 1);

      sort(queue.begin(), queue.end());
    }


    /**
     * Get entry at given index.
     *
     * \param  index             index
     * \return                   pointer to object (may be NULL)
     */
    virtual data_type* getEntry(Long64_t index) 
    {
      if (index >= 0 && index < (Long64_t) queue.size()) {

        this->get()->GetEvent(queue[index].index);

        return this->get()->getAddress();
      }

      return NULL;
    }


    /**
     * Find index of element that is closest in value to given value.
     *
     * A subsequent call to method getEntry() will point to the corresponding object.
     *
     * \param  value             value
     * \return                   index (-1 in case of error)
     */
    virtual Long64_t find(const double value) const
    {
      using namespace std;
      
      if (!queue.empty()) {

        typename queue_type::const_iterator p = lower_bound(queue.begin(), queue.end(), value);
        
        if        (p == queue.end()) {
          
          return queue.size() - 1;

        } else if (p == queue.begin()) {
          
          return 0;

        } else {
          
          typename queue_type::const_iterator q = p--;
          
          if (value - p->value < q->value - value)
            return distance(queue.begin(), p);
          else
            return distance(queue.begin(), q);
        }
      }

      return -1;
    }

  protected:
    /**
     * Set status of branch.
     * 
     * Note that the status of the branch and all sub-branches with names including given pattern will recursively be set.
     *
     * \param  branch            pointer to branch
     * \param  pattern           pattern
     * \param  status            status
     */
    static void setBranchStatus(TBranch* branch, const char* pattern, const bool status)
    {
      using namespace std;

      if (branch != NULL) {

        TObjArray* array  = branch->GetListOfBranches();

        for (Int_t i = 0; i != array->GetEntries(); ++i) {

          TBranch* p = (TBranch*) array->At(i);

          if (p != NULL && string(p->GetName()).find(pattern) != string::npos) {

            if (p->GetSplitLevel() == 0) {

              NOTICE("Set status of branch " << p->GetName() << " to " << status << endl);
            
              p->SetStatus(status);
            }
          } else {

            setBranchStatus(p, pattern, status);
          }
        }
      }
    }


    queue_type   queue;

  private:
    pointer_type ps;
  };


  /**
   * Specialiation of class JTreeScanner for ordered direct access of elements in ROOT TChain.
   */
  template<class JClass_t, class JEvaluator_t>
  struct JTreeScanner :
    public JTreeScanner<JAssertConversion<JClass_t, JClass_t>, JEvaluator_t>
  {
    /**
     * Default constructor.
     */
    JTreeScanner() 
    {}


    /**
     * Constructor.
     *
     * \param  file_list         file list
     * \param  evaluator         evaluator
     */
    JTreeScanner(const JMultipleFileScanner_t& file_list,
                 const JEvaluator_t&           evaluator = JEvaluator_t()) :
      JTreeScanner<JAssertConversion<JClass_t, JClass_t>, JEvaluator_t>(file_list, evaluator)
    {}
    
    
    /**
     * Constructor.
     *
     * \param  file_list         file list
     * \param  limit             limit
     * \param  evaluator         evaluator
     */
    JTreeScanner(const JMultipleFileScanner_t& file_list,
                 const JLimit&                 limit,
                 const JEvaluator_t&           evaluator = JEvaluator_t()) :
      JTreeScanner<JAssertConversion<JClass_t, JClass_t>, JEvaluator_t>(file_list, limit, evaluator)
    {}
  };
}

#endif