statspai.dml

dml

Double/Debiased Machine Learning module for StatsPAI.

Implements the Chernozhukov et al. (2018) framework with separate per-model estimator classes (PLR / IRM / PLIV / IIVM) sharing a common cross-fitting infrastructure.

Public entry points:

• :func:dml — dispatcher, selects the model via model= string.
• :class:DoubleML — legacy façade, delegates to a per-model class.
• :class:DoubleMLPLR, :class:DoubleMLIRM, :class:DoubleMLPLIV, :class:DoubleMLIIVM — direct per-model entry points.

References

Chernozhukov, V., Chetverikov, D., Demirer, M., Duflo, E., Hansen, C., Newey, W., and Robins, J. (2018). "Double/Debiased Machine Learning for Treatment and Structural Parameters." Econometrics Journal, 21(1), C1-C68. [@chernozhukov2018double]

DMLDiagnostics dataclass

Bases: ResultProtocolMixin

Bundled DML diagnostics returned by :func:dml_diagnostics.

Holds the overlap table, residual-balance table, residual-density moments, and the orthogonality-test availability for a fitted DML estimator. Use .summary() for the text report and .plot() for the 2x2 diagnostic panel.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 400
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> d = 0.5 * x1 + rng.normal(size=n)
>>> y = 1.0 * d + x1 + 0.5 * x2 + rng.normal(size=n)
>>> df = pd.DataFrame({'y': y, 'd': d, 'x1': x1, 'x2': x2})
>>> fit = sp.dml(df, y='y', treat='d', covariates=['x1', 'x2'],
...              model='plr', ml_g='linear', ml_m='linear',
...              n_folds=2)
>>> diag = sp.dml_diagnostics(fit)
>>> isinstance(diag, sp.DMLDiagnostics)
True
>>> diag.method
'PLR'

plot

plot(figsize: Any = (10.0, 8.0), bins: int = 30) -> Any

Render a 2×2 publication-style diagnostic panel.

Top-left : overlap histogram (propensity for IRM, |D-resid| for PLR). Top-right : score-density histogram with N(0,σ̂²) overlay. Bottom-left: balance bar chart (corr with each residualised covariate). Bottom-right: Q-Q plot of standardised score vs N(0,1).

DMLSensitivityResult dataclass

Bases: ResultProtocolMixin

Output of :func:dml_sensitivity.

Attributes:

Name	Type	Description
estimate	float	Original DML point estimate.
se	float	Original DML standard error.
rv_q	float	Robustness value at the user's q-threshold (default q=1 ⇒ strength of confounder needed to shrink estimate to zero).
rv_qa	float	Confounder strength needed to push the (1-α)·100% lower CI across zero. Strictly less than or equal to rv_q.
bias_bound	float	Maximum |bias| under the user-specified cf_y, cf_d.
adjusted_estimate_low	float
adjusted_estimate_high	float	Bias-adjusted estimate range under the (cf_y, cf_d) scenario.
se_low, se_high	float	Standard errors of the two ends of that range.
ci_low, ci_high	float	Interval that covers the range with probability 1 - alpha: adjusted_estimate_low - z * se_low and adjusted_estimate_high + z * se_high with the one-sided quantile z = Phi^{-1}(1 - alpha). NaN for sample-weighted fits, which do not store the score elements.
benchmarks	DataFrame	For each benchmark covariate X_k: cf_y_bench, cf_d_bench, and the implied bias / adjusted estimate when a confounder is assumed to be k_y, k_d times as strong as X_k in the residualised regression. Empty if benchmark not provided.
s	float	Scaling factor :math:S = \sqrt{\sigma^2\nu^2} used in the bias formula. For PLR it is the root-mean-square structural outcome residual :math:Y-\ell(X)-\theta(D-m(X)) over the root-mean-square treatment residual :math:D-m(X); for IRM see the module notes.
q	float
alpha	float
method	str	Always "DML-OVB (Chernozhukov-Cinelli-Newey 2022)".

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 400
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> d = 0.5 * x1 + rng.normal(size=n)
>>> y = 1.0 * d + x1 + 0.5 * x2 + rng.normal(size=n)
>>> df = pd.DataFrame({'y': y, 'd': d, 'x1': x1, 'x2': x2})
>>> fit = sp.dml(df, y='y', treat='d', covariates=['x1', 'x2'],
...              model='plr', ml_g='linear', ml_m='linear', n_folds=2)
>>> sens = sp.dml_sensitivity(fit, q=1.0, cf_y=0.05, cf_d=0.05)
>>> isinstance(sens, sp.DMLSensitivityResult)
True
>>> round(sens.rv_q, 2)  # confounder strength needed to zero out theta
0.59

plot

plot(ax: Any = None, levels: Sequence[float] = (0.0, 0.5, 1.0), figsize: Any = (6.0, 5.0)) -> Any

Plot bias-contour grid for hypothetical (cf_y, cf_d) pairs.

Plots the |bias|/|θ̂| contour as a function of the confounder strength on Y and D. The user can read off the RV directly. Mirrors the sensemakr plotting convention.

DoubleML

Legacy façade. Prefer the per-model classes directly for new code: :class:DoubleMLPLR, :class:DoubleMLIRM, :class:DoubleMLPLIV, :class:DoubleMLIIVM. Kept for backward compatibility.

For most workflows call the functional entry point :func:dml instead; this class wraps the same dispatcher behind a .fit() method for code that predates the functional API.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 400
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> d = 0.5 * x1 + rng.normal(size=n)
>>> y = 1.0 * d + x1 + 0.5 * x2 + rng.normal(size=n)
>>> df = pd.DataFrame({'y': y, 'd': d, 'x1': x1, 'x2': x2})
>>> est = sp.DoubleML(df, y='y', treat='d', covariates=['x1', 'x2'],
...                   model='plr', ml_g='linear', ml_m='linear',
...                   n_folds=2)
>>> est.model
'plr'
>>> res = est.fit()
>>> isinstance(res, sp.CausalResult)
True

References

[@chernozhukov2018double]

fit

fit(*, external_predictions: Optional[OOFPredictions] = None, store_oof: bool = False, observation_ids: Optional[Sequence[str]] = None) -> CausalResult

Fit, with optional Python-only retained IRM ATE records.

store_oof=True enables get_oof() and get_residuals(); ordinary serialization omits the individual-level records. The supported path is unweighted, unnormalised binary-treatment IRM ATE, and external caller_declared records remain declarations. Retained internal fits require at least 10 rows per treatment arm in each fold or raise DataInsufficient instead of using a subgroup-mean fallback; internal explicit fold_indices require n_rep=1, while external repeated IRM predictions are allowed when the explicit partition is equivalent for every repeat. external_predictions, store_oof, and observation_ids are Python-only controls.

Parameters:

