diff --git a/dte_adj/base.py b/dte_adj/base.py index cee191c..dba2477 100644 --- a/dte_adj/base.py +++ b/dte_adj/base.py @@ -60,36 +60,34 @@ def predict_dte( - Upper bounds (np.ndarray): Upper confidence interval bounds Example: - .. code-block:: python - - import numpy as np - from dte_adj import SimpleDistributionEstimator - - # Generate sample data - X = np.random.randn(1000, 5) - D = np.random.binomial(1, 0.5, 1000) - Y = X[:, 0] + 2 * D + np.random.randn(1000) - - # Fit estimator - estimator = SimpleDistributionEstimator() - estimator.fit(X, D, Y) - - # Compute DTE - locations = np.linspace(Y.min(), Y.max(), 20) - dte, lower, upper = estimator.predict_dte( - target_treatment_arm=1, - control_treatment_arm=0, - locations=locations, - variance_type="moment" - ) + ```python + import numpy as np + from dte_adj import SimpleDistributionEstimator + + # Generate sample data + X = np.random.randn(1000, 5) + D = np.random.binomial(1, 0.5, 1000) + Y = X[:, 0] + 2 * D + np.random.randn(1000) + + # Fit estimator + estimator = SimpleDistributionEstimator() + estimator.fit(X, D, Y) + + # Compute DTE + locations = np.linspace(Y.min(), Y.max(), 20) + dte, lower, upper = estimator.predict_dte( + target_treatment_arm=1, + control_treatment_arm=0, + locations=locations, + variance_type="moment" + ) - print(f"DTE shape: {dte.shape}") # Should match locations.shape - print(f"Average DTE: {dte.mean():.3f}") + print(f"DTE shape: {dte.shape}") # Should match locations.shape + print(f"Average DTE: {dte.mean():.3f}") + ``` """ if locations is None: - locations = _infer_default_locations( - self.outcomes, for_intervals=False - ) + locations = _infer_default_locations(self.outcomes, for_intervals=False) self.last_locations = locations return self._compute_dtes( target_treatment_arm, @@ -143,38 +141,36 @@ def predict_pte( - Upper bounds (np.ndarray): Upper confidence interval bounds Example: - .. code-block:: python - - import numpy as np - from dte_adj import SimpleDistributionEstimator - - # Generate sample data - X = np.random.randn(1000, 5) - D = np.random.binomial(1, 0.5, 1000) - Y = X[:, 0] + 2 * D + np.random.randn(1000) - - # Fit estimator - estimator = SimpleDistributionEstimator() - estimator.fit(X, D, Y) - - # Define interval boundaries - locations = np.array([-2, -1, 0, 1, 2]) # Creates intervals: (-2,-1], (-1,0], (0,1], (1,2] - - # Compute PTE - pte, lower, upper = estimator.predict_pte( - target_treatment_arm=1, - control_treatment_arm=0, - locations=locations, - variance_type="moment" - ) + ```python + import numpy as np + from dte_adj import SimpleDistributionEstimator + + # Generate sample data + X = np.random.randn(1000, 5) + D = np.random.binomial(1, 0.5, 1000) + Y = X[:, 0] + 2 * D + np.random.randn(1000) + + # Fit estimator + estimator = SimpleDistributionEstimator() + estimator.fit(X, D, Y) + + # Define interval boundaries + locations = np.array([-2, -1, 0, 1, 2]) # Creates intervals: (-2,-1], (-1,0], (0,1], (1,2] + + # Compute PTE + pte, lower, upper = estimator.predict_pte( + target_treatment_arm=1, + control_treatment_arm=0, + locations=locations, + variance_type="moment" + ) - print(f"PTE shape: {pte.shape}") # Should be (4,) for 4 intervals - print(f"Interval effects: {pte}") + print(f"PTE shape: {pte.shape}") # Should be (4,) for 4 intervals + print(f"Interval effects: {pte}") + ``` """ if locations is None: - locations = _infer_default_locations( - self.outcomes, for_intervals=True - ) + locations = _infer_default_locations(self.outcomes, for_intervals=True) self.last_locations = locations return self._compute_ptes( target_treatment_arm, @@ -222,32 +218,32 @@ def predict_qte( - Upper bounds (np.ndarray): Upper confidence interval bounds Example: - .. code-block:: python - - import numpy as np - from dte_adj import SimpleStratifiedDistributionEstimator - - # Generate stratified sample data - X = np.random.randn(1000, 5) - strata = np.random.choice([0, 1, 2], size=1000) - D = np.random.binomial(1, 0.5, 1000) - Y = X[:, 0] + 2 * D + 0.5 * strata + np.random.randn(1000) - - # Fit stratified estimator - estimator = SimpleStratifiedDistributionEstimator() - estimator.fit(X, D, Y, strata) - - # Compute QTE at specific quantiles - quantiles = np.array([0.25, 0.5, 0.75]) # 25th, 50th, 75th percentiles - qte, lower, upper = estimator.predict_qte( - target_treatment_arm=1, - control_treatment_arm=0, - quantiles=quantiles, - n_bootstrap=100 - ) + ```python + import numpy as np + from dte_adj import SimpleStratifiedDistributionEstimator + + # Generate stratified sample data + X = np.random.randn(1000, 5) + strata = np.random.choice([0, 1, 2], size=1000) + D = np.random.binomial(1, 0.5, 1000) + Y = X[:, 0] + 2 * D + 0.5 * strata + np.random.randn(1000) + + # Fit stratified estimator + estimator = SimpleStratifiedDistributionEstimator() + estimator.fit(X, D, Y, strata) + + # Compute QTE at specific quantiles + quantiles = np.array([0.25, 0.5, 0.75]) # 25th, 50th, 75th percentiles + qte, lower, upper = estimator.predict_qte( + target_treatment_arm=1, + control_treatment_arm=0, + quantiles=quantiles, + n_bootstrap=100 + ) - print(f"QTE at quantiles {quantiles}: {qte}") - print(f"Median effect (50th percentile): {qte[1]:.3f}") + print(f"QTE at quantiles {quantiles}: {qte}") + print(f"Median effect (50th percentile): {qte[1]:.3f}") + ``` """ if quantiles is None: quantiles = np.arange(1, 10) / 10 diff --git a/dte_adj/local.py b/dte_adj/local.py index 5d1fdd2..6e65917 100644 --- a/dte_adj/local.py +++ b/dte_adj/local.py @@ -94,39 +94,37 @@ def predict_ldte( - Upper bounds (np.ndarray): Upper confidence interval bounds Example: - .. code-block:: python - - import numpy as np - from sklearn.linear_model import LogisticRegression - from dte_adj import AdjustedLocalDistributionEstimator - - # Generate sample data with strata - np.random.seed(42) - X = np.random.randn(1000, 5) - strata = np.random.choice([0, 1], size=1000) # Binary strata - D = np.random.binomial(1, 0.3 + 0.4 * strata, 1000) # Treatment depends on strata - Y = X[:, 0] + 2 * D + strata + np.random.randn(1000) - - # Fit local estimator - base_model = LogisticRegression() - estimator = AdjustedLocalDistributionEstimator(base_model) - estimator.fit(X, D, D, Y, strata) # treatment_arms = treatment_indicator for binary case - - # Compute LDTE - locations = np.linspace(Y.min(), Y.max(), 20) - ldte, lower, upper = estimator.predict_ldte( - target_treatment_arm=1, - control_treatment_arm=0, - locations=locations - ) - - print(f"LDTE shape: {ldte.shape}") # Should match locations.shape - print(f"Average LDTE: {ldte.mean():.3f}") + ```python + import numpy as np + from sklearn.linear_model import LogisticRegression + from dte_adj import AdjustedLocalDistributionEstimator + + # Generate sample data with strata + np.random.seed(42) + X = np.random.randn(1000, 5) + strata = np.random.choice([0, 1], size=1000) # Binary strata + D = np.random.binomial(1, 0.3 + 0.4 * strata, 1000) # Treatment depends on strata + Y = X[:, 0] + 2 * D + strata + np.random.randn(1000) + + # Fit local estimator + base_model = LogisticRegression() + estimator = AdjustedLocalDistributionEstimator(base_model) + estimator.fit(X, D, D, Y, strata) # treatment_arms = treatment_indicator for binary case + + # Compute LDTE + locations = np.linspace(Y.min(), Y.max(), 20) + ldte, lower, upper = estimator.predict_ldte( + target_treatment_arm=1, + control_treatment_arm=0, + locations=locations + ) + + print(f"LDTE shape: {ldte.shape}") # Should match locations.shape + print(f"Average LDTE: {ldte.mean():.3f}") + ``` """ if locations is None: - locations = _infer_default_locations( - self.outcomes, for_intervals=False - ) + locations = _infer_default_locations(self.outcomes, for_intervals=False) self.last_locations = locations return compute_ldte( self, @@ -173,40 +171,38 @@ def predict_lpte( - Upper bounds (np.ndarray): Upper confidence interval bounds Example: - .. code-block:: python - - import numpy as np - from dte_adj import SimpleLocalDistributionEstimator - - # Generate sample data with strata - np.random.seed(42) - X = np.random.randn(1000, 5) - strata = np.random.choice([0, 1], size=1000) # Binary strata - Z = np.random.binomial(1, 0.5, 1000) # Treatment assignment - D = np.random.binomial(1, 0.3 + 0.4 * Z, 1000) # Treatment receipt - Y = X[:, 0] + 2 * D + strata + np.random.randn(1000) - - # Fit local estimator - estimator = SimpleLocalDistributionEstimator() - estimator.fit(X, Z, D, Y, strata) - - # Define interval boundaries - locations = np.array([-2, -1, 0, 1, 2]) # Creates intervals: (-2,-1], (-1,0], (0,1], (1,2] - - # Compute LPTE - lpte, lower, upper = estimator.predict_lpte( - target_treatment_arm=1, - control_treatment_arm=0, - locations=locations - ) - - print(f"LPTE shape: {lpte.shape}") # Should be (4,) for 4 intervals - print(f"Interval effects: {lpte}") + ```python + import numpy as np + from dte_adj import SimpleLocalDistributionEstimator + + # Generate sample data with strata + np.random.seed(42) + X = np.random.randn(1000, 5) + strata = np.random.choice([0, 1], size=1000) # Binary strata + Z = np.random.binomial(1, 0.5, 1000) # Treatment assignment + D = np.random.binomial(1, 0.3 + 0.4 * Z, 1000) # Treatment receipt + Y = X[:, 0] + 2 * D + strata + np.random.randn(1000) + + # Fit local estimator + estimator = SimpleLocalDistributionEstimator() + estimator.fit(X, Z, D, Y, strata) + + # Define interval boundaries + locations = np.array([-2, -1, 0, 1, 2]) # Creates intervals: (-2,-1], (-1,0], (0,1], (1,2] + + # Compute LPTE + lpte, lower, upper = estimator.predict_lpte( + target_treatment_arm=1, + control_treatment_arm=0, + locations=locations + ) + + print(f"LPTE shape: {lpte.shape}") # Should be (4,) for 4 intervals + print(f"Interval effects: {lpte}") + ``` """ if locations is None: - locations = _infer_default_locations( - self.outcomes, for_intervals=True - ) + locations = _infer_default_locations(self.outcomes, for_intervals=True) self.last_locations = locations return compute_lpte( self, @@ -287,41 +283,39 @@ def predict_ldte( - Upper bounds (np.ndarray): Upper confidence interval bounds Example: - .. code-block:: python - - import numpy as np - from sklearn.ensemble import RandomForestClassifier - from dte_adj import AdjustedLocalDistributionEstimator - - # Generate confounded data with strata - np.random.seed(42) - X = np.random.randn(1000, 5) - strata = np.random.choice([0, 1], size=1000) - # Treatment assignment depends on covariates - Z_prob = 1 / (1 + np.exp(-(X[:, 0] + X[:, 1] + strata))) - Z = np.random.binomial(1, Z_prob, 1000) - D = np.random.binomial(1, 0.3 + 0.4 * Z, 1000) - Y = X.sum(axis=1) + 2 * D + strata + np.random.randn(1000) - - # Fit adjusted local estimator - base_model = RandomForestClassifier(n_estimators=100) - estimator = AdjustedLocalDistributionEstimator(base_model, folds=3) - estimator.fit(X, Z, D, Y, strata) - - # Compute LDTE with ML adjustment - locations = np.linspace(Y.min(), Y.max(), 20) - ldte, lower, upper = estimator.predict_ldte( - target_treatment_arm=1, - control_treatment_arm=0, - locations=locations - ) - - print(f"Adjusted LDTE: {ldte.mean():.3f}") + ```python + import numpy as np + from sklearn.ensemble import RandomForestClassifier + from dte_adj import AdjustedLocalDistributionEstimator + + # Generate confounded data with strata + np.random.seed(42) + X = np.random.randn(1000, 5) + strata = np.random.choice([0, 1], size=1000) + # Treatment assignment depends on covariates + Z_prob = 1 / (1 + np.exp(-(X[:, 0] + X[:, 1] + strata))) + Z = np.random.binomial(1, Z_prob, 1000) + D = np.random.binomial(1, 0.3 + 0.4 * Z, 1000) + Y = X.sum(axis=1) + 2 * D + strata + np.random.randn(1000) + + # Fit adjusted local estimator + base_model = RandomForestClassifier(n_estimators=100) + estimator = AdjustedLocalDistributionEstimator(base_model, folds=3) + estimator.fit(X, Z, D, Y, strata) + + # Compute LDTE with ML adjustment + locations = np.linspace(Y.min(), Y.max(), 20) + ldte, lower, upper = estimator.predict_ldte( + target_treatment_arm=1, + control_treatment_arm=0, + locations=locations + ) + + print(f"Adjusted LDTE: {ldte.mean():.3f}") + ``` """ if locations is None: - locations = _infer_default_locations( - self.outcomes, for_intervals=False - ) + locations = _infer_default_locations(self.outcomes, for_intervals=False) self.last_locations = locations return compute_ldte( self, @@ -367,43 +361,41 @@ def predict_lpte( - Upper bounds (np.ndarray): Upper confidence interval bounds Example: - .. code-block:: python - - import numpy as np - from sklearn.linear_model import LogisticRegression - from dte_adj import AdjustedLocalDistributionEstimator - - # Generate confounded data with strata - np.random.seed(42) - X = np.random.randn(1000, 5) - strata = np.random.choice([0, 1], size=1000) - # Treatment assignment depends on covariates - Z_prob = 1 / (1 + np.exp(-(X[:, 0] + strata))) - Z = np.random.binomial(1, Z_prob, 1000) - D = np.random.binomial(1, 0.3 + 0.4 * Z, 1000) - Y = X.sum(axis=1) + 2 * D + strata + np.random.randn(1000) - - # Fit adjusted local estimator - base_model = LogisticRegression() - estimator = AdjustedLocalDistributionEstimator(base_model, folds=3) - estimator.fit(X, Z, D, Y, strata) - - # Define interval boundaries - locations = np.array([-2, -1, 0, 1, 2]) - - # Compute LPTE with ML adjustment - lpte, lower, upper = estimator.predict_lpte( - target_treatment_arm=1, - control_treatment_arm=0, - locations=locations - ) - - print(f"Adjusted LPTE: {lpte}") + ```python + import numpy as np + from sklearn.linear_model import LogisticRegression + from dte_adj import AdjustedLocalDistributionEstimator + + # Generate confounded data with strata + np.random.seed(42) + X = np.random.randn(1000, 5) + strata = np.random.choice([0, 1], size=1000) + # Treatment assignment depends on covariates + Z_prob = 1 / (1 + np.exp(-(X[:, 0] + strata))) + Z = np.random.binomial(1, Z_prob, 1000) + D = np.random.binomial(1, 0.3 + 0.4 * Z, 1000) + Y = X.sum(axis=1) + 2 * D + strata + np.random.randn(1000) + + # Fit adjusted local estimator + base_model = LogisticRegression() + estimator = AdjustedLocalDistributionEstimator(base_model, folds=3) + estimator.fit(X, Z, D, Y, strata) + + # Define interval boundaries + locations = np.array([-2, -1, 0, 1, 2]) + + # Compute LPTE with ML adjustment + lpte, lower, upper = estimator.predict_lpte( + target_treatment_arm=1, + control_treatment_arm=0, + locations=locations + ) + + print(f"Adjusted LPTE: {lpte}") + ``` """ if locations is None: - locations = _infer_default_locations( - self.outcomes, for_intervals=True - ) + locations = _infer_default_locations(self.outcomes, for_intervals=True) self.last_locations = locations return compute_lpte( self, diff --git a/dte_adj/simple.py b/dte_adj/simple.py index 37fc93c..1e1a6ca 100644 --- a/dte_adj/simple.py +++ b/dte_adj/simple.py @@ -19,24 +19,24 @@ class SimpleDistributionEstimator(SimpleStratifiedDistributionEstimator): covariate adjustment is not needed. Example: - .. code-block:: python - - import numpy as np - from dte_adj import SimpleDistributionEstimator - - # Generate sample data - X = np.random.randn(1000, 5) - D = np.random.binomial(1, 0.5, 1000) # Random treatment - Y = X[:, 0] + 2 * D + np.random.randn(1000) - - # Fit simple estimator - estimator = SimpleDistributionEstimator() - estimator.fit(X, D, Y) - - # Compute treatment effects - locations = np.linspace(Y.min(), Y.max(), 20) - dte, lower, upper = estimator.predict_dte(1, 0, locations) - pte, pte_lower, pte_upper = estimator.predict_pte(1, 0, locations) + ```python + import numpy as np + from dte_adj import SimpleDistributionEstimator + + # Generate sample data + X = np.random.randn(1000, 5) + D = np.random.binomial(1, 0.5, 1000) # Random treatment + Y = X[:, 0] + 2 * D + np.random.randn(1000) + + # Fit simple estimator + estimator = SimpleDistributionEstimator() + estimator.fit(X, D, Y) + + # Compute treatment effects + locations = np.linspace(Y.min(), Y.max(), 20) + dte, lower, upper = estimator.predict_dte(1, 0, locations) + pte, pte_lower, pte_upper = estimator.predict_pte(1, 0, locations) + ``` """ def __init__(self): @@ -89,26 +89,26 @@ class AdjustedDistributionEstimator(AdjustedStratifiedDistributionEstimator): assignment depends on observed covariates. Example: - .. code-block:: python - - import numpy as np - from sklearn.ensemble import RandomForestClassifier - from dte_adj import AdjustedDistributionEstimator - - # Generate confounded data - X = np.random.randn(1000, 5) - treatment_prob = 1 / (1 + np.exp(-(X[:, 0] + X[:, 1]))) - D = np.random.binomial(1, treatment_prob, 1000) - Y = X.sum(axis=1) + 2 * D + np.random.randn(1000) - - # Fit adjusted estimator - base_model = RandomForestClassifier(n_estimators=100) - estimator = AdjustedDistributionEstimator(base_model, folds=3) - estimator.fit(X, D, Y) - - # Compute adjusted treatment effects - locations = np.linspace(Y.min(), Y.max(), 20) - dte, lower, upper = estimator.predict_dte(1, 0, locations, variance_type="moment") + ```python + import numpy as np + from sklearn.ensemble import RandomForestClassifier + from dte_adj import AdjustedDistributionEstimator + + # Generate confounded data + X = np.random.randn(1000, 5) + treatment_prob = 1 / (1 + np.exp(-(X[:, 0] + X[:, 1]))) + D = np.random.binomial(1, treatment_prob, 1000) + Y = X.sum(axis=1) + 2 * D + np.random.randn(1000) + + # Fit adjusted estimator + base_model = RandomForestClassifier(n_estimators=100) + estimator = AdjustedDistributionEstimator(base_model, folds=3) + estimator.fit(X, D, Y) + + # Compute adjusted treatment effects + locations = np.linspace(Y.min(), Y.max(), 20) + dte, lower, upper = estimator.predict_dte(1, 0, locations, variance_type="moment") + ``` """ def fit(