feat: add bar chart

Author: TomeHirata

dte_adj/__init__.py

@@ -3,6 +3,7 @@
 from scipy.stats import norm
 import math
 from copy import deepcopy
+from .util import compute_confidence_intervals, find_le
 
 
 class DistributionFunctionMixin(object):
@@ -33,15 +34,15 @@ def predict_dte(
             target_treatment_arm (int): The index of the treatment arm of the treatment group.
             control_treatment_arm (int): The index of the treatment arm of the control group.
             locations (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.
+            alpha (float, optional): Significance level of the confidence bound. Defaults to 0.05.
             variance_type (str, optional): Variance type to be used to compute confidence intervals. Available values are moment, simple, and uniform.
             n_bootstrap (int, optional): Number of bootstrap samples. Defaults to 500.
 
         Returns:
             Tuple[np.ndarray, np.ndarray, np.ndarray]: A tuple containing:
                 - Expected DTEs
-                - Upper bands
-                - Lower bands
+                - Upper bounds
+                - Lower bounds
         """
         return self._compute_dtes(
             target_treatment_arm,
@@ -68,14 +69,14 @@ def predict_pte(
             control_treatment_arm (int): The index of the treatment arm of the control group.
             locations (np.ndarray): Scalar values to be used for computing the cumulative distribution.
             width (float): The width of each outcome interval.
-            alpha (float, optional): Significance level of the confidence band. Defaults to 0.05.
+            alpha (float, optional): Significance level of the confidence bound. Defaults to 0.05.
             variance_type (str, optional): Variance type to be used to compute confidence intervals. Available values are moment, simple, and uniform.
 
         Returns:
             Tuple[np.ndarray, np.ndarray, np.ndarray]: A tuple containing:
                 - Expected PTEs
-                - Upper bands
-                - Lower bands
+                - Upper bounds
+                - Lower bounds
         """
         return self._compute_ptes(
             target_treatment_arm,
@@ -102,14 +103,14 @@ def predict_qte(
             target_treatment_arm (int): The index of the treatment arm of the treatment group.
             control_treatment_arm (int): The index of the treatment arm of the control group.
             quantiles (np.ndarray, optional): Quantiles used for QTE. Defaults to [0.1 * i for i in range(1, 10)].
-            alpha (float, optional): Significance level of the confidence band. Defaults to 0.05.
+            alpha (float, optional): Significance level of the confidence bound. Defaults to 0.05.
             n_bootstrap (int, optional): Number of bootstrap samples. Defaults to 500.
 
         Returns:
             Tuple[np.ndarray, np.ndarray, np.ndarray]: A tuple containing:
                 - Expected QTEs
-                - Upper bands
-                - Lower bands
+                - Upper bounds
+                - Lower bounds
         """
         qte = self._compute_qtes(
             target_treatment_arm,
@@ -170,7 +171,7 @@ def _compute_dtes(
 
         mat_indicator = (self.outcome[:, np.newaxis] <= locations).astype(int)
 
-        lower_band, upper_band = compute_confidence_intervals(
+        lower_bound, upper_bound = compute_confidence_intervals(
             vec_y=self.outcome,
             vec_d=self.treatment_arm,
             vec_loc=locations,
@@ -188,8 +189,8 @@ def _compute_dtes(
 
         return (
             dte,
-            lower_band,
-            upper_band,
+            lower_bound,
+            upper_bound,
         )
 
     def _compute_ptes(
@@ -248,7 +249,7 @@ def _compute_ptes(
             int
         )
 
-        lower_band, upper_band = compute_confidence_intervals(
+        lower_bound, upper_bound = compute_confidence_intervals(
             vec_y=self.outcome,
             vec_d=self.treatment_arm,
             vec_loc=locations,
@@ -266,8 +267,8 @@ def _compute_ptes(
 
         return (
             pte,
-            lower_band,
-            upper_band,
+            lower_bound,
+            upper_bound,
         )
 
     def _compute_qtes(
@@ -326,28 +327,6 @@ def _compute_cumulative_distribution(
         raise NotImplementedError()
 
 
-def find_le(array: np.ndarray, threshold):
-    """Find the rightmost value less than or equal to threshold in a sorted array
-
-    Args:
-        array (np.ndarray): The sorted array to search in.
-        threshold (float): The threshold value.
-
-    Returns:
-        int: The index where the value first exceeds the threshold.
-    """
-    low, high = 0, array.shape[0] - 1
-    result = -1
-    while low <= high:
-        mid = (low + high) // 2
-        if array[mid] <= threshold:
-            result = mid
-            low = mid + 1
-        else:
-            high = mid - 1
-    return result
-
-
 class SimpleDistributionEstimator(DistributionFunctionMixin):
     """A class for computing the empirical distribution function and the distributional parameters
     based on the distribution function.
@@ -564,108 +543,3 @@ def _compute_cumulative_distribution(
             )
         return cumulative_distribution, superset_prediction
 
-
-def compute_confidence_intervals(
-    vec_y: np.ndarray,
-    vec_d: np.ndarray,
-    vec_loc: np.ndarray,
-    mat_y_u: np.ndarray,
-    vec_prediction_target: np.ndarray,
-    vec_prediction_control: np.ndarray,
-    mat_entire_predictions_target: np.ndarray,
-    mat_entire_predictions_control: np.ndarray,
-    ind_target: int,
-    ind_control: int,
-    alpha: 0.05,
-    variance_type="moment",
-    n_bootstrap=500,
-):
-    """Computes the confidence intervals of distribution parameters.
-
-    Args:
-        vec_y (np.ndarray): Outcome variable vector.
-        vec_d (np.ndarray): Treatment indicator vector.
-        vec_loc (np.ndarray): Locations where the distribution parameters are estimated.
-        mat_y_u (np.ndarray): Indicator function for 1{Y⩽y}. Shape is n_obs * n_loc.
-        vec_prediction_target (np.ndarray): Estimated values from the conditional model for the treatment group.
-        vec_prediction_control (np.ndarray): Estimated values from the conditional model for the control group.
-        mat_entire_predictions_target (np.ndarray): Prediction of the conditional distribution estimator for target group.
-        mat_entire_predictions_control (np.ndarray): Prediction of the conditional distribution estimator for control group.
-        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 moment, simple, and uniform.
-        n_bootstrap (int, optional): Number of bootstrap samples. Defaults to 500.
-
-    Returns:
-        Tuple[np.ndarray, np.ndarray]: A tuple containing:
-            - np.ndarray: lower band.
-            - np.ndarray: upper band.
-    """
-    num_obs = vec_y.shape[0]
-    n_loc = vec_loc.shape[0]
-    mat_d = np.tile(vec_d, (n_loc, 1)).T
-    vec_dte = vec_prediction_target - vec_prediction_control
-    mat_dte = np.tile(vec_dte, (num_obs, 1))
-
-    num_target = (vec_d == ind_target).sum()
-    num_control = (vec_d == ind_control).sum()
-    influence_function = (
-        num_obs
-        / num_target
-        * (mat_d == ind_target)
-        * (mat_y_u - mat_entire_predictions_target)
-        + mat_entire_predictions_target
-        - num_obs
-        / num_control
-        * (mat_d == ind_control)
-        * (mat_y_u - mat_entire_predictions_control)
-        - mat_entire_predictions_control
-        - mat_dte
-    )
-
-    omega = (influence_function**2).mean(axis=0)
-
-    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
-        )
-        return vec_dte_lower_moment, vec_dte_upper_moment
-    elif variance_type == "uniform":
-        tstats = np.zeros((n_bootstrap, len(vec_loc)))
-        boot_draw = np.zeros((n_bootstrap, len(vec_loc)))
-
-        for b in range(n_bootstrap):
-            eta1 = np.random.normal(0, 1, num_obs)
-            eta2 = np.random.normal(0, 1, num_obs)
-            xi = eta1 / np.sqrt(2) + (eta2**2 - 1) / 2
-
-            boot_draw[b, :] = (
-                1 / num_obs * np.sum(xi[:, np.newaxis] * influence_function, axis=0)
-            )
-
-        tstats = np.abs(boot_draw)[:, :-1] / np.sqrt(omega[:-1] / num_obs)
-        max_tstats = np.max(tstats, axis=1)
-        quantile_max_tstats = np.quantile(max_tstats, 1 - alpha)
-
-        vec_dte_lower_boot = vec_dte - quantile_max_tstats * np.sqrt(omega / num_obs)
-        vec_dte_upper_boot = vec_dte + quantile_max_tstats * np.sqrt(omega / num_obs)
-        return vec_dte_lower_boot, vec_dte_upper_boot
-    elif variance_type == "simple":
-        w_target = num_obs / num_target
-        w_control = num_obs / num_control
-        vec_dte_var = w_target * (
-            vec_prediction_target * (1 - vec_prediction_target)
-        ) + w_control * vec_prediction_control * (1 - vec_prediction_control)
-
-        vec_dte_lower_simple = vec_dte + norm.ppf(alpha / 2) / np.sqrt(
-            num_obs
-        ) * np.sqrt(vec_dte_var)
-        vec_dte_upper_simple = vec_dte + norm.ppf(1 - alpha / 2) / np.sqrt(
-            num_obs
-        ) * np.sqrt(vec_dte_var)
-
-        return vec_dte_lower_simple, vec_dte_upper_simple
-    else:
-        raise RuntimeError(f"Invalid variance type was speficied: {variance_type}")

dte_adj/plot/__init__.py

@@ -5,10 +5,11 @@
 
 
 def plot(
-    x_values: np.ndarray,
-    y_values: np.ndarray,
-    upper_bands: np.ndarray,
-    lower_bands: np.ndarray,
+    X: np.ndarray,
+    means: np.ndarray,
+    upper_bounds: np.ndarray,
+    lower_bounds: np.ndarray,
+    chart_type="line",
     ax: Optional[axis.Axis] = None,
     title: Optional[str] = None,
     xlabel: Optional[str] = None,
@@ -17,10 +18,11 @@ def plot(
     """Visualize distributional parameters and their confidence intervals.
 
     Args:
-        x_values (np.Array): values to be used for x axis.
-        y_values (np.Array): Expected distributional parameters.
-        upper_bands (np.Array): Upper band for the distributional parameters.
-        lower_bands (np.Array): Lower band for the distributional parameters.
+        X (np.Array): values to be used for x axis.
+        means (np.Array): Expected distributional parameters.
+        upper_bounds (np.Array): Upper bound for the distributional parameters.
+        lower_bounds (np.Array): Lower bound for the distributional parameters.
+        chart_type (str): Chart type of the plotting. Available values are line or bar.
         ax (matplotlib.axes.Axes, optional): Target axes instance. If None, a new figure and axes will be created.
         title (str, optional): Axes title.
         xlabel (str, optional): X-axis title label.
@@ -32,15 +34,28 @@ def plot(
     if ax is None:
         fig, ax = plt.subplots()
 
-    ax.plot(x_values, y_values, label="Values", color="blue")
-    ax.fill_between(
-        x_values,
-        lower_bands,
-        upper_bands,
-        color="gray",
-        alpha=0.3,
-        label="Confidence Interval",
-    )
+    if chart_type == "line":
+        ax.plot(X, means, label="Values", color="blue")
+        ax.fill_between(
+            X,
+            lower_bounds,
+            upper_bounds,
+            color="gray",
+            alpha=0.3,
+            label="Confidence Interval",
+        )
+    elif chart_type == "bar":
+        ax.bar(
+            X,
+            means,
+            yerr=[
+                np.clip(means - lower_bounds, 0, None),
+                np.clip(upper_bounds - means, 0, None),
+            ],
+            capsize=5,
+        )
+    else:
+        raise ValueError(f"Chart type {chart_type} is not supported")
 
     if title is not None:
         ax.set_title(title)