Enable internal bootstrap on all estimators

Author: kbattocchi

econml/cate_estimator.py

@@ -5,14 +5,63 @@
 
 import abc
 import numpy as np
+from .bootstrap import BootstrapEstimator
+from .inference import BootstrapOptions
 from .utilities import tensordot, ndim, reshape, shape
 
 
-class BaseCateEstimator:
-    """Base class for all CATE estimators in this package."""
+class BaseCateEstimator(metaclass=abc.ABCMeta):
+    """
+    Base class for all CATE estimators in this package.
+
+    Parameters
+    ----------
+    inference: string, inference method, or None
+        Method for performing inference.  All estimators support 'bootstrap'
+        (or an instance of `BootstrapOptions`), some support other methods as well.
+
+    """
+
+    _inference_options = {'bootstrap': BootstrapOptions()}
+    _bootstrap_whitelist = {'effect', 'marginal_effect'}
+
+    @abc.abstractmethod
+    def __init__(self, inference):
+        """
+        Initialize the estimator.
+
+        All subclass overrides should complete with a call to this method on the super class,
+        since it enables bootstrapping.
+
+        """
+        if inference in self._inference_options:
+            inference = self._inference_options[inference]
+        if isinstance(inference, BootstrapOptions):
+            # Note that fit (and other methods) check for the presence of a _bootstrap attribute
+            # to determine whether to delegate to that object or not;
+            # The clones wrapped inside the BootstrapEstimator will not have that attribute since
+            # it's assigned *after* creating the estimator
+            self._bootstrap = BootstrapEstimator(self, inference.n_bootstrap_samples, inference.n_jobs)
+        self._inference = inference
+
+    def __getattr__(self, name):
+        suffix = '_interval'
+        if name.endswith(suffix) and name[: - len(suffix)] in self._bootstrap_whitelist:
+            if hasattr(self, '_bootstrap'):
+                return getattr(self._bootstrap, name)
+            else:
+                raise AttributeError('\'%s\' object does not support attribute \'%s\'; '
+                                     'consider passing inference=\'bootstrap\' when initializing'
+                                     % (type(self).__name__, name))
+        else:
+            raise AttributeError('\'%s\' object has no attribute \'%s\''
+                                 % (type(self).__name__, name))
 
     @abc.abstractmethod
-    def fit(self, Y, T, X=None, W=None, Z=None):
+    def _fit_impl(self, Y, T, X=None, W=None, Z=None):
+        pass
+
+    def fit(self, *args, **kwargs):
         """
         Estimate the counterfactual model from data, i.e. estimates functions τ(·,·,·), ∂τ(·,·).
 
@@ -37,7 +86,9 @@ def fit(self, Y, T, X=None, W=None, Z=None):
         self
 
         """
-        pass
+        if hasattr(self, '_bootstrap'):
+            self._bootstrap.fit(*args, **kwargs)
+        return self._fit_impl(*args, **kwargs)
 
     @abc.abstractmethod
     def effect(self, X=None, T0=0, T1=1):
@@ -92,7 +143,20 @@ def marginal_effect(self, T, X=None):
 
 
 class LinearCateEstimator(BaseCateEstimator):
-    """Base class for all CATE estimators with linear treatment effects in this package."""
+    """
+    Base class for all CATE estimators with linear treatment effects in this package.
+
+    Parameters
+    ----------
+    inference: string, inference method, or None
+        Method for performing inference.  All estimators support 'bootstrap'
+        (or an instance of `BootstrapOptions`), some support other methods as well.
+
+    """
+
+    @abc.abstractmethod
+    def __init__(self, inference):
+        super().__init__(inference=inference)
 
     @abc.abstractmethod
     def const_marginal_effect(self, X=None):

econml/deepiv.py

@@ -268,13 +268,19 @@ class DeepIVEstimator(BaseCateEstimator):
     second_stage_options : dictionary, optional
         The keyword arguments to pass to Keras's `fit` method when training the second stage model.
         Defaults to `{"epochs": 100}`.
