#775: Support multiple treatments in CausalTreeRegressor and CausalRandomForestRegressor (#852)

* Add outcome vector y preparation for multiple treatment groups
* Update CausalTreeRegressor class
* Add multiple groups support in cython part of causal trees building
* Add min_group_size parameter to control tree building
* Update python part of causal trees, keep consistent namings for tree builder
* Add an option to pass custom matplotlib axes in charts
* Update Jupyter notebooks with causal trees and forests

* Keep consistent codestyle with black

* Update validate_data arguments support for sklearn <1.6 & >=1.6

* Extend causal tree and forest tests for multiple treatment groups

* Keep consistent codestyle with black

* Fix description in causal trees notebok

* Keep consistent codestyle in tests

* Remove unused cython var in criterion header

* Add separate function for check_y_params

Author: alexander-pv

causalml/inference/tree/_tree/_criterion.pxd

@@ -24,7 +24,6 @@ cdef class Criterion:
 
     # Internal structures
     cdef const float64_t[:, ::1] y         # Values of y
-    cdef const int32_t[:] treatment         # Treatment assignment: 1 for treatment, 0 for control
     cdef const float64_t[:] sample_weight  # Sample weights
 
     cdef const intp_t[:] sample_indices    # Sample indices in X, y
@@ -50,7 +49,6 @@ cdef class Criterion:
     cdef int init(
         self,
         const float64_t[:, ::1] y,
-        const int32_t[:] treatment,
         const float64_t[:] sample_weight,
         float64_t weighted_n_samples,
         const intp_t[:] sample_indices,

causalml/inference/tree/_tree/_criterion.pyx

@@ -49,7 +49,6 @@ cdef class Criterion:
     cdef int init(
         self,
         const float64_t[:, ::1] y,
-        const int32_t[:] treatment,
         const float64_t[:] sample_weight,
         float64_t weighted_n_samples,
         const intp_t[:] sample_indices,
@@ -357,7 +356,6 @@ cdef class ClassificationCriterion(Criterion):
     cdef int init(
         self,
         const float64_t[:, ::1] y,
-        const int32_t[:] treatment,
         const float64_t[:] sample_weight,
         float64_t weighted_n_samples,
         const intp_t[:] sample_indices,
@@ -871,7 +869,6 @@ cdef class RegressionCriterion(Criterion):
     cdef int init(
         self,
         const float64_t[:, ::1] y,
-        const int32_t[:] treatment,
         const float64_t[:] sample_weight,
         float64_t weighted_n_samples,
         const intp_t[:] sample_indices,
@@ -1250,7 +1247,6 @@ cdef class MAE(RegressionCriterion):
     cdef int init(
         self,
         const float64_t[:, ::1] y,
-        const int32_t[:] treatment,
         const float64_t[:] sample_weight,
         float64_t weighted_n_samples,
         const intp_t[:] sample_indices,

causalml/inference/tree/_tree/_splitter.pxd

@@ -70,7 +70,6 @@ cdef class Splitter:
     #   +1: monotonic increase
     cdef const int8_t[:] monotonic_cst
     cdef bint with_monotonic_cst
-    cdef const int32_t[:] treatment
     cdef const float64_t[:] sample_weight
 
     # The samples vector `samples` is maintained by the Splitter object such
@@ -94,7 +93,6 @@ cdef class Splitter:
         self,
         object X,
         const float64_t[:, ::1] y,
-        const int32_t[:] treatment,
         const float64_t[:] sample_weight,
         const unsigned char[::1] missing_values_in_feature_mask,
     ) except -1

causalml/inference/tree/_tree/_splitter.pyx

@@ -125,7 +125,6 @@ cdef class Splitter:
         self,
         object X,
         const float64_t[:, ::1] y,
-        const int32_t[:] treatment,
         const float64_t[:] sample_weight,
         const unsigned char[::1] missing_values_in_feature_mask,
     ) except -1:
@@ -145,9 +144,6 @@ cdef class Splitter:
             This is the vector of targets, or true labels, for the samples represented
             as a Cython memoryview.
 
-        treatment : ndarray, dtype=int32_t
-            The treatment labels for each sample, represented as a Cython memoryview.
-
         sample_weight : ndarray, dtype=float64_t
             The weights of the samples, where higher weighted samples are fit
             closer than lower weight samples. If not provided, all samples
@@ -194,7 +190,6 @@ cdef class Splitter:
 
         self.y = y
 
-        self.treatment = treatment
         self.sample_weight = sample_weight
         if missing_values_in_feature_mask is not None:
             self.criterion.init_sum_missing()
@@ -226,7 +221,6 @@ cdef class Splitter:
 
         self.criterion.init(
             self.y,
-            self.treatment,
             self.sample_weight,
             self.weighted_n_samples,
             self.samples,
@@ -1515,11 +1509,10 @@ cdef class BestSplitter(Splitter):
         self,
         object X,
         const float64_t[:, ::1] y,
-        const int32_t[:] treatment,
         const float64_t[:] sample_weight,
         const unsigned char[::1] missing_values_in_feature_mask,
     ) except -1:
-        Splitter.init(self, X, y, treatment, sample_weight, missing_values_in_feature_mask)
+        Splitter.init(self, X, y, sample_weight, missing_values_in_feature_mask)
         self.partitioner = DensePartitioner(
             X, self.samples, self.feature_values, missing_values_in_feature_mask
         )
@@ -1546,11 +1539,10 @@ cdef class BestSparseSplitter(Splitter):
         self,
         object X,
         const float64_t[:, ::1] y,
-        const int32_t[:] treatment,
         const float64_t[:] sample_weight,
         const unsigned char[::1] missing_values_in_feature_mask,
     ) except -1:
-        Splitter.init(self, X, y, treatment, sample_weight, missing_values_in_feature_mask)
+        Splitter.init(self, X, y, sample_weight, missing_values_in_feature_mask)
         self.partitioner = SparsePartitioner(
             X, self.samples, self.n_samples, self.feature_values, missing_values_in_feature_mask
         )
@@ -1577,11 +1569,10 @@ cdef class RandomSplitter(Splitter):
         self,
         object X,
         const float64_t[:, ::1] y,
-        const int32_t[:] treatment,
         const float64_t[:] sample_weight,
         const unsigned char[::1] missing_values_in_feature_mask,
     ) except -1:
-        Splitter.init(self, X, y, treatment, sample_weight, missing_values_in_feature_mask)
+        Splitter.init(self, X, y, sample_weight, missing_values_in_feature_mask)
         self.partitioner = DensePartitioner(
             X, self.samples, self.feature_values, missing_values_in_feature_mask
         )
@@ -1608,11 +1599,10 @@ cdef class RandomSparseSplitter(Splitter):
         self,
         object X,
         const float64_t[:, ::1] y,
-        const int32_t[:] treatment,
         const float64_t[:] sample_weight,
         const unsigned char[::1] missing_values_in_feature_mask,
     ) except -1:
