From 61e88e6b3c914807bf23da8e39d239f9102364dd Mon Sep 17 00:00:00 2001
From: Rishabh Garg <56191449+RishabhGarg108@users.noreply.github.com>
Date: Wed, 30 Jun 2021 23:00:50 +0530
Subject: [PATCH] Add OutputValue and SimilarityScore methods

---
 .../xgboost/loss_functions/sse_loss.hpp       | 27 +++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/src/mlpack/methods/xgboost/loss_functions/sse_loss.hpp b/src/mlpack/methods/xgboost/loss_functions/sse_loss.hpp
index a06c7123b8..fc4f941620 100644
--- a/src/mlpack/methods/xgboost/loss_functions/sse_loss.hpp
+++ b/src/mlpack/methods/xgboost/loss_functions/sse_loss.hpp
@@ -92,6 +92,33 @@ class SSELoss
   {
     return - Gradients(observed, f);
   }
+
+  /**
+   * Returns the output value for the leaf in the tree.
+   */
+  template<typename VecType>
+  typename VecType::elem_type
+  OutputValue(const VecType& gradients, const VecType& hessians,
+              const double lambda)
+  {
+    return - arma::accu(gradients) / (arma::accu(hessians) + lambda);
+  }
+
+  /**
+   * Calculates the similarity score for evaluating the splits.
+   */
+  template<typename VecType>
+  double SimilarityScore(const VecType& observed, const VecType& residuals,
+      const size_t begin, const size_t end, const double lambda)
+  {
+    VecType gradients = Gradients(observed.subvec(begin, end),
+        residuals.subvec(begin, end));
+    VecType hessians = Hessians(observed.subvec(begin, end),
+        residuals.subvec(begin, end));
+
+    return std::pow(arma::accu(gradients), 2) /
+        (arma::accu(hessians) + lambda);
+  }
 }
 
 } // namespace ensemble