Python API¶
All classes are available from the top-level package, e.g. dml_rs.PLR.
Models¶
PLR(
ml_l: Any,
ml_m: Any,
ml_g: Any = None,
score: Union[
Literal["partialling out", "IV-type"],
Callable[..., Any],
] = "partialling out",
n_folds: int = 5,
n_rep: int = 1,
random_state: int = 0,
)
Bases: _HeterogeneityMixin, _DMLEstimator
Partially linear regression Y = θ·D + g(X) + ε.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
ml_l
|
regressor for ``E[Y | X]`` — a ``dml_rs`` learner or any object with
|
|
required |
ml_m
|
regressor for ``E[D | X]``.
|
|
required |
ml_g
|
regressor for ``g(X)``, used by ``score="IV-type"`` (default: a clone of ``ml_l``).
|
|
None
|
score
|
``"partialling out"``, ``"IV-type"`` or a callable ``score(y, d, l_hat, m_hat, g_hat,
|
smpls) -> (psi_a, psi_b) |
'partialling out'
|
n_folds
|
cross-fitting folds and repetitions of the sample split.
|
|
5
|
n_rep
|
cross-fitting folds and repetitions of the sample split.
|
|
5
|
random_state
|
seed for the sample splits.
|
|
0
|
fit(
X: ArrayLike,
y: ArrayLike,
d: ArrayLike,
folds: Optional[Folds] = None,
cluster: Optional[ArrayLike] = None,
external_predictions: Optional[dict[Any, Any]] = None,
store_models: bool = False,
) -> Self
Fit on covariates X, outcome y and treatment(s) d.
folds optionally fixes the sample splits: a list of (train, test)
index pairs, or one such list per repetition.
cluster gives cluster labels, one column (one-way) or two columns (two-way
clustering, as Chiang et al. 2022): whole clusters are then assigned to folds and
the standard errors are cluster-robust. Accepted by PLR, IRM, PLIV, IIVM, DID and
DIDCS; the multiplier bootstrap is not available with clusters (as in DoubleML).
store_models=True keeps the fitted learners in models_ ({learner: {treatment:
[[model per fold] per repetition]}}, as DoubleML's models); folds are then fitted one
after the other.
LPLR(
ml_M: Any,
ml_t: Any,
ml_m: Any,
ml_a: Any = None,
score: Literal[
"nuisance_space", "instrument"
] = "nuisance_space",
n_folds: int = 5,
n_folds_inner: int = 5,
n_rep: int = 1,
random_state: int = 0,
)
Bases: _DMLEstimator
Logistic partially linear regression for a binary outcome, as DoubleML's DoubleMLLPLR
(Liu, Zhang and Zhou, 2021): P(Y = 1 | D, X) = expit(θ·D + r(X)), so θ is a log odds
ratio.
ml_M and ml_a are double cross-fitted (inner folds within every outer training fold),
and the score is solved with Newton's method from a preliminary estimate.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
ml_M
|
classifier for ``P(Y = 1 | D, X)``.
|
|
required |
ml_t
|
regressor for the log odds ``E[logit M(D, X) | X]``.
|
|
required |
ml_m
|
regressor for ``E[D | X]``, or a classifier for a binary treatment.
|
|
required |
ml_a
|
learner for ``E[D | X]`` used for the preliminary estimate (default: a copy of ``ml_m``).
|
|
None
|
score
|
``"nuisance_space"`` (``ml_m`` fitted on ``Y = 0``) or ``"instrument"`` (``ml_m`` fitted
|
with weights |
'nuisance_space'
|
n_folds_inner
|
inner folds of the double cross-fitting.
|
|
5
|
fit(
X: ArrayLike,
y: ArrayLike,
d: ArrayLike,
folds: Optional[Folds] = None,
inner_folds: Optional[Sequence[Any]] = None,
external_predictions: Any = None,
store_models: bool = False,
) -> Self
inner_folds (with folds) fixes the inner splits in DoubleML's smpls_inner
layout: per repetition, per outer fold, a list of (train, test) index pairs within that
fold's training set. external_predictions as DoubleML: external ml_M and ml_a
also need their inner predictions ml_M_inner_k / ml_a_inner_k for every outer fold
k (a fit reports them in predictions_).
store_models=True keeps the fitted learners in models_ ({learner: {treatment:
[[model per fold] per repetition]}}, as DoubleML's models); folds are then fitted one
after the other. ml_M and ml_a are fitted on every inner fold and
on the outer fold, so their entries are lists.
PLPR(
ml_l: Any,
ml_m: Any,
ml_g: Any = None,
score: Literal[
"partialling out", "IV-type"
] = "partialling out",
approach: Literal[
"cre_general", "cre_normal", "fd_exact", "wg_approx"
] = "fd_exact",
n_folds: int = 5,
n_rep: int = 1,
random_state: int = 0,
)
Bases: _DMLEstimator
Partially linear regression for static panels, as DoubleML's DoubleMLPLPR (Clarke and
Polselli, 2025): Y_it = θ·D_it + g(X_it) + α_i + U_it with unit effects correlated with the
covariates and the treatment.
The panel (long format, one row per unit id and period t) is transformed according to
approach, and a partially linear regression is fitted on the transformed rows with folds of
whole units and unit-clustered standard errors.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
ml_l
|
regressor for ``E[Y | X]`` on the transformed panel.
|
|
required |
ml_m
|
regressor for ``E[D | X]``, or a classifier for a binary treatment with the correlated
|
random effects approaches. |
required |
ml_g
|
regressor for ``E[Y - θD | X]``, used by ``score="IV-type"`` (default: a copy of ``ml_l``).
|
|
None
|
score
|
``"partialling out"`` or ``"IV-type"``.
|
|
'partialling out'
|
approach
|
``"fd_exact"`` (first differences, covariates and their lags), ``"wg_approx"``
|
(within-unit demeaning), |
'fd_exact'
|
fit(
X: ArrayLike,
y: ArrayLike,
d: ArrayLike,
id: ArrayLike,
t: ArrayLike,
folds: Optional[Folds] = None,
external_predictions: Optional[dict[Any, Any]] = None,
store_models: bool = False,
) -> Self
folds (optional) index the rows of the transformed panel, sorted by unit and period,
as DoubleML's smpls; each test set should hold whole units.
store_models=True keeps the fitted learners in models_ ({learner: {treatment:
[[model per fold] per repetition]}}, as DoubleML's models); folds are then fitted one
after the other.
IRM(
ml_g: Any,
ml_m: Any,
score: Union[
Literal["ATE", "ATTE"], Callable[..., Any]
] = "ATE",
clipping_threshold: float = 0.01,
calibration: Optional[
Literal["isotonic", "isotonic_cv"]
] = None,
normalize_ipw: bool = False,
n_folds: int = 5,
n_rep: int = 1,
random_state: int = 0,
weights: Any = None,
)
Bases: _HeterogeneityMixin, _DMLEstimator
Interactive regression model for a binary treatment (ATE or ATTE).
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
ml_g
|
regressor for ``E[Y | D, X]`` (fitted separately on treated and controls).
|
|
required |
ml_m
|
classifier for the propensity ``P(D = 1 | X)`` with ``predict_proba``.
|
|
required |
score
|
``"ATE"``, ``"ATTE"`` or a callable ``score(y, d, g_hat0, g_hat1, m_hat, smpls)
|
-> (psi_a, psi_b) |
'ATE'
|
clipping_threshold
|
propensities are clipped to ``[t, 1 - t]``.
|
|
0.01
|
normalize_ipw
|
normalise the inverse-probability weights.
|
|
False
|
weights
|
weighted average effect, as DoubleML: an array ``w`` (ATE score: ``E[w (g1 - g0)] /
|
E[w] |
None
|
fit(
X: ArrayLike,
y: ArrayLike,
d: ArrayLike,
folds: Optional[Folds] = None,
cluster: Optional[ArrayLike] = None,
external_predictions: Optional[dict[Any, Any]] = None,
store_models: bool = False,
) -> Self
Fit on covariates X, outcome y and binary treatment d.
store_models=True keeps the fitted learners in models_ ({learner: {treatment:
[[model per fold] per repetition]}}, as DoubleML's models); folds are then fitted one
after the other.
policy_tree(
features: ArrayLike,
depth: int = 2,
min_samples_leaf: int = 8,
ccp_alpha: float = 0.01,
) -> PolicyTree
Policy tree recommending treatment where the estimated effect is positive (ATE score, one
repetition), as DoubleML's policy_tree: a decision tree of the sign of the orthogonal
signal weighted by its size, pruned with cost-complexity ccp_alpha.
APO(
ml_g: Any,
ml_m: Any,
treatment_level: float,
clipping_threshold: float = 0.01,
calibration: Optional[
Literal["isotonic", "isotonic_cv"]
] = None,
normalize_ipw: bool = False,
n_folds: int = 5,
n_rep: int = 1,
random_state: int = 0,
weights: Any = None,
)
Bases: _DMLEstimator
Average potential outcome E[Y(d)] of one level d of a discrete treatment.
Default sample splits are stratified by the treatment, as in DoubleML.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
ml_g
|
regressor for ``E[Y | D = d, X]``.
|
|
required |
ml_m
|
classifier for the generalised propensity ``P(D = d | X)``.
|
|
required |
treatment_level
|
the level ``d``.
|
|
required |
clipping_threshold
|
propensities are clipped to ``[t, 1 - t]``.
|
|
0.01
|
normalize_ipw
|
normalise the inverse-probability weights.
|
|
False
|
fit(
X: ArrayLike,
y: ArrayLike,
d: ArrayLike,
folds: Optional[Folds] = None,
cluster: Optional[ArrayLike] = None,
external_predictions: Optional[dict[Any, Any]] = None,
store_models: bool = False,
) -> Self
Fit on covariates X, outcome y and discrete treatment d.
store_models=True keeps the fitted learners in models_ ({learner: {treatment:
[[model per fold] per repetition]}}, as DoubleML's models); folds are then fitted one
after the other.
capo(basis: ArrayLike) -> BLP
Best linear predictor of the conditional average potential outcome, as DoubleML's capo.
gapo(groups: ArrayLike) -> BLP
Group average potential outcomes for 0/1 group dummies, as DoubleML's gapo.
APOS(
ml_g: Any,
ml_m: Any,
treatment_levels: Sequence[float],
clipping_threshold: float = 0.01,
calibration: Optional[
Literal["isotonic", "isotonic_cv"]
] = None,
normalize_ipw: bool = False,
n_folds: int = 5,
n_rep: int = 1,
random_state: int = 0,
weights: Any = None,
)
Bases: _DMLEstimator
Average potential outcomes of several levels of a discrete treatment, and contrasts.
All levels are fitted on the same sample splits, so their estimates and contrasts have joint
inference (bootstrap, confint(joint=True)). Parameters as APO, with
treatment_levels the levels to estimate.
fit(
X: ArrayLike,
y: ArrayLike,
d: ArrayLike,
folds: Optional[Folds] = None,
cluster: Optional[ArrayLike] = None,
) -> Self
causal_contrast(
reference_levels: Union[float, Sequence[float]],
) -> CausalContrast
Average effects of each level against each reference level, as DoubleML's causal_contrast.
For every reference level in turn, all levels not yet used as a reference are compared with
it ("d vs d_ref"). Inference accounts for the correlation between the estimates; with
clustered data the standard errors are cluster-robust.
CausalContrast(result: Any)
Bases: Framework
Contrasts between average potential outcomes, from APOS.causal_contrast.
Offers the fitted-model API (summary, confint, bootstrap, p_adjust,
sensitivity_analysis); there is nothing to refit.
QTE(
ml_g: Any,
ml_m: Any = None,
quantiles: Union[float, Sequence[float]] = 0.5,
score: Literal["PQ", "LPQ", "CVaR"] = "PQ",
clipping_threshold: float = 0.01,
calibration: Optional[
Literal["isotonic", "isotonic_cv"]
] = None,
normalize_ipw: bool = True,
n_folds: int = 5,
n_rep: int = 1,
random_state: int = 0,
)
Bases: _DMLEstimator
Quantile treatment effects q(Y(1)) - q(Y(0)) at several quantiles, as DoubleML's
DoubleMLQTE. score is "PQ" (quantiles), "LPQ" (compliers' quantiles, needs
z) or "CVaR" (effects on the conditional value at risk).
Both potential quantiles are fitted on the same sample splits at every quantile, so the effects
have joint inference (bootstrap, confint(joint=True)). Other parameters as PQ.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
ml_g
|
classifier (``"PQ"``, ``"LPQ"``) or regressor (``"CVaR"``), as in `PQ`, `LPQ`, `CVAR`.
|
|
required |
ml_m
|
propensity classifier (default: ``ml_g``, which must then be a classifier).
|
|
None
|
quantiles
|
the quantile levels.
|
|
0.5
|
PotentialQuantiles(result: Any)
Bases: Framework
Potential quantiles (or CVaRs) of both treatment levels, from QTE.potential, ordered
q1 (d=0), q1 (d=1), q2 (d=0), ... and with joint inference.
PQ(
ml_g: Any,
ml_m: Any,
treatment: int = 1,
quantile: float = 0.5,
clipping_threshold: float = 0.01,
calibration: Optional[
Literal["isotonic", "isotonic_cv"]
] = None,
normalize_ipw: bool = True,
n_folds: int = 5,
n_rep: int = 1,
random_state: int = 0,
)
Bases: _QuantileBase
Potential quantile q with P(Y(d) <= q) = quantile for a binary treatment, as
DoubleML's DoubleMLPQ.
Default sample splits are stratified by the treatment. Each fold runs a nested split for a preliminary inverse-probability-weighted estimate, at which the conditional distribution function is learned; the score is then solved with Brent's method and the standard error uses a kernel density estimate.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
ml_g
|
classifier for ``P(Y <= q | D = d, X)`` with ``predict_proba``.
|
|
required |
ml_m
|
classifier for the propensity ``P(D = 1 | X)``.
|
|
required |
treatment
|
the treatment level ``d`` (0 or 1).
|
|
1
|
quantile
|
the quantile level in ``(0, 1)``.
|
|
0.5
|
clipping_threshold
|
propensities are clipped to ``[t, 1 - t]``.
|
|
0.01
|
normalize_ipw
|
normalise the inverse-probability weights (default True, as DoubleML).
|
|
True
|
fit(
X: ArrayLike,
y: ArrayLike,
d: ArrayLike,
folds: Optional[Folds] = None,
external_predictions: Any = None,
) -> Self
external_predictions for ml_g and / or ml_m as DoubleML. A supplied ml_m
also serves the preliminary estimate (on the rows of the first half of each training fold).
LPQ(
ml_g: Any,
ml_m: Any,
treatment: int = 1,
quantile: float = 0.5,
clipping_threshold: float = 0.01,
calibration: Optional[
Literal["isotonic", "isotonic_cv"]
] = None,
normalize_ipw: bool = True,
n_folds: int = 5,
n_rep: int = 1,
random_state: int = 0,
)
Bases: _QuantileBase
Local potential quantile of the compliers, with a binary instrument z, as DoubleML's
DoubleMLLPQ. Default sample splits are stratified by treatment and instrument. Other
parameters as PQ.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
ml_g
|
classifier for ``P(Y <= q, D = d | Z = z, X)``.
|
|
required |
ml_m
|
classifier for the instrument propensity ``P(Z = 1 | X)``.
|
|
required |
fit(
X: ArrayLike,
y: ArrayLike,
d: ArrayLike,
z: ArrayLike,
folds: Optional[Folds] = None,
external_predictions: Any = None,
) -> Self
external_predictions must cover all five learners (ml_m_z, ml_m_d_z0,
ml_m_d_z1, ml_g_du_z0, ml_g_du_z1) or none, as DoubleML.
CVAR(
ml_g: Any,
ml_m: Any,
treatment: int = 1,
quantile: float = 0.5,
clipping_threshold: float = 0.01,
calibration: Optional[
Literal["isotonic", "isotonic_cv"]
] = None,
normalize_ipw: bool = True,
n_folds: int = 5,
n_rep: int = 1,
random_state: int = 0,
)
Bases: _QuantileBase
Conditional value at risk E[Y(d) | Y(d) >= q] above the potential quantile q, as
DoubleML's DoubleMLCVAR. Other parameters as PQ.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
ml_g
|
regressor for ``E[max(q, (Y - quantile * q) / (1 - quantile)) | D = d, X]``.
|
|
required |
ml_m
|
classifier for the propensity ``P(D = 1 | X)``.
|
|
required |
SSM(
ml_g: Any,
ml_pi: Any,
ml_m: Any,
score: Literal[
"missing-at-random", "nonignorable"
] = "missing-at-random",
clipping_threshold: float = 0.01,
calibration: Optional[
Literal["isotonic", "isotonic_cv"]
] = None,
normalize_ipw: bool = False,
n_folds: int = 5,
n_rep: int = 1,
random_state: int = 0,
)
Bases: _DMLEstimator
Average treatment effect of a binary treatment under sample selection (attrition or
non-response), as DoubleML's DoubleMLSSM (Bia, Huber and Lafférs, 2024). The outcome is
observed only where the selection indicator s is 1; other values of y are ignored but
must be finite.
score="missing-at-random": selection depends on the treatment and the covariates only.
score="nonignorable": selection may also depend on unobservables; needs an instrument z
that shifts selection but not the outcome. Each training fold is then split in halves: the
selection probability is learned on the first, the propensity and outcome models (with the
predicted selection probability as a covariate) on the second.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
ml_g
|
regressor for ``E[Y | D, X]`` among the selected.
|
|
required |
ml_pi
|
classifier for the selection probability ``P(S = 1 | D, X)`` (and ``Z``).
|
|
required |
ml_m
|
classifier for the propensity ``P(D = 1 | X)``.
|
|
required |
score
|
``"missing-at-random"`` or ``"nonignorable"``.
|
|
'missing-at-random'
|
clipping_threshold
|
propensities are clipped to ``[t, 1 - t]`` (selection probabilities are not,
|
as in DoubleML). |
0.01
|
normalize_ipw
|
normalise the inverse-probability weights.
|
|
False
|
fit(
X: ArrayLike,
y: ArrayLike,
d: ArrayLike,
s: ArrayLike,
z: Optional[ArrayLike] = None,
folds: Optional[Folds] = None,
cluster: Optional[ArrayLike] = None,
store_models: bool = False,
) -> Self
Fit on covariates X, outcome y, binary treatment d, selection indicator s
and (nonignorable score) instrument z.
store_models=True keeps the fitted learners in models_ ({learner: {treatment:
[[model per fold] per repetition]}}, as DoubleML's models); folds are then fitted one
after the other.
PLIV(
ml_l: Any,
ml_m: Any,
ml_r: Any,
ml_g: Any = None,
score: Union[
Literal["partialling out", "IV-type"],
Callable[..., Any],
] = "partialling out",
partial: Literal["X", "Z", "XZ"] = "X",
n_folds: int = 5,
n_rep: int = 1,
random_state: int = 0,
)
Bases: _DMLEstimator
Partially linear IV regression with one or several instruments z (one column each), as
DoubleML's DoubleMLPLIV.
With partial="X" (default) X is partialled out of Y, D and every instrument;
several instruments are combined by projecting the treatment residual on the instrument
residuals (two-stage least squares). partial="Z" uses E[D | X, Z] (ml_r) as
instrument, and partial="XZ" partials out both (DoubleML's _partialZ / _partialXZ).
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
ml_l
|
regressor for ``E[Y | X]`` (not used with ``partial="Z"``).
|
|
required |
ml_m
|
regressor for ``E[Z | X]``; with ``partial="XZ"`` for ``E[D | X, Z]`` (not used with ``"Z"``).
|
|
required |
ml_r
|
regressor for ``E[D | X]``; with ``partial="Z"`` for ``E[D | X, Z]``; with ``"XZ"`` for
|
the projection of |
required |
ml_g
|
regressor for ``E[Y - θD | X]``, used by ``score="IV-type"`` (default: a copy of ``ml_l``).
|
|
None
|
score
|
``"partialling out"``, ``"IV-type"`` (one instrument, ``partial="X"``), or a callable
|
|
'partialling out'
|
partial
|
``"X"``, ``"Z"`` or ``"XZ"``.
|
|
'X'
|
fit(
X: ArrayLike,
y: ArrayLike,
d: ArrayLike,
z: ArrayLike,
folds: Optional[Folds] = None,
cluster: Optional[ArrayLike] = None,
external_predictions: Optional[dict[Any, Any]] = None,
store_models: bool = False,
) -> Self
Fit on covariates X, outcome y, treatment(s) d and instrument(s) z.
store_models=True keeps the fitted learners in models_ ({learner: {treatment:
[[model per fold] per repetition]}}, as DoubleML's models); folds are then fitted one
after the other.
IIVM(
ml_g: Any,
ml_m: Any,
ml_r: Any,
subgroups: Optional[dict[str, bool]] = None,
clipping_threshold: float = 0.01,
calibration: Optional[
Literal["isotonic", "isotonic_cv"]
] = None,
normalize_ipw: bool = False,
n_folds: int = 5,
n_rep: int = 1,
random_state: int = 0,
score: Union[
Literal["LATE"], Callable[..., Any]
] = "LATE",
)
Bases: _DMLEstimator
Interactive IV model: local average treatment effect (LATE) of a binary
treatment d using a binary instrument z.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
ml_g
|
regressor for ``E[Y | Z, X]`` (fitted separately for ``Z = 0`` and ``Z = 1``).
|
|
required |
ml_m
|
classifier for the instrument propensity ``P(Z = 1 | X)``.
|
|
required |
ml_r
|
classifier for ``P(D = 1 | Z, X)`` (fitted separately for ``Z = 0`` and ``Z = 1``).
|
|
required |
subgroups
|
dict with ``always_takers`` / ``never_takers`` (default both True). Set one
|
to False under one-sided non-compliance. |
None
|
clipping_threshold
|
instrument propensities are clipped to ``[t, 1 - t]``.
|
|
0.01
|
normalize_ipw
|
normalise the inverse-probability weights within instrument arms.
|
|
False
|
score
|
``"LATE"`` or a callable ``score(y, z, d, g_hat0, g_hat1, m_hat, r_hat0, r_hat1, smpls)
|
-> (psi_a, psi_b)`` returning the elements of a linear score, as DoubleML. |
'LATE'
|
fit(
X: ArrayLike,
y: ArrayLike,
d: ArrayLike,
z: ArrayLike,
folds: Optional[Folds] = None,
cluster: Optional[ArrayLike] = None,
external_predictions: Optional[dict[Any, Any]] = None,
store_models: bool = False,
) -> Self
Fit on covariates X, outcome y, binary treatment d and binary instrument z.
store_models=True keeps the fitted learners in models_ ({learner: {treatment:
[[model per fold] per repetition]}}, as DoubleML's models); folds are then fitted one
after the other.
DID(
ml_g: Any,
ml_m: Any = None,
score: Literal[
"observational", "experimental"
] = "observational",
in_sample_normalization: bool = True,
clipping_threshold: float = 0.01,
calibration: Optional[
Literal["isotonic", "isotonic_cv"]
] = None,
n_folds: int = 5,
n_rep: int = 1,
random_state: int = 0,
)
Bases: _DIDBase
Doubly robust difference-in-differences (ATT) with panel data over two periods.
y passed to fit is each unit's outcome change Y_post - Y_pre.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
ml_g
|
regressor for ``E[ΔY | D = d, X]`` (fitted per treatment group).
|
|
required |
ml_m
|
classifier for the propensity ``P(D = 1 | X)``; required for ``score="observational"``.
|
|
None
|
score
|
``"observational"`` (doubly robust) or ``"experimental"`` (randomised treatment).
|
|
'observational'
|
in_sample_normalization
|
normalise the weights by their sample means (DoubleML default).
|
|
True
|
fit(
X: ArrayLike,
y: ArrayLike,
d: ArrayLike,
folds: Optional[Folds] = None,
cluster: Optional[ArrayLike] = None,
external_predictions: Optional[dict[Any, Any]] = None,
store_models: bool = False,
) -> Self
Fit on covariates X, outcome change y and binary treatment d.
store_models=True keeps the fitted learners in models_ ({learner: {treatment:
[[model per fold] per repetition]}}, as DoubleML's models); folds are then fitted one
after the other.
DIDCS(
ml_g: Any,
ml_m: Any = None,
score: Literal[
"observational", "experimental"
] = "observational",
in_sample_normalization: bool = True,
clipping_threshold: float = 0.01,
calibration: Optional[
Literal["isotonic", "isotonic_cv"]
] = None,
n_folds: int = 5,
n_rep: int = 1,
random_state: int = 0,
)
Bases: _DIDBase
Doubly robust difference-in-differences (ATT) with repeated cross-sections.
Each row is observed once; t passed to fit is 1 in the post-period, 0 before.
Parameters as DID, with ml_g fitted in each treatment-group x period cell.
fit(
X: ArrayLike,
y: ArrayLike,
d: ArrayLike,
t: ArrayLike,
folds: Optional[Folds] = None,
cluster: Optional[ArrayLike] = None,
external_predictions: Optional[dict[Any, Any]] = None,
store_models: bool = False,
) -> Self
Fit on covariates X, outcome y, binary treatment d and period indicator t.
store_models=True keeps the fitted learners in models_ ({learner: {treatment:
[[model per fold] per repetition]}}, as DoubleML's models); folds are then fitted one
after the other.
StaggeredDID(
ml_g: Any,
ml_m: Any = None,
control_group: Literal[
"never_treated", "not_yet_treated"
] = "never_treated",
anticipation_periods: int = 0,
gt_combinations: Literal[
"standard", "all"
] = "standard",
score: Literal[
"observational", "experimental"
] = "observational",
in_sample_normalization: bool = True,
clipping_threshold: float = 0.01,
calibration: Optional[
Literal["isotonic", "isotonic_cv"]
] = None,
n_folds: int = 5,
n_rep: int = 1,
random_state: int = 0,
panel: bool = True,
)
Bases: _DMLEstimator
Difference-in-differences with staggered adoption (Callaway & Sant'Anna, 2021),
as DoubleML's DoubleMLDIDMulti for panel data.
Estimates every group-time effect ATT(g, t) by a doubly robust two-period DID and
aggregates them with aggregate.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
ml_g
|
regressor for the control group's outcome change given ``X``.
|
|
required |
ml_m
|
classifier for the cohort propensity; required for ``score="observational"``.
|
|
None
|
control_group
|
``"never_treated"`` or ``"not_yet_treated"``.
|
|
'never_treated'
|
anticipation_periods
|
periods before treatment in which units may already react.
|
|
0
|
gt_combinations
|
``"standard"`` (varying base), ``"all"`` or ``"universal"`` (fixed base).
|
|
'standard'
|
score
|
as `DID`.
|
|
'observational'
|
in_sample_normalization
|
as `DID`.
|
|
'observational'
|
clipping_threshold
|
as `DID`.
|
|
'observational'
|
panel
|
``True`` for panel data (rows linked by ``id``); ``False`` for repeated
|
cross-sections, where every row is a separate observation. |
True
|
fit(
X: ArrayLike,
y: ArrayLike,
id: ArrayLike,
t: ArrayLike,
g: ArrayLike,
cluster: Optional[ArrayLike] = None,
) -> Self
Fit on long-format data: one row per observation with period t and cohort
g (first treated period; np.inf or NaN for never treated). For panels, rows
are linked by id and X must be time-invariant; for repeated cross-sections
(panel=False) id is ignored.
cluster (one label per row, of any type; constant within a panel unit) makes the
standard errors, aggregations and sensitivity bounds cluster-robust, and each ATT(g, t)
assigns whole clusters to folds. DoubleML's DoubleMLDIDMulti has no clustered
version; the multiplier bootstrap (joint bands) is not available with clusters.
fit_frame(
df: DataFrame,
y_col: str,
id_col: str,
t_col: str,
g_col: str,
x_cols: Sequence[str],
cluster_col: Optional[str] = None,
) -> Self
Convenience wrapper for a pandas DataFrame in long format.
plot_effects(
level: float = 0.95,
result_type: Literal[
"effect", "rv", "est_bounds", "ci_bounds"
] = "effect",
joint: bool = True,
figsize: tuple[float, float] = (12, 8),
y_label: Optional[str] = None,
title: Optional[str] = None,
) -> tuple[Any, list[Any]]
The ATT(g, t) over evaluation periods, one panel per cohort, as DoubleML's
plot_effects: pre-treatment, anticipation and post-treatment periods in different
colours, with confidence intervals (simultaneous with joint; runs bootstrap if
needed), or the sensitivity bounds / robustness values of the last
sensitivity_analysis. Needs matplotlib. Returns (fig, axes).
aggregate(
aggregation: Literal[
"group", "time", "eventstudy"
] = "group",
) -> DIDAggregation
Aggregate the ATT(g, t): "group", "time" or "eventstudy".
DIDAggregation(
how: str,
effects: Any,
overall: Any,
weights: Any,
overall_weights: Any,
)
Aggregated staggered-DID effects (by cohort, calendar time or event time).
property
¶
aggregated_frameworks: Framework
The aggregated effects as a Framework (summary, bootstrap,
sensitivity_analysis, ...), as DoubleML's aggregated_frameworks.
property
¶
overall_aggregated_framework: Framework
The overall effect as a Framework, as DoubleML's overall_aggregated_framework.
class-attribute
instance-attribute
¶
names = property(lambda self: self._effects.treatment_names)
class-attribute
instance-attribute
¶
overall_coef = property(
lambda self: float(self._overall.coef[0])
)
class-attribute
instance-attribute
¶
overall_se = property(
lambda self: float(self._overall.se[0])
)
confint(
level: float = 0.95,
joint: bool = False,
n_boot: int = 500,
random_state: int = 0,
) -> NDArray[np.float64]
Intervals for the aggregated effects; joint=True gives simultaneous bands.
plot_effects(
level: float = 0.95,
joint: bool = True,
figsize: tuple[float, float] = (12, 6),
sort_by: Optional[Literal["estimate", "name"]] = None,
title: str = "Aggregated Treatment Effects",
y_label: str = "Effect",
) -> tuple[Any, Any]
Aggregated effects with confidence intervals (simultaneous with joint), as DoubleML's
plot_effects; event-study effects before treatment are drawn in a second colour.
Needs matplotlib. Returns (fig, ax).
RDFlex(
ml_g: Any,
ml_m: Any = None,
fuzzy: bool = False,
cutoff: float = 0.0,
h_fs: Optional[float] = None,
fs_specification: Literal[
"cutoff",
"cutoff and score",
"interacted cutoff and score",
] = "cutoff",
fs_kernel: Literal[
"triangular", "uniform", "epanechnikov"
] = "triangular",
n_iterations: int = 2,
n_folds: int = 5,
n_rep: int = 1,
random_state: int = 0,
**rdrobust_kwargs: Any,
)
Bases: _Params
Regression discontinuity design with flexible covariate adjustment, as DoubleML's RDFlex
(Noack, Olma and Rothe, 2024).
The outcome (and in a fuzzy design the treatment) is adjusted by a cross-fitted prediction from
the covariates at the cutoff, learned with kernel weights around it; a local linear RD estimator
(a port of rdrobust: bias correction, robust intervals) is then run on the adjusted outcome.
The first-stage bandwidth is updated from the RD bandwidth over n_iterations rounds. Fitted
attributes hold three rows: conventional, bias-corrected and robust estimates.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
ml_g
|
regressor for the outcome; must accept ``sample_weight``.
|
|
required |
ml_m
|
classifier for the treatment (fuzzy design only); must accept ``sample_weight``.
|
|
None
|
fuzzy
|
fuzzy design (the treatment may differ from ``score >= cutoff``).
|
|
False
|
cutoff
|
the cutoff in the score.
|
|
0.0
|
h_fs
|
initial first-stage bandwidth (default: ``rdbwselect`` without covariates).
|
|
None
|
fs_specification
|
``"cutoff"``, ``"cutoff and score"`` or ``"interacted cutoff and score"``.
|
|
'cutoff'
|
fs_kernel
|
``"triangular"``, ``"uniform"`` or ``"epanechnikov"``.
|
|
'triangular'
|
n_iterations
|
rounds of first-stage fitting and bandwidth updates.
|
|
2
|
**rdrobust_kwargs
|
options of the RD estimator, as DoubleML passes them to ``rdrobust``:
|
|
{}
|
property
¶
summary: DataFrame
Coefficient table (rows: conventional, bias-corrected, robust).
fit(
X: ArrayLike,
y: ArrayLike,
d: ArrayLike,
score: ArrayLike,
folds: Optional[Folds] = None,
) -> Self
Fit on covariates X, outcome y, treatment d and running variable score. A
classifier ml_g (with predict_proba) learns a binary outcome, as in DoubleML.
confint(level: float = 0.95) -> NDArray[np.float64]
Intervals at level (rows: conventional, bias-corrected, robust), as DoubleML's confint.
AutoDML(
ml_g: Any,
functional: Union[str, Functional] = "ate",
riesz: Optional[RieszLasso] = None,
n_folds: int = 5,
n_rep: int = 1,
random_state: int = 0,
)
Bases: _HeterogeneityMixin, _DMLEstimator
Automatic debiased machine learning (Chernozhukov, Newey and Singh, 2022) of
θ = E[m(W; γ)], a mean of a linear functional m of the regression γ(D, X) = E[Y | D, X]:
the average treatment effect, average derivatives, effects of shifting the treatment or of a
policy, or any functional you define.
The score m(W; γ) + α(W)(Y - γ(W)) - θ needs the Riesz representer α of the
functional (for the ATE the inverse propensity weights); here it is learned from the functional
itself, so no formula has to be derived. Inference, repeated sample splitting, clustering,
sensitivity analysis and cate / gate work as for the other models.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
ml_g
|
regressor for ``E[Y | D, X]``, fitted on ``W = (D, X)``.
|
|
required |
functional
|
``"ate"``, ``"average_derivative"`` or a `Functional` (`ATE`,
|
|
'ate'
|
riesz
|
learner of the Riesz representer (default `RieszLasso()`).
|
|
None
|
n_folds
|
sample splitting.
|
|
5
|
n_rep
|
sample splitting.
|
|
5
|
random_state
|
sample splitting.
|
|
5
|
fit(
X: ArrayLike,
y: ArrayLike,
d: ArrayLike,
folds: Optional[Folds] = None,
cluster: Optional[ArrayLike] = None,
store_models: bool = False,
) -> Self
Fit on covariates X, outcome y and treatment(s) d (each treatment gets the
functional; other treatments act as covariates). predictions_ holds ml_g (the
regression at the data), riesz (the Riesz representer) and functional
(m(W; γ̂)).
Functionals and Riesz representers¶
Bases: Functional
Average treatment effect of a binary treatment: m = f(1, X) - f(0, X).
AverageDerivative(step: Optional[float] = None)
Bases: Functional
Average derivative E[dγ(D, X)/dD] of a continuous treatment, by a central difference
with step step (default 1e-4 times the standard deviation of the treatment). Needs a
regression that is smooth in the treatment (linear models, polynomials, neural networks); for
trees and forests use ShiftEffect.
ShiftEffect(delta: float, per_unit: bool = False)
PolicyEffect(
policy: Union[ArrayLike, Callable[[Any], Any]],
)
Bases: Functional
Average effect of a policy that assigns each unit a treatment: m = f(d_new, X) - f(D, X).
policy is the array of new treatments (one per observation) or a function policy(W)
of the data matrix W = (D, X) returning them.
CustomFunctional(
terms: Callable[[Any], Sequence[tuple[Any, Any]]],
name: str = "custom functional",
)
Bases: Functional
Any linear functional, given by terms(W): a list of (W_k, weight_k) pairs with
m(W; f) = sum_k weight_k * f(W_k). W is the data matrix (treatment in column 0,
covariates after it); each W_k has its shape, and a weight is a number or one value per
row. For example the ATE is [(W with D = 1, 1.0), (W with D = 0, -1.0)].
RieszLasso(
degree: int = 2,
treatment_degree: int = 1,
penalty: Optional[float] = None,
n_alphas: int = 30,
eps: float = 0.001,
cv: int = 5,
dictionary: Optional[Callable[[Any], Any]] = None,
)
The Lasso Riesz representer of Chernozhukov, Newey and Singh (2022): α(w) = b(w)'ρ for a
dictionary b, minimising the sample Riesz loss ρ'Gρ - 2M'ρ plus a weighted l1 penalty.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
degree
|
the dictionary holds all monomials of the covariates up to this total degree ...
|
|
2
|
treatment_degree
|
... each times the powers ``1, d, ..., d^treatment_degree`` of the treatment.
|
|
1
|
penalty
|
fixed penalty ``r`` (0: unpenalised); ``None`` chooses it by ``cv``-fold
|
cross-validation of the Riesz loss over |
None
|
dictionary
|
a function ``dictionary(W)`` returning the basis matrix (rows of ``W``: treatment
|
first, then covariates), replacing the polynomial dictionary. |
None
|
Tuning and learner parameters¶
PLR, LPLR, PLPR, IRM, APO, SSM, PLIV, IIVM, DID and DIDCS share these methods, and their fit
takes store_models=True to keep the fitted learners in models_.
set_ml_nuisance_params, tune, tune_ml_models and models_ for the models with
named nuisance learners (DoubleML's params_names).
property
¶
params_names: list[str]
Names of the nuisance models (DoubleML's params_names).
property
¶
learner_names: list[str]
Names of the learner arguments the nuisance models use (DoubleML's learner_names).
property
¶
params: dict[
str, dict[Optional[str], list[list[dict[str, Any]]]]
]
Parameters set per nuisance model and treatment: {name: {treatment: [[params per fold]
per repetition]}} (treatment None: every treatment).
set_ml_nuisance_params(
learner: str, treat_var: Optional[str], params: Any
) -> Any
Set hyperparameters of one nuisance model, as DoubleML.
learner is a name from params_names (e.g. "ml_g1"), treat_var a treatment name
(None: every treatment), and params a dict used in every fold, or a nested list with
one dict per fold for each repetition (n_rep x n_folds). New values update earlier
ones. Learners need the scikit-learn API (set_params). Fold-specific parameters need a
model that fits the learner once per fold; folds are then fitted one after the other.
tune(
*args: Any,
param_grids: dict[str, Any],
tune_on_folds: bool = False,
scoring_methods: Optional[dict[str, Any]] = None,
n_folds_tune: int = 5,
search_mode: str = "grid_search",
n_iter_randomized_search: int = 100,
n_jobs_cv: Optional[int] = None,
set_as_params: bool = True,
return_tune_res: bool = False,
**kwargs: Any,
) -> Any
Tune the nuisance learners by grid or randomized search, as DoubleML's tune.
Takes the arguments of fit (the data, and optionally folds) plus DoubleML's tuning
arguments: param_grids maps every name in learner_names (e.g. ml_g, ml_m) to a
scikit-learn parameter grid; scoring_methods optionally maps learner names to scorers.
Each learner is tuned on the training data of its nuisance model (e.g. ml_g for
ml_g0 on the untreated) with scikit-learn's GridSearchCV (or RandomizedSearchCV)
on n_folds_tune shuffled folds drawn from NumPy's global random state, as DoubleML.
With tune_on_folds=True the search runs on every cross-fitting training fold (the
folds fit draws, or folds) and the parameters are fold-specific.
The best parameters are set with set_ml_nuisance_params (set_as_params); returns
self, or the search results per treatment with return_tune_res=True.
tune_ml_models(
*args: Any,
ml_param_space: dict[
str, Callable[[Any], dict[str, Any]]
],
scoring_methods: Optional[dict[str, Any]] = None,
cv: Any = 5,
set_as_params: bool = True,
return_tune_res: bool = False,
optuna_settings: Optional[dict[str, Any]] = None,
**kwargs: Any,
) -> Any
Tune the nuisance learners with Optuna, as DoubleML's tune_ml_models (needs optuna).
Takes the arguments of fit plus DoubleML's tuning arguments. ml_param_space maps
learner names (learner_names, applied to every nuisance model of that learner) or
nuisance names (params_names, taking precedence) to functions trial -> dict of
parameters. Each nuisance learner is tuned once on its training data in the whole sample
by cross-validation (cv: folds, a splitter or (train, test) pairs; an int means
shuffled K-fold with random_state=42), maximising scoring_methods[name] (default:
negative RMSE for regressors, negative log loss for classifiers). optuna_settings as
DoubleML (n_trials default 100, sampler, study_kwargs, optimize_kwargs, ...;
also per learner or nuisance name). Returns self, or per treatment {name:
OptunaResult} with return_tune_res=True.
dataclass
¶
OptunaResult(
learner_name: str,
params_name: str,
best_estimator: Any,
best_params: dict[str, Any],
best_score: float,
scoring_method: Any,
study: Any,
tuned: bool,
)
Result of tuning one nuisance learner with Optuna (DoubleML's DMLOptunaResult).
Data¶
DmlData(
data: DataFrame,
y_col: str,
d_cols: Columns,
x_cols: Columns = None,
z_cols: Columns = None,
cluster_cols: Columns = None,
t_col: Optional[str] = None,
s_col: Optional[str] = None,
id_col: Optional[str] = None,
score_col: Optional[str] = None,
)
Columns of a DataFrame in their roles, as DoubleML's DoubleMLData.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
data
|
DataFrame
|
The data. |
required |
y_col
|
str
|
Outcome column. |
required |
d_cols
|
str or list of str
|
Treatment column(s); for |
required |
x_cols
|
str or list of str
|
Covariates; by default all columns not used in another role. |
None
|
z_cols
|
str or list of str
|
Instrument column(s). |
None
|
cluster_cols
|
str or list of str
|
One or two cluster columns. |
None
|
t_col
|
str
|
Time: the period of a panel or staggered DiD, or the 0/1 post-period indicator of
repeated cross-sections ( |
None
|
s_col
|
str
|
Selection indicator ( |
None
|
id_col
|
str
|
Unit identifier of a panel ( |
None
|
score_col
|
str
|
Running variable of a regression discontinuity design ( |
None
|
classmethod
¶
from_arrays(
x: ArrayLike,
y: ArrayLike,
d: ArrayLike,
z: Optional[ArrayLike] = None,
cluster_vars: Optional[ArrayLike] = None,
t: Optional[ArrayLike] = None,
s: Optional[ArrayLike] = None,
id: Optional[ArrayLike] = None,
score: Optional[ArrayLike] = None,
) -> DmlData
Build from arrays with DoubleML's default column names (X1, ..., y, d or
d1, ..., z or z1, ..., cluster_var or cluster_var1, cluster_var2,
t, s, id, score).
Simulated data sets: the data generating processes of DoubleML's make_* functions.
The functions keep DoubleML's names, parameters and defaults. The data are drawn in Rust from a
seeded generator: pass random_state for reproducible data (without it, the seed is drawn from
NumPy's global generator, so np.random.seed makes the data reproducible as well). The draws
are not identical to DoubleML's, but the distributions are the same.
return_type is "array" (a tuple of arrays, the default), "DataFrame" (pandas, with
DoubleML's column names) or "DmlData" (a DmlData with the columns in their roles, as
DoubleML's data classes). Generators that DoubleML returns as a dict or a DataFrame return the
same here.
make_plr_CCDDHNR2018(
n_obs: int = 500,
dim_x: int = 20,
alpha: float = 0.5,
return_type: Literal[
"array", "DataFrame", "DmlData"
] = "array",
random_state: Optional[int] = None,
**kwargs: float,
)
Partially linear regression of Chernozhukov et al. (2018):
d = a_0 x_0 + a_1 expit(x_2) + s_1 e, y = alpha d + b_0 expit(x_0) + b_1 x_2 + s_2 u
with Toeplitz-correlated covariates. kwargs: a_0, a_1, s_1, b_0, b_1,
s_2. Returns (x, y, d).
make_plr_turrell2018(
n_obs: int = 100,
dim_x: int = 20,
theta: float = 0.5,
return_type: Literal[
"array", "DataFrame", "DmlData"
] = "array",
random_state: Optional[int] = None,
**kwargs: float,
)
Partially linear regression of Turrell (2018) with a random covariance matrix.
kwargs: nu, gamma. Returns (x, y, d).
make_confounded_plr_data(
n_obs: int = 500,
theta: float = 5.0,
cf_y: float = 0.04,
cf_d: float = 0.04,
random_state: Optional[int] = None,
**kwargs: Any,
) -> dict[str, Any]
Partially linear model with an unobserved confounder of strength (cf_y, cf_d), for
sensitivity analysis. kwargs: c, dim_x. Returns {"x", "y", "d",
"oracle_values"}.
make_pliv_CHS2015(
n_obs: int,
alpha: float = 1.0,
dim_x: int = 200,
dim_z: int = 150,
return_type: Literal[
"array", "DataFrame", "DmlData"
] = "array",
random_state: Optional[int] = None,
)
Partially linear IV of Chernozhukov, Hansen and Spindler (2015) with dim_z
instruments. Returns (x, y, d, z).
make_pliv_multiway_cluster_CKMS2021(
N: int = 25,
M: int = 25,
dim_X: int = 100,
theta: float = 1.0,
return_type: Literal[
"array", "DataFrame", "DmlData"
] = "array",
random_state: Optional[int] = None,
**kwargs: Any,
)
Partially linear IV with two-way clustering (Chiang, Kato, Ma and Sasaki, 2021) on an
N x M grid. Returns (x, y, d, cluster_vars, z).
make_lplr_LZZ2020(
n_obs: int = 500,
dim_x: int = 20,
alpha: float = 0.5,
return_type: Literal[
"array", "DataFrame", "DmlData"
] = "array",
balanced_r0: bool = True,
treatment: Literal[
"continuous", "binary", "binary_unbalanced"
] = "continuous",
random_state: Optional[int] = None,
)
Logistic partially linear model of Liu, Zhang and Zhou (2020): a binary outcome with
P(y = 1) = expit(alpha d + r_0(x)). Returns (x, y, d, p) with p the outcome
probabilities.
make_plpr_CP2025(
num_id: int = 250,
num_t: int = 10,
dim_x: int = 30,
theta: float = 0.5,
dgp_type: Literal["dgp1", "dgp2", "dgp3"] = "dgp1",
time_type: Literal["int", "float", "datetime"] = "int",
random_state: Optional[int] = None,
)
Static panel of Clarke and Polselli (2025) with correlated unit effects, as a long-format
DataFrame with id, time, y, d and x1, x2, ....
make_irm_data(
n_obs: int = 500,
dim_x: int = 20,
theta: float = 0.0,
R2_d: float = 0.5,
R2_y: float = 0.5,
return_type: Literal[
"array", "DataFrame", "DmlData"
] = "array",
random_state: Optional[int] = None,
)
Interactive regression with a binary treatment (Belloni et al., 2017), with the strength of
the propensity and outcome models set by R2_d and R2_y. Returns (x, y, d).
make_confounded_irm_data(
n_obs: int = 500,
theta: float = 0.0,
gamma_a: float = 0.127,
beta_a: float = 0.58,
linear: bool = False,
random_state: Optional[int] = None,
**kwargs: Any,
) -> dict[str, Any]
Binary treatment with an unobserved confounder, with oracle confounding strengths for
sensitivity analysis. kwargs: dim_x, trimming_threshold, var_eps_y. Returns
{"x", "y", "d", "oracle_values"}.
make_irm_data_discrete_treatments(
n_obs: int = 200,
n_levels: int = 3,
linear: bool = False,
random_state: Optional[int] = None,
**kwargs: Any,
) -> dict[str, Any]
Multi-valued treatment d in {0, 1, ..., n_levels}. kwargs: xi, c,
dim_x. Returns {"x", "y", "d", "oracle_values"}.
make_iivm_data(
n_obs: int = 500,
dim_x: int = 20,
theta: float = 1.0,
alpha_x: float = 0.2,
return_type: Literal[
"array", "DataFrame", "DmlData"
] = "array",
random_state: Optional[int] = None,
)
Binary treatment shifted by a binary instrument (LATE). Returns (x, y, d, z).
make_ssm_data(
n_obs: int = 8000,
dim_x: int = 100,
theta: float = 1.0,
mar: bool = True,
return_type: Literal[
"array", "DataFrame", "DmlData"
] = "array",
random_state: Optional[int] = None,
)
Sample selection: the outcome is observed where s = 1 (0 elsewhere); missing at random
(mar) or nonignorable with an instrument z for selection. Returns
(x, y, d, z, s).
make_heterogeneous_data(
n_obs: int = 200,
p: int = 30,
support_size: int = 5,
n_x: int = 1,
binary_treatment: bool = False,
random_state: Optional[int] = None,
) -> dict[str, Any]
Heterogeneous effects exp(2 x_0) + 3 sin(4 x_0) (n_x=2: 3 sin(4 x_1)) on
uniform covariates. Returns {"data", "effects", "treatment_effect"}: the data as a
DataFrame (columns y, d, X_0, ...; a dict without pandas), the individual effects
and the effect function.
make_did_SZ2020(
n_obs: int = 500,
dgp_type: int = 1,
cross_sectional_data: bool = False,
return_type: Literal[
"array", "DataFrame", "DmlData"
] = "array",
random_state: Optional[int] = None,
**kwargs: float,
)
Two-period difference-in-differences of Sant'Anna and Zhao (2020), dgp_type 1 to 6.
Panel data give the outcome change; cross_sectional_data=True repeated cross-sections
with a time indicator. kwargs: xi, c, lambda_t. Returns (z, y, d, t)
(t is None for panels).
make_did_CS2021(
n_obs: int = 1000,
dgp_type: int = 1,
include_never_treated: bool = True,
time_type: Literal["float", "datetime"] = "float",
random_state: Optional[int] = None,
**kwargs: Any,
)
Staggered adoption of Callaway and Sant'Anna (2021): a long-format panel DataFrame with
id, y, y0, y1, d (first treated period; inf or NaT never),
t and Z1–Z4. kwargs: c, xi, n_periods, anticipation_periods,
n_pre_treat_periods, start_date (for time_type="datetime").
make_did_cs_CS2021(
n_obs: int = 1000,
dgp_type: int = 1,
include_never_treated: bool = True,
lambda_t: float = 0.5,
time_type: Literal["float", "datetime"] = "float",
random_state: Optional[int] = None,
**kwargs: Any,
)
Repeated cross-sections of the Callaway and Sant'Anna (2021) design: rows of a panel of
ceil(n_obs / lambda_t) units, each kept with probability lambda_t. Columns as
make_did_CS2021.
make_simple_rdd_data(
n_obs: int = 5000,
p: int = 4,
fuzzy: bool = True,
binary_outcome: bool = False,
random_state: Optional[int] = None,
**kwargs: float,
) -> dict[str, Any]
Regression discontinuity with a normal score and polynomial covariate effects. kwargs:
cutoff, dim_x, a, tau. Returns {"score", "Y", "D", "X",
"oracle_values": {"Y0", "Y1"}}.
fetch_401K(
return_type: Literal[
"DmlData", "DataFrame"
] = "DmlData",
polynomial_features: bool = False,
)
The 401(k) data of Chernozhukov and Hansen (2004), as DoubleML's fetch_401K: net
financial assets (net_tfa) of 9,915 households, eligibility for a 401(k) plan (e401)
and nine covariates. Downloaded from the DoubleML authors' repository (needs internet
access). Returns a DmlData (outcome net_tfa, treatment e401) or the DataFrame.
fetch_bonus(
return_type: Literal[
"DmlData", "DataFrame"
] = "DmlData",
polynomial_features: bool = False,
)
The Pennsylvania reemployment bonus experiment (Bilias, 2000), as DoubleML's
fetch_bonus: log unemployment duration (inuidur1) of the control group and treatment
group 4 (tg), with 15 covariates (degree-2 polynomial features with
polynomial_features). Downloaded from the DoubleML authors' repository (needs internet
access). Returns a DmlData or the DataFrame.
Combining estimates¶
Framework(result: Any)
Bases: _ResultView
Estimates with their scores, for inference on combinations, as DoubleML's DoubleMLFramework.
a + b and a - b combine the estimates of two fits on the same observations (and, for
meaningful inference, the same sample splits) treatment by treatment; c * a scales them.
The standard errors use the combined scores, so the correlation between the estimates is taken
into account; the sensitivity bounds add up with absolute weights. linear_combination gives
any weighted sums and concat stacks frameworks for joint inference. Offers the fitted-model
API (summary, confint, bootstrap, p_adjust, sensitivity_analysis).
linear_combination(
weights: ArrayLike,
names: Optional[Sequence[str]] = None,
) -> Framework
Weighted sums of the estimates, one per row of weights (shape
(n_combinations, n_coefs); a vector gives one combination), with joint inference.
concat(objs: Sequence[Any]) -> Framework
Stack fitted estimators or frameworks on the same observations and sample splits into one
Framework with all their estimates, for joint inference (DoubleML's concat).
Heterogeneous effects¶
BLP(native: Any, names: list[str])
Best linear predictor of the CATE (or GATEs) with HC0 inference.
Mirrors DoubleML's DoubleMLBLP: confint() gives coefficient (or, for
GATEs, group) intervals; confint(basis) gives intervals for the CATE at
the rows of a new basis, optionally simultaneous with joint=True.
predict(
basis: ArrayLike,
) -> tuple[NDArray[np.float64], NDArray[np.float64]]
(effect, se) at the rows of basis.
confint(
basis: Optional[ArrayLike] = None,
joint: bool = False,
level: float = 0.95,
n_boot: int = 500,
random_state: int = 0,
) -> pd.DataFrame
Intervals with columns [lower, effect, upper].
SLearner(model: Any)
Bases: _MetaLearner
One model μ(x, d) on [X, d]; τ̂(x) = μ(x, 1) − μ(x, 0).
TLearner(model: Any)
Bases: _MetaLearner
Separate outcome models for treated and controls; τ̂ = μ₁ − μ₀.
XLearner(
outcome_model: Any,
effect_model: Any,
propensity_model: Any,
)
RLearner(
outcome_model: Any,
propensity_model: Any,
effect_model: Any,
n_folds: int = 5,
clipping_threshold: float = 0.01,
random_state: int = 0,
)
Bases: _CrossFitted
R-learner (Nie & Wager, 2021). effect_model must accept sample_weight.
DRLearner(
outcome_model: Any,
propensity_model: Any,
effect_model: Any,
n_folds: int = 5,
clipping_threshold: float = 0.01,
random_state: int = 0,
)
Bases: _CrossFitted
DR-learner (Kennedy, 2023): regress cross-fitted AIPW pseudo-outcomes on X.
Supports inference on effect heterogeneity via blp and gate.
PolicyTree(native: Any, features: Any)
A fitted policy tree, as DoubleML's DoubleMLPolicyTree.
predict recommends treatment (1) or not (0); str(tree) shows the decision rules.
instance-attribute
¶
feature_names = (
[str(c) for c in columns]
if columns is not None
else None
)
predict(features: ArrayLike) -> Any
Recommended treatment for each row; for a DataFrame, the frame with a pred_treatment
column (as DoubleML), otherwise an integer array.
plot_tree(ax: Any = None, fontsize: int = 9) -> Any
Draw the tree with matplotlib (installed separately), as DoubleML's plot_tree:
splits show their rule, leaves the recommendation and the share of positive effects
(blue: treat, orange: do not treat); left branches are the rule's "True" side. Returns the
axes.
Learners¶
Ordinary least squares with intercept; ridge=α adds an L2 penalty as scikit-learn's Ridge(alpha=α).
Lasso by coordinate descent: the objective, stopping rule (tol) and defaults of scikit-learn's Lasso.
Lasso with the penalty chosen by cross-validation, as scikit-learn's LassoCV.
alphas is the number of penalties on the automatic grid (from alpha_max down to
eps * alpha_max) or a list of penalties; cv is the number of contiguous folds.
__new__(
alphas: Union[None, int, Sequence[float]] = None,
eps: float = 0.001,
cv: int = 5,
tol: float = 0.0001,
max_iter: int = 100000,
) -> Self
L2-penalised logistic regression (Newton's method); scikit-learn's LogisticRegression(C=C).
predict_proba(x: NDArray[float64]) -> NDArray[np.float64]
Probability of class 1 (1-D, unlike scikit-learn's two columns).
Logistic regression with C chosen by stratified cross-validation, as scikit-learn's
LogisticRegressionCV.
Cs is the number of values between 1e-4 and 1e4 or a list of values; scoring is
"neg_log_loss" (default) or "accuracy".
__new__(
Cs: Union[None, int, Sequence[float]] = None,
cv: int = 5,
scoring: Literal[
"neg_log_loss", "log_loss", "accuracy"
] = "neg_log_loss",
tol: float = 1e-10,
max_iter: int = 100,
) -> Self
Grid search over regression learners: the candidate with the smallest cross-validated mean
squared error (cv-fold, no shuffling) is refitted on all data. As a nuisance learner the
search runs inside every cross-fitting fold.
Grid search over classifiers by cross-validated log loss (stratified folds).
Random forest regressor with scikit-learn's parameters, fitted in parallel.
__new__(
n_estimators: int = 100,
max_depth: Optional[int] = None,
min_samples_split: int = 2,
min_samples_leaf: int = 1,
max_features: _MaxFeatures = None,
bootstrap: bool = True,
max_samples: Optional[Union[int, float]] = None,
random_state: int = 0,
) -> Self
Random forest classifier with scikit-learn's parameters, fitted in parallel.
__new__(
n_estimators: int = 100,
max_depth: Optional[int] = None,
min_samples_split: int = 2,
min_samples_leaf: int = 1,
max_features: _MaxFeatures = None,
bootstrap: bool = True,
max_samples: Optional[Union[int, float]] = None,
random_state: int = 0,
) -> Self
Gradient boosting (squared error) that grows the same trees as scikit-learn's
GradientBoostingRegressor. max_depth defaults to 3.
__new__(
n_estimators: int = 100,
learning_rate: float = 0.1,
max_depth: Optional[int] = ...,
min_samples_split: int = 2,
min_samples_leaf: int = 1,
max_features: _MaxFeatures = None,
subsample: float = 1.0,
random_state: int = 0,
) -> Self
Gradient boosting (log loss) that grows the same trees as scikit-learn's
GradientBoostingClassifier. max_depth defaults to 3.
__new__(
n_estimators: int = 100,
learning_rate: float = 0.1,
max_depth: Optional[int] = ...,
min_samples_split: int = 2,
min_samples_leaf: int = 1,
max_features: _MaxFeatures = None,
subsample: float = 1.0,
random_state: int = 0,
) -> Self