Vasilis/drlearner

doc/reference.rst

@@ -9,6 +9,7 @@ Public Module Reference
     econml.deepiv
     econml.dgp
     econml.dml
+    econml.drlearner
     econml.inference
     econml.ortho_forest
     econml.selective_regularization

econml/_ortho_learner.py

@@ -524,7 +524,7 @@ def const_marginal_effect_interval(self, X=None, *, alpha=0.1):
         return super().const_marginal_effect_interval(X, alpha=alpha)
     const_marginal_effect_interval.__doc__ = LinearCateEstimator.const_marginal_effect_interval.__doc__
 
-    def effect_interval(self, X=None, T0=0, T1=1, *, alpha=0.1):
+    def effect_interval(self, X=None, *, T0=0, T1=1, alpha=0.1):
         self._check_fitted_dims(X)
         return super().effect_interval(X, T0=T0, T1=T1, alpha=alpha)
     effect_interval.__doc__ = LinearCateEstimator.effect_interval.__doc__

econml/cate_estimator.py

@@ -11,7 +11,7 @@
 from .bootstrap import BootstrapEstimator
 from .inference import BootstrapInference
 from .utilities import tensordot, ndim, reshape, shape
-from .inference import StatsModelsInference
+from .inference import StatsModelsInference, StatsModelsInferenceDiscrete
 
 
 class BaseCateEstimator(metaclass=abc.ABCMeta):
@@ -173,10 +173,50 @@ def call(self, *args, **kwargs):
 
     @_defer_to_inference
     def effect_interval(self, X=None, *, T0=0, T1=1, alpha=0.1):
+        """ Confidence intervals for the quantities :math:`\\tau(X, T0, T1)` produced
+        by the model. Available only when :code:`inference` is not :code:`None`, when
+        calling the fit method.
+
+        Parameters
+        ----------
+        X: optional (m, d_x) matrix
+            Features for each sample
+        T0: optional (m, d_t) matrix or vector of length m (Default=0)
+            Base treatments for each sample
+        T1: optional (m, d_t) matrix or vector of length m (Default=1)
+            Target treatments for each sample
+        alpha: optional float in [0, 1] (Default=0.1)
+            The overall level of confidence of the reported interval.
+            The alpha/2, 1-alpha/2 confidence interval is reported.
+
+        Returns
+        -------
+        lower, upper : tuple(type of effect(X, T0, T1), type of effect(X, T0, T1))
+            The lower and the upper bounds of the confidence interval for each quantity.
+        """
         pass
 
     @_defer_to_inference
     def marginal_effect_interval(self, T, X=None, *, alpha=0.1):
+        """ Confidence intervals for the quantities :math:`\\partial \\tau(T, X)` produced
+        by the model. Available only when :code:`inference` is not :code:`None`, when
+        calling the fit method.
+
+        Parameters
+        ----------
+        T: (m, d_t) matrix
+            Base treatments for each sample
+        X: optional (m, d_x) matrix or None (Default=None)
+            Features for each sample
+        alpha: optional float in [0, 1] (Default=0.1)
+            The overall level of confidence of the reported interval.
+            The alpha/2, 1-alpha/2 confidence interval is reported.
+
+        Returns
+        -------
+        lower, upper : tuple(type of marginal_effect(T, X), type of marginal_effect(T, X))
+            The lower and the upper bounds of the confidence interval for each quantity.
+        """
         pass
 
 
@@ -279,9 +319,27 @@ def marginal_effect_interval(self, T, X=None, *, alpha=0.1):
         effs = self.const_marginal_effect_interval(X=X, alpha=alpha)
         return tuple(np.repeat(eff, shape(T)[0], axis=0) if X is None else eff
                      for eff in effs)
+    marginal_effect_interval.__doc__ = BaseCateEstimator.marginal_effect_interval.__doc__
 
     @BaseCateEstimator._defer_to_inference
     def const_marginal_effect_interval(self, X=None, *, alpha=0.1):
