mse_gain.hpp

Go to the documentation of this file.

#ifndef MLPACK_METHODS_DECISION_TREE_MSE_GAIN_HPP
#define MLPACK_METHODS_DECISION_TREE_MSE_GAIN_HPP

#include <mlpack/prereqs.hpp>
#include "utils.hpp"

namespace mlpack {
namespace tree {

class MSEGain
{
 public:
  template<bool UseWeights, typename VecType, typename WeightVecType>
  static double Evaluate(const VecType& values,
                         const WeightVecType& weights,
                         const size_t begin,
                         const size_t end)
  {
    double mse = 0.0;

    if (UseWeights)
    {
      double accWeights = 0.0;
      double weightedMean = 0.0;
      WeightedSum(values, weights, begin, end, accWeights, weightedMean);

      // Catch edge case: if there are no weights, the impurity is zero.
      if (accWeights == 0.0)
        return 0.0;

      weightedMean /= accWeights;

      for (size_t i = begin; i < end; ++i)
        mse += weights[i] * std::pow(values[i] - weightedMean, 2);

      mse /= accWeights;
    }
    else
    {
      double mean = 0.0;
      Sum(values, begin, end, mean);
      mean /= (double) (end - begin);

      mse = arma::accu(arma::square(values.subvec(begin, end - 1) - mean));
      mse /= (double) (end - begin);
    }

    return -mse;
  }

  template<bool UseWeights, typename VecType, typename WeightVecType>
  static double Evaluate(const VecType& values,
                         const WeightVecType& weights)
  {
    // Corner case: if there are no elements, the impurity is zero.
    if (values.n_elem == 0)
      return 0.0;

    return Evaluate<UseWeights>(values, weights, 0, values.n_elem);
  }

  std::tuple<double, double> BinaryGains()
  {
    double mseLeft = leftSumSquares / leftSize - leftMean * leftMean;
    double mseRight = (totalSumSquares - leftSumSquares) / rightSize
          - rightMean * rightMean;

    return std::make_tuple(-mseLeft, -mseRight);
  }

  template<bool UseWeights, typename ResponsesType, typename WeightVecType>
  void BinaryScanInitialize(const ResponsesType& responses,
                            const WeightVecType& weights,
                            const size_t minimum)
  {
    typedef typename ResponsesType::elem_type RType;
    typedef typename WeightVecType::elem_type WType;

    // Initializing data members to cache statistics.
    leftMean = 0.0;
    rightMean = 0.0;
    leftSize = 0.0;
    rightSize = 0.0;
    leftSumSquares = 0.0;
    totalSumSquares = 0.0;

    if (UseWeights)
    {
      totalSumSquares = arma::accu(weights % arma::square(responses));
      for (size_t i = 0; i < minimum - 1; ++i)
      {
        const WType w = weights[i];
        const RType x = responses[i];

        // Calculating initial weighted mean of responses for the left child.
        leftSize += w;
        leftMean += w * x;
        leftSumSquares += w * x * x;
      }
      if (leftSize > 1e-9)
        leftMean /= leftSize;

      for (size_t i = minimum - 1; i < responses.n_elem; ++i)
      {
        const WType w = weights[i];
        const RType x = responses[i];

        // Calculating initial weighted mean of responses for the right child.
        rightSize += w;
        rightMean += w * x;
      }
      if (rightSize > 1e-9)
        rightMean /= rightSize;
    }
    else
    {
      totalSumSquares = arma::accu(arma::square(responses));
      for (size_t i = 0; i < minimum - 1; ++i)
      {
        const RType x = responses[i];

        // Calculating the initial mean of responses for the left child.
        ++leftSize;
        leftMean += x;
        leftSumSquares += x * x;
      }
      if (leftSize > 1e-9)
        leftMean /= leftSize;

      for (size_t i = minimum - 1; i < responses.n_elem; ++i)
      {
        const RType x = responses[i];

        // Calculating the initial mean of responses for the right child.
        ++rightSize;
        rightMean += x;
      }
      if (rightSize > 1e-9)
        rightMean /= rightSize;
    }
  }

  template<bool UseWeights, typename ResponsesType, typename WeightVecType>
  void BinaryStep(const ResponsesType& responses,
                  const WeightVecType& weights,
                  const size_t index)
  {
    typedef typename ResponsesType::elem_type RType;
    typedef typename WeightVecType::elem_type WType;

    if (UseWeights)
    {
      const WType w = weights[index];
      const RType x = responses[index];

      // Update weighted sum of squares for left child.
      leftSumSquares += w * x * x;

      // Update weighted mean for both childs.
      leftMean = (leftMean * leftSize + w * x) / (leftSize + w);
      leftSize += w;

      rightMean = (rightMean * rightSize - w * x) / (rightSize - w);
      rightSize -= w;
    }
    else
    {
      const RType x = responses[index];

      // Update sum of squares for left child.
      leftSumSquares += x * x;

      // Update mean for both childs.
      leftMean = (leftMean * leftSize + x) / (leftSize + 1);
      ++leftSize;

      rightMean = (rightMean * rightSize - x) / (rightSize - 1);
      --rightSize;
    }
  }

 private:
  // Stores the sum of squares / weighted sum of squares for the left child.
  double leftSumSquares;
  // For unweighted data, stores the number of elements in each child.
  // For weighted data, stores the sum of weights of elements in each
  // child.
  double leftSize;
  double rightSize;
  // Stores the mean / weighted mean.
  double leftMean;
  double rightMean;
  // Stores the total sum of squares / total weighted sum of squares.
  double totalSumSquares;
};

} // namespace tree
} // namespace mlpack

#endif