ReferenceEffect estimation

Effect estimation

Causal inference: estimate causal effects from data and a causal model. Import from ergodic.inference.

An estimator factors into independent pieces: the estimand family (from identification), the strategy that turns it into a number (a CausalEstimator), and the learner in each nuisance slot. estimate_effect runs identification then estimation and returns an EffectEstimate.

>>> from ergodic import dag, tabular
>>> from ergodic.inference import estimate_effect
>>> est = estimate_effect(graph, data, "T", "Y")          # doctest: +SKIP
>>> est.value, est.interval()                              # doctest: +SKIP

The learner is structurally a sklearn estimator, so estimate_effect(..., learner="knn") or a slot mapping like learner={"outcome": GradientBoostingRegressor()} drops any sklearn or xgboost model into the slots with no adapter. The result is distribution-shaped: a point estimate with an analytic standard error, or bootstrap draws when there is no closed form.

The quasi-experimental designs read panels and time series instead of a graph: did() and event_study() over a PanelData (parallel trends), synthetic_control() (a convex donor combination), and interrupted_time_series() over a TimeSeriesData (a projected pre-trend). Each states its identification in a Design record and returns the same EffectEstimate surface.

Also exported from ergodic: CATEResult, CausalForest, EffectEstimate, EffectTarget, ITEResult, MediationEstimate, PerUnitEffect, PolicyTree, Refutation, SelectionReport, VariableImportance, cate, did, estimate, estimate_effect, estimate_mediation, event_study, interrupted_time_series, policy_tree, refute_effect, select_cate, synthetic_control.

AIPW#

class AIPW(
    learners: Any = None,
    *,
    clip: float = 0.001,
    bayesian_bootstrap: bool = True,
    **kwargs: Any,
) -> None

Combine an outcome model and a propensity model.

Consistent if either model is right. The ATE standard error comes from the efficient influence function (Neyman-orthogonal, so first-order insensitive to nuisance error); the ATT standard error falls back to bootstrap.

AIPW.fit#

def AIPW.fit(identification: Identification, data: TabularData) -> Self

Read the estimand and the data, resolve learners, store the fit state.

AIPW.effect#

def AIPW.effect(
    *,
    target: EffectTarget | str | None = None,
) -> EffectEstimate

Report the effect, with an analytic standard error or a bootstrap one.

DML#

class DML(learners: Any = None, *, n_folds: int = 5, **kwargs: Any) -> None

Double machine learning, the partially linear model (partialling-out).

Residualize the outcome and the treatment against the covariates with cross-fitting, then regress one residual on the other. The treatment may be binary or continuous; a continuous treatment reports the average partial effect. The ATE standard error comes from the orthogonal moment.

DML.fit#

def DML.fit(identification: Identification, data: TabularData) -> Self

Read the estimand and the data, resolve learners, store the fit state.

DML.cate#

def DML.cate(
    data: TabularData | None = None,
    *,
    uncertainty: bool = True,
) -> CATEResult

The constant per-unit effect for the partially-linear model: theta.

The partially-linear DML assumes a homogeneous effect, so every unit's CATE is the single slope theta, with the same analytic standard error from the orthogonal moment. Returned as a CATEResult for a uniform surface; a binary treatment additionally caches the doubly-robust scores so explain can test for the heterogeneity the model assumes away. uncertainty=False returns the point theta with kind="point" and no standard error.

DML.effect#

def DML.effect(
    *,
    target: EffectTarget | str | None = None,
) -> EffectEstimate

Report the effect, with an analytic standard error or a bootstrap one.

IPW#

class IPW(
    learners: Any = None,
    *,
    scheme: str = 'hajek',
    clip: float = 0.001,
    **kwargs: Any,
) -> None

Weight by the inverse propensity of treatment.

scheme="hajek" (the default) self-normalizes the weights; scheme="ht" is the raw Horvitz-Thompson estimator. Propensities are clipped to [clip, 1 - clip] to bound the weights.

IPW.fit#

def IPW.fit(identification: Identification, data: TabularData) -> Self

Read the estimand and the data, resolve learners, store the fit state.

IPW.effect#

def IPW.effect(
    *,
    target: EffectTarget | str | None = None,
) -> EffectEstimate

Report the effect, with an analytic standard error or a bootstrap one.

OLS#

class OLS(*, fit_intercept: bool = True, ridge: float = 0.0) -> None

Ordinary least squares, with optional ridge penalty and sample weights.

OLS.fit#

def OLS.fit(
    X: npt.NDArray[np.float64],
    y: npt.NDArray[np.float64],
    sample_weight: npt.NDArray[np.float64] | None = None,
) -> OLS

Fit the coefficients by (weighted, optionally ridge-penalized) least squares.

The intercept is never penalized.

Returns: The fitted learner itself.

OLS.predict#

def OLS.predict(X: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]

Return the fitted linear prediction for the rows of X.

OLS.get_params#

def OLS.get_params(deep: bool = True) -> dict[str, Any]

Return the constructor parameters, fit_intercept and ridge.

OLS.set_params#

def OLS.set_params(**params: Any) -> OLS

Set parameters by name and return the learner itself.

Raises ValueError on a name get_params() does not list.

PLIV#

class PLIV(learners: Any = None, *, n_folds: int = 5, **kwargs: Any) -> None

Orthogonal instrumental variables (the partially linear IV model).

Residualize the outcome, treatment, and instrument against the covariates with cross-fitting, then the IV moment on the residuals. The flexible-learner instrument estimator that two-stage least squares is not. Reports LATE with an analytic standard error.

PLIV.fit#

def PLIV.fit(identification: Identification, data: TabularData) -> Self

Read the estimand and the data, resolve learners, store the fit state.

PLIV.effect#

def PLIV.effect(
    *,
    target: EffectTarget | str | None = None,
) -> EffectEstimate

Report the effect, with an analytic standard error or a bootstrap one.

BayesianLearner#

class BayesianLearner(*args, **kwargs)

A learner that also carries a posterior over its mean function.

Structurally a Learner (fit then predict, where predict is the posterior predictive mean), plus predict_draws returning (S, n) draws of the mean function. A classifier variant exposes predict_proba and predict_proba_draws (draws of the class-1 probability) instead. Detection is duck-typed in _is_bayesian(), the same way a classifier is spotted by predict_proba, so any object that quacks this way counts.

Because predict stays the posterior mean, a Bayesian learner is also a valid point learner in every estimator; the posterior path is an addition, not a fork.

BayesianLearner.fit#

def BayesianLearner.fit(
    X: npt.NDArray[np.float64],
    y: npt.NDArray[np.float64],
    sample_weight: npt.NDArray[np.float64] | None = None,
) -> Any

Fit the learner to features X and target y.

BayesianLearner.predict#

def BayesianLearner.predict(X: npt.NDArray[np.float64]) -> Any

Return the posterior predictive mean for the rows of X.

BayesianLearner.predict_draws#

def BayesianLearner.predict_draws(X: npt.NDArray[np.float64]) -> Any

Return (S, n) posterior draws of the mean function at the rows of X.

CATEResult#

class CATEResult(
    treatment: str,
    outcome: str,
    covariates: tuple[str, ...],
    n: int,
    kind: str,
    estimand: str = 'cate',
    source: str = '',
    diagnostics: Mapping[str, float] = <factory>,
    _values: npt.NDArray[np.float64] = <factory>,
    _index: npt.NDArray[Any] | None = None,
    _std_error: npt.NDArray[np.float64] | None = None,
    _draws: npt.NDArray[np.float64] | None = None,
    identification: Identification | None = None,
    _scores: DRScores | None = None,
) -> None

Conditional average treatment effects, with an explainability hook.

A conditional average over the unobserved noise (the inference-pillar object). Carries the Identification it was fit under and the cached doubly-robust scores + covariate design, so explain() projects honest summaries (BLP, GATES, CLAN, the calibration test, RATE) off the same orthogonal signal.

CATEResult.explain#

def CATEResult.explain() -> Heterogeneity

Tier-B characterization of the heterogeneity, on the doubly-robust scores.

Builds a Heterogeneity record whose methods (blp, gates, clan, calibration_test, rate) all consume the cached cross-fit scores, the covariate design, and the fitted per-unit tau_hat. Needs the scores; raises when they were not cached (an estimator that does not build them, or a result without them).

CATEResult.values#

property CATEResult.values: npt.NDArray[np.float64]

The per-unit point estimates, (n,), ordered to the input rows.

When only draws are present, the posterior or bootstrap mean per unit.

CATEResult.shape#

property CATEResult.shape: tuple[int, ...]

The shape of the per-unit point vector, (n,).

CATEResult.index#

property CATEResult.index: npt.NDArray[Any]

The row ids for the units, positional arange(n) when none were given.

CATEResult.std_error#

property CATEResult.std_error: npt.NDArray[np.float64] | None

Per-unit standard errors, or the per-unit draw standard deviations.

CATEResult.interval#

def CATEResult.interval(
    level: float = 0.9,
) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]

Per-unit confidence or credible intervals, (lo[n], hi[n]) at level.

Draw quantiles when draws are present, else the normal interval from the per-unit point and standard error. Raises when there is no uncertainty.

CATEResult.draws#

def CATEResult.draws() -> npt.NDArray[np.float64] | None

The (S, n) bootstrap or posterior draws, or None for a point result.

CATEResult.ate#

def CATEResult.ate() -> EffectEstimate

