Covariate Adaptive Randomization Estimators

This page documents estimators that work with stratified experimental designs, particularly for covariate-adaptive randomization (CAR) within strata.

These estimators are designed to handle stratified block randomization where participants are grouped into strata based on baseline covariates before treatment assignment. The key methodological contribution is leveraging additional covariates beyond strata indicators using machine learning methods to enhance the precision of distributional treatment effect estimates.

Byambadalai et al. (2025) [1] propose a flexible distribution regression framework that achieves the semiparametric efficiency bound for distributional treatment effects under CAR, demonstrating that regression-adjusted estimators can optimally utilize covariate information in stratified designs.

SimpleStratifiedDistributionEstimator

class dte_adj.SimpleStratifiedDistributionEstimator[source]

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[source]

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

AdjustedStratifiedDistributionEstimator

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

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[source]

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