diff --git a/CHANGELOG.md b/CHANGELOG.md index 38516d22..72e8962a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,47 @@ 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:** 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` 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 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. + `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 be7532e3..cb69a528 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/TODO.md b/TODO.md index 31ec3697..dd834799 100644 --- a/TODO.md +++ b/TODO.md @@ -22,6 +22,9 @@ 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 | ### Performance diff --git a/diff_diff/__init__.py b/diff_diff/__init__.py index fc8f8a15..139e9d22 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 00000000..6340c995 --- /dev/null +++ b/diff_diff/aggregation.py @@ -0,0 +1,538 @@ +"""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 N_KIND_VOCABULARY, 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 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? + + 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 + :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 - + 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) + # 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) + 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: + # 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." + ) + 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 + + # 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) + # ------------------------------------------------------------------ # + + 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. + + 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})", + "=" * 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={self.alpha}.") + 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/guides/llms-full.txt b/diff_diff/guides/llms-full.txt index 5c92b3aa..da6452c8 100644 --- a/diff_diff/guides/llms-full.txt +++ b/diff_diff/guides/llms-full.txt @@ -225,11 +225,34 @@ 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', 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 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 +# 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) ``` ### ChaisemartinDHaultfoeuille @@ -1518,10 +1541,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 cc9ae69f..1465a94d 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/honest_did.py b/diff_diff/honest_did.py index 3e2edc4a..10b376a9 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. " - "Re-run CallawaySantAnna.fit() with aggregate='event_study' to compute " - "event study effects." + "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 11a3d44f..79c94c21 100644 --- a/diff_diff/pretrends.py +++ b/diff_diff/pretrends.py @@ -1140,7 +1140,9 @@ 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'." + "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 2c5a308c..2ad7586e 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 @@ -351,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): @@ -515,7 +542,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/diff_diff/staggered.py b/diff_diff/staggered.py index 397b9012..82ff94ee 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, @@ -456,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): @@ -1795,8 +1821,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: """ @@ -1820,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 @@ -1847,15 +1885,44 @@ 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}") - # 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 +2714,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 +2728,7 @@ def fit( unit, precomputed, ) + event_study_effects = es_aggregation.effects if aggregate in ["group", "all"]: group_effects = self._aggregate_by_group( @@ -2772,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 @@ -2814,15 +2890,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 @@ -2904,6 +2980,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_ @@ -5000,3 +5085,77 @@ 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] + + # 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, + 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 871174c5..e37fa4ba 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]]: @@ -67,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]]: """ @@ -371,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, @@ -502,7 +538,25 @@ 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. 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 @@ -564,16 +618,37 @@ 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: 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: + 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: + unit_groups_array[idx] = unit_first_treat # Vectorized WIF computation groups_for_gt_array = np.array(groups_for_gt) @@ -652,8 +727,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]]: ... @@ -666,8 +741,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], @@ -680,8 +755,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]]]": @@ -815,7 +890,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). @@ -944,8 +1019,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, @@ -969,13 +1051,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 +1081,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 +1156,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 +1190,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, @@ -1198,6 +1285,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] = { @@ -1207,6 +1305,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_bootstrap.py b/diff_diff/staggered_bootstrap.py index 326e950d..180f65a7 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 98357d7b..66a523fa 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, resolve_inference_df 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. @@ -213,6 +234,13 @@ class CallawaySantAnnaResults(BaseResults): # 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: @@ -253,6 +281,153 @@ 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(kit) + 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, 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", + 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), + # 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=n_kind, + 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), + # 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, + estimator=type(self).__name__.replace("Results", ""), + ) + def summary(self, alpha: Optional[float] = None) -> str: """ Generate formatted summary of estimation results. @@ -504,7 +679,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)) @@ -525,7 +704,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/diff_diff/staggered_triple_diff.py b/diff_diff/staggered_triple_diff.py index a838ce69..0508e78c 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"] @@ -814,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/index.rst b/docs/api/index.rst index fd6d32ae..13ddf823 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 36c4c046..593372a8 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/api/staggered.rst b/docs/api/staggered.rst index 1ceb0e15..8bf5fbd0 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/docs/doc-deps.yaml b/docs/doc-deps.yaml index 823fb87e..448745db 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: diff --git a/docs/v4-deprecations.yaml b/docs/v4-deprecations.yaml index 8acdeed7..6600b866 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 cf38b98e..8198d72e 100644 --- a/docs/v4-design.md +++ b/docs/v4-design.md @@ -412,10 +412,25 @@ 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, 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 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 00000000..1594bf66 --- /dev/null +++ b/tests/test_aggregate_contract.py @@ -0,0 +1,837 @@ +"""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) + + 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 +# --------------------------------------------------------------------------- # + + +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 + + +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) + + @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.""" + 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" + + +def _survey_panel(): + 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), + } + ) + 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(_survey_panel(), survey_design=_survey_design(), **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}" + + @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 + + @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.""" + 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 + + 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 + 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 + 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 +# --------------------------------------------------------------------------- # + + +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 35aba6f6..5d504a6f 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():