Modules

class dte_adj.AdjustedDistributionEstimator(base_model: Any, folds=3, is_multi_task=False)

Bases: AdjustedStratifiedDistributionEstimator

A class for computing distribution treatment effects using machine learning adjustment.

This estimator uses cross-fitting with ML models to adjust for confounding when computing Distribution Treatment Effects (DTE), Probability Treatment Effects (PTE), and Quantile Treatment Effects (QTE). It provides more precise estimates when treatment assignment depends on observed covariates.

Example

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")

fit(covariates: ndarray, treatment_arms: ndarray, outcomes: ndarray) → AdjustedDistributionEstimator

Set parameters.

Parameters:

• covariates (np.ndarray) – Pre-treatment covariates.

• treatment_arms (np.ndarray) – The index of the treatment arm.

• outcomes (np.ndarray) – Scalar-valued observed outcome.

Returns:

The fitted estimator.

Return type:

AdjustedDistributionEstimator

predict(treatment_arm: int, locations: ndarray) → ndarray

Compute cumulative distribution values.

Parameters:

• treatment_arm (int) – The index of the treatment arm.

• outcomes (np.ndarray) – Scalar values to be used for computing the cumulative distribution.

Returns:

Estimated cumulative distribution values for the input.

Return type:

np.ndarray

predict_dte(target_treatment_arm: int, control_treatment_arm: int, locations: ndarray, alpha: float = 0.05, variance_type='moment', n_bootstrap=500) → Tuple[ndarray, ndarray, ndarray]

Compute Distribution Treatment Effects (DTE) based on the estimator for the distribution function.

The DTE measures the difference in cumulative distribution functions between treatment groups at specified locations. It quantifies how treatment affects the probability of observing outcomes below each threshold.

Parameters:

• 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 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”. Defaults to “moment”.

• n_bootstrap (int, optional) – Number of bootstrap samples. Defaults to 500.

Returns:

A tuple containing:

• Expected DTEs (np.ndarray): Treatment effect estimates at each location

• Lower bounds (np.ndarray): Lower confidence interval bounds

• Upper bounds (np.ndarray): Upper confidence interval bounds

Return type:

Tuple[np.ndarray, np.ndarray, np.ndarray]

Example

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}")

predict_pte(target_treatment_arm: int, control_treatment_arm: int, locations: ndarray, alpha: float = 0.05, variance_type='moment', n_bootstrap=500) → Tuple[ndarray, ndarray, ndarray]

Compute Probability Treatment Effects (PTE) based on the estimator for the distribution function.

The PTE measures the difference in probability mass between treatment groups for intervals defined by consecutive location pairs. It quantifies how treatment affects the probability of observing outcomes within specific ranges.

Parameters:

• 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 defining interval boundaries for probability computation. For each interval (locations[i], locations[i+1]], the PTE is computed.

• 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”. Defaults to “moment”.

• n_bootstrap (int, optional) – Number of bootstrap samples. Defaults to 500.

Returns:

A tuple containing:

• Expected PTEs (np.ndarray): Treatment effect estimates for each interval, shape (len(locations)-1,)

• Lower bounds (np.ndarray): Lower confidence interval bounds

• Upper bounds (np.ndarray): Upper confidence interval bounds

Return type:

Tuple[np.ndarray, np.ndarray, np.ndarray]

Example

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}")

predict_qte(target_treatment_arm: int, control_treatment_arm: int, quantiles: ndarray | None = None, alpha: float = 0.05, n_bootstrap=500) → Tuple[ndarray, ndarray, ndarray]

Compute Quantile Treatment Effects (QTE) based on the estimator for the distribution function.

The QTE measures the difference in quantiles between treatment groups, providing insights into how treatment affects different parts of the outcome distribution. For stratified estimators, the computation properly accounts for strata.

Parameters:

• 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, 0.2, …, 0.9].

• 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:

A tuple containing:

• Expected QTEs (np.ndarray): Treatment effect estimates at each quantile

• Lower bounds (np.ndarray): Lower confidence interval bounds

• Upper bounds (np.ndarray): Upper confidence interval bounds