-        Splitter.init(self, X, y, treatment, sample_weight, missing_values_in_feature_mask)
+        Splitter.init(self, X, y, sample_weight, missing_values_in_feature_mask)
         self.partitioner = SparsePartitioner(
             X, self.samples, self.n_samples, self.feature_values, missing_values_in_feature_mask
         )

causalml/inference/tree/_tree/_tree.pxd

@@ -106,20 +106,19 @@ cdef class TreeBuilder:
     # This class controls the various stopping criteria and the node splitting
     # evaluation order, e.g. depth-first or best-first.
 
-    cdef Splitter splitter              # Splitting algorithm
+    cdef Splitter splitter                 # Splitting algorithm
 
-    cdef intp_t min_samples_split       # Minimum number of samples in an internal node
-    cdef intp_t min_samples_leaf        # Minimum number of samples in a leaf
+    cdef intp_t min_samples_split          # Minimum number of samples in an internal node
+    cdef intp_t min_samples_leaf           # Minimum number of samples in a leaf
     cdef float64_t min_weight_leaf         # Minimum weight in a leaf
-    cdef intp_t max_depth               # Maximal tree depth
+    cdef intp_t max_depth                  # Maximal tree depth
     cdef float64_t min_impurity_decrease   # Impurity threshold for early stopping
 
     cpdef build(
         self,
         Tree tree,
         object X,
         const float64_t[:, ::1] y,
-        const int32_t[:] treatment,
         const float64_t[:] sample_weight=*,
         const unsigned char[::1] missing_values_in_feature_mask=*,
     )
@@ -128,7 +127,6 @@ cdef class TreeBuilder:
         self,
         object X,
         const float64_t[:, ::1] y,
-        const int32_t[:] treatment,
         const float64_t[:] sample_weight,
     )
 

causalml/inference/tree/_tree/_tree.pyx

@@ -106,7 +106,6 @@ cdef class TreeBuilder:
         Tree tree,
         object X,
         const float64_t[:, ::1] y,
-        const int32_t[:] treatment,
         const float64_t[:] sample_weight=None,
         const unsigned char[::1] missing_values_in_feature_mask=None,
     ):
@@ -117,7 +116,6 @@ cdef class TreeBuilder:
         self,
         object X,
         const float64_t[:, ::1] y,
