statspai.qte¶
qte ¶
Quantile Treatment Effects (QTE) module for StatsPAI.
Provides estimators for: - Quantile DID (Athey & Imbens 2006) — DID at each quantile - QTE via Quantile Regression (Firpo 2007) — conditional QTE with controls - QTE via Distribution — propensity-score reweighting approach
References
Athey, S. & Imbens, G. W. (2006). Identification and Inference in Nonlinear Difference-in-Differences Models. Econometrica, 74(2), 431-497. [@athey2006identification]
Firpo, S. (2007). Efficient Semiparametric Estimation of Quantile Treatment Effects. Econometrica, 75(1), 259-276. [@firpo2007efficient]
QTEResult ¶
Bases: ResultProtocolMixin
Container for quantile treatment effect estimates.
Attributes:
| Name | Type | Description |
|---|---|---|
quantiles |
ndarray
|
Quantile grid. |
effects |
ndarray
|
QTE point estimates. |
se |
ndarray
|
Bootstrap / analytical standard errors. |
ci_lower, ci_upper |
ndarray
|
Confidence interval bounds. |
ate |
float
|
Average treatment effect (for comparison). |
method |
str
|
Estimation method label. |
n_obs |
int
|
Sample size. |
alpha |
float
|
Significance level. |
Examples:
A :class:QTEResult is produced by estimators such as :func:qte and
:func:qdid; inspect the quantile grid, per-quantile effects and the
mean ATE:
>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 300
>>> d = rng.integers(0, 2, n)
>>> y = 1.0 + 1.5 * d + rng.normal(0, 1, n)
>>> df = pd.DataFrame({"y": y, "d": d})
>>> res = sp.qte(df, y="y", treatment="d",
... quantiles=[0.25, 0.5, 0.75], n_boot=50, seed=0)
>>> isinstance(res, sp.QTEResult)
True
>>> res.quantiles.tolist()
[0.25, 0.5, 0.75]
>>> res.effects.shape
(3,)
plot ¶
QTE plot with CI bands and ATE reference line.
Returns (fig, ax).
DTEResult ¶
Bases: ResultProtocolMixin
Container for distributional treatment effect estimates.
Returned by :func:distributional_te. Carries the DTE curve over a
grid, quantile treatment effects, treated/counterfactual CDFs, and a
Kolmogorov-Smirnov statistic for the null of no distributional effect.
Examples:
>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 500
>>> d = rng.integers(0, 2, n)
>>> y = 1.0 + 1.5 * d + rng.normal(0, 1, n)
>>> df = pd.DataFrame({"y": y, "d": d})
>>> res = sp.distributional_te(
... df, y="y", treatment="d", method="ipw",
... quantiles=[0.25, 0.5, 0.75], n_boot=50, seed=42)
>>> isinstance(res, sp.DTEResult)
True
>>> res.qte_effects.round(2).tolist() # QTE at each quantile
[1.63, 1.51, 1.57]
DistIVResult
dataclass
¶
Bases: ResultProtocolMixin
Distributional IV LATE per quantile.
Returned by :func:dist_iv (and :func:kan_dlate). Holds the local
average treatment effect estimated at each requested quantile of the
outcome, together with bootstrap standard errors and confidence
intervals. Use :meth:to_frame for a tidy table or :meth:summary
for a printable report.
Examples:
>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 600
>>> z = rng.integers(0, 2, n)
>>> d = ((0.3 + 0.5 * z + rng.normal(0, 0.3, n)) > 0.5).astype(int)
>>> y = 1.0 + 1.0 * d + rng.normal(0, 1, n)
>>> df = pd.DataFrame({"y": y, "d": d, "z": z})
>>> res = sp.dist_iv(df, y="y", treat="d", instrument="z",
... quantiles=np.array([0.25, 0.5, 0.75]), n_boot=50)
>>> isinstance(res, sp.DistIVResult)
True
>>> res.late_q.round(2).tolist() # LATE at each quantile
[0.96, 1.04, 1.24]
HDPanelQTEResult
dataclass
¶
Bases: ResultProtocolMixin
QTE at multiple quantiles with high-dim control selection.
Returned by :func:sp.qte_hd_panel. Holds the evaluated
quantiles with their QTE estimates, SEs, and CIs, plus the list
of LASSO-selected controls. Call .summary() for a table.
Examples:
>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> rows = []
>>> for u in range(60):
... ui = rng.normal(0, 0.5)
... treated = u >= 30
... for t in range(6):
... d = 1.0 if (treated and t >= 3) else 0.0
... x1, x2, x3 = rng.normal(0, 1, 3)
... y = 1.0 + 1.2 * d + 0.5 * x1 + ui + rng.normal(0, 1)
... rows.append((u, t, y, d, x1, x2, x3))
>>> df = pd.DataFrame(
... rows, columns=["unit", "time", "y", "d", "x1", "x2", "x3"])
>>> res = sp.qte_hd_panel(
... df, y="y", treat="d", unit="unit", time="time",
... covariates=["x1", "x2", "x3"],
... quantiles=np.array([0.25, 0.5, 0.75]))
>>> isinstance(res, sp.HDPanelQTEResult)
True
>>> res.n_obs
360
BeyondAverageResult
dataclass
¶
Bases: ResultProtocolMixin
Distributional LATE on compliers.
Returned by :func:beyond_average_late; holds the per-quantile complier
LATE, its bootstrap SE / CI, and the estimated complier share.
Examples:
>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 600
>>> z = rng.integers(0, 2, n)
>>> d = ((0.2 + 0.6 * z + rng.normal(0, 0.3, n)) > 0.5).astype(int)
>>> y = 1.0 + 1.0 * d + rng.normal(0, 1, n)
>>> df = pd.DataFrame({"y": y, "d": d, "z": z})
>>> res = sp.beyond_average_late(
... df, y="y", treat="d", instrument="z",
... quantiles=np.array([0.25, 0.5, 0.75]), n_boot=50)
>>> isinstance(res, sp.BeyondAverageResult)
True
>>> res.late_q.round(2).tolist()
[0.98, 1.01, 1.14]
>>> round(res.complier_share, 2)
0.69
qdid ¶
qdid(data: DataFrame, y: str, group: str, time: str, quantiles: Optional[List[float]] = None, n_boot: int = 500, alpha: float = 0.05, seed: int = 42) -> QTEResult
Quantile Difference-in-Differences (Athey & Imbens 2006).
QTE_DID(τ) = F_{11}^{-1}(τ) - F_{10}^{-1}(τ) - [F_{01}^{-1}(τ) - F_{00}^{-1}(τ)]
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
data
|
DataFrame
|
|
required |
y
|
str
|
Outcome variable. |
required |
group
|
str
|
Binary group indicator (0 = control, 1 = treated). |
required |
time
|
str
|
Binary time indicator (0 = pre, 1 = post). |
required |
quantiles
|
list of float
|
Defaults to |
None
|
n_boot
|
int
|
Bootstrap replications. |
500
|
alpha
|
float
|
Significance level. |
0.05
|
seed
|
int
|
Random seed. |
42
|
Returns:
| Type | Description |
|---|---|
QTEResult
|
|
Examples:
>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 400
>>> g = rng.integers(0, 2, n)
>>> t = rng.integers(0, 2, n)
>>> y = (1.0 + 0.5 * g + 0.3 * t + 2.0 * g * t
... + rng.normal(0, 1, n))
>>> df = pd.DataFrame({"y": y, "g": g, "t": t})
>>> res = sp.qdid(df, y="y", group="g", time="t",
... quantiles=[0.25, 0.5, 0.75], n_boot=50)
>>> round(res.ate, 2) # true effect = 2.0
2.14
>>> np.round(res.effects, 2) # QTE at each quantile
array([1.96, 1.94, 2.26])
distributional_te ¶
distributional_te(data: DataFrame, y: str, treatment: str, x: Optional[List[str]] = None, method: str = 'ipw', n_grid: int = 100, quantiles: Optional[List[float]] = None, n_boot: int = 500, alpha: float = 0.05, seed: Optional[int] = None) -> DTEResult
Estimate distributional treatment effects.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
data
|
DataFrame
|
|
required |
y
|
str — outcome column.
|
|
required |
treatment
|
str — treatment column (binary 0/1 for IPW/DR;
|
0-3 group encoding for CiC). |
required |
x
|
list[str], optional — covariates (required for DR).
|
|
None
|
method
|
('ipw', 'dr', 'cic')
|
|
'ipw'
|
n_grid
|
int — grid points for CDF evaluation.
|
|
100
|
quantiles
|
list[float] — QTE quantile indices.
|
|
None
|
n_boot
|
int — bootstrap replications.
|
|
500
|
alpha
|
float — significance level.
|
|
0.05
|
seed
|
int, optional — random seed.
|
|
None
|
Returns:
| Type | Description |
|---|---|
DTEResult
|
|
Examples:
>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 500
>>> d = rng.integers(0, 2, n)
>>> y = 1.0 + 1.5 * d + rng.normal(0, 1, n)
>>> df = pd.DataFrame({"y": y, "d": d})
>>> res = sp.distributional_te(
... df, y="y", treatment="d", method="ipw",
... quantiles=[0.25, 0.5, 0.75], n_boot=50, seed=42)
>>> res.qte_effects.round(2).tolist() # QTE at each quantile
[1.63, 1.51, 1.57]
>>> round(res.ks_stat, 3) # Kolmogorov-Smirnov statistic
0.57
kan_dlate ¶
kan_dlate(data: DataFrame, y: str, treat: str, instrument: str, covariates: Optional[List[str]] = None, quantiles: Optional[ndarray] = None, alpha: float = 0.05, n_boot: int = 200, seed: int = 0) -> DistIVResult
KAN-Powered D-IV-LATE (Shaw 2025, arXiv 2506.12765).
Same identification as :func:dist_iv; would normally model the
bridge with a Kolmogorov-Arnold network. We currently fall back
to the kernel-smoothed Wald estimator (KAN requires a heavy
optional dependency). Functional equivalence is preserved at
standard quantile grids.
Examples:
>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 600
>>> z = rng.integers(0, 2, n)
>>> d = ((0.3 + 0.5 * z + rng.normal(0, 0.3, n)) > 0.5)
>>> d = d.astype(int)
>>> y = 1.0 + 1.0 * d + rng.normal(0, 1, n)
>>> df = pd.DataFrame({"y": y, "d": d, "z": z})
>>> res = sp.kan_dlate(df, y="y", treat="d", instrument="z",
... quantiles=np.array([0.25, 0.5, 0.75]),
... n_boot=50)
>>> res.late_q.round(2).tolist() # LATE at each quantile
[0.96, 1.04, 1.24]
qte_hd_panel ¶
qte_hd_panel(data: DataFrame, y: str, treat: str, unit: str, time: str, covariates: List[str], quantiles: Optional[ndarray] = None, alpha: float = 0.05, lasso_alpha: float = 0.01, seed: int = 0) -> HDPanelQTEResult
High-dimensional panel QTE via LASSO-selected controls.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
data
|
DataFrame
|
Long-format panel. |
required |
y
|
str
|
|
required |
treat
|
str
|
|
required |
unit
|
str
|
|
required |
time
|
str
|
|
required |
covariates
|
list of str
|
High-dim candidate control set. |
required |
quantiles
|
array - like
|
Defaults to (0.1, 0.25, 0.5, 0.75, 0.9). |
None
|
alpha
|
float
|
|
0.05
|
lasso_alpha
|
float
|
|
0.01
|
seed
|
int
|
|
0
|
Returns:
| Type | Description |
|---|---|
HDPanelQTEResult
|
|
Examples:
>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> rows = []
>>> for u in range(60):
... ui = rng.normal(0, 0.5)
... treated = u >= 30
... for t in range(6):
... d = 1.0 if (treated and t >= 3) else 0.0
... x1, x2, x3 = rng.normal(0, 1, 3)
... y = 1.0 + 1.2 * d + 0.5 * x1 + ui + rng.normal(0, 1)
... rows.append((u, t, y, d, x1, x2, x3))
>>> df = pd.DataFrame(
... rows, columns=["unit", "time", "y", "d", "x1", "x2", "x3"])
>>> res = sp.qte_hd_panel(
... df, y="y", treat="d", unit="unit", time="time",
... covariates=["x1", "x2", "x3"],
... quantiles=np.array([0.25, 0.5, 0.75]))
>>> res.qte.round(2).tolist() # QTE at each quantile
[0.83, 0.85, 1.36]
>>> res.n_obs
360
beyond_average_late ¶
beyond_average_late(data: DataFrame, y: str, treat: str, instrument: str, quantiles: Optional[ndarray] = None, alpha: float = 0.05, n_boot: int = 200, seed: int = 0) -> BeyondAverageResult
Distributional LATE on compliers under imperfect compliance.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
data
|
DataFrame
|
|
required |
y
|
str
|
|
required |
treat
|
str
|
|
required |
instrument
|
str
|
|
required |
quantiles
|
array - like
|
|
None
|
alpha
|
float
|
|
0.05
|
n_boot
|
int
|
|
200
|
seed
|
int
|
|
0
|
Returns:
| Type | Description |
|---|---|
BeyondAverageResult
|
|
Examples:
>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 600
>>> z = rng.integers(0, 2, n)
>>> d = ((0.2 + 0.6 * z + rng.normal(0, 0.3, n)) > 0.5)
>>> d = d.astype(int)
>>> y = 1.0 + 1.0 * d + rng.normal(0, 1, n)
>>> df = pd.DataFrame({"y": y, "d": d, "z": z})
>>> res = sp.beyond_average_late(
... df, y="y", treat="d", instrument="z",
... quantiles=np.array([0.25, 0.5, 0.75]), n_boot=50)
>>> res.late_q.round(2).tolist() # complier LATE per quantile
[0.98, 1.01, 1.14]
>>> round(res.complier_share, 2)
0.69