Return type:

Tuple[np.ndarray, np.ndarray, np.ndarray]

Example

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}")

class dte_adj.AdjustedLocalDistributionEstimator(base_model: Any, folds=3, is_multi_task=False)

Bases: AdjustedStratifiedDistributionEstimator

A class for computing Local Distribution Treatment Effects (LDTE) and Local Probability Treatment Effects (LPTE) using ML-adjusted estimation.

This estimator combines local treatment effect estimation with machine learning adjustment, providing treatment effects that are both locally robust to treatment assignment heterogeneity and adjusted for confounding through observed covariates. It uses cross-fitting for more precise estimates in complex treatment assignment scenarios.

fit(covariates: ndarray, treatment_arms: ndarray, treatment_indicator: ndarray, outcomes: ndarray, strata: ndarray) → AdjustedLocalDistributionEstimator

Train the AdjustedLocalDistributionEstimator.

Parameters:

• covariates (np.ndarray) – Pre-treatment covariates.

• treatment_arms (np.ndarray) – Treatment assignment variable (Z).

• treatment_indicator (np.ndarray) – Treatment indicator variable (D).

• outcomes (np.ndarray) – Scalar-valued observed outcome.

• strata (np.ndarray) – Stratum indicators.

Returns:

The fitted estimator.

Return type:

AdjustedLocalDistributionEstimator

predict(treatment_arm: int, locations: ndarray) → ndarray

Compute cumulative distribution values.

Parameters:

• treatment_arm (int) – The index of the treatment arm.

• outcomes (np.ndarray) – Scalar values to be used for computing the cumulative distribution.

Returns:

Estimated cumulative distribution values for the input.

Return type:

np.ndarray

predict_dte(target_treatment_arm: int, control_treatment_arm: int, locations: ndarray, alpha: float = 0.05, variance_type='moment', n_bootstrap=500) → Tuple[ndarray, ndarray, ndarray]

Compute Distribution Treatment Effects (DTE) based on the estimator for the distribution function.

The DTE measures the difference in cumulative distribution functions between treatment groups at specified locations. It quantifies how treatment affects the probability of observing outcomes below each threshold.

Parameters:

• 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 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”. Defaults to “moment”.

• n_bootstrap (int, optional) – Number of bootstrap samples. Defaults to 500.

Returns:

A tuple containing:

• Expected DTEs (np.ndarray): Treatment effect estimates at each location

• Lower bounds (np.ndarray): Lower confidence interval bounds

• Upper bounds (np.ndarray): Upper confidence interval bounds

Return type:

Tuple[np.ndarray, np.ndarray, np.ndarray]

Example

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}")

predict_ldte(target_treatment_arm: int, control_treatment_arm: int, locations: ndarray, alpha: float = 0.05) → Tuple[ndarray, ndarray, ndarray]

Compute Local Distribution Treatment Effects (LDTE) with ML-adjusted estimation. LDTE measures the difference in cumulative distribution functions between treatment groups weighted by treatment propensity within each stratum, using ML models for adjustment. Currently, this API only supports analytical confidence interval.

Parameters:

• 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 bound. Defaults to 0.05.

Returns:

A tuple containing:

• Expected LDTEs (np.ndarray): ML-adjusted local treatment effect estimates

• Lower bounds (np.ndarray): Lower confidence interval bounds

• Upper bounds (np.ndarray): Upper confidence interval bounds

Return type:

Tuple[np.ndarray, np.ndarray, np.ndarray]

Example

import numpy as np
from sklearn.ensemble import RandomForestClassifier
from dte_adj import AdjustedLocalDistributionEstimator

# Generate sample data with complex treatment assignment
np.random.seed(42)
X = np.random.randn(1000, 5)
strata = np.random.choice([0, 1], size=1000)
# Treatment depends on covariates and strata
treatment_prob = 0.2 + 0.3 * (X[:, 0] > 0) + 0.2 * strata
D = np.random.binomial(1, treatment_prob, 1000)
Y = X.sum(axis=1) + 2 * D + strata + np.random.randn(1000)

