VIVisitor.hpp

Go to the documentation of this file.

#pragma once
 
#include "models/TreeBranch.hpp"
#include "utils/Types.hpp"
 
#include <vector>
 
namespace ppforest2 {
  class VIVisitor : public TreeNode::Visitor {
  public:
    std::vector<types::Feature> vi2_contributions;

    std::vector<types::Feature> vi3_contributions;

    types::FeatureVector const* scale = nullptr;

    explicit VIVisitor(int n_vars, types::FeatureVector const* scale = nullptr)
        : vi2_contributions(static_cast<std::size_t>(n_vars), 0.0f)
        , vi3_contributions(static_cast<std::size_t>(n_vars), 0.0f)
        , scale(scale) {}

    void visit(TreeBranch const& node) override {
      int const n_vars         = static_cast<int>(vi2_contributions.size());
      int const G_s            = static_cast<int>(node.groups.size());
      types::Feature const I_s = node.pp_index_value;
 
      for (int j = 0; j < n_vars && j < node.projector.size(); ++j) {
        types::Feature coeff = std::abs(node.projector(j));
 
        if (scale != nullptr && j < scale->size()) {
          coeff *= (*scale)(j);
        }
 
        if (G_s > 0) {
          vi2_contributions[static_cast<std::size_t>(j)] +=
              coeff / static_cast<types::Feature>(G_s); // <- not std::abs(coeff)
        }
 
        vi3_contributions[static_cast<std::size_t>(j)] += I_s * coeff;
      }
 
      node.lower->accept(*this);
      node.upper->accept(*this);
    }

    void visit(TreeLeaf const&) override {}
  };
}