The population summary: the average per-unit effect, as an EffectEstimate.

The bridge back to the scalar family. With draws, the per-unit draws average to a draw vector of the mean effect (so the population interval reflects the joint uncertainty). With neither draws nor standard errors, a point.

Per-unit analytic standard errors cannot be combined into the mean's standard error: the units share one fitted model, so their errors are correlated, and treating them as independent understates the uncertainty by about a factor of sqrt(n). So, with analytic standard errors:

• a homogeneous result (one effect and one standard error repeated over the units, as DML.cate() returns) reports that effect and its standard error;
• a result carrying its cross-fit doubly-robust scores reports the AIPW average off them: the mean score with standard error std(scores) / sqrt(n) (the same scores explain() reads). For the DR-learner's linear final stage the point equals the mean per-unit effect; for a non-linear final stage such as the causal forest it can differ slightly from the plain mean of the per-unit effects (diagnostics["ate_from_dr_scores"] marks this path);
• otherwise it reports the mean with no standard error (kind="point"). For an interval, use estimate_effect().

CATEResult.to_frame#

def CATEResult.to_frame() -> pd.DataFrame

A per-unit table (index, effect, and lo/hi/se when available). Pandas only here.

CausalEstimator#

class CausalEstimator(
    learners: Any = None,
    *,
    random_state: int | None = None,
    n_boot: int = 1000,
    treated: Any | None = None,
) -> None

Base for the strategies. Stateful: fit then effect.

CausalEstimator.fit#

def CausalEstimator.fit(
    identification: Identification,
    data: TabularData,
) -> Self

Read the estimand and the data, resolve learners, store the fit state.

CausalEstimator.effect#

def CausalEstimator.effect(
    *,
    target: EffectTarget | str | None = None,
) -> EffectEstimate

Report the effect, with an analytic standard error or a bootstrap one.

CausalForest#

class CausalForest(
    learners: Any = None,
    *,
    n_trees: int = 200,
    min_leaf: int = 5,
    max_depth: int = 8,
    honest: bool = True,
    subsample: float = 0.5,
    max_features: float | int | None = None,
    alpha: float = 0.05,
    n_folds: int = 5,
    clip: float = 0.001,
    **kwargs: Any,
) -> None

An honest causal forest (GRF) for the conditional average treatment effect.

The forest is the CATE learner. The outcome and propensity nuisance slots take any learner (they residualize Y and W on the covariates before the forest grows); the forest machinery, honest splitting and the kernel readout, is native. The ATE is the average per-unit effect, and cate returns a CATEResult with per-unit infinitesimal-jackknife standard errors and intervals.

Hyperparameters: n_trees (forest size), min_leaf (the minimum estimate-half rows per leaf), max_depth, honest (split/estimate sample-splitting, on by default), subsample (the per-tree subsample fraction), and random_state. Binary treatment only; a continuous treatment via the same local moment is not supported and raises.

The per-unit intervals need many trees. The infinitesimal-jackknife variance subtracts a Monte-Carlo correction that shrinks only as n_trees grows, and where the correction exceeds the raw variance the standard error is clipped to zero (a zero-width band). Measured on a 4,000-row world with a linear effect (min_leaf=15, nominal 90%): 400 trees covered the true effect for 68% of units with 25% zero-width bands; 1,600 trees covered 87% with 1% zero-width. An independent check (1,500 rows) still found 7% zero-width bands at 750 trees and none at 1,500. So use at least as many trees as rows when the per-unit intervals matter; more trees reduce the clipping but not the smoothing bias below. cate warns, and records diagnostics["ij_zero_se_share"], when more than 5% of the bands are zero-width. Coverage stays somewhat below nominal where the effect changes steeply (the forest's smoothing bias), so treat the bands as approximate.

CausalForest.cate#

def CausalForest.cate(
    data: TabularData | None = None,
    *,
    uncertainty: bool = True,
) -> CATEResult

Per-unit CATE with infinitesimal-jackknife standard errors and intervals.

Overrides the base meta-learner cate (which would bootstrap-refit the whole forest). The point is the forest-kernel local effect; the per-unit variance is the infinitesimal jackknife over the trees with the Wager-Hastie-Efron Monte-Carlo correction, valid because the forest is honest and subsampled. The result carries the cached doubly-robust scores for explain().

uncertainty=False returns the point effects with kind="point" and no standard error, skipping the infinitesimal-jackknife pass over the trees.

CausalForest.variable_importance#

def CausalForest.variable_importance() -> VariableImportance

Depth-weighted, normalized split frequencies over the covariates.

A driver signal: which covariates the forest splits on, weighted toward the shallow splits that move the effect most. Available after fit.

CausalForest.fit#

def CausalForest.fit(
    identification: Identification,
    data: TabularData,
) -> Self

Read the estimand and the data, resolve learners, store the fit state.

CausalForest.effect#

def CausalForest.effect(
    *,
    target: EffectTarget | str | None = None,
) -> EffectEstimate

Report the effect, with an analytic standard error or a bootstrap one.

DRLearner#

class DRLearner(
    learners: Any = None,
    *,
    n_folds: int = 5,
    clip: float = 0.001,
    **kwargs: Any,
) -> None

Doubly-robust learning of the ATE and the CATE.

Cross-fit the AIPW score into a per-unit pseudo-outcome, average it for the ATE (with the influence-function standard error), and regress it on the covariates for the CATE through cate. The CATE result carries the cached doubly-robust scores, so its explain runs the Tier-B characterization.

DRLearner.fit#

def DRLearner.fit(identification: Identification, data: TabularData) -> Self

Read the estimand and the data, resolve learners, store the fit state.

DRLearner.cate#

def DRLearner.cate(
    data: TabularData | None = None,
    *,
    uncertainty: bool = True,
) -> CATEResult

Per-unit CATE as a CATEResult.

In-sample when data is None, else scored on its covariates. The result carries the per-unit point, the uncertainty (posterior draws for a Bayesian slot, the bootstrap draws otherwise, the analytic standard error for the linear DR final), the identification, and the cached doubly-robust scores for explain().

Pass uncertainty=False for a fast point-only result (kind="point", no standard error or draws): it skips the uncertainty step, which for a non-linear final means a full bootstrap refit of the cross-fit pipeline. Use it when only the effect curve is needed (ranking, targeting, plots); the per-unit points are identical either way. explain still works (it reads the cached scores, not the draws).

DRLearner.effect#

def DRLearner.effect(
    *,
    target: EffectTarget | str | None = None,
) -> EffectEstimate

Report the effect, with an analytic standard error or a bootstrap one.

DRScores#

class DRScores(
    scores: npt.NDArray[np.float64],
    mu0: npt.NDArray[np.float64],
    mu1: npt.NDArray[np.float64],
    propensity: npt.NDArray[np.float64],
    covariates: npt.NDArray[np.float64],
    treatment: npt.NDArray[np.float64],
    outcome: npt.NDArray[np.float64],
    covariate_names: tuple[str, ...],
    tau_hat: npt.NDArray[np.float64] | None = None,
    covariate_variables: tuple[str, ...] = (),
) -> None

The cross-fit doubly-robust object: the AIPW signal plus its nuisances.

scores is the per-unit Gamma_i (its mean is the AIPW ATE). mu0, mu1, and propensity are the cross-fit out-of-fold nuisances, kept so the calibration and projection tools reuse them without refitting. treatment, outcome, and covariates carry the design back to the result.

DRScores.column_variables#

property DRScores.column_variables: tuple[str, ...]

The variable each design column encodes (the column name when not given).

Design#

class Design(
    strategy: str,
    outcome: str,
    treatment: str | None,
    assumptions: tuple[str, ...],
    estimand: str = '',
) -> None

How a quasi-experimental estimator claims identification.

strategy is one of "did", "synthetic_control", or "interrupted_time_series". assumptions are the untestable claims the design rests on, stated in words (parallel trends, convex donor span, no co-timed shock); they are what a reader should challenge, so they ride with every result.

DiDEffect#

class DiDEffect(
    target: EffectTarget,
    estimator: str,
    learners: Mapping[str, str],
    n: int,
    kind: str,
    diagnostics: Mapping[str, float] = <factory>,
    _point: float | None = None,
    _std_error: float | None = None,
    _draws: npt.NDArray[np.float64] | None = None,
    design: Design | None = None,
    flavor: str = '',
    comparison: str = '',
    by_group_time: tuple[GroupTimeEffect, ...] = (),
    by_event_time: tuple[EventTimeEffect, ...] = (),
) -> None

A difference-in-differences result: the estimate plus the design it rests on.

The EffectEstimate surface (value, std_error, interval, draws) reads the headline ATT. flavor records which estimator ran, comparison which pool stood in for the counterfactual, by_group_time the per-cell detail of the group-time flavor, and by_event_time the event-study aggregation when event_study() built the result.

DiDEffect.value#

property DiDEffect.value: float

The point estimate, or the mean of the draws.

DiDEffect.std_error#

property DiDEffect.std_error: float | None

The analytic standard error, or the standard deviation of the draws.

DiDEffect.interval#

def DiDEffect.interval(level: float = 0.95) -> tuple[float, float]

A confidence or credible interval at level.

From the draw quantiles when draws are present, else the normal interval around the point estimate.

DiDEffect.draws#

property DiDEffect.draws: npt.NDArray[np.float64] | None

The bootstrap or posterior draws, or None for a pure analytic result.

EffectEstimate#

class EffectEstimate(
    target: EffectTarget,
    estimator: str,
    learners: Mapping[str, str],
    n: int,
    kind: str,
    diagnostics: Mapping[str, float] = <factory>,
    _point: float | None = None,
    _std_error: float | None = None,
    _draws: npt.NDArray[np.float64] | None = None,
) -> None

A causal effect with its uncertainty.

Built by a CausalEstimator. Either _point/_std_error are set (an analytic result) or _draws is set (bootstrap or posterior), and the accessors read whichever is present.

EffectEstimate.value#

property EffectEstimate.value: float

The point estimate, or the mean of the draws.

EffectEstimate.std_error#

property EffectEstimate.std_error: float | None

The analytic standard error, or the standard deviation of the draws.

EffectEstimate.interval#

def EffectEstimate.interval(level: float = 0.95) -> tuple[float, float]

A confidence or credible interval at level.

From the draw quantiles when draws are present, else the normal interval around the point estimate.

EffectEstimate.draws#

property EffectEstimate.draws: npt.NDArray[np.float64] | None

The bootstrap or posterior draws, or None for a pure analytic result.

EffectTarget#

class EffectTarget(enum.Enum)

Which effect a CausalEstimator reports.

The estimators return ATE and ATT for the adjustment and front-door families and LATE for instruments; CATE and ITE tag the per-unit results' summaries. ATC is reserved and no estimator returns it.

EventTimeEffect#

class EventTimeEffect(
    event_time: int,
    estimate: float,
    std_error: float | None,
    interval: tuple[float, float],
    n_treated: int,
) -> None

The effect at one event time, aggregated over cohorts with cohort-size weights.

std_error and interval come from the unit bootstrap: the standard deviation of the draws and their 95% percentile interval. n_treated counts the treated units contributing a cell at this event time.

GroupTimeEffect#

class GroupTimeEffect(
    group: Any,
    period: Any,
    event_time: int,
    estimate: float,
    n_treated: int,
    n_comparison: int,
) -> None

One ATT(g, t) cell: the cohort adopting at group, read at period.

event_time is the position of period minus the position of group in the sorted period sequence, so -1 is the base period and never appears. n_treated and n_comparison count the units observed at both the base and the target period on each side of the comparison.

Heterogeneity#

class Heterogeneity(
    scores: DRScores,
    tau_hat: npt.NDArray[np.float64],
    treatment: str,
    outcome: str,
) -> None

Tier-B heterogeneity characterization on shared doubly-robust scores.

Carries the cross-fit DRScores and the fitted per-unit tau_hat. The honesty register: the rankings and proxies are fixed (computed cross-fit), so the orthogonal summaries below inherit valid inference while the CATE point predictions themselves never carry a confidence interval.

Heterogeneity.blp#

def Heterogeneity.blp(*, level: float = 0.9) -> BLPResult

Best Linear Projection of the CATE on the covariates (HC3 inference).

Heterogeneity.gates#

def Heterogeneity.gates(
    *,
    n_groups: int = 5,
    level: float = 0.9,
) -> GATESResult

Sorted Group ATEs across n_groups predicted-effect quantiles.

Heterogeneity.clan#

def Heterogeneity.clan(
    *,
    n_groups: int = 5,
    level: float = 0.9,
) -> CLANResult

Covariate-mean differences between the most- and least-affected groups.

Heterogeneity.calibration_test#

def Heterogeneity.calibration_test() -> CalibrationResult

The omnibus calibration / heterogeneity test (one-sided HTE p-value).

Heterogeneity.rate#

def Heterogeneity.rate(
    metric: str = 'autoc',
    *,
    n_boot: int = 200,
    random_state: int | None = None,
) -> RATEResult

The Rank-Weighted ATE of the tau_hat ranking: "autoc" or "qini".

Heterogeneity.policy_tree#

def Heterogeneity.policy_tree(
    covariates: str | tuple[str, ...] | list[str] | None = None,
    *,
    depth: int = 2,
) -> PolicyTree

The depth-limited, DR-welfare-optimal policy tree off these scores.

A convenience that forwards the shared doubly-robust scores to policy_tree(). covariates selects the split columns by name (None uses all of them). The "act" companion to rate(): where rate scores a ranking, this returns an interpretable treat-or-not rule and its estimated policy value.

ITEResult#

class ITEResult(
    treatment: str,
    outcome: str,
    covariates: tuple[str, ...],
    n: int,
    kind: str,
    estimand: str = 'ite',
    source: str = '',
    diagnostics: Mapping[str, float] = <factory>,
    _values: npt.NDArray[np.float64] = <factory>,
    _index: npt.NDArray[Any] | None = None,
    _std_error: npt.NDArray[np.float64] | None = None,
    _draws: npt.NDArray[np.float64] | None = None,
    framing: str = '',
    note: str = '',
    _factual: npt.NDArray[np.float64] | None = None,
) -> None

Individual treatment effects from a structural model (the twin network).

The individual's own counterfactual contrast (the SCM-pillar object), not a conditional average. Carries the factual outcome column, the contrast framing (levels= or shift=), and a structural provenance note.

explain() is not available: the Tier-B machinery is defined over the CATE's doubly-robust scores, and there is no inference-valid descriptive regression of abducted ITEs, so it raises with a pointer rather than faking inference.

ITEResult.factual#

property ITEResult.factual: npt.NDArray[np.float64] | None

The per-unit factual outcome the contrast was taken against, when carried.

ITEResult.explain#

def ITEResult.explain() -> Heterogeneity

Not available: the Tier-B tools need the CATE's doubly-robust scores.

ITEResult.values#

property ITEResult.values: npt.NDArray[np.float64]

The per-unit point estimates, (n,), ordered to the input rows.

When only draws are present, the posterior or bootstrap mean per unit.

ITEResult.shape#

property ITEResult.shape: tuple[int, ...]

The shape of the per-unit point vector, (n,).

ITEResult.index#

property ITEResult.index: npt.NDArray[Any]

The row ids for the units, positional arange(n) when none were given.

ITEResult.std_error#

property ITEResult.std_error: npt.NDArray[np.float64] | None

Per-unit standard errors, or the per-unit draw standard deviations.

ITEResult.interval#

def ITEResult.interval(
    level: float = 0.9,
) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]