# Fit adjusted local estimator
base_model = RandomForestClassifier(n_estimators=50, random_state=42)
estimator = AdjustedLocalDistributionEstimator(base_model, folds=3)
estimator.fit(X, D, D, Y, strata)

# Compute LDTE
locations = np.linspace(Y.min(), Y.max(), 15)
ldte, lower, upper = estimator.predict_ldte(
    target_treatment_arm=1,
    control_treatment_arm=0,
    locations=locations
)

print(f"ML-adjusted LDTE shape: {ldte.shape}")
print(f"Average LDTE: {ldte.mean():.3f}")

predict_lpte(target_treatment_arm: int, control_treatment_arm: int, locations: ndarray, alpha: float = 0.05) → Tuple[ndarray, ndarray, ndarray]

Compute Local Probability Treatment Effects (LPTE) with ML-adjusted estimation. LPTE measures the difference in probability mass between treatment groups for intervals weighted by treatment propensity within each stratum, using ML models for adjustment. Currently, this API only supports analytical confidence interval.

Parameters:

• 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 defining interval boundaries for probability computation. For each interval (locations[i], locations[i+1]], the LPTE is computed.

• alpha (float, optional) – Significance level of the confidence bound. Defaults to 0.05.

Returns:

A tuple containing:

• Expected LPTEs (np.ndarray): ML-adjusted local treatment effect estimates, shape (len(locations)-1,)

• Lower bounds (np.ndarray): Lower confidence interval bounds

• Upper bounds (np.ndarray): Upper confidence interval bounds

Return type:

Tuple[np.ndarray, np.ndarray, np.ndarray]

Example

import numpy as np
from sklearn.ensemble import GradientBoostingClassifier
from dte_adj import AdjustedLocalDistributionEstimator

# Generate sample data with confounding
np.random.seed(42)
X = np.random.randn(1000, 5)
strata = np.random.choice([0, 1, 2], size=1000)
# Complex treatment assignment mechanism
logit_score = X[:, 0] + 0.5 * X[:, 1] + strata
treatment_prob = 1 / (1 + np.exp(-logit_score))
D = np.random.binomial(1, treatment_prob, 1000)
Y = X.sum(axis=1) + 1.5 * D + 0.3 * strata + np.random.randn(1000)

# Fit adjusted estimator with gradient boosting
base_model = GradientBoostingClassifier(n_estimators=100, random_state=42)
estimator = AdjustedLocalDistributionEstimator(base_model, folds=5)
estimator.fit(X, D, D, Y, strata)

# Define intervals and compute LPTE
locations = np.array([-3, -1, 0, 1, 3])  # 4 intervals
lpte, lower, upper = estimator.predict_lpte(
    target_treatment_arm=1,
    control_treatment_arm=0,
    locations=locations
)

print(f"ML-adjusted LPTE shape: {lpte.shape}")  # Should be (4,)
print(f"Interval effects: {lpte}")

predict_pte(target_treatment_arm: int, control_treatment_arm: int, locations: ndarray, alpha: float = 0.05, variance_type='moment', n_bootstrap=500) → Tuple[ndarray, ndarray, ndarray]

Compute Probability Treatment Effects (PTE) based on the estimator for the distribution function.

The PTE measures the difference in probability mass between treatment groups for intervals defined by consecutive location pairs. It quantifies how treatment affects the probability of observing outcomes within specific ranges.

Parameters:

• 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 defining interval boundaries for probability computation. For each interval (locations[i], locations[i+1]], the PTE is computed.

• 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”. Defaults to “moment”.

• n_bootstrap (int, optional) – Number of bootstrap samples. Defaults to 500.

Returns:

A tuple containing:

• Expected PTEs (np.ndarray): Treatment effect estimates for each interval, shape (len(locations)-1,)

• Lower bounds (np.ndarray): Lower confidence interval bounds

• Upper bounds (np.ndarray): Upper confidence interval bounds

Return type:

Tuple[np.ndarray, np.ndarray, np.ndarray]

Example

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}")

