From 79f884261cad8e8975e4a5c050da28278a8e641e Mon Sep 17 00:00:00 2001 From: TomeHirata Date: Sun, 28 Sep 2025 17:50:27 +0900 Subject: [PATCH 1/2] support target treatment indicator other than 1 --- dte_adj/stratified.py | 17 ++++++++--------- dte_adj/util.py | 18 ++++++++++++------ example/example.ipynb | 4 +--- 3 files changed, 21 insertions(+), 18 deletions(-) diff --git a/dte_adj/stratified.py b/dte_adj/stratified.py index e7e3d30..f5bb3bb 100644 --- a/dte_adj/stratified.py +++ b/dte_adj/stratified.py @@ -73,10 +73,9 @@ def _compute_cumulative_distribution( for s in s_list: s_mask = strata == s w_s[s] = (s_mask & treatment_mask).sum() / s_mask.sum() - for i, outcome in enumerate(locations): - for j in range(n_records): - s = strata[j] - prediction[j, i] = (outcomes[j] <= outcome) / w_s[s] * treatment_mask[j] + for j in range(n_records): + s = strata[j] + prediction[j] = (outcomes[j] <= locations) / w_s[s] * treatment_mask[j] unconditional_pred = {s: prediction[s == strata].mean(axis=0) for s in s_list} conditional_prediction = np.array([unconditional_pred[s] for s in strata]) @@ -263,9 +262,9 @@ def _compute_cumulative_distribution( covariates_train = covariates[fold_mask] binomial_train = binomial[fold_mask] # Pool the records across strata and train the model - if len(np.unique(binomial_train)) > 1: - self.model = deepcopy(self.base_model) - self.model.fit(covariates_train, binomial_train) + # if len(np.unique(binomial_train)) > 1: + # self.model = deepcopy(self.base_model) + # self.model.fit(covariates_train, binomial_train) for s in s_list: s_mask = strata == s weight = (s_mask & treatment_mask).sum() / s_mask.sum() @@ -275,8 +274,8 @@ def _compute_cumulative_distribution( binomial_train = binomial[subset_train_mask] # TODO: revisit the logic here if len(np.unique(binomial_train)) > 1: - # self.model = deepcopy(self.base_model) - # self.model.fit(covariates_train, binomial_train) + self.model = deepcopy(self.base_model) + self.model.fit(covariates_train, binomial_train) pass else: pred = binomial_train[0] diff --git a/dte_adj/util.py b/dte_adj/util.py index c0265b3..6b52f13 100644 --- a/dte_adj/util.py +++ b/dte_adj/util.py @@ -1,6 +1,12 @@ import numpy as np from scipy.stats import norm -from typing import Tuple +from typing import Tuple, TYPE_CHECKING + +if TYPE_CHECKING: + from dte_adj.local import ( + SimpleStratifiedDistributionEstimator, + AdjustedLocalDistributionEstimator, + ) def compute_confidence_intervals( @@ -110,7 +116,7 @@ def compute_confidence_intervals( def _compute_local_treatment_effects_core( - estimator, + estimator: "SimpleStratifiedDistributionEstimator | AdjustedLocalDistributionEstimator", target_treatment_arm: int, control_treatment_arm: int, locations: np.ndarray, @@ -149,10 +155,10 @@ def _compute_local_treatment_effects_core( # Compute treatment propensity (probability of treatment) d_t_prediction, d_t_psi, d_t_eta = estimator._compute_cumulative_distribution( - target_treatment_arm, np.zeros(1), X, Z, 1 - D + target_treatment_arm, np.zeros(1), X, Z, 1 - (target_treatment_arm == D) ) d_c_prediction, d_c_psi, d_c_eta = estimator._compute_cumulative_distribution( - control_treatment_arm, np.zeros(1), X, Z, 1 - D + control_treatment_arm, np.zeros(1), X, Z, 1 - (target_treatment_arm == D) ) # Compute outcome distributions (different for LDTE vs LPTE) @@ -257,7 +263,7 @@ def xi(s): def compute_ldte( - estimator, + estimator: "SimpleStratifiedDistributionEstimator | AdjustedLocalDistributionEstimator", target_treatment_arm: int, control_treatment_arm: int, locations: np.ndarray, @@ -290,7 +296,7 @@ def compute_ldte( def compute_lpte( - estimator, + estimator: "SimpleStratifiedDistributionEstimator | AdjustedLocalDistributionEstimator", target_treatment_arm: int, control_treatment_arm: int, locations: np.ndarray, diff --git a/example/example.ipynb b/example/example.ipynb index e81d582..a47b590 100644 --- a/example/example.ipynb +++ b/example/example.ipynb @@ -278,9 +278,7 @@ ], "source": [ "pte, lower_bound, upper_bound = estimator.predict_pte(\n", - " target_treatment_arm=1,\n", - " control_treatment_arm=0,\n", - " locations=locations\n", + " target_treatment_arm=1, control_treatment_arm=0, locations=locations\n", ")\n", "plot(\n", " locations[:-1],\n", From 76d06c02178da940de72a9178551d17d24ff25fc Mon Sep 17 00:00:00 2001 From: TomeHirata Date: Sun, 28 Sep 2025 17:51:42 +0900 Subject: [PATCH 2/2] revert --- dte_adj/stratified.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/dte_adj/stratified.py b/dte_adj/stratified.py index f5bb3bb..880f371 100644 --- a/dte_adj/stratified.py +++ b/dte_adj/stratified.py @@ -262,9 +262,9 @@ def _compute_cumulative_distribution( covariates_train = covariates[fold_mask] binomial_train = binomial[fold_mask] # Pool the records across strata and train the model - # if len(np.unique(binomial_train)) > 1: - # self.model = deepcopy(self.base_model) - # self.model.fit(covariates_train, binomial_train) + if len(np.unique(binomial_train)) > 1: + self.model = deepcopy(self.base_model) + self.model.fit(covariates_train, binomial_train) for s in s_list: s_mask = strata == s weight = (s_mask & treatment_mask).sum() / s_mask.sum() @@ -274,8 +274,8 @@ def _compute_cumulative_distribution( binomial_train = binomial[subset_train_mask] # TODO: revisit the logic here if len(np.unique(binomial_train)) > 1: - self.model = deepcopy(self.base_model) - self.model.fit(covariates_train, binomial_train) + # self.model = deepcopy(self.base_model) + # self.model.fit(covariates_train, binomial_train) pass else: pred = binomial_train[0]