feat: change param names and add tests

Author: TomeHirata

dte_adj/__init__.py

@@ -24,7 +24,7 @@ def predict_dte(
         control_treatment_arm: int,
         outcomes: np.ndarray,
         alpha: float = 0.05,
-        variance_type="pointwise",
+        variance_type="moment",
         n_bootstrap=500,
     ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
         """Compute DTE based on the estimator for the distribution function.
@@ -34,7 +34,7 @@ def predict_dte(
             control_treatment_arm (int): The index of the treatment arm of the control group.
             outcomes (np.ndarray): Scalar values to be used for computing the cumulative distribution.
             alpha (float, optional): Significance level of the confidence band. Defaults to 0.05.
-            variance_type (str, optional): Variance type to be used to compute confidence intervals. Available values are pointwise, analytic, and uniform.
+            variance_type (str, optional): Variance type to be used to compute confidence intervals. Available values are moment, analytic, and uniform.
             n_bootstrap (int, optional): Number of bootstrap samples. Defaults to 500.
 
         Returns:
@@ -196,12 +196,18 @@ def _compute_expected_qtes(
     ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
         """Compute expected QTEs."""
         result = np.zeros(quantiles.shape)
+        treatment_cumulative = self.predict(
+            np.full(self.outcome.shape, target_treatment_arm), self.outcome
+        )
+        control_cumulative = self.predict(
+            np.full(self.outcome.shape, control_treatment_arm), self.outcome
+        )
         for i, q in enumerate(quantiles):
-            treatment_quantile = self.outcome[target_treatment_arm][
-                math.floor(self.outcome[i].shape[0] * q)
+            treatment_quantile = treatment_cumulative[
+                math.floor(treatment_cumulative.shape[0] * q)
             ]
-            control_quantile = self.outcome[control_treatment_arm][
-                math.floor(self.outcome[i].shape[0] * q)
+            control_quantile = control_cumulative[
+                math.floor(control_cumulative.shape[0] * q)
             ]
             result[i] = treatment_quantile - control_quantile
 
@@ -481,7 +487,7 @@ def compute_dte_confidence_intervals(
     ind_target: int,
     ind_control: int,
     alpha: 0.05,
-    variance_type="pointwise",
+    variance_type="moment",
     n_bootstrap=500,
 ):
     """Computes the confidence intervals of DTE.
@@ -497,7 +503,7 @@ def compute_dte_confidence_intervals(
         ind_target (int): Index of the target treatment indicator.
         ind_control (int): Index of the control treatment indicator.
         alpha (float, optional): Significance level of the confidence band. Defaults to 0.05.
-        variance_type (str, optional): Variance type to be used to compute confidence intervals. Available values are pointwise, analytic, and uniform.
+        variance_type (str, optional): Variance type to be used to compute confidence intervals. Available values are moment, analytic, and uniform.
         n_bootstrap (int, optional): Number of bootstrap samples. Defaults to 500.
 
     Returns:
@@ -530,7 +536,7 @@ def compute_dte_confidence_intervals(
 
     omega = (influence_function**2).mean(axis=0)
 
-    if variance_type == "pointwise":
+    if variance_type == "moment":
         vec_dte_lower_moment = vec_dte + norm.ppf(alpha / 2) * np.sqrt(omega / num_obs)
         vec_dte_upper_moment = vec_dte + norm.ppf(1 - alpha / 2) * np.sqrt(
             omega / num_obs

example/example.ipynb

@@ -0,0 +1,43 @@
+import unittest
+import numpy as np
+from dte_adj import AdjustedDistributionEstimator
+from unittest.mock import MagicMock
+
+
+class TestAdjustedEstimator(unittest.TestCase):
+    def test_prediction_success(self):
+        # TODO!
+        return
+
+    def test_prediction_fail_before_fit(self):
+        # Arrange
+        D = np.zeros(20)
+        D[:10] = 1
+        Y = np.arange(20)
+        base_model = MagicMock()
+        subject = AdjustedDistributionEstimator(base_model)
+
+        # Act, Assert
+        with self.assertRaises(RuntimeError) as cm:
+            subject.predict(D, Y)
+        self.assertEqual(
+            str(cm.exception),
+            "This estimator has not been trained yet. Please call fit first",
+        )
+
+    def test_fit_fail_invalid_input(self):
+        # Arrange
+        X = np.arange(20)
+        D = np.zeros(10)
+        D[:10] = 1
+        Y = np.arange(20)
+        base_model = MagicMock()
+        subject = AdjustedDistributionEstimator(base_model)
+
+        # Act, Assert
+        with self.assertRaises(RuntimeError) as cm:
+            subject.fit(X, D, Y)
+        self.assertEqual(
+            str(cm.exception),
+            "The shape of confounding and treatment_arm should be same",
+        )