predict_qte(target_treatment_arm: int, control_treatment_arm: int, quantiles: ndarray | None = None, alpha: float = 0.05, n_bootstrap=500) → Tuple[ndarray, ndarray, ndarray]

Compute Quantile Treatment Effects (QTE) based on the estimator for the distribution function.

The QTE measures the difference in quantiles between treatment groups, providing insights into how treatment affects different parts of the outcome distribution. For stratified estimators, the computation properly accounts for strata.

Parameters:

• 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, 0.2, …, 0.9].

• 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:

A tuple containing:

• Expected QTEs (np.ndarray): Treatment effect estimates at each quantile

• Lower bounds (np.ndarray): Lower confidence interval bounds

• Upper bounds (np.ndarray): Upper confidence interval bounds

Return type:

Tuple[np.ndarray, np.ndarray, np.ndarray]

Example

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}")

class dte_adj.AdjustedStratifiedDistributionEstimator(base_model: Any, folds=3, is_multi_task=False)

Bases: DistributionEstimatorBase

A class is for estimating the adjusted distribution function and computing the Distributional parameters for CAR.

fit(covariates: ndarray, treatment_arms: ndarray, outcomes: ndarray, strata: ndarray) → DistributionEstimatorBase

Train the DistributionEstimatorBase.

Parameters:

• covariates (np.ndarray) – Pre-treatment covariates.

• treatment_arms (np.ndarray) – The index of the treatment arm.

• outcomes (np.ndarray) – Scalar-valued observed outcome.

Returns:

The fitted estimator.

Return type:

DistributionEstimatorBase

predict(treatment_arm: int, locations: ndarray) → ndarray

Compute cumulative distribution values.

Parameters:

• treatment_arm (int) – The index of the treatment arm.

• outcomes (np.ndarray) – Scalar values to be used for computing the cumulative distribution.

Returns:

Estimated cumulative distribution values for the input.

Return type:

np.ndarray

predict_dte(target_treatment_arm: int, control_treatment_arm: int, locations: ndarray, alpha: float = 0.05, variance_type='moment', n_bootstrap=500) → Tuple[ndarray, ndarray, ndarray]

Compute Distribution Treatment Effects (DTE) based on the estimator for the distribution function.

The DTE measures the difference in cumulative distribution functions between treatment groups at specified locations. It quantifies how treatment affects the probability of observing outcomes below each threshold.

Parameters:

• 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 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”. Defaults to “moment”.

• n_bootstrap (int, optional) – Number of bootstrap samples. Defaults to 500.

Returns:

A tuple containing:

• Expected DTEs (np.ndarray): Treatment effect estimates at each location

• Lower bounds (np.ndarray): Lower confidence interval bounds

• Upper bounds (np.ndarray): Upper confidence interval bounds

Return type:

Tuple[np.ndarray, np.ndarray, np.ndarray]

Example

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}")

predict_pte(target_treatment_arm: int, control_treatment_arm: int, locations: ndarray, alpha: float = 0.05, variance_type='moment', n_bootstrap=500) → Tuple[ndarray, ndarray, ndarray]

Compute Probability Treatment Effects (PTE) based on the estimator for the distribution function.

The PTE measures the difference in probability mass between treatment groups for intervals defined by consecutive location pairs. It quantifies how treatment affects the probability of observing outcomes within specific ranges.

Parameters:

• 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 defining interval boundaries for probability computation. For each interval (locations[i], locations[i+1]], the PTE is computed.

• 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”. Defaults to “moment”.

• n_bootstrap (int, optional) – Number of bootstrap samples. Defaults to 500.

Returns:

A tuple containing:

• Expected PTEs (np.ndarray): Treatment effect estimates for each interval, shape (len(locations)-1,)

• Lower bounds (np.ndarray): Lower confidence interval bounds

• Upper bounds (np.ndarray): Upper confidence interval bounds

Return type:

Tuple[np.ndarray, np.ndarray, np.ndarray]

Example

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}")

predict_qte(target_treatment_arm: int, control_treatment_arm: int, quantiles: ndarray | None = None, alpha: float = 0.05, n_bootstrap=500) → Tuple[ndarray, ndarray, ndarray]