Name	Type	Description	Default
external_predictions	OOFPredictions or None	In-memory predictions and declared training scopes, aligned to every analysis row, fold, and repeat.	None
store_oof	bool	Retain complete IRM ATE records for explicit audit export.	False
observation_ids	sequence of str or None	Unique input row identifiers; generated ordinals otherwise.	None

Returns:

Type	Description
CausalResult	The configured model's effect and uncertainty.

Examples:

>>> import statspai as sp
>>> callable(sp.DoubleML.fit)
True

References

See docs/dml_oof_audit.md for the audit contract.

DynamicDMLResult dataclass

Bases: ResultProtocolMixin

Output of :func:dynamic_dml.

Attributes:

Name	Type	Description
periods	DataFrame	One row per period, indexed by the period label, with lag (periods before the outcome), estimate, se, z, p, ci_low and ci_high.
estimate, se, ci_lower, ci_upper, p_value	float	The total sequence effect -- treating in every period rather than none -- and its inference from the joint covariance.
vcov	ndarray	Joint (m, m) covariance of the per-period effects, the thing that makes any contrast across periods reportable.
coef	DataFrame or None	Heterogeneity coefficients when modifiers= was given: one block per period, the linear projection of that period's effect on the modifiers.
n_units, n_periods, n_folds	int
period_labels	list
lags	int
method	str	Always "dynamic_dml".
diagnostics	dict	first_stage_r2_outcome / first_stage_r2_treatment per period, state_columns, lagged_treatment_included, n_units_dropped, independent_sum_se, modifiers.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(1)
>>> rows = []
>>> for i in range(250):
...     state, prev = rng.normal(), 0.0
...     hist = []
...     for period in range(3):
...         if period:
...             state = 0.6 * state + 0.7 * prev + rng.normal()
...         dose = 0.8 * state + 0.5 * prev + rng.normal()
...         hist.append((state, dose))
...         prev = dose
...     out = sum(0.5 * s + 0.4 * d for s, d in hist) + rng.normal()
...     for period, (s, d) in enumerate(hist):
...         rows.append({"id": i, "t": period, "y": out, "d": d, "w": s})
>>> df = pd.DataFrame(rows)
>>> res = sp.dynamic_dml(df, y="y", treat="d", id="id", time="t",
...                      covariates=["w"], n_folds=3)
>>> list(res.periods.columns)
['lag', 'estimate', 'se', 'z', 'p', 'ci_low', 'ci_high']
>>> res.vcov.shape
(3, 3)
>>> # Any contrast uses the joint covariance, so summing periods does
>>> # not overstate the uncertainty.
>>> bool(res.se < res.diagnostics["independent_sum_se"])
True
>>> sorted(res.contrast([1, 1, 1]))
['ci_high', 'ci_low', 'estimate', 'p', 'se', 'z']
>>> res.cumulative().shape
(3, 7)

contrast

contrast(weights: Any) -> Dict[str, float]

Inference for w' theta, using the joint covariance.

weights=[1, 1, 1] is the total sequence effect; [0, 0, 1] the last period alone; [1, -1, 0] asks whether the first two periods differ.

cumulative

cumulative() -> DataFrame

Effect of treating from each period to the end, with joint SEs.

DoubleMLIIVM

Bases: _DoubleMLBase

Interactive IV DML — binary D, binary Z, LATE via Wald.

Direct entry point for the interactive IV model with a binary treatment and binary instrument; estimates the LATE (complier effect) via the Neyman-orthogonal ratio of two doubly-robust scores. Usually reached through sp.dml(..., model='iivm').

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 400
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> z = rng.integers(0, 2, size=n)
>>> d = ((0.6 * z + 0.3 * x1 + rng.normal(size=n)) > 0.3).astype(int)
>>> y = 1.0 * d + x1 + 0.5 * x2 + rng.normal(size=n)
>>> df = pd.DataFrame({"y": y, "d": d, "z": z, "x1": x1, "x2": x2})
>>> est = sp.DoubleMLIIVM(
...     df, y="y", treat="d", covariates=["x1", "x2"],
...     instrument="z", n_folds=3,
... )
>>> res = est.fit()
>>> bool(np.isfinite(res.estimate))
True

DoubleMLIRM

Bases: _DoubleMLBase

Interactive regression DML — binary D, ATE via AIPW.

Direct entry point for the interactive regression model with a binary treatment; estimates the ATE via the cross-fitted doubly-robust (AIPW) score. Usually reached through the dispatcher sp.dml(..., model='irm').

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 400
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> ps = 1 / (1 + np.exp(-(0.5 * x1)))
>>> d = (rng.uniform(size=n) < ps).astype(int)
>>> y = 1.0 * d + x1 + 0.5 * x2 + rng.normal(size=n)
>>> df = pd.DataFrame({"y": y, "d": d, "x1": x1, "x2": x2})
>>> est = sp.DoubleMLIRM(
...     df, y="y", treat="d", covariates=["x1", "x2"], n_folds=3,
... )
>>> res = est.fit()
>>> bool(np.isfinite(res.estimate))
True

DMLAveragingResult

Bases: CausalResult

CausalResult extended with per-candidate and weight details.

Attributes stored in model_info:

• candidates — list of candidate labels.
• theta_k — per-candidate :math:\hat\theta (only meaningful for the non-stacking weight rules).
• se_k — per-candidate SE.
• mse_k — per-candidate nuisance risk (g + m).
• weights — averaging or stacking weights.
• weights_g / weights_m — separate stacking weights per nuisance under weight_rule="short_stacking".
• weight_rule — how the weights were computed.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> from sklearn.linear_model import LassoCV, RidgeCV
>>> rng = np.random.default_rng(0)
>>> n = 300
>>> X = rng.normal(size=(n, 5))
>>> d = X[:, 0] + rng.normal(size=n)
>>> y = 1.0 * d + X[:, 0] + 0.5 * X[:, 1] + rng.normal(size=n)
>>> df = pd.DataFrame({"y": y, "d": d,
...                    **{f"x{j}": X[:, j] for j in range(5)}})
>>> candidates = [(LassoCV(), LassoCV(), "lasso"),
...               (RidgeCV(), RidgeCV(), "ridge")]
>>> r = sp.dml_model_averaging(df, y="y", treat="d",
...                            covariates=[f"x{j}" for j in range(5)],
...                            candidates=candidates, n_folds=3)
>>> isinstance(r, sp.DMLAveragingResult)
True
>>> sorted(r.model_info["weights_m"])  # CLS stacking weights
['lasso', 'ridge']

OOFBundle dataclass

Validated StatsPAI binary-ATE score output with bounded ULP slack.

Returns:

Type	Description
OOFBundle	Immutable schema-v1 predictions, scores, and repeat aggregation.

Examples:

