"""
Heterogeneous Adoption Difference-in-Differences (HAD) estimator (Phase 2a).
Implements the de Chaisemartin, Ciccia, D'Haultfoeuille, and Knau (2026)
Weighted-Average-Slope (WAS) estimator with three design-dispatch paths.
All three paths produce a beta-scale point estimate of the form
``(mean(Delta Y) - [boundary limit]) / [expected dose gap]`` (Design 1
family) or the Wald-IV ratio (mass-point), then route inference through
:func:`diff_diff.utils.safe_inference`.
1. Design 1' (``continuous_at_zero``): ``d_lower = 0``, boundary density
continuous at zero, Assumption 3. Theorem 1 / Equation 3
(identification); Equation 7 (sample estimator):
beta = (E[Delta Y] - lim_{d v 0} E[Delta Y | D_2 0``, continuous boundary density, Assumption 5 or 6.
Theorem 3 / Equation 11 (``WAS_{d_lower}`` under Assumption 6;
Theorem 4 is the QUG null test, not this estimand):
beta = (E[Delta Y] - lim_{d v d_lower} E[Delta Y | D_2
native Python via ``.item()``
- ``pd.Timestamp`` / ``pd.Timedelta`` -> ISO 8601 string via
``.isoformat()``
- Everything else returned as-is.
The ``to_dict`` methods use this to keep the returned dict
serializable via ``json.dumps`` regardless of the underlying
pandas/numpy dtype of the time / first_treat columns.
"""
if isinstance(x, (pd.Timestamp, pd.Timedelta)):
return x.isoformat()
if hasattr(x, "item") and callable(getattr(x, "item")):
try:
return x.item()
except (AttributeError, ValueError, TypeError):
return x
return x
def _json_safe_filter_info(
filter_info: Optional[Dict[str, Any]],
) -> Optional[Dict[str, Any]]:
"""Normalize a ``filter_info`` dict to JSON-safe scalars.
Returns ``None`` unchanged; otherwise coerces ``F_last`` and each
entry in ``dropped_cohorts`` via :func:`_json_safe_scalar`. Int
counts are cast to ``int`` for stability.
"""
if filter_info is None:
return None
return {
"F_last": _json_safe_scalar(filter_info.get("F_last")),
"n_kept": int(filter_info.get("n_kept", 0)),
"n_dropped": int(filter_info.get("n_dropped", 0)),
"dropped_cohorts": [_json_safe_scalar(c) for c in filter_info.get("dropped_cohorts", [])],
}
# =============================================================================
# Results dataclass
# =============================================================================
@dataclass
class HeterogeneousAdoptionDiDResults:
"""Estimator output for :class:`HeterogeneousAdoptionDiD`.
NaN-safe inference: the three downstream fields ``t_stat``,
``p_value``, and ``conf_int`` are routed through
:func:`diff_diff.utils.safe_inference`, which returns NaN on all
three whenever ``se`` is non-finite, zero, or negative. ``att`` and
``se`` themselves are RAW estimator outputs from the chosen fit
path and are NOT gated by ``safe_inference``:
- On the degenerate fit configurations (constant outcome on the
continuous paths, all-units-at-d_lower / no-dose-variation on the
mass-point path), the fit path explicitly returns
``(att=nan, se=nan)``, which combined with the safe-inference
gate yields all five fields NaN together.
- On the degenerate CR1 cluster configuration (mass-point path
with a single cluster), ``_fit_mass_point_2sls`` returns
``(att=beta_hat, se=nan)`` - ``att`` stays finite because the
Wald-IV ratio is well defined, but the cluster-robust SE is
not, so ``se`` is NaN and the downstream triple becomes NaN
via the safe-inference gate.
So the guaranteed NaN coupling is on the downstream triple
(``t_stat``, ``p_value``, ``conf_int``), not on ``att``. The
``assert_nan_inference`` fixture in ``tests/conftest.py`` checks
the downstream triple against the gate contract and does not
assume ``att`` is NaN.
Attributes
----------
att : float
Point estimate of the WAS parameter on the beta-scale.
- Design 1' (paper Theorem 1 / Equation 3 identification;
Equation 7 sample estimator):
``att = (mean(Y) - tau_bc) / D_bar``
where ``tau_bc`` is the bias-corrected local-linear estimate
of ``lim_{d v 0} E[Y | D_2 d_lower``.
``None`` on continuous paths.
inference_method : str
``"analytical_nonparametric"`` (continuous designs) or
``"analytical_2sls"`` (mass-point).
vcov_type : str or None
Effective variance-covariance family used. ``None`` on continuous
paths (they use the CCT-2014 robust SE from Phase 1c, not the
library's ``vcov_type`` enum). Mass-point: ``"classical"`` or
``"hc1"`` when ``cluster`` is not supplied, and ``"cr1"``
whenever ``cluster`` is supplied (cluster-robust CR1 is computed
regardless of the requested ``vcov_type`` because
classical/hc1 + cluster collapses to the same CR1 sandwich).
Downstream consumers reading ``result.to_dict()`` can inspect
this field directly to determine the effective SE family.
cluster_name : str or None
Column name of the cluster variable on the mass-point path when
cluster-robust SE is requested. ``None`` otherwise.
survey_metadata : SurveyMetadata or None
Repo-standard survey metadata dataclass from
:class:`diff_diff.survey.SurveyMetadata`. ``None`` when ``fit()``
was called without ``survey=`` or ``weights=``; populated on the
continuous-dose weighted paths via
:func:`diff_diff.survey.compute_survey_metadata`. Exposes
``weight_type``, ``effective_n``, ``design_effect``,
``sum_weights``, ``n_strata``, ``n_psu``, ``weight_range``, and
``df_survey`` for downstream reporting consumers (BusinessReport,
DiagnosticReport) that read these fields via attribute access.
HAD-specific inference-method info (pweight vs Binder-TSL) is
carried on ``inference_method`` and ``variance_formula``.
bandwidth_diagnostics : BandwidthResult or None
Full Phase 1b MSE-DPI selector output on the continuous paths
(when bandwidths were auto-selected). ``None`` on the mass-point
path (parametric, no bandwidth).
bias_corrected_fit : BiasCorrectedFit or None
Full Phase 1c bias-corrected local-linear fit on the continuous
paths. ``None`` on the mass-point path.
"""
# Point estimate + inference (safe_inference-gated)
att: float
se: float
t_stat: float
p_value: float
conf_int: Tuple[float, float]
alpha: float
# Design metadata
design: str
target_parameter: str
d_lower: float
dose_mean: float
# Sample counts
n_obs: int
n_treated: int
n_control: int
n_mass_point: Optional[int]
n_above_d_lower: Optional[int]
# Inference metadata
inference_method: str
vcov_type: Optional[str]
cluster_name: Optional[str]
survey_metadata: Optional[SurveyMetadata]
# Nonparametric-only diagnostics
bandwidth_diagnostics: Optional[BandwidthResult]
bias_corrected_fit: Optional[BiasCorrectedFit]
# Phase 4.5 weighted-path extras (optional so unweighted fits stay unchanged)
variance_formula: Optional[str] = None
"""HAD-specific label for the SE formula on weighted fits, populated
on BOTH continuous and mass-point designs (Phase 4.5 A / B):
``"pweight"`` (continuous, weighted-robust CCT 2014 under the
``weights=`` shortcut), ``"survey_binder_tsl"`` (continuous, Binder
1983 TSL with PSU/strata/FPC under ``survey_design=SurveyDesign(...)``),
``"pweight_2sls"`` (mass-point, weighted 2SLS HC1/CR1 sandwich
under the ``weights=`` shortcut), or ``"survey_binder_tsl_2sls"``
(mass-point, Binder 1983 TSL under ``survey_design=``). ``None`` on
unweighted fits. Orthogonal to ``survey_metadata`` which is the
repo-standard :class:`diff_diff.survey.SurveyMetadata` shared with
downstream report/diagnostic consumers (no HAD-specific leakage)."""
effective_dose_mean: Optional[float] = None
"""Weighted denominator used by the beta-scale rescaling, populated
on weighted fits across all designs: ``sum(w_g D_g) / sum(w_g)``
on ``continuous_at_zero``, ``sum(w_g (D_g - d_lower)) / sum(w_g)``
on ``continuous_near_d_lower``, and the weighted Wald-IV dose gap
``mean(D | Z=1, w) - mean(D | Z=0, w)`` on ``mass_point`` (where
``Z = 1{D > d_lower}``). On the continuous designs reduces
bit-exactly to ``dose_mean`` / ``mean(D - d_lower)`` when weights
are uniform or absent. ``None`` when ``fit()`` was called without
``survey_design=`` / ``survey=`` / ``weights=`` (use ``dose_mean``
there). Exists because ``dose_mean`` is the raw sample mean of the
dose column; under weighted fits the estimator's actual denominator
is the weighted form above, and users reconstructing the -scale
value by hand need the weighted one."""
def __repr__(self) -> str:
base = (
f"HeterogeneousAdoptionDiDResults("
f"att={self.att:.4f}, se={self.se:.4f}, "
f"design={self.design!r}, n_obs={self.n_obs}"
)
# Surface weighted-path identity when the fit was weighted, so the
# one-line repr makes it unambiguous which inference family was
# used (pweight-shortcut vs full Binder-TSL survey) and the
# effective denominator flows into ad-hoc log output.
if self.variance_formula is not None:
base += f", variance_formula={self.variance_formula!r}"
if self.effective_dose_mean is not None:
base += f", effective_dose_mean={self.effective_dose_mean:.4g}"
return base + ")"
def summary(self) -> str:
"""Formatted summary table."""
width = 72
conf_level = int((1 - self.alpha) * 100)
lines = [
"=" * width,
"HeterogeneousAdoptionDiD Estimation Results".center(width),
"=" * width,
"",
f"{'Design:':20}",
f"{'Target parameter:':20}",
f"{'d_lower:':20.6g}",
f"{'D_bar (dose mean):':20.6g}",
f"{'Observations (units):':20}",
f"{'Above d_lower:':20}",
f"{'At/below d_lower:':20}",
]
if self.n_mass_point is not None:
lines.append(f"{'At d_lower (mass point):':20}")
if self.n_above_d_lower is not None:
lines.append(f"{'Strictly above d_lower:':20}")
lines.append(f"{'Inference method:':20}")
if self.vcov_type is not None:
if self.cluster_name is not None:
# Cluster-robust (CR1): the stored vcov_type is already "cr1",
# but render with the cluster column for clarity.
label = f"CR1 at {self.cluster_name}"
else:
label = self.vcov_type
lines.append(f"{'Variance:':20}")
if self.bandwidth_diagnostics is not None:
bw = self.bandwidth_diagnostics
lines.append(f"{'Bandwidth h (MSE-DPI):':20.6g}")
lines.append(f"{'Bandwidth b (bias):':20.6g}")
if self.bias_corrected_fit is not None:
bc = self.bias_corrected_fit
lines.append(f"{'Bandwidth h used:':20.6g}")
lines.append(f"{'Obs in window (n_used):':20}")
if self.survey_metadata is not None:
sm = self.survey_metadata
vf_label = self.variance_formula or "unknown"
lines.append(f"{'Variance formula:':20}")
lines.append(f"{'Effective sample size:':20.6g}")
if self.effective_dose_mean is not None:
lines.append(
f"{'Weighted D (denominator):':20.6g}"
)
if sm.df_survey is not None:
lines.append(f"{'Survey df:':20}")
param_label = self.target_parameter
lines.extend(
[
"",
"-" * width,
(
f"{'Parameter':12} {'Std. Err.':>12} "
f"{'t-stat':>10} {'P>|t|':>10}"
),
"-" * width,
(
f"{param_label:12.4f} {self.se:>12.4f} "
f"{self.t_stat:>10.3f} {self.p_value:>10.4f}"
),
"-" * width,
"",
(
f"{conf_level}% Confidence Interval: "
f"[{self.conf_int[0]:.4f}, {self.conf_int[1]:.4f}]"
),
"=" * width,
]
)
return "\n".join(lines)
def print_summary(self) -> None:
"""Print the summary to stdout."""
print(self.summary())
def to_dict(self) -> Dict[str, Any]:
"""Return results as a dict of scalars + weighted-path surfaces.
Always-present keys mirror the dataclass fields: ``att``, ``se``,
``t_stat``, ``p_value``, ``conf_int_lower`` / ``conf_int_upper``,
``alpha``, ``design``, ``target_parameter``, ``d_lower``,
``dose_mean``, ``n_obs`` / ``n_treated`` / ``n_control`` /
``n_mass_point`` / ``n_above_d_lower``, ``inference_method``,
``vcov_type``, ``cluster_name``.
Weighted-path keys (``None`` on unweighted fits):
- ``survey_metadata``: repo-standard
:class:`diff_diff.survey.SurveyMetadata` dataclass (object, not
dict) carrying ``weight_type`` / ``effective_n`` /
``design_effect`` / ``sum_weights`` / ``weight_range`` +
``n_strata`` / ``n_psu`` / ``df_survey`` (latter three
``None`` on the ``weights=`` shortcut).
- ``variance_formula``: HAD-specific SE label, populated on BOTH
continuous and mass-point designs (Phase 4.5 A / B):
``"pweight"`` (continuous, weighted-robust CCT 2014 under
``weights=``), ``"survey_binder_tsl"`` (continuous, Binder
1983 TSL under ``survey_design=``), ``"pweight_2sls"``
(mass-point, weighted 2SLS HC1/CR1 sandwich under ``weights=``),
or ``"survey_binder_tsl_2sls"`` (mass-point, Binder 1983 TSL
under ``survey_design=``). See the field docstring above for
the full contract.
- ``effective_dose_mean``: weighted denominator used by the
beta-scale rescaling - weighted ``mean(D)`` on
``continuous_at_zero``, weighted ``mean(D - d_lower)`` on
``continuous_near_d_lower``, or the weighted Wald-IV dose gap
``mean(D | Z=1, w) - mean(D | Z=0, w)`` on ``mass_point``."""
return {
"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,
"design": self.design,
"target_parameter": self.target_parameter,
"d_lower": self.d_lower,
"dose_mean": self.dose_mean,
"n_obs": self.n_obs,
"n_treated": self.n_treated,
"n_control": self.n_control,
"n_mass_point": self.n_mass_point,
"n_above_d_lower": self.n_above_d_lower,
"inference_method": self.inference_method,
"vcov_type": self.vcov_type,
"cluster_name": self.cluster_name,
"survey_metadata": self.survey_metadata,
"variance_formula": self.variance_formula,
"effective_dose_mean": self.effective_dose_mean,
}
def to_dataframe(self) -> pd.DataFrame:
"""Return a one-row DataFrame of the result dict."""
return pd.DataFrame([self.to_dict()])
@dataclass
class HeterogeneousAdoptionDiDEventStudyResults:
"""Event-study results for :class:`HeterogeneousAdoptionDiD` (Phase 2b).
Per-horizon arrays align with ``event_times`` by index; all per-horizon
arrays have shape ``(n_horizons,)``. The anchor horizon ``e = -1``
(i.e., ``t = F - 1``) is NOT included because
``Y_{g, F-1} - Y_{g, F-1} = 0`` trivially and the WAS is not identified
there.
Per-horizon inference fields (``t_stat``, ``p_value``, ``conf_int_low``,
``conf_int_high``) are NaN-coupled to the per-horizon ``se`` via
:func:`diff_diff.utils.safe_inference`; ``att`` and ``se`` themselves
are raw estimator outputs from the chosen design path on each
horizon's first differences.
Design resolution is SHARED across horizons: the design, ``d_lower``,
``target_parameter``, and ``inference_method`` are single scalars
determined once from the post-period dose distribution ``D_{g, F}``
(paper Appendix B.2 convention the dose regressor is invariant
across event-time horizons).
Attributes
----------
event_times : np.ndarray, shape (n_horizons,)
Integer event-time labels ``e = t - F``, sorted ascending.
Excludes ``e = -1`` (the anchor). Post-period horizons have
``e >= 0``; pre-period placebos have ``e 12} "
f"{'t-stat':>10} {'P>|t|':>10} "
f"{str(conf_level) + '% CI':>22}"
),
"-" * width,
]
)
for i, e in enumerate(self.event_times):
ci_str = f"[{self.conf_int_low[i]:.4f}, {self.conf_int_high[i]:.4f}]"
# Default float formatting renders non-finite values as "nan";
# we do not override this here since the column width is fixed
# and lowercase "nan" is unambiguous.
se_i = self.se[i]
t_i = self.t_stat[i]
p_i = self.p_value[i]
lines.append(
f"{int(e):>10} {self.att[i]:>12.4f} "
f"{se_i:>12.4f} {t_i:>10.3f} {p_i:>10.4f} {ci_str:>22}"
)
lines.extend(
[
"-" * width,
"",
"=" * width,
]
)
return "\n".join(lines)
def print_summary(self) -> None:
"""Print the summary to stdout."""
print(self.summary())
def to_dict(self) -> Dict[str, Any]:
"""Return results as a dict with per-horizon arrays and scalars.
Per-horizon arrays are converted to Python lists via
``ndarray.tolist()`` (which unwraps NumPy scalar elements to
native ``int`` / ``float``); scalar fields are coerced to
native Python types via ``_json_safe_scalar`` where relevant
(NumPy scalars -> ``.item()``, pandas ``Timestamp`` -> ISO
string, ``Timedelta`` -> ISO string). The returned dict is
JSON-serializable directly via ``json.dumps``.
"""
return {
"event_times": self.event_times.tolist(),
"att": self.att.tolist(),
"se": self.se.tolist(),
"t_stat": self.t_stat.tolist(),
"p_value": self.p_value.tolist(),
"conf_int_low": self.conf_int_low.tolist(),
"conf_int_high": self.conf_int_high.tolist(),
"n_obs_per_horizon": self.n_obs_per_horizon.tolist(),
"alpha": float(self.alpha),
"design": self.design,
"target_parameter": self.target_parameter,
"d_lower": float(self.d_lower),
"dose_mean": float(self.dose_mean),
"F": _json_safe_scalar(self.F),
"n_units": int(self.n_units),
"inference_method": self.inference_method,
"vcov_type": self.vcov_type,
"cluster_name": self.cluster_name,
"filter_info": _json_safe_filter_info(self.filter_info),
# Phase 4.5 B weighted/survey-path surfaces (None on
# unweighted fits). The full SurveyMetadata dataclass is
# carried as an object, matching the static-path ``to_dict``
# contract consumers read attributes uniformly.
"survey_metadata": self.survey_metadata,
"variance_formula": self.variance_formula,
"effective_dose_mean": self.effective_dose_mean,
"cband_low": (self.cband_low.tolist() if self.cband_low is not None else None),
"cband_high": (self.cband_high.tolist() if self.cband_high is not None else None),
"cband_crit_value": self.cband_crit_value,
"cband_method": self.cband_method,
"cband_n_bootstrap": self.cband_n_bootstrap,
}
def to_dataframe(self) -> pd.DataFrame:
"""Return a tidy per-horizon DataFrame.
Columns: ``event_time, att, se, t_stat, p_value, conf_int_low,
conf_int_high, n_obs``. One row per event-time horizon. On the
weighted event-study path with ``cband=True``, also includes
``cband_low`` and ``cband_high`` columns.
"""
data: Dict[str, Any] = {
"event_time": self.event_times,
"att": self.att,
"se": self.se,
"t_stat": self.t_stat,
"p_value": self.p_value,
"conf_int_low": self.conf_int_low,
"conf_int_high": self.conf_int_high,
"n_obs": self.n_obs_per_horizon,
}
if self.cband_low is not None:
data["cband_low"] = self.cband_low
if self.cband_high is not None:
data["cband_high"] = self.cband_high
return pd.DataFrame(data)
# =============================================================================
# Panel validation and aggregation
# =============================================================================
def _validate_had_panel(
data: pd.DataFrame,
outcome_col: str,
dose_col: str,
time_col: str,
unit_col: str,
first_treat_col: Optional[str],
) -> Tuple[int, int]:
"""Validate a HAD panel and return ``(t_pre, t_post)``.
Enforces the Phase 2a panel contract:
- All required columns present.
- Exactly two distinct time periods. Staggered timing (``>2`` periods)
with ``first_treat_col=None`` raises; with ``first_treat_col`` it
also raises (multi-period reduction is Phase 2b).
- Balanced panel (all units observed at both periods).
- ``D_{g, t_pre} == 0`` for all units (HAD no-unit-untreated pre-period).
- No NaN in outcome, dose, or unit columns.
Parameters
----------
data : pd.DataFrame
outcome_col, dose_col, time_col, unit_col : str
first_treat_col : str or None
Optional column for cross-validation. Supplied column must contain
``0`` for never-treated and the post-period value for treated units.
Returns
-------
tuple[Any, Any]
``(t_pre, t_post)`` - the two period identifiers identified by
the HAD dose invariant (``t_pre`` is the period with dose == 0
for all units; ``t_post`` is the other period). Supports
arbitrary-dtype period labels (int, str, datetime, etc.) rather
than relying on ordinal / lexicographic sort.
Raises
------
ValueError
"""
required = [outcome_col, dose_col, time_col, unit_col]
if first_treat_col is not None:
required.append(first_treat_col)
missing = [c for c in required if c not in data.columns]
if missing:
raise ValueError(f"Missing column(s) in data: {missing}. Required: {required}.")
periods_list = list(data[time_col].unique())
if len(periods_list) < 2:
raise ValueError(
f"HAD requires a two-period panel; got {len(periods_list)} distinct "
f"period(s) in column {time_col!r}."
)
if len(periods_list) > 2:
raise ValueError(
f"HAD with aggregate='overall' requires exactly two time "
f"periods (got {len(periods_list)} in {time_col!r}). For "
f"multi-period panels, pass aggregate='event_study' (paper "
f"Appendix B.2 multi-period event-study extension) which "
f"produces per-event-time WAS estimates."
)
# Balanced-panel check: every unit appears exactly once per period.
counts = data.groupby([unit_col, time_col]).size()
if (counts != 1).any():
n_bad = int((counts != 1).sum())
raise ValueError(
f"Unbalanced panel: {n_bad} unit-period cells have != 1 "
f"observation. HAD requires a balanced two-period panel "
f"(each unit observed exactly once at each period)."
)
unit_counts = data.groupby(unit_col)[time_col].nunique()
incomplete = unit_counts[unit_counts != 2]
if len(incomplete) > 0:
raise ValueError(
f"Unbalanced panel: {len(incomplete)} unit(s) do not appear "
f"in both periods. HAD requires a balanced two-period panel."
)
# NaN checks on key columns.
for col in [outcome_col, dose_col, unit_col]:
if bool(data[col].isna().any()):
n_nan = int(data[col].isna().sum())
raise ValueError(
f"{n_nan} NaN value(s) found in column {col!r}. HAD "
f"does not silently drop NaN rows; drop or impute before "
f"calling fit()."
)
# Identify t_pre and t_post by the HAD invariant rather than by
# lexicographic sort on the time labels: D_{g, t_pre} = 0 for all
# units (paper Section 2 no-unit-untreated pre-period convention).
# Sorting labels alphabetically reverses valid chronologies like
# ("pre", "post") where ordering is semantic, not alphabetic.
per_period_nonzero: Dict[Any, int] = {}
for p in periods_list:
p_doses = np.asarray(data.loc[data[time_col] == p, dose_col], dtype=np.float64)
per_period_nonzero[p] = int((p_doses != 0).sum())
all_zero_periods = [p for p, nz in per_period_nonzero.items() if nz == 0]
if len(all_zero_periods) == 0:
# Neither period has all-zero dose: HAD pre-period contract violated.
stats_str = ", ".join(f"{p!r}: {nz} nonzero" for p, nz in per_period_nonzero.items())
raise ValueError(
f"HAD requires D_{{g,1}} = 0 for all units (pre-period "
f"untreated). Neither period in column {time_col!r} has "
f"all-zero dose ({stats_str}). Exactly one period must be "
f"the pre-treatment period with D_{{g,1}} = 0 for every unit; "
f"drop rows with nonzero pre-period dose or verify the dose "
f"column."
)
if len(all_zero_periods) == 2:
raise ValueError(
f"HAD requires variation in D_{{g,2}} for estimation. Both "
f"periods in column {time_col!r} have all-zero dose, so "
f"there is no treatment assignment to estimate."
)
t_pre = all_zero_periods[0]
t_post = [p for p in periods_list if p != t_pre][0]
# Post-period nonnegative-dose check on the ORIGINAL (unshifted) dose
# scale. Front-door rejection per paper Assumption (dose definition
# Section 2) which treats D_{g,2} as nonnegative. Without this
# check, negative original doses would only surface after the
# regressor shift in ``_fit_continuous`` via Phase 1c's
# ``_validate_had_inputs``, which references the shifted values
# and would confuse users about which column is malformed.
post_mask = data[time_col] == t_post
post_doses = np.asarray(data.loc[post_mask, dose_col], dtype=np.float64)
neg_post = post_doses < 0
if neg_post.any():
n_neg = int(neg_post.sum())
min_neg = float(post_doses[neg_post].min())
raise ValueError(
f"HAD requires D_{{g,2}} >= 0 for all units (paper Section "
f"2 dose definition). {n_neg} unit(s) have negative post-"
f"period dose at t_post={t_post} (min={min_neg!r}). Drop "
f"these units or verify the dose column."
)
# Optional value-domain validation via first_treat_col: if supplied,
# every unit's first_treat value must be in {0, t_post} (0 = never
# treated, t_post = treated in the second period). This is a value-
# domain check that catches typos and staggered-timing mix-ups; it
# does NOT cross-validate first_treat against post-period dose
# (D_{g, t_post} remains the primary signal). Extended cross-checks
# are queued for a follow-up PR. The check is DTYPE-AGNOSTIC: it uses
# pd.isna() for missingness and raw-value membership against
# {0, t_post} so that string-labelled periods (e.g., ("A", "B")) with
# first_treat in {0, "B"} are supported.
if first_treat_col is not None:
# Row-level NaN check: `groupby().first()` skips NaNs silently, so a
# unit with rows [valid, NaN] would pass a collapsed check. Validate
# raw per-row values first, then verify per-unit constancy with
# `nunique(dropna=False)` so within-unit NaN variation is caught.
ft_raw = data[first_treat_col]
if bool(ft_raw.isna().any()):
n_nan = int(ft_raw.isna().sum())
raise ValueError(
f"first_treat_col={first_treat_col!r} contains "
f"{n_nan} NaN value(s) at the row level. Use 0 for "
f"never-treated units and t_post for treated, and drop "
f"or impute any NaN rows before calling fit()."
)
# Row-level domain check: every row (not just the collapsed first())
# must be in {0, t_post}. Catches mixed-row malformed inputs where
# a unit has [valid, invalid].
valid_values = {0, t_post}
observed_raw = set(ft_raw.unique().tolist())
bad = sorted(observed_raw - valid_values, key=lambda x: str(x))
if bad:
raise ValueError(
f"first_treat_col={first_treat_col!r} contains value(s) "
f"{bad} outside the allowed set {{0, {t_post!r}}} for a "
f"two-period HAD panel. Staggered timing with multiple "
f"cohorts is Phase 2b."
)
# Within-unit consistency: every unit must have a single
# first_treat value across its rows. Uses dropna=False so a unit
# with [value, NaN] counts as 2 unique values (caught above by
# the NaN check anyway, but this is belt-and-suspenders).
ft_per_unit_nunique = data.groupby(unit_col)[first_treat_col].nunique(dropna=False)
if (ft_per_unit_nunique > 1).any():
n_bad = int((ft_per_unit_nunique > 1).sum())
raise ValueError(
f"first_treat_col={first_treat_col!r} is not constant "
f"within unit for {n_bad} unit(s). Each unit must have "
f"a single first_treat value across both observed periods."
)
return t_pre, t_post
def _validate_had_panel_event_study(
data: pd.DataFrame,
outcome_col: str,
dose_col: str,
time_col: str,
unit_col: str,
first_treat_col: Optional[str],
) -> Tuple[Any, List[Any], List[Any], pd.DataFrame, Optional[Dict[str, Any]]]:
"""Validate a HAD panel for multi-period event-study mode (Phase 2b).
Implements paper Appendix B.2 contract: a common treatment date ``F``
where ``D_{g,t} = 0`` for all units at ``t < F`` and some units have
``D_{g,t} > 0`` for ``t >= F``. Requires ``len(periods) > 2`` with at
least one pre-period (``t < F``, all D=0) and at least one post-period
(``t >= F``, some D > 0).
Staggered-timing handling: when ``first_treat_col`` is supplied and
indicates more than one nonzero cohort, the panel is auto-filtered to
the LAST cohort (``F_last = max(cohorts)``) per paper Appendix B.2
prescription "did_had may be used only for the last treatment cohort
in a staggered design". A ``UserWarning`` is emitted with drop-counts.
Parameters
----------
data, outcome_col, dose_col, time_col, unit_col, first_treat_col
As in :func:`_validate_had_panel`.
Returns
-------
F : period label
First-treatment period (the earliest period where any unit has
``D > 0`` in the filtered data).
t_pre_list : list
Pre-period labels (``t < F``, all D=0), sorted by natural ordering
on the column dtype.
t_post_list : list
Post-period labels (``t >= F``, some D > 0), sorted.
data_filtered : pd.DataFrame
Input with earlier cohorts (``first_treat`` in ``dropped_cohorts``)
dropped if staggered; never-treated units (``first_treat = 0``)
are RETAINED per paper Appendix B.2's "there must be an untreated
group" requirement. Identical to input when no staggered filter
applies.
filter_info : dict or None
Populated on staggered filter with keys ``F_last`` (kept cohort
label), ``n_kept`` (last-cohort units PLUS never-treated units),
``n_dropped`` (earlier-cohort units removed), ``dropped_cohorts``
(list of earlier cohort labels). ``None`` otherwise.
Raises
------
ValueError
On missing columns, NaN in key columns, malformed panel, dose-
invariant violations, or no-treatment detected.
"""
required = [outcome_col, dose_col, time_col, unit_col]
if first_treat_col is not None:
required.append(first_treat_col)
missing = [c for c in required if c not in data.columns]
if missing:
raise ValueError(f"Missing column(s) in data: {missing}. Required: {required}.")
periods_list = list(data[time_col].unique())
if len(periods_list) < 3:
raise ValueError(
f"HAD with aggregate='event_study' requires more than two "
f"time periods (got {len(periods_list)} in {time_col!r}). "
f"For two-period panels, pass aggregate='overall' (Phase 2a "
f"single-period WAS)."
)
# Ordered-time-type check. Paper Appendix B.2 event-time horizons
# require chronological ordering of periods (anchor at F-1, horizons
# e = t - F relative to F). Phase 2a two-period panels can use the
# dose invariant alone to distinguish pre from post without needing
# chronological order, so string labels ("pre", "post") work there.
# For multi-period event-study, multiple pre-periods all have D=0
# and multiple post-periods may both have D>0, so dose alone cannot
# recover chronology: we must trust the time column's natural order.
# Raw lexicographic sort on object/string labels silently misorders
# panels like "pre1"/"pre2"/"post1"/"post2" or month-name labels.
# Require an explicitly-ordered time representation.
time_dtype = data[time_col].dtype
if not (
pd.api.types.is_numeric_dtype(time_dtype)
or pd.api.types.is_datetime64_any_dtype(time_dtype)
or (isinstance(time_dtype, pd.CategoricalDtype) and bool(time_dtype.ordered))
):
raise ValueError(
f"HAD aggregate='event_study' requires an ordered time "
f"column. time_col={time_col!r} has dtype={time_dtype!r}, "
f"which has no defined chronological order; raw sort would "
f"fall back to lexicographic ordering and silently misindex "
f"event-time horizons (e.g., 'pre1'/'pre2'/'post1'/'post2' "
f"sorts lexicographically but not chronologically). "
f"Convert time_col to numeric (e.g., integer year), "
f"datetime, or ordered categorical "
f"(``pd.Categorical(..., ordered=True, categories=[...])``) "
f"before calling fit() with aggregate='event_study'."
)
# Construct the chronological sort key once, shared across every
# downstream ordering: cohort ranking, pre/post period sorting, and
# contiguity checks. Ordered categoricals use their declared
# category index (``list(categorical)`` strips the ordering and
# falls back to string comparison); numeric / datetime use natural
# Python order. Reused by ``_aggregate_multi_period_first_differences``
# via a parallel construction in that helper (both read the same
# ``time_dtype``).
if isinstance(time_dtype, pd.CategoricalDtype) and time_dtype.ordered:
_cat_order = {c: i for i, c in enumerate(time_dtype.categories)}
def _sort_key(x: Any) -> Tuple[bool, Any]:
return (x is None, _cat_order.get(x, len(_cat_order)))
else:
def _sort_key(x: Any) -> Tuple[bool, Any]:
return (x is None, x)
# NaN checks on key columns (before any filter).
for col in [outcome_col, dose_col, unit_col]:
if bool(data[col].isna().any()):
n_nan = int(data[col].isna().sum())
raise ValueError(
f"{n_nan} NaN value(s) found in column {col!r}. HAD "
f"does not silently drop NaN rows; drop or impute before "
f"calling fit()."
)
# Cohort detection and staggered-timing auto-filter.
filter_info: Optional[Dict[str, Any]] = None
data_filtered = data
if first_treat_col is not None:
ft_raw = data[first_treat_col]
if bool(ft_raw.isna().any()):
n_nan = int(ft_raw.isna().sum())
raise ValueError(
f"first_treat_col={first_treat_col!r} contains "
f"{n_nan} NaN value(s) at the row level. Use 0 for "
f"never-treated units and the treatment-start period "
f"for treated units. Drop or impute any NaN rows "
f"before calling fit()."
)
# Within-unit constancy check.
ft_per_unit_nunique = data.groupby(unit_col)[first_treat_col].nunique(dropna=False)
if (ft_per_unit_nunique > 1).any():
n_bad = int((ft_per_unit_nunique > 1).sum())
raise ValueError(
f"first_treat_col={first_treat_col!r} is not constant "
f"within unit for {n_bad} unit(s). Each unit must have "
f"a single first_treat value across all observed periods."
)
# Cross-validate first_treat_col against observed first-positive-
# dose period for every unit. A mislabeled cohort column would
# otherwise silently select the wrong cohort as F_last and return
# event-study estimates for the wrong units. Contract:
# - declared first_treat == 0: unit must have D == 0 at all t
# (never-treated)
# - declared first_treat == F_g > 0: unit's first period with
# D > 0 must equal F_g
df_for_check = data.sort_values([unit_col, time_col])
pos_rows = df_for_check.loc[df_for_check[dose_col] > 0]
actual_first_pos = pos_rows.groupby(unit_col)[time_col].first()
declared_ft = df_for_check.groupby(unit_col)[first_treat_col].first()
n_mismatch = 0
example_mismatch: Optional[Tuple[Any, Any, Any]] = None
for u, declared in declared_ft.items():
actual = actual_first_pos.get(u, None)
if declared == 0:
if actual is not None:
n_mismatch += 1
if example_mismatch is None:
example_mismatch = (u, declared, actual)
else:
if actual is None or actual != declared:
n_mismatch += 1
if example_mismatch is None:
example_mismatch = (u, declared, actual)
if n_mismatch > 0:
u, declared, actual = example_mismatch # type: ignore[misc]
raise ValueError(
f"first_treat_col={first_treat_col!r} disagrees with the "
f"observed dose path for {n_mismatch} unit(s). Example: "
f"unit={u!r} declares first_treat={declared!r} but the "
f"unit's first period with D>0 is {actual!r} "
f"(None means never-treated). A mislabeled cohort column "
f"would silently select the wrong cohort as F_last in the "
f"last-cohort auto-filter. Fix the first_treat_col values "
f"to equal each unit's first positive-dose period (or 0 "
f"for never-treated) before calling fit()."
)
# Identify cohorts (nonzero first_treat values). Sort using
# ``_sort_key`` (chronological order from ``time_dtype``), NOT
# raw Python sort: first_treat values are period labels and
# must rank chronologically so ``F_last = cohorts[-1]`` is the
# chronologically latest cohort. Under ordered-categorical time
# labels (e.g. month names), raw Python sort is lexicographic
# and would silently pick the wrong ``F_last``.
ft_unique = list(pd.unique(ft_raw))
cohorts = sorted(
[v for v in ft_unique if v != 0 and not (isinstance(v, float) and np.isnan(v))],
key=_sort_key,
)
if len(cohorts) == 0:
raise ValueError(
f"first_treat_col={first_treat_col!r} has no nonzero "
f"cohort values (all units appear never-treated). HAD "
f"requires at least one treated cohort with "
f"first_treat > 0 to identify a WAS effect."
)
if len(cohorts) > 1:
F_last = cohorts[-1]
dropped_cohorts = cohorts[:-1]
# Filter: keep last-cohort AND never-treated (first_treat == 0).
# Paper Appendix B.2: "in designs with variation in treatment
# timing, there must be an untreated group, at least till the
# period where the last cohort gets treated". Never-treated
# units (first_treat=0) satisfy the dose invariant for every
# period (D=0 throughout) and serve as the untreated-group
# comparison at every pre-period horizon. Keeping them matches
# the paper's "there must be an untreated group" language and
# preserves Design 1' identifiability (boundary at 0) when the
# last-cohort doses are uniformly positive. Only earlier-treated
# cohorts (first_treat in dropped_cohorts) are dropped.
keep_mask = (data[first_treat_col] == F_last) | (data[first_treat_col] == 0)
dropped_unit_ids = set(data.loc[~keep_mask, unit_col].unique())
kept_unit_ids = set(data.loc[keep_mask, unit_col].unique())
data_filtered = data.loc[keep_mask].copy()
n_dropped = len(dropped_unit_ids - kept_unit_ids)
n_kept = len(kept_unit_ids)
if n_kept == 0:
raise ValueError(
f"Staggered auto-filter to last cohort "
f"(F_last={F_last!r}) left 0 units. Verify "
f"first_treat_col={first_treat_col!r} contains the "
f"expected cohort labels."
)
filter_info = {
"F_last": F_last,
"n_kept": n_kept,
"n_dropped": n_dropped,
"dropped_cohorts": dropped_cohorts,
}
warnings.warn(
f"Staggered-timing panel detected: {len(cohorts)} distinct "
f"nonzero cohorts in first_treat_col={first_treat_col!r} "
f"({cohorts!r}). Auto-filtering to the last cohort "
f"(F_last={F_last!r}) plus never-treated units "
f"(first_treat=0): {n_kept} units kept, {n_dropped} "
f"earlier-cohort units dropped (from cohorts "
f"{dropped_cohorts!r}). HAD applies only to the last "
f"treatment cohort in staggered designs (paper Appendix "
f"B.2); never-treated units are retained as the untreated-"
f"group comparison per the paper's \"there must be an "
f'untreated group" requirement. For earlier-cohort '
f"effects, use ChaisemartinDHaultfoeuille "
f"(did_multiplegt_dyn).",
UserWarning,
stacklevel=3,
)
# After filter, re-read periods_list (cohort filter may have
# dropped some periods if earlier cohorts contributed uniquely).
periods_list = list(data_filtered[time_col].unique())
if len(periods_list) < 3:
raise ValueError(
f"After staggered auto-filter to last cohort "
f"(F_last={F_last!r}), only {len(periods_list)} "
f"distinct time periods remain in {time_col!r}. "
f"Event-study requires >2 periods; the filtered "
f"panel is too small. Pass aggregate='overall' on "
f"a two-period subset, or supply data with more "
f"pre- or post-periods for the last cohort."
)
# Balanced panel on the (possibly-filtered) data: every unit appears
# exactly once per period. ``observed=True`` tells categorical
# groupby to count only OBSERVED unit-period cells. Without it, a
# time_col with an ordered-categorical dtype carrying extra unused
# category levels (beyond the periods actually present in the data)
# would expand to zero-count cells and the balance check would
# falsely reject valid panels. The rest of the validator is keyed
# to ``periods_list`` (observed unique values) so this stays
# consistent.
counts = data_filtered.groupby([unit_col, time_col], observed=True).size()
if (counts != 1).any():
n_bad = int((counts != 1).sum())
raise ValueError(
f"Unbalanced panel: {n_bad} unit-period cells have != 1 "
f"observation. HAD requires a balanced panel (each unit "
f"observed exactly once at each period)."
)
unit_counts = data_filtered.groupby(unit_col)[time_col].nunique()
incomplete = unit_counts[unit_counts != len(periods_list)]
if len(incomplete) > 0:
raise ValueError(
f"Unbalanced panel: {len(incomplete)} unit(s) do not appear "
f"in all {len(periods_list)} periods. HAD requires a balanced "
f"panel (each unit observed at every period)."
)
# Dose-invariant period classification on filtered data.
per_period_nonzero: Dict[Any, int] = {}
for p in periods_list:
p_doses = np.asarray(
data_filtered.loc[data_filtered[time_col] == p, dose_col], dtype=np.float64
)
per_period_nonzero[p] = int((p_doses != 0).sum())
t_pre_list_unsorted = [p for p, nz in per_period_nonzero.items() if nz == 0]
t_post_list_unsorted = [p for p, nz in per_period_nonzero.items() if nz > 0]
if len(t_pre_list_unsorted) == 0:
stats_str = ", ".join(f"{p!r}: {nz} nonzero" for p, nz in per_period_nonzero.items())
raise ValueError(
f"HAD requires D_{{g,t}} = 0 for all units in at least one "
f"pre-period. No period in column {time_col!r} has all-zero "
f"dose ({stats_str}). The panel has no identifiable baseline."
)
if len(t_post_list_unsorted) == 0:
raise ValueError(
f"HAD requires at least one period with nonzero dose for "
f"some unit. All periods in column {time_col!r} have all-"
f"zero dose; there is no treatment to estimate."
)
# Sort using the same ``_sort_key`` already constructed for cohorts
# (ordered-categorical uses declared category order; numeric /
# datetime use natural Python order).
t_pre_list = sorted(t_pre_list_unsorted, key=_sort_key)
t_post_list = sorted(t_post_list_unsorted, key=_sort_key)
# Contiguity check: all pre < all post in the natural ordering.
# The HAD dose invariant requires a single transition from all-zero
# to any-nonzero; interleaved pre/post periods indicate a malformed
# panel (e.g., dose going back to zero after treatment, or mixing
# never-treated units with out-of-order labels). Uses ``_sort_key``
# so ordered categoricals respect their declared category order.
if t_pre_list and t_post_list:
max_pre = t_pre_list[-1]
min_post = t_post_list[0]
contiguous = _sort_key(max_pre) < _sort_key(min_post)
if not contiguous:
raise ValueError(
f"HAD dose invariant violated: pre-periods (all D=0) "
f"and post-periods (some D>0) are not contiguous. "
f"Pre-periods: {t_pre_list!r}; post-periods: "
f"{t_post_list!r}. The dose sequence must transition "
f"from all-zero to nonzero exactly once. For panels "
f"where dose varies non-monotonically (e.g., reversed "
f"treatment, switching), use "
f"ChaisemartinDHaultfoeuille (did_multiplegt_dyn)."
)
F = t_post_list[0] # earliest post-period
# Post-period nonnegative-dose check on the filtered data.
post_mask = data_filtered[time_col].isin(t_post_list)
post_doses = np.asarray(data_filtered.loc[post_mask, dose_col], dtype=np.float64)
neg_post = post_doses < 0
if neg_post.any():
n_neg = int(neg_post.sum())
min_neg = float(post_doses[neg_post].min())
raise ValueError(
f"HAD requires D_{{g,t}} >= 0 for all units in post-periods "
f"(paper Section 2 dose definition). {n_neg} unit-period "
f"cell(s) have negative dose at t >= F={F!r} "
f"(min={min_neg!r}). Drop these units or verify the dose "
f"column."
)
# Staggered-without-``first_treat_col`` detection. When cohort metadata
# is not supplied, the dose-invariant period classification still
# declares t=F=min-post-period based on "any unit has nonzero dose".
# That silently accepts staggered panels where units have DIFFERENT
# first-positive-dose periods: the later-treated cohorts enter
# ``d_arr`` as zero-dose "controls" at the inferred F, violating
# paper Appendix B.2's last-cohort-only contract. Compute per-unit
# first-positive-dose period directly from the dose path and raise
# if multiple cohorts are present, directing users to pass
# ``first_treat_col`` (which activates the last-cohort auto-filter)
# or to use ChaisemartinDHaultfoeuille for full staggered support.
if first_treat_col is None:
df_sorted = data_filtered.sort_values([unit_col, time_col])
# For each unit, the first period at which dose > 0.
pos_mask_global = df_sorted[dose_col] > 0
first_pos_per_unit = df_sorted.loc[pos_mask_global].groupby(unit_col)[time_col].first()
cohort_labels = list(first_pos_per_unit.unique())
if len(cohort_labels) > 1:
# Sort chronologically via the validated time-column order.
distinct_cohorts = sorted(cohort_labels, key=_sort_key)
raise ValueError(
f"Staggered-timing panel detected (first_treat_col is "
f"None): {len(distinct_cohorts)} distinct first-positive-"
f"dose periods {distinct_cohorts!r} across units. HAD's "
f"last-cohort auto-filter (paper Appendix B.2) only runs "
f"when first_treat_col is supplied so the estimator can "
f"identify cohorts. Pass first_treat_col= to "
f"enable the auto-filter to the last cohort, or use "
f"ChaisemartinDHaultfoeuille (did_multiplegt_dyn) for "
f"full staggered support."
)
# Constant post-period dose check. Paper Appendix B.2 assumes
# "once treated, stay treated with the same dose"; the event-study
# aggregation uses ``D_{g, F}`` as the single regressor for every
# event-time horizon. Panels where a unit's dose varies across
# post-periods (e.g., phased adoption, dose changes after F) would
# silently misattribute later-horizon effects to the period-F dose.
# Reject front-door with a redirect to ChaisemartinDHaultfoeuille
# for genuinely time-varying post-treatment doses.
if len(t_post_list) > 1:
post_data = data_filtered.loc[post_mask]
dose_spread_per_unit = post_data.groupby(unit_col)[dose_col].agg(
lambda x: float(x.max() - x.min())
)
abs_max_dose = float(np.max(np.abs(post_doses))) if post_doses.size else 0.0
tol = 1e-12 * max(1.0, abs_max_dose)
bad_mask = dose_spread_per_unit > tol
if bool(bad_mask.any()):
n_bad = int(bad_mask.sum())
max_spread = float(dose_spread_per_unit.max())
raise ValueError(
f"HAD event-study requires constant dose within unit for "
f"all post-treatment periods t >= F={F!r}. {n_bad} unit(s) "
f"have time-varying doses across post-periods "
f"{t_post_list!r} (max within-unit spread={max_spread!r}, "
f"tolerance={tol!r}). The aggregation uses D_{{g, F}} as "
f"the single regressor for every event-time horizon "
f"(paper Appendix B.2 constant-dose convention), so "
f"silently accepting time-varying post-treatment doses "
f"would misattribute later-horizon effects. For genuinely "
f"time-varying post-treatment doses use "
f"ChaisemartinDHaultfoeuille (did_multiplegt_dyn)."
)
return F, t_pre_list, t_post_list, data_filtered, filter_info
def _aggregate_first_difference(
data: pd.DataFrame,
outcome_col: str,
dose_col: str,
time_col: str,
unit_col: str,
t_pre: Any,
t_post: Any,
cluster_col: Optional[str],
) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray], np.ndarray]:
"""Reduce a balanced two-period panel to unit-level first differences.
Returns
-------
d_arr : np.ndarray, shape (G,)
Post-period dose ``D_{g,2}`` per unit, ordered by sorted unit id.
dy_arr : np.ndarray, shape (G,)
First-difference outcome ``Y_{g,2} - Y_{g,1}`` per unit.
cluster_arr : np.ndarray or None, shape (G,)
Cluster IDs per unit (must be unit-constant); ``None`` when
``cluster_col`` is ``None``.
unit_ids : np.ndarray, shape (G,)
Sorted unit identifiers (useful for debugging / test inspection).
"""
df = data.sort_values([unit_col, time_col]).reset_index(drop=True)
pre = df[df[time_col] == t_pre].set_index(unit_col).sort_index()
post = df[df[time_col] == t_post].set_index(unit_col).sort_index()
if bool(pre.index.equals(post.index)) is False:
# Should not reach here after _validate_had_panel balanced-panel
# check, but belt-and-suspenders:
raise ValueError(
"Internal error: pre and post period unit indices do not "
"match after balanced-panel validation. Please report this."
)
unit_ids = np.asarray(pre.index)
d_arr = post[dose_col].to_numpy(dtype=np.float64)
dy_arr = post[outcome_col].to_numpy(dtype=np.float64) - pre[outcome_col].to_numpy(
dtype=np.float64
)
cluster_arr: Optional[np.ndarray] = None
if cluster_col is not None:
if cluster_col not in data.columns:
raise ValueError(f"cluster column {cluster_col!r} not found in data.")
# Row-level NaN check: `groupby().first()` skips NaNs silently, so a
# unit with rows [valid, NaN] would pass a collapsed check. Validate
# raw per-row values first so any NaN in the cluster column is
# rejected up-front, not masked by the per-unit collapse.
cluster_raw = data[cluster_col]
if bool(cluster_raw.isna().any()):
n_nan = int(cluster_raw.isna().sum())
raise ValueError(
f"cluster column {cluster_col!r} contains {n_nan} NaN "
f"value(s) at the row level. Silent row dropping is "
f"disabled; drop or impute cluster ids before calling fit()."
)
# Cluster must be unit-constant. `nunique(dropna=False)` counts NaN
# as a distinct value so rows like [value, NaN] register as 2
# uniques (also caught by the row-level NaN check above).
cluster_per_unit = df.groupby(unit_col)[cluster_col].nunique(dropna=False)
if (cluster_per_unit > 1).any():
n_bad = int((cluster_per_unit > 1).sum())
raise ValueError(
f"cluster column {cluster_col!r} is not constant within "
f"unit for {n_bad} unit(s). Cluster must be unit-level "
f"(e.g., cluster_col=unit_col, or a coarser grouping "
f"where each unit belongs to a single cluster)."
)
cluster_arr = df.groupby(unit_col)[cluster_col].first().sort_index().to_numpy()
return d_arr, dy_arr, cluster_arr, unit_ids
def _aggregate_unit_weights(
data: pd.DataFrame,
weights_arr: np.ndarray,
unit_col: str,
) -> np.ndarray:
"""Aggregate per-row weights to per-unit, enforcing constant-within-unit.
HAD's continuous path operates at the per-unit level (G rows, one per
unit) so survey weights which typically arrive per-row in a long
panel must collapse to a single value per unit. The paper's weighted
extension assumes sampling weights assigned at the sampling-unit level
(standard BRFSS/CPS/NHANES convention), so any within-unit weight
variance is interpreted as user error and rejected front-door rather
than silently mean-rolled (per ``feedback_no_silent_failures``).
Parameters
----------
data : pd.DataFrame
Long panel; the same frame passed to ``_aggregate_first_difference``.
weights_arr : np.ndarray, shape (n_rows,)
Row-aligned weights.
unit_col : str
Unit identifier column.
Returns
-------
w_unit : np.ndarray, shape (G,)
Per-unit weights sorted by unit id to align with the d/dy arrays
returned by ``_aggregate_first_difference``.
Raises
------
ValueError
Shape mismatch, non-finite weights, negative weights, zero-sum
weights, or weights that vary within a unit.
"""
n_rows = int(data.shape[0])
w = np.asarray(weights_arr, dtype=np.float64).ravel()
if w.shape[0] != n_rows:
raise ValueError(
f"weights length ({w.shape[0]}) does not match number of " f"rows in data ({n_rows})."
)
if not np.all(np.isfinite(w)):
raise ValueError("weights contains non-finite values (NaN or Inf).")
if np.any(w < 0):
raise ValueError("weights must be non-negative.")
if np.sum(w) Any: # ResolvedSurveyDesign at unit level
"""Collapse a row-level ResolvedSurveyDesign to a unit-level analogue.
HAD's continuous path operates at G per-unit rows, so the survey
design (weights / strata / PSU / FPC) must collapse the same way.
Each design column is required to be constant-within-unit (same
invariant as ``_aggregate_unit_weights``); a within-unit inconsistency
raises ``ValueError`` per ``feedback_no_silent_failures``.
Parameters
----------
data : pd.DataFrame
Long panel.
resolved : ResolvedSurveyDesign
Resolved design with ``(n_rows,)`` arrays.
unit_col : str
Returns
-------
ResolvedSurveyDesign
New resolved design with ``(G,)`` arrays aligned to sorted unit ids.
Raises
------
ValueError
Strata / PSU / FPC vary within unit, or replicate-weight designs
are passed (replicate-weight HAD is Phase 4.5 C scope, not this
commit).
NotImplementedError
``resolved.replicate_weights is not None`` replicate-weight HAD
is deferred.
"""
from diff_diff.survey import ResolvedSurveyDesign
if resolved.replicate_weights is not None:
raise NotImplementedError(
"Replicate-weight SurveyDesign on HAD is deferred to Phase 4.5 C. "
"Pass a SurveyDesign with weights/strata/psu/fpc (Taylor-series "
"linearization path) for this PR."
)
n_rows = int(data.shape[0])
if resolved.weights.shape[0] != n_rows:
raise ValueError(
f"ResolvedSurveyDesign.weights length ({resolved.weights.shape[0]}) "
f"does not match data rows ({n_rows}). The SurveyDesign must "
f"have been resolved against the same DataFrame passed to fit()."
)
def _collapse(arr: Optional[np.ndarray], name: str) -> Optional[np.ndarray]:
if arr is None:
return None
if arr.shape[0] != n_rows:
raise ValueError(
f"ResolvedSurveyDesign.{name} length does not match " f"data rows ({n_rows})."
)
df = data.reset_index(drop=True).copy()
# Use a stable per-row column so we can take first() per unit and
# verify constancy. For numeric designs, use np.isclose; for
# object/string (rare), require equality.
df["_tmp__"] = arr
unit_min = df.groupby(unit_col)["_tmp__"].min()
unit_max = df.groupby(unit_col)["_tmp__"].max()
# nunique=1 is the robust check across dtypes.
nunique = df.groupby(unit_col)["_tmp__"].nunique(dropna=False)
if (nunique > 1).any():
n_bad = int((nunique > 1).sum())
raise ValueError(
f"ResolvedSurveyDesign.{name} varies within {n_bad} unit(s). "
f"Survey design columns must be constant within each unit "
f"(sampling-unit-level assignment convention). If your "
f"panel uses unit-varying design columns, aggregate them "
f"to unit-level before calling fit()."
)
# The unit-level min and max coincide by the nunique check; take min.
_ = unit_max # referenced for symmetry of intent; min suffices
return unit_min.sort_index().to_numpy(dtype=arr.dtype)
w_unit = _collapse(resolved.weights, "weights")
assert w_unit is not None # resolved.weights is non-Optional
strata_unit = _collapse(resolved.strata, "strata")
psu_unit = _collapse(resolved.psu, "psu")
fpc_unit = _collapse(resolved.fpc, "fpc")
# Recompute n_strata and n_psu at the unit level.
n_strata_unit = int(np.unique(strata_unit).shape[0]) if strata_unit is not None else 1
n_psu_unit = int(np.unique(psu_unit).shape[0]) if psu_unit is not None else int(w_unit.shape[0])
return ResolvedSurveyDesign(
weights=w_unit,
weight_type=resolved.weight_type,
strata=strata_unit,
psu=psu_unit,
fpc=fpc_unit,
n_strata=n_strata_unit,
n_psu=n_psu_unit,
lonely_psu=resolved.lonely_psu,
replicate_weights=None,
replicate_method=None,
fay_rho=0.0,
n_replicates=0,
replicate_strata=None,
combined_weights=resolved.combined_weights,
replicate_scale=None,
replicate_rscales=None,
mse=resolved.mse,
)
def _aggregate_multi_period_first_differences(
data: pd.DataFrame,
outcome_col: str,
dose_col: str,
time_col: str,
unit_col: str,
F: Any,
t_pre_list: List[Any],
t_post_list: List[Any],
cluster_col: Optional[str],
) -> Tuple[np.ndarray, Dict[int, np.ndarray], Optional[np.ndarray], np.ndarray, Any]:
"""Reduce a multi-period HAD panel to per-horizon first differences.
For each period ``t`` other than the anchor ``t_anchor = F - 1`` (the
last pre-period), computes the unit-level first difference
``Y_{g,t} = Y_{g,t} - Y_{g, t_anchor}`` and stores it under the event
time ``e = rank(t) - rank(F)`` where ``rank`` is the natural ordering
on the period column (so ``e = 0`` at ``t = F``, ``e = 1`` at the next
post-period, etc., and ``e 0`` since otherwise the "at d_lower" subset is always
the "zero-dose" subset and the continuous-at-zero path is the
paper-endorsed handling.)
2. Otherwise, if ``d.min() < 0.01 * median(|d|)``, resolve to
``continuous_at_zero`` (small-share-of-treated samples that
effectively satisfy Design 1').
3. Else, if the modal fraction at ``d.min()`` exceeds 2%, resolve to
``mass_point``.
4. Else, resolve to ``continuous_near_d_lower``.
The rule is strict first-match: when both the `` 0 and d_min < _CONTINUOUS_AT_ZERO_THRESHOLD * median_abs:
return "continuous_at_zero"
# Modal-fraction rule for mass-point detection.
eps = 1e-12 * max(1.0, abs(d_min))
at_d_min = np.abs(d_arr - d_min) _MASS_POINT_THRESHOLD:
return "mass_point"
return "continuous_near_d_lower"
# =============================================================================
# Sup-t multiplier bootstrap (Phase 4.5 B event-study simultaneous CI)
# =============================================================================
def _sup_t_multiplier_bootstrap(
influence_matrix: np.ndarray,
att_per_horizon: np.ndarray,
se_per_horizon: np.ndarray,
resolved_survey: Any, # Optional[ResolvedSurveyDesign]
*,
n_bootstrap: int,
alpha: float,
seed: Optional[int],
bootstrap_weights: str = "rademacher",
) -> Tuple[float, Optional[np.ndarray], Optional[np.ndarray], int]:
"""Compute sup-t simultaneous CI via PSU-level multiplier bootstrap.
Reuses :func:`diff_diff.bootstrap_utils.generate_survey_multiplier_weights_batch`
(survey path) / :func:`generate_bootstrap_weights_batch` (weights=
shortcut) to draw ``n_bootstrap`` replicates of multiplier weights
shaped ``(n_bootstrap, n_units)``; the helpers handle stratum
centering, lonely-PSU, and FPC scaling so this function only
composes them with the per-unit influence function.
Construction (mirrors `staggered_bootstrap.py:354-373` perturbation
idiom and `:497-533` sup-t quantile idiom; NO ``(1/n)`` prefactor
``psi`` is already on the -scale per the Phase 4.5 B IF scale
convention):
1. Draw multiplier weights ``xi`` shape ``(n_bootstrap, n_units)``.
2. Perturbations: ``delta[b, e] = sum_g xi[b, g] * psi[g, e]``.
3. t-statistics: ``t[b, e] = delta[b, e] / se[e]``.
4. Sup-t: ``sup_t[b] = max_e |t[b, e]|``.
5. Critical value: ``q = quantile(sup_t[isfinite], 1 - alpha)``.
Under ``H=1`` the sup reduces to the marginal, so
``q -> Phi^{-1}(1 - alpha/2) 1.96`` at ``alpha=0.05`` (locked by
the ``TestSupTReducesToNormalAtH1`` reduction-invariant test).
Parameters
----------
influence_matrix : np.ndarray, shape (n_units, n_horizons)
Per-unit per-horizon influence function on the -scale. NaN
columns are treated as degenerate-horizon placeholders and drop
out of the sup via the finite mask.
att_per_horizon : np.ndarray, shape (n_horizons,)
Per-horizon point estimates (used to assemble the simultaneous
band: ``att q se``).
se_per_horizon : np.ndarray, shape (n_horizons,)
Per-horizon analytical SE (Binder-TSL on survey path, HC1
sandwich on weights= shortcut).
resolved_survey : Optional[ResolvedSurveyDesign]
``None`` unit-level Rademacher draw via
``generate_bootstrap_weights_batch``. Otherwise
PSU-level draw via ``generate_survey_multiplier_weights_batch``
(stratum-centered, FPC-scaled, lonely-PSU-aware).
n_bootstrap : int
Number of multiplier replicates.
alpha : float
CI level (``0.05`` for 95% simultaneous band).
seed : int or None
RNG seed for reproducibility.
bootstrap_weights : str
Passed through to the helper: ``"rademacher"``, ``"mammen"``, or
``"webb"``. Default ``"rademacher"`` (binary 1 multipliers).
Returns
-------
(cband_crit_value, cband_low, cband_high, n_valid) : tuple
``cband_crit_value`` is the sup-t quantile (float); ``cband_low``
/ ``cband_high`` are simultaneous-band endpoints shape
``(n_horizons,)``; ``n_valid`` is the count of finite sup-t
draws ( d_lower}``
gives a 2SLS estimator that collapses to the Wald-IV sample-average
ratio
``beta_hat = (Ybar_{Z=1} - Ybar_{Z=0}) / (Dbar_{Z=1} - Dbar_{Z=0})``.
The STANDARD ERROR is computed via the 2SLS sandwich
``V = [Z'WX]^{-1} * Omega * [Z'WX]^{-T}`` where ``Omega`` is built
from the STRUCTURAL residuals
``u = dy - alpha_hat - beta_hat * d`` (NOT the reduced-form
residuals). This is the canonical 2SLS inference path and matches
what ``estimatr::iv_robust`` / Stata ``ivregress`` would produce.
Supported ``vcov_type``:
- ``"classical"``: constant variance ``sigma_hat^2 = sum(w u) /
(sum(w) - k)`` (weighted) or ``sum(u^2) / (n-k)`` (unweighted);
sandwich form with ``w`` in the meat when weighted.
- ``"hc1"``: heteroskedasticity-robust with small-sample DOF scaling
``n / (n - k)``; meat = ``Z' diag(w u) Z`` (pweight convention,
Wooldridge 2010 Eq. 12.37; matches ``estimatr::iv_robust(...,
weights=..., se_type="HC1")`` bit-exactly). With ``cluster``
supplied, switches to CR1 (Liang-Zeger) with cluster score
``Z'_c (w u)_c``.
``"hc2"`` and ``"hc2_bm"`` raise ``NotImplementedError`` (2SLS-specific
leverage derivation pending).
Parameters
----------
d : np.ndarray, shape (n,)
Post-period doses ``D_{g,2}``.
dy : np.ndarray, shape (n,)
First-difference outcome ``Delta Y_g``.
d_lower : float
Support infimum / evaluation point.
cluster : np.ndarray or None, shape (n,)
Cluster ids per unit (``None`` for no clustering).
vcov_type : str
One of ``"classical"``, ``"hc1"``.
weights : np.ndarray or None, shape (n,)
Per-unit sampling weights (pweight convention). ``None`` for the
unweighted path that branch is numerically bit-exact with
pre-Phase 4.5 B output. Zero-weight units contribute zero to
every sum-weighted expression and zero to the returned IF.
return_influence : bool
When True, returns the per-unit influence function (IF) on the
-scale, shape ``(n,)`` with zeros at zero-weight rows. The IF
is scaled so that under a trivial ``ResolvedSurveyDesign``
(single stratum, each unit its own PSU, no FPC),
``compute_survey_if_variance(IF, trivial)`` ``V_HC1[1, 1]``
at ``atol=1e-10`` (PR #359 convention; see the "IF scale
convention" section of the Phase 4.5 B plan for derivation).
Returns
-------
tuple[float, float, np.ndarray or None]
``(beta_hat, se_beta, psi)``. ``psi`` is the per-unit IF when
``return_influence=True``, else ``None``. NaN for SE when the
dose-gap vanishes (``Dbar_{Z=1} == Dbar_{Z=0}``) or the
sandwich is singular; in those cases ``psi`` is returned as a
length-``n`` zero array when ``return_influence=True``.
"""
d = np.asarray(d, dtype=np.float64)
dy = np.asarray(dy, dtype=np.float64)
n = d.shape[0]
# Weight validation / normalization. The unweighted branch preserves
# numerical bit-parity vs pre-Phase 4.5 B by skipping np.average and
# using plain sums / means exactly as before. `w_arr`, `w_sum`, and
# `pos_mask` are initialized as sentinels so static typing flows
# cleanly; only the weighted branch populates them with real values.
weighted = weights is not None
w_arr: np.ndarray = np.ones(n, dtype=np.float64)
w_sum: float = float(n)
pos_mask: np.ndarray = np.ones(n, dtype=bool)
if weighted:
w_arr = np.asarray(weights, dtype=np.float64).ravel()
if w_arr.shape[0] != n:
raise ValueError(
f"weights length ({w_arr.shape[0]}) does not match d / dy length ({n})."
)
if not np.all(np.isfinite(w_arr)):
raise ValueError("weights contains non-finite values (NaN or Inf).")
if np.any(w_arr < 0):
raise ValueError("weights must be non-negative.")
w_sum = float(w_arr.sum())
if w_sum 0
Z = (d > d_lower).astype(np.float64)
# Degeneracy checks on the positive-weight subset (zero-weight units
# do not drive design resolution; same subpopulation convention as
# PR #359). Under unweighted fits the "positive" subset is the full
# sample so behavior is unchanged.
if weighted:
n_above = int(((Z == 1) & pos_mask).sum())
n_at_or_below = int(((Z == 0) & pos_mask).sum())
else:
n_above = int(Z.sum())
n_at_or_below = n - n_above
_null_psi = np.zeros(n, dtype=np.float64) if return_influence else None
if n_above == 0 or n_at_or_below == 0:
return float("nan"), float("nan"), _null_psi
# Point estimate: weighted Wald-IV ratio (reduces to unweighted at w=1).
if weighted:
Z1_idx = (Z == 1) & pos_mask
Z0_idx = (Z == 0) & pos_mask
w_Z1 = float(w_arr[Z1_idx].sum())
w_Z0 = float(w_arr[Z0_idx].sum())
if w_Z1 >> import numpy as np
>>> import pandas as pd
>>> from diff_diff import HeterogeneousAdoptionDiD
>>> rng = np.random.default_rng(42) # doctest: +SKIP
>>> G = 500 # doctest: +SKIP
>>> dose_post = rng.uniform(0.0, 1.0, G) # doctest: +SKIP
>>> dose_post[0] = 0.0 # at least one zero-dose unit for Design 1'
>>> delta_y = 0.3 * dose_post + 0.1 * rng.standard_normal(G) # doctest: +SKIP
>>> data = pd.DataFrame({ # doctest: +SKIP
... "unit": np.repeat(np.arange(G), 2),
... "period": np.tile([1, 2], G),
... "dose": np.column_stack([np.zeros(G), dose_post]).ravel(),
... "outcome": np.column_stack([np.zeros(G), delta_y]).ravel(),
... })
>>> est = HeterogeneousAdoptionDiD(design="auto") # doctest: +SKIP
>>> result = est.fit( # doctest: +SKIP
... data, outcome_col="outcome", dose_col="dose",
... time_col="period", unit_col="unit",
... )
>>> result.design # doctest: +SKIP
'continuous_at_zero'
"""
def __init__(
self,
design: str = "auto",
d_lower: Optional[float] = None,
kernel: str = "epanechnikov",
alpha: float = 0.05,
vcov_type: Optional[str] = None,
robust: bool = False,
cluster: Optional[str] = None,
n_bootstrap: int = 999,
seed: Optional[int] = None,
) -> None:
self.design = design
self.d_lower = d_lower
self.kernel = kernel
self.alpha = alpha
self.vcov_type = vcov_type
self.robust = robust
self.cluster = cluster
# Phase 4.5 B: event-study survey sup-t simultaneous-CI support.
# ``n_bootstrap`` = number of multiplier-bootstrap replicates for
# the sup-t band on the event-study + weighted path. ``seed`` =
# reproducibility seed for the multiplier draws. Both are
# consulted only when ``aggregate="event_study"`` AND a
# ``survey=`` / ``weights=`` is passed to ``fit()`` with
# ``cband=True`` (default). Unweighted event-study skips the
# bootstrap entirely pre-Phase 4.5 B numerical output preserved.
self.n_bootstrap = n_bootstrap
self.seed = seed
self._validate_constructor_args()
def _validate_constructor_args(self) -> None:
if self.design not in _VALID_DESIGNS:
raise ValueError(
f"Invalid design={self.design!r}. Must be one of " f"{_VALID_DESIGNS}."
)
if not (0.0 < float(self.alpha) < 1.0):
raise ValueError(f"alpha must be in (0, 1); got {self.alpha!r}.")
# d_lower must be None or a finite scalar. NaN / +/-inf would
# bypass every downstream comparison-based guard (`>`, ` Dict[str, Any]:
"""Return the raw constructor parameters (sklearn-compatible).
Matches the :meth:`sklearn.base.BaseEstimator.get_params`
signature. Preserves the user's original inputs - in particular,
``design`` returns ``"auto"`` when the user set it to ``"auto"``
(even after fit), so ``sklearn.base.clone(est)`` round-trips
exactly.
Parameters
----------
deep : bool, default=True
Accepted for sklearn-contract compatibility. This estimator
has no nested sub-estimator parameters, so ``deep=False``
and ``deep=True`` return the same dict.
"""
del deep # accepted for compat; this estimator has no nested params
return {
"design": self.design,
"d_lower": self.d_lower,
"kernel": self.kernel,
"alpha": self.alpha,
"vcov_type": self.vcov_type,
"robust": self.robust,
"cluster": self.cluster,
"n_bootstrap": self.n_bootstrap,
"seed": self.seed,
}
def set_params(self, **params: Any) -> "HeterogeneousAdoptionDiD":
"""Set estimator parameters and return self (sklearn-compatible).
Only keys returned by :meth:`get_params` are accepted. Passing
any other attribute name (including method names like ``fit``)
raises ``ValueError`` so the estimator cannot be silently
corrupted by a mistyped or attacker-supplied key.
Mutation is ATOMIC: validation runs on a proposed merged
parameter dict before any attribute is overwritten. A failing
call (invalid key, or an otherwise valid key whose value
violates the constructor constraints) leaves ``self`` unchanged
and safe to reuse.
"""
valid_keys = set(self.get_params().keys())
invalid = [k for k in params if k not in valid_keys]
if invalid:
raise ValueError(
f"Invalid parameter: {invalid[0]!r}. Valid parameters: " f"{sorted(valid_keys)}."
)
# Dry-run validation by constructing a fresh instance with the
# merged state. If the constructor raises, self is not mutated.
merged = self.get_params()
merged.update(params)
type(self)(**merged) # raises ValueError on invalid combination
# All checks passed; apply atomically.
for key, value in params.items():
setattr(self, key, value)
return self
# ------------------------------------------------------------------
# Main fit entry point
# ------------------------------------------------------------------
def fit(
self,
data: pd.DataFrame,
outcome_col: str,
dose_col: str,
time_col: str,
unit_col: str,
first_treat_col: Optional[str] = None,
aggregate: str = "overall",
# PR #376 R4 P1: preserve pre-PR positional-or-keyword status of
# `survey`, `weights`, `cband` for back-compat with positional
# callers. `survey_design=` is the only new addition and is
# keyword-only. PR #389 (Phase 4 R-parity): `trends_lin=` is
# likewise keyword-only (additive; no positional callers can
# exist for it pre-PR).
survey: Any = None,
weights: Optional[np.ndarray] = None,
cband: bool = True,
*,
survey_design: Any = None,
trends_lin: bool = False,
) -> HeterogeneousAdoptionDiDResults:
"""Fit the HAD estimator.
``aggregate="overall"`` (default) fits on a two-period panel and
returns a :class:`HeterogeneousAdoptionDiDResults` with the
single-period WAS estimate. ``aggregate="event_study"`` fits on
a multi-period panel (``T > 2``) and returns a
:class:`HeterogeneousAdoptionDiDEventStudyResults` with per-
event-time WAS estimates using a uniform ``F-1`` anchor (paper
Appendix B.2).
Both the overall and event-study paths are **panel-only**: the paper
(Section 2) defines HAD on panel or repeated-cross-section data,
but this implementation requires a balanced panel with a unit
identifier so that unit-level first differences
``Y_{g,t} = Y_{g,t} - Y_{g,t_anchor}`` can be formed.
Repeated-cross-section inputs (disjoint unit IDs between
periods) are rejected by the balanced-panel validator.
Repeated-cross-section support is queued for a follow-up PR
(tracked in ``TODO.md``); it requires a separate identification
path based on pre/post cell means rather than unit-level
differences.
Parameters
----------
data : pd.DataFrame
outcome_col, dose_col, time_col, unit_col : str
Column names.
first_treat_col : str or None
Optional first-treatment column (the period at which each
unit first receives treatment; ``0`` for never-treated).
Required on the event-study path when the panel has more
than two distinct first-treat values (staggered timing):
the estimator auto-filters to the last-treatment cohort
with a ``UserWarning`` per paper Appendix B.2 prescription.
For common-adoption panels the column is optional; when
omitted, the event-study path infers the first-treatment
period ``F`` from the dose invariant.
aggregate : {"overall", "event_study"}
``"overall"`` (default): returns a single-period
:class:`HeterogeneousAdoptionDiDResults` (Phase 2a). Requires
exactly two time periods.
``"event_study"`` (Phase 2b): returns a
:class:`HeterogeneousAdoptionDiDEventStudyResults` with per-
event-time WAS estimates on the multi-period panel (paper
Appendix B.2). Requires more than two time periods. Pointwise
CIs per horizon; joint cross-horizon covariance is deferred
to a follow-up PR. Staggered-timing panels are auto-filtered
to the last-treatment cohort with a ``UserWarning``.
survey_design : SurveyDesign or None, keyword-only
Survey design (sampling weights + optional strata / PSU / FPC)
for design-based inference. Supported on ALL design aggregate
combinations after Phase 4.5 B: continuous paths
(``continuous_at_zero``, ``continuous_near_d_lower``) on both
``aggregate="overall"`` and ``aggregate="event_study"``, AND
the ``mass_point`` design on both aggregates. Continuous paths
compose the SE via :func:`compute_survey_if_variance` (Binder
1983 TSL); weights propagate pointwise into the lprobust
kernel. Mass-point composes the per-unit 2SLS IF on the
HC1-scale and Binder-TSL-aggregates that requires
``vcov_type='hc1'`` (the classical default raises
``NotImplementedError`` on the survey path). Event-study fits
with ``cband=True`` add a multiplier-bootstrap simultaneous
confidence band. Only ``weight_type="pweight"`` is supported
(``aweight`` / ``fweight`` raise ``NotImplementedError``).
Survey design columns (strata / PSU / FPC) must be constant
within unit (sampling-unit-level assignment); within-unit
variance raises ``ValueError``. Replicate-weight designs raise
``NotImplementedError``. Mutually exclusive with the deprecated
``survey=`` and ``weights=`` aliases. See
``docs/methodology/REGISTRY.md`` HeterogeneousAdoptionDiD
"Note (HAD survey-design API consolidation)" for the full
dispatch matrix.
survey : SurveyDesign or None
DEPRECATED alias of ``survey_design=``. Remains positional-or-
keyword for one minor cycle to preserve pre-PR call shapes;
will be removed in the next minor release. Prefer
``survey_design=``.
weights : np.ndarray or None
DEPRECATED alias for the per-row pweight shortcut. Remains
positional-or-keyword for one minor cycle. Prefer adding the
weights as a column on ``data`` and passing
``survey_design=SurveyDesign(weights='col_name')`` instead.
Will be removed in the next minor release. Currently
preserved as the analytical-HC1-sandwich shortcut (continuous:
CCT-2014 weighted-robust; mass-point: pweight 2SLS sandwich)
with the per-row per-unit aggregation invariant intact.
Mutually exclusive with ``survey_design=`` and ``survey=``.
cband : bool, default True
Phase 4.5 B: controls the multiplier-bootstrap simultaneous
confidence band on the weighted event-study path. When
``True`` (default) and ``aggregate="event_study"`` AND any of
``survey_design=`` / ``survey=`` / ``weights=`` is supplied,
the fit populates ``cband_low`` / ``cband_high`` /
``cband_crit_value`` / ``cband_method`` / ``cband_n_bootstrap``
on the result. When ``False`` those fields stay ``None``. No
effect on ``aggregate="overall"`` or on unweighted event-
study. ``n_bootstrap`` and ``seed`` (constructor params)
control replicate count and RNG; defaults are 999 / ``None``.
trends_lin : bool, default False, keyword-only
When ``True``, applies paper Eq 17 linear-trend detrending
to per-event-time outcome evolutions. Mirrors R
``DIDHAD::did_had(..., trends_lin=TRUE)``. Per-group slope
is estimated as ``Y[g, F-1] - Y[g, F-2]``; each event-time
``e`` evolution is replaced by ``dy_dict[e] - (e+1)
slope`` (uniform formula that absorbs both effect-side
detrending and placebo-side anchor swap). Requires
``aggregate="event_study"`` AND ``F >= 3`` (panel must
include both ``F-1`` and ``F-2``); raises
``NotImplementedError`` on ``aggregate="overall"`` and
``ValueError`` on ``F < 3``. The "consumed" placebo at
event time ``e=-2`` is auto-dropped (R reduces max
placebo lag by 1 with the same effect). Mutually
exclusive with survey weighting (``survey_design`` /
``survey`` / ``weights``); raises ``NotImplementedError``
if combined. Default ``False`` preserves bit-exact
backcompat with all pre-PR fits.
Returns
-------
HeterogeneousAdoptionDiDResults
"""
# ---- aggregate / survey_design / survey / weights validation ----
if aggregate not in _VALID_AGGREGATES:
raise ValueError(
f"Invalid aggregate={aggregate!r}. Must be one of " f"{_VALID_AGGREGATES}."
)
# Three-way mutex on survey_design / survey / weights (data-in pattern).
n_set = sum(x is not None for x in (survey_design, survey, weights))
if n_set > 1:
raise ValueError(HAD_DUAL_KNOB_MUTEX_MSG_DATA_IN)
# ---- trends_lin scope gates (PR #389 / Phase 4 R-parity).
# `trends_lin=True` implements paper Eq 17 linear-trend detrending
# (per-group slope from Y[F-1]-Y[F-2], applied to per-event-time
# outcome evolutions). Requires F >= 3 (need both F-1 and F-2 in
# panel) and is currently event-study-only the overall path is
# 2-period and cannot accommodate the F-2 row.
if trends_lin:
if aggregate != "event_study":
raise NotImplementedError(
"HAD.fit(trends_lin=True) requires "
"aggregate='event_study' (the linear-trend slope "
"estimator needs Y at F-2, which a 2-period panel "
"does not contain). Pass a panel with at least 3 "
"periods and aggregate='event_study'; the per-"
"horizon arrays in the resulting "
"HeterogeneousAdoptionDiDEventStudyResults provide "
"the same single-effect / per-effect estimates as "
"the overall path."
)
if survey_design is not None or survey is not None or weights is not None:
raise NotImplementedError(
"HAD.fit(trends_lin=True) is not yet supported with "
"survey weighting (`survey_design=` / `survey=` / "
"`weights=`). The per-group slope estimator's "
"weighted variant (weighted-OLS slope? per-PSU slope?) "
"is not derived from the paper. Use trends_lin=True "
"WITHOUT survey weights, or use survey weights "
"WITHOUT trends_lin. Tracked in TODO.md as a follow-"
"up if user demand emerges."
)
# Soft deprecation: route legacy survey=/weights= aliases to
# survey_design=. The internal back-end paths (legacy weights= and
# survey= routing below) are unchanged; only the entry signature
# wraps them. The bit-exact back-compat invariant is preserved
# because we only rebind names, not values, and the legacy `survey`
# / `weights` variables are re-derived from `survey_design` for
# downstream consumption.
if survey is not None:
warnings.warn(HAD_DEPRECATION_MSG_SURVEY_KWARG, DeprecationWarning, stacklevel=2)
survey_design = survey
elif weights is not None:
warnings.warn(
HAD_DEPRECATION_MSG_WEIGHTS_KWARG_HAD_FIT,
DeprecationWarning,
stacklevel=2,
)
# weights= shortcut preserved as-is on the back end (the
# downstream `if weights is not None:` branch consumes the
# raw array directly via _aggregate_unit_weights). Don't
# rebind survey_design here the array is not a
# SurveyDesign and survey_design= cannot accept arrays.
else:
# Canonical path: survey_design= may be None or a SurveyDesign
# instance. Map back to the internal `survey` variable name
# so downstream code (legacy `if survey is not None:` branch)
# consumes the input transparently.
survey = survey_design
# Type guard on the data-in surface (PR #376 R8 P1): HAD.fit()
# accepts a SurveyDesign that gets resolved against `data` at fit
# time; a pre-resolved ResolvedSurveyDesign (or its
# make_pweight_design factory output) goes to the array-in pretest
# helpers, NOT to fit(). Reject explicitly with migration guidance
# rather than letting `survey.resolve(data)` AttributeError or
# `survey.weights` (a numpy array on Resolved) be misinterpreted as
# a column name. Mirrors the array-in helpers' isinstance-SurveyDesign
# rejection in stute_test/yatchew_hr_test/stute_joint_pretest.
if survey is not None and not hasattr(survey, "resolve"):
raise TypeError(
"HeterogeneousAdoptionDiD.fit: `survey_design=` accepts a "
"SurveyDesign instance (column-referencing, gets "
"`.resolve(data)`'d at fit time) on the data-in estimator "
"surface. Got "
f"{type(survey).__name__} (no `.resolve()` method). "
"If you have a pre-resolved ResolvedSurveyDesign or used "
"`make_pweight_design(arr)`, that pattern is for the "
"array-in pretest helpers (`stute_test`, `yatchew_hr_test`, "
"`stute_joint_pretest`). On HAD.fit, add the weights as a "
"column on `data` and pass "
"`survey_design=SurveyDesign(weights='col_name', ...)`."
)
# Dispatch the event-study path to a dedicated method so the
# single-period path stays unchanged (Phase 2a contract preserved).
# Note: event_study returns HeterogeneousAdoptionDiDEventStudyResults
# (distinct type from the overall path's HeterogeneousAdoptionDiDResults);
# the static return-type annotation reflects the common "overall" case
# to keep Phase 2a call-sites type-clean. Users explicitly passing
# aggregate="event_study" should annotate the result as
# HeterogeneousAdoptionDiDEventStudyResults.
if aggregate == "event_study":
return self._fit_event_study( # type: ignore[return-value]
data=data,
outcome_col=outcome_col,
dose_col=dose_col,
time_col=time_col,
unit_col=unit_col,
first_treat_col=first_treat_col,
survey=survey,
weights=weights,
cband=cband,
trends_lin=trends_lin,
)
# ---- Resolve effective fit-time state (local vars only, per
# feedback_fit_does_not_mutate_config; do not mutate self.*) ----
design_arg = self.design
d_lower_arg = self.d_lower
vcov_type_arg = self.vcov_type
robust_arg = self.robust
cluster_arg = self.cluster
# ---- Validate panel contract ----
t_pre, t_post = _validate_had_panel(
data, outcome_col, dose_col, time_col, unit_col, first_treat_col
)
# ---- Aggregate to unit-level first differences (no cluster yet) ----
# Defer cluster validation/extraction until after the design is
# resolved: the continuous paths ignore cluster= with a warning,
# so a malformed or irrelevant cluster column must not abort a
# valid continuous fit. Cluster extraction is re-run below only
# when resolved_design == "mass_point".
d_arr, dy_arr, _, _ = _aggregate_first_difference(
data,
outcome_col,
dose_col,
time_col,
unit_col,
t_pre,
t_post,
None,
)
# Resolve survey/weights into per-unit weights + optional
# ResolvedSurveyDesign (for PSU/strata/FPC composition).
# - `weights=` per-row array, no PSU/strata composition.
# - `survey=SurveyDesign(weights="col", ...)` resolve the design
# and produce a unit-level analogue for per-unit IF composition.
# - neither unweighted path.
weights_unit: Optional[np.ndarray] = None
raw_weights_unit: Optional[np.ndarray] = None
resolved_survey_unit: Any = None # ResolvedSurveyDesign (G,) when survey=
if weights is not None:
weights_unit = _aggregate_unit_weights(data, weights, unit_col)
# On the ``weights=`` shortcut, the passed array IS the raw
# pre-normalization weights no SurveyDesign.resolve() scaling.
raw_weights_unit = weights_unit
elif survey is not None:
if not hasattr(survey, "weights"):
raise TypeError(
f"survey= must be a SurveyDesign-like object with a "
f".weights attribute; got {type(survey).__name__}. "
f"Construct a SurveyDesign via diff_diff.survey."
)
if getattr(survey, "weights", None) is None:
raise NotImplementedError(
"survey= without weights is not yet supported. Pass "
"survey=SurveyDesign(weights='', ...) with a "
"per-row weight column."
)
# HAD's weighted local-linear treats ``weights`` as sampling
# (probability) weights: the kernel-composition formula
# ``W_combined = k((D-d)/h) w`` is the inverse-probability
# weighting convention. Frequency weights (``fweight``)
# would imply replicating observations, and analytic weights
# (``aweight``, inverse-variance) would imply a different
# inferential target. Reject those up front rather than
# silently reinterpreting.
weight_type = getattr(survey, "weight_type", "pweight")
if weight_type != "pweight":
raise NotImplementedError(
f"survey=SurveyDesign(weight_type={weight_type!r}) is "
f"not supported on HeterogeneousAdoptionDiD's "
f"continuous path. Only ``weight_type='pweight'`` "
f"(sampling / inverse-probability weights) is "
f"implemented in Phase 4.5 A. Frequency weights "
f"(fweight) and analytic weights (aweight) would "
f"imply different estimands and are not yet derived."
)
# Capture the RAW pre-normalization weight column before
# ``resolve()`` rescales pweights/aweights to mean=1. Needed
# below so ``compute_survey_metadata`` receives raw weights
# (per its contract ``sum_weights`` / ``weight_range`` are
# raw-scale quantities; passing normalized weights would make
# the metadata disagree with the ``weights=`` shortcut and
# drift from the docstring/test contract in ``survey.py``).
weights_col_name = survey.weights # known non-None from guard above
if weights_col_name not in data.columns:
raise ValueError(f"survey.weights column {weights_col_name!r} not found in data.")
raw_weights_row = np.asarray(data[weights_col_name].values, dtype=np.float64)
raw_weights_unit = _aggregate_unit_weights(data, raw_weights_row, unit_col)
# Resolve the SurveyDesign against the long-panel data. This
# validates column names, applies pweight/aweight normalization
# to mean=1, and extracts numpy arrays for all design columns.
resolved_survey_row = survey.resolve(data)
# Collapse design columns to unit-level (constant-within-unit
# invariant) so the IF-based variance composition operates at
# the G-unit scale that matches the local-linear fit.
resolved_survey_unit = _aggregate_unit_resolved_survey(
data, resolved_survey_row, unit_col
)
weights_unit = np.asarray(resolved_survey_unit.weights, dtype=np.float64)
# Zero-weight units (e.g. SurveyDesign.subpopulation() output, or
# a user-supplied pweight column with excluded observations) must
# not drive design resolution ``_detect_design`` / ``d_lower``
# / mass-point threshold / cohort counts run on the POSITIVE-
# weight subset. But the survey VARIANCE and ``SurveyMetadata``
# preserve the FULL ResolvedSurveyDesign (zero-weight PSUs /
# strata kept in the design with zero in-domain mass) that is
# the standard subpopulation / domain-estimation convention in
# ``diff_diff.survey``: keep the sampling frame, zero the
# contributions. The weighted kernel in ``lprobust`` drops
# zero-weight observations via its ``w > 0`` selector, and
# ``bias_corrected_local_linear`` zero-pads the returned IF back
# to the full unit ordering so the survey composition at the
# HAD level sees IF=0 for zero-weight units on the FULL design.
# (CI review PR #359 round 5 P0 + round 6 P1 cascade.)
d_arr_full = d_arr # unfiltered (G units); passed to _fit_continuous
dy_arr_full = dy_arr
weights_unit_full = weights_unit # may contain zeros; used for FIT
resolved_survey_unit_full = resolved_survey_unit # full design for VARIANCE
raw_weights_unit_full = raw_weights_unit # full for SurveyMetadata
if weights_unit is not None:
positive_mask = weights_unit > 0.0
if not bool(positive_mask.all()):
n_dropped = int((~positive_mask).sum())
warnings.warn(
f"HAD continuous path: {n_dropped} unit(s) have "
f"weight == 0 and are excluded from design resolution "
f"(auto-detect design, d_lower, mass-point threshold, "
f"cohort counts). They are RETAINED in the survey "
f"design for variance + SurveyMetadata (subpopulation "
f"convention: zero-weight contributions but full "
f"sampling frame), and their IF is 0 on the full "
f"design.",
UserWarning,
stacklevel=2,
)
# Filter arrays used for DESIGN-RESOLUTION ONLY.
d_arr = d_arr[positive_mask]
dy_arr = dy_arr[positive_mask]
weights_unit = weights_unit[positive_mask]
n_obs = int(d_arr.shape[0])
if n_obs < 3:
raise ValueError(
f"HAD requires at least 3 units for inference; got "
f"n_obs={n_obs} after aggregation."
)
# ---- Resolve design ----
if design_arg == "auto":
resolved_design = _detect_design(d_arr)
else:
resolved_design = design_arg
# ---- Extract cluster IDs (mass-point path only) ----
# Continuous paths ignore cluster= with a warning emitted later in
# the dispatch block; the cluster column is not read for them. On
# the mass-point path we now re-run the aggregation with
# cluster_col so validation (missing column / NaN / within-unit
# variance) fires only when cluster is actually going to be used.
cluster_arr: Optional[np.ndarray] = None
if resolved_design == "mass_point" and cluster_arg is not None:
_, _, cluster_arr, _ = _aggregate_first_difference(
data,
outcome_col,
dose_col,
time_col,
unit_col,
t_pre,
t_post,
cluster_arg,
)
# ---- Resolve d_lower ----
if resolved_design == "continuous_at_zero":
# Design 1' regime (paper Section 3.2) is defined at d_lower = 0.
# Reject explicit nonzero d_lower overrides front-door rather
# than silently coerce to zero. Tolerance family matches the
# Design 1 d_lower guards below.
if d_lower_arg is not None:
scale = max(1.0, float(np.max(np.abs(d_arr))))
if abs(float(d_lower_arg)) > 1e-12 * scale:
raise ValueError(
f"design='continuous_at_zero' (Design 1') requires "
f"d_lower == 0 within float tolerance (paper Section "
f"3.2 Design 1' regime). Got d_lower="
f"{float(d_lower_arg)!r}. For d_lower > 0 use "
f"design='continuous_near_d_lower' (local-linear "
f"boundary-limit estimator) or design='mass_point' "
f"(2SLS) as appropriate, or design='auto' which "
f"auto-detects the correct path from the dose "
f"distribution."
)
d_lower_val = 0.0
elif d_lower_arg is None:
d_lower_val = float(d_arr.min())
else:
d_lower_val = float(d_lower_arg)
# ---- Regime partition: d_lower > 0 for Design 1 paths ----
# Paper Section 3.2 partitions HAD into d_lower = 0 (Design 1',
# continuous_at_zero) and d_lower > 0 (Design 1, continuous_near
# _d_lower or mass_point). The auto-detect rule already enforces
# this partition; explicit overrides must respect it too, otherwise
# `design="mass_point", d_lower=0` returns a finite but
# paper-incompatible 2SLS result and `design="continuous_near_d_lower"`
# with d_lower=0 reduces to Design 1' algebra while mislabeling the
# estimand as `WAS_d_lower` and emitting the wrong Assumption 5/6
# warning. Use the same float-tolerance family as _detect_design's
# d.min()==0 tie-break.
if resolved_design in ("mass_point", "continuous_near_d_lower"):
scale = max(1.0, float(np.max(np.abs(d_arr))))
if abs(d_lower_val) d_lower}).
# The auto-detect rule already enforces this; explicit overrides
# must too, otherwise the wrong estimand is returned silently.
# Both guards are symmetric around the 2% threshold used in
# _detect_design() and Phase 1c's _validate_had_inputs().
if resolved_design in ("continuous_near_d_lower", "mass_point"):
d_min_orig = float(d_arr.min())
if d_min_orig > 0:
eps_mp = 1e-12 * max(1.0, abs(d_min_orig))
at_d_min_mask_orig = np.abs(d_arr - d_min_orig) _MASS_POINT_THRESHOLD
):
raise ValueError(
f"design='continuous_near_d_lower' cannot be used on a "
f"mass-point sample (modal fraction {modal_fraction_orig:.4f} "
f"at d.min()={d_min_orig!r} exceeds the "
f"{_MASS_POINT_THRESHOLD:.2f} threshold from paper Section "
f"3.2.4). Use design='mass_point' (Wald-IV / 2SLS) or "
f"design='auto' which will auto-detect. Forcing the "
f"continuous path on a mass-point sample would produce "
f"the wrong estimand."
)
if resolved_design == "mass_point" and modal_fraction_orig d_lower} requires d_lower to be
# the lower-support mass point. For the continuous-near-d_lower
# path, evaluating the local-linear fit after the regressor
# shift (D - d_lower) at boundary=0 only makes sense when the
# shift anchors to the realized sample minimum; otherwise the
# boundary evaluation is off-support (no observations near zero
# on the shifted scale) and Phase 1c's 5% plausibility heuristic
# may fail to catch the mismatch. We enforce d_lower == d.min()
# within float tolerance on both Design 1 paths; mismatched
# overrides raise with a clear pointer to the unsupported
# estimand.
if resolved_design in ("mass_point", "continuous_near_d_lower") and d_lower_arg is not None:
d_min = float(d_arr.min())
tol = 1e-12 * max(1.0, abs(d_min))
if abs(d_lower_val - d_min) > tol:
raise ValueError(
f"design={resolved_design!r} requires d_lower to equal "
f"the support infimum float(d.min())={d_min!r}; got "
f"d_lower={d_lower_val!r}. The paper's Design 1 "
f"estimators (Sections 3.2.2-3.2.4) identify at the "
f"lower-support boundary, not at an arbitrary "
f"threshold. Pass d_lower=None to auto-resolve, or "
f"d_lower=float(d.min()) explicitly. Non-support-"
f"infimum thresholds identify a different (LATE-like "
f"for mass_point, off-support for continuous_near_"
f"d_lower) estimand that is out of Phase 2a scope."
)
# Snap tolerance-accepted overrides back to the exact support
# infimum. Float-rounding drift matters downstream: on the
# mass-point path, `Z = d > d_lower` with d_lower = d.min() -
# puts the mass-point units into Z=1 (control group empties);
# on the continuous-near-d_lower path, d_lower = d.min() +
# makes `d - d_lower` negative and trips Phase 1c's
# _validate_had_inputs negative-dose guard. Snapping preserves
# the "within tolerance" contract while keeping downstream
# algebra exact.
d_lower_val = d_min
# ---- Compute cohort counts ----
if resolved_design == "mass_point":
eps = 1e-12 * max(1.0, abs(d_lower_val))
at_d_min_mask = np.abs(d_arr - d_lower_val) d_lower_val
n_mass_point: Optional[int] = int(at_d_min_mask.sum())
n_above_d_lower: Optional[int] = int(above_mask.sum())
n_treated = n_above_d_lower
n_control = int(at_d_min_mask.sum())
else:
n_mass_point = None
n_above_d_lower = None
above_mask = d_arr > d_lower_val
n_treated = int(above_mask.sum())
n_control = n_obs - n_treated
dose_mean = float(d_arr.mean())
# ---- Assumption 5/6 warning on Design 1 paths ----
# Paper Sections 3.2.2-3.2.4: when d_lower > 0 (Design 1 family),
# point identification of WAS_{d_lower} requires Assumption 6 in
# addition to parallel trends (Assumption 1-3); Assumption 5 gives
# only sign identification. These extra assumptions are NOT
# testable via pre-trends. Surface this to the user front-door so
# results are not silently interpreted as full point identification.
if resolved_design in ("continuous_near_d_lower", "mass_point"):
warnings.warn(
f"design={resolved_design!r} (Design 1, d_lower > 0) requires "
f"Assumption 6 from de Chaisemartin et al. (2026) for point "
f"identification of WAS_{{d_lower}}, or Assumption 5 for "
f"sign identification only. Neither is testable via "
f"pre-trends. Confirm the extra assumption is defensible "
f"for your setting before interpreting the returned "
f"point estimate as the WAS.",
UserWarning,
stacklevel=2,
)
# ---- Dispatch ----
if resolved_design in ("continuous_at_zero", "continuous_near_d_lower"):
# Warn when the user set a mass-point-only knob that's ignored
# on the continuous path. (Emitted per fit call; this is not
# suppressed after the first call.)
if vcov_type_arg is not None:
warnings.warn(
f"vcov_type={vcov_type_arg!r} is ignored on the "
f"'{resolved_design}' path (the continuous designs "
f"use the CCT-2014 robust SE from Phase 1c). "
f"vcov_type applies only to design='mass_point'.",
UserWarning,
stacklevel=2,
)
if robust_arg:
warnings.warn(
f"robust=True is ignored on the '{resolved_design}' "
f"path (the continuous designs use the CCT-2014 "
f"robust SE from Phase 1c unconditionally; the "
f"robust flag is a mass-point-only backward-compat "
f"alias for vcov_type).",
UserWarning,
stacklevel=2,
)
if cluster_arg is not None:
warnings.warn(
f"cluster={cluster_arg!r} is ignored on the "
f"'{resolved_design}' path in Phase 2a. Cluster-"
f"robust SE on the nonparametric path is exposed "
f"via diff_diff.bias_corrected_local_linear directly "
f"but not yet threaded through the estimator-level "
f"knob.",
UserWarning,
stacklevel=2,
)
# Fit on FULL (unfiltered) arrays so the IF aligns with the
# full survey design. bias_corrected_local_linear drops
# zero-weight rows internally for its validation + selector +
# fit, then zero-pads the IF back to full length. Survey
# composition below runs on the full design, preserving
# domain-estimation semantics.
att, se, bc_fit, bw_diag = self._fit_continuous(
d_arr_full,
dy_arr_full,
resolved_design,
d_lower_val,
weights_arr=weights_unit_full,
resolved_survey_unit=resolved_survey_unit_full,
)
inference_method = "analytical_nonparametric"
vcov_label: Optional[str] = None
cluster_label: Optional[str] = None
elif resolved_design == "mass_point":
# Review R4 P1: narrow the cluster+weighted rejection. Only
# survey= + cluster= is a silent-mismatch case (the
# Binder-TSL override would overwrite the CR1 sandwich while
# result metadata still advertises vcov_type='cr1'). The
# weights= shortcut + cluster= path just returns the
# weighted-CR1 sandwich from _fit_mass_point_2sls directly
# (no survey composition) and matches estimatr::iv_robust
# (se_type="stata") bit-exactly see
# tests/test_estimatr_iv_robust_parity.py::TestEstimatrIVRobustCR1Parity.
if cluster_arg is not None and resolved_survey_unit_full is not None:
raise NotImplementedError(
f"cluster={cluster_arg!r} + survey= on "
f"design='mass_point' is not yet supported: the "
f"survey path composes Binder-TSL variance via "
f"compute_survey_if_variance and would silently "
f"override the CR1 cluster-robust sandwich while "
f"result metadata still advertises "
f"vcov_type='cr1'. Pass cluster= alone "
f"(unweighted CR1), or weights= + cluster= "
f"(weighted-CR1 pweight sandwich; parity-tested vs "
f"estimatr::iv_robust se_type='stata'), or "
f"survey= alone (Binder-TSL). Combined cluster-"
f"robust + survey inference is deferred to a "
f"follow-up PR."
)
# Resolve the EFFECTIVE vcov family first (vcov_type_arg,
# with default-mapping from robust= when unset). Reject the
# classical-on-IF-consumption combination against the
# resolved value, NOT the raw kwarg, so the default
# `vcov_type=None, robust=False` case (which maps to
# classical) hits the guard too (review R5 P1 previous
# fix only fired on explicit vcov_type='classical').
if vcov_type_arg is None:
# Backward-compat: robust=True -> hc1, robust=False -> classical.
vcov_requested = "hc1" if robust_arg else "classical"
else:
vcov_requested = vcov_type_arg.lower()
# Review R3 P1 / R5 P1: the weighted mass-point path returns
# an HC1-scaled influence function (IF scale convention
# locks compute_survey_if_variance(psi, trivial) V_HC1[1,1]
# via the sqrt((n-1)/(n-k)) factor in _fit_mass_point_2sls).
# On the survey= path the analytical SE is ALWAYS overwritten
# with that HC1-scale Binder-TSL composition, so effective
# classical + survey= would silently report an HC1-target
# SE under a classical label. Reject until a
# classical-aligned IF is derived.
if vcov_requested == "classical" and resolved_survey_unit_full is not None:
raise NotImplementedError(
"vcov_type='classical' (resolved either explicit or "
"from the default robust=False mapping) + survey= on "
"design='mass_point' is not yet supported: the "
"survey path composes Binder-TSL variance via the "
"HC1-scale influence function, which targets V_HC1 "
"rather than the classical sandwich "
"V_cl = (Z'WX)^{-1}(Z'WZ)(X'WZ)^{-1}. Use "
"vcov_type='hc1' (or leave vcov_type unset with "
"robust=True) on the weighted path; a classical-"
"aligned IF derivation is queued for a follow-up PR."
)
# Phase 4.5 B: accept weights_unit (None on unweighted fits).
# return_influence=True only on the survey= path because
# Binder-TSL composition consumes the IF; the weights=
# shortcut and unweighted paths use the analytical sandwich
# SE directly. ``psi_mp`` is per-unit IF on -scale or None.
# Fit on FULL (unfiltered) arrays so the IF aligns with the
# full survey design (subpopulation convention: zero-weight
# units contribute 0 to all sums; IF zero-padded back to full
# length). Under unweighted fits d_arr_full == d_arr and
# weights_unit_full is None, so behavior is unchanged.
att, se, psi_mp = _fit_mass_point_2sls(
d_arr_full,
dy_arr_full,
d_lower_val,
cluster_arr,
vcov_requested,
weights=weights_unit_full,
return_influence=resolved_survey_unit_full is not None,
)
# Survey path: compose Binder-TSL variance from per-unit IF
# (replaces analytical sandwich SE). Mirrors continuous-path
# branch at lines 3082-3099. Under trivial resolved (single
# stratum, no PSU/FPC, uniform w), this reduces to analytical
# HC1 within the IF-scale-convention tolerance (atol=1e-10).
if resolved_survey_unit_full is not None and psi_mp is not None:
from diff_diff.survey import compute_survey_if_variance
v_survey = compute_survey_if_variance(psi_mp, resolved_survey_unit_full)
if np.isfinite(v_survey) and v_survey > 0.0:
se = float(np.sqrt(v_survey))
else:
se = float("nan")
bc_fit = None
bw_diag = None
inference_method = "analytical_2sls"
# Store the EFFECTIVE variance family so downstream consumers
# (to_dict, to_dataframe, summary) see the actual SE type that
# was computed. When cluster is supplied, _fit_mass_point_2sls
# unconditionally computes CR1 regardless of vcov_requested
# (e.g. classical+cluster -> CR1), so we surface that here.
vcov_label = "cr1" if cluster_arg is not None else vcov_requested
cluster_label = cluster_arg if cluster_arg is not None else None
else:
raise ValueError(f"Internal error: unhandled design={resolved_design!r}.")
# ---- Route all inference fields through safe_inference ----
# Survey path: use t-distribution with ``df_survey = n_psu -
# n_strata`` (or replicate-QR rank 1) so small-PSU designs
# don't get Normal-theory inference that overstates precision.
# Non-survey path (``weights=`` shortcut or unweighted): use
# the existing Normal-theory default.
df_infer: Optional[int] = None
if resolved_survey_unit is not None:
df_infer = resolved_survey_unit.df_survey
t_stat, p_value, conf_int = safe_inference(att, se, alpha=float(self.alpha), df=df_infer)
# Build survey metadata (repo-standard SurveyMetadata from
# diff_diff.survey.compute_survey_metadata) when weights/survey
# were supplied, so downstream report/diagnostic consumers can
# read attributes uniformly. HAD-specific extras (variance-
# formula label, effective-denominator value) live on dedicated
# result fields rather than being folded into the survey dict.
survey_metadata: Optional[SurveyMetadata] = None
variance_formula_label: Optional[str] = None
effective_dose_mean_value: Optional[float] = None
if weights_unit_full is not None:
if resolved_survey_unit_full is not None:
# survey= path: build metadata from the FULL
# ResolvedSurveyDesign (pre-zero-weight-filter), so
# ``n_strata`` / ``n_psu`` / ``df_survey`` / weight sums
# reflect the sampling frame, not the in-domain subset.
# Pass the RAW pre-normalization per-unit weights
# (captured before survey.resolve() rescaled pweights/
# aweights to mean=1) so ``sum_weights`` / ``weight_range``
# reflect the user-supplied scale matching both the
# ``weights=`` shortcut and ``compute_survey_metadata``'s
# contract.
assert raw_weights_unit_full is not None # set in survey= branch
survey_metadata = compute_survey_metadata(
resolved_survey_unit_full, raw_weights_unit_full
)
# Design-specific label continuous uses bias-corrected CCT
# IF, mass-point uses 2SLS IF; both route through Binder TSL.
variance_formula_label = (
"survey_binder_tsl_2sls"
if resolved_design == "mass_point"
else "survey_binder_tsl"
)
else:
# weights= shortcut: construct a minimal resolved
# SurveyDesign with the FULL user-supplied weights
# (including zero-weight units) so SurveyMetadata
# summarizes the full sample. No strata / PSU / FPC
# structure the shortcut is pweight-only by contract.
from diff_diff.survey import ResolvedSurveyDesign
minimal_resolved = ResolvedSurveyDesign(
weights=weights_unit_full,
weight_type="pweight",
strata=None,
psu=None,
fpc=None,
n_strata=1,
n_psu=int(weights_unit_full.shape[0]),
lonely_psu="remove",
combined_weights=True,
mse=False,
)
# weights_unit_full is already the raw user-supplied
# array (no SurveyDesign.resolve() normalization here).
survey_metadata = compute_survey_metadata(minimal_resolved, weights_unit_full)
# On the ``weights=`` shortcut, inference stays Normal
# (df=None in safe_inference) no PSU / strata / FPC
# composition. Clear the survey-only fields that
# ``compute_survey_metadata`` derives from the synthetic
# minimal design (``n_psu = G``, ``n_strata = 1``,
# ``df_survey = G 1``) so ``summary()`` / BusinessReport
# do not misdescribe the fit as a finite-df survey
# result. ``weight_type``, ``effective_n``,
# ``design_effect``, ``sum_weights``, and
# ``weight_range`` stay populated they describe the
# weighted sample regardless of inference family.
survey_metadata.n_strata = None
survey_metadata.n_psu = None
survey_metadata.df_survey = None
variance_formula_label = (
"pweight_2sls" if resolved_design == "mass_point" else "pweight"
)
# Expose the effective weighted denominator used by the
# beta-scale rescaling. Continuous paths use a weighted sample
# mean of d (or d d_lower). Mass-point uses the weighted
# Wald-IV dose-gap (the denominator of =
# dy_gap_w / dose_gap_w), computed from the FULL arrays
# (zero-weight units contribute 0 to both subgroup sums).
if resolved_design == "continuous_at_zero":
effective_dose_mean_value = float(np.average(d_arr_full, weights=weights_unit_full))
elif resolved_design == "continuous_near_d_lower":
effective_dose_mean_value = float(
np.average(d_arr_full - d_lower_val, weights=weights_unit_full)
)
elif resolved_design == "mass_point":
# Weighted Wald-IV dose gap: mean(d | Z=1, w) - mean(d | Z=0, w).
# Surface this as the "effective denominator" so downstream
# reporting displays the -scale denominator consistently
# across designs. Guard against empty subgroups (handled
# upstream by _fit_mass_point_2sls returning NaN, so we
# only reach here on a successful fit with positive mass
# on both sides).
Z_mp = (d_arr_full > d_lower_val).astype(np.float64)
pos_mp = weights_unit_full > 0
Z1_mp = (Z_mp == 1) & pos_mp
Z0_mp = (Z_mp == 0) & pos_mp
w_Z1_mp = float(weights_unit_full[Z1_mp].sum())
w_Z0_mp = float(weights_unit_full[Z0_mp].sum())
if w_Z1_mp > 0.0 and w_Z0_mp > 0.0:
effective_dose_mean_value = float(
(weights_unit_full[Z1_mp] * d_arr_full[Z1_mp]).sum() / w_Z1_mp
- (weights_unit_full[Z0_mp] * d_arr_full[Z0_mp]).sum() / w_Z0_mp
)
return HeterogeneousAdoptionDiDResults(
att=float(att),
se=float(se),
t_stat=float(t_stat),
p_value=float(p_value),
conf_int=(float(conf_int[0]), float(conf_int[1])),
alpha=float(self.alpha),
design=resolved_design,
target_parameter=_TARGET_PARAMETER[resolved_design],
d_lower=d_lower_val,
dose_mean=dose_mean,
n_obs=n_obs,
n_treated=n_treated,
n_control=n_control,
n_mass_point=n_mass_point,
n_above_d_lower=n_above_d_lower,
inference_method=inference_method,
vcov_type=vcov_label,
cluster_name=cluster_label,
survey_metadata=survey_metadata,
bandwidth_diagnostics=bw_diag,
bias_corrected_fit=bc_fit,
variance_formula=variance_formula_label,
effective_dose_mean=effective_dose_mean_value,
)
# ------------------------------------------------------------------
# Continuous-design dispatch (Design 1' + Design 1 continuous-near-d_lower)
# ------------------------------------------------------------------
def _fit_continuous(
self,
d_arr: np.ndarray,
dy_arr: np.ndarray,
resolved_design: str,
d_lower_val: float,
weights_arr: Optional[np.ndarray] = None,
resolved_survey_unit: Any = None, # ResolvedSurveyDesign (G,) or None
force_return_influence: bool = False,
) -> Tuple[float, float, Optional[BiasCorrectedFit], Optional[BandwidthResult]]:
"""Fit Phase 1c ``bias_corrected_local_linear`` and form the WAS estimate.
Implements de Chaisemartin, Ciccia, D'Haultfoeuille, and Knau
(2026) continuous-design estimators:
- Design 1' (``continuous_at_zero``), paper Theorem 1 /
Equation 3 (identification); Equation 7 (sample estimator):
beta = (E[Delta Y] - lim_{d v 0} E[Delta Y | D_2 0.0
if not bool(positive_mask.all()):
n_dropped = int((~positive_mask).sum())
warnings.warn(
f"HAD event-study: {n_dropped} unit(s) have weight == 0 "
f"and are excluded from design resolution (auto-detect, "
f"d_lower, mass-point threshold). Retained in the survey "
f"design for variance + SurveyMetadata (subpopulation "
f"convention).",
UserWarning,
stacklevel=2,
)
d_arr = d_arr[positive_mask]
dy_dict = {e: v[positive_mask] for e, v in dy_dict.items()}
weights_unit = weights_unit[positive_mask]
n_units = int(d_arr.shape[0])
if n_units < 3:
raise ValueError(
f"HAD event-study requires at least 3 positive-"
f"weight units for inference; got n_units={n_units} "
f"after the zero-weight filter."
)
# ---- Resolve design (once, from D_{g, F} distribution) ----
if design_arg == "auto":
resolved_design = _detect_design(d_arr)
else:
resolved_design = design_arg
# ---- Resolve d_lower ----
if resolved_design == "continuous_at_zero":
if d_lower_arg is not None:
scale = max(1.0, float(np.max(np.abs(d_arr))))
if abs(float(d_lower_arg)) > 1e-12 * scale:
raise ValueError(
f"design='continuous_at_zero' (Design 1') requires "
f"d_lower == 0 within float tolerance (paper Section "
f"3.2 Design 1' regime). Got d_lower="
f"{float(d_lower_arg)!r}. For d_lower > 0 use "
f"design='continuous_near_d_lower' or "
f"design='mass_point', or design='auto'."
)
d_lower_val = 0.0
elif d_lower_arg is None:
d_lower_val = float(d_arr.min())
else:
d_lower_val = float(d_lower_arg)
# ---- Regime-partition guards (mirror Phase 2a) ----
if resolved_design in ("mass_point", "continuous_near_d_lower"):
scale = max(1.0, float(np.max(np.abs(d_arr))))
if abs(d_lower_val) 0:
eps_mp = 1e-12 * max(1.0, abs(d_min_orig))
at_d_min_mask_orig = np.abs(d_arr - d_min_orig) _MASS_POINT_THRESHOLD
):
raise ValueError(
f"design='continuous_near_d_lower' cannot be used on a "
f"mass-point sample (modal fraction {modal_fraction_orig:.4f} "
f"at d.min()={d_min_orig!r} exceeds the "
f"{_MASS_POINT_THRESHOLD:.2f} threshold)."
)
if resolved_design == "mass_point" and modal_fraction_orig tol:
raise ValueError(
f"design={resolved_design!r} requires d_lower to equal "
f"the support infimum float(d.min())={d_min!r}; got "
f"d_lower={d_lower_val!r}."
)
d_lower_val = d_min # snap
dose_mean = float(d_arr.mean())
# ---- Extract cluster IDs on mass-point path only ----
cluster_arr: Optional[np.ndarray] = None
if resolved_design == "mass_point" and cluster_arg is not None:
# Review R4 P1: narrow the cluster+weighted guard (mirrors
# the static-path narrowing). Incompatible cases on the
# event-study path:
# (a) survey= + cluster=: Binder-TSL override would
# silently overwrite CR1.
# (b) weights= shortcut + cluster= + cband=True: the
# sup-t bootstrap normalizes HC1-scale perturbations
# by the CR1 analytical SE, producing an inconsistent
# variance family in the bootstrap t-distribution.
# weights= shortcut + cluster= + cband=False is fine: the
# per-horizon CR1 sandwich is returned as-is and no IF is
# consumed. Unweighted + cluster= also unchanged.
if resolved_survey_unit_full is not None:
raise NotImplementedError(
f"cluster={cluster_arg!r} + survey= on "
f"design='mass_point' event-study is not yet "
f"supported: the survey path composes Binder-TSL "
f"variance per horizon and would silently override "
f"the CR1 cluster-robust sandwich. Pass cluster= "
f"alone (unweighted CR1), or weights= + cluster= "
f"+ cband=False (weighted-CR1 per horizon), or "
f"survey= alone (Binder-TSL). Combined cluster-"
f"robust + survey event-study inference is deferred."
)
if weights_unit_full is not None and cband:
raise NotImplementedError(
f"cluster={cluster_arg!r} + weights= + cband=True "
f"on design='mass_point' event-study is not yet "
f"supported: the sup-t bootstrap uses an HC1-scale "
f"influence function and normalizes by the CR1 "
f"analytical SE, mixing variance families in the "
f"bootstrap t-distribution. Pass cband=False to "
f"disable the simultaneous band (pointwise CIs "
f"still use the weighted-CR1 sandwich per horizon), "
f"or drop cluster= to use the weighted-HC1 sandwich "
f"with sup-t."
)
_, _, cluster_arr, _, _ = _aggregate_multi_period_first_differences(
data_filtered,
outcome_col,
dose_col,
time_col,
unit_col,
F,
t_pre_list,
t_post_list,
cluster_arg,
)
# ---- One-time warnings (per fit call, not per horizon) ----
if resolved_design in ("continuous_near_d_lower", "mass_point"):
warnings.warn(
f"design={resolved_design!r} (Design 1, d_lower > 0) requires "
f"Assumption 6 from de Chaisemartin et al. (2026) for point "
f"identification of WAS_{{d_lower}}, or Assumption 5 for "
f"sign identification only. Neither is testable via "
f"pre-trends.",
UserWarning,
stacklevel=3,
)
if resolved_design in ("continuous_at_zero", "continuous_near_d_lower"):
if vcov_type_arg is not None:
warnings.warn(
f"vcov_type={vcov_type_arg!r} is ignored on the "
f"'{resolved_design}' path (continuous designs use the "
f"CCT-2014 robust SE from Phase 1c).",
UserWarning,
stacklevel=3,
)
if robust_arg:
warnings.warn(
f"robust=True is ignored on the '{resolved_design}' " f"path.",
UserWarning,
stacklevel=3,
)
if cluster_arg is not None:
warnings.warn(
f"cluster={cluster_arg!r} is ignored on the "
f"'{resolved_design}' path in Phase 2b (estimator-"
f"level cluster threading on the nonparametric path "
f"is queued for a follow-up PR).",
UserWarning,
stacklevel=3,
)
# ---- Resolve vcov label for mass-point ----
if resolved_design == "mass_point":
# Resolve the EFFECTIVE vcov family first (review R5 P1
# previous fix only fired on explicit vcov_type='classical'
# and missed the default vcov_type=None, robust=False
# 'classical' mapping).
if vcov_type_arg is None:
vcov_requested = "hc1" if robust_arg else "classical"
else:
vcov_requested = vcov_type_arg.lower()
# Review R3/R5 P1 (event-study arm): reject effective
# classical when the weighted path will compute the IF
# (always on survey= path; on weights= shortcut when
# cband=True the bootstrap divides HC1-scale perturbations
# by per-horizon analytical SE, so classical SE would
# give wrong t-stats). Matches the static-path rejection
# weighted mass-point paths use the HC1-scale IF
# convention uniformly.
_uses_if_matrix = resolved_survey_unit_full is not None or (
weights_unit_full is not None and cband
)
if vcov_requested == "classical" and _uses_if_matrix:
raise NotImplementedError(
"vcov_type='classical' (resolved either explicit "
"or from the default robust=False mapping) + "
"weights/survey= on design='mass_point' event-study "
"is not yet supported: the per-horizon IF matrix is "
"HC1-scale (targets V_HC1 via "
"compute_survey_if_variance) and mixing it with a "
"classical analytical SE either through the "
"survey Binder-TSL override or the sup-t bootstrap "
"normalization would produce an inconsistent "
"variance family. Use vcov_type='hc1' (or leave "
"vcov_type unset with robust=True) on the weighted "
"event-study path, or pass cband=False to skip the "
"bootstrap on the weights= shortcut."
)
inference_method = "analytical_2sls"
vcov_label: Optional[str] = "cr1" if cluster_arg is not None else vcov_requested
cluster_label: Optional[str] = cluster_arg if cluster_arg is not None else None
else:
vcov_requested = ""
inference_method = "analytical_nonparametric"
vcov_label = None
cluster_label = None
# ---- Per-horizon loop ----
# On the weighted path, every horizon uses the FULL arrays
# (zero-weight units padded to 0 contribution) so the stacked IF
# matrix aligns with the full survey design. On unweighted fits,
# `d_arr_full == d_arr` and `dy_dict_full == dy_dict`, so this
# branch is a no-op.
event_times_sorted = sorted(dy_dict.keys())
n_horizons = len(event_times_sorted)
# Use the full arrays when weighted so the IF matrix aligns with
# the survey design; unweighted uses the same d_arr either way.
weighted_es = weights_unit_full is not None
d_arr_loop = d_arr_full if weighted_es else d_arr
dy_dict_loop = dy_dict_full if weighted_es else dy_dict
G_full = int(d_arr_full.shape[0])
att_arr = np.full(n_horizons, np.nan, dtype=np.float64)
se_arr = np.full(n_horizons, np.nan, dtype=np.float64)
t_arr = np.full(n_horizons, np.nan, dtype=np.float64)
p_arr = np.full(n_horizons, np.nan, dtype=np.float64)
ci_lo_arr = np.full(n_horizons, np.nan, dtype=np.float64)
ci_hi_arr = np.full(n_horizons, np.nan, dtype=np.float64)
# Review R4 P2: report the POSITIVE-WEIGHT contributing sample
# size, not the full pre-filter design size. Matches the
# static-path n_obs contract where zero-weight units are
# excluded from the reported count (survey_metadata still
# carries the full-design effective_n / n_psu / etc.).
n_obs_arr = np.full(n_horizons, n_units, dtype=np.int64)
# Two IF-consumption flags (review R6 P2): the PER-HORIZON IF is
# needed when the survey= path composes Binder-TSL variance (via
# compute_survey_if_variance inside _fit_continuous or the
# mass-point override below); the STACKED (G, H) IF matrix is
# needed only when the sup-t multiplier bootstrap runs
# (``cband=True`` on the weighted path). Splitting them avoids
# allocating / filling Psi on the common opt-out path
# ``cband=False`` + weights= shortcut, where no IF consumer
# exists.
needs_per_horizon_if = resolved_survey_unit_full is not None or (weighted_es and cband)
needs_stacked_if_matrix = weighted_es and cband
if needs_stacked_if_matrix:
Psi = np.full((G_full, n_horizons), np.nan, dtype=np.float64)
else:
Psi = np.zeros((0, 0), dtype=np.float64) # sentinel, not used
# Collect per-horizon diagnostics on continuous paths. Entries may be
# None for horizons where ``_fit_continuous`` caught a degenerate
# bandwidth-selector failure (constant/perfectly-linear outcome).
bc_fits: Optional[List[Optional[BiasCorrectedFit]]] = (
[] if resolved_design in ("continuous_at_zero", "continuous_near_d_lower") else None
)
bw_diags: Optional[List[Optional[BandwidthResult]]] = (
[] if resolved_design in ("continuous_at_zero", "continuous_near_d_lower") else None
)
# df_survey for t-inference on survey= path (mirrors static path).
df_infer: Optional[int] = None
if resolved_survey_unit_full is not None:
df_infer = resolved_survey_unit_full.df_survey
# On the weighted event-study path, the sup-t multiplier bootstrap
# operates on the per-horizon IF matrix, so we must force the IF
# computation even on the ``weights=`` shortcut (no survey
# structure _fit_continuous normally skips IF). Pass through
# the actual ``resolved_survey_unit_full`` (None on shortcut) so
# the per-horizon analytical SE still matches the static-path
# convention (bc_fit.se_robust on shortcut; Binder-TSL on
# survey=). IF return is gated on `force_return_influence=True`.
# Track the Binder-TSL den for continuous paths so we can
# reconstruct the per-unit IF (psi / den) for the sup-t bootstrap
# where both numerator IF and denominator divide are needed.
for i, e in enumerate(event_times_sorted):
dy_e = dy_dict_loop[e]
if resolved_design in ("continuous_at_zero", "continuous_near_d_lower"):
att_e, se_e, bc_fit_e, bw_diag_e = self._fit_continuous(
d_arr_loop,
dy_e,
resolved_design,
d_lower_val,
weights_arr=weights_unit_full,
resolved_survey_unit=resolved_survey_unit_full,
# Force IF return only when the sup-t bootstrap
# needs the stacked matrix AND the survey= gate
# won't already produce it. Under survey= path,
# _fit_continuous returns the IF automatically
# (resolved_survey_unit_full != None); under the
# weights= shortcut + cband=True, force it here;
# otherwise skip the O(G) IF work (review R6 P2).
force_return_influence=(
needs_stacked_if_matrix and resolved_survey_unit_full is None
),
)
if bc_fits is not None:
bc_fits.append(bc_fit_e)
if bw_diags is not None:
bw_diags.append(bw_diag_e)
# Collect per-unit IF on -scale (psi_bc / den) into
# Psi ONLY when the sup-t bootstrap will consume it.
if (
needs_stacked_if_matrix
and bc_fit_e is not None
and bc_fit_e.influence_function is not None
):
if resolved_design == "continuous_at_zero":
den_e = float(np.average(d_arr_full, weights=weights_unit_full))
else:
den_e = float(
np.average(
d_arr_full - d_lower_val,
weights=weights_unit_full,
)
)
if abs(den_e) > 1e-12:
Psi[:, i] = bc_fit_e.influence_function / abs(den_e)
elif resolved_design == "mass_point":
att_e, se_e, psi_e = _fit_mass_point_2sls(
d_arr_loop,
dy_e,
d_lower_val,
cluster_arr,
vcov_requested,
weights=weights_unit_full,
# Return IF only when a consumer exists: survey=
# path needs it for per-horizon Binder-TSL override;
# weights= shortcut + cband=True needs it for the
# bootstrap. weights= shortcut + cband=False skips
# IF computation entirely (review R6 P2).
return_influence=needs_per_horizon_if,
)
# Survey path: override analytical sandwich SE with
# Binder-TSL via compute_survey_if_variance (matches
# continuous-path convention from PR #359).
if resolved_survey_unit_full is not None and psi_e is not None:
from diff_diff.survey import compute_survey_if_variance
v_survey = compute_survey_if_variance(psi_e, resolved_survey_unit_full)
if np.isfinite(v_survey) and v_survey > 0.0:
se_e = float(np.sqrt(v_survey))
else:
se_e = float("nan")
if needs_stacked_if_matrix and psi_e is not None:
Psi[:, i] = psi_e
else:
raise ValueError(f"Internal error: unhandled design={resolved_design!r}.")
t_stat_e, p_value_e, conf_int_e = safe_inference(
att_e, se_e, alpha=float(self.alpha), df=df_infer
)
att_arr[i] = float(att_e)
se_arr[i] = float(se_e)
t_arr[i] = float(t_stat_e)
p_arr[i] = float(p_value_e)
ci_lo_arr[i] = float(conf_int_e[0])
ci_hi_arr[i] = float(conf_int_e[1])
# ---- Sup-t simultaneous confidence band (weighted + cband only) ----
cband_low_arr: Optional[np.ndarray] = None
cband_high_arr: Optional[np.ndarray] = None
cband_crit_value: Optional[float] = None
cband_method_label: Optional[str] = None
cband_n_bootstrap_eff: Optional[int] = None
if weighted_es and cband and n_horizons >= 1:
# Review R7 P0: the per-unit influence function returned by
# _fit_continuous / _fit_mass_point_2sls is HC1-scaled per
# the PR #359 convention compute_survey_if_variance(psi,
# trivial_resolved) V_HC1. Routing the weights= shortcut
# through the unit-level ``resolved_survey=None`` branch of
# _sup_t_multiplier_bootstrap would normalize against raw
# sum(psi) = ((n-1)/n) V_HC1, producing silently too-
# narrow simultaneous bands. Construct a synthetic trivial
# ResolvedSurveyDesign on the weights= shortcut so the
# bootstrap always fires the survey-aware branch (centered
# + sqrt(n/(n-1))-corrected), matching the variance family
# of the analytical per-horizon SE.
if resolved_survey_unit_full is not None:
resolved_for_bootstrap: Any = resolved_survey_unit_full
else:
from diff_diff.survey import ResolvedSurveyDesign
assert weights_unit_full is not None # weighted_es invariant
resolved_for_bootstrap = ResolvedSurveyDesign(
weights=weights_unit_full,
weight_type="pweight",
strata=None,
psu=None,
fpc=None,
n_strata=1,
n_psu=int(weights_unit_full.shape[0]),
lonely_psu="remove",
combined_weights=True,
mse=False,
)
q, cband_low_arr, cband_high_arr, _n_valid = _sup_t_multiplier_bootstrap(
influence_matrix=Psi,
att_per_horizon=att_arr,
se_per_horizon=se_arr,
resolved_survey=resolved_for_bootstrap,
n_bootstrap=n_bootstrap_eff,
alpha=float(self.alpha),
seed=seed_eff,
)
cband_crit_value = q
cband_method_label = "multiplier_bootstrap"
cband_n_bootstrap_eff = n_bootstrap_eff
# ---- Build survey metadata + variance_formula + effective_dose_mean
# (mirrors static-path branch). ----
survey_metadata: Optional[SurveyMetadata] = None
variance_formula_label: Optional[str] = None
effective_dose_mean_value: Optional[float] = None
if weights_unit_full is not None:
if resolved_survey_unit_full is not None:
assert raw_weights_unit_full is not None
survey_metadata = compute_survey_metadata(
resolved_survey_unit_full, raw_weights_unit_full
)
variance_formula_label = (
"survey_binder_tsl_2sls"
if resolved_design == "mass_point"
else "survey_binder_tsl"
)
else:
from diff_diff.survey import ResolvedSurveyDesign
minimal_resolved = ResolvedSurveyDesign(
weights=weights_unit_full,
weight_type="pweight",
strata=None,
psu=None,
fpc=None,
n_strata=1,
n_psu=int(weights_unit_full.shape[0]),
lonely_psu="remove",
combined_weights=True,
mse=False,
)
survey_metadata = compute_survey_metadata(minimal_resolved, weights_unit_full)
survey_metadata.n_strata = None
survey_metadata.n_psu = None
survey_metadata.df_survey = None
variance_formula_label = (
"pweight_2sls" if resolved_design == "mass_point" else "pweight"
)
if resolved_design == "continuous_at_zero":
effective_dose_mean_value = float(np.average(d_arr_full, weights=weights_unit_full))
elif resolved_design == "continuous_near_d_lower":
effective_dose_mean_value = float(
np.average(d_arr_full - d_lower_val, weights=weights_unit_full)
)
elif resolved_design == "mass_point":
Z_mp = (d_arr_full > d_lower_val).astype(np.float64)
pos_mp = weights_unit_full > 0
Z1_mp = (Z_mp == 1) & pos_mp
Z0_mp = (Z_mp == 0) & pos_mp
w_Z1_mp = float(weights_unit_full[Z1_mp].sum())
w_Z0_mp = float(weights_unit_full[Z0_mp].sum())
if w_Z1_mp > 0.0 and w_Z0_mp > 0.0:
effective_dose_mean_value = float(
(weights_unit_full[Z1_mp] * d_arr_full[Z1_mp]).sum() / w_Z1_mp
- (weights_unit_full[Z0_mp] * d_arr_full[Z0_mp]).sum() / w_Z0_mp
)
return HeterogeneousAdoptionDiDEventStudyResults(
event_times=np.asarray(event_times_sorted, dtype=np.int64),
att=att_arr,
se=se_arr,
t_stat=t_arr,
p_value=p_arr,
conf_int_low=ci_lo_arr,
conf_int_high=ci_hi_arr,
n_obs_per_horizon=n_obs_arr,
alpha=float(self.alpha),
design=resolved_design,
target_parameter=_TARGET_PARAMETER[resolved_design],
d_lower=d_lower_val,
dose_mean=dose_mean,
F=F,
# Review R4 P2: report positive-weight contributing count
# (matches n_obs_per_horizon; full-design size surfaces
# through survey_metadata.n_psu / effective_n / etc.).
n_units=n_units,
inference_method=inference_method,
vcov_type=vcov_label,
cluster_name=cluster_label,
survey_metadata=survey_metadata,
bandwidth_diagnostics=bw_diags,
bias_corrected_fit=bc_fits,
filter_info=filter_info,
variance_formula=variance_formula_label,
effective_dose_mean=effective_dose_mean_value,
cband_low=cband_low_arr,
cband_high=cband_high_arr,
cband_crit_value=cband_crit_value,
cband_method=cband_method_label,
cband_n_bootstrap=cband_n_bootstrap_eff,
)