Per-unit confidence or credible intervals, (lo[n], hi[n]) at level.

Draw quantiles when draws are present, else the normal interval from the per-unit point and standard error. Raises when there is no uncertainty.

ITEResult.draws#

def ITEResult.draws() -> npt.NDArray[np.float64] | None

The (S, n) bootstrap or posterior draws, or None for a point result.

ITEResult.ate#

def ITEResult.ate() -> EffectEstimate

The population summary: the average per-unit effect, as an EffectEstimate.

The bridge back to the scalar family. With draws, the per-unit draws average to a draw vector of the mean effect (so the population interval reflects the joint uncertainty). With neither draws nor standard errors, a point.

Per-unit analytic standard errors cannot be combined into the mean's standard error: the units share one fitted model, so their errors are correlated, and treating them as independent understates the uncertainty by about a factor of sqrt(n). So, with analytic standard errors:

• a homogeneous result (one effect and one standard error repeated over the units, as DML.cate() returns) reports that effect and its standard error;
• a result carrying its cross-fit doubly-robust scores reports the AIPW average off them: the mean score with standard error std(scores) / sqrt(n) (the same scores explain() reads). For the DR-learner's linear final stage the point equals the mean per-unit effect; for a non-linear final stage such as the causal forest it can differ slightly from the plain mean of the per-unit effects (diagnostics["ate_from_dr_scores"] marks this path);
• otherwise it reports the mean with no standard error (kind="point"). For an interval, use estimate_effect().

ITEResult.to_frame#

def ITEResult.to_frame() -> pd.DataFrame

A per-unit table (index, effect, and lo/hi/se when available). Pandas only here.

ITSEffect#

class ITSEffect(
    target: EffectTarget,
    estimator: str,
    learners: Mapping[str, str],
    n: int,
    kind: str,
    diagnostics: Mapping[str, float] = <factory>,
    _point: float | None = None,
    _std_error: float | None = None,
    _draws: npt.NDArray[np.float64] | None = None,
    design: Design | None = None,
    at: Any = None,
    level_change: float = nan,
    level_change_se: float = nan,
    slope_change: float = nan,
    slope_change_se: float = nan,
    times: tuple[Any, ...] = (),
    observed: tuple[float, ...] = (),
    counterfactual: tuple[float, ...] = (),
    fitted: tuple[float, ...] = (),
) -> None

A segmented-regression fit around one break.

level_change is the jump at the break and slope_change the change in trend per period, each with a Newey-West standard error. value is the ATT: the average post-period gap between fitted and counterfactual, where counterfactual projects the pre-period trend over every period with the break terms zeroed and the seasonal terms kept.

ITSEffect.value#

property ITSEffect.value: float

The point estimate, or the mean of the draws.

ITSEffect.std_error#

property ITSEffect.std_error: float | None

The analytic standard error, or the standard deviation of the draws.

ITSEffect.interval#

def ITSEffect.interval(level: float = 0.95) -> tuple[float, float]

A confidence or credible interval at level.

From the draw quantiles when draws are present, else the normal interval around the point estimate.

ITSEffect.draws#

property ITSEffect.draws: npt.NDArray[np.float64] | None