+        """ Confidence intervals for the quantities :math:`\\theta(X)` produced
+        by the model. Available only when :code:`inference` is not :code:`None`, when
+        calling the fit method.
+
+        Parameters
+        ----------
+        X: optional (m, d_x) matrix or None (Default=None)
+            Features for each sample
+        alpha: optional float in [0, 1] (Default=0.1)
+            The overall level of confidence of the reported interval.
+            The alpha/2, 1-alpha/2 confidence interval is reported.
+
+        Returns
+        -------
+        lower, upper : tuple(type of const_marginal_effect(X), type of const_marginal_effect(X))
+            The lower and the upper bounds of the confidence interval for each quantity.
+        """
         pass
 
 
@@ -314,7 +372,7 @@ def _expand_treatments(self, X=None, *Ts):
         return (X,) + tuple(outTs)
 
     # override effect to set defaults, which works with the new definition of _expand_treatments
-    def effect(self, X=None, T0=0, T1=1):
+    def effect(self, X=None, *, T0=0, T1=1):
         # NOTE: don't explicitly expand treatments here, because it's done in the super call
         return super().effect(X, T0=T0, T1=T1)
     effect.__doc__ = BaseCateEstimator.effect.__doc__
@@ -348,3 +406,95 @@ def coef__interval(self, *, alpha=0.1):
     @BaseCateEstimator._defer_to_inference
     def intercept__interval(self, *, alpha=0.1):
         pass
+
+
+class StatsModelsCateEstimatorDiscreteMixin(BaseCateEstimator):
+
+    def _get_inference_options(self):
+        # add statsmodels to parent's options
+        options = super()._get_inference_options()
+        options.update(statsmodels=StatsModelsInferenceDiscrete)
+        return options
+
+    @property
+    @abc.abstractmethod
+    def statsmodels(self):
+        pass
+
+    def coef_(self, T):
+        """ The coefficients in the linear model of the constant marginal treatment
+        effect associated with treatment T.
+
+        Parameters
+        ----------
+        T: alphanumeric
+            The input treatment for which we want the coefficients.
+
+        Returns
+        -------
+        coef_: (n_x,) or (n_y, n_x) array like
+            Where n_x is the number of features that enter the final model (either the
+            dimension of X or the dimension of featurizer.fit_transform(X) if the CATE
+            estimator has a featurizer.)
+        """
+        _, T = self._expand_treatments(None, T)
+        ind = (T @ np.arange(1, T.shape[1] + 1)).astype(int)[0] - 1
+        return self.statsmodels[ind].coef_
+
+    @property
+    def intercept_(self, T):
+        """ The intercept in the linear model of the constant marginal treatment
+        effect associated with treatment T.
+
+        Parameters
+        ----------
+        T: alphanumeric
+            The input treatment for which we want the coefficients.
+
+        Returns
+        -------
+        intercept_: float or (n_y,) array like
+        """
+        _, T = self._expand_treatments(None, T)
+        ind = (T @ np.arange(1, T.shape[1] + 1)).astype(int)[0] - 1
+        return self.statsmodels[ind].intercept_
+
+    @BaseCateEstimator._defer_to_inference
+    def coef__interval(self, T, *, alpha=0.1):
+        """ The confidence interval for the coefficients in the linear model of the
+        constant marginal treatment effect associated with treatment T.
+
+        Parameters
+        ----------
+        T: alphanumeric
+            The input treatment for which we want the coefficients.
+        alpha: optional float in [0, 1] (Default=0.1)
+            The overall level of confidence of the reported interval.
+            The alpha/2, 1-alpha/2 confidence interval is reported.
+
+        Returns
+        -------
+        lower, upper: tuple(type of coef_, type of coef_)
+            The lower and upper bounds of the confidence interval for each quantity.
+        """
+        pass
+
+    @BaseCateEstimator._defer_to_inference
+    def intercept__interval(self, T, *, alpha=0.1):
+        """ The intercept in the linear model of the constant marginal treatment
+        effect associated with treatment T.
+
+        Parameters
+        ----------
+        T: alphanumeric
+            The input treatment for which we want the coefficients.
+        alpha: optional float in [0, 1] (Default=0.1)
+            The overall level of confidence of the reported interval.
+            The alpha/2, 1-alpha/2 confidence interval is reported.
+
+        Returns
+        -------
+        lower, upper: tuple(type of intercept_, type of intercept_)
+            The lower and upper bounds of the confidence interval.
+        """
+        pass