Compute Quantile Treatment Effects (QTE) based on the estimator for the distribution function.

The QTE measures the difference in quantiles between treatment groups, providing insights into how treatment affects different parts of the outcome distribution. For stratified estimators, the computation properly accounts for strata.

Parameters:

• 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, 0.2, …, 0.9].

• 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:

A tuple containing:

• Expected QTEs (np.ndarray): Treatment effect estimates at each quantile

• Lower bounds (np.ndarray): Lower confidence interval bounds

• Upper bounds (np.ndarray): Upper confidence interval bounds

Return type:

Tuple[np.ndarray, np.ndarray, np.ndarray]

Example

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}")

class dte_adj.SimpleDistributionEstimator

Bases: SimpleStratifiedDistributionEstimator

A class for computing the empirical distribution function and distributional treatment effects using simple (unadjusted) estimation methods.

This estimator computes Distribution Treatment Effects (DTE), Probability Treatment Effects (PTE), and Quantile Treatment Effects (QTE) without using machine learning models for adjustment. It provides a baseline approach suitable when treatment assignment is random or when covariate adjustment is not needed.

Example

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)

fit(covariates: ndarray, treatment_arms: ndarray, outcomes: ndarray) → SimpleDistributionEstimator

Set parameters.

Parameters:

• covariates (np.ndarray) – Pre-treatment covariates.

• treatment_arms (np.ndarray) – The index of the treatment arm.

• outcomes (np.ndarray) – Scalar-valued observed outcome.

Returns:

The fitted estimator.

Return type:

SimpleDistributionEstimator

predict(treatment_arm: int, locations: ndarray) → ndarray

Compute cumulative distribution values.

Parameters:

• treatment_arm (int) – The index of the treatment arm.

• outcomes (np.ndarray) – Scalar values to be used for computing the cumulative distribution.

Returns:

Estimated cumulative distribution values for the input.

Return type:

np.ndarray

predict_dte(target_treatment_arm: int, control_treatment_arm: int, locations: ndarray, alpha: float = 0.05, variance_type='moment', n_bootstrap=500) → Tuple[ndarray, ndarray, ndarray]

Compute Distribution Treatment Effects (DTE) based on the estimator for the distribution function.

The DTE measures the difference in cumulative distribution functions between treatment groups at specified locations. It quantifies how treatment affects the probability of observing outcomes below each threshold.

Parameters:

• 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 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”. Defaults to “moment”.

• n_bootstrap (int, optional) – Number of bootstrap samples. Defaults to 500.

Returns:

A tuple containing:

• Expected DTEs (np.ndarray): Treatment effect estimates at each location

• Lower bounds (np.ndarray): Lower confidence interval bounds

• Upper bounds (np.ndarray): Upper confidence interval bounds

Return type:

Tuple[np.ndarray, np.ndarray, np.ndarray]

Example

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}")

predict_pte(target_treatment_arm: int, control_treatment_arm: int, locations: ndarray, alpha: float = 0.05, variance_type='moment', n_bootstrap=500) → Tuple[ndarray, ndarray, ndarray]

Compute Probability Treatment Effects (PTE) based on the estimator for the distribution function.

The PTE measures the difference in probability mass between treatment groups for intervals defined by consecutive location pairs. It quantifies how treatment affects the probability of observing outcomes within specific ranges.

Parameters:

• 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 defining interval boundaries for probability computation. For each interval (locations[i], locations[i+1]], the PTE is computed.

• 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”. Defaults to “moment”.

• n_bootstrap (int, optional) – Number of bootstrap samples. Defaults to 500.

Returns:

A tuple containing:

• Expected PTEs (np.ndarray): Treatment effect estimates for each interval, shape (len(locations)-1,)

• Lower bounds (np.ndarray): Lower confidence interval bounds

• Upper bounds (np.ndarray): Upper confidence interval bounds

Return type:

Tuple[np.ndarray, np.ndarray, np.ndarray]

Example

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}")

predict_qte(target_treatment_arm: int, control_treatment_arm: int, quantiles: ndarray | None = None, alpha: float = 0.05, n_bootstrap=500) → Tuple[ndarray, ndarray, ndarray]