The bootstrap or posterior draws, or None for a pure analytic result.

KNNClassifier#

class KNNClassifier(*, k: int = 5) -> None

K-nearest-neighbors classification for a binary 0/1 label.

KNNClassifier.fit#

def KNNClassifier.fit(
    X: npt.NDArray[np.float64],
    y: npt.NDArray[np.float64],
    sample_weight: npt.NDArray[np.float64] | None = None,
) -> KNNClassifier

Store the standardized training rows, their 0/1 labels, and their weights.

sample_weight (finite, non-negative, one per row) weights each row in the neighbor vote; without it every row counts once.

Returns: The fitted learner itself.

KNNClassifier.predict_proba#

def KNNClassifier.predict_proba(
    X: npt.NDArray[np.float64],
) -> npt.NDArray[np.float64]

Return class probabilities as an (n, 2) array.

Column 1 is the (weighted) share of 1s among the k nearest training rows. Raises RuntimeError before fit().

KNNClassifier.predict#

def KNNClassifier.predict(
    X: npt.NDArray[np.float64],
) -> npt.NDArray[np.float64]

Return the 0/1 label, 1 where the neighbor share of 1s is at least 0.5.

KNNClassifier.get_params#

def KNNClassifier.get_params(deep: bool = True) -> dict[str, Any]

Return the constructor parameter k.

KNNClassifier.set_params#

def KNNClassifier.set_params(**params: Any) -> _KNNBase

Set parameters by name and return the learner itself.

Raises ValueError on a name get_params() does not list.

KNNRegressor#

class KNNRegressor(*, k: int = 5) -> None

K-nearest-neighbors regression: the prediction is the neighbor mean.

KNNRegressor.fit#

def KNNRegressor.fit(
    X: npt.NDArray[np.float64],
    y: npt.NDArray[np.float64],
    sample_weight: npt.NDArray[np.float64] | None = None,
) -> KNNRegressor

Store the standardized training rows, their targets, and their weights.

sample_weight (finite, non-negative, one per row) weights each row in the neighbor average; without it every row counts once.

Returns: The fitted learner itself.

KNNRegressor.predict#

def KNNRegressor.predict(
    X: npt.NDArray[np.float64],
) -> npt.NDArray[np.float64]

Return the (weighted) mean target of the k nearest training rows, per row.

Raises RuntimeError before fit().

KNNRegressor.get_params#

def KNNRegressor.get_params(deep: bool = True) -> dict[str, Any]

Return the constructor parameter k.

KNNRegressor.set_params#

def KNNRegressor.set_params(**params: Any) -> _KNNBase

Set parameters by name and return the learner itself.

Raises ValueError on a name get_params() does not list.

Learner#

class Learner(*args, **kwargs)

The sklearn estimator surface the estimators rely on.

Learner.fit#

def Learner.fit(
    X: npt.NDArray[np.float64],
    y: npt.NDArray[np.float64],
    sample_weight: npt.NDArray[np.float64] | None = None,
) -> Any

Fit the learner to features X and target y.

Learner.predict#

def Learner.predict(X: npt.NDArray[np.float64]) -> Any

Return predictions for the rows of X.

LogisticRegression#

class LogisticRegression(
    *,
    fit_intercept: bool = True,
    l2: float = 1e-06,
    max_iter: int = 100,
    tol: float = 1e-08,
) -> None

Binary logistic regression by Newton-Raphson (IRLS), ridge-stabilized.

LogisticRegression.fit#

def LogisticRegression.fit(
    X: npt.NDArray[np.float64],
    y: npt.NDArray[np.float64],
    sample_weight: npt.NDArray[np.float64] | None = None,
) -> LogisticRegression

Fit the coefficients by Newton-Raphson with an L2 ridge on the slopes.

Iteration stops when the largest step falls under tol or after max_iter steps; stopping on max_iter warns with ConvergenceWarning.

Returns: The fitted learner itself.

LogisticRegression.predict_proba#

def LogisticRegression.predict_proba(
    X: npt.NDArray[np.float64],
) -> npt.NDArray[np.float64]

Return class probabilities as an (n, 2) array: column 1 is P(y = 1).

LogisticRegression.predict#

def LogisticRegression.predict(
    X: npt.NDArray[np.float64],
) -> npt.NDArray[np.float64]

Return the 0/1 label, 1 where P(y = 1) is at least 0.5.

LogisticRegression.get_params#

def LogisticRegression.get_params(deep: bool = True) -> dict[str, Any]

Return the constructor parameters: fit_intercept, l2, max_iter, tol.

LogisticRegression.set_params#

def LogisticRegression.set_params(**params: Any) -> LogisticRegression

Set parameters by name and return the learner itself.

Raises ValueError on a name get_params() does not list.

MediationEstimate#

class MediationEstimate(
    nde: EffectEstimate,
    nie: EffectEstimate,
    total: EffectEstimate,
    cde: dict[Any, EffectEstimate],
    design: Design,
    mediators: frozenset[str],
    covariates: frozenset[str],
) -> None

A mediation decomposition with its uncertainty.

nde, nie, and total are the (pure) natural direct effect, the (pure) natural indirect effect, and the total effect E[Y(1, M(1))] - E[Y(0, M(0))], each a distribution-shaped EffectEstimate with a bootstrap standard error and interval. The total equals nde + nie when the treatment does not interact with the mediator in the outcome; under an interaction the two parts do not add up to it, and total - nde - nie is the mediated interaction. cde maps each mediator level to the controlled direct effect at that level (also an EffectEstimate). design names the (untestable) assumptions the decomposition rests on, the same role Design plays for the quasi-experimental estimators.

proportion_mediated is nie / total, the pure indirect effect's share of the total effect (under an interaction it and nde / total need not sum to one). It is unstable when the total effect is near zero (a small denominator), so it is offered with that caveat and a flagged diagnostic.

MediationEstimate.proportion_mediated#

property MediationEstimate.proportion_mediated: float

NIE / TE, with the pure NIE: the share explained by mediation.

Pearl's (2001; 2012, "The Mediation Formula") fraction of the response explained by the mediator. The other common convention, VanderWeele's total NIE over TE (the share owed to mediation, which CMAverse and paramed report), is 1 - nde / total. The two agree without a treatment-mediator interaction and differ by mediated_interaction / total with one (on Berkeley admissions, 1.01 against 1.50). Unstable when the total effect is near zero; see proportion_mediated_flagged.

MediationEstimate.mediated_interaction#

property MediationEstimate.mediated_interaction: float

total - nde - nie: the part of the total effect the pure effects miss.

VanderWeele's mediated interaction (INTmed): zero without a treatment-by-mediator interaction in the outcome. A point value, from the three estimates.

MediationEstimate.proportion_mediated_flagged#

property MediationEstimate.proportion_mediated_flagged: bool

True when the total effect is too small to trust the proportion.

The proportion mediated has the total effect in its denominator, so a total effect close to zero makes it numerically unstable: a tiny shift in either piece swings the ratio. Flagged when |TE| is below one bootstrap standard error (or essentially zero).

NullEffect#

class NullEffect(
    learners: Any = None,
    *,
    random_state: int | None = None,
    n_boot: int = 1000,
    treated: Any | None = None,
) -> None

Report the exact zero effect for a null estimand.

Identification returns strategy="null" when the graph has no directed path from the treatment to the outcome, so the intervention cannot move the outcome and the effect is zero by construction. No nuisance slots, no uncertainty: the point is 0.0 with a zero standard error.

NullEffect.fit#

def NullEffect.fit(
    identification: Identification,
    data: TabularData,
) -> Self

Read the estimand and the data, resolve learners, store the fit state.

NullEffect.effect#

def NullEffect.effect(
    *,
    target: EffectTarget | str | None = None,
) -> EffectEstimate

Report the effect, with an analytic standard error or a bootstrap one.

OutcomeRegression#

class OutcomeRegression(
    learners: Any = None,
    *,
    per_arm: bool = False,
    **kwargs: Any,
) -> None

Fit a model for the outcome and average the predicted contrast.

With per_arm=False (the default) one model carries the treatment as a feature (the S-learner / g-formula). With per_arm=True a separate model is fit on each arm (the T-learner). A Bayesian outcome learner flows its posterior into the per-unit CATE result and into cate_draws/cate_interval.

OutcomeRegression.fit#

def OutcomeRegression.fit(
    identification: Identification,
    data: TabularData,
) -> Self

Read the estimand and the data, resolve learners, store the fit state.

OutcomeRegression.cate_draws#

def OutcomeRegression.cate_draws(
    data: TabularData | None = None,
) -> npt.NDArray[np.float64]

Per-unit CATE posterior (mu1 - mu0), shape (S, n).

Needs a Bayesian outcome learner; in-sample when data is None.

OutcomeRegression.cate_interval#

def OutcomeRegression.cate_interval(
    data: TabularData | None = None,
    *,
    level: float = 0.95,
) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]

Per-unit credible interval for the CATE: the column quantiles of the draws.

OutcomeRegression.cate#

def OutcomeRegression.cate(
    data: TabularData | None = None,
    *,
    uncertainty: bool = True,
) -> CATEResult

Per-unit CATE as a CATEResult.

In-sample when data is None, else scored on its covariates. The result carries the per-unit point, the uncertainty (posterior draws for a Bayesian slot, the bootstrap draws otherwise, the analytic standard error for the linear DR final), the identification, and the cached doubly-robust scores for explain().