+
+    inference: string, inference method, or None
+        Method for performing inference.  This estimator supports 'bootstrap'
+        (or an instance of `BootstrapOptions`)
+
     """
 
     def __init__(self, n_components, m, h,
                  n_samples, use_upper_bound_loss=False, n_gradient_samples=0,
                  optimizer='adam',
                  first_stage_options={"epochs": 100},
-                 second_stage_options={"epochs": 100}):
+                 second_stage_options={"epochs": 100},
+                 inference=None):
         self._n_components = n_components
         self._m = m
         self._h = h
@@ -284,8 +290,9 @@ def __init__(self, n_components, m, h,
         self._optimizer = optimizer
         self._first_stage_options = first_stage_options
         self._second_stage_options = second_stage_options
+        super().__init__(inference=inference)
 
-    def fit(self, Y, T, X, Z):
+    def _fit_impl(self, Y, T, X, Z):
         """Estimate the counterfactual model from data.
 
         That is, estimate functions τ(·, ·, ·), ∂τ(·, ·).

econml/dml.py

@@ -50,18 +50,23 @@ class _RLearner(LinearCateEstimator):
         If RandomState instance, random_state is the random number generator;
         If None, the random number generator is the RandomState instance used
         by `np.random`.
+
+    inference: string, inference method, or None
+        Method for performing inference.  This estimator supports 'bootstrap'
+        (or an instance of `BootstrapOptions`)
     """
 
     def __init__(self, model_y, model_t, model_final,
-                 discrete_treatment, n_splits, random_state):
+                 discrete_treatment, n_splits, random_state, inference):
         self._models_y = [clone(model_y, safe=False) for _ in range(n_splits)]
         self._models_t = [clone(model_t, safe=False) for _ in range(n_splits)]
         self._model_final = clone(model_final, safe=False)
         self._n_splits = n_splits
         self._discrete_treatment = discrete_treatment
         self._random_state = check_random_state(random_state)
+        super().__init__(inference=inference)
 
-    def fit(self, Y, T, X=None, W=None):
+    def _fit_impl(self, Y, T, X=None, W=None):
         if X is None:
             X = np.ones((shape(Y)[0], 1))
         if W is None:
@@ -203,6 +208,10 @@ class _DMLCateEstimatorBase(_RLearner):
         If RandomState instance, random_state is the random number generator;
         If None, the random number generator is the RandomState instance used
         by `np.random`.
+
+    inference: string, inference method, or None
+        Method for performing inference.  This estimator supports 'bootstrap'
+        (or an instance of `BootstrapOptions`).
     """
 
     def __init__(self,
@@ -211,7 +220,8 @@ def __init__(self,
                  sparseLinear,
                  discrete_treatment,
                  n_splits,
-                 random_state):
+                 random_state,
+                 inference):
 
         class FirstStageWrapper:
             def __init__(self, model, is_Y):
@@ -274,7 +284,8 @@ def coef_(self):
                          model_final=FinalWrapper(),
                          discrete_treatment=discrete_treatment,
                          n_splits=n_splits,
-                         random_state=random_state)
+                         random_state=random_state,
+                         inference=inference)
 
     @property
     def coef_(self):
@@ -321,22 +332,28 @@ class DMLCateEstimator(_DMLCateEstimatorBase):
         If RandomState instance, random_state is the random number generator;
         If None, the random number generator is the RandomState instance used
         by `np.random`.
