Skip to content

statspai.principal_strat

principal_strat

Principal Stratification (Frangakis & Rubin 2002).

Principal strata classify units by the joint potential values of a post-treatment variable (e.g. compliance type or survival status). Methods provided:

  • principal_strat(..., method='monotonicity') — sharp bounds on complier/always-taker/never-taker ATEs under monotonicity (Angrist, Imbens & Rubin 1996; Abadie 2002) + point-identified LATE.
  • principal_strat(..., method='principal_score') — covariate- based weighting estimator (Jo & Stuart 2009; Ding & Lu 2017) that point-identifies stratum-specific effects when principal ignorability holds.

Also ships :func:survivor_average_causal_effect (SACE bounds — Zhang & Rubin 2003) as a specialized entry point for the classical truncation-by-death problem.

PrincipalStratResult dataclass

Bases: ResultProtocolMixin

Principal stratification result.

Attributes:

Name Type Description
method str

'monotonicity' or 'principal_score'.
strata_proportions dict

Estimated proportion in each stratum.
effects DataFrame

Point estimate / SE / CI for each stratum-specific causal effect.
bounds DataFrame or None

For 'monotonicity' method, sharp Zhang-Rubin bounds on SACE.
n_obs int
alpha float
model_info dict

Examples:

>>> import statspai as sp
>>> import numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)
>>> n = 400
>>> D = rng.integers(0, 2, n)
>>> S = (rng.uniform(size=n) < 0.4 + 0.3 * D).astype(int)
>>> Y = 1.0 + 0.5 * D + 0.8 * S + rng.normal(0, 1.0, n)
>>> df = pd.DataFrame({"y": Y, "treat": D, "surv": S})
>>> res = sp.principal_strat(
...     df, y="y", treat="treat", strata="surv",
...     method="monotonicity", n_boot=100, seed=0,
... )
>>> isinstance(res, sp.PrincipalStratResult)
True
>>> res.n_obs
400
>>> bool("Principal Stratification" in res.summary())
True

survivor_average_causal_effect 

survivor_average_causal_effect(data: DataFrame, y: str, treat: str, survival: str, alpha: float = 0.05, n_boot: int = 500, seed: Optional[int] = None) -> CausalResult

Zhang-Rubin (2003) sharp bounds on the Survivor Average Causal Effect.

Returns a :class:CausalResult with estimate set to the midpoint of the SACE bounds and the endpoints stored in model_info.

References

zhang2003estimation

Examples:

>>> import statspai as sp
>>> import numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)
>>> n = 400
>>> D = rng.integers(0, 2, n)                              # treatment
>>> S = (rng.uniform(size=n) < 0.4 + 0.3 * D).astype(int)  # survival
>>> Y = 1.0 + 0.5 * D + 0.8 * S + rng.normal(0, 1.0, n)    # outcome
>>> df = pd.DataFrame({"y": Y, "treat": D, "surv": S})
>>> res = sp.survivor_average_causal_effect(
...     df, y="y", treat="treat", survival="surv",
...     n_boot=100, seed=0,
... )
>>> res.estimand
'SACE'
>>> bool(res.model_info["sace_lower"] <= res.model_info["sace_upper"])
True