NodeImpl.h

Go to the documentation of this file.

/*
  Copyright (C) 2002-2020 CERN for the benefit of the ATLAS collaboration
*/
 
#ifndef MVAUtils_NodeImpl_H
#define MVAUtils_NodeImpl_H
 
#include "MVAUtils/MVAUtilsDefs.h"
#include <vector>
#include <cstdint>
#include <cmath>
 
namespace TMVA{
  class MethodBDT;
  class DecisionTreeNode;
}
 
namespace MVAUtils
{
 
  // classes do not inherit from a common note. Different Node classes
  // are used as template argument for Forest classes. Forest classes will
  // hold vector of Nodes, so it is important to avoid overhead when calling
  // Node methods -> no polymoriphis
 
  // The main differences between Node classes is the logic to get the next
  // node, depending on the input values

  class NodeTMVA
  {
  public:
    NodeTMVA(const int ivar, const float val, const index_t right):
      m_cut(val), m_right(right), m_var(ivar) { }

    void Print(index_t index) const;

    index_t GetNext(const float value, index_t index) const;
 
    bool IsLeaf() const { return m_var < 0; } 

    var_t GetVar() const { return m_var; }

    float GetVal() const { return m_cut; }

    index_t GetLeft(index_t index) const { return index + 1; }

    index_t GetRight(index_t index) const { return index + m_right; }
 
  private:
    // the order is important to have the optimal memory usage
    float m_cut; 
    int16_t m_right; 
    var_t m_var; 
  };
 
  // inline speedup ~15%
  inline index_t NodeTMVA::GetNext(const float value, index_t index) const {
    return (value >= m_cut) ? GetLeft(index) : GetRight(index);
  }
 

  class NodeLGBMSimple
  {
  public:
    NodeLGBMSimple(const int ivar, const float val, const index_t right)
    : m_cut(val), m_right(right), m_var(ivar) { }
    void Print(index_t index) const;
    index_t GetNext(const float value, index_t index) const;
    bool IsLeaf() const { return m_var < 0; } 

    var_t GetVar() const { return m_var; }

    float GetVal() const { return m_cut; }

    index_t GetLeft(index_t index) const { return index + 1; }

    index_t GetRight(index_t index) const { return index + m_right; }
 
  private:
    float m_cut; 
    int16_t m_right; 
    var_t m_var; 
  };
 
  inline index_t NodeLGBMSimple::GetNext(const float value, index_t index) const {
    // note that this is different from TMVA (and not the opposite, e.g. cannot simply invert left/right)
    return (value <= m_cut) ? GetLeft(index) : GetRight(index);
  }

  class NodeLGBM
  {
  public:
    NodeLGBM(const int ivar, const float val, const index_t right, const int8_t default_left)
    : m_cut(val), m_right(right), m_var(ivar), m_default_left(default_left) { }
    void Print(index_t index) const;
    index_t GetNext(const float value, index_t index) const;
    bool GetDefaultLeft() const { return m_default_left; }
    bool IsLeaf() const { return m_var < 0; } 
 
    var_t GetVar() const { return m_var; }
    float GetVal() const { return m_cut; }
    index_t GetLeft(index_t index) const { return index + 1; }
    index_t GetRight(index_t index) const { return index + m_right; }
 
  private:
    float m_cut; 
    int16_t m_right; 
    var_t m_var; 
    int8_t m_default_left;  //< if to go left in case of nan input
  };
 
 
  inline index_t NodeLGBM::GetNext(const float value, index_t index) const {
    if (not std::isnan(value)) {
      return (value <= m_cut) ? GetLeft(index) : GetRight(index);
    }
    else {
      return (m_default_left) ? GetLeft(index) : GetRight(index);
    }
  }

  class NodeXGBoost
  {
  public:
    NodeXGBoost(const int ivar, const float val, const index_t right, const int8_t default_left)
    : m_cut(val), m_right(right), m_var(ivar), m_default_left(default_left) { }
    void Print(index_t index) const;
    index_t GetNext(const float value, index_t index) const;
    bool GetDefaultLeft() const { return m_default_left; }
    bool IsLeaf() const { return m_var < 0; } 
 
    var_t GetVar() const { return m_var; }
    float GetVal() const { return m_cut; }
    index_t GetLeft(index_t index) const { return index + 1; }
    index_t GetRight(index_t index) const { return index + m_right; }
 
  private:
    float m_cut; // cut value for internal nodes or response for leaf nodes
    int16_t m_right; // right relative index (to be added to current) (left is always current + 1)
    var_t m_var; // index of the variable to cut for internal nodes, -1 for leaf nodes
    int8_t m_default_left;  //default if to go left in case of nan input
  };
 
 
  inline index_t NodeXGBoost::GetNext(const float value, index_t index) const {
    if (not std::isnan(value)) {
      return (value < m_cut) ? GetLeft(index) : GetRight(index);
    }
    else {
      return (m_default_left) ? GetLeft(index) : GetRight(index);
    }
  }
 
}
#endif