>>> import statspai as sp
>>> sp.OOFBundle.__name__
'OOFBundle'

Notes

Obtain from a retained result or construct with :meth:from_arrays.

References

See docs/dml_oof_audit.md for score and variance conventions.

from_arrays classmethod

from_arrays(*, predictions: OOFPredictions, ps_used: Any, psi_b: Any, psi: Any, theta: Any, se: Any, input_positions: Any, dropped_positions: Any, aggregation: Mapping[str, Any], metadata: Mapping[str, Any]) -> 'OOFBundle'

Validate and copy canonical binary-ATE scoring output.

Parameters:

Name	Type	Description	Default
predictions	OOFPredictions	Complete validated analysis snapshot.	required
ps_used	array-like of shape (n_rep, n_obs)	Clipped propensities, AIPW pseudo-outcomes, and centered scores.	required
psi_b	array-like of shape (n_rep, n_obs)	Clipped propensities, AIPW pseudo-outcomes, and centered scores.	required
psi	array-like of shape (n_rep, n_obs)	Clipped propensities, AIPW pseudo-outcomes, and centered scores.	required
theta	array-like of shape (n_rep,)	Per-repeat effects and standard errors.	required
se	array-like of shape (n_rep,)	Per-repeat effects and standard errors.	required
input_positions	array-like of int	Disjoint mappings covering the original input row positions.	required
dropped_positions	array-like of int	Disjoint mappings covering the original input row positions.	required
aggregation	mapping	Validated median-effect and median-variance aggregation record.	required
metadata	mapping	Trimming, fit, and provenance metadata required by schema v1.	required

Returns:

Type	Description
OOFBundle	Independent snapshot; inconsistent scores or mappings raise.

to_frame

to_frame() -> DataFrame

Return repeat-major score rows without individual covariate values.

Returns:

Type	Description
DataFrame	Fresh row/fold identity, outcomes, treatments, predictions and scores, with one row per repeat and analysis observation.

score_concentration

score_concentration(*, top_fraction: float = 0.01) -> tuple[dict, ...]

Return per-repeat concentration diagnostics for centered scores.

Parameters:

Name	Type	Description	Default
top_fraction	float	Fraction in (0, 1] used for the largest squared-score share sum.	0.01

Returns:

Type	Description
tuple of dict	One independent JSON-ready record per repeat, with cmax, z_c, top_share, h and effective_score_count. Undefined diagnostics have status='diagnostic_unavailable' and explicit failure_reason.

References

See docs/dml_oof_audit.md for definitions; these are diagnostics, not tests of causal identification.

to_json

to_json(path: Union[str, Path]) -> None

Write full individual-level IDs, Y, D, X, nuisance predictions, and scores.

Choose the path, permissions, and sharing scope under the applicable user-data policy; the file carries one whole-bundle hash.

Parameters:

Name	Type	Description	Default
path	str or Path	Destination for full individual-level schema-v1 records.	required

Returns:

Type	Description
None	Writes deterministic UTF-8 JSON.

from_json classmethod

from_json(path: Union[str, Path]) -> 'OOFBundle'

Read only after checking schema, duplicate keys, and stored hash.

Parameters:

Name	Type	Description	Default
path	str or Path	Schema-v1 bundle JSON file.	required

Returns:

Type	Description
OOFBundle	Validated independent snapshot, including reconstructed scores.

OOFPredictions dataclass

Validated predictions and declared cross-fitting provenance.

Records marked caller_declared describe claims StatsPAI can check for structural consistency, not training behavior StatsPAI observed. Public array access returns a fresh read-only NumPy header over immutable bytes.

Returns:

Type	Description
OOFPredictions	Immutable schema-v1 data, predictions, folds, and declared scopes.

Examples:

>>> import statspai as sp
>>> sp.OOFPredictions.__name__
'OOFPredictions'

Notes

Construct with :meth:from_arrays, not the class constructor.

References

See docs/dml_oof_audit.md for the data contract.

from_arrays classmethod

from_arrays(*, ids: Sequence[str], y: Any, d: Any, x: Any, covariate_names: Sequence[str], g0: Any, g1: Any, ps_raw: Any, fold_ids: Any, training_records: Sequence[Mapping[str, Any]], source: Mapping[str, Any]) -> 'OOFPredictions'

Validate and defensively copy a complete analysis-sample snapshot.

Parameters:

Name	Type	Description	Default
ids	sequence of str	Unique row identifiers in analysis order.	required
y	array-like of shape (n_obs,)	Finite outcomes and binary treatment containing both arms.	required
d	array-like of shape (n_obs,)	Finite outcomes and binary treatment containing both arms.	required
x	array-like of shape (n_obs, n_covariates)	Finite covariate matrix in the declared column order.	required
covariate_names	sequence of str	Unique covariate names.	required
g0	array-like of shape (n_rep, n_obs)	Outcome predictions by arm and raw propensities in [0, 1].	required
g1	array-like of shape (n_rep, n_obs)	Outcome predictions by arm and raw propensities in [0, 1].	required
ps_raw	array-like of shape (n_rep, n_obs)	Outcome predictions by arm and raw propensities in [0, 1].	required
fold_ids	array-like of int of shape (n_rep, n_obs)	Complete held-out fold assignments for each repeat.	required
training_records	sequence of mapping	One record per repeat/fold declaring disjoint train/test IDs and nuisance and preprocessing training scopes. See the audit guide.	required
source	mapping	Engine, recipe, seed, and software version provenance.	required

Returns:

Type	Description
OOFPredictions	Validated independent snapshot with corruption hashes.

Examples:

>>> import statspai as sp
>>> callable(sp.OOFPredictions.from_arrays)
True

References

See docs/dml_oof_audit.md for training-record fields.

validate_alignment

validate_alignment(*, ids: Sequence[str], y: Any, d: Any, x: Any, covariate_names: Sequence[str]) -> None

Raise unless current analysis inputs exactly match this row order.

Parameters:

Name	Type	Description	Default
ids	sequence of str	Ordered identifiers to compare with the snapshot.	required
covariate_names	sequence of str	Ordered identifiers to compare with the snapshot.	required
y	array - like	Observed arrays to compare exactly, including shape.	required
d	array - like	Observed arrays to compare exactly, including shape.	required
x	array - like	Observed arrays to compare exactly, including shape.	required

Returns:

Type	Description
None	Raises ValueError on any mismatch.

to_json

to_json(path: Union[str, Path]) -> None

Write deterministic UTF-8 JSON with three partition hashes.

Parameters:

Name	Type	Description	Default
path	str or Path	Destination for individual-level data, predictions, and scopes.	required

Returns:

Type	Description
None	Writes schema-v1 JSON to the requested path.

from_json classmethod

from_json(path: Union[str, Path]) -> 'OOFPredictions'

Read only after checking schema, duplicate keys, and stored hashes.