-        const int32_t[:] treatment,
         const float64_t[:] sample_weight,
     ):
         """Check input dtype, layout and format"""
@@ -142,9 +140,6 @@ cdef class TreeBuilder:
         if y.base.dtype != DOUBLE or not y.base.flags.contiguous:
             y = np.ascontiguousarray(y, dtype=DOUBLE)
 
-        if treatment.base.dtype != INT or not treatment.base.flags.contiguous:
-            treatment = np.ascontiguousarray(treatment, dtype=INT)
-
         if (
             sample_weight is not None and
             (
@@ -154,15 +149,16 @@ cdef class TreeBuilder:
         ):
             sample_weight = np.asarray(sample_weight, dtype=DOUBLE, order="C")
 
-        return X, y, treatment, sample_weight
+        return X, y, sample_weight
 
 # Depth first builder ---------------------------------------------------------
 cdef class DepthFirstTreeBuilder(TreeBuilder):
     """Build a decision tree in depth-first fashion."""
 
     def __cinit__(self, Splitter splitter, intp_t min_samples_split,
                   intp_t min_samples_leaf, float64_t min_weight_leaf,
-                  intp_t max_depth, float64_t min_impurity_decrease):
+                  intp_t max_depth, float64_t min_impurity_decrease,
+                  *args, **kwargs):
         self.splitter = splitter
         self.min_samples_split = min_samples_split
         self.min_samples_leaf = min_samples_leaf
@@ -175,14 +171,13 @@ cdef class DepthFirstTreeBuilder(TreeBuilder):
         Tree tree,
         object X,
         const float64_t[:, ::1] y,
-        const int32_t[:] treatment,
         const float64_t[:] sample_weight=None,
         const unsigned char[::1] missing_values_in_feature_mask=None,
     ):
         """Build a decision tree from the training set (X, y)."""
 
         # check input
-        X, y, treatment, sample_weight = self._check_input(X, y, treatment, sample_weight)
+        X, y, sample_weight = self._check_input(X, y, sample_weight)
 
         # Initial capacity
         cdef intp_t init_capacity
@@ -203,7 +198,7 @@ cdef class DepthFirstTreeBuilder(TreeBuilder):
         cdef float64_t min_impurity_decrease = self.min_impurity_decrease
 
         # Recursive partition (without actual recursion)
-        splitter.init(X, y, treatment, sample_weight, missing_values_in_feature_mask)
+        splitter.init(X, y, sample_weight, missing_values_in_feature_mask)
 
         cdef intp_t start
         cdef intp_t end
@@ -399,7 +394,8 @@ cdef class BestFirstTreeBuilder(TreeBuilder):
     def __cinit__(self, Splitter splitter, intp_t min_samples_split,
                   intp_t min_samples_leaf,  min_weight_leaf,
                   intp_t max_depth, intp_t max_leaf_nodes,
-                  float64_t min_impurity_decrease):
+                  float64_t min_impurity_decrease,
+                  *args, **kwargs):
         self.splitter = splitter
         self.min_samples_split = min_samples_split
         self.min_samples_leaf = min_samples_leaf
@@ -413,21 +409,20 @@ cdef class BestFirstTreeBuilder(TreeBuilder):
         Tree tree,
         object X,
         const float64_t[:, ::1] y,
-        const int32_t[:] treatment,
         const float64_t[:] sample_weight=None,
         const unsigned char[::1] missing_values_in_feature_mask=None,
     ):
         """Build a decision tree from the training set (X, y)."""
 
         # check input
-        X, y, treatment, sample_weight = self._check_input(X, y, treatment, sample_weight)
+        X, y, sample_weight = self._check_input(X, y, sample_weight)
 
         # Parameters
         cdef Splitter splitter = self.splitter
         cdef intp_t max_leaf_nodes = self.max_leaf_nodes
 
         # Recursive partition (without actual recursion)
-        splitter.init(X, y, treatment, sample_weight, missing_values_in_feature_mask)
+        splitter.init(X, y, sample_weight, missing_values_in_feature_mask)
 
         cdef vector[FrontierRecord] frontier
         cdef FrontierRecord record

causalml/inference/tree/causal/_builder.pxd

@@ -6,6 +6,6 @@
 
 from .._tree._tree cimport Node, Tree, TreeBuilder
 from .._tree._splitter cimport Splitter, SplitRecord
-from .._tree._tree cimport intp_t, int32_t, float64_t
+from .._tree._typedefs cimport intp_t, int32_t, int64_t, float32_t, float64_t
 from .._tree._tree cimport FrontierRecord, StackRecord
 from .._tree._tree cimport ParentInfo, _init_parent_record