Compute Quantile Treatment Effects (QTE) based on the estimator for the distribution function.

The QTE measures the difference in quantiles between treatment groups, providing insights into how treatment affects different parts of the outcome distribution. For stratified estimators, the computation properly accounts for strata.

Parameters:

• 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, 0.2, …, 0.9].

• 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:

A tuple containing:

• Expected QTEs (np.ndarray): Treatment effect estimates at each quantile

• Lower bounds (np.ndarray): Lower confidence interval bounds

• Upper bounds (np.ndarray): Upper confidence interval bounds

Return type:

Tuple[np.ndarray, np.ndarray, np.ndarray]

Example

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}")

class dte_adj.SimpleLocalDistributionEstimator

Bases: SimpleStratifiedDistributionEstimator

A class for computing Local Distribution Treatment Effects (LDTE) and Local Probability Treatment Effects (LPTE) using simple empirical estimation.

This estimator computes treatment effects that are weighted by treatment propensity within each stratum, providing estimates that are locally robust to treatment assignment heterogeneity across strata. It uses empirical methods without ML adjustment.

fit(covariates: ndarray, treatment_arms: ndarray, treatment_indicator: ndarray, outcomes: ndarray, strata: ndarray) → SimpleLocalDistributionEstimator

Train the SimpleLocalDistributionEstimator.

Parameters:

• covariates (np.ndarray) – Pre-treatment covariates.

• treatment_arms (np.ndarray) – Treatment assignment variable (Z).

• treatment_indicator (np.ndarray) – Treatment indicator variable (D).

• outcomes (np.ndarray) – Scalar-valued observed outcome.

• strata (np.ndarray) – Stratum indicators.

Returns:

The fitted estimator.

Return type:

SimpleLocalDistributionEstimator

predict(treatment_arm: int, locations: ndarray) → ndarray

Compute cumulative distribution values.

Parameters:

• treatment_arm (int) – The index of the treatment arm.

• outcomes (np.ndarray) – Scalar values to be used for computing the cumulative distribution.

Returns:

Estimated cumulative distribution values for the input.

Return type:

np.ndarray

predict_dte(target_treatment_arm: int, control_treatment_arm: int, locations: ndarray, alpha: float = 0.05, variance_type='moment', n_bootstrap=500) → Tuple[ndarray, ndarray, ndarray]

Compute Distribution Treatment Effects (DTE) based on the estimator for the distribution function.

The DTE measures the difference in cumulative distribution functions between treatment groups at specified locations. It quantifies how treatment affects the probability of observing outcomes below each threshold.

Parameters:

• 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 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”. Defaults to “moment”.

• n_bootstrap (int, optional) – Number of bootstrap samples. Defaults to 500.

Returns:

A tuple containing:

• Expected DTEs (np.ndarray): Treatment effect estimates at each location

• Lower bounds (np.ndarray): Lower confidence interval bounds

• Upper bounds (np.ndarray): Upper confidence interval bounds

Return type:

Tuple[np.ndarray, np.ndarray, np.ndarray]

Example

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}")

predict_ldte(target_treatment_arm: int, control_treatment_arm: int, locations: ndarray, alpha: float = 0.05) → Tuple[ndarray, ndarray, ndarray]

Compute Local Distribution Treatment Effects (LDTE).

LDTE measures the difference in cumulative distribution functions between treatment groups weighted by treatment propensity within each stratum. This provides estimates that are locally robust to treatment assignment heterogeneity across strata.

Parameters:

• 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 bound. Defaults to 0.05.

Returns:

A tuple containing:

• Expected LDTEs (np.ndarray): Local treatment effect estimates at each location

• Lower bounds (np.ndarray): Lower confidence interval bounds

• Upper bounds (np.ndarray): Upper confidence interval bounds

Return type:

Tuple[np.ndarray, np.ndarray, np.ndarray]

Example

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}")

predict_lpte(target_treatment_arm: int, control_treatment_arm: int, locations: ndarray, alpha: float = 0.05) → Tuple[ndarray, ndarray, ndarray]