+
+    inference: string, inference method, or None
+        Method for performing inference.  This estimator supports 'bootstrap'
+        (or an instance of `BootstrapOptions`)
     """
 
     def __init__(self,
                  model_y, model_t, model_final=LinearRegression(fit_intercept=False),
                  featurizer=PolynomialFeatures(degree=1, include_bias=True),
                  discrete_treatment=False,
                  n_splits=2,
-                 random_state=None):
+                 random_state=None,
+                 inference=None):
         super().__init__(model_y=model_y,
                          model_t=model_t,
                          model_final=model_final,
                          featurizer=featurizer,
                          sparseLinear=False,
                          discrete_treatment=discrete_treatment,
                          n_splits=n_splits,
-                         random_state=random_state)
+                         random_state=random_state,
+                         inference=inference)
 
 
 class SparseLinearDMLCateEstimator(_DMLCateEstimatorBase):
@@ -376,22 +393,28 @@ class SparseLinearDMLCateEstimator(_DMLCateEstimatorBase):
         If RandomState instance, random_state is the random number generator;
         If None, the random number generator is the RandomState instance used
         by `np.random`.
+
+    inference: string, inference method, or None
+        Method for performing inference.  This estimator supports 'bootstrap'
+        (or an instance of `BootstrapOptions`)
     """
 
     def __init__(self,
                  linear_model_y=LassoCV(), linear_model_t=LassoCV(), model_final=LinearRegression(fit_intercept=False),
                  featurizer=PolynomialFeatures(degree=1, include_bias=True),
                  discrete_treatment=False,
                  n_splits=2,
-                 random_state=None):
+                 random_state=None,
+                 inference=None):
         super().__init__(model_y=linear_model_y,
                          model_t=linear_model_t,
                          model_final=model_final,
                          featurizer=featurizer,
                          sparseLinear=True,
                          discrete_treatment=discrete_treatment,
                          n_splits=n_splits,
-                         random_state=random_state)
+                         random_state=random_state,
+                         inference=inference)
 
 
 class KernelDMLCateEstimator(DMLCateEstimator):
@@ -421,15 +444,19 @@ class KernelDMLCateEstimator(DMLCateEstimator):
     n_splits: int, optional (default is 2)
         The number of splits to use when fitting the first-stage models.
 
-     random_state: int, RandomState instance or None, optional (default=None)
+    random_state: int, RandomState instance or None, optional (default=None)
         If int, random_state is the seed used by the random number generator;
         If RandomState instance, random_state is the random number generator;
         If None, the random number generator is the RandomState instance used
         by `np.random`.
-   """
+
+    inference: string, inference method, or None
+        Method for performing inference.  This estimator supports 'bootstrap'
+        (or an instance of `BootstrapOptions`)
+    """
 
     def __init__(self, model_y, model_t, model_final=LinearRegression(fit_intercept=False),
-                 dim=20, bw=1.0, n_splits=2, random_state=None):
+                 dim=20, bw=1.0, n_splits=2, random_state=None, inference=None):
         class RandomFeatures(TransformerMixin):
             def fit(innerself, X):
                 innerself.omegas = self._random_state.normal(0, 1 / bw, size=(shape(X)[1], dim))
@@ -440,4 +467,5 @@ def transform(innerself, X):
                 return np.sqrt(2 / dim) * np.cos(np.matmul(X, innerself.omegas) + innerself.biases)
 
         super().__init__(model_y=model_y, model_t=model_t, model_final=model_final,
-                         featurizer=RandomFeatures(), n_splits=n_splits, random_state=random_state)
+                         featurizer=RandomFeatures(), n_splits=n_splits, random_state=random_state,
+                         inference=inference)

econml/inference.py

@@ -0,0 +1,35 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Options for performing inference in estimators."""
+
+
+class BootstrapOptions:
+    """
+    Wrapper storing bootstrap options.
+
+    This class can be used for inference with any CATE estimator.
+
+    Parameters
+    ----------
+    n_bootstrap_samples : int, optional (default 100)
+        How many draws to perform.
+
+    n_jobs: int, optional (default -1)
+        The maximum number of concurrently running jobs, as in joblib.Parallel.
+
+    """
+
+    def __init__(self, n_bootstrap_samples=100, n_jobs=-1):
+        self._n_bootstrap_samples = n_bootstrap_samples
+        self._n_jobs = n_jobs
+
+    @property
+    def n_bootstrap_samples(self):
+        """Get how many draws to perform."""
+        return self._n_bootstrap_samples
+
+    @property
+    def n_jobs(self):
+        """Get the maximum number of concurrently running jobs, as in joblib.Parallel."""
+        return self._n_jobs