Pass uncertainty=False for a fast point-only result (kind="point", no standard error or draws): it skips the uncertainty step, which for a non-linear final means a full bootstrap refit of the cross-fit pipeline. Use it when only the effect curve is needed (ranking, targeting, plots); the per-unit points are identical either way. explain still works (it reads the cached scores, not the draws).

OutcomeRegression.effect#

def OutcomeRegression.effect(
    *,
    target: EffectTarget | str | None = None,
) -> EffectEstimate

Report the effect, with an analytic standard error or a bootstrap one.

PerUnitEffect#

class PerUnitEffect(
    treatment: str,
    outcome: str,
    covariates: tuple[str, ...],
    n: int,
    kind: str,
    estimand: str = '',
    source: str = '',
    diagnostics: Mapping[str, float] = <factory>,
    _values: npt.NDArray[np.float64] = <factory>,
    _index: npt.NDArray[Any] | None = None,
    _std_error: npt.NDArray[np.float64] | None = None,
    _draws: npt.NDArray[np.float64] | None = None,
) -> None

Per-unit treatment effects with their uncertainty.

_values is the per-unit point vector (n,), ordered to the input rows. Either _std_error (per-unit analytic SEs) or _draws (S, n) (bootstrap or posterior) carries the uncertainty, and the accessors read whichever is present. index are explicit row ids; estimand is set by the subclass.

PerUnitEffect.values#

property PerUnitEffect.values: npt.NDArray[np.float64]

The per-unit point estimates, (n,), ordered to the input rows.

When only draws are present, the posterior or bootstrap mean per unit.

PerUnitEffect.shape#

property PerUnitEffect.shape: tuple[int, ...]

The shape of the per-unit point vector, (n,).

PerUnitEffect.index#

property PerUnitEffect.index: npt.NDArray[Any]

The row ids for the units, positional arange(n) when none were given.

PerUnitEffect.std_error#

property PerUnitEffect.std_error: npt.NDArray[np.float64] | None

Per-unit standard errors, or the per-unit draw standard deviations.

PerUnitEffect.interval#

def PerUnitEffect.interval(
    level: float = 0.9,
) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]

Per-unit confidence or credible intervals, (lo[n], hi[n]) at level.

Draw quantiles when draws are present, else the normal interval from the per-unit point and standard error. Raises when there is no uncertainty.

PerUnitEffect.draws#

def PerUnitEffect.draws() -> npt.NDArray[np.float64] | None

The (S, n) bootstrap or posterior draws, or None for a point result.

PerUnitEffect.ate#

def PerUnitEffect.ate() -> EffectEstimate

The population summary: the average per-unit effect, as an EffectEstimate.

The bridge back to the scalar family. With draws, the per-unit draws average to a draw vector of the mean effect (so the population interval reflects the joint uncertainty). With neither draws nor standard errors, a point.

Per-unit analytic standard errors cannot be combined into the mean's standard error: the units share one fitted model, so their errors are correlated, and treating them as independent understates the uncertainty by about a factor of sqrt(n). So, with analytic standard errors:

• a homogeneous result (one effect and one standard error repeated over the units, as DML.cate() returns) reports that effect and its standard error;
• a result carrying its cross-fit doubly-robust scores reports the AIPW average off them: the mean score with standard error std(scores) / sqrt(n) (the same scores explain() reads). For the DR-learner's linear final stage the point equals the mean per-unit effect; for a non-linear final stage such as the causal forest it can differ slightly from the plain mean of the per-unit effects (diagnostics["ate_from_dr_scores"] marks this path);
• otherwise it reports the mean with no standard error (kind="point"). For an interval, use estimate_effect().

PerUnitEffect.to_frame#

def PerUnitEffect.to_frame() -> pd.DataFrame

A per-unit table (index, effect, and lo/hi/se when available). Pandas only here.

PolicyTree#

class PolicyTree(
    root: PolicyNode,
    covariates: tuple[str, ...],
    n: int,
    depth: int,
    value: float,
    std_error: float,
    treat_all_value: float,
    treat_none_value: float,
    exhaustive: bool,
    treatment: str = '',
    outcome: str = '',
    _scores: npt.NDArray[np.float64] = <factory>,
    _design: npt.NDArray[np.float64] = <factory>,
    _blocks: Mapping[str, tuple[str, ...]] = <factory>,
) -> None

A shallow doubly-robust optimal-policy tree with its value and regret.

root is the learned tree over the named covariates; predict() and assign() evaluate it. value/std_error are the estimated policy value V(pi_hat) = mean[pi_hat(X) Gamma] and its standard error from the per-unit influence terms, in outcome units per unit and measured against treating no one (the mean outcome gain of following the rule rather than treating nobody). On that scale treat_none_value is 0 and treat_all_value is the doubly-robust ATE, and regret() is the value gap against either, also in outcome units per unit. These are measured on the units the tree was learned on, so they are optimistic (see evaluate() for the value on held-out units). depth is the requested depth and exhaustive records whether the globally optimal search was used.

PolicyTree.predict#

def PolicyTree.predict(
    covariates: npt.NDArray[np.float64],
) -> npt.NDArray[np.int64]

The per-row action (0/1) for a covariate matrix aligned to covariates.

PolicyTree.assign#

def PolicyTree.assign() -> npt.NDArray[np.int64]

The fitted policy's action for each unit it was learned on.

PolicyTree.interval#

def PolicyTree.interval(level: float = 0.9) -> tuple[float, float]

A normal interval for the in-sample policy value at level.

It ignores that the tree was chosen to maximize this value, so it is optimistic; evaluate() gives a valid interval on held-out units.

PolicyTree.regret#

def PolicyTree.regret(baseline: str = 'treat_all') -> float

The learned policy's welfare GAIN over a trivial baseline (the baseline's regret).

baseline is "treat_all", "treat_none", or "best_trivial" (the better of the two). Returns value - baseline_value: how much the mean outcome per unit rises under the learned rule compared with treating everyone / no one, in outcome units. The trivial constant policies are depth-0 leaves inside the tree's class, so the optimal tree is at least as good and this gain is non-negative up to sampling noise (it is the baseline's regret against the learned optimum, not the optimum's own regret, which would need the unknown oracle value).

PolicyTree.evaluate#

def PolicyTree.evaluate(
    result_or_scores: CATEResult | DRScores,
) -> PolicyValue

The learned rule's value on other units: the honest, out-of-sample value.

result_or_scores carries doubly-robust scores for units the tree was not learned on (a CATEResult from a held-out sample, or a DRScores); its covariates must include the tree's split covariates by name, encoded the same way: a categorical with three or more levels is split on per-level columns ("C==level"), and held-out scores whose data has different categories for it (a level missing or new) raise, since the same column would mean a different group. Build both datasets with the same categories (for example one pandas.Categorical dtype). Passing the learning sample's own result reproduces the optimistic in-sample value. Returns a PolicyValue: the mean of pi(X_i) Gamma_i there (outcome units, against treating no one), its standard error, and the treat-all baseline on the same units. Because the rule was fixed before these scores were seen, the interval is valid, unlike the in-sample value the search maximized.

PolicyTree.leaves#

def PolicyTree.leaves() -> list[PolicyNode]

The leaf nodes, left to right (each carries its action, count, and value).

PolicyTree.describe#

def PolicyTree.describe() -> str

A readable if-else rendering of the tree (the interpretable rule).

PolicyValue#

class PolicyValue(
    value: float,
    std_error: float,
    treat_all_value: float,
    n: int,
    regret_std_error: float = nan,
) -> None

A fixed policy's value on a set of units, from evaluate().

value is the mean outcome gain over treating no one (outcome units per unit), std_error its standard error, and treat_all_value the treat-everyone baseline (the ATE) on the same units. regret_std_error is the standard error of regret("treat_all"), from the paired per-unit terms (pi(X_i) - 1) Gamma_i.

PolicyValue.interval#

def PolicyValue.interval(level: float = 0.9) -> tuple[float, float]

A normal confidence interval for the value at level.

PolicyValue.regret#

def PolicyValue.regret(baseline: str = 'treat_all') -> float

value minus the "treat_all", "treat_none", or "best_trivial" value.

PyMCClassifier#

class PyMCClassifier(
    *,
    draws: int = 1000,
    tune: int = 1000,
    chains: int = 2,
    cores: int = 1,
    target_accept: float = 0.9,
    random_state: int | None = None,
    progressbar: bool = False,
    **build_kwargs: Any,
) -> None

A Bayesian classification learner over a PyMC model.

Name the probability node "p" (the per-row P(class=1)). Exposes predict_proba and predict_proba_draws, so it fills a propensity or a mediator slot and drives the posterior path there.

PyMCClassifier.predict_proba_draws#

def PyMCClassifier.predict_proba_draws(
    X: npt.NDArray[np.float64],
) -> npt.NDArray[np.float64]

Return (S, n) posterior draws of P(y = 1) at the rows of X.

PyMCClassifier.predict_proba#

def PyMCClassifier.predict_proba(
    X: npt.NDArray[np.float64],
) -> npt.NDArray[np.float64]

Return class probabilities as an (n, 2) array from the posterior mean of P(y = 1).

PyMCClassifier.predict#

def PyMCClassifier.predict(
    X: npt.NDArray[np.float64],
) -> npt.NDArray[np.float64]