Compute Local Probability Treatment Effects (LPTE).

LPTE measures the difference in probability mass between treatment groups for intervals weighted by treatment propensity within each stratum. This provides interval-based treatment effect estimates that are locally robust to treatment assignment heterogeneity.

Parameters:

• 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 defining interval boundaries for probability computation. For each interval (locations[i], locations[i+1]], the LPTE is computed.

• alpha (float, optional) – Significance level of the confidence bound. Defaults to 0.05.

Returns:

A tuple containing:

• Expected LPTEs (np.ndarray): Local treatment effect estimates for each interval, shape (len(locations)-1,)

• Lower bounds (np.ndarray): Lower confidence interval bounds

• Upper bounds (np.ndarray): Upper confidence interval bounds

Return type:

Tuple[np.ndarray, np.ndarray, np.ndarray]

Example

import numpy as np
from sklearn.linear_model import LogisticRegression
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, 2], size=1000)  # Multiple strata
D = np.random.binomial(1, 0.2 + 0.3 * (strata == 1) + 0.4 * (strata == 2), 1000)
Y = X[:, 0] + 1.5 * D + 0.5 * strata + np.random.randn(1000)

# Fit simple local estimator
estimator = SimpleLocalDistributionEstimator()
estimator.fit(X, D, D, Y, strata)

# Define interval boundaries
locations = np.array([-2, -1, 0, 1, 2])  # Creates 4 intervals

# 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}")

predict_pte(target_treatment_arm: int, control_treatment_arm: int, locations: ndarray, alpha: float = 0.05, variance_type='moment', n_bootstrap=500) → Tuple[ndarray, ndarray, ndarray]

Compute Probability Treatment Effects (PTE) based on the estimator for the distribution function.

The PTE measures the difference in probability mass between treatment groups for intervals defined by consecutive location pairs. It quantifies how treatment affects the probability of observing outcomes within specific ranges.

Parameters:

• 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 defining interval boundaries for probability computation. For each interval (locations[i], locations[i+1]], the PTE is computed.

• 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”. Defaults to “moment”.

• n_bootstrap (int, optional) – Number of bootstrap samples. Defaults to 500.

Returns:

A tuple containing:

• Expected PTEs (np.ndarray): Treatment effect estimates for each interval, shape (len(locations)-1,)

• Lower bounds (np.ndarray): Lower confidence interval bounds

• Upper bounds (np.ndarray): Upper confidence interval bounds

Return type:

Tuple[np.ndarray, np.ndarray, np.ndarray]

Example

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}")

predict_qte(target_treatment_arm: int, control_treatment_arm: int, quantiles: ndarray | None = None, alpha: float = 0.05, n_bootstrap=500) → Tuple[ndarray, ndarray, ndarray]

Compute Quantile Treatment Effects (QTE) based on the estimator for the distribution function.

The QTE measures the difference in quantiles between treatment groups, providing insights into how treatment affects different parts of the outcome distribution. For stratified estimators, the computation properly accounts for strata.

Parameters:

• 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, 0.2, …, 0.9].

• 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:

A tuple containing:

• Expected QTEs (np.ndarray): Treatment effect estimates at each quantile

• Lower bounds (np.ndarray): Lower confidence interval bounds

• Upper bounds (np.ndarray): Upper confidence interval bounds

Return type:

Tuple[np.ndarray, np.ndarray, np.ndarray]

Example

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}")

class dte_adj.SimpleStratifiedDistributionEstimator

Bases: DistributionEstimatorBase

A class is for estimating the empirical distribution function and computing the Distributional parameters for CAR.

fit(covariates: ndarray, treatment_arms: ndarray, outcomes: ndarray, strata: ndarray) → DistributionEstimatorBase

Train the DistributionEstimatorBase.

Parameters:

• covariates (np.ndarray) – Pre-treatment covariates.

• treatment_arms (np.ndarray) – The index of the treatment arm.

• outcomes (np.ndarray) – Scalar-valued observed outcome.

Returns:

The fitted estimator.

Return type:

DistributionEstimatorBase

