Add in an association function

Author: HackToHell

sklearn/tree/_criterion.pxd

@@ -62,6 +62,8 @@ cdef class Criterion:
     cdef double node_impurity(self) nogil
     cdef void children_impurity(self, double* impurity_left,
                                 double* impurity_right) nogil
+    cdef void children_af(self, double* impurity_left,
+                                double* impurity_right) nogil
     cdef void node_value(self, double* dest) nogil
     cdef double impurity_improvement(self, double impurity) nogil
     cdef double proxy_impurity_improvement(self) nogil

sklearn/tree/_criterion.pyx

@@ -142,6 +142,10 @@ cdef class Criterion:
 
         pass
 
+    cdef void children_af(self, double* impurity_left,
+                                double* impurity_right) nogil:
+        pass
+
     cdef void node_value(self, double* dest) nogil:
         """Placeholder for storing the node value.
 
@@ -176,6 +180,7 @@ cdef class Criterion:
 
     cdef double impurity_improvement(self, double impurity) nogil:
         """Compute the improvement in impurity
+        gain
 
         This method computes the improvement in impurity when a split occurs.
         The weighted impurity improvement equation is the following:
@@ -199,14 +204,17 @@ cdef class Criterion:
 
         cdef double impurity_left
         cdef double impurity_right
+        cdef double af_left
+        cdef double af_right
 
         self.children_impurity(&impurity_left, &impurity_right)
+        self.children_af(&af_left,&af_right)
 
         return ((self.weighted_n_node_samples / self.weighted_n_samples) *
                 (impurity - (self.weighted_n_right / 
-                             self.weighted_n_node_samples * impurity_right)
+                             self.weighted_n_node_samples * impurity_right * af_right)
                           - (self.weighted_n_left / 
-                             self.weighted_n_node_samples * impurity_left)))
+                             self.weighted_n_node_samples * impurity_left * af_left)))
 
 
 cdef class ClassificationCriterion(Criterion):
@@ -597,6 +605,53 @@ cdef class Entropy(ClassificationCriterion):
         impurity_left[0] = entropy_left / self.n_outputs
         impurity_right[0] = entropy_right / self.n_outputs
 
+    cdef void children_af(self, double* impurity_left,
+                                double* impurity_right) nogil:
+        """Evaluate the impurity in children nodes
+
+        i.e. the impurity of the left child (samples[start:pos]) and the
+        impurity the right child (samples[pos:end]).
+
+        Parameters
+        ----------
+        impurity_left : double pointer
+            The memory address to save the impurity of the left node
+        impurity_right : double pointer
+            The memory address to save the impurity of the right node
+        """
+
+        cdef SIZE_t* n_classes = self.n_classes
+        cdef double* sum_left = self.sum_left
+        cdef double* sum_right = self.sum_right
+        cdef double entropy_left = 0.0
+        cdef double entropy_right = 0.0
+        cdef double af_left = 0.0
+        cdef double af_right = 0.0
+        cdef double count_k
+        cdef double two = 2
+        cdef double one = 1
+        cdef SIZE_t k
+        cdef SIZE_t c
+
+        for k in range(self.n_outputs):
+            for c in range(n_classes[k]):
+                count_k = sum_left[c]
+                if count_k > 0.0:
+                    count_k /= self.weighted_n_left
+                    entropy_left -= count_k * log(count_k)
+                    af_left = fabs(count_k*two - one)
+
+                count_k = sum_right[c]
+                if count_k > 0.0:
+                    count_k /= self.weighted_n_right
+                    entropy_right -= count_k * log(count_k)
+                    af_right = fabs(count_k*two - one)
+
+            sum_left += self.sum_stride
+            sum_right += self.sum_stride
+
+        impurity_left[0] = af_left / self.n_outputs
+        impurity_right[0] = af_right / self.n_outputs
 
 cdef class Gini(ClassificationCriterion):
     """Gini Index impurity criterion.