Return the 0/1 label, 1 where the posterior mean of P(y = 1) is at least 0.5.

PyMCClassifier.build#

def PyMCClassifier.build(
    X: npt.NDArray[np.float64],
    y: npt.NDArray[np.float64],
) -> Any

Return a pymc.Model. Subclasses implement this.

PyMCClassifier.fit#

def PyMCClassifier.fit(
    X: npt.NDArray[np.float64],
    y: npt.NDArray[np.float64],
    sample_weight: npt.NDArray[np.float64] | None = None,
) -> PyMCLearner

Build the model with build() and sample its posterior with NUTS.

Raises TypeError when sample_weight is given: a PyMC model has no generic place for row weights (supports_sample_weight is False). Needs pymc.

Returns: The fitted learner itself.

PyMCClassifier.predict_draws#

def PyMCClassifier.predict_draws(
    X: npt.NDArray[np.float64],
) -> npt.NDArray[np.float64]

Return (S, n) posterior draws of the mean node at the rows of X.

S is chains * draws. For a PyMCClassifier the mean node is "p", so these are probability draws.

PyMCClassifier.get_params#

def PyMCClassifier.get_params(deep: bool = True) -> dict[str, Any]

Return the sampler settings and the extra build keyword arguments.

PyMCClassifier.set_params#

def PyMCClassifier.set_params(**params: Any) -> PyMCLearner

Set sampler settings or build keyword arguments by name; return the learner.

PyMCLearner#

class PyMCLearner(
    *,
    draws: int = 1000,
    tune: int = 1000,
    chains: int = 2,
    cores: int = 1,
    target_accept: float = 0.9,
    random_state: int | None = None,
    progressbar: bool = False,
    **build_kwargs: Any,
) -> None

A Bayesian regression learner over a PyMC model.

Implement build() to return a pymc.Model. predict is the posterior predictive mean and predict_draws the (S, n) posterior of the mean node, so the learner drops into a regression slot and also drives the posterior path. cores (default 1) caps the processes PyMC may use for the chains, so a fit runs its chains one after another unless you raise it.

PyMCLearner.build#

def PyMCLearner.build(
    X: npt.NDArray[np.float64],
    y: npt.NDArray[np.float64],
) -> Any

Return a pymc.Model. Subclasses implement this.

PyMCLearner.fit#

def PyMCLearner.fit(
    X: npt.NDArray[np.float64],
    y: npt.NDArray[np.float64],
    sample_weight: npt.NDArray[np.float64] | None = None,
) -> PyMCLearner

Build the model with build() and sample its posterior with NUTS.

Raises TypeError when sample_weight is given: a PyMC model has no generic place for row weights (supports_sample_weight is False). Needs pymc.

Returns: The fitted learner itself.

PyMCLearner.predict_draws#

def PyMCLearner.predict_draws(
    X: npt.NDArray[np.float64],
) -> npt.NDArray[np.float64]

Return (S, n) posterior draws of the mean node at the rows of X.

S is chains * draws. For a PyMCClassifier the mean node is "p", so these are probability draws.

PyMCLearner.predict#

def PyMCLearner.predict(
    X: npt.NDArray[np.float64],
) -> npt.NDArray[np.float64]

Return the posterior mean of the mean node at the rows of X.

PyMCLearner.get_params#

def PyMCLearner.get_params(deep: bool = True) -> dict[str, Any]

Return the sampler settings and the extra build keyword arguments.

PyMCLearner.set_params#

def PyMCLearner.set_params(**params: Any) -> PyMCLearner

Set sampler settings or build keyword arguments by name; return the learner.

RLearner#

class RLearner(
    learners: Any = None,
    *,
    n_folds: int = 5,
    clip: float = 0.001,
    **kwargs: Any,
) -> None

The R-learner: residualize the outcome and treatment, then fit the CATE.

Robinson's residual-on-residual moment with a heterogeneous final model, weighted by the squared treatment residual. The ATE is the average of the fitted CATE, with a bootstrap standard error.

RLearner.fit#

def RLearner.fit(identification: Identification, data: TabularData) -> Self

Read the estimand and the data, resolve learners, store the fit state.

RLearner.cate#

def RLearner.cate(
    data: TabularData | None = None,
    *,
    uncertainty: bool = True,
) -> CATEResult

Per-unit CATE as a CATEResult.

In-sample when data is None, else scored on its covariates. The result carries the per-unit point, the uncertainty (posterior draws for a Bayesian slot, the bootstrap draws otherwise, the analytic standard error for the linear DR final), the identification, and the cached doubly-robust scores for explain().

Pass uncertainty=False for a fast point-only result (kind="point", no standard error or draws): it skips the uncertainty step, which for a non-linear final means a full bootstrap refit of the cross-fit pipeline. Use it when only the effect curve is needed (ranking, targeting, plots); the per-unit points are identical either way. explain still works (it reads the cached scores, not the draws).

RLearner.effect#

def RLearner.effect(
    *,
    target: EffectTarget | str | None = None,
) -> EffectEstimate

Report the effect, with an analytic standard error or a bootstrap one.

Refutation#

class Refutation(
    method: str,
    estimand_value: float,
    refuted_value: float,
    refuted_std: float,
    p_value: float | None,
    passed: bool,
    n_simulations: int,
    diagnostics: Mapping[str, float] = <factory>,
) -> None

The result of one refutation check.