predict(treatment_arm: int, locations: ndarray) → ndarray

Compute cumulative distribution values.

Parameters:

• treatment_arm (int) – The index of the treatment arm.

• outcomes (np.ndarray) – Scalar values to be used for computing the cumulative distribution.

Returns:

Estimated cumulative distribution values for the input.

Return type:

np.ndarray

predict_dte(target_treatment_arm: int, control_treatment_arm: int, locations: ndarray, alpha: float = 0.05, variance_type='moment', n_bootstrap=500) → Tuple[ndarray, ndarray, ndarray]

Compute Distribution Treatment Effects (DTE) based on the estimator for the distribution function.

The DTE measures the difference in cumulative distribution functions between treatment groups at specified locations. It quantifies how treatment affects the probability of observing outcomes below each threshold.

Parameters:

• 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 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”. Defaults to “moment”.

• n_bootstrap (int, optional) – Number of bootstrap samples. Defaults to 500.

Returns:

A tuple containing:

• Expected DTEs (np.ndarray): Treatment effect estimates at each location

• Lower bounds (np.ndarray): Lower confidence interval bounds

• Upper bounds (np.ndarray): Upper confidence interval bounds

Return type:

Tuple[np.ndarray, np.ndarray, np.ndarray]

Example

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}")

predict_pte(target_treatment_arm: int, control_treatment_arm: int, locations: ndarray, alpha: float = 0.05, variance_type='moment', n_bootstrap=500) → Tuple[ndarray, ndarray, ndarray]

Compute Probability Treatment Effects (PTE) based on the estimator for the distribution function.

The PTE measures the difference in probability mass between treatment groups for intervals defined by consecutive location pairs. It quantifies how treatment affects the probability of observing outcomes within specific ranges.

Parameters:

• 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 defining interval boundaries for probability computation. For each interval (locations[i], locations[i+1]], the PTE is computed.

• 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”. Defaults to “moment”.

• n_bootstrap (int, optional) – Number of bootstrap samples. Defaults to 500.

Returns:

A tuple containing:

• Expected PTEs (np.ndarray): Treatment effect estimates for each interval, shape (len(locations)-1,)

• Lower bounds (np.ndarray): Lower confidence interval bounds

• Upper bounds (np.ndarray): Upper confidence interval bounds

Return type:

Tuple[np.ndarray, np.ndarray, np.ndarray]

Example

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}")

predict_qte(target_treatment_arm: int, control_treatment_arm: int, quantiles: ndarray | None = None, alpha: float = 0.05, n_bootstrap=500) → Tuple[ndarray, ndarray, ndarray]

Compute Quantile Treatment Effects (QTE) based on the estimator for the distribution function.

The QTE measures the difference in quantiles between treatment groups, providing insights into how treatment affects different parts of the outcome distribution. For stratified estimators, the computation properly accounts for strata.

Parameters:

• 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, 0.2, …, 0.9].

• 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:

A tuple containing:

• Expected QTEs (np.ndarray): Treatment effect estimates at each quantile

• Lower bounds (np.ndarray): Lower confidence interval bounds

• Upper bounds (np.ndarray): Upper confidence interval bounds

Return type:

Tuple[np.ndarray, np.ndarray, np.ndarray]

Example

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}")

dte_adj.plot.plot(X: ndarray, means: ndarray, lower_bounds: ndarray, upper_bounds: ndarray, chart_type: str = 'line', color: str = 'green', ax: Axis | None = None, title: str | None = None, xlabel: str | None = None, ylabel: str | None = None)

Visualize distributional parameters and their confidence intervals.

Parameters:

• X (np.Array) – values to be used for x axis.

• means (np.Array) – Expected distributional parameters.

• lower_bounds (np.Array) – Lower bound for the distributional parameters.

• upper_bounds (np.Array) – Upper bound for the distributional parameters.

• chart_type (str) – Chart type of the plotting. Available values are line or bar.

• color (str) – The color of lines or bars.

• 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.

• ylabel (str, optional) – Y-axis title label.

Returns:

The axes with the plot.

Return type:

matplotlib.axes.Axes