Modules

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

Bases: AdjustedStratifiedDistributionEstimator

A class is for estimating the adjusted distribution function and computing the Distributional parameters based on the trained conditional estimator.

fit(covariates: ndarray, treatment_arms: ndarray, outcomes: ndarray) → AdjustedDistributionEstimator[source]

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 DTE based on the estimator for the distribution function.

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.

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

Returns:

A tuple containing:

  • Expected DTEs

  • Upper bounds

  • Lower bounds

Return type:

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

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

Compute PTE based on the estimator for the distribution function.

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.

  • width (float) – The width of each outcome interval.

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

  • variance_type (str, optional) – Variance type to be used to compute confidence intervals. Available values are moment, simple, and uniform.

Returns:

A tuple containing:

  • Expected PTEs

  • Upper bounds

  • Lower bounds

Return type:

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

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 QTE based on the estimator for the distribution function.

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 * i for i in range(1, 10)].

  • 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

  • Upper bounds

  • Lower bounds

Return type:

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

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 DTE based on the estimator for the distribution function.

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.

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

Returns:

A tuple containing:

  • Expected DTEs

  • Upper bounds

  • Lower bounds

Return type:

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

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

Compute PTE based on the estimator for the distribution function.

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.

  • width (float) – The width of each outcome interval.

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

  • variance_type (str, optional) – Variance type to be used to compute confidence intervals. Available values are moment, simple, and uniform.

Returns:

A tuple containing:

  • Expected PTEs

  • Upper bounds

  • Lower bounds

Return type:

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

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 QTE based on the estimator for the distribution function.

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 * i for i in range(1, 10)].

  • 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

  • Upper bounds

  • Lower bounds

Return type:

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

class dte_adj.SimpleDistributionEstimator[source]

Bases: SimpleStratifiedDistributionEstimator

A class for computing the empirical distribution function and the distributional parameters based on the distribution function.

fit(covariates: ndarray, treatment_arms: ndarray, outcomes: ndarray) → SimpleDistributionEstimator[source]

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 DTE based on the estimator for the distribution function.

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.

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

Returns:

A tuple containing:

  • Expected DTEs

  • Upper bounds

  • Lower bounds

Return type:

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

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

Compute PTE based on the estimator for the distribution function.

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.

  • width (float) – The width of each outcome interval.

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

  • variance_type (str, optional) – Variance type to be used to compute confidence intervals. Available values are moment, simple, and uniform.

Returns:

A tuple containing:

  • Expected PTEs

  • Upper bounds

  • Lower bounds

Return type:

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

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 QTE based on the estimator for the distribution function.

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 * i for i in range(1, 10)].

  • 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

  • Upper bounds

  • Lower bounds

Return type:

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

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 DTE based on the estimator for the distribution function.

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.

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

Returns:

A tuple containing:

  • Expected DTEs

  • Upper bounds

  • Lower bounds

Return type:

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

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

Compute PTE based on the estimator for the distribution function.

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.

  • width (float) – The width of each outcome interval.

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

  • variance_type (str, optional) – Variance type to be used to compute confidence intervals. Available values are moment, simple, and uniform.

Returns:

A tuple containing:

  • Expected PTEs

  • Upper bounds

  • Lower bounds

Return type:

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

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 QTE based on the estimator for the distribution function.

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 * i for i in range(1, 10)].

  • 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

  • Upper bounds

  • Lower bounds

Return type:

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

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

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