<commit_before>package io.citrine.lolo.trees.regression

import io.citrine.lolo.trees.{ModelLeaf, ModelNode, TrainingNode}
import io.citrine.lolo.{Learner, Model, PredictionResult}

import scala.collection.mutable

/**
  * Created by maxhutch on 3/8/17.
  */
class RegressionTrainingLeaf(
                              trainingData: Seq[(Vector[AnyVal], Double, Double)],
                              leafLearner: Learner,
                              depth: Int
                            ) extends TrainingNode(trainingData, depth) {

  /**
    * Average the training data
    *
    * @return lightweight prediction node
    */
  def getNode(): ModelNode[PredictionResult[Double]] = {
    new ModelLeaf(leafLearner.train(trainingData).getModel().asInstanceOf[Model[PredictionResult[Double]]], depth)
  }

  def getFeatureImportance(): scala.collection.mutable.ArraySeq[Double] = {
    importance match {
      case Some(x) =>
        val expectations: (Double, Double) = trainingData.map{ case (v, l, w) =>
          (l * w, l * l * w)
        }.reduce((u: (Double, Double), v: (Double, Double)) => (u._1 + v._1, u._2 + v._2))
        val impurity = expectations._2 / trainingData.size - Math.pow(expectations._1 / trainingData.size, 2.0)
        mutable.ArraySeq(x: _*).map(_ * impurity)
      case None => mutable.ArraySeq.fill(trainingData.head._1.size)(0.0)
    }
  }

  val leafTrainingResult = leafLearner.train(trainingData)
  val model = leafTrainingResult.getModel().asInstanceOf[Model[PredictionResult[Double]]]
  val importance = leafTrainingResult.getFeatureImportance()
}<commit_msg>Adding comments; fixing duplicate train<commit_after>package io.citrine.lolo.trees.regression

import io.citrine.lolo.trees.{ModelLeaf, ModelNode, TrainingNode}
import io.citrine.lolo.{Learner, Model, PredictionResult}

import scala.collection.mutable

/**
  * Training leaf node for regression trees
  * Created by maxhutch on 3/8/17.
  */
class RegressionTrainingLeaf(
                              trainingData: Seq[(Vector[AnyVal], Double, Double)],
                              leafLearner: Learner,
                              depth: Int
                            ) extends TrainingNode(trainingData, depth) {

  /**
    * Wrap the leaf model (previously trained) in a lightweight leaf node
    *
    * @return lightweight prediction node
    */
  def getNode(): ModelNode[PredictionResult[Double]] = {
    new ModelLeaf(model.asInstanceOf[Model[PredictionResult[Double]]], depth)
  }

  /**
    * Pull the leaf model's feature importance and rescale it by the remaining impurity
    * @return feature importance as a vector
    */
  def getFeatureImportance(): scala.collection.mutable.ArraySeq[Double] = {
    importance match {
      case Some(x) =>
        val expectations: (Double, Double) = trainingData.map{ case (v, l, w) =>
          (l * w, l * l * w)
        }.reduce((u: (Double, Double), v: (Double, Double)) => (u._1 + v._1, u._2 + v._2))
        val impurity = expectations._2 / trainingData.size - Math.pow(expectations._1 / trainingData.size, 2.0)
        mutable.ArraySeq(x: _*).map(_ * impurity)
      case None => mutable.ArraySeq.fill(trainingData.head._1.size)(0.0)
    }
  }

  /** Train the leaf learner on the training data */
  val leafTrainingResult = leafLearner.train(trainingData)
  /** Pull out the model for future use */
  val model = leafTrainingResult.getModel().asInstanceOf[Model[PredictionResult[Double]]]
  /** Pull out the importance for future use */
  val importance = leafTrainingResult.getFeatureImportance()
}
{'fraction_non_alphanumeric': 0.1199887860947575, 'fraction_numerical': 0.008410428931875526, 'mean_word_length': 2.8698481561822127, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '378493', 'n_tokens_mistral': 1091, 'n_tokens_neox': 1043, 'n_words': 252}