estimand_value is the effect on the real data. refuted_value is the mean effect over the perturbed runs, with refuted_std their spread. passed is the convenience verdict (the rule is in the method's section of refute_effect()); read the numbers, not just the flag. p_value is the permutation tail for the placebo check and None for the others.

SelectionReport#

class SelectionReport(
    ranking: tuple[CandidateScore, ...],
    selected_method: str,
    selected: CATEResult,
    aggregate: CATEResult,
    aggregate_methods: tuple[str, ...],
    weights: npt.NDArray[np.float64] = <factory>,
    aggregate_tau_risk: float = nan,
) -> None

The ranked CATE model-selection result with the chosen and aggregated models.

ranking is the candidates sorted by ascending DR tau-risk (best first). selected is the fitted CATEResult of the lowest risk; selected_method names it. aggregate is the Q-aggregation ensemble result and weights its convex weights (aligned to aggregate_methods), with aggregate_tau_risk the ensemble's DR tau-risk.

SelectionReport.risk#

def SelectionReport.risk(method: str) -> float

The DR tau-risk of a candidate by name (raises when it was not scored).

SelectionReport.best_single_tau_risk#

property SelectionReport.best_single_tau_risk: float

The DR tau-risk of the best single candidate (the first in the ranking).

SelectionReport.frame#

def SelectionReport.frame() -> dict[str, float]

A method-to-risk mapping (the aggregate included), handy for tables/tests.

Slot#

class Slot(name: str, task: Task, weighted: bool = False) -> None

A named nuisance position with the learning task it requires.

weighted marks a slot whose learner is fit with sample_weight. A learner that declares supports_sample_weight = False (a remote learner whose fits take no row weights) is refused there when the slots resolve, before any fit.

Stratification#

class Stratification(
    learners: Any = None,
    *,
    n_strata: int = 5,
    clip: float = 0.001,
    **kwargs: Any,
) -> None

Propensity stratification: bin on the propensity and difference within strata.

The estimate is the stratum-size-weighted (ATE) or treated-size-weighted (ATT) average of the within-stratum mean difference. Standard error by bootstrap.

Stratification.fit#

def Stratification.fit(
    identification: Identification,
    data: TabularData,
) -> Self

Read the estimand and the data, resolve learners, store the fit state.

Stratification.effect#

def Stratification.effect(
    *,
    target: EffectTarget | str | None = None,
) -> EffectEstimate

Report the effect, with an analytic standard error or a bootstrap one.

SyntheticControlEffect#

class SyntheticControlEffect(
    target: EffectTarget,
    estimator: str,
    learners: Mapping[str, str],
    n: int,
    kind: str,
    diagnostics: Mapping[str, float] = <factory>,
    _point: float | None = None,
    _std_error: float | None = None,
    _draws: npt.NDArray[np.float64] | None = None,
    design: Design | None = None,
    treated_unit: Any = None,
    at: Any = None,
    weights: Mapping[Any, float] = <factory>,
    periods: tuple[Any, ...] = (),
    observed: tuple[float, ...] = (),
    synthetic: tuple[float, ...] = (),
    pre_rmspe: float = nan,
    post_rmspe: float = nan,
    p_value: float | None = None,
    placebo_gaps: Mapping[Any, tuple[float, ...]] = <factory>,
) -> None

A synthetic-control fit: the weights, the two paths, and placebo inference.

observed and synthetic run over periods, and gaps is their difference. pre_rmspe is the tracking error before at, the fit diagnostic to judge before trusting the effect. p_value ranks the treated unit's post-to-pre RMSPE ratio among the donor placebos, and placebo_gaps holds each donor's post-period placebo gaps. There is no sampling story for one treated unit, so std_error stays None on purpose, interval is a placebo band instead, and kind is "placebo".

SyntheticControlEffect.gaps#

property SyntheticControlEffect.gaps: tuple[float, ...]

Observed minus synthetic, one value per period.

SyntheticControlEffect.interval#

def SyntheticControlEffect.interval(
    level: float = 0.95,
) -> tuple[float, float]

A placebo-permutation band around the estimate, not a sampling CI.

The half width is the level quantile of the absolute placebo post-gap means: the gap a unit with no treatment can show. An estimate whose band excludes zero clears the placebo bar at that level.

SyntheticControlEffect.value#

property SyntheticControlEffect.value: float

The point estimate, or the mean of the draws.

SyntheticControlEffect.std_error#

property SyntheticControlEffect.std_error: float | None

The analytic standard error, or the standard deviation of the draws.

SyntheticControlEffect.draws#

property SyntheticControlEffect.draws: npt.NDArray[np.float64] | None

The bootstrap or posterior draws, or None for a pure analytic result.

Task#

class Task(enum.Enum)

What a nuisance slot asks of its learner.

TwoStageLeastSquares#

class TwoStageLeastSquares(learners: Any = None, **kwargs: Any) -> None

Two-stage least squares. OLS in both stages, robust standard errors.

The first stage regresses the treatment on the instruments and any conditioning covariates; the second regresses the outcome on the fitted treatment. Reports the local average treatment effect (LATE).

TwoStageLeastSquares.fit#

def TwoStageLeastSquares.fit(
    identification: Identification,
    data: TabularData,
) -> Self

Read the estimand and the data, resolve learners, store the fit state.

TwoStageLeastSquares.effect#

def TwoStageLeastSquares.effect(
    *,
    target: EffectTarget | str | None = None,
) -> EffectEstimate

Report the effect, with an analytic standard error or a bootstrap one.

TwoStageRegression#

class TwoStageRegression(learners: Any = None, **kwargs: Any) -> None

Estimate the front-door functional from a mediator model and an outcome model.

For binary treatment A and binary mediator M: a classifier gives P(M | A), a regression gives E[Y | M, A], and the effect is the front-door sum E[Y | do(A=1)] - E[Y | do(A=0)]. Continuous and multi-level mediators are not supported.

TwoStageRegression.fit#

def TwoStageRegression.fit(
    identification: Identification,
    data: TabularData,
) -> Self

Read the estimand and the data, resolve learners, store the fit state.

TwoStageRegression.effect#

def TwoStageRegression.effect(
    *,
    target: EffectTarget | str | None = None,
) -> EffectEstimate

Report the effect, with an analytic standard error or a bootstrap one.

VariableImportance#

class VariableImportance(
    names: tuple[str, ...],
    values: npt.NDArray[np.float64] = <factory>,
) -> None

Depth-weighted, normalized split frequencies: a driver signal over covariates.

names are the covariate columns and values the matching importances, which sum to one (or are all zero for a forest that never split). A split at depth d contributes 1 / (d + 1), so shallow splits, the ones that move the effect most, weigh more, following the GRF variable-importance convention.

VariableImportance.top#

def VariableImportance.top(k: int | None = None) -> list[tuple[str, float]]

The covariates by descending importance (all of them when k is None).

VariableImportance.to_dict#

def VariableImportance.to_dict() -> dict[str, float]

Return {covariate: importance} for every covariate, in the stored order.

Wald#

class Wald(learners: Any = None, **kwargs: Any) -> None

The Wald ratio: two-stage least squares with a single instrument.

Wald.fit#

def Wald.fit(identification: Identification, data: TabularData) -> Self

Read the estimand and the data, resolve learners, store the fit state.

Wald.effect#

def Wald.effect(
    *,
    target: EffectTarget | str | None = None,
) -> EffectEstimate

Report the effect, with an analytic standard error or a bootstrap one.

XLearner#

class XLearner(
    learners: Any = None,
    *,
    clip: float = 0.001,
    **kwargs: Any,
) -> None

The X-learner: per-arm outcome models, imputed effects, propensity-weighted.

Fit one outcome model per arm, impute each unit's effect against the other arm's model, regress those imputed effects per arm, and combine the two with the propensity. The ATE is the average of the fitted CATE, with a bootstrap standard error.

XLearner.fit#

def XLearner.fit(identification: Identification, data: TabularData) -> Self

Read the estimand and the data, resolve learners, store the fit state.

XLearner.cate#

def XLearner.cate(
    data: TabularData | None = None,
    *,
    uncertainty: bool = True,
) -> CATEResult

Per-unit CATE as a CATEResult.

In-sample when data is None, else scored on its covariates. The result carries the per-unit point, the uncertainty (posterior draws for a Bayesian slot, the bootstrap draws otherwise, the analytic standard error for the linear DR final), the identification, and the cached doubly-robust scores for explain().

Pass uncertainty=False for a fast point-only result (kind="point", no standard error or draws): it skips the uncertainty step, which for a non-linear final means a full bootstrap refit of the cross-fit pipeline. Use it when only the effect curve is needed (ranking, targeting, plots); the per-unit points are identical either way. explain still works (it reads the cached scores, not the draws).

XLearner.effect#

def XLearner.effect(
    *,
    target: EffectTarget | str | None = None,
) -> EffectEstimate

Report the effect, with an analytic standard error or a bootstrap one.

candidates#

def candidates(
    identification: Identification,
) -> list[type[CausalEstimator]]

Every estimator that consumes the identified family, default first.

cate#

def cate(
    graph: MixedGraph,
    data: TabularData,
    treatment: str | Iterable[str],
    outcome: str | Iterable[str],
    *,
    method: str = 'auto',
    learner: Any = None,
    treated: Any | None = None,
    random_state: int | None = None,
    n_boot: int = 1000,
    uncertainty: bool = True,
    **kwargs: Any,
) -> CATEResult

Identify the effect of treatment on outcome, then estimate the CATE.

Identification-gated exactly like estimate_effect(): identify the effect, select a CATE-capable estimator for the family, fit it, and return a CATEResult. method names the meta-learner ("dr_learner" by default, the doubly-robust choice that shares the explainability score helper; also "r_learner", "x_learner", "t_learner", "s_learner", "dml"). learner fills the nuisance slots (an alias, a learner, or a slot mapping). The result carries per-unit uncertainty (a posterior with a Bayesian slot, the bootstrap otherwise, the analytic standard error for the linear DR final) and the cached doubly-robust scores, so explain() runs the Tier-B characterization. Pass uncertainty=False to skip that step for a fast point-only result (explain still works off the cached scores).

did#

def did(
    data: PanelData,
    outcome: str,
    treatment: str,
    *,
    controls: Iterable[str] = (),
    comparison: str = 'never_treated',
    strategy: str = 'auto',
    n_boot: int = 1000,
    random_state: int | None = None,
) -> DiDEffect

Difference-in-differences: the average effect of adoption on the treated (ATT).

The panel is read through adoption timing: a unit's adoption period is the first period its treatment equals 1, units that never show a 1 are never treated, and the treatment must be absorbing (never switching from 1 back to 0 within a unit). Periods compare by their position in the sorted period sequence, so integer and datetime periods both work.

Flavors. strategy="canonical" needs a single adoption period g: each unit's outcome change from its pre-period mean (periods before g) to its post-period mean (periods at or after g), treated minus comparison. Units lacking a pre or a post observation are dropped with a warning and counted in diagnostics["n_dropped_units"]. Controls enter through a unit-level regression of that change on the treated flag and the unit's pre-period control means. strategy="group_time" handles staggered adoption in the style of Callaway and Sant'Anna: one ATT(g, t) per adoption cohort and period, each differenced against the period just before adoption, then aggregated over post periods with cohort-size weights; the cells land in by_group_time. strategy="twfe" is the static two-way fixed-effects regression (iterated within transform, so unbalanced panels work). Under staggered adoption with effects that change over time it is a weighted mix of cell effects with possibly negative weights (Goodman-Bacon), so prefer group_time there. strategy="auto" picks canonical for one adoption period and group_time for several, never twfe.

Comparison pools. comparison="never_treated" uses units that never adopt and raises when none exist. comparison="not_yet_treated" also admits units that adopt only after the periods a cell reads, the pool a panel without never-treated units needs. In the canonical flavor the two pools coincide because no later adopters exist, and twfe ignores the pool.

Inference is a unit bootstrap: resample units with replacement, recompute, and keep n_boot draws (a failed resample is skipped, with a budget of five attempts per draw). The result is kind="bootstrap"; the canonical no-controls fit also stores the two-sample analytic standard error in diagnostics["analytic_se"]. The parallel-trends and no-anticipation assumptions ride along in design.

Raises ValueError for an unknown column, a missing value in the outcome, treatment, or controls, a treatment not coded 0/1 or boolean, a non-absorbing treatment, no treated units, or an empty comparison pool.

dr_scores#

def dr_scores(
    identification: Identification,
    data: TabularData,
    *,
    learners: object = None,
    folds: int = 5,
    clip: float = 0.001,
    treated: object = None,
    random_state: int | None = None,
) -> DRScores

Cross-fit doubly-robust scores for an adjustment-identified effect.

learners is an alias, a learner, or a slot mapping over outcome and propensity (the same surface every estimator takes); the defaults are the native OLS (or logistic for a binary outcome) and logistic propensity. folds sets the cross-fit count. Returns a DRScores carrying the signal and the out-of-fold nuisances.

estimate#

def estimate(
    identification: Identification,
    data: TabularData,
    *,
    strategy: str | None = None,
    learner: Any = None,
    target: EffectTarget | str | None = None,
    treated: Any | None = None,
    random_state: int | None = None,
    n_boot: int = 1000,
    **kwargs: Any,
) -> EffectEstimate

Estimate an identified effect over tabular data.

strategy names the estimator (else the family default). learner is an alias, a learner, or a slot-to-learner mapping (else the slot defaults). treated names the treatment level that codes to 1 for a two-level treatment (else level order decides, with a warning when the levels are not 0/1). Extra keyword arguments pass to the estimator (per_arm, scheme, clip).

estimate_effect#

def estimate_effect(
    graph: MixedGraph,
    data: TabularData,
    treatment: str | Iterable[str],
    outcome: str | Iterable[str],
    *,
    strategy: str | None = None,
    learner: Any = None,
    target: EffectTarget | str | None = None,
    treated: Any | None = None,
    random_state: int | None = None,
    n_boot: int = 1000,
    **kwargs: Any,
) -> EffectEstimate

Identify the effect of treatment on outcome in graph, then estimate it.

treated names the treatment level that codes to 1 (see estimate()).

estimate_mediation#

def estimate_mediation(
    graph: MixedGraph,
    data: TabularData,
    treatment: str,
    outcome: str,
    mediators: str | frozenset[str],
    *,
    learner: Any = None,
    treated: Any | None = None,
    random_state: int | None = None,
    n_boot: int = 1000,
) -> MediationEstimate

Identify and estimate a mediation decomposition by g-computation.

Identifies the decomposition with identify_mediation() (raising if the natural direct and indirect effects are not identifiable), then estimates NDE, NIE, the total effect, and the controlled direct effect at each mediator level by g-computation: fit mu(x, m, w) = E[Y | x, m, w] and p(m | x, w), plug them into the mediation formula, and average over the empirical covariates.

For a binary treatment and a discrete (binary or low-cardinality) mediator, with a continuous or binary outcome. learner is an alias, a learner, or a {"outcome": ..., "mediator": ...} mapping (the mediator slot needs a classifier). treated names the treatment level coded to 1. The standard errors are a nonparametric bootstrap over the rows at random_state, with the models averaged over the full sample's covariates (so they are for the sample-average effect; see the module notes).

A continuous mediator raises (it needs a density estimator, not available); a non-binary treatment raises. Returns a MediationEstimate with a Design record naming the (untestable) cross-world assumption.

event_study#

def event_study(
    data: PanelData,
    outcome: str,
    treatment: str,
    *,
    controls: Iterable[str] = (),
    comparison: str = 'never_treated',
    leads: int | None = None,
    lags: int | None = None,
    n_boot: int = 1000,
    random_state: int | None = None,
) -> DiDEffect

The effect by event time: leads as the placebo check, lags as the dynamic path.

Runs the group-time machinery of did() and aggregates the ATT(g, t) cells that share an event time, the position of the period minus the position of the adoption period, weighting cohorts by size. Event time -1 is the reference: every cell differences against the period just before adoption, so the effect there is zero by construction and it is omitted. leads and lags bound the reported window (leads=2 keeps event times from -2, lags=3 up to 3); None keeps everything available.

The headline value is the overall post-adoption ATT, the same aggregate did() reports under strategy="group_time", and by_event_time carries one EventTimeEffect per event time with a bootstrap standard error and 95% percentile interval from the same unit bootstrap that drives the headline draws. A lead whose interval sits away from zero is evidence against parallel trends.

Validation and the comparison pools match did().

interrupted_time_series#

def interrupted_time_series(
    data: TimeSeriesData,
    outcome: str,
    *,
    at: Any,
    seasonality: int | None = None,
    hac_lags: int | None = None,
) -> ITSEffect

Fit an interrupted time series for outcome with a break at at.

The break index is the first row whose time value is at least at; both sides need at least 3 points. The design matrix over row positions t = 0..T-1 is [1, t, post, post * (t - t0)], plus, with seasonality=m (between 2 and T // 2), Fourier pairs sin/cos(2 pi k t / m) for k = 1..min(3, m // 2). OLS fits it, and the coefficient covariance is Newey-West with Bartlett weights over hac_lags lags, default floor(0.75 * T ** (1/3)).

level_change is the post coefficient and slope_change the interaction. The headline value is the ATT, the mean post-period gap between the fit and the pre-trend projection, which equals level_change + slope_change * mean(t - t0 | post); its standard error comes from the matching linear combination of the HAC covariance. The counterfactual path zeroes the break terms but keeps the trend and seasonal terms, all at the fitted coefficients.

A NaN outcome raises, as does a break too close to either end.

policy_tree#

def policy_tree(
    result_or_scores: CATEResult | DRScores,
    covariates: str | tuple[str, ...] | list[str] | None = None,
    *,
    depth: int = 2,
) -> PolicyTree

Fit a shallow doubly-robust optimal-policy tree of depth depth.

result_or_scores is a CATEResult (its cached cross-fit scores are read) or a DRScores directly. covariates names the split covariates (a name, a list of names, or None for every covariate carried by the scores). A categorical covariate with three or more levels is carried as one indicator column per non-first level, named "C==level"; naming the variable "C" selects its whole block, and the tree's covariates lists the columns it can split on. The tree maximizes the doubly-robust welfare V(pi) = mean[pi(X) Gamma_i] (the mean outcome gain over treating no one, in outcome units); at depth <= 2 the search is exhaustive over thresholds and so globally optimal, and a deeper tree uses a greedy recursive best-split (documented, not guaranteed optimal).

Returns a frozen PolicyTree carrying the rule, the estimated policy value with a standard error from the per-unit influence terms, and the regret against treating everyone / no one.

refute_effect#

def refute_effect(
    graph: MixedGraph,
    data: TabularData,
    treatment: str,
    outcome: str,
    *,
    method: str = 'placebo_treatment',
    strategy: str | None = None,
    learner: Any = None,
    treated: Any | None = None,
    estimand_value: float | None = None,
    n_simulations: int = 100,
    subset_fraction: float = 0.8,
    random_state: int | None = None,
    **kwargs: Any,
) -> Refutation

Stress-test the effect of treatment on outcome under a perturbation.

Runs the same identify-then-estimate path as estimate_effect(), so pass the same strategy, learner, and treated you used for the estimate. estimand_value skips re-computing the effect on the real data when you already have it.

The three methods:

• "placebo_treatment" shuffles the treatment column, breaking its link to everything, and re-estimates. The placebo effect should sit near zero. passed when the mean placebo effect is under 20 percent of the real one; p_value is the share of placebo runs whose effect matches or beats the real one in size.
• "random_common_cause" adds an independent random covariate as a parent of both treatment and outcome and re-estimates. The effect should not move. passed when the mean stays within 10 percent of the original.
• "data_subset" re-estimates on a random subset_fraction of the rows. The effect should hold. passed when the original sits within two standard deviations of the subset spread.

Flexible or cross-fit estimators make each run a full fit, so lower n_simulations (or pick a faster strategy) when the estimator is heavy.

select_cate#

def select_cate(
    graph: MixedGraph,
    data: TabularData,
    treatment: str | Iterable[str],
    outcome: str | Iterable[str],
    *,
    candidates: Sequence[str] | None = None,
    learner: Any = None,
    treated: Any | None = None,
    random_state: int | None = None,
    n_boot: int = 1000,
    folds: int = 5,
    nu: float = 0.5,
    q_iters: int = 200,
) -> SelectionReport

Fit and rank candidate CATE learners by the doubly-robust tau-risk.

Identifies the effect (the same back-door gate as cate()), fits each candidate in candidates (default the registered meta-learners dr_learner/r_learner/x_learner/t_learner) through the cate facade, and scores each against one shared cross-fit AIPW signal Gamma by the DR tau-risk mean[(Gamma - tau_hat)^2]. Returns a SelectionReport with the ranked candidates (best risk first), the selected fitted CATEResult, and a Q-aggregation convex-ensemble result.

learner fills every candidate's nuisance slots; nu is the Q-aggregation mix (the penalty pulling weight toward lower-risk models). The shared Gamma is computed once with dr_scores(), so the comparison is fair across candidates and each candidate's own out-of-fold tau_hat keeps the ranking honest.

synthetic_control#

def synthetic_control(
    data: PanelData,
    outcome: str,
    *,
    treatment: str | None = None,
    treated: Any | None = None,
    at: Any | None = None,
    donors: Iterable[Any] | None = None,
) -> SyntheticControlEffect

Fit a synthetic control for one treated unit in a panel.

Name the treated unit in exactly one of two ways. Pass treatment, an absorbing 0/1 column (once a unit turns 1 it stays 1) that marks exactly one unit as treated; its first treated period becomes at. Or pass treated and at directly. donors defaults to every other unit.

The donor weights minimize the squared pre-period tracking error over the simplex (nonnegative, summing to one) by projected gradient descent with the Duchi projection; the run is deterministic. The headline value is the ATT, the mean of observed - synthetic over the post periods (t >= at), and pre_rmspe reports the fit it rests on.

Inference is by placebo permutation, with the treated unit excluded from every placebo run: each donor is treated as pseudo-treated, fit on the remaining donors, and scored by its post-to-pre RMSPE ratio (skipped when its pre RMSPE falls below 1e-12). p_value is (1 + count(ratio_j >= ratio_treated)) / (1 + n_placebos). There is no sampling standard error for one treated unit, so std_error stays None and interval reads the placebo gaps instead.

Validation: at must be one of the panel's periods with at least 2 pre periods and 1 post period; the treated unit must be observed at every period; donors missing any period are dropped with a warning and at least 2 complete donors must remain; a NaN outcome among the used rows raises.