Parameters:

Name	Type	Description	Default
path	str or Path	Schema-v1 prediction JSON file.	required

Returns:

Type	Description
OOFPredictions	Validated independent snapshot.

DMLPanelResult dataclass

Bases: ResultProtocolMixin

Output of :func:dml_panel.

Attributes:

Name	Type	Description
estimate	float	Debiased treatment effect β̂.
se	float	Cluster-robust SE at the unit level.
ci_lower, ci_upper	float
p_value	float
t_stat	float
n_units	int
n_obs	int
n_folds	int
include_time_fe	bool
ml_g_name	str	Short name of the outcome nuisance learner.
ml_m_name	str	Short name of the treatment nuisance learner.
method	str	Always "dml_panel".
diagnostics	dict	Populated with {'y_resid_std', 'd_resid_std', 'corr_yd_resid', 'within_r2', 'omega_cluster'}.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> rows = []
>>> for i in range(50):
...     alpha_i = rng.normal()
...     for t in range(5):
...         x1, x2 = rng.normal(), rng.normal()
...         union = 0.4 * x1 + alpha_i + rng.normal()
...         lw = (0.8 * union + x1 + 0.5 * x2 + alpha_i
...               + 0.1 * t + rng.normal())
...         rows.append({'pid': i, 'year': t, 'log_wage': lw,
...                      'union': union, 'exper': x1, 'educ': x2})
>>> df = pd.DataFrame(rows)
>>> res = sp.dml_panel(df, y='log_wage', treat='union',
...                    covariates=['exper', 'educ'], unit='pid',
...                    time='year', include_time_fe=True, n_folds=3)
>>> isinstance(res, sp.DMLPanelResult)
True
>>> res.method
'dml_panel'

DoubleMLPLIV

Bases: _DoubleMLBase

Partially linear IV DML — endogenous D with continuous/binary Z.

Direct entry point for the partially linear IV model Y = theta * D + g(X) + eps with E[eps | Z, X] = 0. Usually reached through the dispatcher sp.dml(..., model='pliv'); an instrument column must be supplied.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 400
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> z = rng.normal(size=n)
>>> d = 0.8 * z + 0.5 * x1 + rng.normal(size=n)
>>> y = 1.5 * d + x1 + 0.5 * x2 + rng.normal(size=n)
>>> df = pd.DataFrame({"y": y, "d": d, "z": z, "x1": x1, "x2": x2})
>>> est = sp.DoubleMLPLIV(
...     df, y="y", treat="d", covariates=["x1", "x2"],
...     instrument="z", n_folds=3,
... )
>>> res = est.fit()
>>> bool(np.isfinite(res.estimate))
True

DoubleMLPLR

Bases: _DoubleMLBase

Partially linear regression DML (continuous or binary D, no IV).

Direct entry point for the partially linear model Y = theta * D + g(X) + eps. Usually reached through the dispatcher sp.dml(..., model='plr'); the class is exposed for callers who want to construct the estimator explicitly.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 300
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> d = 0.5 * x1 + rng.normal(size=n)
>>> y = 1.5 * d + x1 + 0.5 * x2 + rng.normal(size=n)
>>> df = pd.DataFrame({"y": y, "d": d, "x1": x1, "x2": x2})
>>> est = sp.DoubleMLPLR(
...     df, y="y", treat="d", covariates=["x1", "x2"], n_folds=3,
... )
>>> res = est.fit()
>>> bool(np.isfinite(res.estimate))
True

dml_diagnostics

dml_diagnostics(result: Any, clip: float = 0.02) -> DMLDiagnostics

Build a :class:DMLDiagnostics report from a DML CausalResult.

Parameters:

Name	Type	Description	Default
result	CausalResult	Result returned by :func:statspai.dml.dml. Must include the post-fit residuals (model_info['_y_resid'], model_info['_d_resid']); for IRM, additionally the propensity model_info['diagnostics']['pscore_min'] etc. are surfaced.	required
clip	float	For IRM-style overlap: count units with propensity within [0, clip] ∪ [1-clip, 1] as overlap-violating.	0.02

Returns:

Type	Description
DMLDiagnostics

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 400
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> d = 0.5 * x1 + rng.normal(size=n)
>>> y = 1.0 * d + x1 + 0.5 * x2 + rng.normal(size=n)
>>> df = pd.DataFrame({'y': y, 'd': d, 'x1': x1, 'x2': x2})
>>> fit = sp.dml(df, y='y', treat='d', covariates=['x1', 'x2'],
...              model='plr', ml_g='linear', ml_m='linear',
...              n_folds=2)
>>> diag = sp.dml_diagnostics(fit)
>>> diag.n_obs
400
>>> diag.method
'PLR'
>>> diag.orth_pvalue is None
True
>>> text = diag.summary()  # overlap / score / balance report

dml_sensitivity

dml_sensitivity(result: Any, q: float = 1.0, cf_y: Optional[float] = None, cf_d: Optional[float] = None, benchmark_covariates: Optional[Sequence[str]] = None, k_y: float = 1.0, k_d: float = 1.0) -> DMLSensitivityResult

Compute DML-OVB sensitivity for a fitted DML CausalResult.

Parameters:

Name	Type	Description	Default
result	CausalResult	From :func:statspai.dml.dml with model='plr' or model='irm' (ATE or ATTE score), including clustered and repeated cross-fitting fits. IV models are refused: the bound is for estimands identified by conditional exogeneity.	required
q	float	Bias threshold as a fraction of |θ̂|. q=1 ⇒ confounder needed to shrink estimate to zero; q=0.5 ⇒ half the estimate.	1.0
cf_y	float	Hypothesized partial-R² of an unobserved confounder with the residualised outcome and treatment. If both are given, the report includes a bias bound, the adjusted-estimate range, the standard errors of its two ends and an interval for the range.	None
cf_d	float	Hypothesized partial-R² of an unobserved confounder with the residualised outcome and treatment. If both are given, the report includes a bias bound, the adjusted-estimate range, the standard errors of its two ends and an interval for the range.	None
benchmark_covariates	list of str	Subset of the original covariates to benchmark against. For each X_k, the benchmark sets cf_y_bench, cf_d_bench to the partial R² that X_k itself contributes (multiplied by k_y, k_d to express "what if a confounder were k× as strong as X_k?").	None
k_y	float	Multipliers for the benchmark strengths.	1.0
k_d	float	Multipliers for the benchmark strengths.	1.0

Returns:

Type	Description
DMLSensitivityResult

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 400
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> d = 0.5 * x1 + rng.normal(size=n)
>>> y = 1.0 * d + x1 + 0.5 * x2 + rng.normal(size=n)
>>> df = pd.DataFrame({'y': y, 'd': d, 'x1': x1, 'x2': x2})
>>> fit = sp.dml(df, y='y', treat='d', covariates=['x1', 'x2'],
...              model='plr', ml_g='linear', ml_m='linear',
...              n_folds=2)
>>> sens = sp.dml_sensitivity(fit, q=1.0, cf_y=0.05, cf_d=0.05)
>>> round(sens.estimate, 2)
0.98
>>> round(sens.rv_q, 2)  # confounder strength to zero out theta
0.59
>>> round(sens.bias_bound, 2)  # |bias| under cf_y = cf_d = 0.05
0.05

dml_did

dml_did(data: DataFrame, y: str, treat: str, covariates: Sequence[str], *, time: Optional[str] = None, id: Optional[str] = None, ml_g: Any = None, ml_m: Any = None, score: str = 'observational', in_sample_normalization: bool = True, n_folds: int = 5, n_rep: int = 1, trimming_threshold: float = 0.01, fold_indices: Any = None, random_state: int = 42, alpha: float = 0.05) -> CausalResult

Double machine learning difference-in-differences (two periods).

Estimates the average treatment effect on the treated under parallel trends that hold conditionally on covariates, with the nuisance functions fitted by machine learners and cross-fitted.

Parameters:

Name	Type	Description	Default
data	DataFrame		required
y	str	Outcome. With time=None this is the change in the outcome between the two periods, one row per unit. Otherwise it is the outcome level.	required
treat	str	0/1 indicator of the treated group. On a long panel it may equal 1 only in the post-period rows of treated units; a unit is treated if it is ever 1.	required
covariates	sequence of str	Pre-treatment covariates. On a long panel the first-period values are used.	required
time	str	Period column with two values; the larger is the post-period. Omit it when y is already the outcome change.	None
id	str	Unit identifier. With time and id the data are a two-period panel in long form and the outcome is differenced within unit. With time alone the rows are repeated cross-sections.	None
ml_g	estimator or str	Regressor for the outcome nuisance: E[dY | X, D=0] on a panel, E[Y | X, D=d, T=t] in each of the four cells for repeated cross-sections. A scikit-learn estimator or an alias of :func:statspai.dml ('lasso', 'rf', 'linear' ...). Default: gradient boosting, as in :func:statspai.dml.	None
ml_m	estimator or str	Classifier for the propensity score P(D=1 | X). Not used with score='experimental'.	None
score	('observational', 'experimental')	'observational' is the doubly robust score for conditional parallel trends. 'experimental' assumes treatment is independent of the covariates (a randomised design), uses no propensity score, and uses the covariates only to reduce variance.	'observational'
in_sample_normalization	bool	Normalise the weights to average one in the sample (Hajek-type, the DoubleML default). False gives the score of Chang (2020): (D - m(X)) / (p (1 - m(X))) * (dY - g0(X)).	True
n_folds	int		5
n_rep	int	Repetitions of the sample split; estimates are aggregated by the median, and the variance adds the spread across splits.	1
trimming_threshold	float	Propensity scores are clipped to [c, 1 - c]. The number clipped is reported and a warning is issued when it is positive.	0.01
fold_indices	array - like or str	One fold label per analysis row (per unit on a panel), or a column name. Requires n_rep=1.	None
random_state	int		42
alpha	float		0.05

Returns:

Type	Description
CausalResult	estimate is the ATT. model_info holds the score, the nuisance fit (rmse_g0, logloss_m), the propensity range, the number of clipped scores, the per-repetition estimates and _psi, the influence values of the last repetition.

Notes

The estimand is the ATT in the post-period. The identifying assumption is that, given X, the untreated outcome of the treated group would have changed as that of the comparison group did; X must be measured before treatment.

The standard error is sqrt(mean(psi ** 2) / n) for the normalised influence function psi. On a panel each row is one unit, so this is robust to serial correlation within unit.

With a correctly specified parametric model the estimator is :func:statspai.drdid; this function is for when the outcome trend or the propensity score needs a flexible learner. Staggered adoption is :func:statspai.callaway_santanna.

Cross-fitting removes overfitting bias, not the bias of a learner that cannot fit the nuisance at the sample size at hand. Compare model_info['rmse_g0'] and model_info['logloss_m'] across learners, as one would for any double machine learning estimate. On a design with a quadratic trend and propensity and 600 units, a random forest leaves the estimate 0.7 above a true effect of 2 (0.1 above at 20,000 units), while a learner with quadratic terms is centred on it (tests/reference_parity/test_dml_did_doubleml_parity.py).

Examples:

>>> import numpy as np, pandas as pd, statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 2000
>>> x = rng.normal(size=(n, 3))
>>> d = rng.binomial(1, 1 / (1 + np.exp(-0.5 * x[:, 0])))
>>> dy = 1.0 + x[:, 0] ** 2 + 2.0 * d + rng.normal(size=n)
>>> df = pd.DataFrame(x, columns=["x1", "x2", "x3"]).assign(d=d, dy=dy)
>>> fit = sp.dml_did(df, "dy", "d", ["x1", "x2", "x3"])
>>> bool(abs(fit.estimate - 2.0) < 0.3)
True

References

[@chang2020double], [@santanna2020doubly], [@chernozhukov2018double]

dml

dml(data: DataFrame, y: str, treat: str, covariates: List[str], model: str = 'plr', instrument: Optional[Union[str, List[str]]] = None, ml_g: Optional[Any] = None, ml_m: Optional[Any] = None, ml_r: Optional[Any] = None, n_folds: int = 5, n_rep: int = 1, alpha: float = 0.05, random_state: int = 42, sample_weight: Optional[Any] = None, fold_indices: Optional[Any] = None, score: Optional[str] = None, normalize_ipw: bool = False, trimming_threshold: float = 0.01, *, cluster: Optional[Union[str, Sequence[str]]] = None, external_predictions: Optional[OOFPredictions] = None, store_oof: bool = False, observation_ids: Optional[Sequence[str]] = None, design: Optional[Any] = None) -> CausalResult

Estimate causal effect using Double/Debiased Machine Learning.

Parameters:

Name	Type	Description	Default
data	DataFrame		required
y	str	Outcome variable.	required
treat	str	Treatment variable.	required
covariates	list of str	High-dimensional controls X.	required
model	str	DML model: 'plr' : partially linear (continuous or binary D) 'irm' : interactive regression (binary D; ATE via AIPW) 'pliv' : partially linear IV (endogenous D with instrument Z) 'iivm' : interactive IV (binary D, binary Z; LATE via Wald ratio)	'plr'
instrument	str	Scalar instrument Z. Required for model in {'pliv', 'iivm'}.	None
ml_g	sklearn estimator or str	Learners for outcome / treatment / instrument nuisance. Pass any scikit-learn-compatible estimator, or one of the convenience string aliases (case-insensitive): 'rf' — random forest (200 trees) 'gbm' — gradient boosting (100 estimators, depth 3) 'lasso' / 'ridge' — penalised linear (CV-tuned) 'linear' / 'ols' — plain linear regression 'logistic' — logistic regression (classifier roles only) 'xgb' / 'lgbm' — XGBoost / LightGBM (optional install) For binary treatment (model='irm') ml_m is auto-coerced to the classifier variant; same for ml_r under 'iivm'. Under 'plr' and 'pliv' a classifier may be passed for a 0/1 outcome, treatment or instrument; its predict_proba is the conditional mean (as in DoubleML's PLR). A classifier for a target that is not 0/1 is refused. Up to 1.38.0 these two models called predict on a classifier and residualised on the hard label. None falls back to the per-model gradient-boosting default.	None
ml_m	sklearn estimator or str	Learners for outcome / treatment / instrument nuisance. Pass any scikit-learn-compatible estimator, or one of the convenience string aliases (case-insensitive): 'rf' — random forest (200 trees) 'gbm' — gradient boosting (100 estimators, depth 3) 'lasso' / 'ridge' — penalised linear (CV-tuned) 'linear' / 'ols' — plain linear regression 'logistic' — logistic regression (classifier roles only) 'xgb' / 'lgbm' — XGBoost / LightGBM (optional install) For binary treatment (model='irm') ml_m is auto-coerced to the classifier variant; same for ml_r under 'iivm'. Under 'plr' and 'pliv' a classifier may be passed for a 0/1 outcome, treatment or instrument; its predict_proba is the conditional mean (as in DoubleML's PLR). A classifier for a target that is not 0/1 is refused. Up to 1.38.0 these two models called predict on a classifier and residualised on the hard label. None falls back to the per-model gradient-boosting default.	None
ml_r	sklearn estimator or str	Learners for outcome / treatment / instrument nuisance. Pass any scikit-learn-compatible estimator, or one of the convenience string aliases (case-insensitive): 'rf' — random forest (200 trees) 'gbm' — gradient boosting (100 estimators, depth 3) 'lasso' / 'ridge' — penalised linear (CV-tuned) 'linear' / 'ols' — plain linear regression 'logistic' — logistic regression (classifier roles only) 'xgb' / 'lgbm' — XGBoost / LightGBM (optional install) For binary treatment (model='irm') ml_m is auto-coerced to the classifier variant; same for ml_r under 'iivm'. Under 'plr' and 'pliv' a classifier may be passed for a 0/1 outcome, treatment or instrument; its predict_proba is the conditional mean (as in DoubleML's PLR). A classifier for a target that is not 0/1 is refused. Up to 1.38.0 these two models called predict on a classifier and residualised on the hard label. None falls back to the per-model gradient-boosting default.	None
n_folds	int		5
n_rep	int	Repeated cross-fitting splits (median aggregation).	1
alpha	float		0.05
random_state	int	Seed for the cross-fitting fold assignment. Repeat splits use random_state + rep so a single random_state fully determines the result.	42
sample_weight	ndarray | Series | str	Per-observation weights (e.g., survey/probability weights). Pass either a 1-D array of length len(data) or a column name present in data. Supported for model in {'plr', 'irm', 'pliv', 'iivm'}.	None
fold_indices	array - like or str	Explicit cross-fit fold labels. Pass either a 1-D array of length len(data) or a column name. Scikit-learn style splits, a list of (train, test) positional index pairs such as list(KFold(5).split(X)), are also read, provided the test sets partition the rows and each training set is its complement. For internal fitting, labels must define exactly n_folds non-empty folds and n_rep must be 1. External repeated IRM predictions may use the same labels when that partition is equivalent to every caller repeat.	None
score	str	Orthogonal score variant (DoubleML-compatible). Model-specific: model='plr' : 'partialling out' (default) or 'IV-type'. The IV-type score fits a third nuisance g(X) = E[Y - theta~ D | X] (cross-fitted on the same partition with a clone of ml_g, theta~ the preliminary partialling-out estimate) and solves E[(Y - D theta - g(X))(D - m(X))] = 0, exactly as both DoubleML ports do. With linear nuisance learners it coincides with partialling out to machine precision; with regularised or nonlinear learners it is a distinct estimator. Versions before 1.29.0 used l(X) = E[Y|X] in place of g(X) and did not reproduce DoubleML (see MIGRATION.md). model='irm' : 'ATE' (default) or 'ATTE' (effect on the treated) None selects each model's default; the defaults reproduce the historical StatsPAI output exactly. Invalid for pliv / iivm beyond their single native score.	None
normalize_ipw	bool	Self-normalize the inverse-propensity weights (Hájek-style) for the IPW-based models irm / iivm. Matches DoubleML's normalize_ipw. Ignored (and rejected) for plr / pliv.	False
trimming_threshold	float	Symmetric propensity-score trimming for irm / iivm: the estimated propensity is clipped to [t, 1 - t]. The default 0.01 reproduces the historical clip and matches DoubleML's trimming_threshold with trimming_rule='truncate'.	0.01
cluster	str or list of two str	Column identifying clusters, or two columns for two-way clustering. Cross-fitting folds are formed over whole clusters and the standard error is the cluster-robust DML variance of Chiang, Kato, Ma and Sasaki (2022) [@chiang2022multiway], as DoubleML computes it for DoubleMLClusterData; the estimate solves the fold-weighted score. Two-way: each dimension is split into n_folds folds of clusters, the n_folds**2 cells test the units in fold k of the first and fold l of the second dimension and train on units sharing no cluster with them, and the variance sums the two dimensions' within-cell cluster scores (equal to DoubleML 0.11.3 to 1e-15 given the same folds, tests/reference_parity/test_dml_twoway_cluster_doubleml_parity.py). Pass fold_indices as an n x 2 array (fold of each unit's cluster in each dimension) to fix the cells. With two-way clusters sp.dml_sensitivity reports the bound without its sampling error (no two-way variance for the bound's score elements). Not available with external_predictions or two-way with design=.	None
external_predictions	OOFPredictions	Python-only in-memory predictions for unweighted, unnormalised binary-treatment IRM ATE scoring. caller_declared training records are auditable declarations, not proof that training avoided leakage.	None
store_oof	bool	Python-only retention switch for the supported IRM ATE scope. With store_oof=True, use get_oof() or get_residuals() on the result. Ordinary result serialization omits these individual-level records.	False
observation_ids	sequence of str	Python-only stable row IDs. Internal IDs are checked before complete-case filtering; external IDs must match the prediction record exactly.	None
design	SurveyDesign	A complex survey design from :func:sp.svydesign declared on data (same rows). The design weights are the sample weights (nuisance fits and the score), cross-fitting folds are formed over whole PSUs when the design has them, and the standard error is the survey linearisation of the weighted score, u_i = -w_i psi_i / sum(w psi_a), over strata, PSUs and the finite population correction (R survey::svyrecvar), with t inference on the design degrees of freedom (#PSU - #strata). The nuisance fits are held fixed, as in every DML variance. Not with sample_weight=, cluster=, external_predictions or store_oof; with model='irm' only the unnormalised ATE score is weighted.	None

Returns:

Type

Description

CausalResult

For model='pliv' (unweighted, no clusters), model_info['anderson_rubin'] holds a confidence set for the coefficient that stays valid when the instrument is weak: the values of theta at which the Anderson-Rubin statistic of the cross-fitted score is below the chi-squared(1) critical value. kind is 'interval', or 'unbounded' / 'disjoint' when the first stage is not significant, and intervals lists the pieces. With n_rep > 1 it refers to the last repetition.

Notes

Numerical parity: all four models are pinned against the upstream doubleml-for-py reference implementation in tests/external_parity/test_dml_python_parity.py — the partialling-out models (plr, pliv) agree to machine precision under shared learners and folds. See the guide "sp.dml and the DoubleML reference implementation" in the docs for the full parity table and methodology.

Declared scope boundaries (each fails loudly with a workaround in the error message — see the guide's "Scope and known limitations"):

• pliv / iivm accept a single scalar instrument; for multiple excluded instruments build a first-stage index with sp.scalar_iv_projection first.
• One treatment column per call (no multi-treatment joint inference; combine per-treatment calls with sp.romano_wolf for simultaneous inference).
• DML2 (pooled-moment) procedure only; PLR exposes the partialling-out score only.
• Retained internal IRM fits require at least 10 rows per treatment arm in every training fold or raise DataInsufficient instead of using the legacy subgroup-mean fallback; internal explicit fold_indices require n_rep=1, while external repeated IRM predictions are accepted when the explicit fold partition is equivalent to every repeat.
• Cluster-robust DML inference lives in sp.dml_panel (unit-clustered panel PLR), not in this cross-sectional entry point.

Default learners and coverage. When no learners are passed, the nuisances are fitted by gradient boosting (100 trees, depth 3). That default is a convenience, not a validated choice: on the Track B PLR design (n = 500, nonlinear nuisances) its regularisation bias is about half a Monte Carlo SD and the 95% interval covers 0.89 (0.94 at n = 2,000), while the true nuisance functions give 0.95 and a random forest or a spline Lasso about 0.94 (tests/coverage_monte_carlo/mechanisms/dml_plr_learners.py). The orthogonal score makes nuisance error second order, not zero; pass learners suited to the problem, and check sp.validation_scope for which learner configurations carry evidence. A PLR fit that used the defaults records this in model_info["learner_note"], which summary() prints.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 500
>>> df = pd.DataFrame(rng.normal(size=(n, 3)), columns=['age', 'edu', 'exp'])
>>> X_cols = ['age', 'edu', 'exp']

>>> # Partially Linear Regression (true effect 1.0)
>>> df['training'] = 0.5 * df['age'] + rng.normal(size=n)
>>> df['wage'] = 1.0 * df['training'] + df['edu'] + rng.normal(size=n)
>>> result = sp.dml(df, y='wage', treat='training',
...                 covariates=['age', 'edu', 'exp'])
>>> result.method, round(float(result.estimate), 1)
('Double ML (PLR)', 1.0)

>>> # Interactive Regression (binary treatment, ATE; true effect 2.0)
>>> df['D'] = rng.binomial(1, 0.5, n)
>>> df['y_bin'] = 2.0 * df['D'] + df['edu'] + rng.normal(size=n)
>>> result = sp.dml(df, y='y_bin', treat='D', covariates=X_cols,
...                 model='irm')
>>> result.method, result.estimand
('Double ML (IRM)', 'ATE')

>>> # Partially Linear IV — endogenous schooling, instrument Z
>>> df['Z'] = rng.normal(size=n)
>>> u = rng.normal(size=n)                     # unobserved confounder
>>> df['schooling'] = df['Z'] + u + rng.normal(size=n)
>>> df['earnings'] = 0.5 * df['schooling'] + df['age'] + u + rng.normal(size=n)
>>> result = sp.dml(df, y='earnings', treat='schooling',
...                 covariates=['age', 'edu'],
...                 model='pliv', instrument='Z')
>>> result.method, result.estimand
('Double ML (PLIV)', 'LATE')

>>> # Interactive IV — binary D, binary Z (LATE)
>>> df['lottery_win'] = rng.binomial(1, 0.5, n)
>>> df['college'] = ((df['lottery_win'] + u + rng.normal(size=n)) > 0.5)
>>> df['college'] = df['college'].astype(int)
>>> df['earnings2'] = 1.0 * df['college'] + df['age'] + u + rng.normal(size=n)
>>> result = sp.dml(df, y='earnings2', treat='college',
...                 covariates=['age', 'edu'],
...                 model='iivm', instrument='lottery_win')
>>> result.method, result.estimand
('Double ML (IIVM)', 'LATE')

dml_model_averaging

dml_model_averaging(data: DataFrame, y: str, treat: str, covariates: Sequence[str], candidates: Optional[List[Tuple[Any, Any, str]]] = None, n_folds: int = 5, seed: int = 0, weight_rule: str = 'short_stacking', alpha: float = 0.05, sample_weight: Optional[Any] = None, fold_indices: Optional[Any] = None) -> DMLAveragingResult

Model-averaging / stacking DML-PLR estimator.

Parameters:

Name	Type	Description	Default
data	DataFrame		required
y	str	Outcome column.	required
treat	str	Continuous-or-binary treatment column.	required
covariates	list of str	Covariate columns X.	required
candidates	list of (ml_g, ml_m, label)	Candidate nuisance learners. ml_g regresses y on X; ml_m regresses treat on X. Defaults to a Lasso/Ridge/ RandomForest/GradientBoosting roster.	None
n_folds	int	Cross-fitting folds per candidate.	5
seed	int		0
weight_rule	('short_stacking', 'single_best', 'inverse_risk', 'equal')	How to combine candidate nuisance predictions or estimates. "short_stacking" (default; Ahrens et al. 2025 eq. 7) — solve constrained least squares on cross-fitted predictions for each nuisance separately (ŷ and D̂), produce stacked nuisances, plug into the PLR moment equation. "single_best" — Ahrens et al. (2025, fn. 8): pick the candidate with lowest joint nuisance MSE. "inverse_risk" — :math:w_k \propto 1/(\text{MSE}_g + \text{MSE}_m). Convenience baseline; not in the paper. "equal" — :math:w_k = 1/K. Convenience baseline; not in the paper. For the non-stacking rules (inverse_risk / equal / single_best) the function computes per-candidate :math:\hat\theta_k and reports the weighted average with a between-candidate-covariance-corrected SE; for "short_stacking" it reports the standard PLR sandwich SE on the stacked-nuisance score (Neyman orthogonality is preserved).	"short_stacking"
alpha	float	Two-sided CI level.	0.05
sample_weight	ndarray | Series | str	Per-observation weights. If supplied, every nuisance fit uses sample_weight= (with a graceful fallback warning if the learner does not accept it), the CLS stacking objective becomes weighted least squares, and the PLR moment + sandwich variance use weighted sums. The MSE used for inverse_risk / single_best weighting is also the weighted MSE.	None
fold_indices	array - like or str	Explicit cross-fitting fold per row (column name or array of length len(data)). Overrides n_folds / seed; every candidate uses the same partition.	None

Returns:

Type	Description
DMLAveragingResult	With the weighted :math:\hat\theta, SE, CI, and per-candidate diagnostics under result.model_info.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> from sklearn.linear_model import LassoCV, RidgeCV
>>> rng = np.random.default_rng(0)
>>> n = 300
>>> X = rng.normal(size=(n, 5))
>>> d = X[:, 0] + rng.normal(size=n)
>>> y = 1.0 * d + X[:, 0] + 0.5 * X[:, 1] + rng.normal(size=n)
>>> df = pd.DataFrame({"y": y, "d": d,
...                    **{f"x{j}": X[:, j] for j in range(5)}})
>>> candidates = [(LassoCV(), LassoCV(), "lasso"),
...               (RidgeCV(), RidgeCV(), "ridge")]
>>> r = sp.dml_model_averaging(df, y="y", treat="d",
...                            covariates=[f"x{j}" for j in range(5)],
...                            candidates=candidates, n_folds=3)
>>> type(r).__name__
'DMLAveragingResult'
>>> sorted(r.model_info["weights_g"])  # CLS stacking weights for E[Y|X]
['lasso', 'ridge']
>>> sorted(r.model_info["weights_m"])  # CLS stacking weights for E[D|X]
['lasso', 'ridge']
>>> print(r.summary())

dml_panel

dml_panel(data: DataFrame, y: str, treat: str, covariates: List[str], *, unit: str, time: Optional[str] = None, ml_g: Optional[Any] = None, ml_m: Optional[Any] = None, n_folds: int = 5, alpha: float = 0.05, include_time_fe: bool = False, binary_treatment: bool = False, seed: int = 0, sample_weight: Optional[Any] = None, fold_indices: Optional[Any] = None) -> DMLPanelResult

Long-panel Double/Debiased ML with unit FE and cluster-robust SE.

Estimates β in

.. math::

Y_{it} = \alpha_i + \lambda_t + \beta D_{it} + g(X_{it}) + \varepsilon_{it}

by (1) within-transforming y, treat and covariates to absorb unit (and optionally time) fixed effects; (2) running cross-fit PLR on the demeaned data with folds that split units; (3) computing the Neyman-orthogonal score and cluster-robust SE at the unit level.

Parameters:

Name	Type	Description	Default
data	DataFrame	Long-format panel; must contain y, treat, all covariates, unit, and time (if given).	required
y	str	Column names.	required
treat	str	Column names.	required
covariates	list of str	High-dimensional controls X_it. Pass [] to fit a pure FE model with ML-free residuals.	required
unit	str	Unit identifier column. Used both for FE absorption and cross-fit fold assignment.	required
time	str	Time identifier column; required when include_time_fe=True.	None
ml_g	sklearn-style estimators	Nuisance learners. Default: GradientBoosting with 200 trees, depth 3, lr 0.05 — same convention as :func:sp.dml PLR.	None
ml_m	sklearn-style estimators	Nuisance learners. Default: GradientBoosting with 200 trees, depth 3, lr 0.05 — same convention as :func:sp.dml PLR.	None
n_folds	int	Cross-fit folds over units. Must be >= 2 and <= n_units.	5
alpha	float		0.05
include_time_fe	bool	If True, also subtract time means (two-way within transform).	False
binary_treatment	bool	Deprecated. Previously routed binary D through a classifier path that mixed within-demeaned features with raw {0,1} labels — the resulting propensity had no clean interpretation as :math:E[\tilde D_{it} \mid \tilde X_{it}]. The flag is now ignored: :func:dml_panel always residualises the within-transformed D̃ with a regressor (PLR is agnostic to D's type, and within-transformed binary D is no longer binary). For DR-style ATE on binary D in panels, use :func:sp.dml(..., model='irm') with unit dummies in the covariate set, :func:sp.etwfe, or :func:sp.callaway_santanna. A :class:DeprecationWarning is emitted when this argument is passed; passing D ∈ {0,1} while binary_treatment=True is validated to catch incidental misuse.	False
seed	int	RNG seed for fold assignment.	0
sample_weight	ndarray | Series | str	Per-observation weights (e.g., survey/probability weights). Pass either a 1-D array of length len(data) or a column name in data. The within transform becomes a weighted within transform (subtract weighted unit / time means), and the PLR moment + cluster-robust SE use weighted sums. Required if your survey design carries informative sampling probabilities.	None
fold_indices	array - like or str	Explicit cross-fitting fold per row (a column name or a length- len(data) array), constant within unit. Overrides n_folds and seed; with a shared partition the estimate is reproducible across implementations (e.g. DoubleMLPLR with cluster splits).	None

Returns:

Type	Description
class:`DMLPanelResult`

Notes

The cluster-robust SE follows Liang-Zeger 1986 at the unit level (the coarser of the two dimensions): stacked scores are summed within unit before squaring. This is the appropriate clustering level when shocks at higher frequencies than the unit are plausibly correlated (cf. Cameron-Miller 2015 §3.2).

Identification requires no unobserved time-varying confounders — only time-invariant unit heterogeneity + high-dim observed X_it. Violations of strict exogeneity of D (e.g. dynamic-feedback) are not handled here; use :func:sp.msm or :func:sp.gformula_ice.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> rows = []
>>> for i in range(50):
...     alpha_i = rng.normal()  # time-invariant unit heterogeneity
...     for t in range(5):
...         x1, x2 = rng.normal(), rng.normal()
...         union = 0.4 * x1 + alpha_i + rng.normal()
...         lw = (0.8 * union + x1 + 0.5 * x2 + alpha_i
...               + 0.1 * t + rng.normal())
...         rows.append({'pid': i, 'year': t, 'log_wage': lw,
...                      'union': union, 'exper': x1, 'educ': x2})
>>> df = pd.DataFrame(rows)
>>> res = sp.dml_panel(
...     df, y='log_wage', treat='union',
...     covariates=['exper', 'educ'],
...     unit='pid', time='year', include_time_fe=True, n_folds=3,
... )
>>> res.method
'dml_panel'
>>> print(res.summary())