"""
Regression discontinuity design (RDD) estimation - sharp and fuzzy, with
optional covariate adjustment - and robust bias-corrected inference,
parity-targeting R ``rdrobust`` 4.0.0.
Implements the local-polynomial RD estimators of Calonico, Cattaneo &
Titiunik (2014). SHARP (default): treatment is assigned by
``running >= cutoff``; the effect is the jump in the conditional
expectation of the outcome at the cutoff. FUZZY (pass
``fit(..., treatment_col=...)`` with the OBSERVED take-up column):
crossing the cutoff shifts take-up rather than determining it, and the
estimand is the local Wald ratio - the outcome jump divided by the
take-up jump - which for BINARY take-up under monotonicity is the LATE
for compliers at the cutoff (non-binary take-up keeps the ratio-of-jumps
reading; the ``estimand`` results field says which applies). Both designs
use kernel-weighted polynomial regressions on each side with data-driven
MSE/CER-optimal bandwidths and robust
bias-corrected (RBC) inference; the fuzzy bias correction is the
linearization of the ratio (not per-component), matching CCT 2014
Section 3.2 and rdrobust exactly.
Covariate adjustment (``fit(..., covariates=[...])``; Calonico, Cattaneo,
Farrell & Titiunik 2019, R's ``covs=``): covariates enter ADDITIVELY with
a common coefficient pooled across sides (CCFT 2019 Equation 2 - the only
specification with a clean guarantee; treatment-interacted and demeaned
variants are documented as inconsistent-or-inferior there). UNLIKE the
library's DiD estimators, where ``covariates`` switches identification to
conditional parallel trends, RD covariates DO NOT change the estimand -
the ``att`` still measures the same cutoff jump/ratio and the
``estimand`` label is unchanged; adjustment buys precision (shorter CIs)
when covariates predict the outcome near the cutoff. The operative
requirement is covariate BALANCE at the cutoff (zero RD effect on each
covariate); imbalanced covariates make the adjusted estimator
inconsistent, and adjusting "for" imbalance cannot restore
identification. Balance is testable with the estimator itself::
balance = RegressionDiscontinuity().fit(df, outcome_col="z1",
running_col="x")
balance.p_value # small p = imbalance; do not adjust for z1
Bandwidths are covariate-AWARE (covariates propagate into selection, not
just estimation, as in R). Collinear covariates are dropped with a
warning under ``covs_drop=True`` (R's default; the warning names the
dropped columns).
Canonical inference binding
---------------------------
``RegressionDiscontinuityResults`` binds the library-canonical fields to ONE
internally coherent inference row - the ROBUST row of rdrobust's output:
``att`` is the bias-corrected point estimate ``tau_bc`` (the linearized
bias-corrected RATIO on fuzzy fits), ``se`` its robust standard error, and
``t_stat``/``p_value``/``conf_int`` are computed from that same pair, so
the library-wide identities hold (``t_stat == att/se``, ``conf_int``
centered on ``att``). The ``estimand`` results field names what ``att``
measures for the fit at hand. This deliberately differs from rdrobust's
PRINTED headline, which reports the conventional estimate ``tau_cl`` in the
coefficient column while taking inference from the robust row; ``tau_cl`` is
first-class here as ``att_conventional`` (with its own full inference row),
and ``summary()`` prints the familiar three-row rdrobust table. Fuzzy fits
additionally expose the first stage (take-up jump) as a full three-row
mirror (``first_stage*`` fields) and print it above the treatment effects,
as R does.
rdrobust equivalents
--------------------
======================= ==========================================
diff-diff R rdrobust
======================= ==========================================
``cutoff`` ``c``
``vcov_type`` ``vce``
``alpha`` ``1 - level/100``
``h``, ``b``, ``rho`` ``h``, ``b``, ``rho`` (same semantics)
``p``, ``q`` ``p``, ``q``
``bwselect`` ``bwselect`` (same 10-option menu)
``kernel`` ``kernel`` (accepts "tri"/"epa"/"uni" too)
``masspoints`` ``masspoints`` ("adjust"/"check"/"off")
``nnmatch`` ``nnmatch``
``treatment_col`` (fit) ``fuzzy`` (observed take-up variable)
``sharpbw`` ``sharpbw`` (same default and semantics)
``covariates`` (fit) ``covs`` (column names instead of a matrix)
``covs_drop`` ``covs_drop`` (same default and semantics)
======================= ==========================================
Not in v1 (documented seams, see REGISTRY.md): cluster-robust variance,
weights, ``deriv``/kink estimands, ``scalepar``, ``stdvars``, hc0-hc3
variance modes, weak-IV-robust fuzzy inference (Feir-Lemieux-Marmer),
and a packaged covariate-balance helper (the recipe above covers it).
References
----------
- Calonico, S., Cattaneo, M. D., & Titiunik, R. (2014). Robust Nonparametric
Confidence Intervals for Regression-Discontinuity Designs. *Econometrica*,
82(6), 2295-2326.
- Calonico, S., Cattaneo, M. D., Farrell, M. H., & Titiunik, R. (2017).
rdrobust: Software for regression-discontinuity designs. *Stata Journal*,
17(2), 372-404.
- Calonico, S., Cattaneo, M. D., & Farrell, M. H. (2018). On the Effect of
Bias Estimation on Coverage Accuracy in Nonparametric Inference. *JASA*,
113(522), 767-779.
- Calonico, S., Cattaneo, M. D., Farrell, M. H., & Titiunik, R. (2019).
Regression Discontinuity Designs Using Covariates. *Review of Economics
and Statistics*, 101(3), 442-451.
"""
from __future__ import annotations
import warnings
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple
import numpy as np
import pandas as pd
from diff_diff._rdrobust_port import (
BWSELECT_OPTIONS,
_fuzzy_identification_stop,
_normalize_kernel,
covs_drop_fun,
rdbwselect,
rdrobust_fit,
)
from diff_diff.utils import safe_inference, validate_covariate_names
__all__ = [
"RegressionDiscontinuity",
"RegressionDiscontinuityResults",
]
def _json_safe(value: Any) -> Any:
if isinstance(value, (np.floating, np.integer)):
return value.item()
return value
[docs]
@dataclass
class RegressionDiscontinuityResults:
"""Results of a regression discontinuity fit (sharp or fuzzy; the
``estimand`` field names which one, and ``first_stage*`` fields are
populated on fuzzy fits only).
Canonical inference fields (``att``, ``se``, ``t_stat``, ``p_value``,
``conf_int``) all describe the ROBUST bias-corrected row: ``att`` is the
bias-corrected point estimate and ``conf_int`` is centered on it (see
module docstring for the binding rationale and the deviation from
rdrobust's printed headline). The conventional row is exposed as
explicit ``*_conventional`` fields. The (rarely used) middle
"Bias-Corrected" row shares its coefficient with ``att`` (both are
``tau_bc``) and its standard error with ``se_conventional`` - only its
inference triple carries the ``*_bias_corrected`` suffix
(``t_stat_bias_corrected``, ``p_value_bias_corrected``,
``conf_int_bias_corrected``); there are deliberately no redundant
``att_bias_corrected`` / ``se_bias_corrected`` fields. Together the
three rows mirror rdrobust's output exactly.
"""
# Canonical (robust row; internally coherent)
att: float
se: float
t_stat: float
p_value: float
conf_int: Tuple[float, float]
alpha: float
# Conventional row (rdrobust's printed headline coefficient)
att_conventional: float
se_conventional: float
t_stat_conventional: float
p_value_conventional: float
conf_int_conventional: Tuple[float, float]
# Bias-corrected middle row (tau_bc with the CONVENTIONAL SE; exposed
# for rdrobust parity - prefer the robust row for inference)
t_stat_bias_corrected: float
p_value_bias_corrected: float
conf_int_bias_corrected: Tuple[float, float]
# Explicit duplicates for clarity
se_robust: float
# Bandwidths (rdrobust bws layout)
h_left: float
h_right: float
b_left: float
b_right: float
# Sample composition
n_obs: int
n_left: int
n_right: int
n_h_left: int
n_h_right: int
n_b_left: int
n_b_right: int
n_unique_left: int
n_unique_right: int
n_dropped: int
# Config echoes. ``bwselect`` is the RESOLVED selector label ("Manual"
# when bandwidths were user-supplied or N<20 forced the full-range
# fallback, matching rdrobust's printed "BW type"); ``h_input`` /
# ``b_input`` / ``rho_input`` echo the constructor arguments as supplied
# (None when data-driven; a warned-and-ignored ``b``-without-``h`` still
# echoes here) - together with the other echoes they reconstruct the
# full fit configuration from a saved result. Resolved per-side
# bandwidths live in ``h_left``/``h_right``/``b_left``/``b_right``.
cutoff: float
p: int
q: int
kernel: str
bwselect: str
vcov_type: str
nnmatch: int
masspoints: str
bwcheck: Optional[int]
bwrestrict: bool
scaleregul: float
h_input: Optional[float]
b_input: Optional[float]
rho_input: Optional[float]
# Design echoes: ``estimand`` names what ``att`` measures for THIS fit
# - "sharp (ATE at the cutoff)"; "fuzzy (LATE for compliers at the
# cutoff)" for BINARY take-up; or "fuzzy (local Wald ratio at the
# cutoff; non-binary take-up)" when the take-up column is not {0, 1}
# (the complier-LATE reading does not apply to dose take-up).
# ``treatment_col`` is the fit-time take-up column name
# (None on sharp fits; no ``_input`` suffix - that convention is
# reserved for constructor arguments); ``sharpbw`` and ``covs_drop``
# echo the constructor flags. The estimand label deliberately does NOT
# change under covariate adjustment: CCFT 2019 covariates target the
# SAME estimand (precision only) - see ``covariates`` below.
estimand: str
sharpbw: bool
treatment_col: Optional[str]
covs_drop: bool
# First-stage (take-up jump) three-row mirror - fuzzy fits only, all
# None on sharp fits. Same binding rule as the main estimate: the
# unsuffixed quintet is the coherent ROBUST row (first_stage = the
# bias-corrected first-stage estimate tau_T_bc, first_stage_se = its
# robust SE); the conventional row and the bias-corrected middle-row
# inference triple mirror the main fields' suffix scheme.
first_stage: Optional[float] = None
first_stage_se: Optional[float] = None
first_stage_t_stat: Optional[float] = None
first_stage_p_value: Optional[float] = None
first_stage_conf_int: Optional[Tuple[float, float]] = None
first_stage_conventional: Optional[float] = None
first_stage_se_conventional: Optional[float] = None
first_stage_t_stat_conventional: Optional[float] = None
first_stage_p_value_conventional: Optional[float] = None
first_stage_conf_int_conventional: Optional[Tuple[float, float]] = None
first_stage_t_stat_bias_corrected: Optional[float] = None
first_stage_p_value_bias_corrected: Optional[float] = None
first_stage_conf_int_bias_corrected: Optional[Tuple[float, float]] = None
# Covariate adjustment (CCFT 2019) - all None on unadjusted fits.
# ``covariates`` echoes the fit-time column names AS PASSED;
# ``covariates_dropped`` lists columns removed as collinear by
# covs_drop ([] when nothing was dropped); ``covariate_coefficients``
# maps each RETAINED covariate name to its common (pooled across
# sides) outcome-equation projection coefficient gamma - these are
# nuisance coefficients for the adjustment, NOT causal effects of the
# covariates. Fuzzy fits add ``first_stage_covariate_coefficients``
# (the take-up-equation gamma). Name-keyed dicts make R's internal
# name-length column sort invisible to users.
covariates: Optional[List[str]] = None
covariates_dropped: Optional[List[str]] = None
covariate_coefficients: Optional[Dict[str, float]] = None
first_stage_covariate_coefficients: Optional[Dict[str, float]] = None
# Per-side order-p coefficient vectors (rdplot seam); the outcome pair
# is always populated by fit(), so typed non-Optional despite the
# dataclass default; the take-up pair is fuzzy-only. On
# covariate-adjusted fits these are the ADJUSTED vectors (gamma
# combination applied), matching R's beta_Y_p_* / beta_T_p_*.
beta_p_left: np.ndarray = field(repr=False, default=None)
beta_p_right: np.ndarray = field(repr=False, default=None)
beta_t_p_left: Optional[np.ndarray] = field(repr=False, default=None)
beta_t_p_right: Optional[np.ndarray] = field(repr=False, default=None)
[docs]
def summary(self) -> str:
"""Human-readable summary with the three-row rdrobust table."""
width = 72
conf_level = 100 * (1 - self.alpha)
lines = []
lines.append("=" * width)
design = "Fuzzy" if self.first_stage is not None else "Sharp"
if self.covariates:
# Mirrors R's rdmodel string ("Covariate-adjusted ... RD
# estimates"); the estimand line below is deliberately
# UNCHANGED - covariates buy precision, not a new estimand.
design = f"Covariate-adjusted {design}"
lines.append(f"{design} Regression Discontinuity (rdrobust parity)".center(width))
lines.append("=" * width)
lines.append(f"Cutoff: {self.cutoff:g}")
lines.append(f"Estimand: {self.estimand}")
if self.covariates:
cov_line = f"Covariates ({len(self.covariates)}): " + ", ".join(self.covariates)
if self.covariates_dropped:
cov_line += " [dropped: " + ", ".join(self.covariates_dropped) + "]"
lines.append(cov_line)
lines.append(f"Kernel: {self.kernel:<14} Bandwidth selector: {self.bwselect}")
lines.append(
f"Order (p, q): ({self.p}, {self.q}) VCE: {self.vcov_type} "
f"(nnmatch={self.nnmatch}) Masspoints: {self.masspoints}"
)
lines.append(
f"N = {self.n_obs} ({self.n_left} left / {self.n_right} right); "
f"effective N_h = {self.n_h_left}/{self.n_h_right}, "
f"N_b = {self.n_b_left}/{self.n_b_right}"
)
lines.append(
f"h = [{self.h_left:.4f}, {self.h_right:.4f}] "
f"b = [{self.b_left:.4f}, {self.b_right:.4f}]"
)
lines.append("-" * width)
header = (
f"{'Method':<16}{'Coef.':>11}{'Std. Err.':>11}{'z':>9}"
f"{'P>|z|':>9}{'[' + f'{conf_level:g}% Conf. Int.]':>16}"
)
if self.first_stage is not None:
# Fuzzy: R prints a first-stage block above the treatment
# effects (print.summary.rdrobust); same three-row structure.
lines.append("First-stage estimates (treatment take-up jump)".center(width))
lines.append(header)
lines.append("-" * width)
fs_rows = [
(
"Conventional",
self.first_stage_conventional,
self.first_stage_se_conventional,
self.first_stage_t_stat_conventional,
self.first_stage_p_value_conventional,
self.first_stage_conf_int_conventional,
),
(
"Bias-Corrected",
self.first_stage,
self.first_stage_se_conventional,
self.first_stage_t_stat_bias_corrected,
self.first_stage_p_value_bias_corrected,
self.first_stage_conf_int_bias_corrected,
),
(
"Robust",
self.first_stage,
self.first_stage_se,
self.first_stage_t_stat,
self.first_stage_p_value,
self.first_stage_conf_int,
),
]
for name, coef, se, z, pv, ci in fs_rows:
assert coef is not None and se is not None and ci is not None
assert z is not None and pv is not None
lines.append(
f"{name:<16}{coef:>11.4f}{se:>11.4f}{z:>9.3f}{pv:>9.3f}"
f" [{ci[0]:>7.4f}, {ci[1]:>7.4f}]"
)
lines.append("-" * width)
lines.append("Treatment effect estimates".center(width))
lines.append(header)
lines.append("-" * width)
rows = [
(
"Conventional",
self.att_conventional,
self.se_conventional,
self.t_stat_conventional,
self.p_value_conventional,
self.conf_int_conventional,
),
(
"Bias-Corrected",
self.att,
self.se_conventional,
self.t_stat_bias_corrected,
self.p_value_bias_corrected,
self.conf_int_bias_corrected,
),
(
"Robust",
self.att,
self.se_robust,
self.t_stat,
self.p_value,
self.conf_int,
),
]
for name, coef, se, z, pv, ci in rows:
lines.append(
f"{name:<16}{coef:>11.4f}{se:>11.4f}{z:>9.3f}{pv:>9.3f}"
f" [{ci[0]:>7.4f}, {ci[1]:>7.4f}]"
)
lines.append("-" * width)
lines.append("Note: canonical att/se/t_stat/p_value/conf_int are the ROBUST row")
lines.append("(att = bias-corrected estimate; rdrobust prints the conventional")
lines.append("estimate as its headline coefficient - see att_conventional).")
lines.append("=" * width)
return "\n".join(lines)
[docs]
def print_summary(self) -> None:
print(self.summary())
[docs]
def to_dict(self) -> Dict[str, Any]:
"""Flat scalar dict; confidence intervals split into lower/upper."""
out: Dict[str, Any] = {
"att": self.att,
"se": self.se,
"t_stat": self.t_stat,
"p_value": self.p_value,
"conf_int_lower": self.conf_int[0],
"conf_int_upper": self.conf_int[1],
"alpha": self.alpha,
"att_conventional": self.att_conventional,
"se_conventional": self.se_conventional,
"t_stat_conventional": self.t_stat_conventional,
"p_value_conventional": self.p_value_conventional,
"conf_int_conventional_lower": self.conf_int_conventional[0],
"conf_int_conventional_upper": self.conf_int_conventional[1],
"t_stat_bias_corrected": self.t_stat_bias_corrected,
"p_value_bias_corrected": self.p_value_bias_corrected,
"conf_int_bias_corrected_lower": self.conf_int_bias_corrected[0],
"conf_int_bias_corrected_upper": self.conf_int_bias_corrected[1],
"se_robust": self.se_robust,
"h_left": self.h_left,
"h_right": self.h_right,
"b_left": self.b_left,
"b_right": self.b_right,
"n_obs": self.n_obs,
"n_left": self.n_left,
"n_right": self.n_right,
"n_h_left": self.n_h_left,
"n_h_right": self.n_h_right,
"n_b_left": self.n_b_left,
"n_b_right": self.n_b_right,
"n_unique_left": self.n_unique_left,
"n_unique_right": self.n_unique_right,
"n_dropped": self.n_dropped,
"cutoff": self.cutoff,
"p": self.p,
"q": self.q,
"kernel": self.kernel,
"bwselect": self.bwselect,
"vcov_type": self.vcov_type,
"nnmatch": self.nnmatch,
"masspoints": self.masspoints,
"bwcheck": self.bwcheck,
"bwrestrict": self.bwrestrict,
"scaleregul": self.scaleregul,
"h_input": self.h_input,
"b_input": self.b_input,
"rho_input": self.rho_input,
"estimand": self.estimand,
"sharpbw": self.sharpbw,
"treatment_col": self.treatment_col,
"covs_drop": self.covs_drop,
# List/dict-valued covariate echoes (None on unadjusted fits;
# the lpdid/continuous_did echo convention).
"covariates": self.covariates,
"covariates_dropped": self.covariates_dropped,
"covariate_coefficients": self.covariate_coefficients,
"first_stage_covariate_coefficients": self.first_stage_covariate_coefficients,
"first_stage": self.first_stage,
"first_stage_se": self.first_stage_se,
"first_stage_t_stat": self.first_stage_t_stat,
"first_stage_p_value": self.first_stage_p_value,
"first_stage_conventional": self.first_stage_conventional,
"first_stage_se_conventional": self.first_stage_se_conventional,
"first_stage_t_stat_conventional": self.first_stage_t_stat_conventional,
"first_stage_p_value_conventional": self.first_stage_p_value_conventional,
"first_stage_t_stat_bias_corrected": self.first_stage_t_stat_bias_corrected,
"first_stage_p_value_bias_corrected": self.first_stage_p_value_bias_corrected,
}
# First-stage CIs are None on sharp fits - guard the tuple splits.
for key, ci in (
("first_stage_conf_int", self.first_stage_conf_int),
("first_stage_conf_int_conventional", self.first_stage_conf_int_conventional),
("first_stage_conf_int_bias_corrected", self.first_stage_conf_int_bias_corrected),
):
out[f"{key}_lower"] = None if ci is None else ci[0]
out[f"{key}_upper"] = None if ci is None else ci[1]
return {k: _json_safe(v) for k, v in out.items()}
[docs]
def to_dataframe(self) -> pd.DataFrame:
return pd.DataFrame([self.to_dict()])
[docs]
class RegressionDiscontinuity:
"""Regression discontinuity estimator - sharp and fuzzy, with
optional covariate adjustment (rdrobust 4.0.0 parity).
SHARP (default): treatment is defined by the running variable crossing
a known cutoff (``running >= cutoff`` treated, matching rdrobust:
units exactly at the cutoff are treated). FUZZY: pass the observed
take-up column via ``fit(..., treatment_col=...)`` - the estimand
becomes the local Wald ratio (complier LATE at the cutoff for binary
take-up under monotonicity; the ``estimand`` results field says which
reading applies) and the results gain a first-stage block.
COVARIATE ADJUSTMENT: pass ``fit(..., covariates=[...])`` (R's
``covs=``) for the CCFT 2019 additive common-coefficient adjustment -
the estimand is UNCHANGED (precision only; requires covariate balance
at the cutoff, see the module docstring), bandwidths become
covariate-aware, and collinear columns are dropped with a warning
under ``covs_drop=True``. Point
estimation uses kernel-weighted local polynomials of order ``p`` on
each side; inference is robust bias-corrected per Calonico, Cattaneo &
Titiunik (2014). Defaults reproduce ``rdrobust(y, x)`` /
``rdrobust(y, x, fuzzy=t)`` / ``rdrobust(y, x, covs=Z)``: ``p=1``,
``q=2``, triangular kernel, ``bwselect="mserd"``, nearest-neighbor
variance with 3 matches, ``masspoints="adjust"``, ``covs_drop=True``.
Parameters
----------
cutoff : float, default 0.0
The known threshold ``c`` of the running variable.
p : int, default 1
Local-polynomial order for point estimation; integer in 0..20
(mirroring rdrobust's accepted surface; ``p=0`` is the
local-constant fit).
q : int or None, default None
Order for the bias regression; an explicit ``q`` must satisfy
``p < q <= 20``. ``None`` resolves to ``p + 1`` WITHOUT
re-validation, exactly as R does (rdrobust.R:53-57) - so ``p=20``
with the default ``q`` yields ``q=21`` in both implementations
while an explicit ``q=21`` is rejected.
kernel : str, default "triangular"
"triangular", "epanechnikov", or "uniform" (R spellings
"tri"/"epa"/"uni" accepted).
bwselect : str, default "mserd"
Data-driven bandwidth selector; one of the 10 rdrobust options
(mserd, msetwo, msesum, msecomb1, msecomb2, cerrd, certwo, cersum,
cercomb1, cercomb2). Ignored when ``h`` is supplied.
h, b : float or None
Manual main / bias bandwidths (both sides). ``h`` alone implies
``b = h``; ``h`` with ``rho`` implies ``b = h/rho`` (overriding a
supplied ``b``, as in R); ``b`` without ``h`` is ignored with a
warning (R ignores it silently - documented deviation).
rho : float or None
Bandwidth ratio ``h/b``. Without ``h``, applies to the SELECTED
bandwidths (``b = h_selected/rho``), mirroring rdrobust.
vcov_type : str, default "nn"
Variance estimator. Only "nn" (same-side nearest-neighbor,
rdrobust's default) is implemented in this release; "hc0"-"hc3"
and cluster modes raise ``NotImplementedError``.
nnmatch : int, default 3
Minimum number of nearest neighbors for the NN variance.
masspoints : str, default "adjust"
Mass-point handling: "adjust" (rdrobust default), "check", "off".
bwcheck : int or None, default None
Minimum unique support points forced inside the bandwidth window.
bwrestrict : bool, default True
Clamp bandwidths to the running variable's observed range.
scaleregul : float, default 1.0
Scale of the IK-style regularization in bandwidth selection
(0 removes it).
sharpbw : bool, default False
Fuzzy fits only (``fit(..., treatment_col=...)``): when True,
bandwidths are selected for the SHARP reduced-form estimator on
the outcome (rdrobust's "approach 1") instead of the default
fuzzy-ratio objective. Automatically in effect - regardless of
this flag - under one-sided perfect compliance (zero take-up
variance on either side), exactly as in R. On sharp fits the flag
has no effect and a warning is emitted (R ignores it silently -
documented deviation). Never drops covariates from selection -
with ``covariates`` it selects on the covariate-adjusted sharp
objective, as in R.
covs_drop : bool, default True
Covariate-adjusted fits only (``fit(..., covariates=[...])``):
when True (R's default), redundant (collinear) covariate columns
are dropped with a warning naming them before fitting, and the
covariate projection uses a pseudo-inverse; when False the solve
is strict and collinear covariates raise a clear error. Without
``covariates`` the flag has no effect and setting it to False
emits a warning (same pattern as ``sharpbw`` on sharp fits).
alpha : float, default 0.05
Significance level (rdrobust ``level = 100*(1-alpha)``).
Examples
--------
>>> rd = RegressionDiscontinuity(cutoff=0.0)
>>> results = rd.fit(df, outcome_col="y", running_col="x")
>>> results.att, results.conf_int # robust bias-corrected inference
>>> fuzzy = rd.fit(df, "y", "x", treatment_col="takeup") # fuzzy RD
>>> fuzzy.att, fuzzy.first_stage # local Wald ratio + take-up jump
"""
[docs]
def __init__(
self,
cutoff: float = 0.0,
p: int = 1,
q: Optional[int] = None,
kernel: str = "triangular",
bwselect: str = "mserd",
h: Optional[float] = None,
b: Optional[float] = None,
rho: Optional[float] = None,
vcov_type: str = "nn",
nnmatch: int = 3,
masspoints: str = "adjust",
bwcheck: Optional[int] = None,
bwrestrict: bool = True,
scaleregul: float = 1.0,
sharpbw: bool = False,
covs_drop: bool = True,
alpha: float = 0.05,
):
self.cutoff = cutoff
self.p = p
self.q = q
self.kernel = kernel
self.bwselect = bwselect
self.h = h
self.b = b
self.rho = rho
self.vcov_type = vcov_type
self.nnmatch = nnmatch
self.masspoints = masspoints
self.bwcheck = bwcheck
self.bwrestrict = bwrestrict
self.scaleregul = scaleregul
self.sharpbw = sharpbw
self.covs_drop = covs_drop
self.alpha = alpha
self._validate_constructor_args()
# ------------------------------------------------------------------
# Configuration plumbing (sklearn-like)
# ------------------------------------------------------------------
@staticmethod
def _is_real_scalar(val: Any) -> bool:
# Reject non-numeric types up front so every scalar knob fails with
# the estimator's ValueError, not NumPy's TypeError (bool excluded:
# True is not a bandwidth).
return isinstance(val, (int, float, np.integer, np.floating)) and not isinstance(
val, (bool, np.bool_)
)
@staticmethod
def _is_int_scalar(val: Any) -> bool:
# bool is an int subclass; p=True must not silently become p=1.
return isinstance(val, (int, np.integer)) and not isinstance(val, (bool, np.bool_))
def _validate_constructor_args(self) -> None:
if not (self._is_real_scalar(self.cutoff) and np.isfinite(self.cutoff)):
raise ValueError(f"cutoff must be finite; got {self.cutoff!r}.")
# p/q bounds mirror rdrobust.R:47-57 exactly: integers in 0:20 with
# q > p (p=0 is R's local-constant fit; q caps at 20 like p).
if not (self._is_int_scalar(self.p) and 0 <= self.p <= 20):
raise ValueError(f"p must be an integer in 0..20; got {self.p!r}.")
if self.q is not None and not (self._is_int_scalar(self.q) and self.p < self.q <= 20):
raise ValueError(
f"q must be None (-> p+1) or an integer > p and <= 20; got {self.q!r}."
)
_normalize_kernel(self.kernel) # raises on unknown kernel
if self.bwselect not in BWSELECT_OPTIONS:
raise ValueError(f"bwselect must be one of {BWSELECT_OPTIONS}; got {self.bwselect!r}.")
for name, val in (("h", self.h), ("b", self.b), ("rho", self.rho)):
if val is not None and not (self._is_real_scalar(val) and np.isfinite(val) and val > 0):
raise ValueError(f"{name} must be None or finite and > 0; got {val!r}.")
if self.vcov_type != "nn":
raise NotImplementedError(
"Only vcov_type='nn' (rdrobust's default nearest-neighbor "
"variance) is implemented in this release; 'hc0'-'hc3' and "
"cluster-robust modes are a documented seam."
)
if not (self._is_int_scalar(self.nnmatch) and self.nnmatch >= 1):
raise ValueError(f"nnmatch must be an integer >= 1; got {self.nnmatch!r}.")
if self.masspoints not in ("adjust", "check", "off"):
raise ValueError(
f"masspoints must be 'adjust', 'check', or 'off'; got {self.masspoints!r}."
)
if self.bwcheck is not None and not (
self._is_int_scalar(self.bwcheck) and self.bwcheck >= 1
):
raise ValueError(f"bwcheck must be None or an integer >= 1; got {self.bwcheck!r}.")
if not isinstance(self.bwrestrict, (bool, np.bool_)):
# No silent truthiness: a string like "False" must not coerce
# to bandwidth-restriction ON.
raise ValueError(f"bwrestrict must be a bool; got {self.bwrestrict!r}.")
if not isinstance(self.sharpbw, (bool, np.bool_)):
raise ValueError(f"sharpbw must be a bool; got {self.sharpbw!r}.")
if not isinstance(self.covs_drop, (bool, np.bool_)):
raise ValueError(f"covs_drop must be a bool; got {self.covs_drop!r}.")
if not (
self._is_real_scalar(self.scaleregul)
and np.isfinite(self.scaleregul)
and self.scaleregul >= 0
):
raise ValueError(f"scaleregul must be finite and >= 0; got {self.scaleregul!r}.")
if not (self._is_real_scalar(self.alpha) and 0 < self.alpha < 1):
raise ValueError(f"alpha must be in (0, 1); got {self.alpha!r}.")
[docs]
def get_params(self, deep: bool = True) -> Dict[str, Any]:
"""Return raw constructor parameters (sklearn-compatible)."""
del deep
return {
"cutoff": self.cutoff,
"p": self.p,
"q": self.q,
"kernel": self.kernel,
"bwselect": self.bwselect,
"h": self.h,
"b": self.b,
"rho": self.rho,
"vcov_type": self.vcov_type,
"nnmatch": self.nnmatch,
"masspoints": self.masspoints,
"bwcheck": self.bwcheck,
"bwrestrict": self.bwrestrict,
"scaleregul": self.scaleregul,
"sharpbw": self.sharpbw,
"covs_drop": self.covs_drop,
"alpha": self.alpha,
}
[docs]
def set_params(self, **params: Any) -> "RegressionDiscontinuity":
"""Transactionally update parameters (validate before mutating)."""
valid = set(self.get_params().keys())
unknown = set(params) - valid
if unknown:
raise ValueError(f"Unknown parameter(s): {sorted(unknown)}. Valid: {sorted(valid)}.")
merged = self.get_params()
merged.update(params)
type(self)(**merged) # dry-run: raises before any mutation
for key, value in params.items():
setattr(self, key, value)
return self
# ------------------------------------------------------------------
# Fitting
# ------------------------------------------------------------------
[docs]
def fit(
self,
data: pd.DataFrame,
outcome_col: str,
running_col: str,
treatment_col: Optional[str] = None,
covariates: Optional[List[str]] = None,
) -> RegressionDiscontinuityResults:
"""Estimate the RD effect at the cutoff (sharp or fuzzy, optionally
covariate-adjusted).
Parameters
----------
data : pd.DataFrame
Cross-sectional data.
outcome_col, running_col : str
Column names of the outcome and the running variable.
treatment_col : str or None, default None
``None`` (sharp design): treatment is derived as
``running >= cutoff``; no treatment column is needed. A column
name activates the FUZZY design: the column holds the OBSERVED
treatment take-up (typically binary, any numeric accepted,
matching R's ``fuzzy=``), the estimand becomes the local Wald
ratio, and the results gain the ``first_stage*`` block. The
``estimand`` label is data-dependent: for BINARY take-up it
reads "fuzzy (LATE for compliers at the cutoff)" (the
monotonicity-based complier reading); for non-binary (dose)
take-up it reads "fuzzy (local Wald ratio at the cutoff;
non-binary take-up)" - the complier-LATE interpretation does
not apply there. A take-up column that is deterministic in
the running variable reproduces the sharp fit exactly
(first stage == 1).
covariates : list of str or None, default None
Column names of pre-determined covariates for the additive
common-coefficient adjustment of CCFT (2019) (R's ``covs=``).
The estimand is UNCHANGED - unlike the DiD estimators'
conditional-parallel-trends role, RD covariates buy precision
only, and require covariate BALANCE at the cutoff (zero RD
effect on each covariate; testable by fitting each covariate
as the outcome - imbalanced covariates make the adjusted
estimator inconsistent). Continuous, discrete, or mixed
columns are accepted; covariates propagate into bandwidth
selection (covariate-aware, as in R). Collinear columns are
dropped with a warning under ``covs_drop=True``; see the
``covariates*`` results fields for the echo and the fitted
projection coefficients.
"""
cols = [outcome_col, running_col]
if treatment_col is not None:
cols.append(treatment_col)
if covariates is not None:
if isinstance(covariates, str):
# A bare string would iterate characters; fail closed.
raise ValueError(f"covariates must be a list of column names; got {covariates!r}.")
# Materialize BEFORE validating: a generator would be consumed
# by the all() check and then silently collapse to an empty
# list (disabling adjustment without a whisper).
covariates = list(covariates)
if not all(isinstance(name, str) for name in covariates):
raise ValueError(f"covariates must be a list of column names; got {covariates!r}.")
if not covariates:
covariates = None # empty list == no adjustment
if covariates is not None:
# Duplicate names and collisions with the fit's structural
# columns corrupt the name-keyed coefficient dict.
validate_covariate_names(
covariates,
cols,
estimator="RegressionDiscontinuity",
)
cols.extend(covariates)
for col in cols:
if col not in data.columns:
raise ValueError(f"Column {col!r} not found in data.")
fuzzy_fit = treatment_col is not None
if self.sharpbw and not fuzzy_fit:
# Deviation from R, which silently ignores sharpbw on sharp
# fits (no-silent-failures policy; same pattern as b-without-h).
warnings.warn(
"sharpbw has no effect without treatment_col (sharp design) " "and is ignored.",
UserWarning,
stacklevel=2,
)
if not self.covs_drop and covariates is None:
# Same pattern as sharpbw-on-sharp: a non-default knob that
# cannot apply must not pass silently.
warnings.warn(
"covs_drop=False has no effect without covariates and is ignored.",
UserWarning,
stacklevel=2,
)
y_raw = np.asarray(pd.to_numeric(data[outcome_col], errors="coerce"), dtype=np.float64)
x_raw = np.asarray(pd.to_numeric(data[running_col], errors="coerce"), dtype=np.float64)
ok = np.isfinite(y_raw) & np.isfinite(x_raw)
t_raw: Optional[np.ndarray] = None
if fuzzy_fit:
# R's complete.cases filter includes the fuzzy column
# (rdrobust.R:86-89) - the joint drop must too.
t_raw = np.asarray(
pd.to_numeric(data[treatment_col], errors="coerce"), dtype=np.float64
)
ok = ok & np.isfinite(t_raw)
z_raw: Optional[np.ndarray] = None
if covariates is not None:
# R's complete.cases filter includes the covariate columns
# (rdrobust.R:80-84) - the joint drop must too. Column order
# here is AS PASSED; the R name-length sort applies below.
z_raw = np.column_stack(
[
np.asarray(pd.to_numeric(data[name], errors="coerce"), dtype=np.float64)
for name in covariates
]
)
ok = ok & np.all(np.isfinite(z_raw), axis=1)
n_dropped = int(y_raw.shape[0] - np.sum(ok))
if n_dropped > 0:
# Deviation from R (which drops silently via complete.cases):
dropped_cols = f"{outcome_col!r}/{running_col!r}"
if fuzzy_fit:
dropped_cols += f"/{treatment_col!r}"
if covariates is not None:
dropped_cols += "/covariates"
warnings.warn(
f"Dropping {n_dropped} row(s) with missing or non-numeric "
f"values in {dropped_cols}.",
UserWarning,
stacklevel=2,
)
y = y_raw[ok]
x = x_raw[ok]
t = t_raw[ok] if t_raw is not None else None
z = z_raw[ok] if z_raw is not None else None
N = int(y.shape[0])
if N == 0:
raise ValueError("No complete-case observations to fit on.")
c = float(self.cutoff)
if not (np.min(x) <= c <= np.max(x)):
raise ValueError(
f"cutoff={c:g} lies outside the observed running-variable "
f"range [{np.min(x):g}, {np.max(x):g}]."
)
p = int(self.p)
q = int(self.q) if self.q is not None else p + 1
kernel = _normalize_kernel(self.kernel)
# --- Covariate column sort + redundant-column drop (hoisted from
# rdrobust.R:121-140, like the fuzzy identification hoist below;
# R's order: NaN drop -> covs_drop -> fuzzy stop -> mass points).
# Under covs_drop=True R first sorts columns by NAME LENGTH
# (order(nchar), stable - rdrobust.R:131); the sort decides which
# of a collinear set survives, and all user-facing surfaces are
# name-keyed so the internal order never leaks. The QR runs on
# x-SORTED rows - the row order R (and the port entry points) use
# - so near-threshold rank decisions cannot diverge from the
# downstream calls. Passing the already-reduced matrix down means
# the port's own entry-point drop finds full rank and stays
# silent (no double warning).
model_covariates: Optional[List[str]] = None
covariates_dropped: Optional[List[str]] = None
if covariates is not None:
assert z is not None
model_covariates = list(covariates)
covariates_dropped = []
if self.covs_drop:
model_covariates = sorted(model_covariates, key=len)
z = np.column_stack([z[:, covariates.index(name)] for name in model_covariates])
keep_idx, rank = covs_drop_fun(z[np.argsort(x, kind="stable")])
if rank == 0:
raise ValueError(
"All covariates are numerically zero (rank-0 "
"covariate matrix); remove the covariates instead."
)
if rank < len(model_covariates):
covariates_dropped = [
name
for i, name in enumerate(model_covariates)
if i not in set(keep_idx.tolist())
]
# R's warning is a generic "Multicollinearity issue
# detected in covs." - naming the dropped columns is a
# documented enhancement.
warnings.warn(
"Multicollinearity detected in covariates: "
f"dropped redundant column(s) {covariates_dropped} "
"(covs_drop=True; set covs_drop=False for a strict "
"error instead).",
UserWarning,
stacklevel=2,
)
model_covariates = [
name
for i, name in enumerate(model_covariates)
if i in set(keep_idx.tolist())
]
z = z[:, keep_idx]
# --- Fuzzy identification check (rdrobust.R:164-185) ---
# Hoisted to run immediately after the NaN drop and BEFORE
# mass-point detection, matching R's rdrobust ordering exactly
# (live-verified: R raises this with NO mass-point warning on
# degenerate fuzzy + tied data). The port re-checks defensively.
if t is not None:
_fuzzy_identification_stop(t[x < c], t[x >= c])
# --- Mass points (rdrobust.R:365-380) ---
# R's rdrobust() runs this detection ITSELF, before the manual-vs-
# data-driven bandwidth branch, so the warning fires on manual-h
# fits too (verified against installed 4.0.0). The port's
# rdbwselect-level copy is silenced below (warn_masspoints=False)
# to mirror R's single warning from the estimation call.
n_left_pre = int(np.sum(x < c))
n_right_pre = int(np.sum(x >= c))
n_unique_left = int(np.unique(x[x < c]).shape[0])
n_unique_right = int(np.unique(x[x >= c]).shape[0])
if self.masspoints in ("check", "adjust") and n_left_pre > 0 and n_right_pre > 0:
mass_l = 1.0 - n_unique_left / n_left_pre
mass_r = 1.0 - n_unique_right / n_right_pre
if mass_l >= 0.2 or mass_r >= 0.2:
warnings.warn(
"Mass points detected in the running variable.",
UserWarning,
stacklevel=2,
)
if self.masspoints == "check":
warnings.warn(
"Try using option masspoints='adjust'.",
UserWarning,
stacklevel=2,
)
# --- Bandwidth resolution (rdrobust.R:295-307, 501-504) ---
h_user, b_user, rho = self.h, self.b, self.rho
bwselect_label = self.bwselect
if b_user is not None and h_user is None:
# R silently ignores b without h; we warn (documented deviation).
warnings.warn(
"b= was supplied without h= and is ignored (matching "
"rdrobust's behavior); supply h= to use a manual bias "
"bandwidth.",
UserWarning,
stacklevel=2,
)
b_user = None
if N < 20:
# rdrobust.R:303-307: unconditional override, INCLUDING a
# user-supplied h (the block runs after manual-h resolution).
warnings.warn(
"Not enough observations to perform bandwidth calculations. "
"Estimates computed using entire sample.",
UserWarning,
stacklevel=2,
)
x_min, x_max = float(np.min(x)), float(np.max(x))
full = max(abs(c - x_min), abs(c - x_max))
h_l = h_r = b_l = b_r = full
bwselect_label = "Manual"
elif h_user is not None:
bwselect_label = "Manual"
if rho is None and b_user is None:
b_resolved = h_user # rho = 1 (rdrobust.R:296-299)
elif rho is not None:
if b_user is not None:
warnings.warn(
"Both b= and rho= supplied with h=; rho takes "
"precedence (b = h/rho), matching rdrobust.",
UserWarning,
stacklevel=2,
)
b_resolved = h_user / rho # rdrobust.R:300
else:
# rho is None and b_user is not None (first branch handled
# the both-None case).
assert b_user is not None
b_resolved = b_user
h_l = h_r = float(h_user)
b_l = b_r = float(b_resolved)
else:
bw = rdbwselect(
y,
x,
c=c,
p=p,
q=q,
kernel=kernel,
vce=self.vcov_type,
nnmatch=int(self.nnmatch),
masspoints=self.masspoints,
bwcheck=None if self.bwcheck is None else int(self.bwcheck),
bwrestrict=bool(self.bwrestrict),
scaleregul=float(self.scaleregul),
warn_masspoints=False, # fit() already warned (rdrobust.R:365-380)
fuzzy=t,
sharpbw=bool(self.sharpbw),
covs=z,
covs_drop=bool(self.covs_drop),
)
h_l, h_r, b_l, b_r = bw.bws[self.bwselect]
n_unique_left = bw.M_l if self.masspoints != "off" else n_unique_left
n_unique_right = bw.M_r if self.masspoints != "off" else n_unique_right
if rho is not None:
# rdrobust.R:501-504: rho applies to the SELECTED bandwidths.
b_l = h_l / rho
b_r = h_r / rho
# --- Estimation (port validates h/b finite and positive) ---
fit = rdrobust_fit(
y,
x,
c,
h_l,
h_r,
b_l,
b_r,
p=p,
q=q,
kernel=kernel,
vce=self.vcov_type,
nnmatch=int(self.nnmatch),
t=t,
covs=z,
covs_drop=bool(self.covs_drop),
# The estimator owns the degeneracy warning (with column
# names) - same plumbing pattern as warn_masspoints.
warn_covs_degenerate=False,
)
# --- Degenerate covariate adjustment warning (estimator-level,
# with column names; deviation from R, which silently inverts a
# noise singular value on these systems - see the port's
# _covs_gamma for the guard) ---
if model_covariates is not None and fit.covs_excluded is not None:
excluded_names = [
name for name, flag in zip(model_covariates, fit.covs_excluded) if bool(flag)
]
parts = []
if excluded_names:
parts.append(
f"covariate(s) {excluded_names} are numerically "
"collinear with the local polynomial design (e.g. "
"constant near the cutoff) and were excluded from "
"the adjustment"
)
if fit.covs_set_degenerate:
parts.append(
"the covariate set is numerically rank-deficient "
"after partialling (e.g. a full dummy set); a "
"stabilized pseudo-inverse cut was used - consider "
"dropping a reference category"
)
if parts:
warnings.warn(
"Degenerate covariate adjustment: " + "; ".join(parts) + ".",
UserWarning,
stacklevel=2,
)
# Estimand label: the complier-LATE reading requires BINARY
# take-up (plus monotonicity); non-binary (dose) take-up - accepted,
# matching R's fuzzy= - is the ratio-of-jumps estimand and must not
# be labeled a LATE (the label is what `att` claims to measure).
if not fuzzy_fit:
estimand = "sharp (ATE at the cutoff)"
elif t is not None and bool(np.all(np.isin(t, (0.0, 1.0)))):
estimand = "fuzzy (LATE for compliers at the cutoff)"
else:
estimand = "fuzzy (local Wald ratio at the cutoff; non-binary take-up)"
alpha = float(self.alpha)
# Three inference rows (rdrobust.R:854-863), each through the
# library's joint-NaN gate:
t_rb, p_rb, ci_rb = safe_inference(fit.tau_bc, fit.se_rb, alpha=alpha)
t_cl, p_cl, ci_cl = safe_inference(fit.tau_cl, fit.se_cl, alpha=alpha)
t_bcm, p_bcm, ci_bcm = safe_inference(fit.tau_bc, fit.se_cl, alpha=alpha)
# --- First-stage rows + weak-identification warning (fuzzy) ---
fs: Dict[str, Any] = {}
if fuzzy_fit:
assert fit.tau_T_bc is not None and fit.tau_T_cl is not None
assert fit.se_T_rb is not None and fit.se_T_cl is not None
fs_t_rb, fs_p_rb, fs_ci_rb = safe_inference(fit.tau_T_bc, fit.se_T_rb, alpha=alpha)
fs_t_cl, fs_p_cl, fs_ci_cl = safe_inference(fit.tau_T_cl, fit.se_T_cl, alpha=alpha)
fs_t_bcm, fs_p_bcm, fs_ci_bcm = safe_inference(fit.tau_T_bc, fit.se_T_cl, alpha=alpha)
fs = dict(
first_stage=fit.tau_T_bc,
first_stage_se=fit.se_T_rb,
first_stage_t_stat=fs_t_rb,
first_stage_p_value=fs_p_rb,
first_stage_conf_int=fs_ci_rb,
first_stage_conventional=fit.tau_T_cl,
first_stage_se_conventional=fit.se_T_cl,
first_stage_t_stat_conventional=fs_t_cl,
first_stage_p_value_conventional=fs_p_cl,
first_stage_conf_int_conventional=fs_ci_cl,
first_stage_t_stat_bias_corrected=fs_t_bcm,
first_stage_p_value_bias_corrected=fs_p_bcm,
first_stage_conf_int_bias_corrected=fs_ci_bcm,
beta_t_p_left=fit.beta_t_p_l,
beta_t_p_right=fit.beta_t_p_r,
)
# Deviation from R (verified silent): warn when the take-up
# jump is not distinguishable from zero at the fit's own alpha
# - the ratio is then unreliable (CCT 2014 Theorem 3 requires
# tau_T != 0; weak-IV-robust inference per Feir-Lemieux-Marmer
# is a documented seam). Gate: FINITE robust CI containing 0,
# so NaN-gated first stages (e.g. perfect compliance's se=0)
# correctly do not fire.
if (
np.isfinite(fs_ci_rb[0])
and np.isfinite(fs_ci_rb[1])
and fs_ci_rb[0] <= 0.0 <= fs_ci_rb[1]
):
warnings.warn(
"Weak first stage: the take-up discontinuity "
f"({fit.tau_T_bc:.4g}, robust CI [{fs_ci_rb[0]:.4g}, "
f"{fs_ci_rb[1]:.4g}]) is not distinguishable from zero "
f"at alpha={alpha:g}; the fuzzy (ratio) estimates are "
"unreliable under weak identification.",
UserWarning,
stacklevel=2,
)
return RegressionDiscontinuityResults(
att=fit.tau_bc,
se=fit.se_rb,
t_stat=t_rb,
p_value=p_rb,
conf_int=ci_rb,
alpha=alpha,
att_conventional=fit.tau_cl,
se_conventional=fit.se_cl,
t_stat_conventional=t_cl,
p_value_conventional=p_cl,
conf_int_conventional=ci_cl,
t_stat_bias_corrected=t_bcm,
p_value_bias_corrected=p_bcm,
conf_int_bias_corrected=ci_bcm,
se_robust=fit.se_rb,
h_left=float(h_l),
h_right=float(h_r),
b_left=float(b_l),
b_right=float(b_r),
n_obs=N,
n_left=int(np.sum(x < c)),
n_right=int(np.sum(x >= c)),
n_h_left=fit.N_h_l,
n_h_right=fit.N_h_r,
n_b_left=fit.N_b_l,
n_b_right=fit.N_b_r,
n_unique_left=n_unique_left,
n_unique_right=n_unique_right,
n_dropped=n_dropped,
cutoff=c,
p=p,
q=q,
kernel=kernel,
bwselect=bwselect_label,
vcov_type=self.vcov_type,
nnmatch=int(self.nnmatch),
masspoints=self.masspoints,
bwcheck=None if self.bwcheck is None else int(self.bwcheck),
bwrestrict=bool(self.bwrestrict),
scaleregul=float(self.scaleregul),
h_input=None if self.h is None else float(self.h),
b_input=None if self.b is None else float(self.b),
rho_input=None if self.rho is None else float(self.rho),
estimand=estimand,
sharpbw=bool(self.sharpbw),
treatment_col=treatment_col,
covs_drop=bool(self.covs_drop),
covariates=None if covariates is None else list(covariates),
covariates_dropped=covariates_dropped,
covariate_coefficients=(
None
if model_covariates is None or fit.gamma_p is None
else {name: float(fit.gamma_p[i, 0]) for i, name in enumerate(model_covariates)}
),
first_stage_covariate_coefficients=(
None
if not fuzzy_fit or model_covariates is None or fit.gamma_p is None
else {name: float(fit.gamma_p[i, 1]) for i, name in enumerate(model_covariates)}
),
beta_p_left=fit.beta_p_l,
beta_p_right=fit.beta_p_r,
**fs,
)