From 9db57062bac9d146714b2d266cf3f442ad28bc3f Mon Sep 17 00:00:00 2001 From: igerber Date: Sat, 25 Jul 2026 17:57:31 -0400 Subject: [PATCH 01/11] feat(v4): aggregation contract module + pure CS event-study aggregator Groundwork for Phase 2b PR 1 (spec section 6, ledger rows M-020..M-027). New diff_diff/aggregation.py: - AggregationResult(BaseResults): the tabular post-fit aggregation container with a pinned 12-column schema and summary/to_dict/to_dataframe. Three choices, each grounded in what CS actually produces rather than assumed: "target" is PER-ROW so one container can carry two aligned estimands over the same labels; "n_kind" uses EventStudyResults' existing vocabulary ("groups"/"cells"/"obs"/"clusters") rather than inventing units/obs; and "weight" is nullable, because CS's group aggregation weights (g,t) cells equally WITHIN each cohort and has no cross-cohort mass to normalize - populating it would be a fabricated number. - AggregationKit: the compact fit-time payload (bookkeeping + influence + alpha/anticipation/cband), excluding the data matrices. - BootstrapReplaySpec: a value-bound replay description. Retaining the estimator's ReplayableWeightStream does not work - it holds a function-local closure (unpicklable) whose body reads self.n_bootstrap and self.bootstrap_weights LAZILY, so a post-fit set_params silently truncates the replay. Recording the generator state plus the parameters BY VALUE and rebuilding through a module-level factory replays bit-identically, pickles in ~255 bytes, and is immune to later mutation. - AggregationMixin: shared validation and dispatch, applied PER RESULTS CLASS. Never on BaseResults: resolve_locator uses inspect.getattr_static, which walks the MRO, so a base-class aggregate() would make every still-planned ledger row's "new" locator resolve and fail test_row_matches_reality. Pure CS event-study aggregator: _aggregate_event_study returned only its effects dict and stashed four more values on self (_event_study_overall / _df_used / _vcov / _vcov_index). It now returns EventStudyAggregation carrying all five, so it can run post-fit from a retained kit without mutating its host. StaggeredTripleDifference inherits the same mixin and is the ONLY reader of _event_study_overall (its Eq. 4.14 overall_att_es), so it is updated in the same diff. fit()'s stale-state reset is gone: the cross-fit leak it guarded against is now impossible by construction. Numerically inert. A 10-scenario golden captured BEFORE the refactor (all four aggregate values, a balance_e sweep, repeated cross-sections, a bootstrapped fit and anticipation=1) is reproduced at atol=rtol=1e-14. 980 tests pass across the staggered, SDDD, event-study-surface, pretrends, honest-did, diagnostic-report, R-parity and ledger suites. --- diff_diff/__init__.py | 4 + diff_diff/aggregation.py | 477 +++++++++++++++++++++++++++++ diff_diff/staggered.py | 23 +- diff_diff/staggered_aggregation.py | 69 ++++- diff_diff/staggered_triple_diff.py | 13 +- docs/doc-deps.yaml | 10 + 6 files changed, 568 insertions(+), 28 deletions(-) create mode 100644 diff_diff/aggregation.py diff --git a/diff_diff/__init__.py b/diff_diff/__init__.py index fc8f8a15c..139e9d229 100644 --- a/diff_diff/__init__.py +++ b/diff_diff/__init__.py @@ -30,6 +30,9 @@ ) from diff_diff._guides_api import get_llm_guide from diff_diff.agent_workflow import agent_workflow +from diff_diff.aggregation import ( + AggregationResult, +) from diff_diff.bacon import ( BaconDecomposition, BaconDecompositionResults, @@ -586,6 +589,7 @@ "BaseResults", "Diagnostic", "EventStudyResults", + "AggregationResult", ] # Agent-facing entrypoints surface first in dir(diff_diff). LLM agents diff --git a/diff_diff/aggregation.py b/diff_diff/aggregation.py new file mode 100644 index 000000000..a414c409b --- /dev/null +++ b/diff_diff/aggregation.py @@ -0,0 +1,477 @@ +"""Post-fit aggregation surface (4.0 program Phase 2, spec section 6). + +``results.aggregate(type=...)`` replaces the fit-time ``fit(aggregate=)`` +argument (ledger rows M-020..M-027): estimate once, aggregate as a post-fit +step - the ecosystem's strongest norm (``did::aggte``, ``etwfe::emfx``, Stata +``estat aggregation``). + +Three pieces live here: + +- :class:`AggregationResult` - the tabular container every non-event-study + aggregation returns (row M-122). ``type="event_study"`` returns the + :class:`~diff_diff.results_base.EventStudyResults` container instead, which + is where that surface's public exposure finally lands (row M-092). +- :class:`AggregationKit` - the compact per-estimator payload retained at fit + time so re-aggregation needs neither a refit nor the source frame. +- :class:`AggregationMixin` - shared validation and dispatch, applied PER + RESULTS CLASS. + +The mixin is deliberately NOT mixed into +:class:`~diff_diff.results_base.BaseResults`: ``test_v4_matrix``'s +``resolve_locator`` uses ``inspect.getattr_static``, which walks the MRO, so a +base-class ``aggregate`` would make every still-``planned`` ledger row's +``new`` locator resolve and fail ``test_row_matches_reality``. A results class +gains the mixin in the same diff that flips its row. +""" + +from dataclasses import dataclass, field +from typing import Any, Dict, Optional, Tuple + +import numpy as np +import pandas as pd + +from diff_diff.results_base import BaseResults, _json_safe_label + +__all__ = ["AggregationResult", "AGGREGATION_SCHEMA"] + + +#: Closed aggregation vocabulary (spec section 6). Per-estimator support is a +#: SUBSET of this set - an estimator asked for a type it does not implement +#: raises ``ValueError`` naming what it does support, never silently falling +#: back. Estimator-specific extras (e.g. ContinuousDiD's ``"dose"``) are +#: declared by that estimator's results class, not added here. +AGGREGATION_VOCABULARY: Tuple[str, ...] = ( + "simple", + "event_study", + "group", + "calendar", +) + +#: Pinned column schema of ``AggregationResult.to_dataframe()`` - identical for +#: every producer, mirroring ``EVENT_STUDY_SCHEMA``'s contract (spec section 5). +AGGREGATION_SCHEMA: Tuple[str, ...] = ( + "level", + "label", + "target", + "att", + "se", + "t_stat", + "p_value", + "conf_int_lower", + "conf_int_upper", + "n", + "weight", + "df", +) + + +def _sortable(labels: np.ndarray) -> bool: + """Can ``labels`` be ordered without raising? + + Cohort labels are usually numeric, but the column is object-dtype in + general and may be mixed-type, where a naive ``argsort`` raises + ``TypeError``. Producer order is preserved in that case rather than + guessing an order. + """ + try: + np.argsort(labels, kind="stable") + except TypeError: + return False + return True + + +@dataclass +class AggregationResult(BaseResults): + """One post-fit aggregation, as a table (spec section 6, row M-122). + + Columnar arrays index-aligned to ``label``. Values are computed by the + producing estimator's aggregation machinery and stored here verbatim - + this container never re-derives inference. + + Parameters + ---------- + level : str + The aggregation type that produced this table - one of + :data:`AGGREGATION_VOCABULARY` or a documented per-estimator extra. + label : np.ndarray + Per-row aggregation key: the cohort for ``"group"``, the calendar + period for ``"calendar"``, the dose for ``"dose"``. A single + ``"overall"`` entry for ``"simple"``. + target : np.ndarray + Per-row estimand discriminator, so one container can carry two + aligned estimands over the same labels (ContinuousDiD's ATT(d) and + ACRT(d) become 2N rows). ``"att"`` where an estimator has one. + att, se, t_stat, p_value : np.ndarray + The canonical quintet, per row. On a bootstrapped fit ``t_stat`` / + ``p_value`` / the interval are the producer's percentile-bootstrap + statistics carried through unchanged - NOT recomputed analytically. + conf_int_lower, conf_int_upper : np.ndarray + Interval bounds at the fit's ``alpha``. + n : np.ndarray + Per-row count as float, NaN where the producer records none. Its + SEMANTIC is ``n_kind`` - never assume units. + n_kind : str or None + Semantic of ``n``, from the same vocabulary + :class:`~diff_diff.results_base.EventStudyResults` uses + (``"groups"`` / ``"cells"`` / ``"obs"`` / ``"clusters"``). ``None`` + when the producer records no count. + weight : np.ndarray or None + Normalized aggregation mass per row, summing to 1 within one + ``(level, target)`` group. ``None`` where no per-row mass exists - + CallawaySantAnna's ``"group"`` aggregation weights ``(g, t)`` cells + equally WITHIN each cohort and has no cross-cohort mass, so inventing + one would be a fabricated number. + df : np.ndarray + Per-row inference degrees of freedom, NaN where none governed the + stored p-value. Entirely NaN on bootstrap fits, whose percentile + inference uses no df. + alpha : float + Significance level the interval was computed at. + estimator : str or None + Producing estimator class name, for provenance in ``summary()``. + """ + + level: str + label: np.ndarray + target: np.ndarray + att: np.ndarray + se: np.ndarray + t_stat: np.ndarray + p_value: np.ndarray + conf_int_lower: np.ndarray + conf_int_upper: np.ndarray + n: np.ndarray + df: np.ndarray + alpha: float = 0.05 + n_kind: Optional[str] = None + weight: Optional[np.ndarray] = None + estimator: Optional[str] = None + + _COLUMN_FIELDS: Tuple[str, ...] = field( + default=( + "label", + "target", + "att", + "se", + "t_stat", + "p_value", + "conf_int_lower", + "conf_int_upper", + "n", + "df", + ), + repr=False, + compare=False, + ) + + def __post_init__(self) -> None: + self.label = np.asarray(self.label, dtype=object) + n_rows = self.label.shape[0] + if self.label.ndim != 1: + raise ValueError( + f"AggregationResult label must be one-dimensional; got shape {self.label.shape}." + ) + + for name in ("att", "se", "t_stat", "p_value", "conf_int_lower", "conf_int_upper", "n"): + arr = np.asarray(getattr(self, name), dtype=float) + if arr.shape != (n_rows,): + raise ValueError( + f"AggregationResult field {name!r} has shape {arr.shape}; " + f"expected ({n_rows},) to align with label." + ) + setattr(self, name, arr) + + target = np.asarray(self.target, dtype=object) + if target.shape != (n_rows,): + raise ValueError( + f"AggregationResult target has shape {target.shape}; " + f"expected ({n_rows},) - it is a PER-ROW discriminator, not a scalar." + ) + self.target = target + + # df: scalar broadcasts across rows (the EventStudyResults convention); + # NaN wherever no df governed the stored p-value. + if self.df is None: + df_arr = np.full(n_rows, np.nan) + elif np.ndim(self.df) == 0: + df_arr = np.full(n_rows, float(self.df)) + else: + df_arr = np.asarray(self.df, dtype=float) + if df_arr.shape != (n_rows,): + raise ValueError( + f"AggregationResult df has shape {df_arr.shape}; expected ({n_rows},) or scalar." + ) + df_arr[~np.isfinite(self.p_value)] = np.nan + self.df = df_arr + + if self.weight is not None: + w = np.asarray(self.weight, dtype=float) + if w.shape != (n_rows,): + raise ValueError( + f"AggregationResult weight has shape {w.shape}; expected ({n_rows},) or None." + ) + self.weight = w + + # ------------------------------------------------------------------ # + # Serialization (spec section 5: every main results class has all three) + # ------------------------------------------------------------------ # + + def to_dataframe(self) -> pd.DataFrame: + """Return the pinned :data:`AGGREGATION_SCHEMA` columns, in order. + + Rows are ordered by ``label`` when the labels are homogeneously + sortable, and in producer order otherwise (mixed-type cohort labels + cannot be ordered without raising). + """ + data: Dict[str, Any] = { + "level": self.level, + "label": self.label, + "target": self.target, + "att": self.att, + "se": self.se, + "t_stat": self.t_stat, + "p_value": self.p_value, + "conf_int_lower": self.conf_int_lower, + "conf_int_upper": self.conf_int_upper, + "n": self.n, + "weight": self.weight if self.weight is not None else np.nan, + "df": self.df, + } + frame = pd.DataFrame(data, columns=list(AGGREGATION_SCHEMA)) + if len(frame) > 1 and _sortable(self.label): + order = np.argsort(self.label, kind="stable") + frame = frame.iloc[order].reset_index(drop=True) + return frame + + def to_dict(self) -> Dict[str, Any]: + """Canonical-name mapping (deprecated names never leak into output).""" + out: Dict[str, Any] = { + "level": self.level, + "alpha": self.alpha, + "n_kind": self.n_kind, + "estimator": self.estimator, + "label": [_json_safe_label(v) for v in self.label], + "target": [str(v) for v in self.target], + } + for name in ( + "att", + "se", + "t_stat", + "p_value", + "conf_int_lower", + "conf_int_upper", + "n", + "df", + ): + out[name] = [float(v) for v in getattr(self, name)] + out["weight"] = None if self.weight is None else [float(v) for v in self.weight] + return out + + def summary(self, alpha: Optional[float] = None) -> str: + """Human-readable table. ``alpha`` is display-only - it does not + recompute the stored interval, which was fixed at aggregation time.""" + a = self.alpha if alpha is None else alpha + who = self.estimator or "estimator" + lines = [ + f"{who} - aggregate(type={self.level!r})", + "=" * 64, + ] + if len(self.label) == 0: + lines.append("(no rows - the aggregation selected no cells)") + return "\n".join(lines) + + n_label = "n" if self.n_kind is None else f"n[{self.n_kind}]" + lines.append(f"{'label':>14} {'ATT':>11} {'SE':>10} {'t':>8} {'p':>8} {n_label:>10}") + lines.append("-" * 64) + frame = self.to_dataframe() + for _, row in frame.iterrows(): + n_disp = "" if not np.isfinite(row["n"]) else f"{row['n']:.0f}" + lines.append( + f"{str(row['label']):>14} {row['att']:>11.4f} {row['se']:>10.4f} " + f"{row['t_stat']:>8.3f} {row['p_value']:>8.4f} {n_disp:>10}" + ) + lines.append("-" * 64) + lines.append(f"Confidence intervals at alpha={a}.") + if self.weight is None: + lines.append("Per-row aggregation weights are not defined for this level.") + return "\n".join(lines) + + +@dataclass +class AggregationKit: + """Compact fit-time payload enabling post-fit re-aggregation. + + Built and attached DURING ``fit()`` - neither the estimator's + ``precomputed`` bookkeeping nor its influence-function payload survives + the call, so this cannot be reconstructed afterwards. + + Deliberately excludes the data matrices (``outcome_matrix``, + ``covariate_matrix``, ``obs_outcome``, ``obs_covariates``): re-aggregation + reads only unit-level bookkeeping, so the source panel is never retained. + + Attributes + ---------- + bookkeeping : dict + The aggregation-relevant subset of the estimator's ``precomputed`` + mapping. O(n_units) on panel fits; several entries are + observation-length on repeated cross-sections, where + ``all_units = np.arange(n_obs)`` by construction. + influence : dict + Per-``(g, t)`` influence-function payload. The DOMINANT retained + object, roughly O(n_units x n_gt). + alpha, anticipation : float, int + The only two estimator attributes the aggregation machinery reads. + Carried explicitly so the aggregators need no estimator reference. + cband : bool + Whether the fit requested simultaneous bands. Retained because + ``cband_crit_value`` is ``None`` both when bands were disabled and + when no aggregation ran, so it cannot distinguish the two. + bootstrap : AggregationKit.BootstrapReplaySpec or None + Value-bound bootstrap replay description; ``None`` on analytical fits. + """ + + bookkeeping: Dict[str, Any] + influence: Dict[Any, Any] + alpha: float + anticipation: int + cband: bool + bootstrap: Optional["BootstrapReplaySpec"] = None + + +@dataclass +class BootstrapReplaySpec: + """Value-bound description of a bootstrap weight stream. + + Retaining the estimator's ``ReplayableWeightStream`` directly does not + work: it stores a function-local closure (unpicklable) whose body reads + ``self.n_bootstrap`` / ``self.bootstrap_weights`` LAZILY, so a post-fit + ``set_params(n_bootstrap=...)`` silently changes - and can truncate - the + replayed stream. + + This records the generator state plus the parameters BY VALUE and rebuilds + the stream through a module-level factory, which replays bit-identically, + pickles, and is immune to later mutation of the estimator. + """ + + bitgen_state: Dict[str, Any] + n_bootstrap: int + n_units: int + weight_type: str + block_size: Optional[int] = None + expand_index: Optional[np.ndarray] = None + + def rebuild(self) -> Any: + """Reconstruct the replayable weight stream.""" + from diff_diff.bootstrap_chunking import ReplayableWeightStream + + rng = np.random.default_rng() + rng.bit_generator.state = self.bitgen_state + return ReplayableWeightStream(_make_weight_iter_from_spec(self), rng) + + +def _make_weight_iter_from_spec(spec: BootstrapReplaySpec) -> Any: + """Module-level factory - never a local closure, so the spec stays picklable.""" + from diff_diff.bootstrap_chunking import iter_weight_blocks + + def _factory(rng: np.random.Generator) -> Any: + return iter_weight_blocks( + spec.n_bootstrap, + spec.n_units, + spec.weight_type, + rng, + expand_index=spec.expand_index, + block_size=spec.block_size, + ) + + return _factory + + +class AggregationMixin: + """``aggregate(type=...)`` for one results class (spec section 6). + + Applied PER RESULTS CLASS - never to + :class:`~diff_diff.results_base.BaseResults`, whose MRO position would + make every still-``planned`` ledger row's ``new`` locator resolve. + + A results class opts in by setting :attr:`_AGGREGATE_SUPPORTED` and + implementing ``_aggregate_compute``. + """ + + #: Aggregation types this results class implements. A subset of + #: :data:`AGGREGATION_VOCABULARY` plus any documented per-estimator extra. + _AGGREGATE_SUPPORTED: Tuple[str, ...] = () + + #: Types for which ``balance_e`` is meaningful. CallawaySantAnna threads it + #: only through event-study aggregation, so accepting it elsewhere would + #: silently ignore a user's argument. + _AGGREGATE_BALANCE_E_TYPES: Tuple[str, ...] = ("event_study",) + + def aggregate( + self, + type: str, # noqa: A002 - matches the ecosystem's aggte(type=) vocabulary + weights: Optional[str] = None, + *, + balance_e: Optional[int] = None, + ) -> Any: + """Re-aggregate this fit without refitting. + + Returns a NEW object; ``self`` is never modified. + ``type="event_study"`` returns + :class:`~diff_diff.results_base.EventStudyResults`, every other type an + :class:`AggregationResult`. + + Raises ``ValueError`` - never falls back silently - on an unsupported + type, on a ``weights`` value this estimator does not accept, and on + ``balance_e`` passed with a type that does not use it. + """ + # NB: the parameter `type` shadows the builtin for this method's body - + # the ecosystem's aggte(type=) vocabulary is worth the shadow, but every + # class-name lookup below must go through __class__, not type(self). + cls_name = self.__class__.__name__ + supported = tuple(self._AGGREGATE_SUPPORTED) + if not supported: + raise NotImplementedError( + f"{cls_name} does not implement aggregate(); it is added " + "per results class as each estimator's ledger row flips." + ) + if type not in supported: + known = ", ".join(repr(t) for t in supported) + extra = "" + if type in AGGREGATION_VOCABULARY: + extra = ( + f" {type!r} is part of the library-wide aggregation vocabulary " + "but this estimator does not implement it." + ) + raise ValueError(f"Unsupported aggregation type {type!r}. Supported: {known}.{extra}") + if balance_e is not None and type not in self._AGGREGATE_BALANCE_E_TYPES: + usable = ", ".join(repr(t) for t in self._AGGREGATE_BALANCE_E_TYPES) + raise ValueError( + f"balance_e is not used by aggregate(type={type!r}) and would be " + f"silently ignored. It applies to: {usable}." + ) + self._aggregate_validate_weights(weights) + return self._aggregate_compute(type, weights=weights, balance_e=balance_e) + + # -- hooks for the results class ----------------------------------- # + + def _aggregate_validate_weights(self, weights: Optional[str]) -> None: + """Reject a weighting scheme this estimator does not offer. + + Default: no selector, so anything but ``None`` fails closed. Estimators + with a real scheme (Wooldridge's ``"cell"`` / ``"cohort_share"``) + override. + """ + if weights is not None: + raise ValueError( + f"{type(self).__name__}.aggregate() does not accept a weights " + f"selector (got {weights!r}); its aggregation weights are " + "determined by the estimator." + ) + + def _aggregate_compute( + self, level: str, *, weights: Optional[str], balance_e: Optional[int] + ) -> Any: + raise NotImplementedError( + f"{type(self).__name__} declares _AGGREGATE_SUPPORTED but does not " + "implement _aggregate_compute()." + ) diff --git a/diff_diff/staggered.py b/diff_diff/staggered.py index 397b9012f..17b08d00f 100644 --- a/diff_diff/staggered.py +++ b/diff_diff/staggered.py @@ -1851,11 +1851,12 @@ def fit( if not (0 < self.pscore_trim < 0.5): raise ValueError(f"pscore_trim must be in (0, 0.5), got {self.pscore_trim}") - # Reset stale state from prior fit (prevents leaking event-study VCV - # and its df provenance across fits / non-ES aggregates / the ES - # empty-result early return) - self._event_study_vcov = None - self._event_study_df_used: Optional[float] = None + # NB: the event-study VCV and its df provenance used to be reset here, + # because ``_aggregate_event_study`` stashed them on ``self`` and a + # reused estimator could leak them across fits / non-ES aggregates / + # the ES empty-result early return. The aggregator now RETURNS them + # (``EventStudyAggregation``), so there is no cross-fit state to reset + # and the leak is impossible by construction. # Tracker for _safe_inv lstsq fallbacks across all analytical SE # paths (PS Hessian, OR bread, event-study bread, etc.). Emit ONE @@ -2647,8 +2648,11 @@ def fit( event_study_effects = None group_effects = None + # Aggregation outputs that used to arrive as ``self._event_study_*`` + # side channels; the aggregator is now pure and returns them. + es_aggregation = None if aggregate in ["event_study", "all"]: - event_study_effects = self._aggregate_event_study( + es_aggregation = self._aggregate_event_study( group_time_effects, influence_func_info, treatment_groups, @@ -2658,6 +2662,7 @@ def fit( unit, precomputed, ) + event_study_effects = es_aggregation.effects if aggregate in ["group", "all"]: group_effects = self._aggregate_by_group( @@ -2814,15 +2819,15 @@ def fit( # Retrieve event-study VCV from aggregation mixin (Phase 7d). # Clear it when bootstrap overwrites event-study SEs to prevent # HonestDiD from mixing analytical VCV with bootstrap SEs. - event_study_vcov = getattr(self, "_event_study_vcov", None) - event_study_vcov_index = getattr(self, "_event_study_vcov_index", None) + event_study_vcov = es_aggregation.vcov if es_aggregation is not None else None + event_study_vcov_index = es_aggregation.vcov_index if es_aggregation is not None else None if bootstrap_results is not None and event_study_vcov is not None: event_study_vcov = None event_study_vcov_index = None # Same clearing rule for the ES df provenance: bootstrap replaces the # stored ES se/p/CI with percentile values that never used the # analytical df, so surfacing it would be false provenance. - event_study_df = getattr(self, "_event_study_df_used", None) + event_study_df = es_aggregation.df_used if es_aggregation is not None else None if bootstrap_results is not None: event_study_df = None diff --git a/diff_diff/staggered_aggregation.py b/diff_diff/staggered_aggregation.py index 871174c5a..a46ed3fc4 100644 --- a/diff_diff/staggered_aggregation.py +++ b/diff_diff/staggered_aggregation.py @@ -5,6 +5,7 @@ group-time average treatment effects into summary measures. """ +from dataclasses import dataclass, field from typing import Any, Dict, List, Literal, Optional, Set, Tuple, Union, overload import numpy as np @@ -16,6 +17,41 @@ PrecomputedData = Dict[str, Any] +@dataclass +class EventStudyAggregation: + """Everything one event-study aggregation produces. + + ``_aggregate_event_study`` previously returned only ``effects`` and stashed + the other four values on ``self`` as side channels. Returning them makes + the aggregator PURE, which is what lets it run post-fit from a retained kit + without mutating the results object it was called from (spec section 6). + + Attributes + ---------- + effects : dict + Per-event-time effect records, keyed by event time. + overall : dict or None + Eq. (4.14) overall ATT (``att`` / ``se`` / ``effective_df``) - the + unweighted mean of post-treatment ES(e). Consumed by + ``StaggeredTripleDifference`` as ``overall_att_es``; CallawaySantAnna + leaves it unread. ``None`` when no post-treatment horizon qualifies. + df_used : float or None + The ONE df every ES row's inference actually used, recorded iff it + governs a t-reference (finite and > 0). Provenance for + ``CallawaySantAnnaResults.event_study_df``. + vcov : np.ndarray or None + Full event-study covariance, when computable. + vcov_index : list or None + Event times aligned 1:1 with ``vcov``'s columns. + """ + + effects: Dict[Any, Dict[str, Any]] = field(default_factory=dict) + overall: Optional[Dict[str, Any]] = None + df_used: Optional[float] = None + vcov: Optional[np.ndarray] = None + vcov_index: Optional[List[Any]] = None + + def fixed_cohort_agg_weights( precomputed: Optional["PrecomputedData"], ) -> Optional[Dict[Any, float]]: @@ -815,7 +851,7 @@ def _aggregate_event_study( df: Optional[pd.DataFrame] = None, unit: Optional[str] = None, precomputed: Optional["PrecomputedData"] = None, - ) -> Dict[int, Dict[str, Any]]: + ) -> EventStudyAggregation: """ Aggregate effects by relative time (event study). @@ -969,13 +1005,15 @@ def _aggregate_event_study( _psi_vectors.append(psi_e) _psi_event_times.append(e) - # Reset the Eq. (4.14) overall before any early return so a reused estimator - # instance never reads a stale value from a prior fit. - self._event_study_overall = None + # The Eq. (4.14) overall starts unset. It used to be reset on ``self`` + # before any early return so a reused estimator never read a stale value + # from a prior fit; returning it removes that hazard by construction. + es_overall: Optional[Dict[str, Any]] = None + es_df_used: Optional[float] = None # Batch inference for all relative periods if not agg_effects_list: - return {} + return EventStudyAggregation() # Use per-horizon effective df if any replicate aggregation overrode it; # otherwise fall back to the original df from the survey design. df_survey_val = precomputed.get("df_survey") if precomputed is not None else None @@ -997,10 +1035,10 @@ def _aggregate_event_study( # conservative min under dropped replicates) as provenance for # CallawaySantAnnaResults.event_study_df. Recorded iff it will # govern a t-reference (finite, > 0; the df=0 replicate sentinel - # yields NaN inference, not a t-law). fit() resets this stash - # alongside _event_study_vcov. + # yields NaN inference, not a t-law). Returned on the aggregation + # object alongside the VCV rather than stashed on ``self``. if df_survey_val is not None and np.isfinite(df_survey_val) and df_survey_val > 0: - self._event_study_df_used = float(df_survey_val) + es_df_used = float(df_survey_val) t_stats, p_values, ci_lowers, ci_uppers = safe_inference_batch( np.array(agg_effects_list), np.array(agg_ses_list), @@ -1072,10 +1110,7 @@ def _aggregate_event_study( # in HonestDiD when some event times are filtered out). Uses the # non-empty-psi event times so the index aligns 1:1 with the VCV columns # (reference-only and empty-IF horizons never get a column). - self._event_study_vcov_index = valid_event_times if event_study_vcov is not None else None - - # Attach VCV to self for CallawaySantAnna to pick up - self._event_study_vcov = event_study_vcov + event_study_vcov_index = valid_event_times if event_study_vcov is not None else None # Eq. (4.14) overall ATT: the unweighted mean of the post-treatment # event-study effects ES(e). Stashed on self (mirroring _event_study_vcov) @@ -1109,13 +1144,19 @@ def _aggregate_event_study( if len({len(p) for p in psis}) == 1: psi_es = np.column_stack(psis).mean(axis=1) se_es, eff_df_es = self._se_from_psi(psi_es, precomputed) - self._event_study_overall = { + es_overall = { "att": att_es, "se": float(se_es), "effective_df": eff_df_es, } - return event_study_effects + return EventStudyAggregation( + effects=event_study_effects, + overall=es_overall, + df_used=es_df_used, + vcov=event_study_vcov, + vcov_index=event_study_vcov_index, + ) def _aggregate_by_group( self, diff --git a/diff_diff/staggered_triple_diff.py b/diff_diff/staggered_triple_diff.py index a838ce693..df39a542b 100644 --- a/diff_diff/staggered_triple_diff.py +++ b/diff_diff/staggered_triple_diff.py @@ -619,8 +619,9 @@ def fit( # Aggregations event_study_effects = None group_effects = None + es_aggregation = None if aggregate in ("event_study", "all"): - event_study_effects = self._aggregate_event_study( + es_aggregation = self._aggregate_event_study( group_time_effects, influence_func_info, treatment_groups, @@ -630,6 +631,7 @@ def fit( unit, precomputed_agg, ) + event_study_effects = es_aggregation.effects if aggregate in ("group", "all"): group_effects = self._aggregate_by_group( group_time_effects, @@ -642,9 +644,10 @@ def fit( # Paper Eq. (4.14) overall ATT (event-study average): an opt-in summary # alongside the default CS-simple ``overall_att``. ``_aggregate_event_study`` - # stashes it on ``self._event_study_overall``; populated only when the - # event-study aggregation ran. Analytical inference here; the bootstrap block - # below overrides the SE when ``n_bootstrap > 0`` (mirroring ``overall_se``). + # RETURNS it on its aggregation object (it used to stash it on + # ``self._event_study_overall``); populated only when the event-study + # aggregation ran. Analytical inference here; the bootstrap block below + # overrides the SE when ``n_bootstrap > 0`` (mirroring ``overall_se``). overall_att_es = None overall_se_es = None overall_t_stat_es = None @@ -657,7 +660,7 @@ def fit( # analytical-IF failure (the bootstrap path emits its own warning). analytical_overall_es_se_nonfinite = False if aggregate in ("event_study", "all"): - es_overall = getattr(self, "_event_study_overall", None) + es_overall = es_aggregation.overall if es_aggregation is not None else None if es_overall is not None: overall_att_es = es_overall["att"] overall_se_es = es_overall["se"] diff --git a/docs/doc-deps.yaml b/docs/doc-deps.yaml index 823fb87eb..448745dbf 100644 --- a/docs/doc-deps.yaml +++ b/docs/doc-deps.yaml @@ -1049,6 +1049,16 @@ sources: type: methodology note: "SyntheticDiDResults hosts validation diagnostics (LOO, weight concentration, in-time placebo, zeta sensitivity)" + diff_diff/aggregation.py: + drift_risk: low + docs: + - path: docs/api/results.rst + type: api_reference + - path: docs/v4-design.md + section: "6. Aggregation contract" + type: design_spec + note: "AggregationResult / AggregationKit / AggregationMixin - the post-fit results.aggregate(type=) surface (4.0 program Phase 2; ledger rows M-020..M-027, M-122)" + diff_diff/results_base.py: drift_risk: low docs: From 613751f8d93b7044359c5f7fcc0173f94c3c2047 Mon Sep 17 00:00:00 2001 From: igerber Date: Sat, 25 Jul 2026 18:36:53 -0400 Subject: [PATCH 02/11] feat(v4): results.aggregate() on CallawaySantAnna + fit(aggregate=) shim The post-fit aggregation surface itself (spec section 6, rows M-020 / M-117). res = CallawaySantAnna().fit(df, ...) # no aggregate= argument res.aggregate("event_study") # -> EventStudyResults res.aggregate("group") # -> AggregationResult res.aggregate("simple") # -> AggregationResult This gives EventStudyResults its first public producer: the container shipped in 3.9 (row M-092) has been exported with only package-internal builders. Retention. fit() now attaches an AggregationKit, built there because neither `precomputed` nor `influence_func_info` survives the call - both are locals, so the kit cannot be reconstructed afterwards. It copies only the twelve bookkeeping keys the aggregation machinery actually reads, verified by enumerating every access in staggered_aggregation.py; the data matrices (outcome_matrix, covariate_matrix, obs_outcome, obs_covariates) are excluded, so a results object never holds the source panel. Asserted: no retained attribute is a DataFrame. Reaching the aggregators without an estimator reference. They are mixin methods, but read exactly two attributes from their host - alpha and anticipation - so _KitAggregator supplies those and nothing else. That avoids an _estimator_ref, which would drag the whole fitted estimator, and its frame, onto every result. Frame-free aggregation. Three `df` dereferences in the influence-function path now prefer the precomputed bookkeeping, which the adjacent branches already did: `unit_cohorts` is the per-unit cohort array, so counting its matches is identical to the frame lookup it replaces. The unconditional `assert df is not None` becomes a fail-closed check that accepts EITHER the frame or the bookkeeping. The frame path remains for direct internal callers. Immutability. aggregate() returns a new object and never touches its parent. The event-study builder reads its surface off a results object, so a throwaway carrier (dataclasses.replace) holds the freshly computed values rather than populating self. The `_agg_cache` memo write lands on a shallow copy of the kit's bookkeeping - sharing every array, duplicating no data - mirroring what StaggeredTripleDifference already does when it aggregates through a modified copy. Deprecation shim. fit(aggregate=) and fit(balance_e=) warn (FutureWarning, removed at 4.0) via a sentinel default, so the warning fires only when the caller actually supplies the argument and never on a plain fit(). Routing is untouched: the deprecated path still returns the fully populated legacy surface that honest_did, pretrends and build_event_study_surface read. Fail-closed rather than silently wrong: aggregate("calendar") and aggregate("all") raise naming the supported set; a non-None `weights` raises (CS exposes no weighting selector); balance_e with "simple"/"group" raises rather than being ignored, since the shipped code threads it only through event-study aggregation; and aggregate() on a BOOTSTRAPPED fit raises, because its percentile inference cannot be reproduced from analytical state. Bootstrap replay is deliberately not wired here - BootstrapReplaySpec is the verified mechanism for the follow-up. Verified: post-fit aggregate() reproduces fit-time numbers IDENTICALLY at atol=rtol=1e-14 for all three types; the parent is unchanged after five mixed-order calls and repeated calls agree; the 10-scenario pre-refactor golden still matches; 499 tests pass across the staggered, SDDD, event-study-surface and R-parity suites. --- diff_diff/staggered.py | 117 ++++++++++++++++++++- diff_diff/staggered_aggregation.py | 61 ++++++++--- diff_diff/staggered_results.py | 159 ++++++++++++++++++++++++++++- 3 files changed, 316 insertions(+), 21 deletions(-) diff --git a/diff_diff/staggered.py b/diff_diff/staggered.py index 17b08d00f..89839d751 100644 --- a/diff_diff/staggered.py +++ b/diff_diff/staggered.py @@ -12,6 +12,9 @@ import numpy as np import pandas as pd +from diff_diff.aggregation import ( + AggregationKit, +) from diff_diff.linalg import ( _check_propensity_diagnostics, _detect_rank_deficiency, @@ -51,6 +54,22 @@ PrecomputedData = Dict[str, Any] +class _DeprecatedFitArg: + """Sentinel for `fit(aggregate=)` / `fit(balance_e=)` (rows M-020 / M-117). + + A plain ``None`` default cannot distinguish "not passed" from "passed + None", so a bare ``None`` default would fire the FutureWarning on EVERY + fit. The warning must fire only when the caller actually supplies the + argument. + """ + + def __repr__(self) -> str: # pragma: no cover - debugging aid + return "" + + +_DEPRECATED_FIT_ARG = _DeprecatedFitArg() + + def _linear_regression( X: np.ndarray, y: np.ndarray, @@ -1795,8 +1814,8 @@ def fit( time: str, first_treat: str, covariates: Optional[List[str]] = None, - aggregate: Optional[str] = None, - balance_e: Optional[int] = None, + aggregate: Any = _DEPRECATED_FIT_ARG, + balance_e: Any = _DEPRECATED_FIT_ARG, survey_design: Optional["SurveyDesign"] = None, ) -> CallawaySantAnnaResults: """ @@ -1847,6 +1866,34 @@ def fit( ValueError If required columns are missing or data validation fails. """ + # ---- fit(aggregate=) / fit(balance_e=) shims (rows M-020 / M-117) ---- + # Deprecated in 3.9, removed at 4.0: aggregation moves to the post-fit + # results.aggregate(type=...). The sentinel default means the warning + # fires ONLY when the caller supplies the argument, never on a plain + # fit(). Routing is otherwise untouched, so the deprecated path returns + # exactly the numbers it always did. + _deprecated_passed = [ + name + for name, value in (("aggregate", aggregate), ("balance_e", balance_e)) + if not isinstance(value, _DeprecatedFitArg) + ] + if _deprecated_passed: + _args = " / ".join(f"{n}=" for n in _deprecated_passed) + warnings.warn( + f"CallawaySantAnna.fit({_args}) is deprecated and will be removed " + "in 4.0. Fit once, then aggregate as a post-fit step: " + "results = CallawaySantAnna().fit(...); " + "results.aggregate('event_study') / .aggregate('group') / " + ".aggregate('simple'). balance_e moves onto aggregate() alongside " + "it: results.aggregate('event_study', balance_e=2).", + FutureWarning, + stacklevel=2, + ) + if isinstance(aggregate, _DeprecatedFitArg): + aggregate = None + if isinstance(balance_e, _DeprecatedFitArg): + balance_e = None + # Validate pscore_trim (may have been changed via set_params) if not (0 < self.pscore_trim < 0.5): raise ValueError(f"pscore_trim must be in (0, 0.5), got {self.pscore_trim}") @@ -2909,6 +2956,15 @@ def fit( df_inference=df_inference_for_results, ) + # Attach the post-fit aggregation kit (spec section 6, rows M-020/M-117). + # Built HERE because neither `precomputed` nor `influence_func_info` + # survives fit() - they are locals - so the kit cannot be reconstructed + # later. It deliberately excludes the data matrices: re-aggregation reads + # only unit-level bookkeeping, so the source panel is never retained. + self.results_._aggregation_kit = _build_aggregation_kit( + self, precomputed, influence_func_info, group_time_effects + ) + self.is_fitted_ = True return self.results_ @@ -5005,3 +5061,60 @@ def summary(self) -> str: def print_summary(self) -> None: """Print summary to stdout.""" print(self.summary()) + + +#: `precomputed` keys the aggregation machinery actually reads. Verified by +#: enumerating every `precomputed[...]` / `.get(...)` access in +#: `staggered_aggregation.py`. Everything else - notably `outcome_matrix`, +#: `covariate_matrix`, `obs_outcome`, `obs_covariates` - is DATA and is +#: deliberately not retained, so a results object never holds the source panel. +#: `_agg_cache` is excluded too: it is derived memoization, rebuilt on demand. +_AGGREGATION_BOOKKEEPING_KEYS = ( + "unit_to_idx", + "unit_cohorts", + "all_units", + "obs_per_unit", + "survey_weights", + "resolved_survey_unit", + "df_survey", + "agg_cohort_masses", + "agg_total_weight", + "is_panel", + "canonical_size", + "n_units", +) + + +def _build_aggregation_kit( + estimator: "CallawaySantAnna", + precomputed: Optional[Dict[str, Any]], + influence_func_info: Optional[Dict[str, Any]], + group_time_effects: Optional[Dict[Any, Any]], +) -> Optional["AggregationKit"]: + """Distil the fit-time state post-fit re-aggregation needs. + + Returns ``None`` when the fit produced nothing to re-aggregate, in which + case ``aggregate()`` reports that rather than failing obscurely. + """ + if not influence_func_info or not group_time_effects: + return None + + bookkeeping: Dict[str, Any] = {} + if precomputed is not None: + for key in _AGGREGATION_BOOKKEEPING_KEYS: + if key in precomputed: + bookkeeping[key] = precomputed[key] + + return AggregationKit( + bookkeeping=bookkeeping, + influence=influence_func_info, + alpha=estimator.alpha, + anticipation=estimator.anticipation, + cband=bool(estimator.cband), + # Bootstrap replay is not wired in this PR: a bootstrapped fit's + # percentile inference cannot be reproduced from analytical state, so + # aggregate() fails closed on one rather than silently substituting + # analytical numbers. BootstrapReplaySpec (diff_diff/aggregation.py) is + # the verified mechanism for the follow-up. + bootstrap=None, + ) diff --git a/diff_diff/staggered_aggregation.py b/diff_diff/staggered_aggregation.py index a46ed3fc4..6dbbea9d6 100644 --- a/diff_diff/staggered_aggregation.py +++ b/diff_diff/staggered_aggregation.py @@ -103,8 +103,8 @@ def _aggregate_simple( self, group_time_effects: Dict, influence_func_info: Dict, - df: pd.DataFrame, - unit: str, + df: Optional[pd.DataFrame], + unit: Optional[str], precomputed: Optional["PrecomputedData"] = None, ) -> Tuple[float, float, Optional[int]]: """ @@ -407,8 +407,8 @@ def _compute_combined_influence_function( effects: np.ndarray, groups_for_gt: np.ndarray, influence_func_info: Dict, - df: pd.DataFrame, - unit: str, + df: Optional[pd.DataFrame], + unit: Optional[str], precomputed: Optional["PrecomputedData"] = None, global_unit_to_idx: Optional[Dict[Any, int]] = None, n_global_units: Optional[int] = None, @@ -538,7 +538,18 @@ def _compute_combined_influence_function( for g in unique_groups: group_sizes[g] = int(np.sum(precomputed_cohorts == g)) total_weight = float(n_units) + elif precomputed is not None: + # Panel without survey. ``unit_cohorts`` is the per-unit cohort array + # (``df.groupby(unit)[first_treat].first().values``), so counting its + # matches is IDENTICAL to the frame lookup below - and it lets + # post-fit aggregation run from the retained kit with no frame. + precomputed_cohorts = precomputed["unit_cohorts"] + for g in unique_groups: + group_sizes[g] = int(np.sum(precomputed_cohorts == g)) + total_weight = float(n_units) else: + # No precomputed bookkeeping (direct internal callers only): fall + # back to the fit-time frame. for g in unique_groups: treated_in_g = df[df["first_treat"] == g][unit].nunique() group_sizes[g] = treated_in_g @@ -606,10 +617,21 @@ def _compute_combined_influence_function( unit_groups_array[idx] = unit_first_treat else: idx_uid_pairs = list(enumerate(all_units)) - for idx, uid in idx_uid_pairs: - unit_first_treat = df[df[unit] == uid]["first_treat"].iloc[0] - if unit_first_treat in unique_groups_set: - unit_groups_array[idx] = unit_first_treat + if precomputed is not None and precomputed.get("unit_to_idx") is not None: + # Same per-unit cohort lookup as the branch above, from the kit + # rather than the frame - keeps post-fit aggregation frame-free. + precomputed_cohorts = precomputed["unit_cohorts"] + precomputed_unit_to_idx = precomputed["unit_to_idx"] + for idx, uid in idx_uid_pairs: + if uid in precomputed_unit_to_idx: + cohort = precomputed_cohorts[precomputed_unit_to_idx[uid]] + if cohort in unique_groups_set: + unit_groups_array[idx] = cohort + else: + for idx, uid in idx_uid_pairs: + unit_first_treat = df[df[unit] == uid]["first_treat"].iloc[0] + if unit_first_treat in unique_groups_set: + unit_groups_array[idx] = unit_first_treat # Vectorized WIF computation groups_for_gt_array = np.array(groups_for_gt) @@ -688,8 +710,8 @@ def _compute_aggregated_se_with_wif( effects: np.ndarray, groups_for_gt: np.ndarray, influence_func_info: Dict, - df: pd.DataFrame, - unit: str, + df: Optional[pd.DataFrame], + unit: Optional[str], precomputed: Optional["PrecomputedData"] = None, return_psi: Literal[False] = False, ) -> Tuple[float, Optional[int]]: ... @@ -702,8 +724,8 @@ def _compute_aggregated_se_with_wif( effects: np.ndarray, groups_for_gt: np.ndarray, influence_func_info: Dict, - df: pd.DataFrame, - unit: str, + df: Optional[pd.DataFrame], + unit: Optional[str], precomputed: Optional["PrecomputedData"] = None, *, return_psi: Literal[True], @@ -716,8 +738,8 @@ def _compute_aggregated_se_with_wif( effects: np.ndarray, groups_for_gt: np.ndarray, influence_func_info: Dict, - df: pd.DataFrame, - unit: str, + df: Optional[pd.DataFrame], + unit: Optional[str], precomputed: Optional["PrecomputedData"] = None, return_psi: bool = False, ) -> "Union[Tuple[float, Optional[int]], Tuple[float, np.ndarray, Optional[int]]]": @@ -980,8 +1002,15 @@ def _aggregate_event_study( # reference cells contribute nothing to the variance but their cohort # weight dilutes the real cells, matching R's dynamic aggregation. groups_for_gt = np.array([g for (g, t) in gt_pairs]) - # The wif-SE path requires the fit-time frame (callers pass it). - assert df is not None and unit is not None + # The wif-SE path needs EITHER the fit-time frame or the precomputed + # bookkeeping. Post-fit re-aggregation supplies only the latter (the + # frame is deliberately not retained), and every frame dereference + # inside now prefers `precomputed`. + if precomputed is None and (df is None or unit is None): + raise ValueError( + "Event-study aggregation needs either the fit-time frame " + "(df + unit) or precomputed bookkeeping; got neither." + ) agg_se, psi_e, eff_df = self._compute_aggregated_se_with_wif( gt_pairs, weights, diff --git a/diff_diff/staggered_results.py b/diff_diff/staggered_results.py index 98357d7b9..d219173b1 100644 --- a/diff_diff/staggered_results.py +++ b/diff_diff/staggered_results.py @@ -5,19 +5,40 @@ group-time average treatment effects and their aggregations. """ -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple import numpy as np import pandas as pd +from diff_diff.aggregation import AggregationMixin, AggregationResult from diff_diff.results import _format_survey_block, _get_significance_stars -from diff_diff.results_base import BaseResults +from diff_diff.results_base import BaseResults, build_event_study_surface +from diff_diff.staggered_aggregation import CallawaySantAnnaAggregationMixin if TYPE_CHECKING: from diff_diff.staggered_bootstrap import CSBootstrapResults +class _KitAggregator(CallawaySantAnnaAggregationMixin): + """Runs the CallawaySantAnna aggregation mixin off a retained kit. + + The aggregation methods are ESTIMATOR-bound but read exactly two attributes + from their host - ``alpha`` and ``anticipation`` - so a post-fit caller + needs neither the estimator nor a reference to it, only those two values. + (Verified by enumerating every ``self.`` access in + ``staggered_aggregation.py``: the rest are internal method calls.) + + This is what keeps ``aggregate()`` off an ``_estimator_ref``, which would + otherwise drag the whole fitted estimator - and its source frame - onto + every results object. + """ + + def __init__(self, alpha: float, anticipation: int) -> None: + self.alpha = alpha + self.anticipation = anticipation + + @dataclass class GroupTimeEffect: """ @@ -67,7 +88,7 @@ def significance_stars(self) -> str: @dataclass -class CallawaySantAnnaResults(BaseResults): +class CallawaySantAnnaResults(BaseResults, AggregationMixin): """ Results from Callaway-Sant'Anna (2021) staggered DiD estimation. @@ -253,6 +274,138 @@ def coef_var(self) -> float: return np.nan return self.overall_se / abs(self.overall_att) + # ------------------------------------------------------------------ # + # Post-fit aggregation (spec section 6; ledger rows M-020 / M-117) + # ------------------------------------------------------------------ # + + #: CS implements simple / event-study / group. ``"calendar"`` is part of + #: the library-wide vocabulary but CS has no calendar aggregator (the + #: DEFERRED "Calendar-time aggregation" row), so asking for it raises. + _AGGREGATE_SUPPORTED = ("simple", "event_study", "group") + + def _aggregate_compute( + self, level: str, *, weights: Optional[str], balance_e: Optional[int] + ) -> Any: + kit = getattr(self, "_aggregation_kit", None) + if kit is None: + raise ValueError( + "This result carries no aggregation kit, so it cannot be " + "re-aggregated. Kits are attached by CallawaySantAnna.fit(); a " + "result unpickled from an older release will not have one." + ) + if self.bootstrap_results is not None: + # Fail closed rather than silently handing back analytical numbers: + # a bootstrapped fit's se/p/CI are percentile statistics, and + # reproducing them post-fit needs retained draws (BootstrapReplaySpec). + raise NotImplementedError( + "aggregate() is not yet available on a bootstrapped fit " + "(n_bootstrap > 0): its inference is percentile-bootstrap based " + "and cannot be reproduced from the analytical state retained " + "here. Re-fit with the aggregation you need, or use " + "n_bootstrap=0 for analytical inference." + ) + + # Shallow copy: shares every array (no data is duplicated) but gives the + # aggregators a scratch dict to memoize `_agg_cache` into, so the kit + # itself is never written to. Mirrors what StaggeredTripleDifference + # already does when it aggregates through a modified copy. + precomputed = dict(kit.bookkeeping) + agg = _KitAggregator(kit.alpha, kit.anticipation) + + if level == "simple": + return self._aggregate_simple_result() + if level == "group": + effects = agg._aggregate_by_group( + self.group_time_effects, + kit.influence, + self.groups, + precomputed=precomputed, + df=None, + unit=None, + ) + return self._group_effects_to_aggregation(effects, kit) + + # event_study -> the unified EventStudyResults container (row M-092) + es = agg._aggregate_event_study( + self.group_time_effects, + kit.influence, + self.groups, + self.time_periods, + balance_e, + None, + None, + precomputed, + ) + # `build_event_study_surface` reads the surface off a RESULTS object. + # Under the immutability contract we cannot populate `self`, so a + # throwaway carrier holds the freshly computed values. + carrier = replace( + self, + event_study_effects=es.effects, + event_study_vcov=es.vcov, + event_study_vcov_index=es.vcov_index, + event_study_df=es.df_used, + ) + return build_event_study_surface(carrier) + + def _aggregate_simple_result(self) -> AggregationResult: + """The overall ATT as a one-row table. + + ``_aggregate_simple`` runs unconditionally in ``fit()``, so the numbers + are already stored - this is a view, not a recomputation, and is + therefore bit-identical to the fit by construction. + """ + ci = self.overall_conf_int or (np.nan, np.nan) + return AggregationResult( + level="simple", + label=np.array(["overall"], dtype=object), + target=np.array(["att"], dtype=object), + att=np.array([self.overall_att], dtype=float), + se=np.array([self.overall_se], dtype=float), + t_stat=np.array([self.overall_t_stat], dtype=float), + p_value=np.array([self.overall_p_value], dtype=float), + conf_int_lower=np.array([ci[0]], dtype=float), + conf_int_upper=np.array([ci[1]], dtype=float), + n=np.array([float(self.n_treated_units + self.n_control_units)], dtype=float), + df=self.df_inference, + alpha=self.alpha, + n_kind="units", + weight=np.array([1.0], dtype=float), + estimator=type(self).__name__.replace("Results", ""), + ) + + def _group_effects_to_aggregation( + self, effects: Dict[Any, Dict[str, Any]], kit: Any + ) -> AggregationResult: + """Per-cohort aggregation as a table. + + ``weight`` is left unset: ``_aggregate_by_group`` weights ``(g, t)`` + cells equally WITHIN each cohort and never forms a cross-cohort mass, + so there is no per-row weight to report and inventing one would be a + fabricated number. ``n`` is the cohort's finite-contributing cell + count, hence ``n_kind="cells"`` rather than units. + """ + labels = list(effects.keys()) + rows = [effects[g] for g in labels] + cis = [r.get("conf_int") or (np.nan, np.nan) for r in rows] + return AggregationResult( + level="group", + label=np.array(labels, dtype=object), + target=np.array(["att"] * len(labels), dtype=object), + att=np.array([r["effect"] for r in rows], dtype=float), + se=np.array([r["se"] for r in rows], dtype=float), + t_stat=np.array([r["t_stat"] for r in rows], dtype=float), + p_value=np.array([r["p_value"] for r in rows], dtype=float), + conf_int_lower=np.array([c[0] for c in cis], dtype=float), + conf_int_upper=np.array([c[1] for c in cis], dtype=float), + n=np.array([float(r.get("n_periods", np.nan)) for r in rows], dtype=float), + df=self.event_study_df if self.event_study_df is not None else np.nan, + alpha=kit.alpha, + n_kind="cells", + weight=None, + estimator=type(self).__name__.replace("Results", ""), + ) + def summary(self, alpha: Optional[float] = None) -> str: """ Generate formatted summary of estimation results. From 89879e45ae4efe4efea45fc75651b1e198ff24c9 Mon Sep 17 00:00:00 2001 From: igerber Date: Sat, 25 Jul 2026 18:41:50 -0400 Subject: [PATCH 03/11] feat(v4): ledger rows M-117/M-122, contract tests, CS guidance migration Completes Phase 2b PR 1. Ledger (99 -> 101 rows; floor, id ranges and the snapshot assertion move together): - M-020 flips planned -> shimmed and gains a test_ref (it carried none). Its notes now record the CS SUPPORTED SUBSET: the closed vocabulary is library-wide, but CallawaySantAnna implements simple|event_study|group and has no calendar aggregator, so aggregate('calendar') raises naming what is supported. The DEFERRED "Calendar-time aggregation" row stands. - M-117 (new): balance_e moves from fit() onto aggregate(). It existed only as the prose "balance_e moves to aggregate() in the same PR" inside M-020's notes, which nothing could assert - the same un-rowed-obligation class the gating-completeness amendment closed. The other three balance_e sites get rows in the PRs that migrate them. - M-122 (new): AggregationResult, introduced_in 3.9 with deprecated_in null so the early-flip guard does not fire against the PR that ships it. - Ids M-116 and M-118..M-121 are left UNUSED and the gap is documented: they were drafted for later 2b PRs and are reserved, not deleted, since ids are never reused. tests/test_aggregate_contract.py (33 tests) is the test_ref for all three rows: numerical inertness against fit-time aggregation for every type plus a balance_e sweep; immutability across mixed-order and repeated calls; the kit holding no DataFrame and surviving a pickle round-trip; the whole fail-closed surface (calendar / all / unknown type / weights / balance_e-where-inert / bootstrapped fit); the pinned schema, nullable weight and declared n_kind; the shim warning only when the argument is passed; and a StaggeredTripleDifference regression for the Eq. 4.14 overall that the purity refactor moved off `self`. Spec correction. Section 6 claimed "CallawaySantAnna already stores them" of influence functions. It does not - the field was declared and never assigned - so the text is corrected in place: retention is work in EVERY migrating PR, the kit must be built during fit() because nothing it needs survives the call, and the dominant cost is the per-(g,t) influence dict at ~O(n_units x n_gt), not the O(n_units) bookkeeping. Guidance migration (section 8 rule 11 - a rename's scope includes every reader), CS-scoped: the _ABSENT_SURFACE_HINTS entry, honest_did's and pretrends' CS-gated error messages, and README's step-7 line all told users to re-run fit with aggregate=; they now point at the post-fit call and say no refit is required. AggregationResult is added to docs/api/results.rst and the canonical autosummary in docs/api/index.rst (the CI-enforced IA invariant). 1009 tests pass across the contract, ledger, docs-IA, doc-deps, staggered, SDDD, honest-did, pretrends and event-study-surface suites. --- CHANGELOG.md | 37 +++ README.md | 2 +- diff_diff/honest_did.py | 4 +- diff_diff/pretrends.py | 3 +- diff_diff/results_base.py | 3 +- docs/api/index.rst | 1 + docs/api/results.rst | 6 + docs/v4-deprecations.yaml | 41 +++- docs/v4-design.md | 16 +- tests/test_aggregate_contract.py | 386 +++++++++++++++++++++++++++++++ tests/test_v4_matrix.py | 14 +- 11 files changed, 496 insertions(+), 17 deletions(-) create mode 100644 tests/test_aggregate_contract.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 38516d223..fa2cb1c5d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,43 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added +- **Post-fit aggregation: `results.aggregate(type=...)` on CallawaySantAnna + (4.0 program Phase 2b, ledger rows M-020 / M-117 / M-122).** Estimate once, + aggregate as a post-fit step — the ecosystem norm (`did::aggte`, + `etwfe::emfx`, Stata `estat aggregation`). `results.aggregate("event_study")` + returns the unified `EventStudyResults` container, finally giving [M-092]'s + export a public producer; `"group"` and `"simple"` return the new + `AggregationResult`, a tabular container with a pinned 12-column schema and + `summary()` / `to_dict()` / `to_dataframe()`. **`fit(aggregate=)` and + `fit(balance_e=)` now emit a `FutureWarning`** (removed at 4.0) via a + sentinel default, so the warning fires only when the argument is actually + passed and never on a plain `fit()`; the deprecated path still returns the + fully populated legacy surface. `balance_e` moves onto `aggregate()`, where + it applies to event-study aggregation only — passing it with `"simple"` or + `"group"` raises rather than being silently ignored, matching where the + shipped code actually threads it. + **No numerical change:** a 10-scenario golden (every aggregation type, a + `balance_e` sweep, repeated cross-sections, a bootstrapped fit, + `anticipation=1`) is reproduced at `atol=rtol=1e-14`, and post-fit + `aggregate()` matches the fit-time numbers identically. + Three schema decisions are grounded in what the estimator produces rather + than assumed: `target` is per-row so one container can carry two aligned + estimands; `n_kind` reuses `EventStudyResults`' vocabulary (`"cells"` for + CS's group level, not "units"); and `weight` is `None` for `"group"`, + because that aggregation weights `(g,t)` cells equally *within* each cohort + and forms no cross-cohort mass — reporting one would be fabricated. + Fail-closed elsewhere too: `aggregate("calendar")` raises (CS has no + calendar aggregator; the DEFERRED row stands), a non-`None` `weights` + raises, and `aggregate()` on a **bootstrapped** fit raises rather than + substituting analytical inference for percentile-bootstrap statistics — + bootstrap replay is a tracked follow-up. + Internally, the CS event-study aggregator became **pure**: it returned one + dict and stashed four more values on `self`, and now returns all five, which + is what allows re-aggregation without mutating the result. + `StaggeredTripleDifference` inherits the same mixin and is the only reader of + the Eq. 4.14 overall, so it moved in the same change. Aggregation reads the + retained bookkeeping instead of the fit-time frame, so **no results object + retains the source panel**. - **Internal: 4.0 gating-completeness amendment to the design spec and ledger.** An audit of the public surface against the spec's own contract-rename rules found two renames the initial inventory missed and three Phase 2 obligations diff --git a/README.md b/README.md index be7532e32..cb69a528e 100644 --- a/README.md +++ b/README.md @@ -91,7 +91,7 @@ For rigorous DiD analysis, follow these 8 steps. Skipping diagnostic steps produ 4. **Choose estimator** - staggered adoption -> CS/SA/BJS (NOT plain TWFE); few treated units -> SDiD; factor confounding -> TROP; simple 2x2 -> DiD. Run `BaconDecomposition` to diagnose TWFE bias. 5. **Estimate** - `estimator.fit(data, ...)`. Always print the cluster count first and choose inference method based on the result (cluster-robust if >= 50 clusters, wild bootstrap if fewer). 6. **Sensitivity analysis** - `compute_honest_did(results)` for bounds under PT violations (MultiPeriodDiD, CS, or dCDH), `run_all_placebo_tests()` for 2x2 falsification, specification comparisons for staggered designs. -7. **Heterogeneity** - CS: `aggregate='group'`/`'event_study'`; SA: `results.event_study_effects` / `to_dataframe(level='cohort')`; subgroup re-estimation. +7. **Heterogeneity** - CS: `results.aggregate('group')`/`'event_study'` (post-fit, no refit); SA: `results.event_study_effects` / `to_dataframe(level='cohort')`; subgroup re-estimation. 8. **Robustness** - compare 2-3 estimators (CS vs SA vs BJS), report with and without covariates (shows whether conditioning drives identification), present pre-trends and sensitivity bounds. Full guide: `diff_diff.get_llm_guide("practitioner")`. diff --git a/diff_diff/honest_did.py b/diff_diff/honest_did.py index 3e2edc4a3..9ea413844 100644 --- a/diff_diff/honest_did.py +++ b/diff_diff/honest_did.py @@ -692,8 +692,8 @@ def _extract_event_study_params( if results.event_study_effects is None: raise ValueError( "CallawaySantAnnaResults must have event_study_effects for HonestDiD. " - "Re-run CallawaySantAnna.fit() with aggregate='event_study' to compute " - "event study effects." + "Call results.aggregate('event_study') to compute them post-fit " + "(no refit required)." ) # Warn if not using universal base period (R's HonestDiD requires it) diff --git a/diff_diff/pretrends.py b/diff_diff/pretrends.py index 11a3d44f9..8249f6291 100644 --- a/diff_diff/pretrends.py +++ b/diff_diff/pretrends.py @@ -1140,7 +1140,8 @@ def _extract_pre_period_params( if results.event_study_effects is None: raise ValueError( "CallawaySantAnnaResults must have event_study_effects. " - "Re-run with aggregate='event_study'." + "Call results.aggregate('event_study') to compute them " + "post-fit (no refit required)." ) # Get pre-period effects. Anticipation-aware cutoff per diff --git a/diff_diff/results_base.py b/diff_diff/results_base.py index 2c5a308cd..db3d91fbc 100644 --- a/diff_diff/results_base.py +++ b/diff_diff/results_base.py @@ -515,7 +515,8 @@ def summary(self, alpha: Optional[float] = None) -> str: #: Per-producer remediation for an absent event-study surface. _ABSENT_SURFACE_HINTS: Dict[str, str] = { - "CallawaySantAnnaResults": "refit with aggregate='event_study' (or 'all')", + # Migrated to the post-fit surface (row M-020): no refit needed. + "CallawaySantAnnaResults": "call results.aggregate('event_study')", "ImputationDiDResults": "refit with aggregate='event_study' (or 'all')", "TwoStageDiDResults": "refit with aggregate='event_study' (or 'all')", "StackedDiDResults": "refit with aggregate='event_study'", diff --git a/docs/api/index.rst b/docs/api/index.rst index fd6d32ae1..13ddf823e 100644 --- a/docs/api/index.rst +++ b/docs/api/index.rst @@ -85,6 +85,7 @@ Result containers returned by estimators: diff_diff.BaseResults diff_diff.Diagnostic diff_diff.EventStudyResults + diff_diff.AggregationResult Visualization ------------- diff --git a/docs/api/results.rst b/docs/api/results.rst index 36c4c0466..593372a8f 100644 --- a/docs/api/results.rst +++ b/docs/api/results.rst @@ -115,3 +115,9 @@ representation. :members: :undoc-members: :show-inheritance: + +.. autoclass:: diff_diff.AggregationResult + :no-index: + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/v4-deprecations.yaml b/docs/v4-deprecations.yaml index 8acdeed74..6600b8660 100644 --- a/docs/v4-deprecations.yaml +++ b/docs/v4-deprecations.yaml @@ -220,11 +220,12 @@ rows: introduced_in: "3.9" deprecated_in: "3.9" removed_in: "4.0" - status: planned - phase: 2 + status: shimmed + phase: 5 warning: FutureWarning - code_refs: [diff_diff/staggered.py] - notes: "balance_e moves to aggregate() in the same PR. Vocabulary: simple|event_study|group|calendar." + test_ref: tests/test_aggregate_contract.py + code_refs: [diff_diff/staggered.py, diff_diff/staggered_results.py, diff_diff/aggregation.py] + notes: "Shimmed in 3.9: fit(aggregate=) warns via a sentinel default (so a plain fit() never warns) and still returns the fully populated legacy surface; results.aggregate(type=) is the successor. balance_e moves alongside it as its own row [M-117] - it was previously tracked only as prose here, which nothing asserted. VOCABULARY: the closed set is library-wide (simple|event_study|group|calendar); CallawaySantAnna's SUPPORTED SUBSET is simple|event_study|group - it has no calendar aggregator (the DEFERRED 'Calendar-time aggregation' row), and aggregate('calendar') raises naming what is supported." - id: M-021 kind: param group: aggregate-postfit @@ -1307,3 +1308,35 @@ rows: warning: FutureWarning code_refs: [diff_diff/linalg.py] notes: "Fourth site of the 'robust' drop - spec section 7 says the flag dies 'everywhere it exists' but [M-045]..[M-047] enumerated only three. LinearRegression is a top-level export with its OWN __init__ (the TWFE/MultiPeriodDiD hits inherit DifferenceInDifferences.__init__ and ride [M-045]). Redundant with vcov_type, same as its siblings." + + # ---- Phase 2b PR 1: the post-fit aggregation surface --------------------- + # Ids M-116 and M-118..M-121 are intentionally UNUSED: they were drafted for + # the later 2b PRs (HAD rename, the other three balance_e sites, Wooldridge) + # and are reserved rather than reassigned, since ids are never reused. + - id: M-117 + kind: param + group: aggregate-postfit + old: "diff_diff:CallawaySantAnna.fit[balance_e]" + new: "diff_diff:CallawaySantAnnaResults.aggregate[balance_e]" + introduced_in: "3.9" + deprecated_in: "3.9" + removed_in: "4.0" + status: shimmed + phase: 5 + warning: FutureWarning + test_ref: tests/test_aggregate_contract.py + code_refs: [diff_diff/staggered.py, diff_diff/staggered_results.py, diff_diff/aggregation.py] + notes: "balance_e moves from fit() onto aggregate() with [M-020]. Previously tracked only as the prose 'balance_e moves to aggregate() in the same PR' inside M-020's notes, which no test could assert - the same un-rowed-obligation class the gating-completeness amendment closed. Applies to event-study aggregation ONLY (the shipped code threads it nowhere else), so aggregate(type='simple'|'group', balance_e=...) raises rather than silently ignoring it. The other three balance_e sites (ImputationDiD, TwoStageDiD, EfficientDiD) get their own rows in the PRs that migrate them." + - id: M-122 + kind: behavior + group: aggregate-postfit + old: "diff_diff:AggregationResult" + new: null + introduced_in: "3.9" + deprecated_in: null + removed_in: null + status: done + phase: 2 + test_ref: tests/test_aggregate_contract.py + code_refs: [diff_diff/aggregation.py, diff_diff/__init__.py] + notes: "AggregationResult: the tabular container every non-event-study aggregation returns (aggregate(type='event_study') returns EventStudyResults instead, giving [M-092]'s container its public producer). Pinned 12-column schema with summary/to_dict/to_dataframe. Three schema choices are grounded in what the estimators actually produce: 'target' is PER-ROW so one container can carry two aligned estimands over the same labels; 'n_kind' reuses EventStudyResults' vocabulary rather than inventing units/obs; and 'weight' is NULLABLE because CallawaySantAnna's group aggregation weights (g,t) cells equally within each cohort and forms no cross-cohort mass. introduced_in gates the 3.9 cut ([M-091]/[M-092] pattern); deprecated_in stays null so the early-flip guard does not fire against the PR that ships it." diff --git a/docs/v4-design.md b/docs/v4-design.md index cf38b98e2..23edc62b2 100644 --- a/docs/v4-design.md +++ b/docs/v4-design.md @@ -412,10 +412,18 @@ section's pattern - plus `summary(aggregation=)` and `to_dataframe(aggregation=)` [M-044] [M-086] [M-087]. **Semantics.** `aggregate()` re-aggregates WITHOUT refitting, from influence -functions / bootstrap draws retained on the results object (CallawaySantAnna -already stores them; the Phase 2 PR extends retention where missing - memory -cost documented per estimator). Analytical-vs-bootstrap inference of the -aggregated estimand follows the fit's inference method. Estimators whose +functions retained on the results object. **Correction (Phase 2b PR 1):** this +section previously said "CallawaySantAnna already stores them". It does not - +`CallawaySantAnnaResults.influence_functions` was declared and never assigned, +and the fit-time `precomputed` bookkeeping is a local. Retention is therefore +work in EVERY migrating PR, not just the ones "where missing", and the kit must +be built during `fit()` because nothing it needs survives the call. Memory cost +is documented per estimator: for CallawaySantAnna the dominant payload is the +per-(g,t) influence-function dict at roughly O(n_units x n_gt), not the +O(n_units) bookkeeping. Analytical-vs-bootstrap inference of the aggregated +estimand follows the fit's inference method; where bootstrap draws are not +retained, `aggregate()` on a bootstrapped fit RAISES rather than silently +returning analytical inference. Estimators whose estimand is already a single aggregation (SunAbraham's saturated event study, LPDiD's per-horizon design) expose `aggregate()` where meaningful as additive surface; their native params (`only_event`/`only_pooled` etc.) are documented diff --git a/tests/test_aggregate_contract.py b/tests/test_aggregate_contract.py new file mode 100644 index 000000000..0ff99ac98 --- /dev/null +++ b/tests/test_aggregate_contract.py @@ -0,0 +1,386 @@ +"""Behavioral contract for post-fit ``results.aggregate()`` (spec section 6). + +The ``test_ref`` for ledger rows M-020 (``fit(aggregate=)`` shim), M-117 +(``balance_e`` moves onto ``aggregate()``) and M-122 (``AggregationResult``). + +The headline gate is NUMERICAL INERTNESS: for every supported type, +``fit(aggregate=T)`` and ``fit(); .aggregate(T)`` must agree to 1e-14. The +refactor that enabled post-fit aggregation touched the influence-function +path, so drift here means a regression, not a design change. +""" + +import copy +import pickle +import warnings + +import numpy as np +import pandas as pd +import pytest + +from diff_diff import CallawaySantAnna, StaggeredTripleDifference +from diff_diff.aggregation import AGGREGATION_SCHEMA, AggregationResult +from diff_diff.results_base import EventStudyResults + +FIT_KW = dict(outcome="y", unit="unit", time="time", first_treat="first_treat") + + +def _panel(seed=11, n_units=80, n_periods=7): + rng = np.random.default_rng(seed) + cohorts = [0, 3, 4, 5, 6] + rows = [] + for u in range(n_units): + g = cohorts[u % len(cohorts)] + ui = rng.normal(0, 0.5) + for t in range(1, n_periods + 1): + treated = g != 0 and t >= g + rows.append( + { + "unit": u, + "time": t, + "first_treat": g, + "y": ui + + 0.4 * t + + (1.5 + 0.3 * (t - g) if treated else 0.0) + + rng.normal(0, 0.25), + } + ) + return pd.DataFrame(rows) + + +@pytest.fixture(scope="module") +def panel(): + return _panel() + + +@pytest.fixture(scope="module") +def fitted(panel): + """A plain fit - no aggregate= argument, so no deprecation warning.""" + return CallawaySantAnna().fit(panel, **FIT_KW) + + +@pytest.fixture(scope="module") +def fit_time(panel): + """The DEPRECATED fit-time aggregation, for the inertness comparison.""" + with warnings.catch_warnings(): + warnings.simplefilter("ignore", FutureWarning) + return CallawaySantAnna().fit(panel, aggregate="all", **FIT_KW) + + +# --------------------------------------------------------------------------- # +# Numerical inertness (the headline gate) +# --------------------------------------------------------------------------- # + + +class TestInertness: + def test_simple_matches_fit_time(self, fitted, fit_time): + got = fitted.aggregate("simple") + assert np.allclose(got.att[0], fit_time.overall_att, rtol=1e-14, atol=1e-14) + assert np.allclose(got.se[0], fit_time.overall_se, rtol=1e-14, atol=1e-14) + assert np.allclose(got.p_value[0], fit_time.overall_p_value, rtol=1e-14, atol=1e-14) + + def test_group_matches_fit_time(self, fitted, fit_time): + frame = fitted.aggregate("group").to_dataframe() + assert len(frame) == len(fit_time.group_effects) + for _, row in frame.iterrows(): + native = fit_time.group_effects[row["label"]] + for col, key in ( + ("att", "effect"), + ("se", "se"), + ("t_stat", "t_stat"), + ("p_value", "p_value"), + ): + assert np.allclose( + row[col], native[key], rtol=1e-14, atol=1e-14, equal_nan=True + ), f"group[{row['label']}].{col} drifted" + + def test_event_study_matches_fit_time(self, fitted, fit_time): + frame = fitted.aggregate("event_study").to_dataframe() + compared = 0 + for _, row in frame.iterrows(): + if bool(row["is_reference"]): + continue + native = fit_time.event_study_effects[row["event_time"]] + for col, key in ( + ("att", "effect"), + ("se", "se"), + ("t_stat", "t_stat"), + ("p_value", "p_value"), + ): + assert np.allclose( + row[col], native[key], rtol=1e-14, atol=1e-14, equal_nan=True + ), f"event_study[{row['event_time']}].{col} drifted" + compared += 1 + assert compared > 0, "no non-reference event times compared" + + @pytest.mark.parametrize("balance_e", [0, 1, 2]) + def test_balance_e_matches_fit_time(self, panel, fitted, balance_e): + """M-117: balance_e on aggregate() reproduces fit(balance_e=).""" + with warnings.catch_warnings(): + warnings.simplefilter("ignore", FutureWarning) + native = CallawaySantAnna().fit( + panel, aggregate="event_study", balance_e=balance_e, **FIT_KW + ) + frame = fitted.aggregate("event_study", balance_e=balance_e).to_dataframe() + non_ref = frame[~frame["is_reference"].astype(bool)] + assert len(non_ref) == len( + [e for e in native.event_study_effects if e in set(non_ref["event_time"])] + ) + for _, row in non_ref.iterrows(): + got = native.event_study_effects[row["event_time"]] + assert np.allclose(row["att"], got["effect"], rtol=1e-14, atol=1e-14, equal_nan=True) + assert np.allclose(row["se"], got["se"], rtol=1e-14, atol=1e-14, equal_nan=True) + + +# --------------------------------------------------------------------------- # +# Immutability +# --------------------------------------------------------------------------- # + + +class TestImmutability: + #: The attributes the aggregators used to mutate on their host. + MUTATED = ( + "event_study_effects", + "group_effects", + "event_study_vcov", + "event_study_vcov_index", + "event_study_df", + ) + + def test_parent_unchanged_across_mixed_calls(self, fitted): + before = {f: copy.deepcopy(getattr(fitted, f, None)) for f in self.MUTATED} + for level in ("event_study", "group", "simple", "group", "event_study"): + fitted.aggregate(level) + for f in self.MUTATED: + after = getattr(fitted, f, None) + if before[f] is None: + assert after is None, f"{f} was populated by aggregate()" + else: + assert str(before[f]) == str(after), f"{f} changed" + + def test_repeated_calls_agree(self, fitted): + a = fitted.aggregate("group").to_dataframe() + b = fitted.aggregate("group").to_dataframe() + # assert_frame_equal treats corresponding NaNs as equal and handles the + # mixed object/float dtypes; a bare np.array_equal does neither. + pd.testing.assert_frame_equal(a, b) + + def test_order_independent(self, fitted): + es_first = fitted.aggregate("event_study").to_dataframe() + fitted.aggregate("group") + es_after = fitted.aggregate("event_study").to_dataframe() + assert np.allclose(es_first["att"].to_numpy(), es_after["att"].to_numpy(), equal_nan=True) + + +# --------------------------------------------------------------------------- # +# Retention kit +# --------------------------------------------------------------------------- # + + +class TestRetentionKit: + def test_kit_holds_no_dataframe(self, fitted): + """The source panel must never be retained on a results object.""" + kit = fitted._aggregation_kit + frames = [k for k, v in kit.bookkeeping.items() if isinstance(v, pd.DataFrame)] + assert frames == [], f"kit retained DataFrame(s): {frames}" + assert not isinstance(kit.influence, pd.DataFrame) + + def test_no_dataframe_anywhere_on_results(self, fitted): + frames = [ + f + for f in fitted.__dataclass_fields__ + if isinstance(getattr(fitted, f, None), pd.DataFrame) + ] + assert frames == [], f"results retained DataFrame(s): {frames}" + + def test_aggregate_survives_pickle(self, fitted): + """No live reference to the estimator or its frame.""" + revived = pickle.loads(pickle.dumps(fitted)) + got = revived.aggregate("group").to_dataframe() + want = fitted.aggregate("group").to_dataframe() + assert np.allclose(got["att"].to_numpy(), want["att"].to_numpy(), equal_nan=True) + + +# --------------------------------------------------------------------------- # +# Fail-closed contract +# --------------------------------------------------------------------------- # + + +class TestFailClosed: + def test_calendar_unsupported_by_cs(self, fitted): + with pytest.raises(ValueError, match="calendar"): + fitted.aggregate("calendar") + + def test_all_is_not_a_post_fit_type(self, fitted): + with pytest.raises(ValueError, match="Unsupported aggregation type"): + fitted.aggregate("all") + + def test_error_names_supported_types(self, fitted): + with pytest.raises(ValueError) as exc: + fitted.aggregate("nonsense") + for level in ("simple", "event_study", "group"): + assert level in str(exc.value) + + def test_weights_rejected(self, fitted): + """CS exposes no weighting selector, so anything but None fails closed.""" + with pytest.raises(ValueError, match="weights"): + fitted.aggregate("simple", "cohort_share") + + @pytest.mark.parametrize("level", ["simple", "group"]) + def test_balance_e_rejected_where_inert(self, fitted, level): + """balance_e applies to event-study aggregation ONLY - silently + ignoring it elsewhere would accept a user argument that does nothing.""" + with pytest.raises(ValueError, match="balance_e"): + fitted.aggregate(level, balance_e=2) + + def test_bootstrap_fit_raises(self, panel): + """Percentile-bootstrap inference cannot be reproduced from the + analytical state retained here, so aggregate() must not substitute + analytical numbers silently.""" + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + boot = CallawaySantAnna(n_bootstrap=49, seed=42).fit(panel, **FIT_KW) + with pytest.raises(NotImplementedError, match="bootstrap"): + boot.aggregate("group") + + +# --------------------------------------------------------------------------- # +# Container schema (M-122) +# --------------------------------------------------------------------------- # + + +class TestAggregationResult: + def test_pinned_schema(self, fitted): + assert tuple(fitted.aggregate("group").to_dataframe().columns) == AGGREGATION_SCHEMA + + def test_event_study_returns_the_unified_container(self, fitted): + """M-092's container finally gets a public producer.""" + assert isinstance(fitted.aggregate("event_study"), EventStudyResults) + + def test_non_event_study_returns_aggregation_result(self, fitted): + assert isinstance(fitted.aggregate("group"), AggregationResult) + assert isinstance(fitted.aggregate("simple"), AggregationResult) + + def test_serializers_present(self, fitted): + got = fitted.aggregate("group") + assert isinstance(got.to_dict(), dict) + assert isinstance(got.to_dataframe(), pd.DataFrame) + assert isinstance(got.summary(), str) + + def test_group_weight_is_none_not_fabricated(self, fitted): + """_aggregate_by_group weights (g,t) cells equally WITHIN a cohort and + forms no cross-cohort mass, so there is no per-row weight to report.""" + assert fitted.aggregate("group").weight is None + assert fitted.aggregate("group").to_dict()["weight"] is None + + def test_n_kind_is_declared(self, fitted): + """`n` means different things per level; n_kind is what disambiguates.""" + assert fitted.aggregate("group").n_kind == "cells" + assert fitted.aggregate("simple").n_kind == "units" + + def test_target_is_per_row(self, fitted): + got = fitted.aggregate("group") + assert got.target.shape == got.label.shape + + def test_zero_row_is_a_supported_boundary(self): + empty = AggregationResult( + level="group", + label=np.array([], dtype=object), + target=np.array([], dtype=object), + att=np.array([]), + se=np.array([]), + t_stat=np.array([]), + p_value=np.array([]), + conf_int_lower=np.array([]), + conf_int_upper=np.array([]), + n=np.array([]), + df=np.array([]), + ) + assert len(empty.to_dataframe()) == 0 + assert isinstance(empty.summary(), str) + + def test_non_estimable_row_keeps_point_estimate(self): + """safe_inference NaNs t/p/CI only - att and se are inputs, not + outputs, so NaN-ing the whole quintet would erase valid estimates.""" + got = AggregationResult( + level="group", + label=np.array([1], dtype=object), + target=np.array(["att"], dtype=object), + att=np.array([2.0]), + se=np.array([np.nan]), + t_stat=np.array([np.nan]), + p_value=np.array([np.nan]), + conf_int_lower=np.array([np.nan]), + conf_int_upper=np.array([np.nan]), + n=np.array([3.0]), + df=61.0, + ) + assert got.att[0] == 2.0 + assert np.isnan(got.t_stat[0]) + assert np.isnan(got.df[0]), "df must be NaN where p_value is non-finite" + + +# --------------------------------------------------------------------------- # +# Deprecation shim (M-020 / M-117) +# --------------------------------------------------------------------------- # + + +class TestFitShim: + def test_plain_fit_does_not_warn(self, panel): + """The sentinel default is what makes this true - a bare None default + could not tell 'not passed' from 'passed None'.""" + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + CallawaySantAnna().fit(panel, **FIT_KW) + assert [w for w in caught if issubclass(w.category, FutureWarning)] == [] + + @pytest.mark.parametrize("kwargs", [{"aggregate": "group"}, {"balance_e": 1}]) + def test_deprecated_args_warn(self, panel, kwargs): + with pytest.warns(FutureWarning, match="aggregate"): + CallawaySantAnna().fit(panel, **kwargs, **FIT_KW) + + def test_deprecated_path_still_populates_legacy_surface(self, fit_time): + """Downstream consumers (honest_did, pretrends, + build_event_study_surface) read these off the results object.""" + assert fit_time.event_study_effects is not None + assert fit_time.group_effects is not None + + +# --------------------------------------------------------------------------- # +# Cross-estimator regression +# --------------------------------------------------------------------------- # + + +def test_staggered_triple_diff_overall_att_es_still_works(): + """StaggeredTripleDifference inherits the CS aggregation mixin and is the + ONLY reader of the Eq. 4.14 overall, which the purity refactor converted + from a self-attribute to a returned value.""" + rng = np.random.default_rng(5) + rows = [] + for u in range(80): + g = [0, 3, 4][u % 3] + elig = u % 2 + for t in range(1, 6): + treated = g != 0 and t >= g and elig == 1 + rows.append( + { + "unit": u, + "period": t, + "first_treat": g, + "eligibility": elig, + "outcome": 0.3 * t + (1.2 if treated else 0.0) + rng.normal(0, 0.2), + } + ) + data = pd.DataFrame(rows) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = StaggeredTripleDifference().fit( + data, + "outcome", + "unit", + "period", + "first_treat", + "eligibility", + aggregate="event_study", + ) + assert res.overall_att_es is not None + assert np.isfinite(res.overall_att_es) diff --git a/tests/test_v4_matrix.py b/tests/test_v4_matrix.py index 35aba6f63..5d504a6f4 100644 --- a/tests/test_v4_matrix.py +++ b/tests/test_v4_matrix.py @@ -113,10 +113,11 @@ # plus the 2 Phase 2a results-contract rows (M-092/M-093) = 77, plus the 3 # gating-completeness amendment rows (M-094..M-096) = 80, plus the 19-row # completeness sweep over public FUNCTIONS and the dCDH results mirror -# (M-097..M-115) = 99. Ids are never reused and terminal rows are never +# (M-097..M-115) = 99, plus Phase 2b PR 1's two rows (M-117, M-122) = 101. +# Ids are never reused and terminal rows are never # deleted, so the ledger only grows - raise the floor when rows are added; a # lower parse count means scanner/format drift or an illegal row deletion. -ROW_COUNT_FLOOR = 99 +ROW_COUNT_FLOOR = 101 # Committed snapshot of the shipped id set ("ids are never deleted or reused" # contract - a delete-one-add-one edit keeps the count above the floor but trips @@ -127,7 +128,10 @@ # existing param rows, plus the wild-bootstrap exposure policy); # (97,115) = its completeness sweep over module-level public functions # (twowayfeweights, three HAD pretest entry points, trim_weights) plus the -# dCDH results mirror and the fourth `robust` site. +# dCDH results mirror and the fourth `robust` site; (117,117)/(122,122) = +# Phase 2b PR 1 (balance_e onto aggregate(), AggregationResult). M-116 and +# M-118..M-121 are reserved for the later 2b PRs, not deleted - ids are +# never reused, so a gap here is intentional. _INITIAL_ID_RANGES = [ (1, 8), (10, 16), @@ -140,6 +144,8 @@ (92, 93), (94, 96), (97, 115), + (117, 117), + (122, 122), ] EXPECTED_INITIAL_IDS = frozenset( f"M-{n:03d}" for lo, hi in _INITIAL_ID_RANGES for n in range(lo, hi + 1) @@ -541,7 +547,7 @@ def test_initial_ids_never_deleted(): public-function completeness sweep).""" missing = sorted(EXPECTED_INITIAL_IDS - set(_ROW_IDS)) assert not missing, f"ledger rows deleted (ids are permanent): {missing}" - assert len(EXPECTED_INITIAL_IDS) == 99 + assert len(EXPECTED_INITIAL_IDS) == 101 def test_version_tuple_pads_to_three_components(): From 821dddd3b54ffa215a806dd37d51d70cbb21dc68 Mon Sep 17 00:00:00 2001 From: igerber Date: Sat, 25 Jul 2026 20:00:50 -0400 Subject: [PATCH 04/11] fix(v4): correct aggregate() df provenance, container normalization, stale guidance Local review round 1. Every finding below was reproduced before being fixed. df is per-row PROVENANCE - the degrees of freedom that actually produced that row's stored p-value and interval - so reading whichever df field happened to be populated was wrong twice: - "simple" read df_inference, which the field's OWN docstring documents as staying None on explicit survey_design= fits (the canonical carrier there is survey_metadata.df_survey). On an explicit PSU design the container reported df=NaN - implying normal or undefined inference - while the CI it carried was built on a finite t(16) reference. - "group" read event_study_df: a DIFFERENT aggregation's denominator, and None after a plain fit. Both now resolve from the carrier that governed the statistic. _aggregate_by_group records the df its own safe_inference_batch used (the conservative min under dropped replicates, recorded iff it governs a t-reference), and a shared resolve_inference_df() encodes the survey_metadata-then-df_inference precedence rather than adding a fourth independent copy of it - honest_did.py has three, tracked for consolidation. Container normalization: - df was normalized with asarray, so NaN-ing the non-estimable rows wrote THROUGH to the caller's array (verified: [10., 20.] became [10., nan], shares_memory True) and raised on a read-only buffer. Now np.array, matching EventStudyResults' normalization of the same field. - The 1-d check on label ran AFTER shape[0], so a 0-d label raised IndexError instead of the documented ValueError. - n_kind emitted "units"/"cells", neither in the vocabulary the docs claimed it shared with EventStudyResults. The vocabulary is now a single exported N_KIND_VOCABULARY covering both containers, validated on construction, with "cells" and "units" documented as members rather than silent extras. Stale guidance. The messages this branch rewrote pointed at a path that does not work: HonestDiD, PreTrendsPower AND plot_event_study all dispatch on CallawaySantAnnaResults and reject the EventStudyResults that aggregate() returns, so "call results.aggregate('event_study')" ended in a TypeError. They now say plainly that these consumers read the fit-time surface and do not yet accept the post-fit container. Making them accept it needs an adapter plus base_period/anticipation provenance the unified container does not carry - a scope call, not a message fix. Two CS sites the first sweep missed (to_dataframe's event_study and group branches) are migrated; the estimator docstring no longer demonstrates the deprecated kwarg or a plot call that would raise. Tests (47 in the contract suite, up from 33): repeated-cross-section and anticipation=1 round-trips at 1e-14 - designs the balanced-panel fixture never reached, and which the CHANGELOG had claimed on the strength of a throwaway golden rather than a committed test; explicit-survey df reaching every level; group df independent of event_study_df; the mutation, read-only, 0-d and off-vocabulary regressions. The CHANGELOG claim now describes what ships, and bootstrap replay has a real TODO.md row instead of an untracked "tracked follow-up". --- CHANGELOG.md | 18 ++- TODO.md | 2 + diff_diff/aggregation.py | 61 ++++++++- diff_diff/honest_did.py | 5 +- diff_diff/pretrends.py | 5 +- diff_diff/results_base.py | 32 ++++- diff_diff/staggered.py | 19 ++- diff_diff/staggered_aggregation.py | 12 ++ diff_diff/staggered_results.py | 26 +++- tests/test_aggregate_contract.py | 203 +++++++++++++++++++++++++++++ 10 files changed, 347 insertions(+), 36 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index fa2cb1c5d..72e8962a0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,21 +23,25 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 it applies to event-study aggregation only — passing it with `"simple"` or `"group"` raises rather than being silently ignored, matching where the shipped code actually threads it. - **No numerical change:** a 10-scenario golden (every aggregation type, a - `balance_e` sweep, repeated cross-sections, a bootstrapped fit, - `anticipation=1`) is reproduced at `atol=rtol=1e-14`, and post-fit - `aggregate()` matches the fit-time numbers identically. + **No numerical change:** post-fit `aggregate()` reproduces the fit-time + numbers at `atol=rtol=1e-14` for every aggregation type, across a `balance_e` + sweep, a balanced panel, repeated cross-sections and `anticipation=1`. Three schema decisions are grounded in what the estimator produces rather than assumed: `target` is per-row so one container can carry two aligned - estimands; `n_kind` reuses `EventStudyResults`' vocabulary (`"cells"` for - CS's group level, not "units"); and `weight` is `None` for `"group"`, + estimands; `n_kind` draws on a vocabulary now shared with + `EventStudyResults` and validated on construction (`"cells"` for CS's group + level, `"units"` for the overall); and `weight` is `None` for `"group"`, because that aggregation weights `(g,t)` cells equally *within* each cohort and forms no cross-cohort mass — reporting one would be fabricated. Fail-closed elsewhere too: `aggregate("calendar")` raises (CS has no calendar aggregator; the DEFERRED row stands), a non-`None` `weights` raises, and `aggregate()` on a **bootstrapped** fit raises rather than substituting analytical inference for percentile-bootstrap statistics — - bootstrap replay is a tracked follow-up. + bootstrap replay is tracked in `TODO.md`. + `AggregationResult.df` is per-row provenance: the degrees of freedom that + actually produced that row's stored p-value and interval, read from the + carrier the fit used (`survey_metadata.df_survey` on explicit survey + designs, `df_inference` only as the bare-`cluster=` fallback). Internally, the CS event-study aggregator became **pure**: it returned one dict and stashed four more values on `self`, and now returns all five, which is what allows re-aggregation without mutating the result. diff --git a/TODO.md b/TODO.md index 31ec36971..fbe8124c4 100644 --- a/TODO.md +++ b/TODO.md @@ -22,6 +22,8 @@ Related tracking surfaces: | Issue | Location | Origin | Effort | Priority | |-------|----------|--------|--------|----------| | `SyntheticControl` conformal (CWZ 2021) AR / innovation-permutation path (Lemmas 5-7) for time-series proxies — the residual-permutation shortcut is only valid for time-permutation-invariant proxies (SC/Lasso/DiD); an AR proxy needs innovation permutation. | `diff_diff/conformal.py`, `diff_diff/synthetic_control_results.py` | CWZ-2021 | Heavy | Low | +| Bootstrap re-aggregation for `CallawaySantAnnaResults.aggregate()` — a bootstrapped fit currently RAISES rather than substituting analytical inference for percentile-bootstrap statistics. The value-bound `BootstrapReplaySpec` (bit-identical replay, picklable, immune to post-fit `set_params`) is already in-tree and spike-verified; wiring it needs the per-`(g,t)` / per-event-time draw retention plus `assert_allclose` parity tests against the fit-time bootstrap numbers (NOT bit-identity — the fused GEMM's column count differs post-fit, ~1 ULP reassociation). | `diff_diff/aggregation.py`, `diff_diff/staggered_results.py` | #726 | Mid | Medium | +| Consolidate the inference-df precedence duplicated across `honest_did.py` (3 copies at ~L655/L836/L1004) onto the shared `resolve_inference_df()` helper added in `diff_diff/aggregation.py`. The copies are correct today; the risk is drift if the survey/replicate precedence changes in one place only. | `diff_diff/honest_did.py` | #726 | Quick | Low | | `ContinuousDiD` CGBS-2024 remaining extensions (earlier phases — `covariates=` reg/dr, `treatment_type="discrete"`, single-cohort `control_group="lowest_dose"` with estimand `ATT(d)−ATT(d_L)` — are already supported; see REGISTRY Note #7). Remaining (all deferred `NotImplementedError`, documented): `estimation_method="ipw"` on the dose curve (scalar-adjustment / degenerate); `covariates=` × `survey_design=` (weighted OR + weighted nuisance IF); multi-cohort **heterogeneous-support** discrete aggregation (support-aware: average each dose only over the cohorts that observe it); **multi-cohort `lowest_dose`** (within-cohort `d_L` reference + support-aware cross-cohort aggregation); and **`covariates=` × `lowest_dose`** (conditional-PT-relative-to-`d_L` estimand). Single-cohort / 2-period / shared-support multi-cohort are supported. | `continuous_did.py` | CGBS-2024 | Heavy | Low | ### Performance diff --git a/diff_diff/aggregation.py b/diff_diff/aggregation.py index a414c409b..4c7984e31 100644 --- a/diff_diff/aggregation.py +++ b/diff_diff/aggregation.py @@ -30,7 +30,7 @@ import numpy as np import pandas as pd -from diff_diff.results_base import BaseResults, _json_safe_label +from diff_diff.results_base import N_KIND_VOCABULARY, BaseResults, _json_safe_label __all__ = ["AggregationResult", "AGGREGATION_SCHEMA"] @@ -65,6 +65,38 @@ ) +def resolve_inference_df(results: Any) -> Optional[float]: + """Return the degrees of freedom that governed ``results``' overall statistic. + + Reads the carriers in the order the library documents on + ``CallawaySantAnnaResults.df_inference``: ``survey_metadata.df_survey`` + first - it holds the actual CS-internal df including any post-resolve + tightening for replicate designs - and ``df_inference`` only as the + fallback for bare-``cluster=`` fits, where ``survey_metadata`` is + intentionally ``None`` to preserve the survey/non-survey contract. + Reading ``df_inference`` first would overstate the denominator df on + panel survey fits whose df was tightened during aggregation. + + A replicate design whose df is undefined resolves to the sentinel ``0``, + which yields NaN inference rather than a t-reference. + + ``AggregationResult.df`` is per-row PROVENANCE - the df actually passed to + ``safe_inference`` for that row's stored p-value/CI - so it must come from + here rather than from whichever df field happens to be populated. + """ + if getattr(results, "survey_metadata", None) is not None: + sm = results.survey_metadata + df_survey = getattr(sm, "df_survey", None) + if df_survey is None and getattr(sm, "replicate_method", None) is not None: + return 0.0 # undefined replicate df -> NaN inference + if df_survey is not None: + return float(df_survey) + df_inference = getattr(results, "df_inference", None) + if df_inference is not None: + return float(df_inference) + return None + + def _sortable(labels: np.ndarray) -> bool: """Can ``labels`` be ordered without raising? @@ -111,10 +143,11 @@ class AggregationResult(BaseResults): Per-row count as float, NaN where the producer records none. Its SEMANTIC is ``n_kind`` - never assume units. n_kind : str or None - Semantic of ``n``, from the same vocabulary - :class:`~diff_diff.results_base.EventStudyResults` uses - (``"groups"`` / ``"cells"`` / ``"obs"`` / ``"clusters"``). ``None`` - when the producer records no count. + Semantic of ``n``, from + :data:`~diff_diff.results_base.N_KIND_VOCABULARY` - the SAME closed + vocabulary :class:`~diff_diff.results_base.EventStudyResults` draws + on, so a consumer can route on ``n_kind`` across both containers. + ``None`` when the producer records no count. weight : np.ndarray or None Normalized aggregation mass per row, summing to 1 within one ``(level, target)`` group. ``None`` where no per-row mass exists - @@ -166,11 +199,14 @@ class AggregationResult(BaseResults): def __post_init__(self) -> None: self.label = np.asarray(self.label, dtype=object) - n_rows = self.label.shape[0] + # Shape check BEFORE reading shape[0]: a 0-d label has an empty shape + # tuple, so indexing it first raises IndexError instead of the + # documented ValueError. if self.label.ndim != 1: raise ValueError( f"AggregationResult label must be one-dimensional; got shape {self.label.shape}." ) + n_rows = self.label.shape[0] for name in ("att", "se", "t_stat", "p_value", "conf_int_lower", "conf_int_upper", "n"): arr = np.asarray(getattr(self, name), dtype=float) @@ -196,7 +232,10 @@ def __post_init__(self) -> None: elif np.ndim(self.df) == 0: df_arr = np.full(n_rows, float(self.df)) else: - df_arr = np.asarray(self.df, dtype=float) + # np.array (not asarray) so the NaN-out below cannot write through + # to a caller-owned array or fail on a read-only buffer - matching + # EventStudyResults' normalization of the same field. + df_arr = np.array(self.df, dtype=float) if df_arr.shape != (n_rows,): raise ValueError( f"AggregationResult df has shape {df_arr.shape}; expected ({n_rows},) or scalar." @@ -212,6 +251,14 @@ def __post_init__(self) -> None: ) self.weight = w + # n_kind is a routing key consumers share with EventStudyResults, so + # an off-vocabulary value is a contract break, not a free-form label. + if self.n_kind is not None and self.n_kind not in N_KIND_VOCABULARY: + raise ValueError( + f"AggregationResult n_kind {self.n_kind!r} is not in the shared " + f"vocabulary {N_KIND_VOCABULARY}." + ) + # ------------------------------------------------------------------ # # Serialization (spec section 5: every main results class has all three) # ------------------------------------------------------------------ # diff --git a/diff_diff/honest_did.py b/diff_diff/honest_did.py index 9ea413844..10b376a9d 100644 --- a/diff_diff/honest_did.py +++ b/diff_diff/honest_did.py @@ -692,8 +692,9 @@ def _extract_event_study_params( if results.event_study_effects is None: raise ValueError( "CallawaySantAnnaResults must have event_study_effects for HonestDiD. " - "Call results.aggregate('event_study') to compute them post-fit " - "(no refit required)." + "Fit with aggregate='event_study'. HonestDiD reads the fit-time " + "surface and does not yet accept the post-fit container returned " + "by results.aggregate('event_study')." ) # Warn if not using universal base period (R's HonestDiD requires it) diff --git a/diff_diff/pretrends.py b/diff_diff/pretrends.py index 8249f6291..79c94c212 100644 --- a/diff_diff/pretrends.py +++ b/diff_diff/pretrends.py @@ -1140,8 +1140,9 @@ def _extract_pre_period_params( if results.event_study_effects is None: raise ValueError( "CallawaySantAnnaResults must have event_study_effects. " - "Call results.aggregate('event_study') to compute them " - "post-fit (no refit required)." + "Fit with aggregate='event_study'. PreTrendsPower reads the " + "fit-time surface and does not yet accept the post-fit container " + "returned by results.aggregate('event_study')." ) # Get pre-period effects. Anticipation-aware cutoff per diff --git a/diff_diff/results_base.py b/diff_diff/results_base.py index db3d91fbc..4c533651e 100644 --- a/diff_diff/results_base.py +++ b/diff_diff/results_base.py @@ -129,6 +129,20 @@ def _json_safe_label(value: Any) -> Any: "is_reference", ) +#: Closed vocabulary for the ``n_kind`` field, SHARED by every container that +#: reports a count alongside an estimate (``EventStudyResults`` and +#: ``AggregationResult``), so a consumer can route on ``n_kind`` uniformly +#: across the two. Each value names what one unit of ``n`` counts; never +#: conflate them. +N_KIND_VOCABULARY: Tuple[str, ...] = ( + "groups", + "switcher_cells", + "cells", + "units", + "obs", + "clusters", +) + @dataclass class EventStudyResults(BaseResults): @@ -166,13 +180,17 @@ class EventStudyResults(BaseResults): Per-event-time count as float, NaN where the producer records none (and on the reference row - no estimation happened there). n_kind : str or None - Semantic of ``n`` for this producer: ``"groups"`` (a group-level - count - cohorts for CallawaySantAnna/SunAbraham, eligible switcher - groups per horizon for de Chaisemartin-D'Haultfoeuille with - ``L_max >= 1``), ``"switcher_cells"`` (dCDH legacy ``L_max is None`` - path: switching ``(g, t)`` cells, where one group may contribute - several), ``"obs"`` (observations), ``"clusters"``, or None when no - count is recorded. Never conflate these units. + Semantic of ``n`` for this producer, drawn from the shared + :data:`N_KIND_VOCABULARY`: ``"groups"`` (a group-level count - + cohorts for CallawaySantAnna/SunAbraham, eligible switcher groups + per horizon for de Chaisemartin-D'Haultfoeuille with ``L_max >= 1``), + ``"switcher_cells"`` (dCDH legacy ``L_max is None`` path: switching + ``(g, t)`` cells, where one group may contribute several), + ``"cells"`` (``(g, t)`` cells generally - what CallawaySantAnna's + ``"group"`` aggregation counts), ``"units"`` (distinct units, as in + the overall/simple aggregation), ``"obs"`` (observations), + ``"clusters"``, or None when no count is recorded. Never conflate + these units. reference_period : Any or None Convenience scalar echo of the marked row's ``event_time`` label when there is EXACTLY ONE reference row; None when there are zero or diff --git a/diff_diff/staggered.py b/diff_diff/staggered.py index 89839d751..923d17e7f 100644 --- a/diff_diff/staggered.py +++ b/diff_diff/staggered.py @@ -475,16 +475,23 @@ class CallawaySantAnna( >>> >>> results.print_summary() - With event study aggregation: + With event study aggregation (post-fit - no refit required): >>> cs = CallawaySantAnna() >>> results = cs.fit(data, outcome='outcome', unit='unit', - ... time='time', first_treat='first_treat', - ... aggregate='event_study') - >>> - >>> # Plot event study + ... time='time', first_treat='first_treat') + >>> event_study = results.aggregate('event_study') + >>> event_study.to_dataframe() # doctest: +SKIP + + Plotting and the sensitivity analyses (``plot_event_study``, + ``compute_honest_did``, ``compute_pretrends_power``) still read the + fit-time surface, so they take a fit that requested aggregation: + >>> from diff_diff import plot_event_study - >>> plot_event_study(results) + >>> plotted = cs.fit(data, outcome='outcome', unit='unit', + ... time='time', first_treat='first_treat', + ... aggregate='event_study') # doctest: +SKIP + >>> plot_event_study(plotted) # doctest: +SKIP With covariate adjustment (conditional parallel trends): diff --git a/diff_diff/staggered_aggregation.py b/diff_diff/staggered_aggregation.py index 6dbbea9d6..e34b8940b 100644 --- a/diff_diff/staggered_aggregation.py +++ b/diff_diff/staggered_aggregation.py @@ -1268,6 +1268,17 @@ def _aggregate_by_group( df=df_survey_val, ) + # Provenance: the ONE df every group row's inference just used (the + # conservative min under dropped replicates). Recorded per row iff it + # will govern a t-reference - finite and > 0; the df=0 replicate + # sentinel yields NaN inference, not a t-law. Carried on the effect + # dicts so a post-fit ``aggregate("group")`` reports the df that + # actually produced the stored p-value/CI rather than re-deriving it + # from whichever df field happens to be populated on the results. + group_df_used: Optional[float] = None + if df_survey_val is not None and np.isfinite(df_survey_val) and df_survey_val > 0: + group_df_used = float(df_survey_val) + group_effects = {} for idx, (g, agg_effect, agg_se, n_periods, _eff_df) in enumerate(group_data_list): group_effects[g] = { @@ -1277,6 +1288,7 @@ def _aggregate_by_group( "p_value": float(p_values[idx]), "conf_int": (float(ci_lowers[idx]), float(ci_uppers[idx])), "n_periods": n_periods, + "df_used": group_df_used, } return group_effects diff --git a/diff_diff/staggered_results.py b/diff_diff/staggered_results.py index d219173b1..924793a7f 100644 --- a/diff_diff/staggered_results.py +++ b/diff_diff/staggered_results.py @@ -11,7 +11,7 @@ import numpy as np import pandas as pd -from diff_diff.aggregation import AggregationMixin, AggregationResult +from diff_diff.aggregation import AggregationMixin, AggregationResult, resolve_inference_df from diff_diff.results import _format_survey_block, _get_significance_stars from diff_diff.results_base import BaseResults, build_event_study_surface from diff_diff.staggered_aggregation import CallawaySantAnnaAggregationMixin @@ -367,7 +367,12 @@ def _aggregate_simple_result(self) -> AggregationResult: conf_int_lower=np.array([ci[0]], dtype=float), conf_int_upper=np.array([ci[1]], dtype=float), n=np.array([float(self.n_treated_units + self.n_control_units)], dtype=float), - df=self.df_inference, + # NOT ``df_inference``: that field is documented to stay None on + # explicit ``survey_design=`` fits, where the df that actually + # governed ``overall_p_value`` lives on ``survey_metadata``. + # Reading it directly reported df=NaN for survey fits whose CI + # was built on a finite t-reference. + df=resolve_inference_df(self), alpha=self.alpha, n_kind="units", weight=np.array([1.0], dtype=float), @@ -399,7 +404,10 @@ def _group_effects_to_aggregation( conf_int_lower=np.array([c[0] for c in cis], dtype=float), conf_int_upper=np.array([c[1] for c in cis], dtype=float), n=np.array([float(r.get("n_periods", np.nan)) for r in rows], dtype=float), - df=self.event_study_df if self.event_study_df is not None else np.nan, + # The df the GROUP aggregation's own inference used, recorded by + # ``_aggregate_by_group``. Previously this read ``event_study_df`` + # - a different aggregation's df, and None after a plain fit. + df=rows[0].get("df_used") if rows else None, alpha=kit.alpha, n_kind="cells", weight=None, @@ -657,7 +665,11 @@ def to_dataframe(self, level: str = "group_time") -> pd.DataFrame: elif level == "event_study": if self.event_study_effects is None: - raise ValueError("Event study effects not computed. Use aggregate='event_study'.") + raise ValueError( + "Event study effects not computed. " + "Call results.aggregate('event_study') to compute them post-fit " + "(no refit required)." + ) rows = [] for rel_t, data in sorted(self.event_study_effects.items()): cband_ci = data.get("cband_conf_int", (np.nan, np.nan)) @@ -678,7 +690,11 @@ def to_dataframe(self, level: str = "group_time") -> pd.DataFrame: elif level == "group": if self.group_effects is None: - raise ValueError("Group effects not computed. Use aggregate='group'.") + raise ValueError( + "Group effects not computed. " + "Call results.aggregate('group') to compute them post-fit " + "(no refit required)." + ) rows = [] for group, data in sorted(self.group_effects.items()): rows.append( diff --git a/tests/test_aggregate_contract.py b/tests/test_aggregate_contract.py index 0ff99ac98..2e116386e 100644 --- a/tests/test_aggregate_contract.py +++ b/tests/test_aggregate_contract.py @@ -345,6 +345,209 @@ def test_deprecated_path_still_populates_legacy_surface(self, fit_time): assert fit_time.group_effects is not None +def _rcs(seed=17, n_per_cell=6, n_periods=6): + """Repeated cross-sections: a DIFFERENT set of units each period, so unit + ids must be globally unique rather than recycled across periods.""" + rng = np.random.default_rng(seed) + cohorts = [0, 3, 4, 5] + rows = [] + uid = 0 + for t in range(1, n_periods + 1): + for g in cohorts: + for _ in range(n_per_cell): + treated = g != 0 and t >= g + rows.append( + { + "unit": uid, + "time": t, + "first_treat": g, + "y": 0.4 * t + (1.5 if treated else 0.0) + rng.normal(0, 0.25), + } + ) + uid += 1 + return pd.DataFrame(rows) + + +class TestInertnessAcrossDesigns: + """The inertness gate on designs the balanced-panel fixture does not reach. + + The refactor made the aggregators pure by threading bookkeeping that the + panel path and the repeated-cross-section path populate differently (RCS + makes several kit keys observation-length rather than unit-length), and + ``anticipation`` shifts which ``(g, t)`` cells are eligible - so both need + their own round-trip, not just the panel. + """ + + @pytest.mark.parametrize("level", ["simple", "group", "event_study"]) + def test_repeated_cross_sections_match_fit_time(self, level): + data = _rcs() + post = CallawaySantAnna(panel=False).fit(data, **FIT_KW) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", FutureWarning) + native = CallawaySantAnna(panel=False).fit(data, aggregate="all", **FIT_KW) + _assert_level_matches(post, native, level) + + @pytest.mark.parametrize("level", ["simple", "group", "event_study"]) + def test_anticipation_matches_fit_time(self, panel, level): + post = CallawaySantAnna(anticipation=1).fit(panel, **FIT_KW) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", FutureWarning) + native = CallawaySantAnna(anticipation=1).fit(panel, aggregate="all", **FIT_KW) + _assert_level_matches(post, native, level) + + +def _assert_level_matches(post, native, level): + """Compare a post-fit aggregation against the fit-time surface at 1e-14.""" + if level == "simple": + got = post.aggregate("simple") + assert np.allclose(got.att[0], native.overall_att, rtol=1e-14, atol=1e-14) + assert np.allclose(got.se[0], native.overall_se, rtol=1e-14, atol=1e-14) + return + + frame = post.aggregate(level).to_dataframe() + if level == "group": + expected, key_col = native.group_effects, "label" + else: + expected, key_col = native.event_study_effects, "event_time" + compared = 0 + for _, row in frame.iterrows(): + if level == "event_study" and bool(row["is_reference"]): + continue + ref = expected[row[key_col]] + for col, key in (("att", "effect"), ("se", "se"), ("p_value", "p_value")): + assert np.allclose( + row[col], ref[key], rtol=1e-14, atol=1e-14, equal_nan=True + ), f"{level}[{row[key_col]}].{col} drifted" + compared += 1 + assert compared > 0, f"no {level} rows compared" + + +@pytest.fixture(scope="module") +def survey_fit(): + """An EXPLICIT survey design - the branch where ``df_inference`` is + intentionally None and ``survey_metadata.df_survey`` is the real carrier.""" + from diff_diff import SurveyDesign + + rng = np.random.default_rng(3) + rows = [] + for u in range(80): + g = [0, 4, 6][u % 3] + for t in range(1, 9): + treated = g != 0 and t >= g + rows.append( + { + "unit": u, + "time": t, + "first_treat": g, + "psu": u % 20, + "stratum": u % 4, + "w": 1.0 + (u % 5) * 0.1, + "y": 0.3 * t + (2.0 if treated else 0.0) + rng.normal(0, 0.5), + } + ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + return CallawaySantAnna().fit( + pd.DataFrame(rows), + survey_design=SurveyDesign(weights="w", strata="stratum", psu="psu"), + **FIT_KW, + ) + + +class TestInferenceDfProvenance: + """``df`` is per-row PROVENANCE - the df that actually produced the stored + p-value/CI - so it must come from the carrier that governed inference. + + Regression: ``simple`` read ``df_inference`` (documented to stay None on + explicit ``survey_design=`` fits) and ``group`` read ``event_study_df`` (a + DIFFERENT aggregation's df, and None after a plain fit). Both reported + ``df=NaN`` - implying normal/undefined inference - while the interval they + carried was built on a finite t-reference. + """ + + def test_explicit_survey_df_reaches_every_level(self, survey_fit): + """The canonical carrier is ``survey_metadata.df_survey``; on an + explicit survey design ``df_inference`` is intentionally None.""" + expected = float(survey_fit.survey_metadata.df_survey) + assert survey_fit.df_inference is None, "fixture no longer exercises the bug" + for level in ("simple", "group", "event_study"): + df_col = np.asarray(survey_fit.aggregate(level).df, dtype=float) + finite = df_col[np.isfinite(df_col)] + assert finite.size > 0, f"{level} reported no df at all" + assert np.all(finite == expected), f"{level} df {finite} != {expected}" + + def test_group_df_is_not_the_event_study_df(self, fitted): + """``group`` must not borrow ``event_study_df``: on a plain fit that + field is None, and it is a different aggregation's denominator.""" + assert fitted.event_study_df is None + grp = fitted.aggregate("group") + assert np.asarray(grp.df).shape == np.asarray(grp.att).shape + + +class TestContainerNormalization: + """``AggregationResult`` normalizes its own columns; doing so must not + reach back into caller-owned memory or mask a shape error.""" + + @staticmethod + def _kw(): + return dict( + level="group", + label=np.array(["a", "b"], dtype=object), + target=np.array(["att", "att"], dtype=object), + att=np.array([1.0, 2.0]), + se=np.array([0.1, 0.2]), + t_stat=np.array([10.0, np.nan]), + p_value=np.array([0.01, np.nan]), + conf_int_lower=np.array([0.8, np.nan]), + conf_int_upper=np.array([1.2, np.nan]), + n=np.array([5.0, 6.0]), + ) + + def test_df_input_is_not_mutated(self): + """The NaN-out of non-estimable rows wrote THROUGH to the caller's + array when df was normalized with ``asarray`` instead of ``array``.""" + caller = np.array([10.0, 20.0]) + AggregationResult(df=caller, **self._kw()) + assert np.array_equal(caller, [10.0, 20.0]), "caller's df array was mutated" + + def test_read_only_df_is_accepted(self): + frozen = np.array([10.0, 20.0]) + frozen.flags.writeable = False + got = AggregationResult(df=frozen, **self._kw()) + assert np.isnan(got.df[1]), "non-estimable row should still be NaN'd" + + def test_zero_dim_label_raises_value_error(self): + """Not IndexError: the shape check must precede the shape read.""" + with pytest.raises(ValueError, match="one-dimensional"): + AggregationResult( + level="simple", + label=np.array("x", dtype=object), + target=np.array(["att"], dtype=object), + att=np.array([1.0]), + se=np.array([0.1]), + t_stat=np.array([1.0]), + p_value=np.array([0.1]), + conf_int_lower=np.array([0.0]), + conf_int_upper=np.array([2.0]), + n=np.array([5.0]), + df=None, + ) + + def test_off_vocabulary_n_kind_raises(self): + """``n_kind`` is a routing key shared with EventStudyResults, so an + unknown value is a contract break rather than a free-form label.""" + with pytest.raises(ValueError, match="vocabulary"): + AggregationResult(df=None, n_kind="widgets", **self._kw()) + + @pytest.mark.parametrize("level,expected", [("simple", "units"), ("group", "cells")]) + def test_cs_n_kinds_are_in_the_shared_vocabulary(self, fitted, level, expected): + from diff_diff.results_base import N_KIND_VOCABULARY + + got = fitted.aggregate(level) + assert got.n_kind == expected + assert got.n_kind in N_KIND_VOCABULARY + + # --------------------------------------------------------------------------- # # Cross-estimator regression # --------------------------------------------------------------------------- # From 41049b144500c25ec71f90a15b2c60526ea97f6c Mon Sep 17 00:00:00 2001 From: igerber Date: Sat, 25 Jul 2026 20:07:20 -0400 Subject: [PATCH 05/11] docs: track the downstream EventStudyResults adapter work (honest_did, pretrends, plot_event_study) --- TODO.md | 1 + 1 file changed, 1 insertion(+) diff --git a/TODO.md b/TODO.md index fbe8124c4..dd8347993 100644 --- a/TODO.md +++ b/TODO.md @@ -22,6 +22,7 @@ Related tracking surfaces: | Issue | Location | Origin | Effort | Priority | |-------|----------|--------|--------|----------| | `SyntheticControl` conformal (CWZ 2021) AR / innovation-permutation path (Lemmas 5-7) for time-series proxies — the residual-permutation shortcut is only valid for time-permutation-invariant proxies (SC/Lasso/DiD); an AR proxy needs innovation permutation. | `diff_diff/conformal.py`, `diff_diff/synthetic_control_results.py` | CWZ-2021 | Heavy | Low | +| Make the post-fit `results.aggregate("event_study")` container consumable downstream. `EventStudyResults` is rejected by all THREE consumers that read a CS event study — `compute_honest_did` (`honest_did.py`, dispatches on `CallawaySantAnnaResults` and raises `TypeError`), `compute_pretrends_power` (`pretrends.py`, same), and `plot_event_study` (`visualization`, same) — so `fit(aggregate="event_study")` is still the only route for them and their error messages say so explicitly. Needs an `EventStudyResults` branch in each extraction path (consuming `event_time` / `is_reference` / `vcov` / `vcov_index` / per-row `df`) PLUS `base_period` and `anticipation` provenance, which the unified container does not carry and HonestDiD needs for its universal-base-period warning and pre-period classification. Gate with end-to-end tests: `compute_honest_did(res.aggregate("event_study"))` at `base_period="universal"`, and `compute_pretrends_power(...)` at `anticipation=1`. | `diff_diff/honest_did.py`, `diff_diff/pretrends.py`, `diff_diff/results_base.py` | #726 | Mid | Medium | | Bootstrap re-aggregation for `CallawaySantAnnaResults.aggregate()` — a bootstrapped fit currently RAISES rather than substituting analytical inference for percentile-bootstrap statistics. The value-bound `BootstrapReplaySpec` (bit-identical replay, picklable, immune to post-fit `set_params`) is already in-tree and spike-verified; wiring it needs the per-`(g,t)` / per-event-time draw retention plus `assert_allclose` parity tests against the fit-time bootstrap numbers (NOT bit-identity — the fused GEMM's column count differs post-fit, ~1 ULP reassociation). | `diff_diff/aggregation.py`, `diff_diff/staggered_results.py` | #726 | Mid | Medium | | Consolidate the inference-df precedence duplicated across `honest_did.py` (3 copies at ~L655/L836/L1004) onto the shared `resolve_inference_df()` helper added in `diff_diff/aggregation.py`. The copies are correct today; the risk is drift if the survey/replicate precedence changes in one place only. | `diff_diff/honest_did.py` | #726 | Quick | Low | | `ContinuousDiD` CGBS-2024 remaining extensions (earlier phases — `covariates=` reg/dr, `treatment_type="discrete"`, single-cohort `control_group="lowest_dose"` with estimand `ATT(d)−ATT(d_L)` — are already supported; see REGISTRY Note #7). Remaining (all deferred `NotImplementedError`, documented): `estimation_method="ipw"` on the dose curve (scalar-adjustment / degenerate); `covariates=` × `survey_design=` (weighted OR + weighted nuisance IF); multi-cohort **heterogeneous-support** discrete aggregation (support-aware: average each dose only over the cohorts that observe it); **multi-cohort `lowest_dose`** (within-cohort `d_L` reference + support-aware cross-cohort aggregation); and **`covariates=` × `lowest_dose`** (conditional-PT-relative-to-`d_L` estimand). Single-cohort / 2-period / shared-support multi-cohort are supported. | `continuous_did.py` | CGBS-2024 | Heavy | Low | From a2750c45c82765cd7f32131cd30fc537c5783f08 Mon Sep 17 00:00:00 2001 From: igerber Date: Sat, 25 Jul 2026 20:16:23 -0400 Subject: [PATCH 06/11] fix(v4): stop retaining raw unit ids, refuse to relabel stored intervals Local review round 2 (round 1 findings all cleared; no P0/P1 this round). Data minimization. The aggregation kit retained ``all_units`` and ``unit_to_idx`` verbatim, so a fitted CS result - which is picklable, and which users share as an artifact - became a carrier for raw unit identifiers. Those are routinely names, emails or administrative IDs, and NO pre-existing public field on CallawaySantAnnaResults carried them, so this branch would have been the regression that introduced the disclosure. The kit needs only POSITION: the influence arrays index by treated_idx/control_idx, and the fast aggregation path keys off object identity plus length, never a label lookup. Verified by substituting canonical codes on a fitted result and re-aggregating - identical at atol=rtol=0, not merely close - so the kit now stores 0..n-1 codes. ``unit_cohorts`` is deliberately kept: those are first-treatment PERIODS, reported as group labels, not identifiers. Pinned by a sentinel-identifier test that searches the whole pickled artifact rather than spot-checking the kit. summary(alpha=) printed "Confidence intervals at alpha=" while the interval had been computed at self.alpha and is never recomputed - false provenance on a display path. EventStudyResults.summary already solved this by raising when the passed alpha differs; this was another divergence from the sibling, and now matches it. Inertness coverage extended to the designs the registry gives their own weighting, reference-cell and VCV-index semantics, none of which the balanced-panel fixture reached: base_period="universal", control_group= "not_yet_treated" with anticipation, allow_unbalanced_panel=True, and survey parity on the ESTIMATES rather than only the df metadata. 62 contract tests (up from 47); 1136 pass across the staggered, SDDD, survey, methodology, event-study-surface, honest-did, pretrends, ledger and docs suites, plus 491 across the sibling estimators that share the aggregation module. --- diff_diff/aggregation.py | 22 +++++- diff_diff/staggered.py | 17 ++++ tests/test_aggregate_contract.py | 128 ++++++++++++++++++++++++++++--- 3 files changed, 152 insertions(+), 15 deletions(-) diff --git a/diff_diff/aggregation.py b/diff_diff/aggregation.py index 4c7984e31..6340c9950 100644 --- a/diff_diff/aggregation.py +++ b/diff_diff/aggregation.py @@ -315,9 +315,23 @@ def to_dict(self) -> Dict[str, Any]: return out def summary(self, alpha: Optional[float] = None) -> str: - """Human-readable table. ``alpha`` is display-only - it does not - recompute the stored interval, which was fixed at aggregation time.""" - a = self.alpha if alpha is None else alpha + """Human-readable table. + + Parameters + ---------- + alpha : float, optional + Accepted for signature uniformity (spec section 5). The stored + interval was computed at aggregation time; passing a value + different from the stored ``alpha`` raises rather than silently + recomputing or mislabeling it - re-aggregate instead. Mirrors + :meth:`~diff_diff.results_base.EventStudyResults.summary`. + """ + if alpha is not None and alpha != self.alpha: + raise ValueError( + f"This aggregation stores intervals computed at alpha={self.alpha}; " + f"re-aggregate to obtain alpha={alpha} intervals " + "(summary() never recomputes stored inference)." + ) who = self.estimator or "estimator" lines = [ f"{who} - aggregate(type={self.level!r})", @@ -338,7 +352,7 @@ def summary(self, alpha: Optional[float] = None) -> str: f"{row['t_stat']:>8.3f} {row['p_value']:>8.4f} {n_disp:>10}" ) lines.append("-" * 64) - lines.append(f"Confidence intervals at alpha={a}.") + lines.append(f"Confidence intervals at alpha={self.alpha}.") if self.weight is None: lines.append("Per-row aggregation weights are not defined for this level.") return "\n".join(lines) diff --git a/diff_diff/staggered.py b/diff_diff/staggered.py index 923d17e7f..023794e85 100644 --- a/diff_diff/staggered.py +++ b/diff_diff/staggered.py @@ -5112,6 +5112,23 @@ def _build_aggregation_kit( if key in precomputed: bookkeeping[key] = precomputed[key] + # Data minimization: the results object is picklable and users share + # result artifacts, so the kit must not turn it into a carrier for raw + # unit identifiers (which are routinely names, emails or administrative + # IDs). ``all_units`` and ``unit_to_idx`` exist here only to size and + # position the influence vectors - the influence arrays index by POSITION + # (``treated_idx`` / ``control_idx``), and the fast aggregation path keys + # off object identity plus length, never off a label lookup. Substituting + # canonical 0..n-1 codes is therefore numerically inert (pinned + # bit-identical by TestRetentionKit) while retaining nothing identifying. + # ``unit_cohorts`` is kept as-is: those are first-treatment PERIODS, which + # are reported as group labels and are not unit identifiers. + if "all_units" in bookkeeping and bookkeeping["all_units"] is not None: + n_kit_units = len(bookkeeping["all_units"]) + bookkeeping["all_units"] = np.arange(n_kit_units) + if bookkeeping.get("unit_to_idx") is not None: + bookkeeping["unit_to_idx"] = {i: i for i in range(n_kit_units)} + return AggregationKit( bookkeeping=bookkeeping, influence=influence_func_info, diff --git a/tests/test_aggregate_contract.py b/tests/test_aggregate_contract.py index 2e116386e..06acf5466 100644 --- a/tests/test_aggregate_contract.py +++ b/tests/test_aggregate_contract.py @@ -199,6 +199,48 @@ def test_aggregate_survives_pickle(self, fitted): want = fitted.aggregate("group").to_dataframe() assert np.allclose(got["att"].to_numpy(), want["att"].to_numpy(), equal_nan=True) + def test_no_raw_unit_identifiers_are_retained(self): + """Data minimization: results objects are picklable and get shared, so + the kit must not become a carrier for unit identifiers - which are + routinely names, emails or administrative IDs. The kit needs only + POSITION (influence arrays index by ``treated_idx``/``control_idx``), + so it stores canonical 0..n-1 codes. + + Searches the whole serialized artifact, not just the kit: a recursive + check is what makes this a real guarantee rather than a spot check. + """ + sentinel = "SENTINEL-ID-{}@example.invalid" + rng = np.random.default_rng(4) + rows = [] + for i in range(45): + g = [0, 3, 4][i % 3] + for t in range(1, 7): + treated = g != 0 and t >= g + rows.append( + { + "unit": sentinel.format(i), + "time": t, + "first_treat": g, + "y": 0.4 * t + (1.5 if treated else 0.0) + rng.normal(0, 0.25), + } + ) + res = CallawaySantAnna().fit(pd.DataFrame(rows), **FIT_KW) + assert b"SENTINEL-ID" not in pickle.dumps(res), "raw unit ids reached the artifact" + # And the minimization is inert: aggregation still produces numbers. + for level in ("simple", "group", "event_study"): + assert np.isfinite(np.asarray(res.aggregate(level).att)).any() + + def test_identifier_minimization_is_numerically_inert(self, panel, fitted): + """Canonical codes must not perturb any aggregate: compared at + atol=rtol=0, since substituting positions for labels should not change + a single floating-point operation.""" + with warnings.catch_warnings(): + warnings.simplefilter("ignore", FutureWarning) + native = CallawaySantAnna().fit(panel, aggregate="all", **FIT_KW) + got = fitted.aggregate("simple") + assert got.att[0] == native.overall_att + assert got.se[0] == native.overall_se + # --------------------------------------------------------------------------- # # Fail-closed contract @@ -395,6 +437,56 @@ def test_anticipation_matches_fit_time(self, panel, level): native = CallawaySantAnna(anticipation=1).fit(panel, aggregate="all", **FIT_KW) _assert_level_matches(post, native, level) + @pytest.mark.parametrize("level", ["simple", "group", "event_study"]) + def test_universal_base_period_matches_fit_time(self, panel, level): + """REGISTRY gives universal bases their own reference-cell and + VCV-index semantics (a zero reference cell per cohort), which the + default 'varying' fixture never exercises.""" + kw = dict(base_period="universal") + post = CallawaySantAnna(**kw).fit(panel, **FIT_KW) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", FutureWarning) + native = CallawaySantAnna(**kw).fit(panel, aggregate="all", **FIT_KW) + _assert_level_matches(post, native, level) + + @pytest.mark.parametrize("level", ["simple", "group", "event_study"]) + def test_not_yet_treated_with_anticipation_matches_fit_time(self, panel, level): + """The not-yet-treated control group interacts with anticipation in + picking each (g, t) comparison set - a different code path from the + never-treated default.""" + kw = dict(control_group="not_yet_treated", anticipation=1) + post = CallawaySantAnna(**kw).fit(panel, **FIT_KW) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", FutureWarning) + native = CallawaySantAnna(**kw).fit(panel, aggregate="all", **FIT_KW) + _assert_level_matches(post, native, level) + + @pytest.mark.parametrize("level", ["simple", "group", "event_study"]) + def test_unbalanced_panel_matches_fit_time(self, panel, level): + """allow_unbalanced_panel changes which units survive into the + influence vectors, so the retained bookkeeping must agree with what + the fit actually used.""" + thinned = panel.drop(panel.index[::13]).reset_index(drop=True) + kw = dict(allow_unbalanced_panel=True) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + post = CallawaySantAnna(**kw).fit(thinned, **FIT_KW) + native = CallawaySantAnna(**kw).fit(thinned, aggregate="all", **FIT_KW) + _assert_level_matches(post, native, level) + + @pytest.mark.parametrize("level", ["simple", "group", "event_study"]) + def test_survey_numbers_match_fit_time(self, survey_fit, level): + """Survey parity on the ESTIMATES, not just the df metadata.""" + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + native = CallawaySantAnna().fit( + _survey_panel(), + survey_design=_survey_design(), + aggregate="all", + **FIT_KW, + ) + _assert_level_matches(survey_fit, native, level) + def _assert_level_matches(post, native, level): """Compare a post-fit aggregation against the fit-time surface at 1e-14.""" @@ -422,12 +514,7 @@ def _assert_level_matches(post, native, level): assert compared > 0, f"no {level} rows compared" -@pytest.fixture(scope="module") -def survey_fit(): - """An EXPLICIT survey design - the branch where ``df_inference`` is - intentionally None and ``survey_metadata.df_survey`` is the real carrier.""" - from diff_diff import SurveyDesign - +def _survey_panel(): rng = np.random.default_rng(3) rows = [] for u in range(80): @@ -445,13 +532,22 @@ def survey_fit(): "y": 0.3 * t + (2.0 if treated else 0.0) + rng.normal(0, 0.5), } ) + return pd.DataFrame(rows) + + +def _survey_design(): + from diff_diff import SurveyDesign + + return SurveyDesign(weights="w", strata="stratum", psu="psu") + + +@pytest.fixture(scope="module") +def survey_fit(): + """An EXPLICIT survey design - the branch where ``df_inference`` is + intentionally None and ``survey_metadata.df_survey`` is the real carrier.""" with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) - return CallawaySantAnna().fit( - pd.DataFrame(rows), - survey_design=SurveyDesign(weights="w", strata="stratum", psu="psu"), - **FIT_KW, - ) + return CallawaySantAnna().fit(_survey_panel(), survey_design=_survey_design(), **FIT_KW) class TestInferenceDfProvenance: @@ -547,6 +643,16 @@ def test_cs_n_kinds_are_in_the_shared_vocabulary(self, fitted, level, expected): assert got.n_kind == expected assert got.n_kind in N_KIND_VOCABULARY + def test_summary_refuses_to_relabel_a_stored_interval(self, fitted): + """summary(alpha=) never recomputes, so printing the passed alpha would + assert a confidence level the stored interval was not built at. Raises + instead, matching EventStudyResults.summary.""" + got = fitted.aggregate("group") + assert f"alpha={got.alpha}" in got.summary() + assert f"alpha={got.alpha}" in got.summary(alpha=got.alpha) + with pytest.raises(ValueError, match="never recomputes"): + got.summary(alpha=0.10) + # --------------------------------------------------------------------------- # # Cross-estimator regression From 9bf68f77fb72f161cefb7bf2f193076ce4865a20 Mon Sep 17 00:00:00 2001 From: igerber Date: Sat, 25 Jul 2026 20:27:30 -0400 Subject: [PATCH 07/11] fix(v4): enforce the shared n_kind vocabulary on both containers Local review round 3. The previous round introduced N_KIND_VOCABULARY and documented it as SHARED by every count-bearing container, then validated it on AggregationResult only - a contract declared in one place and enforced in one of its two implementations. A public EventStudyResults could still carry and serialize n_kind="widgets", which defeats exactly the cross-container routing the shared vocabulary exists to enable, and does so through the unchecked side. EventStudyResults.__post_init__ now applies the same membership check. Verified non-breaking before adding it: the values shipped producers actually emit are "groups", "switcher_cells", "cells", "units" and "obs" (plus None), every one of which is already in the vocabulary, so no producer changes behavior. The regression test asserts both directions on both containers - an unknown value raises, and every real producer value still constructs - rather than only the rejection, since a vocabulary that quietly stopped accepting a live producer value would be the worse failure. --- diff_diff/results_base.py | 9 +++++++++ tests/test_aggregate_contract.py | 25 +++++++++++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/diff_diff/results_base.py b/diff_diff/results_base.py index 4c533651e..2ad7586e0 100644 --- a/diff_diff/results_base.py +++ b/diff_diff/results_base.py @@ -369,6 +369,15 @@ def __post_init__(self) -> None: f"expected ({n_rows},) to align with event_time." ) df_arr[~np.isfinite(self.p_value)] = np.nan + # n_kind is a routing key consumers share with AggregationResult, so an + # off-vocabulary value is a contract break, not a free-form label. Both + # containers validate it - enforcing on only one would let an unknown + # value reach a consumer through the unchecked side. + if self.n_kind is not None and self.n_kind not in N_KIND_VOCABULARY: + raise ValueError( + f"EventStudyResults n_kind {self.n_kind!r} is not in the shared " + f"vocabulary {N_KIND_VOCABULARY}." + ) self.df = df_arr if (self.vcov is None) != (self.vcov_index is None): diff --git a/tests/test_aggregate_contract.py b/tests/test_aggregate_contract.py index 06acf5466..038311e61 100644 --- a/tests/test_aggregate_contract.py +++ b/tests/test_aggregate_contract.py @@ -643,6 +643,31 @@ def test_cs_n_kinds_are_in_the_shared_vocabulary(self, fitted, level, expected): assert got.n_kind == expected assert got.n_kind in N_KIND_VOCABULARY + def test_shared_vocabulary_is_enforced_on_both_containers(self, fitted): + """The vocabulary is declared SHARED, so validating it on only one + container would let an unknown value reach a consumer through the + unchecked side. Every value a real producer emits must still pass. + """ + from diff_diff.results_base import N_KIND_VOCABULARY + + base = dict( + event_time=np.array([-1, 0, 1]), + att=np.array([0.0, 1.0, 2.0]), + se=np.array([0.1, 0.2, 0.3]), + t_stat=np.array([0.0, 5.0, 6.0]), + p_value=np.array([1.0, 0.01, 0.01]), + conf_int_lower=np.array([-0.2, 0.6, 1.4]), + conf_int_upper=np.array([0.2, 1.4, 2.6]), + is_reference=np.array([True, False, False]), + n=np.array([np.nan, 4.0, 4.0]), + ) + with pytest.raises(ValueError, match="vocabulary"): + EventStudyResults(n_kind="widgets", **base) + # Every value shipped producers actually emit stays constructible. + for kind in ("groups", "switcher_cells", "cells", "units", "obs", None): + assert EventStudyResults(n_kind=kind, **base).n_kind == kind + assert set(N_KIND_VOCABULARY) >= {"groups", "switcher_cells", "cells", "units", "obs"} + def test_summary_refuses_to_relabel_a_stored_interval(self, fitted): """summary(alpha=) never recomputes, so printing the passed alpha would assert a confidence level the stored interval was not built at. Raises From e8b1f381ed34f6bbbfc91082268a866afa6f32f4 Mon Sep 17 00:00:00 2001 From: igerber Date: Sat, 25 Jul 2026 20:39:12 -0400 Subject: [PATCH 08/11] fix(v4): label RCS counts as obs, document the fit-time deprecation Local review round 4 (no P0/P1; both remaining P2s resolved). n_kind on simple aggregation was hardcoded "units", but fit() only counts units on a panel: with panel=False it counts ROWS, because a repeated cross section has no unit tracking. Reporting that total as "units" misdescribes the sample through the very routing key the shared vocabulary exists to make trustworthy - and it is the conflation N_KIND_VOCABULARY's own documentation warns against. It now reads is_panel from the retained kit, which already carried it. Estimates and inference were never affected; only the label was. The regression pins BOTH directions, so this cannot degrade into a blanket rename: RCS reports "obs" (and the count matches the fitted totals) while the panel fixture still reports "units". The public fit() docstring still presented aggregate= and balance_e= as ordinary parameters. Anyone reading the API reference had no way to know they now emit a FutureWarning and are removed at 4.0. Both are marked deprecated with their ledger rows (M-020, M-117) and point at the post-fit call. The aggregate= entry also records WHY the fit-time path still exists - it is the temporary compatibility surface for compute_honest_did, compute_pretrends_power and plot_event_study - so the docs do not read as recommending a path that is being removed. 64 contract tests; 1200 pass across the staggered, SDDD, survey, methodology, event-study-surface, honest-did, pretrends, ledger and docs suites. Unrelated, NOT introduced here: tests/test_wooldridge.py:: TestMethodologyCorrectness::test_ols_etwfe_att_matches_callaway_santanna fails identically at the pre-branch base cba72583 (ETWFE -0.0431 vs CS -0.0261). It is a pre-existing main-branch parity failure and is left for its own change. --- diff_diff/staggered.py | 18 +++++++++++++++--- diff_diff/staggered_results.py | 13 ++++++++++--- tests/test_aggregate_contract.py | 15 +++++++++++++++ 3 files changed, 40 insertions(+), 6 deletions(-) diff --git a/diff_diff/staggered.py b/diff_diff/staggered.py index 023794e85..e819dfd10 100644 --- a/diff_diff/staggered.py +++ b/diff_diff/staggered.py @@ -1846,15 +1846,27 @@ def fit( covariates : list, optional List of covariate column names for conditional parallel trends. aggregate : str, optional - How to aggregate group-time effects: + DEPRECATED since 3.9, removed in 4.0 (ledger row M-020). Passing + it emits a ``FutureWarning``. Aggregate as a post-fit step + instead - ``results.aggregate("simple" | "event_study" | + "group")`` - which needs no refit and returns a new object, + leaving ``results`` unchanged. Accepted values are unchanged: - None: Only compute ATT(g,t) (default) - "simple": Simple weighted average (overall ATT) - "event_study": Aggregate by relative time (event study) - "group": Aggregate by treatment cohort - "all": Compute all aggregations + + Fit-time aggregation remains ONLY as the temporary compatibility + surface for consumers that still read it (``compute_honest_did``, + ``compute_pretrends_power`` and ``plot_event_study``, tracked in + ``TODO.md``); it is not the recommended path for new code. balance_e : int, optional - For event study, balance the panel at relative time e. - Ensures all groups contribute to each relative period. + DEPRECATED since 3.9, removed in 4.0 (ledger row M-117). Passing + it emits a ``FutureWarning``; it moves onto the post-fit call as + ``results.aggregate("event_study", balance_e=...)``. For event + study, balance the panel at relative time e. Ensures all groups + contribute to each relative period. survey_design : SurveyDesign, optional Survey design specification. Supports pweight with strata/PSU/FPC. Aggregated SEs (overall, event study, group) use design-based diff --git a/diff_diff/staggered_results.py b/diff_diff/staggered_results.py index 924793a7f..991525bb1 100644 --- a/diff_diff/staggered_results.py +++ b/diff_diff/staggered_results.py @@ -313,7 +313,7 @@ def _aggregate_compute( agg = _KitAggregator(kit.alpha, kit.anticipation) if level == "simple": - return self._aggregate_simple_result() + return self._aggregate_simple_result(kit) if level == "group": effects = agg._aggregate_by_group( self.group_time_effects, @@ -348,13 +348,20 @@ def _aggregate_compute( ) return build_event_study_surface(carrier) - def _aggregate_simple_result(self) -> AggregationResult: + def _aggregate_simple_result(self, kit: Any) -> AggregationResult: """The overall ATT as a one-row table. ``_aggregate_simple`` runs unconditionally in ``fit()``, so the numbers are already stored - this is a view, not a recomputation, and is therefore bit-identical to the fit by construction. """ + # n_treated_units / n_control_units are UNITS on a panel fit but + # OBSERVATIONS on a declared repeated cross-section, where fit() counts + # rows because there is no unit tracking (staggered.py, the panel/RCS + # branch). n_kind must say which - conflating the two is exactly what + # the shared vocabulary forbids. + is_panel = kit.bookkeeping.get("is_panel", True) + n_kind = "units" if is_panel else "obs" ci = self.overall_conf_int or (np.nan, np.nan) return AggregationResult( level="simple", @@ -374,7 +381,7 @@ def _aggregate_simple_result(self) -> AggregationResult: # was built on a finite t-reference. df=resolve_inference_df(self), alpha=self.alpha, - n_kind="units", + n_kind=n_kind, weight=np.array([1.0], dtype=float), estimator=type(self).__name__.replace("Results", ""), ) diff --git a/tests/test_aggregate_contract.py b/tests/test_aggregate_contract.py index 038311e61..980e576ef 100644 --- a/tests/test_aggregate_contract.py +++ b/tests/test_aggregate_contract.py @@ -643,6 +643,21 @@ def test_cs_n_kinds_are_in_the_shared_vocabulary(self, fitted, level, expected): assert got.n_kind == expected assert got.n_kind in N_KIND_VOCABULARY + def test_rcs_simple_reports_obs_not_units(self): + """On panel=False, fit() counts ROWS (there is no unit tracking), so + labelling that total "units" would misdescribe the sample - the exact + conflation the shared vocabulary forbids.""" + data = _rcs() + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + res = CallawaySantAnna(panel=False).fit(data, **FIT_KW) + got = res.aggregate("simple") + assert got.n_kind == "obs" + assert got.n[0] == float(res.n_treated_units + res.n_control_units) + # The panel fixture must still say units, or this would be a blanket rename. + panel_res = CallawaySantAnna().fit(_panel(), **FIT_KW) + assert panel_res.aggregate("simple").n_kind == "units" + def test_shared_vocabulary_is_enforced_on_both_containers(self, fitted): """The vocabulary is declared SHARED, so validating it on only one container would let an unknown value reach a consumer through the From 6381b5d017098ff1b62c32e6706a343c5ff04f0a Mon Sep 17 00:00:00 2001 From: igerber Date: Sat, 25 Jul 2026 20:46:13 -0400 Subject: [PATCH 09/11] docs(v4): enumerate every retained buffer in the aggregation memory accounting A resource claim that names only the dominant payload reads as complete when it is not. The accounting now also records the replicate-weight survey matrix (O(n_units x n_replicates)), states that the O(n_units) bookkeeping bound is panel-only because several RCS arrays are observation-length, and records that raw unit identifiers are deliberately not retained. --- docs/v4-design.md | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/docs/v4-design.md b/docs/v4-design.md index 23edc62b2..8198d72ed 100644 --- a/docs/v4-design.md +++ b/docs/v4-design.md @@ -418,9 +418,16 @@ section previously said "CallawaySantAnna already stores them". It does not - and the fit-time `precomputed` bookkeeping is a local. Retention is therefore work in EVERY migrating PR, not just the ones "where missing", and the kit must be built during `fit()` because nothing it needs survives the call. Memory cost -is documented per estimator: for CallawaySantAnna the dominant payload is the -per-(g,t) influence-function dict at roughly O(n_units x n_gt), not the -O(n_units) bookkeeping. Analytical-vs-bootstrap inference of the aggregated +is documented per estimator, enumerating every retained buffer rather than only +the largest: for CallawaySantAnna the dominant payload is the per-(g,t) +influence-function dict at roughly O(n_units x n_gt), on top of the O(n_units) +bookkeeping; a replicate-weight survey design retains a further +O(n_units x n_replicates) matrix through the resolved design object; and on +repeated cross-sections several bookkeeping arrays are observation-length +rather than unit-length, so the O(n_units) figure is a panel-only bound. Raw +unit identifiers are NOT retained - the kit needs only position, so it stores +canonical 0..n-1 codes, keeping a shared results artifact free of names, emails +or administrative IDs. Analytical-vs-bootstrap inference of the aggregated estimand follows the fit's inference method; where bootstrap draws are not retained, `aggregate()` on a bootstrapped fit RAISES rather than silently returning analytical inference. Estimators whose From 09f0ba3001ed899ca0cd8fbd7403537e91343134 Mon Sep 17 00:00:00 2001 From: igerber Date: Sat, 25 Jul 2026 21:52:57 -0400 Subject: [PATCH 10/11] fix(v4): clear group df provenance on bootstrap, fix mypy, document aggregate() CI round 1 (Mypy gate + codex review). P1 - bootstrap group rows kept false analytical df provenance. The bootstrap block replaces each group row's se/p_value/conf_int with percentile values but left the df_used this branch added, so group_effects[g] claimed a t-reference had governed inference that never used one. The event-study path immediately below already clears its df for exactly this reason - the same clearing rule was written for one surface and not its twin, the third instance of that pattern in this branch. Reproduced on a survey bootstrap fit (event_study_df None, group df_used 15.0) before fixing, and fixed in BOTH estimators: CallawaySantAnna and StaggeredTripleDifference, which carries its own copy of the replacement loop. group_effects is a public field, so the wrong value was reachable even though aggregate("group") itself fails closed on bootstrap fits. Regression covers clustered AND explicit-survey bootstrap. Mypy (6 errors, all introduced here; verified against the workflow's exact pins - mypy 2.3.0 / numpy 2.4.5 / pandas 3.0.3 / scipy 1.17.1 - since this worktree's own venv cannot run mypy): - Three "Any | None is not indexable" in staggered_aggregation: relaxing df to Optional left the no-bookkeeping fallback branches indexing a possibly-None frame. Each now raises explicitly when neither the kit nor the frame is available, which both narrows the type and states the invariant. - Two incompatible-base-class errors: the bootstrap mixin's TYPE_CHECKING stub for _compute_combined_influence_function still declared df/unit as required, contradicting the aggregation mixin it shadows wherever both are mixed in. - _aggregation_kit was set dynamically on the results object; it is now a declared field (repr=False, compare=False) and therefore part of the typed contract rather than an undeclared attribute. P2 - the headline method was undocumented. aggregate() is added to the results autosummary in docs/api/staggered.rst and to both agent-facing guides, which still presented fit-time aggregate=/balance_e= with no deprecation context. The guides now lead with the post-fit call and keep the fit-time form only where it is still required (plot_event_study, compute_honest_did, compute_pretrends_power), labelled as the compatibility path. 66 contract tests; 1202 pass across the affected suites; mypy clean. --- diff_diff/guides/llms-full.txt | 32 ++++++++++++++++++---- diff_diff/guides/llms.txt | 2 +- diff_diff/staggered.py | 5 ++++ diff_diff/staggered_aggregation.py | 19 ++++++++++++- diff_diff/staggered_bootstrap.py | 8 ++++-- diff_diff/staggered_results.py | 7 +++++ diff_diff/staggered_triple_diff.py | 4 +++ docs/api/staggered.rst | 1 + tests/test_aggregate_contract.py | 44 ++++++++++++++++++++++++++++++ 9 files changed, 113 insertions(+), 9 deletions(-) diff --git a/diff_diff/guides/llms-full.txt b/diff_diff/guides/llms-full.txt index 5c92b3aa6..6e090bb44 100644 --- a/diff_diff/guides/llms-full.txt +++ b/diff_diff/guides/llms-full.txt @@ -227,9 +227,24 @@ from diff_diff import CallawaySantAnna, plot_event_study cs = CallawaySantAnna(estimation_method="dr", n_bootstrap=999, seed=42) results = cs.fit(data, outcome='outcome', unit='unit', time='period', - first_treat='first_treat', aggregate='event_study') + first_treat='first_treat') results.print_summary() -plot_event_study(results) + +# Aggregate post-fit - no refit. Returns a NEW object; `results` is unchanged. +event_study = results.aggregate('event_study') # -> EventStudyResults +by_cohort = results.aggregate('group') # -> AggregationResult +``` + +Fit-time `aggregate=` / `balance_e=` are DEPRECATED since 3.9 (removed in 4.0, +ledger rows M-020 / M-117) and emit a `FutureWarning`. They remain only as the +compatibility surface for the consumers that still read the fit-time result - +`plot_event_study`, `compute_honest_did` and `compute_pretrends_power` do not +yet accept the post-fit container, so use the fit-time form for those: + +```python +plotted = cs.fit(data, outcome='outcome', unit='unit', time='period', + first_treat='first_treat', aggregate='event_study') +plot_event_study(plotted) ``` ### ChaisemartinDHaultfoeuille @@ -1518,10 +1533,17 @@ Returned by `CallawaySantAnna.fit()`. | `groups` | `list` | Treatment cohorts | | `time_periods` | `list` | All time periods | | `n_obs` | `int` | Number of observations | -| `event_study_effects` | `dict[int, dict]` | Event study effects by relative time | -| `group_effects` | `dict` | Group-level aggregated effects | +| `event_study_effects` | `dict[int, dict]` | Event study effects by relative time. Populated only by the DEPRECATED fit-time `aggregate=`; `None` after a plain fit - use `aggregate("event_study")` instead | +| `group_effects` | `dict` | Group-level aggregated effects. Same: populated only by fit-time `aggregate=` | + +**Methods:** `aggregate(type, weights=None, *, balance_e=None)`, `summary()`, `print_summary()`, `to_dataframe(level="event_study"|"group_time"|"group")` -**Methods:** `summary()`, `print_summary()`, `to_dataframe(level="event_study"|"group_time"|"group")` +`aggregate("simple"|"event_study"|"group")` re-aggregates POST-FIT with no +refit, returning a NEW object and leaving the result unchanged. +`"event_study"` returns `EventStudyResults`; the others return +`AggregationResult`. `balance_e=` applies to `"event_study"` only. Raises on +`"calendar"` (CS has no calendar aggregator) and on a bootstrapped fit, rather +than substituting analytical inference for percentile-bootstrap statistics. ### SunAbrahamResults diff --git a/diff_diff/guides/llms.txt b/diff_diff/guides/llms.txt index cc9ae69fd..1465a94dc 100644 --- a/diff_diff/guides/llms.txt +++ b/diff_diff/guides/llms.txt @@ -21,7 +21,7 @@ diagnostic steps produces unreliable results. 4. **Choose estimator** — staggered adoption → CS/SA/BJS (NOT plain TWFE); few treated units → SDiD; factor confounding → TROP; simple 2x2 → DiD. Run `BaconDecomposition` to diagnose TWFE bias. 5. **Estimate** — `estimator.fit(data, ...)`. Always print the cluster count first and choose inference method based on the result (cluster-robust if >= 50 clusters, wild bootstrap if fewer). 6. **Sensitivity analysis** — `compute_honest_did(results)` for bounds under PT violations (MultiPeriodDiD, CS, or dCDH), `run_all_placebo_tests()` for 2x2 falsification, specification comparisons for staggered designs. -7. **Heterogeneity** — CS: `aggregate='group'`/`'event_study'`; SA: `results.event_study_effects`/`to_dataframe(level='cohort')`; subgroup re-estimation. +7. **Heterogeneity** — CS: `results.aggregate('group')`/`.aggregate('event_study')` post-fit, no refit (fit-time `aggregate=`/`balance_e=` are deprecated since 3.9, removed in 4.0; they remain only for `compute_honest_did` / `compute_pretrends_power` / `plot_event_study`, which still read the fit-time surface); SA: `results.event_study_effects`/`to_dataframe(level='cohort')`; subgroup re-estimation. 8. **Robustness** — compare 2-3 estimators (CS vs SA vs BJS), MUST report with and without covariates (shows whether conditioning drives identification), present pre-trends and sensitivity bounds. After estimation, call `practitioner_next_steps(results)` for context-aware diff --git a/diff_diff/staggered.py b/diff_diff/staggered.py index e819dfd10..82ff94ee1 100644 --- a/diff_diff/staggered.py +++ b/diff_diff/staggered.py @@ -2843,6 +2843,11 @@ def fit( group_effects[g]["conf_int"] = bootstrap_results.group_effect_cis[g] group_effects[g]["p_value"] = bootstrap_results.group_effect_p_values[g] group_effects[g]["t_stat"] = float(grp_t_stats[idx]) + # Same clearing rule the ES df provenance follows below: + # these se/p/CI are now percentile-bootstrap values that + # never used the analytical df, so keeping df_used would + # claim a t-reference that governed nothing. + group_effects[g]["df_used"] = None # Compute simultaneous confidence band CIs if cband is available cband_crit_value = None diff --git a/diff_diff/staggered_aggregation.py b/diff_diff/staggered_aggregation.py index e34b8940b..e37fa4baf 100644 --- a/diff_diff/staggered_aggregation.py +++ b/diff_diff/staggered_aggregation.py @@ -549,7 +549,14 @@ def _compute_combined_influence_function( total_weight = float(n_units) else: # No precomputed bookkeeping (direct internal callers only): fall - # back to the fit-time frame. + # back to the fit-time frame. Reaching here without one is the + # fail-closed case - post-fit callers always carry the kit, so a + # None frame here means neither source is available. + if df is None or unit is None: + raise ValueError( + "Cohort sizes need either precomputed bookkeeping or the fit-time " + "frame; neither was supplied." + ) for g in unique_groups: treated_in_g = df[df["first_treat"] == g][unit].nunique() group_sizes[g] = treated_in_g @@ -611,6 +618,11 @@ def _compute_combined_influence_function( if cohort in unique_groups_set: unit_groups_array[idx] = cohort else: + if df is None or unit is None: + raise ValueError( + "Per-unit cohorts need either precomputed bookkeeping or the " + "fit-time frame; neither was supplied." + ) for idx, uid in idx_uid_pairs: unit_first_treat = df[df[unit] == uid]["first_treat"].iloc[0] if unit_first_treat in unique_groups_set: @@ -628,6 +640,11 @@ def _compute_combined_influence_function( if cohort in unique_groups_set: unit_groups_array[idx] = cohort else: + if df is None or unit is None: + raise ValueError( + "Per-unit cohorts need either precomputed bookkeeping or the " + "fit-time frame; neither was supplied." + ) for idx, uid in idx_uid_pairs: unit_first_treat = df[df[unit] == uid]["first_treat"].iloc[0] if unit_first_treat in unique_groups_set: diff --git a/diff_diff/staggered_bootstrap.py b/diff_diff/staggered_bootstrap.py index 326e950d7..180f65a7e 100644 --- a/diff_diff/staggered_bootstrap.py +++ b/diff_diff/staggered_bootstrap.py @@ -143,8 +143,12 @@ def _compute_combined_influence_function( effects: np.ndarray, groups_for_gt: np.ndarray, influence_func_info: Dict, - df: "pd.DataFrame", - unit: str, + # Optional to match the aggregation mixin, whose implementation this + # stub shadows: post-fit aggregation runs from the retained kit with + # no frame. A narrower stub here makes the two base classes + # incompatible wherever both are mixed in. + df: Optional["pd.DataFrame"], + unit: Optional[str], precomputed: Optional["PrecomputedData"] = None, global_unit_to_idx: Optional[Dict[Any, int]] = None, n_global_units: Optional[int] = None, diff --git a/diff_diff/staggered_results.py b/diff_diff/staggered_results.py index 991525bb1..66a523fad 100644 --- a/diff_diff/staggered_results.py +++ b/diff_diff/staggered_results.py @@ -234,6 +234,13 @@ class CallawaySantAnnaResults(BaseResults, AggregationMixin): # with the original resolved_survey.df_survey. df_inference: Optional[int] = None + # Post-fit re-aggregation payload (spec section 6, rows M-020/M-117), + # attached by fit() because nothing it needs survives the call. Declared + # here rather than set dynamically so it is a typed part of the contract. + # Excluded from repr and equality: it is internal bookkeeping, not a + # reportable result, and its arrays would make `==` raise. + _aggregation_kit: Optional[Any] = field(default=None, repr=False, compare=False) + # --- Inference-field aliases (balance/external-adapter compatibility) --- @property def att(self) -> float: diff --git a/diff_diff/staggered_triple_diff.py b/diff_diff/staggered_triple_diff.py index df39a542b..0508e78c5 100644 --- a/diff_diff/staggered_triple_diff.py +++ b/diff_diff/staggered_triple_diff.py @@ -817,6 +817,10 @@ def fit( df=df_survey, ) group_effects[g_key]["t_stat"] = t_val + # Bootstrap se/p/CI replaced the analytical ones, which + # is what the retained df described - keeping it would + # claim a t-reference that governed nothing. + group_effects[g_key]["df_used"] = None # Eq. 4.14 overall: an ANALYTICAL non-finite SE under a finite point estimate # (a contributing horizon's influence function is non-finite, or the variance is diff --git a/docs/api/staggered.rst b/docs/api/staggered.rst index 1ceb0e150..8bf5fbd04 100644 --- a/docs/api/staggered.rst +++ b/docs/api/staggered.rst @@ -48,6 +48,7 @@ Results container for Callaway-Sant'Anna estimation. .. autosummary:: + ~CallawaySantAnnaResults.aggregate ~CallawaySantAnnaResults.summary ~CallawaySantAnnaResults.to_dataframe diff --git a/tests/test_aggregate_contract.py b/tests/test_aggregate_contract.py index 980e576ef..828a89426 100644 --- a/tests/test_aggregate_contract.py +++ b/tests/test_aggregate_contract.py @@ -572,6 +572,50 @@ def test_explicit_survey_df_reaches_every_level(self, survey_fit): assert finite.size > 0, f"{level} reported no df at all" assert np.all(finite == expected), f"{level} df {finite} != {expected}" + @pytest.mark.parametrize("survey", [False, True]) + def test_bootstrap_clears_group_df_provenance(self, survey): + """When bootstrap replaces a group row's se/p/CI with percentile values, + the retained analytical df described inference that no longer exists. + Leaving it finite would claim a t-reference governed percentile-bootstrap + numbers - the same false provenance the event-study path already clears. + """ + from diff_diff import SurveyDesign + + rng = np.random.default_rng(7) + rows = [] + for u in range(90): + g = [0, 4, 6][u % 3] + for t in range(1, 9): + treated = g != 0 and t >= g + rows.append( + { + "unit": u, + "time": t, + "first_treat": g, + "psu": u % 18, + "stratum": u % 3, + "w": 1.0 + (u % 4) * 0.1, + "y": 0.3 * t + (2.0 if treated else 0.0) + rng.normal(0, 0.5), + } + ) + # cluster= is a CONSTRUCTOR argument; survey_design= is a fit() argument. + ctor = {} if survey else {"cluster": "psu"} + fit_extra = ( + {"survey_design": SurveyDesign(weights="w", strata="stratum", psu="psu")} + if survey + else {} + ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = CallawaySantAnna(n_bootstrap=50, seed=1, **ctor).fit( + pd.DataFrame(rows), aggregate="all", **fit_extra, **FIT_KW + ) + assert res.group_effects, "fixture produced no group rows" + for g, eff in res.group_effects.items(): + assert eff.get("df_used") is None, f"group {g} kept analytical df on a bootstrap fit" + # The event-study path's equivalent clearing must not have regressed. + assert res.event_study_df is None + def test_group_df_is_not_the_event_study_df(self, fitted): """``group`` must not borrow ``event_study_df``: on a plain fit that field is None, and it is a different aggregation's denominator.""" From 25e1ade186059da41558917c81ee658ab5268e86 Mon Sep 17 00:00:00 2001 From: igerber Date: Sun, 26 Jul 2026 05:40:10 -0400 Subject: [PATCH 11/11] docs: fix a guide example that violates the contract it documents; SDDD provenance regression CI round 2 (rerun: no unmitigated P0/P1; both remaining P2s resolved). The post-fit example added last round constructed the estimator with n_bootstrap=999 and then called results.aggregate() - which raises NotImplementedError on a bootstrapped fit, by the very contract documented a few lines below it. Reproduced before fixing. The analytical example now omits n_bootstrap and says why, and the fit-time form is presented as the required path for BOTH cases that need it: bootstrap inference, and the downstream consumers that still read the fit-time surface. SDDD regression for the bootstrap df-provenance fix. The previous round fixed both estimators but tested only CallawaySantAnna - the same fix-one-twin gap that produced the bug. Verified the new test is not vacuous by reverting the fix and confirming it fails. That check also exposed an asymmetry now recorded in the test rather than papered over: only the explicit-survey configuration exercises the clearing. SDDD's plain-cluster path never resolves a finite analytical df, so df_used is already None there and the non-survey parametrization passes either way; it is kept as a cheap guard against that changing, not as a reproduction. The CallawaySantAnna equivalent was checked the same way and DOES fail in both configurations, so its coverage claim stands. 68 contract tests; 1204 pass across the affected suites; mypy clean against the workflow's pins. --- diff_diff/guides/llms-full.txt | 22 ++++++++---- tests/test_aggregate_contract.py | 58 ++++++++++++++++++++++++++++++++ 2 files changed, 73 insertions(+), 7 deletions(-) diff --git a/diff_diff/guides/llms-full.txt b/diff_diff/guides/llms-full.txt index 6e090bb44..da6452c80 100644 --- a/diff_diff/guides/llms-full.txt +++ b/diff_diff/guides/llms-full.txt @@ -225,7 +225,9 @@ cs.fit( ```python from diff_diff import CallawaySantAnna, plot_event_study -cs = CallawaySantAnna(estimation_method="dr", n_bootstrap=999, seed=42) +# NOTE: no n_bootstrap here - post-fit aggregate() is ANALYTICAL-ONLY and +# raises NotImplementedError on a bootstrapped fit (see below). +cs = CallawaySantAnna(estimation_method="dr") results = cs.fit(data, outcome='outcome', unit='unit', time='period', first_treat='first_treat') results.print_summary() @@ -236,14 +238,20 @@ by_cohort = results.aggregate('group') # -> AggregationResult ``` Fit-time `aggregate=` / `balance_e=` are DEPRECATED since 3.9 (removed in 4.0, -ledger rows M-020 / M-117) and emit a `FutureWarning`. They remain only as the -compatibility surface for the consumers that still read the fit-time result - -`plot_event_study`, `compute_honest_did` and `compute_pretrends_power` do not -yet accept the post-fit container, so use the fit-time form for those: +ledger rows M-020 / M-117) and emit a `FutureWarning`. They remain the required +path in two cases: + +1. **Bootstrap inference.** `aggregate()` raises on a bootstrapped fit rather + than substituting analytical inference for percentile-bootstrap statistics. +2. **Downstream consumers.** `plot_event_study`, `compute_honest_did` and + `compute_pretrends_power` still read the fit-time surface and do not yet + accept the post-fit container. ```python -plotted = cs.fit(data, outcome='outcome', unit='unit', time='period', - first_treat='first_treat', aggregate='event_study') +# Both cases: aggregate at fit time. +boot = CallawaySantAnna(estimation_method="dr", n_bootstrap=999, seed=42) +plotted = boot.fit(data, outcome='outcome', unit='unit', time='period', + first_treat='first_treat', aggregate='event_study') plot_event_study(plotted) ``` diff --git a/tests/test_aggregate_contract.py b/tests/test_aggregate_contract.py index 828a89426..1594bf666 100644 --- a/tests/test_aggregate_contract.py +++ b/tests/test_aggregate_contract.py @@ -616,6 +616,64 @@ def test_bootstrap_clears_group_df_provenance(self, survey): # The event-study path's equivalent clearing must not have regressed. assert res.event_study_df is None + @pytest.mark.parametrize("survey", [False, True]) + def test_sddd_bootstrap_clears_group_df_provenance(self, survey): + """StaggeredTripleDifference carries its OWN copy of the bootstrap + group-replacement loop, so the clearing rule has to hold there too - + fixing one twin and testing only the other is how this class of bug + survived in the first place. + + Coverage is ASYMMETRIC and deliberately so: only ``survey=True`` + exercises the clearing (verified by reverting the fix - just that case + fails). SDDD's plain-``cluster`` path never resolves a finite analytical + df, so ``df_used`` is already None there and nothing needs clearing. The + non-survey case is kept as a cheap guard in case that ever changes, not + because it currently reproduces the bug. The CallawaySantAnna + equivalent above DOES fail in both configurations. + """ + from diff_diff import SurveyDesign + + rng = np.random.default_rng(9) + rows = [] + for u in range(96): + g = [0, 3, 4][u % 3] + elig = u % 2 + for t in range(1, 7): + treated = g != 0 and t >= g and elig == 1 + rows.append( + { + "unit": u, + "period": t, + "first_treat": g, + "eligibility": elig, + "psu": u % 16, + "stratum": u % 4, + "w": 1.0 + (u % 3) * 0.1, + "outcome": 0.3 * t + (1.2 if treated else 0.0) + rng.normal(0, 0.2), + } + ) + ctor = {} if survey else {"cluster": "psu"} + fit_extra = ( + {"survey_design": SurveyDesign(weights="w", strata="stratum", psu="psu")} + if survey + else {} + ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = StaggeredTripleDifference(n_bootstrap=40, seed=3, **ctor).fit( + pd.DataFrame(rows), + "outcome", + "unit", + "period", + "first_treat", + "eligibility", + aggregate="all", + **fit_extra, + ) + assert res.group_effects, "fixture produced no SDDD group rows" + for g, eff in res.group_effects.items(): + assert eff.get("df_used") is None, f"SDDD group {g} kept analytical df on bootstrap" + def test_group_df_is_not_the_event_study_df(self, fitted): """``group`` must not borrow ``event_study_df``: on a plain fit that field is None, and it is a different aggregation's denominator."""