Avoid double impurity calculation

Author: gdaiha

econml/tree/_criterion.pyx

@@ -209,7 +209,9 @@ cdef class Criterion:
         return (- self.weighted_n_right * impurity_right
                 - self.weighted_n_left * impurity_left)
 
-    cdef double impurity_improvement(self, double impurity) nogil:
+    cdef double impurity_improvement(self, float64_t impurity_parent,
+                                     float64_t impurity_left,
+                                     float64_t impurity_right) nogil:
         """Compute the improvement in impurity
         This method computes the improvement in impurity when a split occurs.
         The weighted impurity improvement equation is the following:
@@ -218,22 +220,24 @@ cdef class Criterion:
         where N is the total number of samples, N_t is the number of samples
         at the current node, N_t_L is the number of samples in the left child,
         and N_t_R is the number of samples in the right child,
+
         Parameters
         ----------
-        impurity : double
-            The initial impurity of the node before the split
+        impurity_parent : float64_t
+            The initial impurity of the parent node before the split
+
+        impurity_left : float64_t
+            The impurity of the left child
+
+        impurity_right : float64_t
+            The impurity of the right child
+
         Return
         ------
         double : improvement in impurity after the split occurs
         """
-
-        cdef double impurity_left
-        cdef double impurity_right
-
-        self.children_impurity(&impurity_left, &impurity_right)
-
         return ((self.weighted_n_node_samples / self.weighted_n_samples) *
-                (impurity - (self.weighted_n_right / 
+                (impurity_parent - (self.weighted_n_right / 
                              self.weighted_n_node_samples * impurity_right)
                           - (self.weighted_n_left / 
                              self.weighted_n_node_samples * impurity_left)))

econml/tree/_splitter.pyx

@@ -615,11 +615,6 @@ cdef class BestSplitter(Splitter):
             if self.honest:
                 self.criterion_val.reset()
                 self.criterion_val.update(best.pos_val)
-            # Calculate a more accurate version of impurity improvement using the input baseline impurity
-            # passed here by the TreeBuilder. The TreeBuilder uses the proxy_node_impurity() to calculate
-            # this baseline if self.is_children_impurity_proxy(), else uses the call to children_impurity()
-            # on the parent node, when that node was split.
-            best.improvement = self.criterion.impurity_improvement(impurity)
             # if we need children impurities by the builder, then we populate these entries
             # otherwise, we leave them blank to avoid the extra computation.
             if not self.is_children_impurity_proxy():
@@ -630,6 +625,16 @@ cdef class BestSplitter(Splitter):
                 else:
                     best.impurity_left_val = best.impurity_left
                     best.impurity_right_val = best.impurity_right
+                
+                # Calculate a more accurate version of impurity improvement using the input baseline impurity
+                # passed here by the TreeBuilder. The TreeBuilder uses the proxy_node_impurity() to calculate
+                # this baseline if self.is_children_impurity_proxy(), else uses the call to children_impurity()
+                # on the parent node, when that node was split.
+                best.improvement = self.criterion.impurity_improvement(impurity, 
+                best.impurity_left, best.impurity_right)
+
+                if self.honest:
+                    
 
         # Respect invariant for constant features: the original order of
         # element in features[:n_known_constants] must be preserved for sibling