diff --git a/.claude/scripts/premerge_scan.py b/.claude/scripts/premerge_scan.py index 1073e5c0..199d96fc 100644 --- a/.claude/scripts/premerge_scan.py +++ b/.claude/scripts/premerge_scan.py @@ -111,17 +111,24 @@ def _init_self_assigns(cls: ast.ClassDef) -> "set[str]": return names -def _get_params_refs(cls: ast.ClassDef) -> "set[str]": - refs: "set[str]" = set() +def _own_get_params(cls: ast.ClassDef) -> "ast.FunctionDef | None": for node in cls.body: if isinstance(node, ast.FunctionDef) and node.name == "get_params": - for sub in ast.walk(node): - if isinstance(sub, ast.Attribute): - refs.add(sub.attr) - elif isinstance(sub, ast.Name): - refs.add(sub.id) - elif isinstance(sub, ast.Constant) and isinstance(sub.value, str): - refs.add(sub.value) + return node + return None + + +def _get_params_refs(cls: ast.ClassDef) -> "set[str]": + refs: "set[str]" = set() + node = _own_get_params(cls) + if node is not None: + for sub in ast.walk(node): + if isinstance(sub, ast.Attribute): + refs.add(sub.attr) + elif isinstance(sub, ast.Name): + refs.add(sub.id) + elif isinstance(sub, ast.Constant) and isinstance(sub.value, str): + refs.add(sub.value) return refs @@ -130,8 +137,14 @@ def new_params_missing_from_get_params( ) -> "list[str]": """Check C (AST-based): of the ``added_param_names`` (new ``self.X`` from the diff), return those that are assigned in some class ``__init__`` but never referenced in - that same class's ``get_params()``. Restricting to ``__init__`` and to the class's - own ``get_params`` avoids the old substring heuristic's false negatives.""" + that same class's OWN in-body ``get_params()``. Classes WITHOUT an in-body + ``get_params`` are skipped entirely: post-mixin, estimators inherit an + introspective ``get_params`` from ``diff_diff._base.BaseEstimator`` (possibly + indirectly, e.g. SyntheticDiD via DifferenceInDifferences - invisible to this + file-local AST scan), which derives its keys from the ``__init__`` signature, + so a new constructor param is auto-covered and the init-signature<->get_params + sync is enforced by ``tests/test_base_estimator.py``. The check remains for + any class that still hand-rolls ``get_params``.""" try: tree = ast.parse(file_text) except SyntaxError: @@ -139,6 +152,8 @@ def new_params_missing_from_get_params( added = set(added_param_names) missing: "set[str]" = set() for cls in (n for n in ast.walk(tree) if isinstance(n, ast.ClassDef)): + if _own_get_params(cls) is None: + continue assigned = _init_self_assigns(cls) & added if not assigned: continue diff --git a/CHANGELOG.md b/CHANGELOG.md index 8a1c86e3..47f4a557 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,42 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Changed +- **Shared `BaseEstimator` mixin** (`diff_diff/_base.py`, v4 program 2(c)-i): + the 25 hand-rolled `get_params`/`set_params` pairs (6 divergent + implementation shapes) are replaced by one introspective mixin. + `get_params` derives its keys from the `__init__` signature and uniformly + accepts `deep=` (previously 4/25); `set_params` is TRANSACTIONAL via probe + re-init - it validates by constructing `type(self)(**merged)` before any + mutation, so it enforces exactly the constructor's validation, eagerly and + atomically, on every estimator. Behavior changes, all in the + fail-loudly direction: (1) ten estimators whose `set_params` previously + accepted constructor-rejected values until `fit()` (CallawaySantAnna, + TripleDifference, TwoStageDiD, ImputationDiD, SunAbraham, StackedDiD, + StaggeredTripleDifference, SpilloverDiD, TROP, PreTrendsPower) now raise at + `set_params`; (2) `BaconDecomposition.set_params` no longer silently + ignores unknown keys; (3) unknown-key rejection is uniformly strict + (method names / private attributes raise `ValueError: Unknown parameter: + ...`) where 17 classes previously accepted any existing attribute name; + (4) `HonestDiD` gains rollback (a failed call previously left the instance + part-mutated) and `SpilloverDiD` gains batch atomicity; (5) + `TROP.set_params(lambda_*_grid=[])` now resolves the empty grid to the + constructor's default grid (matching `__init__`; the raw-`[]` state was + reachable only through the old setter); (6) list-valued `set_params` + updates to `ContinuousDiD.dvals` / `PreTrendsPower.violation_weights` now + store the constructor-normalized ndarray; (7) `SyntheticDiD.get_params()` + now mirrors its full `__init__` signature - the deprecated + `lambda_reg`/`zeta` appear as always-`None` keys and round-trip silently; + (8) `get_params` key order follows the `__init__` signature everywhere. + Estimator `get_params`/`set_params` signatures are sklearn-compatible + library-wide, so `sklearn.base.clone` round-trips (where scikit-learn is + installed; it is deliberately not a dependency) - except configured + `ChaisemartinDHaultfoeuille` instances, whose pre-existing + `paths_of_interest` canonicalization fails clone's parameter-identity + check (recorded in DEFERRED.md). The cross-estimator contract is pinned by + `tests/test_base_estimator.py` on a dynamically discovered roster; the + pre-merge scanner's Check C now covers hand-rolled `get_params` only. + ### Added - **Naming-completeness guard test** (`tests/test_naming_guard.py`, internal): mechanical enforcement of the three v4-program naming invariants that were diff --git a/CLAUDE.md b/CLAUDE.md index 04e246c2..3a8a0c4c 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -83,23 +83,24 @@ pytest tests/test_rust_backend.py -v together using `safe_inference()` from `diff_diff.utils`. Never compute individually. 7. **Estimator inheritance** — understanding this prevents consistency bugs: ``` - DifferenceInDifferences (base class) - ├── TwoWayFixedEffects (inherits get_params/set_params) - └── MultiPeriodDiD (inherits get_params/set_params) - - Standalone estimators (each has own get_params/set_params): - ├── CallawaySantAnna - ├── SunAbraham - ├── ImputationDiD - ├── TwoStageDiD - ├── TripleDifference - ├── TROP - ├── StackedDiD - ├── SyntheticDiD - └── BaconDecomposition + BaseEstimator (diff_diff/_base.py — shared get_params/set_params mixin) + ├── DifferenceInDifferences + │ ├── TwoWayFixedEffects + │ ├── MultiPeriodDiD + │ └── SyntheticDiD + └── every other estimator class (CallawaySantAnna, SunAbraham, + ImputationDiD, TwoStageDiD, TripleDifference, TROP, StackedDiD, + BaconDecomposition, ... — 25 classes total) ``` - When adding params to `DifferenceInDifferences.get_params()`, subclasses inherit automatically. - Standalone estimators must be updated individually. + `get_params` introspects the `__init__` signature and `set_params` is + TRANSACTIONAL via probe re-init (`type(self)(**merged)` validates before + any mutation) — so a new `__init__` param is automatically in + `get_params`/`set_params` for every class, and `set_params` enforces + exactly the constructor's validation, eagerly. Per-class accommodations + are declarative class attrs (`_PARAM_ATTR_ALIASES`, + `_DERIVED_CONFIG_ATTRS`, `_normalize_set_params`); the cross-estimator + contract is pinned by `tests/test_base_estimator.py` (dynamic roster — + new estimators are automatically enrolled and must mix BaseEstimator in). 8. **Dependencies**: numpy, pandas, and scipy ONLY. No statsmodels. ## Documenting Deviations (AI Review Compatibility) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 379a0962..39f642bc 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -165,13 +165,24 @@ grep -rn "self\." diff_diff/.py A new parameter is only complete when it is: -- stored on `self` and returned by `get_params()` (and accepted by `set_params()`) +- stored on `self` under its own name — `get_params()`/`set_params()` then + cover it automatically (the shared `BaseEstimator` mixin introspects the + `__init__` signature; a param stored under a DIFFERENT attribute name needs + a `_PARAM_ATTR_ALIASES` entry, and constructor-derived state a + `_DERIVED_CONFIG_ATTRS` entry — see `diff_diff/_base.py` and the + inheritance map in `CLAUDE.md`) - applied in **every** aggregation mode — `simple`, `event_study`, and `group` - applied in the **bootstrap/inference** paths, not just the analytical one - reflected on the result object, so `to_dict()`/`summary()` do not misreport it -- propagated to the estimators that inherit it (see the inheritance map in - `CLAUDE.md` — subclasses of `DifferenceInDifferences` inherit automatically, - standalone estimators must each be updated) +- propagated to the estimators that inherit it: `TwoWayFixedEffects` and + `MultiPeriodDiD` define no `__init__` of their own, so they inherit a new + `DifferenceInDifferences` constructor parameter automatically; + `SyntheticDiD` defines its OWN signature (it forwards only + `robust`/`cluster`/`alpha` to `super().__init__`), so a new parent + parameter must be explicitly mirrored — or explicitly rejected, like its + `vcov_type`/`conley_*` guards — there, and `BaseEstimator` introspects the + concrete signature, so an unmirrored parent param simply won't appear in + `SyntheticDiD.get_params()` The recurring bug is a parameter that works in the default aggregation and is silently ignored in one of the others. The same trace applies when adding a new diff --git a/DEFERRED.md b/DEFERRED.md index 9b6976de..864a43d1 100644 --- a/DEFERRED.md +++ b/DEFERRED.md @@ -124,6 +124,8 @@ decisions (refactor waivers, perf trade-offs, test-infrastructure calls) are rec | Decision | Location | Verified | |----------|----------|----------| +| **DCDH `sklearn.base.clone` param-identity failure won't-fix.** `ChaisemartinDHaultfoeuille._validate_paths_of_interest` unconditionally canonicalizes `paths_of_interest` into a fresh `List[Tuple[int, ...]]`, so sklearn `clone()`'s post-construction `param1 is param2` identity check fails for configured instances - a pre-existing normalization the BaseEstimator mixin PR documented rather than changed (get_params/set_params signatures are clone-compatible; the dependency-free `cls(**est.get_params())` config-equality contract is the enforced one, `tests/test_base_estimator.py`). Fixing would mean returning the caller's raw object from a validator whose job is canonicalization. | `chaisemartin_dhaultfoeuille.py` | mixin PR / 2026-08-01 | +| **scikit-learn stays out of dev deps; clone-identity tests remain importorskip-only.** The sklearn-`clone()` round-trip tests (`test_base_estimator.py`, had/rdd/cic suites) run only where scikit-learn happens to be installed - deliberate, matching the numpy/pandas/scipy-only dependency posture; the always-running contract is the dependency-free re-instantiation config-equality test. | `tests/test_base_estimator.py` | mixin PR / 2026-08-01 | | **Plan-review hash gate threat model: accident prevention, NOT adversarial defense.** The ExitPlanMode content-hash gate (hook + `plan_snapshot.py`) exists to stop accidents — stale approvals, concurrent-worktree cross-talk, plans edited mid-review — all of which it closes by construction (snapshot identity + invocation-unique state tokens, 30+ behavioral tests). It does NOT and cannot defend against a malicious local process: nothing verifies review AUTHORSHIP, and such an actor can simply write a matching review file directly — no userland hook can prevent that short of signed reviews, which is out of scope for a personal workflow aid. Review findings that presuppose a hostile local actor against this gate are waived by this decision (2026-07-20, after 7 local AI-review rounds converged on ever-deeper adversarial-model refinements with no reachable fixpoint). Genuine accident vectors remain in scope and are fixed as found. | `.claude/hooks/check-plan-review.py`, `.claude/scripts/plan_snapshot.py` | 2026-07-20 | | **`StackedDiD` survey re-resolution intra-file dedup not warranted.** The cross-estimator ContinuousDiD/EfficientDiD panel-to-unit collapse consolidation landed in #226 (`ResolvedSurveyDesign.subset_to_units_by_row_idx` / `build_unit_first_row_index`); StackedDiD deliberately re-resolves at stacked granularity (control units are duplicated across sub-experiments). The residual intra-file dedup (raw-weight extraction ×3, compose-normalize ×3, resolve-on-stacked ×2) is stacked-specific, low value, and touches the numerically-sensitive composed-weight renormalization; cross-estimator unification was assessed and is not warranted (genuinely different mechanisms; all three already delegate to `_resolve_survey_for_fit` / `compute_survey_metadata`). | `stacked_did.py` | #226 review | | **`HeterogeneousAdoptionDiD` Phase-4 Pierce-Schott (2016) replication harness waived.** R parity at `atol=1e-8` on the same 3 DGPs is a strictly stronger anchor than reproducing Fig 2's pointwise CIs on the LBD-restricted PNTR panel (paper §5.2 self-acknowledges NP estimators too noisy there). Re-open only on user demand. See REGISTRY HAD Deviations Notes #3/#4. | `benchmarks/`, `tests/` | Phase 2a / 2026-05-20 | diff --git a/TODO.md b/TODO.md index 923ac752..63b68b44 100644 --- a/TODO.md +++ b/TODO.md @@ -58,6 +58,7 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m | Issue | Location | Origin | Effort | Priority | |-------|----------|--------|--------|----------| +| Evaluate adding the `BaseEstimator` param surface (get_params/set_params) to the exported classes that never had it - `PowerAnalysis`, `LinearRegression`, `BusinessReport`, `DiagnosticReport`, `TWFEWeightsResult` (a NEW public surface, deliberately out of the 2(c)-i pure-refactor scope; `LinearRegression` is the one `fit`-bearing class excluded from the contract suite's roster-completeness test). | `diff_diff/linalg.py`, `diff_diff/power.py` | mixin PR | Mid | Low | | Tighten the mypy suppressions that back the enforced-zero posture: burn down `prep_dgp`'s per-module `[index]` override (needs a None-vs-array restructure that preserves the seeded RNG stream), and evaluate re-enabling the globally disabled codes (`arg-type`, `return-value`, `var-annotated`, `assignment`) one at a time — `assignment` alone hid several real annotation drifts found during the 2026-07 triage. | `pyproject.toml` `[tool.mypy]`, `diff_diff/prep_dgp.py` | lint-CI | Mid | Low | | MMM interop follow-up: Meridian `roi_calibration_period` mask builder - accept the MMM's time index + channel order and emit the boolean `(n_media_times, n_media_channels)` mask so `.to_code()` scopes the prior to the experiment window automatically (today the caller passes a mask expression / `full_model_window=True`). | `diff_diff/mmm.py` | mmm-interop | Quick | Low | | MMM interop PR-B: calibration tutorial notebook (fit DiD/CS -> scope -> `to_pymc_marketing_lift_test` / `to_meridian_roi_prior`) + a `llms-practitioner.txt` Step 8 pointer to the exporters as the MMM hand-off. | `docs/tutorials/`, `diff_diff/guides/llms-practitioner.txt` | mmm-interop | Mid | Low | diff --git a/diff_diff/_base.py b/diff_diff/_base.py new file mode 100644 index 00000000..3289f4f4 --- /dev/null +++ b/diff_diff/_base.py @@ -0,0 +1,169 @@ +"""Shared estimator parameter surface (4.0 program, Phase 2 PR 2(c)-i). + +This module is a deliberate LEAF: it imports only the standard library, never +other ``diff_diff`` modules, so every estimator module can import it without +cycles (the ``results_base.py`` precedent). + +One public symbol: + +``BaseEstimator`` + Mixin providing the sklearn-compatible ``get_params`` / ``set_params`` + pair for every estimator class, replacing the 25 hand-rolled copies + (see ``docs/v4-design.md`` section 7, "Constructor hygiene"). + +Design contract (the locked transactional rule): + +- Parameter names are introspected from the concrete class's ``__init__`` + signature, so the parameter surface can never drift from the constructor. +- ``get_params`` returns ``{name: getattr(self, alias_or_name)}`` for every + constructor parameter. ``deep`` is accepted for sklearn compatibility + (``sklearn.base.clone`` calls ``get_params(deep=False)``) and ignored: + no diff-diff estimator nests another estimator. +- ``set_params`` is TRANSACTIONAL via probe re-init: it validates the merged + configuration by constructing a throwaway ``type(self)(**merged)`` - which + raises before ``self`` is touched - then adopts the probe's parameter + attributes plus each class's declared derived-config attributes. Fitted + state (attributes ``fit()`` sets) is never touched, because only declared + attributes are copied. The behavioral contract, stated post-normalization: + + est.set_params(**p) == config state of + type(est)(**{**est.get_params(), **cls._normalize_set_params(p)}) + + Validation is therefore exactly ``__init__``'s validation, eagerly, and + can never drift from it. + +Per-class accommodation hooks (declarative; default no-op): + +``_PARAM_ATTR_ALIASES`` + Maps a constructor parameter name to the attribute that stores its RAW + value when the two differ (e.g. ``DifferenceInDifferences`` stores the + raw ``vcov_type`` argument under ``_vcov_type_arg`` and the resolved + value under ``vcov_type``). +``_DERIVED_CONFIG_ATTRS`` + Extra attributes ``__init__`` computes FROM the parameters that must be + re-adopted from the probe after a successful ``set_params`` (e.g. + ``_vcov_type_explicit``). +``_normalize_set_params`` + Classmethod rewriting the user-supplied ``params`` dict before the + merge (e.g. the ``robust``-alone alias re-derivation on + ``DifferenceInDifferences``). +""" + +import inspect +from typing import Any, ClassVar, Dict, Mapping, Tuple, TypeVar + +__all__ = ["BaseEstimator"] + +TSelf = TypeVar("TSelf", bound="BaseEstimator") + +# Signature kinds accepted as constructor parameters. VAR_POSITIONAL / +# VAR_KEYWORD are rejected loudly: an estimator constructor forwarding +# through *args/**kwargs has no introspectable parameter surface. +_ACCEPTED_KINDS = ( + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, +) + + +class BaseEstimator: + """Mixin providing sklearn-compatible ``get_params`` / ``set_params``. + + See the module docstring for the full contract. Subclasses only declare + the accommodation hooks when their constructor stores state under + different names than its parameters; most estimators are plain drop-ins. + """ + + _PARAM_ATTR_ALIASES: ClassVar[Mapping[str, str]] = {} + _DERIVED_CONFIG_ATTRS: ClassVar[Tuple[str, ...]] = () + + @classmethod + def _param_names(cls) -> Tuple[str, ...]: + """Constructor parameter names, introspected once per class. + + Cached via ``cls.__dict__`` (never a plain class attribute, which a + subclass with its own ``__init__`` would inherit; never an instance + attribute, which would break ``vars(est)`` stability across fit). + """ + cached = cls.__dict__.get("_param_names_cache") + if cached is not None: + return cached + names = [] + for name, param in inspect.signature(cls.__init__).parameters.items(): + if name == "self": + continue + if param.kind not in _ACCEPTED_KINDS: + raise TypeError( + f"{cls.__name__}.__init__ uses *args/**kwargs " + f"(parameter {name!r}); its parameter surface cannot be " + "introspected for get_params/set_params." + ) + names.append(name) + result: Tuple[str, ...] = tuple(names) + setattr(cls, "_param_names_cache", result) + return result + + @classmethod + def _normalize_set_params(cls, params: Dict[str, Any]) -> Dict[str, Any]: + """Rewrite the user-supplied params before the merge (default no-op).""" + return params + + def get_params(self, deep: bool = True) -> Dict[str, Any]: + """ + Get estimator parameters (sklearn-compatible). + + Parameters + ---------- + deep : bool, default True + Accepted for sklearn compatibility (``sklearn.base.clone`` + calls ``get_params(deep=False)``) and ignored: no diff-diff + estimator nests another estimator. + + Returns + ------- + Dict[str, Any] + Estimator parameters suitable for passing to ``__init__``. + Keys follow the ``__init__`` signature order. Where a class + stores a parameter's raw value under a different attribute + (``_PARAM_ATTR_ALIASES``), the raw value is returned so that + ``type(est)(**est.get_params())`` reconstructs the same + configuration. + """ + del deep + aliases = self._PARAM_ATTR_ALIASES + return {name: getattr(self, aliases.get(name, name)) for name in self._param_names()} + + def set_params(self: TSelf, **params: Any) -> TSelf: + """ + Set estimator parameters (sklearn-compatible, transactional). + + Validates the merged configuration by constructing a throwaway + instance of ``type(self)`` - so a rejected call raises the same + error ``__init__`` would and leaves this estimator unchanged - then + adopts the probe's parameter attributes and the class's declared + derived-config attributes. Fitted state is never touched. + + Parameters + ---------- + **params + Estimator parameters. Unknown names raise ``ValueError`` + before any validation or mutation. + + Returns + ------- + self + """ + names = self._param_names() + for key in params: + if key not in names: + raise ValueError(f"Unknown parameter: {key}") + cls = type(self) + normalized = cls._normalize_set_params(dict(params)) + merged = {**self.get_params(), **normalized} + probe = cls(**merged) + aliases = self._PARAM_ATTR_ALIASES + for name in names: + attr = aliases.get(name, name) + setattr(self, attr, getattr(probe, attr)) + for attr in self._DERIVED_CONFIG_ATTRS: + setattr(self, attr, getattr(probe, attr)) + return self diff --git a/diff_diff/bacon.py b/diff_diff/bacon.py index ca5f5926..6777f503 100644 --- a/diff_diff/bacon.py +++ b/diff_diff/bacon.py @@ -17,6 +17,7 @@ import numpy as np import pandas as pd +from diff_diff._base import BaseEstimator from diff_diff.results import _format_survey_block from diff_diff.results_base import Diagnostic from diff_diff.utils import pre_demean_norms, snap_absorbed_regressors @@ -358,7 +359,7 @@ def effect_by_type(self) -> Dict[str, Optional[float]]: } -class BaconDecomposition: +class BaconDecomposition(BaseEstimator): """ Goodman-Bacon (2021) decomposition of Two-Way Fixed Effects estimator. @@ -1254,19 +1255,7 @@ def _compute_timing_comparison( time_window=time_window, ) - def get_params(self) -> Dict[str, Any]: - """Get estimator parameters (sklearn-compatible).""" - return {"weights": self.weights} - - def set_params(self, **params) -> "BaconDecomposition": - """Set estimator parameters (sklearn-compatible).""" - if "weights" in params: - if params["weights"] not in ("approximate", "exact"): - raise ValueError( - f"weights must be 'approximate' or 'exact', " f"got '{params['weights']}'" - ) - self.weights = params["weights"] - return self + # get_params/set_params come from BaseEstimator. def summary(self) -> str: """Get summary of decomposition results.""" diff --git a/diff_diff/chaisemartin_dhaultfoeuille.py b/diff_diff/chaisemartin_dhaultfoeuille.py index 2fd34f7e..056d5e85 100644 --- a/diff_diff/chaisemartin_dhaultfoeuille.py +++ b/diff_diff/chaisemartin_dhaultfoeuille.py @@ -36,6 +36,7 @@ import numpy as np import pandas as pd +from diff_diff._base import BaseEstimator from diff_diff.chaisemartin_dhaultfoeuille_bootstrap import ( ChaisemartinDHaultfoeuilleBootstrapMixin, ) @@ -355,7 +356,7 @@ def _validate_paths_of_interest( # ============================================================================= -class ChaisemartinDHaultfoeuille(ChaisemartinDHaultfoeuilleBootstrapMixin): +class ChaisemartinDHaultfoeuille(ChaisemartinDHaultfoeuilleBootstrapMixin, BaseEstimator): """ de Chaisemartin-D'Haultfoeuille (dCDH) estimator. @@ -810,47 +811,7 @@ def __init__( # sklearn-style parameter introspection # ------------------------------------------------------------------ - def get_params(self) -> Dict[str, Any]: - """Return all ``__init__`` parameters as a dictionary.""" - return { - "alpha": self.alpha, - "cluster": self.cluster, - "n_bootstrap": self.n_bootstrap, - "bootstrap_weights": self.bootstrap_weights, - "seed": self.seed, - "placebo": self.placebo, - "twfe_diagnostic": self.twfe_diagnostic, - "drop_larger_lower": self.drop_larger_lower, - "by_path": self.by_path, - "paths_of_interest": self.paths_of_interest, - "rank_deficient_action": self.rank_deficient_action, - } - - def set_params(self, **params: Any) -> "ChaisemartinDHaultfoeuille": - """ - Set estimator parameters (sklearn-compatible). - - **Transactional**: validation runs after the candidate mutations, - and if any rule fails the estimator state is rolled back to its - pre-call values before the exception is re-raised. Callers can - therefore retry with corrected params on the same instance - without repairing inconsistent intermediate state. - """ - # Snapshot current values for the keys we are about to set so - # we can roll back on validation failure (transactional semantics). - for key in params: - if not hasattr(self, key): - raise ValueError(f"Unknown parameter: {key}") - snapshot = {key: getattr(self, key) for key in params} - try: - for key, value in params.items(): - setattr(self, key, value) - self._validate_invariants() - except Exception: - for key, value in snapshot.items(): - setattr(self, key, value) - raise - return self + # get_params/set_params come from BaseEstimator. def _validate_invariants(self) -> None: """Run the post-mutation validation rules. Mirrors `__init__`.""" diff --git a/diff_diff/changes_in_changes.py b/diff_diff/changes_in_changes.py index a93202ce..6ba8e4ac 100644 --- a/diff_diff/changes_in_changes.py +++ b/diff_diff/changes_in_changes.py @@ -53,6 +53,7 @@ from scipy import stats from scipy.optimize import linprog +from diff_diff._base import BaseEstimator from diff_diff.bootstrap_utils import warn_bootstrap_failure_rate from diff_diff.changes_in_changes_results import ChangesInChangesResults from diff_diff.utils import ( @@ -1166,7 +1167,7 @@ def _validate_all_params(params: Dict[str, Any]) -> None: _validate_seed(params["seed"]) -class ChangesInChanges: +class ChangesInChanges(BaseEstimator): """Changes-in-Changes estimator (Athey & Imbens 2006) for the 2x2 design. Estimates the ATT (eq. 36) and quantile treatment effects on the treated @@ -1258,31 +1259,7 @@ def __init__( self.is_fitted_ = False self.results_: Optional[ChangesInChangesResults] = None - def get_params(self, deep: bool = True) -> Dict[str, Any]: - """Return constructor hyperparameters (raw values, round-trips ``__init__``). - - ``deep`` is accepted for sklearn compatibility (``sklearn.base.clone`` - calls ``get_params(deep=False)``) and is ignored - there are no nested - estimators. - """ - return { - "quantiles": self.quantiles, - "n_bootstrap": self.n_bootstrap, - "alpha": self.alpha, - "panel": self.panel, - "seed": self.seed, - } - - def set_params(self, **params: Any) -> "ChangesInChanges": - """Set hyperparameters transactionally (a failing call mutates nothing).""" - valid = set(self.get_params()) - for key in params: - if key not in valid: - raise ValueError(f"Unknown parameter: {key}") - _validate_all_params({**self.get_params(), **params}) - for key, value in params.items(): - setattr(self, key, value) - return self + # get_params/set_params come from BaseEstimator. def fit( self, @@ -1326,7 +1303,7 @@ def fit( ) -class QDiD: +class QDiD(BaseEstimator): """Quantile Difference-in-Differences comparison estimator (2x2 design). Applies DiD quantile-by-quantile: ``qte(tau) = Q(y11, tau) - [Q(y10, tau) @@ -1387,31 +1364,7 @@ def __init__( self.is_fitted_ = False self.results_: Optional[ChangesInChangesResults] = None - def get_params(self, deep: bool = True) -> Dict[str, Any]: - """Return constructor hyperparameters (raw values, round-trips ``__init__``). - - ``deep`` is accepted for sklearn compatibility (``sklearn.base.clone`` - calls ``get_params(deep=False)``) and is ignored - there are no nested - estimators. - """ - return { - "quantiles": self.quantiles, - "n_bootstrap": self.n_bootstrap, - "alpha": self.alpha, - "panel": self.panel, - "seed": self.seed, - } - - def set_params(self, **params: Any) -> "QDiD": - """Set hyperparameters transactionally (a failing call mutates nothing).""" - valid = set(self.get_params()) - for key in params: - if key not in valid: - raise ValueError(f"Unknown parameter: {key}") - _validate_all_params({**self.get_params(), **params}) - for key, value in params.items(): - setattr(self, key, value) - return self + # get_params/set_params come from BaseEstimator. def fit( self, diff --git a/diff_diff/continuous_did.py b/diff_diff/continuous_did.py index 184b5e3a..5dca5247 100644 --- a/diff_diff/continuous_did.py +++ b/diff_diff/continuous_did.py @@ -15,6 +15,7 @@ import numpy as np import pandas as pd +from diff_diff._base import BaseEstimator from diff_diff.bootstrap_utils import ( compute_effect_bootstrap_stats, generate_bootstrap_weights_batch, @@ -49,7 +50,7 @@ __all__ = ["ContinuousDiD", "ContinuousDiDResults", "DoseResponseCurve"] -class ContinuousDiD: +class ContinuousDiD(BaseEstimator): """ Continuous Difference-in-Differences estimator. @@ -232,45 +233,7 @@ def _validate_constrained_params(self) -> None: "lowest-dose fit." ) - def get_params(self) -> Dict[str, Any]: - """Return estimator parameters as a dictionary.""" - return { - "degree": self.degree, - "num_knots": self.num_knots, - "dvals": self.dvals, - "control_group": self.control_group, - "anticipation": self.anticipation, - "base_period": self.base_period, - "alpha": self.alpha, - "n_bootstrap": self.n_bootstrap, - "bootstrap_weights": self.bootstrap_weights, - "seed": self.seed, - "rank_deficient_action": self.rank_deficient_action, - "covariates": self.covariates, - "estimation_method": self.estimation_method, - "pscore_trim": self.pscore_trim, - "epv_threshold": self.epv_threshold, - "pscore_fallback": self.pscore_fallback, - "treatment_type": self.treatment_type, - } - - def set_params(self, **params) -> "ContinuousDiD": - """Set estimator parameters and return self. - - Transactional: the candidate configuration is validated against a clone - before any attribute on ``self`` is mutated, so an invalid update leaves - the estimator unchanged (no half-applied state). - """ - for key in params: - if not hasattr(self, key): - raise ValueError(f"Invalid parameter: {key}") - # Validate the candidate config on a throwaway before mutating self. - candidate = ContinuousDiD(**{**self.get_params(), **params}) - # candidate.__init__ ran _validate_constrained_params() successfully. - del candidate - for key, value in params.items(): - setattr(self, key, value) - return self + # get_params/set_params come from BaseEstimator. # ------------------------------------------------------------------ # Main fit diff --git a/diff_diff/efficient_did.py b/diff_diff/efficient_did.py index 98496c59..7d83c21e 100644 --- a/diff_diff/efficient_did.py +++ b/diff_diff/efficient_did.py @@ -28,6 +28,7 @@ import numpy as np import pandas as pd +from diff_diff._base import BaseEstimator from diff_diff.efficient_did_bootstrap import ( EDiDBootstrapResults, EfficientDiDBootstrapMixin, @@ -209,7 +210,7 @@ def _hausman_quadratic_form( return H, effective_rank, p_value, n_negative, True -class EfficientDiD(EfficientDiDBootstrapMixin): +class EfficientDiD(EfficientDiDBootstrapMixin, BaseEstimator): """Efficient DiD estimator (Chen, Sant'Anna & Xie 2025). Without covariates, achieves the semiparametric efficiency bound for @@ -434,51 +435,7 @@ def _validate_vcov_type(vcov_type: str) -> None: # -- sklearn compatibility ------------------------------------------------ - def get_params(self) -> Dict[str, Any]: - """Get estimator parameters (sklearn-compatible).""" - return { - "pt_assumption": self.pt_assumption, - "anticipation": self.anticipation, - "alpha": self.alpha, - "cluster": self.cluster, - "vcov_type": self.vcov_type, - "control_group": self.control_group, - "n_bootstrap": self.n_bootstrap, - "bootstrap_weights": self.bootstrap_weights, - "seed": self.seed, - "sieve_k_max": self.sieve_k_max, - "sieve_criterion": self.sieve_criterion, - "ratio_clip": self.ratio_clip, - "kernel_bandwidth": self.kernel_bandwidth, - "omega_ridge": self.omega_ridge, - } - - def set_params(self, **params: Any) -> "EfficientDiD": - """Set estimator parameters (sklearn-compatible). - - Atomic: snapshots the original attribute values before applying - mutations, validates the new state via ``_validate_params``, and - rolls every attribute back to its pre-call value if validation - raises. Without this, ``set_params(vcov_type="classical", - alpha=0.1)`` would leave ``self.vcov_type`` partially mutated - even though the call raised, defeating the eager-validation - contract for callers that catch ``ValueError`` and keep using - the estimator. - """ - snapshot: Dict[str, Any] = {} - for key in params: - if not hasattr(self, key): - raise ValueError(f"Unknown parameter: {key}") - snapshot[key] = getattr(self, key) - for key, value in params.items(): - setattr(self, key, value) - try: - self._validate_params() - except Exception: - for key, value in snapshot.items(): - setattr(self, key, value) - raise - return self + # get_params/set_params come from BaseEstimator. # -- Main estimation ------------------------------------------------------ diff --git a/diff_diff/estimators.py b/diff_diff/estimators.py index fbaa5d13..77b35178 100644 --- a/diff_diff/estimators.py +++ b/diff_diff/estimators.py @@ -19,6 +19,7 @@ import numpy as np import pandas as pd +from diff_diff._base import BaseEstimator from diff_diff.linalg import ( LinearRegression, _absorbed_fe_vcov_scale, @@ -46,7 +47,7 @@ ) -class DifferenceInDifferences: +class DifferenceInDifferences(BaseEstimator): """ Difference-in-Differences estimator with sklearn-like interface. @@ -1177,107 +1178,26 @@ def predict(self, data: pd.DataFrame) -> np.ndarray: "training-data predictions." ) - def get_params(self) -> Dict[str, Any]: - """ - Get estimator parameters (sklearn-compatible). - - Returns the *raw* user input for ``vcov_type`` (``None`` when - the value was alias-derived from ``robust``). This preserves - the backward-compat remap semantics across clones: a clone of - ``DifferenceInDifferences(robust=False, cluster="unit")`` must - behave the same as the original on a clustered fit, which - requires the clone's ``__init__`` to see ``vcov_type=None`` (so - it flags ``_vcov_type_explicit=False``) rather than the - alias-resolved ``"classical"`` (which would mark it explicit - and skip the CR1 remap). - - Returns - ------- - Dict[str, Any] - Estimator parameters suitable for passing to ``__init__``. - """ - return { - "robust": self.robust, - "cluster": self.cluster, - "vcov_type": self._vcov_type_arg, # raw, possibly None - "alpha": self.alpha, - "inference": self.inference, - "n_bootstrap": self.n_bootstrap, - "bootstrap_weights": self.bootstrap_weights, - "p_val_type": self.p_val_type, - "seed": self.seed, - "rank_deficient_action": self.rank_deficient_action, - "conley_coords": self.conley_coords, - "conley_cutoff_km": self.conley_cutoff_km, - "conley_metric": self.conley_metric, - "conley_kernel": self.conley_kernel, - "conley_lag_cutoff": self.conley_lag_cutoff, - "df_convention": self.df_convention, - } - - def set_params(self, **params) -> "DifferenceInDifferences": - """ - Set estimator parameters (sklearn-compatible). - - After assignment, the ``robust``/``vcov_type`` pair is re-validated via - the same :func:`diff_diff.linalg.resolve_vcov_type` helper used by - ``__init__``. Invalid combinations (e.g. ``robust=False`` with - ``vcov_type="hc2"``) raise ``ValueError`` instead of leaving the - object in an inconsistent state. - - Parameters - ---------- - **params - Estimator parameters. - - Returns - ------- - self - """ - from diff_diff.linalg import resolve_vcov_type - - # Validate BEFORE mutating `self`. A failing call must leave the - # estimator unchanged so callers that catch `ValueError` can keep - # reasoning about the object; half-mutated state from an earlier - # partial assignment defeats that guarantee. Compute the resolved - # `vcov_type` on local variables, then apply all mutations atomically. - pending_robust = params.get("robust", self.robust) - pending_vcov_type = params.get("vcov_type", self.vcov_type) - pending_df_convention = params.get("df_convention", self.df_convention) - validate_df_convention(pending_df_convention) - - # First pass: validate that every incoming key is a known attribute - # so we don't partially apply a batch that ends in "Unknown parameter". - for key in params: - if not hasattr(self, key): - raise ValueError(f"Unknown parameter: {key}") - - # Second pass: resolve the robust/vcov_type pair. When the user passes - # only `robust=` alongside a previously-set non-aliasing `vcov_type`, - # re-derive `vcov_type` from the new `robust` value for internal - # consistency (matching the prior behavior, but now on locals). - if "vcov_type" in params: - resolved_vcov = resolve_vcov_type(pending_robust, pending_vcov_type) - elif "robust" in params: - resolved_vcov = resolve_vcov_type(pending_robust, None) - else: - resolved_vcov = self.vcov_type # no-op if neither changed - - # All validation passed — apply mutations atomically. - for key, value in params.items(): - setattr(self, key, value) - self.vcov_type = resolved_vcov - # Update the raw-vs-resolved tracking. `vcov_type=` in the call - # updates `_vcov_type_arg` to whatever the user passed (including - # None); `robust=` alone clears the raw arg since the resolution - # re-derives from the alias. The `_vcov_type_explicit` flag is - # True iff the raw arg is non-None. - if "vcov_type" in params: - self._vcov_type_arg = params["vcov_type"] - elif "robust" in params: - self._vcov_type_arg = None - self._vcov_type_explicit = self._vcov_type_arg is not None - return self + # get_params/set_params come from BaseEstimator. `vcov_type`'s RAW arg + # lives at `_vcov_type_arg` (the resolved value at `vcov_type`), and + # get_params must return the raw one: a clone of + # `DifferenceInDifferences(robust=False, cluster="unit")` must behave + # the same as the original on a clustered fit, which requires the + # clone's `__init__` to see `vcov_type=None` (flagging + # `_vcov_type_explicit=False`) rather than the alias-resolved + # "classical" (which would mark it explicit and skip the CR1 remap). + _PARAM_ATTR_ALIASES = {"vcov_type": "_vcov_type_arg"} + _DERIVED_CONFIG_ATTRS = ("vcov_type", "_vcov_type_arg", "_vcov_type_explicit") + + @classmethod + def _normalize_set_params(cls, params: Dict[str, Any]) -> Dict[str, Any]: + # `robust=` alone re-derives `vcov_type` from the alias: the merged + # probe config must carry vcov_type=None so `resolve_vcov_type` + # re-runs on the new `robust` value instead of seeing a conflict + # with the previously stored raw vcov_type. + if "robust" in params and "vcov_type" not in params: + params["vcov_type"] = None + return params def _warn_replicate_vcov_ignored(self) -> bool: """Warn that an explicit ``vcov_type`` has no effect under a diff --git a/diff_diff/had.py b/diff_diff/had.py index 081953cd..a18aa767 100644 --- a/diff_diff/had.py +++ b/diff_diff/had.py @@ -71,6 +71,7 @@ import numpy as np import pandas as pd +from diff_diff._base import BaseEstimator from diff_diff.bootstrap_chunking import ( compute_block_size, iter_survey_multiplier_weight_blocks, @@ -2644,7 +2645,7 @@ def _pack( # ============================================================================= -class HeterogeneousAdoptionDiD: +class HeterogeneousAdoptionDiD(BaseEstimator): """Heterogeneous Adoption Difference-in-Differences estimator. Implements de Chaisemartin, Ciccia, D'Haultfoeuille, and Knau (2026) @@ -2908,64 +2909,7 @@ def _validate_constructor_args(self) -> None: if int(self.seed) < 0: raise ValueError(f"seed must be nonneg; got {self.seed!r}.") - def get_params(self, deep: bool = True) -> Dict[str, Any]: - """Return the raw constructor parameters (sklearn-compatible). - - Matches the :meth:`sklearn.base.BaseEstimator.get_params` - signature. Preserves the user's original inputs - in particular, - ``design`` returns ``"auto"`` when the user set it to ``"auto"`` - (even after fit), so ``sklearn.base.clone(est)`` round-trips - exactly. - - Parameters - ---------- - deep : bool, default=True - Accepted for sklearn-contract compatibility. This estimator - has no nested sub-estimator parameters, so ``deep=False`` - and ``deep=True`` return the same dict. - """ - del deep # accepted for compat; this estimator has no nested params - return { - "design": self.design, - "d_lower": self.d_lower, - "kernel": self.kernel, - "alpha": self.alpha, - "vcov_type": self.vcov_type, - "robust": self.robust, - "cluster": self.cluster, - "n_bootstrap": self.n_bootstrap, - "seed": self.seed, - } - - def set_params(self, **params: Any) -> "HeterogeneousAdoptionDiD": - """Set estimator parameters and return self (sklearn-compatible). - - Only keys returned by :meth:`get_params` are accepted. Passing - any other attribute name (including method names like ``fit``) - raises ``ValueError`` so the estimator cannot be silently - corrupted by a mistyped or attacker-supplied key. - - Mutation is ATOMIC: validation runs on a proposed merged - parameter dict before any attribute is overwritten. A failing - call (invalid key, or an otherwise valid key whose value - violates the constructor constraints) leaves ``self`` unchanged - and safe to reuse. - """ - valid_keys = set(self.get_params().keys()) - invalid = [k for k in params if k not in valid_keys] - if invalid: - raise ValueError( - f"Invalid parameter: {invalid[0]!r}. Valid parameters: " f"{sorted(valid_keys)}." - ) - # Dry-run validation by constructing a fresh instance with the - # merged state. If the constructor raises, self is not mutated. - merged = self.get_params() - merged.update(params) - type(self)(**merged) # raises ValueError on invalid combination - # All checks passed; apply atomically. - for key, value in params.items(): - setattr(self, key, value) - return self + # get_params/set_params come from BaseEstimator. # ------------------------------------------------------------------ # Main fit entry point diff --git a/diff_diff/honest_did.py b/diff_diff/honest_did.py index 10b376a9..e5aa61af 100644 --- a/diff_diff/honest_did.py +++ b/diff_diff/honest_did.py @@ -25,6 +25,7 @@ import pandas as pd from scipy import optimize +from diff_diff._base import BaseEstimator from diff_diff.results import ( MultiPeriodDiDResults, ) @@ -2369,7 +2370,7 @@ def _arp_confidence_set( # ============================================================================= -class HonestDiD: +class HonestDiD(BaseEstimator): """ Honest DiD sensitivity analysis (Rambachan & Roth 2023). @@ -2445,24 +2446,7 @@ def _validate_params(self): if not 0 < self.alpha < 1: raise ValueError(f"alpha must be between 0 and 1, got alpha={self.alpha}") - def get_params(self) -> Dict[str, Any]: - """Get parameters for this estimator.""" - return { - "method": self.method, - "M": self.M, - "alpha": self.alpha, - "l_vec": self.l_vec, - } - - def set_params(self, **params) -> "HonestDiD": - """Set parameters for this estimator.""" - for key, value in params.items(): - if hasattr(self, key): - setattr(self, key, value) - else: - raise ValueError(f"Invalid parameter: {key}") - self._validate_params() - return self + # get_params/set_params come from BaseEstimator. def fit( self, diff --git a/diff_diff/imputation.py b/diff_diff/imputation.py index 7c089294..d8e2e8b2 100644 --- a/diff_diff/imputation.py +++ b/diff_diff/imputation.py @@ -31,6 +31,7 @@ from scipy import sparse, stats from scipy.sparse.linalg import factorized as sparse_factorized +from diff_diff._base import BaseEstimator from diff_diff.imputation_bootstrap import ImputationDiDBootstrapMixin, _compute_target_weights from diff_diff.imputation_results import ( # noqa: F401 (re-export) ImputationBootstrapResults, @@ -145,7 +146,7 @@ def _lsmr_minnorm_normal_solve(A0tA0_csc, rhs: np.ndarray) -> np.ndarray: return z -class ImputationDiD(ImputationDiDBootstrapMixin): +class ImputationDiD(ImputationDiDBootstrapMixin, BaseEstimator): """ Borusyak-Jaravel-Spiess (2024) imputation DiD estimator. @@ -390,9 +391,9 @@ def fit( ValueError If required columns are missing or data validation fails. """ - # Re-validate vcov_type at fit-time so sklearn-style set_params - # mutations (e.g. set_params(vcov_type="classical")) are re-checked - # at use rather than silently accepted by the parameter setter. + # Re-validate vcov_type at fit-time: set_params validates eagerly + # (BaseEstimator probe re-init), so this only catches DIRECT + # attribute mutation (est.vcov_type = ...). self._validate_vcov_type(self.vcov_type) self._validate_leave_one_out(self.leave_one_out) @@ -2748,36 +2749,7 @@ def _pretrend_test(self, n_leads: Optional[int] = None) -> Dict[str, Any]: # sklearn-compatible interface # ========================================================================= - def get_params(self) -> Dict[str, Any]: - """Get estimator parameters (sklearn-compatible).""" - return { - "anticipation": self.anticipation, - "alpha": self.alpha, - "cluster": self.cluster, - "vcov_type": self.vcov_type, - "n_bootstrap": self.n_bootstrap, - "bootstrap_weights": self.bootstrap_weights, - "seed": self.seed, - "rank_deficient_action": self.rank_deficient_action, - "horizon_max": self.horizon_max, - "aux_partition": self.aux_partition, - "pretrends": self.pretrends, - "leave_one_out": self.leave_one_out, - "df_convention": self.df_convention, - } - - def set_params(self, **params) -> "ImputationDiD": - """Set estimator parameters (sklearn-compatible).""" - # Reject unknown keys and validate the pending df_convention BEFORE - # any assignment so a rejected call leaves the estimator unchanged. - for key in params: - if not hasattr(self, key): - raise ValueError(f"Unknown parameter: {key}") - if "df_convention" in params: - validate_df_convention(params["df_convention"]) - for key, value in params.items(): - setattr(self, key, value) - return self + # get_params/set_params come from BaseEstimator. @staticmethod def _validate_leave_one_out(leave_one_out: Any) -> None: @@ -2853,9 +2825,9 @@ def _validate_vcov_type(vcov_type: str) -> None: """Validate ``vcov_type`` membership against ImputationDiD's permanently-narrow influence-function variance contract. - Called from ``__init__`` AND ``fit()`` so sklearn-style - ``set_params(vcov_type=...)`` mutations are re-checked at use - time rather than silently accepted by the parameter setter. + Called from ``__init__`` AND ``fit()``; ``set_params`` validates + eagerly via the BaseEstimator probe re-init, so the fit-time + re-check only catches direct attribute mutation. Mirrors the TripleDifference / CallawaySantAnna pattern (no single design matrix on which hat-matrix leverage or Bell- McCaffrey Satterthwaite DOF can be defined). diff --git a/diff_diff/lpdid.py b/diff_diff/lpdid.py index e98d38ae..dc9fc7eb 100644 --- a/diff_diff/lpdid.py +++ b/diff_diff/lpdid.py @@ -1,9 +1,10 @@ import warnings -from typing import Any, Dict, Iterable, Optional, Union +from typing import Dict, Iterable, Optional, Union import numpy as np import pandas as pd +from diff_diff._base import BaseEstimator from diff_diff.linalg import InvalidClusterKAdjustment, _rank_guarded_inv, solve_ols from diff_diff.lpdid_results import LPDiDResults from diff_diff.utils import ( @@ -17,7 +18,7 @@ __all__ = ["LPDiD", "LPDiDResults"] -class LPDiD: +class LPDiD(BaseEstimator): def __init__( self, pre_window: int = 2, @@ -1693,36 +1694,4 @@ def fit( self.is_fitted_ = True return self.results_ - def get_params(self) -> Dict[str, Any]: - return { - "pre_window": self.pre_window, - "post_window": self.post_window, - "control_group": self.control_group, - "reweight": self.reweight, - "no_composition": self.no_composition, - "pmd": self.pmd, - "alpha": self.alpha, - "cluster": self.cluster, - "rank_deficient_action": self.rank_deficient_action, - "non_absorbing": self.non_absorbing, - "stabilization_window": self.stabilization_window, - "df_convention": self.df_convention, - } - - def set_params(self, **params: Any) -> "LPDiD": - # Reject unknown keys before any assignment (the rollback below - # only covers validator failures on known keys). - for key in params: - if not hasattr(self, key): - raise ValueError(f"Unknown parameter: {key}") - previous_values = {} - for key, value in params.items(): - previous_values[key] = getattr(self, key) - setattr(self, key, value) - try: - self._validate_params() - except ValueError: - for key, value in previous_values.items(): - setattr(self, key, value) - raise - return self + # get_params/set_params come from BaseEstimator. diff --git a/diff_diff/pretrends.py b/diff_diff/pretrends.py index 79c94c21..6e06caab 100644 --- a/diff_diff/pretrends.py +++ b/diff_diff/pretrends.py @@ -33,6 +33,7 @@ import pandas as pd from scipy import optimize, stats +from diff_diff._base import BaseEstimator from diff_diff.results import MultiPeriodDiDResults from diff_diff.results_base import Diagnostic @@ -764,7 +765,7 @@ def plot( # ============================================================================= -class PreTrendsPower: +class PreTrendsPower(BaseEstimator): """ Pre-trends power analysis (Roth 2022). @@ -872,26 +873,8 @@ def __init__( ) self.pretest_form = pretest_form - def get_params(self) -> Dict[str, Any]: - """Get parameters for this estimator.""" - return { - "alpha": self.alpha, - "power": self.target_power, - "violation_type": self.violation_type, - "violation_weights": self.violation_weights, - "pretest_form": self.pretest_form, - } - - def set_params(self, **params) -> "PreTrendsPower": - """Set parameters for this estimator.""" - for key, value in params.items(): - if key == "power": - self.target_power = value - elif hasattr(self, key): - setattr(self, key, value) - else: - raise ValueError(f"Invalid parameter: {key}") - return self + # get_params/set_params come from BaseEstimator. + _PARAM_ATTR_ALIASES = {"power": "target_power"} def _get_violation_weights( self, diff --git a/diff_diff/rdd.py b/diff_diff/rdd.py index 202ae816..fe926a79 100644 --- a/diff_diff/rdd.py +++ b/diff_diff/rdd.py @@ -111,6 +111,7 @@ import numpy as np import pandas as pd +from diff_diff._base import BaseEstimator from diff_diff._rdrobust_port import ( BWSELECT_OPTIONS, _fuzzy_identification_stop, @@ -488,7 +489,7 @@ def to_dataframe(self) -> pd.DataFrame: return pd.DataFrame([self.to_dict()]) -class RegressionDiscontinuity: +class RegressionDiscontinuity(BaseEstimator): """Regression discontinuity estimator - sharp and fuzzy, with optional covariate adjustment (rdrobust 4.0.0 parity). @@ -693,41 +694,7 @@ def _validate_constructor_args(self) -> None: if not (self._is_real_scalar(self.alpha) and 0 < self.alpha < 1): raise ValueError(f"alpha must be in (0, 1); got {self.alpha!r}.") - def get_params(self, deep: bool = True) -> Dict[str, Any]: - """Return raw constructor parameters (sklearn-compatible).""" - del deep - return { - "cutoff": self.cutoff, - "p": self.p, - "q": self.q, - "kernel": self.kernel, - "bwselect": self.bwselect, - "h": self.h, - "b": self.b, - "rho": self.rho, - "vcov_type": self.vcov_type, - "nnmatch": self.nnmatch, - "masspoints": self.masspoints, - "bwcheck": self.bwcheck, - "bwrestrict": self.bwrestrict, - "scaleregul": self.scaleregul, - "sharpbw": self.sharpbw, - "covs_drop": self.covs_drop, - "alpha": self.alpha, - } - - def set_params(self, **params: Any) -> "RegressionDiscontinuity": - """Transactionally update parameters (validate before mutating).""" - valid = set(self.get_params().keys()) - unknown = set(params) - valid - if unknown: - raise ValueError(f"Unknown parameter(s): {sorted(unknown)}. Valid: {sorted(valid)}.") - merged = self.get_params() - merged.update(params) - type(self)(**merged) # dry-run: raises before any mutation - for key, value in params.items(): - setattr(self, key, value) - return self + # get_params/set_params come from BaseEstimator. # ------------------------------------------------------------------ # Fitting diff --git a/diff_diff/rdplot.py b/diff_diff/rdplot.py index 865a1731..42e2ef4c 100644 --- a/diff_diff/rdplot.py +++ b/diff_diff/rdplot.py @@ -36,6 +36,7 @@ import pandas as pd from scipy import stats +from ._base import BaseEstimator from ._rdrobust_port import ( _covs_gamma, _normalize_kernel, @@ -310,7 +311,7 @@ def plot( return ax -class RDPlot: +class RDPlot(BaseEstimator): """Data-driven RD plot builder (CCT 2015; rdrobust 4.0.0 ``rdplot()`` parity). @@ -466,34 +467,7 @@ def _validate_constructor_args(self) -> None: if not isinstance(self.covs_drop, (bool, np.bool_)): raise ValueError(f"covs_drop must be a bool; got {self.covs_drop!r}.") - def get_params(self) -> Dict[str, Any]: - """Get parameters of this plot builder.""" - return { - "cutoff": self.cutoff, - "p": self.p, - "nbins": self.nbins, - "binselect": self.binselect, - "scale": self.scale, - "kernel": self.kernel, - "h": self.h, - "support": self.support, - "masspoints": self.masspoints, - "ci": self.ci, - "covs_drop": self.covs_drop, - } - - def set_params(self, **params: Any) -> "RDPlot": - """Set parameters (transactional: validates the full candidate - configuration before mutating this instance).""" - current = self.get_params() - unknown = set(params) - set(current) - if unknown: - raise ValueError(f"Invalid parameter(s) for RDPlot: {sorted(unknown)}") - candidate = {**current, **params} - RDPlot(**candidate) # full validation on a throwaway instance - for key, value in candidate.items(): - setattr(self, key, value) - return self + # get_params/set_params come from BaseEstimator. # ------------------------------------------------------------------ # Fit diff --git a/diff_diff/spillover.py b/diff_diff/spillover.py index b9d9b1b2..e9b13f74 100644 --- a/diff_diff/spillover.py +++ b/diff_diff/spillover.py @@ -32,6 +32,7 @@ import pandas as pd from scipy import sparse +from diff_diff._base import BaseEstimator from diff_diff.conley import ( _CONLEY_EARTH_RADIUS_KM, _CONLEY_SPARSE_N_THRESHOLD, @@ -1645,7 +1646,7 @@ def _build(mask: Optional[np.ndarray]) -> sparse.csr_matrix: # ============================================================================= -class SpilloverDiD: +class SpilloverDiD(BaseEstimator): """Ring-indicator spillover-aware DiD (Butts 2021). Standalone estimator implementing two-stage Gardner (2022) methodology @@ -1774,32 +1775,7 @@ def __init__( self.is_fitted_ = False self.results_: Optional[Any] = None - def get_params(self) -> Dict[str, Any]: - return { - "rings": self.rings, - "d_bar": self.d_bar, - "vcov_type": self.vcov_type, - "conley_coords": self.conley_coords, - "conley_metric": self.conley_metric, - "conley_cutoff_km": self.conley_cutoff_km, - "conley_lag_cutoff": self.conley_lag_cutoff, - "cluster": self.cluster, - "alpha": self.alpha, - "anticipation": self.anticipation, - "event_study": self.event_study, - "horizon_max": self.horizon_max, - "rank_deficient_action": self.rank_deficient_action, - } - - def set_params(self, **params: Any) -> "SpilloverDiD": - valid = set(self.get_params().keys()) - for key, value in params.items(): - if key not in valid: - raise ValueError( - f"Unknown parameter: {key!r}. Valid parameters: " f"{sorted(valid)}." - ) - setattr(self, key, value) - return self + # get_params/set_params come from BaseEstimator. # ------------------------------------------------------------------------- # Fit-time validators (Step 2) diff --git a/diff_diff/stacked_did.py b/diff_diff/stacked_did.py index 9060366b..de9f9631 100644 --- a/diff_diff/stacked_did.py +++ b/diff_diff/stacked_did.py @@ -24,6 +24,7 @@ import numpy as np import pandas as pd +from diff_diff._base import BaseEstimator from diff_diff.balancing import BalanceError, entropy_balance from diff_diff.linalg import effective_cluster_count, solve_ols from diff_diff.stacked_did_results import StackedDiDResults # noqa: F401 (re-export) @@ -36,7 +37,7 @@ ] -class StackedDiD: +class StackedDiD(BaseEstimator): """ Stacked Difference-in-Differences estimator. @@ -1546,45 +1547,7 @@ def _expected_units(a_val: Any) -> set: # sklearn-compatible interface # ========================================================================= - def get_params(self) -> Dict[str, Any]: - """Get estimator parameters (sklearn-compatible).""" - return { - "kappa_pre": self.kappa_pre, - "kappa_post": self.kappa_post, - "weighting": self.weighting, - "clean_control": self.clean_control, - "cluster": self.cluster, - "alpha": self.alpha, - "anticipation": self.anticipation, - "rank_deficient_action": self.rank_deficient_action, - "vcov_type": self.vcov_type, - "balance": self.balance, - "df_convention": self.df_convention, - } - - def set_params(self, **params: Any) -> "StackedDiD": - """Set estimator parameters (sklearn-compatible). - - Re-validates `vcov_type` via the shared `_validate_vcov_type` - helper so sklearn-style mutation hits the estimator-level guard - before fit() (avoids a later, less-informative failure in the - linalg layer). - """ - # Reject unknown keys, then validate pending values up-front (same - # error surface as __init__), so a rejected call leaves the - # estimator unchanged. - for key in params: - if not hasattr(self, key): - raise ValueError(f"Unknown parameter: {key}") - if "vcov_type" in params: - self._validate_vcov_type(params["vcov_type"]) - if "balance" in params: - self._validate_balance(params["balance"]) - if "df_convention" in params: - validate_df_convention(params["df_convention"]) - for key, value in params.items(): - setattr(self, key, value) - return self + # get_params/set_params come from BaseEstimator. def summary(self) -> str: """Get summary of estimation results.""" diff --git a/diff_diff/staggered.py b/diff_diff/staggered.py index 82ff94ee..8cc31528 100644 --- a/diff_diff/staggered.py +++ b/diff_diff/staggered.py @@ -12,6 +12,7 @@ import numpy as np import pandas as pd +from diff_diff._base import BaseEstimator from diff_diff.aggregation import ( AggregationKit, ) @@ -291,6 +292,7 @@ def _nan_gt_entry( class CallawaySantAnna( CallawaySantAnnaBootstrapMixin, CallawaySantAnnaAggregationMixin, + BaseEstimator, ): """ Callaway-Sant'Anna (2021) estimator for staggered Difference-in-Differences. @@ -595,8 +597,9 @@ def __init__( # per-(g,t) doubly-robust / IPW / outcome-regression structure # doesn't have. See REGISTRY.md "IF-based variance estimators vs # analytical-sandwich estimators" for the structural taxonomy. - # Factored out so fit() can re-run it after sklearn-style - # set_params bypasses __init__ validation. + # Factored out so fit() can re-run it: set_params now validates + # eagerly via the BaseEstimator probe re-init, but DIRECT attribute + # mutation (est.vcov_type = ...) still bypasses validation until fit. self._validate_vcov_type(vcov_type) self.control_group = control_group @@ -1930,10 +1933,10 @@ def fit( # cell. Sibling of PR #9 finding #17. self._safe_inv_tracker: List[float] = [] - # Re-validate vcov_type at fit-time so sklearn-style set_params - # mutations are caught before they propagate to Results metadata. - # __init__ already validated the constructor argument; this is the - # second layer for the post-construction mutation path. + # Re-validate vcov_type at fit-time: __init__ and set_params (via + # the BaseEstimator probe re-init) both validate eagerly, so this + # second layer only catches DIRECT attribute mutation + # (est.vcov_type = ...) before it propagates to Results metadata. self._validate_vcov_type(self.vcov_type) # --- allow_unbalanced_panel routing (RC-on-panel = R's allow_unbalanced_panel) --- @@ -5033,47 +5036,8 @@ def _validate_vcov_type(vcov_type: str) -> None: "structure; see REGISTRY.md." ) - def get_params(self) -> Dict[str, Any]: - """Get estimator parameters (sklearn-compatible).""" - return { - "control_group": self.control_group, - "anticipation": self.anticipation, - "estimation_method": self.estimation_method, - "alpha": self.alpha, - "cluster": self.cluster, - "vcov_type": self.vcov_type, - "n_bootstrap": self.n_bootstrap, - "bootstrap_weights": self.bootstrap_weights, - "seed": self.seed, - "rank_deficient_action": self.rank_deficient_action, - "base_period": self.base_period, - "cband": self.cband, - "pscore_trim": self.pscore_trim, - "panel": self.panel, - "allow_unbalanced_panel": self.allow_unbalanced_panel, - "epv_threshold": self.epv_threshold, - "pscore_fallback": self.pscore_fallback, - } - - def set_params(self, **params) -> "CallawaySantAnna": - """Set estimator parameters (sklearn-compatible). - - Mirrors SA pattern at ``sun_abraham.py:2150-2161``: setattr first, - then refresh ``_vcov_type_explicit`` if ``vcov_type`` changed. - Membership validation of ``vcov_type`` is deferred to next - ``fit()`` call (sklearn-style ``set_params`` is documented as - mutate-then-validate-at-use). Bad values like - ``set_params(vcov_type="hc4")`` surface at the next ``__init__``- - style validation call. - """ - for key, value in params.items(): - if hasattr(self, key): - setattr(self, key, value) - else: - raise ValueError(f"Unknown parameter: {key}") - if "vcov_type" in params: - self._vcov_type_explicit = self.vcov_type != "hc1" - return self + # get_params/set_params come from BaseEstimator. + _DERIVED_CONFIG_ATTRS = ("_vcov_type_explicit",) def summary(self) -> str: """Get summary of estimation results.""" diff --git a/diff_diff/staggered_triple_diff.py b/diff_diff/staggered_triple_diff.py index 0508e78c..107ffac0 100644 --- a/diff_diff/staggered_triple_diff.py +++ b/diff_diff/staggered_triple_diff.py @@ -15,6 +15,7 @@ import numpy as np import pandas as pd +from diff_diff._base import BaseEstimator from diff_diff.linalg import ( _check_propensity_diagnostics, _rank_guarded_inv, @@ -44,6 +45,7 @@ class StaggeredTripleDifference( CallawaySantAnnaBootstrapMixin, CallawaySantAnnaAggregationMixin, + BaseEstimator, ): """ Staggered Triple Difference (DDD) estimator. @@ -161,35 +163,7 @@ def __init__( self.is_fitted_ = False self.results_: Optional[StaggeredTripleDiffResults] = None - def get_params(self) -> Dict[str, Any]: - """Get estimator parameters (sklearn-compatible).""" - return { - "estimation_method": self.estimation_method, - "control_group": self.control_group, - "alpha": self.alpha, - "anticipation": self.anticipation, - "base_period": self.base_period, - "n_bootstrap": self.n_bootstrap, - "bootstrap_weights": self.bootstrap_weights, - "seed": self.seed, - "cband": self.cband, - "pscore_trim": self.pscore_trim, - "cluster": self.cluster, - "rank_deficient_action": self.rank_deficient_action, - "epv_threshold": self.epv_threshold, - "pscore_fallback": self.pscore_fallback, - } - - def set_params(self, **params) -> "StaggeredTripleDifference": - """Set estimator parameters (sklearn-compatible).""" - valid_params = self.get_params() - for key, value in params.items(): - if key not in valid_params: - raise ValueError(f"Unknown parameter: {key}") - setattr(self, key, value) - if "bootstrap_weights" in params: - self.bootstrap_weights = params["bootstrap_weights"] - return self + # get_params/set_params come from BaseEstimator. # ------------------------------------------------------------------ # fit() diff --git a/diff_diff/sun_abraham.py b/diff_diff/sun_abraham.py index 784722d3..19cac581 100644 --- a/diff_diff/sun_abraham.py +++ b/diff_diff/sun_abraham.py @@ -27,6 +27,7 @@ import numpy as np import pandas as pd +from diff_diff._base import BaseEstimator from diff_diff.bootstrap_utils import compute_effect_bootstrap_stats if TYPE_CHECKING: @@ -504,7 +505,7 @@ class SABootstrapResults: bootstrap_distribution: Optional[np.ndarray] = field(default=None, repr=False) -class SunAbraham: +class SunAbraham(BaseEstimator): """ Sun-Abraham (2021) interaction-weighted estimator for staggered DiD. @@ -2534,41 +2535,8 @@ def _run_rao_wu_bootstrap( bootstrap_distribution=bootstrap_overall, ) - def get_params(self) -> Dict[str, Any]: - """Get estimator parameters (sklearn-compatible).""" - return { - "control_group": self.control_group, - "anticipation": self.anticipation, - "alpha": self.alpha, - "cluster": self.cluster, - "n_bootstrap": self.n_bootstrap, - "seed": self.seed, - "rank_deficient_action": self.rank_deficient_action, - "vcov_type": self.vcov_type, - "conley_coords": self.conley_coords, - "conley_cutoff_km": self.conley_cutoff_km, - "conley_metric": self.conley_metric, - "conley_kernel": self.conley_kernel, - "conley_lag_cutoff": self.conley_lag_cutoff, - "df_convention": self.df_convention, - } - - def set_params(self, **params) -> "SunAbraham": - """Set estimator parameters (sklearn-compatible).""" - # Reject unknown keys and validate the pending df_convention BEFORE - # any assignment so a rejected call leaves the estimator unchanged. - for key in params: - if not hasattr(self, key): - raise ValueError(f"Unknown parameter: {key}") - if "df_convention" in params: - validate_df_convention(params["df_convention"]) - for key, value in params.items(): - setattr(self, key, value) - # Refresh the explicit-vcov-type flag if vcov_type changed, so the - # auto-cluster guard at fit time uses the updated value. - if "vcov_type" in params: - self._vcov_type_explicit = self.vcov_type != "hc1" - return self + # get_params/set_params come from BaseEstimator. + _DERIVED_CONFIG_ATTRS = ("_vcov_type_explicit",) def summary(self) -> str: """Get summary of estimation results.""" diff --git a/diff_diff/synthetic_control.py b/diff_diff/synthetic_control.py index 731c3a75..7bf7374a 100644 --- a/diff_diff/synthetic_control.py +++ b/diff_diff/synthetic_control.py @@ -40,6 +40,7 @@ import pandas as pd from scipy.optimize import minimize +from diff_diff._base import BaseEstimator from diff_diff.synthetic_control_results import ( SyntheticControlResults, _SyntheticControlFitSnapshot, @@ -77,7 +78,7 @@ def _softmax(theta: np.ndarray) -> np.ndarray: return e / np.sum(e) -class SyntheticControl: +class SyntheticControl(BaseEstimator): """ Classic Synthetic Control Method estimator (Abadie-Diamond-Hainmueller 2010). @@ -229,44 +230,7 @@ def _validate_config(self) -> None: if self.v_cv_t0 < 1: raise ValueError(f"v_cv_t0 must be >= 1, got {self.v_cv_t0!r}") - def get_params(self) -> Dict[str, Any]: - """Get estimator parameters.""" - return { - "v_method": self.v_method, - "custom_v": self.custom_v, - "optimizer_options": self.optimizer_options, - "n_starts": self.n_starts, - "inner_max_iter": self.inner_max_iter, - "inner_min_decrease": self.inner_min_decrease, - "standardize": self.standardize, - "alpha": self.alpha, - "seed": self.seed, - "v_cv_t0": self.v_cv_t0, - } - - def set_params(self, **params) -> "SyntheticControl": - """Set estimator parameters. - - Applies updates transactionally: if ``_validate_config()`` rejects the - post-update state, the instance is rolled back to its pre-call values so - a raised ``ValueError`` leaves the object consistent. - """ - _rollback: Dict[str, Any] = {} - for key in params: - if hasattr(self, key): - _rollback[key] = getattr(self, key) - try: - for key, value in params.items(): - if hasattr(self, key): - setattr(self, key, value) - else: - raise ValueError(f"Unknown parameter: {key}") - self._validate_config() - except (ValueError, TypeError): - for key, prev in _rollback.items(): - setattr(self, key, prev) - raise - return self + # get_params/set_params come from BaseEstimator. # ========================================================================= # fit diff --git a/diff_diff/synthetic_did.py b/diff_diff/synthetic_did.py index 5a1b8acd..7f7978cb 100644 --- a/diff_diff/synthetic_did.py +++ b/diff_diff/synthetic_did.py @@ -235,6 +235,18 @@ def __init__( self.variance_method = variance_method self.n_bootstrap = n_bootstrap self.seed = seed + # Raw-arg storage for the introspected param surface (BaseEstimator). + # conley_metric/conley_kernel would otherwise read back the DiD + # defaults stored by super().__init__ ("haversine"/"bartlett"), + # poisoning every set_params probe against the non-None TypeError + # guard above; both are None in any legal construction. The + # deprecated lambda_reg/zeta are warn-and-discard: stored normalized + # to None so a probe re-init warns exactly once for the call that + # passes them and never again on unrelated set_params/clone calls. + self.conley_metric = conley_metric + self.conley_kernel = conley_kernel + self.lambda_reg = None + self.zeta = None self._validate_config() @@ -2724,103 +2736,11 @@ def _jackknife_se_survey( variance_nonneg = max(total_variance, 0.0) return float(np.sqrt(variance_nonneg)), tau_loo_arr - def get_params(self) -> Dict[str, Any]: - """Get estimator parameters.""" - return { - "zeta_omega": self.zeta_omega, - "zeta_lambda": self.zeta_lambda, - "alpha": self.alpha, - "variance_method": self.variance_method, - "n_bootstrap": self.n_bootstrap, - "seed": self.seed, - # Conley kwargs are inherited from DifferenceInDifferences.__init__ - # but rejected by SyntheticDiD's __init__ / set_params (Conley uses - # the analytical sandwich, SyntheticDiD uses bootstrap variance). - # Surface them here as None for sklearn-style API consistency; any - # non-None value is rejected by set_params/__init__. - "vcov_type": None, - "conley_coords": None, - "conley_cutoff_km": None, - "conley_metric": None, - "conley_kernel": None, - "conley_lag_cutoff": None, - } - - def set_params(self, **params) -> "SyntheticDiD": - """Set estimator parameters. - - Applies updates transactionally: if ``_validate_config()`` rejects the - post-update state, the instance is rolled back to the pre-call values - so a raised ``ValueError`` leaves the object consistent with its - pre-call configuration. - - Mirrors ``__init__``'s defensive rejection of ``vcov_type`` / - ``conley_*`` non-None values: SyntheticDiD uses bootstrap/jackknife/ - placebo variance, not the analytical sandwich, so any Conley kwarg - would be silently ignored otherwise (forbidden by - ``feedback_no_silent_failures``). Tracked in DEFERRED.md for a follow-up - that wires Conley to a non-bootstrap variance path. - """ - # Reject Conley kwargs / non-None vcov_type before any mutation — - # mirrors __init__'s contract. Empty/None values are permitted so - # round-tripping get_params() back through set_params() is a no-op. - _conley_keys = ( - "conley_coords", - "conley_cutoff_km", - "conley_metric", - "conley_kernel", - "conley_lag_cutoff", - ) - if params.get("vcov_type") is not None and params["vcov_type"] != "conley": - raise TypeError( - f"SyntheticDiD does not accept vcov_type={params['vcov_type']!r}. " - "SyntheticDiD's variance is bootstrap/jackknife/placebo based; " - "configure via variance_method=..." - ) - if params.get("vcov_type") == "conley" or any( - k in params and params[k] is not None for k in _conley_keys - ): - raise TypeError( - "SyntheticDiD does not yet support vcov_type='conley' or any " - "conley_* kwargs. SyntheticDiD uses bootstrap/jackknife/placebo " - "variance (variance_method=...), not the analytical sandwich " - "routed through compute_robust_vcov. Tracked in DEFERRED.md as " - "a follow-up." - ) - # Deprecated parameter names — emit warning and ignore - _deprecated = {"lambda_reg", "zeta"} - # Conley kwargs are not stored as instance attributes; surfacing them - # in get_params() returns None unconditionally. set_params() with None - # values for these keys is a no-op (the rejection above only fires on - # non-None values). - _silent_conley_passthrough = {"vcov_type", *_conley_keys} - # Snapshot original values for transactional rollback on validation failure. - _rollback: Dict[str, Any] = {} - for key in params: - if key in _silent_conley_passthrough: - continue - if key not in _deprecated and hasattr(self, key): - _rollback[key] = getattr(self, key) - try: - for key, value in params.items(): - if key in _silent_conley_passthrough: - # No-op: explicitly None passthrough for round-trip - # get_params() -> set_params() consistency. - continue - if key in _deprecated: - warnings.warn( - f"{key} is deprecated and ignored. Use zeta_omega/zeta_lambda " - f"instead. Will be removed in v4.0.0.", - DeprecationWarning, - stacklevel=2, - ) - elif hasattr(self, key): - setattr(self, key, value) - else: - raise ValueError(f"Unknown parameter: {key}") - self._validate_config() - except (ValueError, TypeError): - for key, prev in _rollback.items(): - setattr(self, key, prev) - raise - return self + # get_params/set_params come from BaseEstimator (via + # DifferenceInDifferences). The probe re-init reproduces every guard in + # __init__: the conley/vcov_type non-None TypeError, the lambda_reg/zeta + # DeprecationWarnings (emitted exactly once per call that passes them - + # the raw-arg storage in __init__ normalizes both to None, so unrelated + # set_params/clone calls stay silent), and _validate_config. The + # inherited vcov_type -> _vcov_type_arg alias makes get_params report + # the raw (always-None) vcov_type, not the resolved "hc1". diff --git a/diff_diff/triple_diff.py b/diff_diff/triple_diff.py index 1d08f5ee..e35be170 100644 --- a/diff_diff/triple_diff.py +++ b/diff_diff/triple_diff.py @@ -34,6 +34,7 @@ import numpy as np import pandas as pd +from diff_diff._base import BaseEstimator from diff_diff.linalg import _rank_guarded_inv, solve_logit, solve_ols from diff_diff.results import _format_survey_block, _get_significance_stars from diff_diff.results_base import BaseResults @@ -377,7 +378,7 @@ def epv_summary(self, show_all: bool = False) -> pd.DataFrame: # ============================================================================= -class TripleDifference: +class TripleDifference(BaseEstimator): """ Triple Difference (DDD) estimator. @@ -551,8 +552,9 @@ def __init__( # TripleDifference's 3-pairwise-DiD influence-function decomposition # doesn't have. See REGISTRY.md "IF-based variance estimators vs # analytical-sandwich estimators" for the structural taxonomy. - # Factored out so fit() can re-run it after sklearn-style - # set_params bypasses __init__ validation. + # Factored out so fit() can re-run it: set_params now validates + # eagerly via the BaseEstimator probe re-init, but DIRECT attribute + # mutation (est.vcov_type = ...) still bypasses validation until fit. self._validate_vcov_type(vcov_type) self.estimation_method = estimation_method @@ -621,10 +623,10 @@ def fit( NotImplementedError If survey_design is used with wild_bootstrap inference. """ - # Re-validate vcov_type at fit-time so sklearn-style set_params - # mutations are caught before they propagate to Results metadata. - # __init__ already validated the constructor argument; this is the - # second layer for the post-construction mutation path. + # Re-validate vcov_type at fit-time: __init__ and set_params (via + # the BaseEstimator probe re-init) both validate eagerly, so this + # second layer only catches DIRECT attribute mutation + # (est.vcov_type = ...) before it propagates to Results metadata. self._validate_vcov_type(self.vcov_type) # Reset replicate state from any previous fit @@ -2059,46 +2061,7 @@ def _validate_vcov_type(vcov_type: str) -> None: "IF-based variance structure; see REGISTRY.md." ) - def get_params(self) -> Dict[str, Any]: - """ - Get estimator parameters (sklearn-compatible). - - Returns - ------- - Dict[str, Any] - Estimator parameters. - """ - return { - "estimation_method": self.estimation_method, - "robust": self.robust, - "cluster": self.cluster, - "vcov_type": self.vcov_type, - "alpha": self.alpha, - "pscore_trim": self.pscore_trim, - "rank_deficient_action": self.rank_deficient_action, - "epv_threshold": self.epv_threshold, - "pscore_fallback": self.pscore_fallback, - } - - def set_params(self, **params) -> "TripleDifference": - """ - Set estimator parameters (sklearn-compatible). - - Parameters - ---------- - **params - Estimator parameters. - - Returns - ------- - self - """ - for key, value in params.items(): - if hasattr(self, key): - setattr(self, key, value) - else: - raise ValueError(f"Unknown parameter: {key}") - return self + # get_params/set_params come from BaseEstimator. def summary(self) -> str: """ diff --git a/diff_diff/trop.py b/diff_diff/trop.py index 968b8adf..1e80b3b1 100644 --- a/diff_diff/trop.py +++ b/diff_diff/trop.py @@ -19,7 +19,7 @@ import logging import warnings -from typing import Any, Dict, List, Optional, Tuple +from typing import Dict, List, Optional, Tuple import numpy as np import pandas as pd @@ -30,6 +30,7 @@ HAS_RUST_BACKEND, _rust_loocv_grid_search, ) +from diff_diff._base import BaseEstimator from diff_diff.trop_global import TROPGlobalMixin from diff_diff.trop_local import ( TROPLocalMixin, @@ -44,7 +45,7 @@ from diff_diff.utils import safe_inference, warn_if_not_converged -class TROP(TROPLocalMixin, TROPGlobalMixin): +class TROP(TROPLocalMixin, TROPGlobalMixin, BaseEstimator): """ Triply Robust Panel (TROP) estimator. @@ -946,33 +947,7 @@ def fit( # sklearn-like API # ========================================================================= - def get_params(self) -> Dict[str, Any]: - """Get estimator parameters.""" - return { - "method": self.method, - "lambda_time_grid": self.lambda_time_grid, - "lambda_unit_grid": self.lambda_unit_grid, - "lambda_nn_grid": self.lambda_nn_grid, - "max_iter": self.max_iter, - "tol": self.tol, - "alpha": self.alpha, - "n_bootstrap": self.n_bootstrap, - "seed": self.seed, - "non_absorbing": self.non_absorbing, - } - - def set_params(self, **params) -> "TROP": - """Set estimator parameters.""" - for key, value in params.items(): - if key == "method" and value not in ("local", "global"): - raise ValueError(f"method must be one of ('local', 'global'), got '{value}'") - if key == "non_absorbing" and not isinstance(value, bool): - raise ValueError(f"non_absorbing must be a bool, got {type(value).__name__}") - if hasattr(self, key): - setattr(self, key, value) - else: - raise ValueError(f"Unknown parameter: {key}") - return self + # get_params/set_params come from BaseEstimator. def trop( diff --git a/diff_diff/two_stage.py b/diff_diff/two_stage.py index e3c73d88..19aa1f6a 100644 --- a/diff_diff/two_stage.py +++ b/diff_diff/two_stage.py @@ -29,6 +29,7 @@ from scipy import sparse from scipy.sparse.linalg import factorized as sparse_factorized +from diff_diff._base import BaseEstimator from diff_diff.conley import ( ConleyMetric, _compute_conley_meat, @@ -1290,7 +1291,7 @@ def _compute_stratified_serial_bartlett_meat( # ============================================================================= -class TwoStageDiD(TwoStageDiDBootstrapMixin): +class TwoStageDiD(TwoStageDiDBootstrapMixin, BaseEstimator): """ Gardner (2022) two-stage Difference-in-Differences estimator. @@ -1478,9 +1479,9 @@ def fit( ValueError If required columns are missing or data validation fails. """ - # Re-validate vcov_type at fit-time so sklearn-style set_params - # mutations (e.g. set_params(vcov_type="classical")) are re-checked - # rather than silently accepted by the attribute setter. + # Re-validate vcov_type at fit-time: set_params validates eagerly + # (BaseEstimator probe re-init), so this only catches DIRECT + # attribute mutation (est.vcov_type = ...). self._validate_vcov_type(self.vcov_type) # ---- Data validation ---- @@ -3454,29 +3455,7 @@ def _validate_vcov_type(vcov_type: str) -> None: f"Accepted: {sorted(_accepted_vcov)}." ) - def get_params(self) -> Dict[str, Any]: - """Get estimator parameters (sklearn-compatible).""" - return { - "anticipation": self.anticipation, - "alpha": self.alpha, - "cluster": self.cluster, - "vcov_type": self.vcov_type, - "n_bootstrap": self.n_bootstrap, - "bootstrap_weights": self.bootstrap_weights, - "seed": self.seed, - "rank_deficient_action": self.rank_deficient_action, - "horizon_max": self.horizon_max, - "pretrends": self.pretrends, - } - - def set_params(self, **params) -> "TwoStageDiD": - """Set estimator parameters (sklearn-compatible).""" - for key, value in params.items(): - if hasattr(self, key): - setattr(self, key, value) - else: - raise ValueError(f"Unknown parameter: {key}") - return self + # get_params/set_params come from BaseEstimator. def summary(self) -> str: """Get summary of estimation results.""" diff --git a/diff_diff/wooldridge.py b/diff_diff/wooldridge.py index 664c5398..72d6822b 100644 --- a/diff_diff/wooldridge.py +++ b/diff_diff/wooldridge.py @@ -20,6 +20,7 @@ import numpy as np import pandas as pd +from diff_diff._base import BaseEstimator from diff_diff.linalg import ( compute_robust_vcov, effective_cluster_count, @@ -828,7 +829,7 @@ def _prepare_covariates( return np.hstack([p if p.ndim == 2 else p.reshape(-1, 1) for p in parts]) -class WooldridgeDiD: +class WooldridgeDiD(BaseEstimator): """Extended Two-Way Fixed Effects (ETWFE) DiD estimator. Implements the Wooldridge (2025) saturated cohort×time regression @@ -1092,66 +1093,8 @@ def results_(self) -> WooldridgeDiDResults: raise RuntimeError("Call fit() before accessing results_") return self._results - def get_params(self) -> Dict[str, Any]: - """Return estimator parameters (sklearn-compatible).""" - return { - "method": self.method, - "control_group": self.control_group, - "anticipation": self.anticipation, - "demean_covariates": self.demean_covariates, - "alpha": self.alpha, - "cluster": self.cluster, - "n_bootstrap": self.n_bootstrap, - "bootstrap_weights": self.bootstrap_weights, - "seed": self.seed, - "rank_deficient_action": self.rank_deficient_action, - "vcov_type": self.vcov_type, - "cohort_trends": self.cohort_trends, - "conley_coords": self.conley_coords, - "conley_cutoff_km": self.conley_cutoff_km, - "conley_metric": self.conley_metric, - "conley_kernel": self.conley_kernel, - "conley_lag_cutoff": self.conley_lag_cutoff, - "df_convention": self.df_convention, - } - - def set_params(self, **params: Any) -> "WooldridgeDiD": - """Set estimator parameters (sklearn-compatible). Returns self. - - Atomic: if validation rejects the incoming combination (unknown - parameter, invalid value, or the ``method`` × ``vcov_type`` - interaction guard fires), ``self`` is unchanged so a caller that - catches ``ValueError`` / ``NotImplementedError`` can keep using - the estimator with its previous configuration. Mirrors the - ``DifferenceInDifferences.set_params`` pattern at - ``estimators.py:995-1023``. - """ - # First pass: validate all incoming keys are known attributes so - # we don't partially apply a batch that ends in "Unknown parameter". - for key in params: - if not hasattr(self, key): - raise ValueError(f"Unknown parameter: {key!r}") - - # Compute pending values by overlaying ``params`` on the current - # configuration; validate on those locals (catches invalid sets + - # the method × vcov_type interaction) BEFORE mutating ``self``. - pending = { - "method": params.get("method", self.method), - "control_group": params.get("control_group", self.control_group), - "anticipation": params.get("anticipation", self.anticipation), - "bootstrap_weights": params.get("bootstrap_weights", self.bootstrap_weights), - "vcov_type": params.get("vcov_type", self.vcov_type), - "cohort_trends": params.get("cohort_trends", self.cohort_trends), - "df_convention": params.get("df_convention", self.df_convention), - } - self._validate_constructor_args(**pending) - - # All validation passed — apply mutations atomically. - for key, value in params.items(): - setattr(self, key, value) - # Recompute the explicit-vcov flag after any vcov_type mutation. - self._vcov_type_explicit = self.vcov_type != "hc1" - return self + # get_params/set_params come from BaseEstimator. + _DERIVED_CONFIG_ATTRS = ("_vcov_type_explicit",) def fit( self, diff --git a/docs/api/changes_in_changes.rst b/docs/api/changes_in_changes.rst index 1d28c0b8..b2070103 100644 --- a/docs/api/changes_in_changes.rst +++ b/docs/api/changes_in_changes.rst @@ -77,6 +77,7 @@ Main changes-in-changes estimator class. :members: :undoc-members: :show-inheritance: + :inherited-members: .. rubric:: Methods @@ -96,6 +97,7 @@ Quantile difference-in-differences comparison estimator. :members: :undoc-members: :show-inheritance: + :inherited-members: .. rubric:: Methods diff --git a/docs/api/had.rst b/docs/api/had.rst index 040fbba8..6e42be85 100644 --- a/docs/api/had.rst +++ b/docs/api/had.rst @@ -127,6 +127,7 @@ HeterogeneousAdoptionDiD :members: :undoc-members: :show-inheritance: + :inherited-members: HeterogeneousAdoptionDiDResults ------------------------------- diff --git a/docs/api/honest_did.rst b/docs/api/honest_did.rst index 6fd000cd..a02b148f 100644 --- a/docs/api/honest_did.rst +++ b/docs/api/honest_did.rst @@ -31,6 +31,7 @@ Main class for computing honest bounds and confidence intervals. :members: :undoc-members: :show-inheritance: + :inherited-members: .. rubric:: Methods diff --git a/docs/api/lpdid.rst b/docs/api/lpdid.rst index 71fb404f..48fbf725 100644 --- a/docs/api/lpdid.rst +++ b/docs/api/lpdid.rst @@ -75,6 +75,7 @@ Main estimator class for Local Projections Difference-in-Differences. :members: :undoc-members: :show-inheritance: + :inherited-members: .. rubric:: Methods diff --git a/docs/api/pretrends.rst b/docs/api/pretrends.rst index 61a2716e..0407e7b6 100644 --- a/docs/api/pretrends.rst +++ b/docs/api/pretrends.rst @@ -32,6 +32,7 @@ Main class for pre-trends power analysis. :members: :undoc-members: :show-inheritance: + :inherited-members: .. rubric:: Methods diff --git a/docs/api/regression_discontinuity.rst b/docs/api/regression_discontinuity.rst index 362cc8a4..84f6caae 100644 --- a/docs/api/regression_discontinuity.rst +++ b/docs/api/regression_discontinuity.rst @@ -68,6 +68,7 @@ RegressionDiscontinuity :members: :undoc-members: :show-inheritance: + :inherited-members: RegressionDiscontinuityResults ------------------------------ @@ -86,6 +87,7 @@ RDPlot :members: :undoc-members: :show-inheritance: + :inherited-members: RDPlotResult ------------ diff --git a/docs/api/staggered.rst b/docs/api/staggered.rst index 8bf5fbd0..78542cb7 100644 --- a/docs/api/staggered.rst +++ b/docs/api/staggered.rst @@ -142,6 +142,7 @@ with group-time ATT identification under heterogeneous treatment timing. :members: :undoc-members: :show-inheritance: + :inherited-members: StaggeredTripleDiffResults -------------------------- diff --git a/docs/api/wooldridge_etwfe.rst b/docs/api/wooldridge_etwfe.rst index faf72ad4..440c79a9 100644 --- a/docs/api/wooldridge_etwfe.rst +++ b/docs/api/wooldridge_etwfe.rst @@ -48,6 +48,7 @@ Main estimator class for Wooldridge ETWFE. :members: :undoc-members: :show-inheritance: + :inherited-members: .. rubric:: Methods diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index 09da6b83..d381ed61 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -1519,7 +1519,7 @@ where `q_{g,e} = pi_g / sum_{g' in G_{trt,e}} pi_{g'}`. - **Note:** Last-cohort-as-control (`control_group="last_cohort"`) reclassifies the latest treatment cohort as pseudo-never-treated and drops time periods at `t >= last_g - anticipation`, excluding anticipation-contaminated periods from the pseudo-control's pre-treatment window. This is distinct from CallawaySantAnna's `not_yet_treated` option which dynamically selects not-yet-treated units per (g,t) pair. - **Note:** `vcov_type` is permanently narrow to `{"hc1"}` per the Chen-Sant'Anna-Xie (2025) EIF-based variance achieving the semiparametric efficiency bound. Analytical-sandwich families `{classical, hc2, hc2_bm}` are rejected at `__init__` — the per-unit EIF aggregation has no equivalent single design matrix on which hat-matrix leverage or Bell-McCaffrey Satterthwaite DOF can be defined. `cluster=` invokes Liang-Zeger CR1 on cluster-aggregated EIF (`_compute_se_from_eif` with `cluster_indices` at `diff_diff/efficient_did.py:124-127`); `survey_design=` invokes TSL on the combined IF (`_compute_survey_eif_se` at `diff_diff/efficient_did.py:1151-1176`). `vcov_type='conley'` deferred to the EfficientDiD Conley follow-up row in DEFERRED.md. - **Note:** Default `cluster=None` (no survey design) renders summary label "HC1 heteroskedasticity-robust" because the per-unit EIF SE `sqrt(mean(EIF²)/n)` is methodologically HC1-style (no Liang-Zeger G/(G-1) finite-sample correction). `EfficientDiDResults.cluster_name` and `n_clusters` stay None under unclustered fits. This diverges from `ImputationDiD` which auto-clusters at unit per Borusyak-Jaravel-Spiess (2024) Theorem 3 — there the default summary renders the CR1 unit-clustered label. -- **Note (deviation from sibling estimators):** `EfficientDiD.set_params(vcov_type=bad)` raises immediately rather than deferring to `fit()`. EfficientDiD's `set_params` calls `_validate_params()` which invokes `_validate_vcov_type`, matching the existing eager-validation pattern for `pt_assumption`, `control_group`, `bootstrap_weights`, etc. This is intentional — `ImputationDiD`/`TripleDifference`/`CallawaySantAnna` use sklearn mutate-then-validate-at-use, so the same set_params + bad vcov_type sequence is silently accepted there until `fit()` is called. +- **Note:** `set_params(vcov_type=bad)` raises immediately on EVERY estimator: since the shared `BaseEstimator` mixin (`diff_diff/_base.py`, v4 2(c)-i), `set_params` validates transactionally by constructor probe re-init, so it enforces exactly `__init__`'s validation, eagerly, library-wide. The former split — EfficientDiD eager vs `ImputationDiD`/`TripleDifference`/`CallawaySantAnna` (and six more: SunAbraham, StackedDiD, StaggeredTripleDifference, SpilloverDiD, TROP, PreTrendsPower) accepting constructor-rejected values until `fit()` — is retired; the fit-time re-validation layers remain as a second check against DIRECT attribute mutation (`est.vcov_type = ...`), which no setter can see. --- diff --git a/docs/methodology/papers/athey-imbens-2006-review.md b/docs/methodology/papers/athey-imbens-2006-review.md index 3be9cabb..007e8c64 100644 --- a/docs/methodology/papers/athey-imbens-2006-review.md +++ b/docs/methodology/papers/athey-imbens-2006-review.md @@ -376,7 +376,7 @@ Aggregation `tau_Lambda = Lambda' tau^CIC_I` (columns of Lambda sum to 1): sampl ### Relation to Existing diff-diff Estimators - `DifferenceInDifferences` is the nested special case (eqs. 1, 4-6): when the linear model holds with full independence, CiC and DiD probability limits coincide (p. 463) - a natural cross-check test. Note neither estimator dominates on efficiency: `Var(tau_hat^CIC)` can equal, exceed, or fall below `Var(tau_hat^DID)` (p. 463; footnote 30's scale example `Y_g1 ~d sigma*Y_g0`: CiC strictly more efficient iff `sigma^2 < 1`). - QDiD's mean effect reproduces the standard DiD ATT under continuity (p. 447) - reuse `DifferenceInDifferences` output as an oracle in tests. -- Reuse: 2x2 cell validation patterns from `DifferenceInDifferences`; bootstrap utilities and `safe_inference()` joint-NaN contract from `diff_diff.utils`; results-dataclass conventions (flat `att`/`se`/`conf_int` family). CiC is a standalone estimator (own `get_params`/`set_params` per the estimator-inheritance map in CLAUDE.md). +- Reuse: 2x2 cell validation patterns from `DifferenceInDifferences`; bootstrap utilities and `safe_inference()` joint-NaN contract from `diff_diff.utils`; results-dataclass conventions (flat `att`/`se`/`conf_int` family). CiC is a standalone estimator (mixing in `BaseEstimator` from `diff_diff/_base.py` for `get_params`/`set_params` per the estimator-inheritance map in CLAUDE.md). - The wage-decomposition estimand `E[Y_B1] - E[F_Y,W1^{-1}(F_Y,W0(Y_B0))]` (Altonji-Blank / Juhn-Murphy-Pierce, Section 3.5, pp. 448-449) is algebraically identical to `tau^CIC`, so the same code serves decomposition use cases (interpretation differs). - The "selection on observables" panel alternative `F_{Y^N,11}(y) = E[F_{Y_01|Y_00}(y | Y_10)]` generally differs from the CiC counterfactual; they coincide iff `U_i0 = U_i1` (perfect rank correlation over time, p. 448) - do not confuse CiC with conditioning-on-lagged-outcome estimators. - Future: the Section 6 multi-group machinery (overidentified `kappa` quadruples + minimum-distance combination + Theorem 6.4 specification test) is structurally similar to `CallawaySantAnna`'s (g,t) cell architecture; the Appendix B shared-cell covariance rule would slot into an influence-function aggregation layer. diff --git a/docs/methodology/papers/melly-santangelo-2015-review.md b/docs/methodology/papers/melly-santangelo-2015-review.md index 0c0e6ede..96fcccb9 100644 --- a/docs/methodology/papers/melly-santangelo-2015-review.md +++ b/docs/methodology/papers/melly-santangelo-2015-review.md @@ -493,7 +493,7 @@ All quantiles and covariate values must be considered to detect deviations; Kolm - `ChangesInChanges` (v1): the exact special case X = constant. The paper's remark 3 on p. 18 notes that even without covariates it contributes the limiting distribution of the whole quantile and distribution processes with a simpler proof than AI - a useful citation when documenting process-level (uniform-band) inference for the v1 estimator. - `QDiD` (v1): the pp. 22-23 DiD-on-indicators analysis is the paper's cautionary tale about additive-in-distribution shortcuts for distributional targets; relevant background for the choosing-an-estimator docs alongside the AI (2006) QDiD discussion, though the object analyzed (DiD applied to `1(Y <= y)`) is not the QDiD transform itself. - Callaway-Li-Oka (2018) review (sibling, same initiative): shared inference stack - exchangeable bootstrap weights, functional delta method over Hadamard-differentiable ECDF/quantile compositions, KS-based uniform bands. Both ultimately lean on Van der Vaart-Wellner chapter 3.9; this paper routes through CFM (2013) Corollary 5.2 for the QR first stage. -- A future covariates implementation should reuse the `safe_inference()` joint-NaN contract from `diff_diff.utils`, the results-dataclass conventions, and the bootstrap utilities, and would be a standalone estimator (own `get_params`/`set_params` per the estimator-inheritance map in CLAUDE.md). +- A future covariates implementation should reuse the `safe_inference()` joint-NaN contract from `diff_diff.utils`, the results-dataclass conventions, and the bootstrap utilities, and would be a standalone estimator (mixing in `BaseEstimator` from `diff_diff/_base.py` for `get_params`/`set_params` per the estimator-inheritance map in CLAUDE.md). --- diff --git a/docs/v4-design.md b/docs/v4-design.md index c5454183..a6d618a7 100644 --- a/docs/v4-design.md +++ b/docs/v4-design.md @@ -517,9 +517,11 @@ domain vocabulary, not drift. suites' expectations and adds the migration-guide entry. - Constructor hygiene rides Phase 2: ContinuousDiD's `covariates` moves from `__init__` to `fit()` [M-084] (the one estimator with a data column in the - constructor), and the shared `BaseEstimator` mixin replaces the 24 + constructor), and the shared `BaseEstimator` mixin replaces the 25 hand-rolled `get_params`/`set_params` copies (transactional set_params per - the locked rule; `deep=` supported uniformly). + the locked rule; `deep=` supported uniformly; the implementation measured + 25 defining classes - `changes_in_changes.py` holds two - correcting this + document's earlier count of 24). ## 8. Contract-rename rules @@ -654,15 +656,15 @@ re-enumerate the cells' M-id lists: which carries the amended spec - surface sweep + phase-table agreement + consumer coverage - in its module docstring) - lands BEFORE 2(c), so the rename PR works from a mechanically-verified list. -3. 2(c)-i: the BaseEstimator mixin, front-loaded. Scope is section 7's - normative statement verbatim: replace the 24 hand-rolled +3. 2(c)-i: the BaseEstimator mixin, front-loaded (shipped: + `diff_diff/_base.py`, with the cross-estimator contract suite + `tests/test_base_estimator.py` on a dynamic roster). Scope was section 7's + normative statement verbatim: replace the hand-rolled `get_params`/`set_params` copies library-wide (3.9-cut checklist item 1) - - not merely the standalone estimator classes. Rationale: the tail-df PR's - review exposed the non-atomic `set_params` pattern and fixed the five - estimators in its scope; others remain non-atomic today (verified examples: - CallawaySantAnna, TwoStageDiD, TripleDifference, TROP, SpilloverDiD - the - mixin PR's first task is the exhaustive inventory of all 24 copies), and - the renames should build on the transactional contract. + not merely the standalone estimator classes; the exhaustive inventory + found 25 defining classes (not 24). set_params is transactional via probe + re-init - `type(self)(**merged)` validates before any mutation - so the + renames build on a contract that can never drift from `__init__`. 4. 2(c)-ii: the rename sweep (the phase-2(c) cell); may split by rename group. 5. 2(b): post-fit `aggregate()` + `fit(aggregate=)` shims (the (b) cell), claiming reserved ids M-116/M-118..M-121 for any new rows; the @@ -685,7 +687,10 @@ not expressible as ledger rows, so the 3.9 release PR asserts them by hand: 1. The shared `BaseEstimator` mixin has replaced the hand-rolled `get_params`/`set_params` copies (section 7), with `deep=` uniform and set_params transactional per the locked rule. A pure refactor - no symbol - changes - so no row can see it. + changes - so no row can see it. DONE: `diff_diff/_base.py` + + `tests/test_base_estimator.py` (25 classes converted; the ten + formerly-lazy set_params surfaces now validate eagerly, REGISTRY + EfficientDiD note updated in the same diff). 2. The R-equivalents mapping table (section 8 rule 8) ships in the docs. 3. The migration guide exists (section 10) and its TL;DR table has a row per breaking change, generated against the matrix rather than hand-listed. diff --git a/tests/test_base_estimator.py b/tests/test_base_estimator.py new file mode 100644 index 00000000..978c51c3 --- /dev/null +++ b/tests/test_base_estimator.py @@ -0,0 +1,277 @@ +"""Cross-estimator contract suite for the shared BaseEstimator mixin. + +Every estimator's ``get_params``/``set_params`` comes from +``diff_diff._base.BaseEstimator`` (v4-design section 7, 2(c)-i). This suite +pins the shared contract ONCE, parametrized over a DYNAMICALLY discovered +roster (never a hand-listed set, so a future estimator cannot silently skip +the contract): + +- roster completeness: every ``diff_diff.__all__`` class with a public + ``fit`` method must have ``BaseEstimator`` in its MRO (explicit exclusion + tuple below, reasons attached); +- init-signature <-> get_params sync; +- ``cls(**est.get_params())`` re-instantiation round-trip; +- strict unknown-key rejection (method names and private attrs included) + with batch atomicity; +- value-level transactional rollback (constructor-validated classes); +- ``deep=`` triple-equality; returns-self; +- ``set_params`` == fresh-construction equivalence, stated POST the + ``_normalize_set_params`` hook (so DiD's robust-alone re-derivation is + inside the contract, not an exemption); +- fitted-state attrs are never touched by ``set_params``; +- ``sklearn.base.clone`` round-trip under ``importorskip`` (config-equality + is the always-running contract; clone identity is opportunistic local + coverage - scikit-learn is deliberately not a dev dependency). +""" + +import inspect +import warnings + +import pytest + +import diff_diff +from diff_diff._base import BaseEstimator + +# --------------------------------------------------------------------------- +# Dynamic roster discovery +# --------------------------------------------------------------------------- + +# __all__ classes with a public `fit` that deliberately do NOT adopt the +# estimator param contract. Reasons required. +NON_PARTICIPANTS = { + # Internal OLS engine exported for power users; predates the estimator + # contract and carries fit-time (not constructor) configuration. Adding + # get_params/set_params is a NEW surface tracked in TODO.md, not part of + # the 2(c)-i refactor. + "LinearRegression", +} + + +def _discover(): + mixin, fit_classes, seen = [], [], set() + for name in diff_diff.__all__: + obj = getattr(diff_diff, name) + if not inspect.isclass(obj) or id(obj) in seen: + continue + seen.add(id(obj)) + if issubclass(obj, BaseEstimator): + mixin.append(obj) + fit = inspect.getattr_static(obj, "fit", None) + if fit is not None and ( + inspect.isfunction(fit) or isinstance(fit, (staticmethod, classmethod)) + ): + fit_classes.append(obj) + return mixin, fit_classes + + +MIXIN_CLASSES, FIT_CLASSES = _discover() + +# Required constructor kwargs for classes that cannot default-construct. +DEFAULT_KWARGS = { + "SpilloverDiD": {"rings": [0.0, 50.0, 100.0]}, +} + +# One known constructor-rejected value per class WITH constructor validation +# (value-level rollback lane). Classes absent here are still covered by the +# unknown-key atomicity lane. +BAD_VALUES = { + "DifferenceInDifferences": {"vcov_type": "hc99"}, + "TwoWayFixedEffects": {"vcov_type": "hc99"}, + "MultiPeriodDiD": {"vcov_type": "hc99"}, + "SyntheticDiD": {"variance_method": "not_a_method"}, + "CallawaySantAnna": {"vcov_type": "hc4"}, + "SunAbraham": {"vcov_type": "hc4"}, + "ImputationDiD": {"vcov_type": "hc4"}, + "TwoStageDiD": {"vcov_type": "hc4"}, + "TripleDifference": {"vcov_type": "hc4"}, + "EfficientDiD": {"vcov_type": "hc4"}, + "StackedDiD": {"clean_control": "not_a_mode"}, + "LPDiD": {"alpha": 5.0}, + "ChangesInChanges": {"alpha": 5.0}, + "QDiD": {"alpha": 5.0}, + "HeterogeneousAdoptionDiD": {"design": "not_a_design"}, + "RegressionDiscontinuity": {"kernel": "not_a_kernel"}, + "ChaisemartinDHaultfoeuille": {"cluster": "unit"}, # NotImplementedError gate + "TROP": {"method": "not_a_method"}, + "PreTrendsPower": {"alpha": 5.0}, + "WooldridgeDiD": {"method": "not_a_method"}, + "SpilloverDiD": {"rank_deficient_action": "not_an_action"}, + "StaggeredTripleDifference": {"estimation_method": "not_a_method"}, + "BaconDecomposition": {"weights": "not_a_weighting"}, +} + +# A safe single-param mutation per class for the equivalence/returns-self +# lanes. Defaults to alpha/seed when present. +MUTATIONS = { + "BaconDecomposition": {"weights": "approximate"}, +} + + +def _make(cls): + return cls(**DEFAULT_KWARGS.get(cls.__name__, {})) + + +def _mutation_for(cls, params): + if cls.__name__ in MUTATIONS: + return MUTATIONS[cls.__name__] + if "alpha" in params: + return {"alpha": 0.11} + if "seed" in params: + return {"seed": 1234} + return None + + +def _params_id(cls): + return cls.__name__ + + +estimators = pytest.mark.parametrize("cls", MIXIN_CLASSES, ids=_params_id) + + +# --------------------------------------------------------------------------- +# Roster completeness +# --------------------------------------------------------------------------- + + +def test_roster_discovered_nontrivially(): + # 25 defining classes + TwoWayFixedEffects/MultiPeriodDiD inheritors. + assert len(MIXIN_CLASSES) >= 27 + + +def test_every_fit_class_is_a_base_estimator(): + missing = [ + cls.__name__ + for cls in FIT_CLASSES + if cls.__name__ not in NON_PARTICIPANTS and not issubclass(cls, BaseEstimator) + ] + assert missing == [], ( + f"exported fit-classes without BaseEstimator in their MRO: {missing}; " + "either mix it in or add a reasoned NON_PARTICIPANTS entry" + ) + + +def test_non_participant_entries_are_live(): + # A dead exclusion reads as load-bearing; every entry must name an + # exported fit-class that really lacks the mixin. + fit_names = {cls.__name__ for cls in FIT_CLASSES} + for name in NON_PARTICIPANTS: + assert name in fit_names, f"stale NON_PARTICIPANTS entry: {name}" + assert not issubclass(getattr(diff_diff, name), BaseEstimator) + + +# --------------------------------------------------------------------------- +# The shared contract +# --------------------------------------------------------------------------- + + +@estimators +def test_init_signature_matches_get_params(cls): + est = _make(est_cls := cls) + sig = set(inspect.signature(est_cls.__init__).parameters) - {"self"} + assert set(est.get_params()) == sig + + +@estimators +def test_reinstantiation_round_trip(cls): + est = _make(cls) + params = est.get_params() + with warnings.catch_warnings(): + warnings.simplefilter("error") + clone = cls(**params) + assert clone.get_params() == params + + +@estimators +def test_deep_triple_equality(cls): + est = _make(cls) + assert est.get_params() == est.get_params(deep=True) == est.get_params(deep=False) + + +@estimators +def test_unknown_key_raises_and_batch_is_atomic(cls): + est = _make(cls) + before = est.get_params() + first_key = next(iter(before)) + with pytest.raises(ValueError, match="Unknown parameter"): + est.set_params(**{first_key: before[first_key], "definitely_not_a_param": 1}) + assert est.get_params() == before + + +@estimators +def test_method_names_and_private_attrs_rejected(cls): + est = _make(cls) + with pytest.raises(ValueError, match="Unknown parameter"): + est.set_params(fit=lambda: None) + with pytest.raises(ValueError, match="Unknown parameter"): + est.set_params(_private_attr=42) + + +@estimators +def test_bad_value_rolls_back(cls): + bad = BAD_VALUES.get(cls.__name__) + if bad is None: + pytest.skip("no constructor-validated bad value catalogued") + est = _make(cls) + before = est.get_params() + with pytest.raises((ValueError, TypeError, NotImplementedError)): + est.set_params(**bad) + assert est.get_params() == before + + +@estimators +def test_set_params_returns_self_and_matches_reinit(cls): + est = _make(cls) + params = est.get_params() + mutation = _mutation_for(cls, params) + if mutation is None: + pytest.skip("no safe single-param mutation catalogued") + result = est.set_params(**mutation) + assert result is est + # Post-normalization-hook equivalence: config state equals a fresh + # construction from the merged params. + normalized = type(est)._normalize_set_params(dict(mutation)) + expected = cls(**{**params, **normalized}) + assert est.get_params() == expected.get_params() + for attr in type(est)._DERIVED_CONFIG_ATTRS: + assert getattr(est, attr) == getattr(expected, attr), attr + + +@estimators +def test_fitted_state_untouched_by_set_params(cls): + est = _make(cls) + sentinels = {} + # `results_` is a read-only property on some classes (WooldridgeDiD keeps + # fitted state at `_results`); only instance-dict attrs can host a sentinel. + for attr in ("is_fitted_", "results_", "_results"): + if attr in vars(est): + sentinels[attr] = object() + setattr(est, attr, sentinels[attr]) + if not sentinels: + pytest.skip("class holds no pre-declared fitted-state attrs") + mutation = _mutation_for(cls, est.get_params()) or {} + est.set_params(**mutation) + for attr, sentinel in sentinels.items(): + assert getattr(est, attr) is sentinel, attr + + +@estimators +def test_sklearn_clone_round_trip_if_available(cls): + base = pytest.importorskip("sklearn.base") + est = _make(cls) + cloned = base.clone(est) + assert cloned is not est + assert type(cloned) is type(est) + assert cloned.get_params() == est.get_params() + + +# --------------------------------------------------------------------------- +# DiD robust-alone normalization hook (H1) - the one non-identity hook +# --------------------------------------------------------------------------- + + +def test_did_robust_alone_rederives_vcov_type(): + est = diff_diff.DifferenceInDifferences(vcov_type="hc2") + est.set_params(robust=False) + assert est.vcov_type == "classical" + assert est._vcov_type_arg is None + assert est._vcov_type_explicit is False diff --git a/tests/test_continuous_did.py b/tests/test_continuous_did.py index bfe80e9b..9ba395dd 100644 --- a/tests/test_continuous_did.py +++ b/tests/test_continuous_did.py @@ -274,7 +274,7 @@ def test_set_params(self): def test_set_invalid_param(self): est = ContinuousDiD() - with pytest.raises(ValueError, match="Invalid parameter"): + with pytest.raises(ValueError, match="Unknown parameter"): est.set_params(nonexistent_param=42) diff --git a/tests/test_efficient_did.py b/tests/test_efficient_did.py index 7c8612c7..d2f058ac 100644 --- a/tests/test_efficient_did.py +++ b/tests/test_efficient_did.py @@ -3113,10 +3113,9 @@ def test_bootstrap_survey_hc1_bit_equal(self, ci_params): # ---- Surface 6: set_params eager revalidation ------------------------- def test_set_params_bad_vcov_caught_immediately(self): - # EfficientDiD's set_params calls _validate_params (which now invokes - # _validate_vcov_type), so the check fires NOW (not at fit-time). - # This intentionally diverges from ImputationDiD/TripleDifference - # (which defer to fit-time per sklearn mutate-then-validate-at-use). + # set_params validates eagerly via the BaseEstimator probe re-init - + # the uniform contract across all estimators since the 2(c)-i mixin + # (ImputationDiD/TripleDifference/CallawaySantAnna flipped with it). ed = EfficientDiD() with pytest.raises(ValueError, match="influence-function"): ed.set_params(vcov_type="classical") diff --git a/tests/test_had.py b/tests/test_had.py index 5e983f72..3b27483c 100644 --- a/tests/test_had.py +++ b/tests/test_had.py @@ -854,7 +854,7 @@ def test_set_params_updates_and_returns_self(self): def test_set_params_invalid_key_raises(self): est = HeterogeneousAdoptionDiD() - with pytest.raises(ValueError, match="Invalid parameter"): + with pytest.raises(ValueError, match="Unknown parameter"): est.set_params(not_a_param=True) def test_set_params_rejects_method_names(self): @@ -863,7 +863,7 @@ def test_set_params_rejects_method_names(self): else they would silently overwrite the method. """ est = HeterogeneousAdoptionDiD() - with pytest.raises(ValueError, match="Invalid parameter"): + with pytest.raises(ValueError, match="Unknown parameter"): est.set_params(fit="not_a_method") # sanity: fit is still callable on the class assert callable(est.fit) @@ -871,7 +871,7 @@ def test_set_params_rejects_method_names(self): def test_set_params_rejects_private_attrs(self): """Internal-looking attribute names must also raise.""" est = HeterogeneousAdoptionDiD() - with pytest.raises(ValueError, match="Invalid parameter"): + with pytest.raises(ValueError, match="Unknown parameter"): est.set_params(_internal=42) def test_get_params_accepts_deep_keyword(self): diff --git a/tests/test_imputation.py b/tests/test_imputation.py index 09533c7f..24decca3 100644 --- a/tests/test_imputation.py +++ b/tests/test_imputation.py @@ -2563,33 +2563,22 @@ def test_bootstrap_survey_hc1_bit_equal(self, ci_params, aggregate): for g, se in r_default.bootstrap_results.group_ses.items(): assert se == r_explicit.bootstrap_results.group_ses[g] - # ---- Surface 6: fit()-time revalidation ------------------------------- + # ---- Surface 6: eager transactional validation (BaseEstimator) -------- - def test_set_params_bad_vcov_caught_at_fit_time_classical(self): - data = generate_test_data(seed=11) + def test_set_params_bad_vcov_raises_eagerly_classical(self): + # set_params validates via constructor probe (transactional per the + # locked v4 rule): the bad value raises at set_params and the + # estimator is unchanged. imp = ImputationDiD() - imp.set_params(vcov_type="classical") # mutate-then-validate-at-use with pytest.raises(ValueError, match="influence-function"): - imp.fit( - data, - outcome="outcome", - unit="unit", - time="time", - first_treat="first_treat", - ) + imp.set_params(vcov_type="classical") + assert imp.vcov_type == "hc1" - def test_set_params_bad_vcov_caught_at_fit_time_unknown(self): - data = generate_test_data(seed=11) + def test_set_params_bad_vcov_raises_eagerly_unknown(self): imp = ImputationDiD() - imp.set_params(vcov_type="hc4") with pytest.raises(ValueError, match="hc4"): - imp.fit( - data, - outcome="outcome", - unit="unit", - time="time", - first_treat="first_treat", - ) + imp.set_params(vcov_type="hc4") + assert imp.vcov_type == "hc1" # ---- Surface 7: bootstrap n_psu/n_clusters<2 NaN propagation ---------- diff --git a/tests/test_methodology_imputation.py b/tests/test_methodology_imputation.py index 80d6a1cb..45e54a54 100644 --- a/tests/test_methodology_imputation.py +++ b/tests/test_methodology_imputation.py @@ -751,18 +751,16 @@ def test_loo_single_unit_group_warns_and_returns_finite(self): def test_loo_param_validation_and_roundtrip(self): """leave_one_out is in get/set_params; a non-bool is rejected (TypeError) - in __init__ AND at fit-time (closing the naive-setattr set_params bypass).""" + in __init__ AND eagerly at set_params (the BaseEstimator probe re-init + runs constructor validation before any mutation).""" assert ImputationDiD().get_params()["leave_one_out"] is False assert ImputationDiD(leave_one_out=True).get_params()["leave_one_out"] is True with pytest.raises(TypeError, match="leave_one_out must be a bool"): ImputationDiD(leave_one_out="yes") # type: ignore[arg-type] - # set_params is a naive setattr; the fit-time re-check must catch it - rng = np.random.default_rng(_BASE_SEED_EQ8 + 14) - panel = _make_staggered_panel(rng, cohorts=[3], n_per_cohort=20, n_periods=5) est = ImputationDiD() - est.set_params(leave_one_out="yes") # type: ignore[arg-type] with pytest.raises(TypeError, match="leave_one_out must be a bool"): - est.fit(panel, **self._COMMON) + est.set_params(leave_one_out="yes") # type: ignore[arg-type] + assert est.leave_one_out is False def test_loo_fit_is_idempotent_on_config(self): """Repeated fits with leave_one_out=True give identical SE (no config mutation).""" diff --git a/tests/test_methodology_sdid.py b/tests/test_methodology_sdid.py index 99133867..ff5089a1 100644 --- a/tests/test_methodology_sdid.py +++ b/tests/test_methodology_sdid.py @@ -2083,12 +2083,18 @@ def test_get_params_includes_new_names(self): assert params["zeta_omega"] == 1.0 assert params["zeta_lambda"] == 0.5 - def test_get_params_excludes_old_names(self): - """get_params should NOT include lambda_reg or zeta.""" + def test_get_params_old_names_always_none(self): + """get_params mirrors the __init__ signature, so the deprecated + lambda_reg/zeta appear - but always as None (warn-and-discard is + normalized at construction), keeping cls(**get_params()) silent.""" sdid = SyntheticDiD() params = sdid.get_params() - assert "lambda_reg" not in params - assert "zeta" not in params + assert params["lambda_reg"] is None + assert params["zeta"] is None + with warnings.catch_warnings(): + warnings.simplefilter("error") + clone = SyntheticDiD(**params) + assert clone.get_params() == params def test_set_params_new_names(self): """set_params with new names should work.""" diff --git a/tests/test_premerge_scan.py b/tests/test_premerge_scan.py index e7212315..b5abe07b 100644 --- a/tests/test_premerge_scan.py +++ b/tests/test_premerge_scan.py @@ -132,6 +132,38 @@ def test_check_c_ignores_assignments_outside_init(mod): assert mod.new_params_missing_from_get_params(["real", "other"], text) == [] +def test_check_c_skips_mixin_classes_without_inbody_get_params(mod): + # Post-mixin shape: get_params is inherited (possibly indirectly, e.g. + # SyntheticDiD via DifferenceInDifferences), invisible to the file-local + # AST scan. A new self.X in such a class must NOT be flagged - the + # introspective BaseEstimator.get_params auto-covers new __init__ params + # and tests/test_base_estimator.py enforces the signature sync. + text = ( + "from diff_diff._base import BaseEstimator\n" + "class E(BaseEstimator):\n" + " def __init__(self):\n" + " self.new_param = 1\n" + "class F(SomeOtherBase):\n" # indirect inheritance: base name is opaque + " def __init__(self):\n" + " self.other_param = 2\n" + ) + assert mod.new_params_missing_from_get_params(["new_param", "other_param"], text) == [] + + +def test_check_c_still_flags_handrolled_get_params(mod): + # A class that hand-rolls get_params keeps the original contract: a new + # __init__ self-assign absent from its own get_params is flagged. + text = ( + "class E:\n" + " def __init__(self):\n" + " self.new_param = 1\n" + " self.kept = 2\n" + " def get_params(self):\n" + " return {'kept': self.kept}\n" + ) + assert mod.new_params_missing_from_get_params(["new_param", "kept"], text) == ["new_param"] + + def test_load_groups(mod): yaml = ( "groups:\n" diff --git a/tests/test_staggered.py b/tests/test_staggered.py index e01a6a68..00a83584 100644 --- a/tests/test_staggered.py +++ b/tests/test_staggered.py @@ -4202,31 +4202,15 @@ def test_set_params_pscore_trim(self): cs.set_params(pscore_trim=0.1) assert cs.pscore_trim == 0.1 - def test_set_params_invalid_pscore_trim_rejected_at_fit(self): - """Invalid pscore_trim via set_params() raises ValueError at fit().""" - np.random.seed(42) - n_units, n_periods = 50, 6 - units = np.repeat(np.arange(n_units), n_periods) - times = np.tile(np.arange(n_periods), n_units) - first_treat = np.zeros(n_units) - first_treat[n_units // 2 :] = 3 - first_treat_expanded = np.repeat(first_treat, n_periods) - post = (times >= first_treat_expanded) & (first_treat_expanded > 0) - outcomes = 1.0 + 2.0 * post + np.random.randn(len(units)) * 0.5 - data = pd.DataFrame( - { - "unit": units, - "time": times, - "outcome": outcomes, - "first_treat": first_treat_expanded.astype(int), - } - ) - + def test_set_params_invalid_pscore_trim_rejected_eagerly(self): + """Invalid pscore_trim raises AT set_params (BaseEstimator probe + re-init runs constructor validation transactionally); the + estimator is unchanged.""" for bad_val in [0.0, -0.1, 0.5]: cs = CallawaySantAnna(estimation_method="ipw") - cs.set_params(pscore_trim=bad_val) with pytest.raises(ValueError, match="pscore_trim must be in"): - cs.fit(data, outcome="outcome", unit="unit", time="time", first_treat="first_treat") + cs.set_params(pscore_trim=bad_val) + assert cs.pscore_trim == 0.01 def test_default_pscore_trim(self): """Default pscore_trim is 0.01.""" @@ -5514,25 +5498,15 @@ def test_get_params_includes_vcov_type(self): assert "vcov_type" in params assert params["vcov_type"] == "hc1" - def test_set_params_bad_vcov_caught_at_fit_time(self): - """set_params is strict-mirror SA (no atomic validation), but - fit() re-validates so a bad set_params(vcov_type='hc4') - surfaces a clear error at fit-time rather than silently - propagating a bad value to Results metadata.""" + def test_set_params_bad_vcov_raises_eagerly(self): + """set_params validates via constructor probe (transactional per + the locked v4 rule): a bad vcov_type raises AT set_params with the + same message __init__ gives, and the estimator is unchanged. The + fit-time re-validation stays in place as belt-and-suspenders.""" cs = CallawaySantAnna() - # set_params succeeds (sklearn-style mutate-then-validate-at-use) - cs.set_params(vcov_type="hc4") - assert cs.vcov_type == "hc4" - # fit() re-validates and raises - data = _generate_clustered_staggered_data(seed=37) with pytest.raises(ValueError, match="hc4"): - cs.fit( - data, - outcome="outcome", - unit="unit", - time="time", - first_treat="first_treat", - ) + cs.set_params(vcov_type="hc4") + assert cs.vcov_type == "hc1" def test_results_carries_vcov_type(self): data = _generate_clustered_staggered_data(seed=41) diff --git a/tests/test_triple_diff.py b/tests/test_triple_diff.py index 163d70ce..ed76a86c 100644 --- a/tests/test_triple_diff.py +++ b/tests/test_triple_diff.py @@ -1395,44 +1395,24 @@ def test_reject_unknown_vcov_type(self): with pytest.raises(ValueError, match="invalid"): TripleDifference(vcov_type="hc4") - # -- Surface 5: fit()-time revalidation (set_params can't bypass) ---- - - def test_set_params_bad_vcov_caught_at_fit_time(self): - """set_params is strict-mirror sklearn (no atomic validation), but - fit() re-validates so a bad set_params(vcov_type='hc4') surfaces a - clear error at fit-time rather than silently propagating a bad - value to Results metadata. Mirrors CS - tests/test_staggered.py::test_set_params_bad_vcov_caught_at_fit_time.""" + # -- Surface 5: eager transactional validation (BaseEstimator) ---- + + def test_set_params_bad_vcov_raises_eagerly(self): + """set_params validates via constructor probe (transactional per the + locked v4 rule), so a bad vcov_type raises AT set_params with the + same message __init__ gives - and the estimator is unchanged.""" td = TripleDifference() - # set_params succeeds (sklearn-style mutate-then-validate-at-use) - td.set_params(vcov_type="hc4") - assert td.vcov_type == "hc4" - # fit() re-validates and raises - data = generate_ddd_data(n_per_cell=40, true_att=2.0, seed=37) with pytest.raises(ValueError, match="hc4"): - td.fit( - data, - outcome="outcome", - group="group", - partition="partition", - time="time", - ) + td.set_params(vcov_type="hc4") + assert td.vcov_type == "hc1" - def test_set_params_bad_vcov_classical_caught_at_fit_time(self): - """Same as above but with an IF-incompatible family (classical). - Catches the silent-propagation path on the methodology-rooted - rejection branch.""" + def test_set_params_bad_vcov_classical_raises_eagerly(self): + """Same for the IF-incompatible family (classical): the eager probe + surfaces the methodology-rooted rejection branch at set_params.""" td = TripleDifference() - td.set_params(vcov_type="classical") - data = generate_ddd_data(n_per_cell=40, true_att=2.0, seed=39) with pytest.raises(ValueError, match="influence-function"): - td.fit( - data, - outcome="outcome", - group="group", - partition="partition", - time="time", - ) + td.set_params(vcov_type="classical") + assert td.vcov_type == "hc1" # -- Introspection contract ------------------------------------------- diff --git a/tests/test_two_stage.py b/tests/test_two_stage.py index 5d51b155..8b0333ae 100644 --- a/tests/test_two_stage.py +++ b/tests/test_two_stage.py @@ -2350,13 +2350,14 @@ def test_invalid_vcov_type_rejected_at_init(self, bad): TwoStageDiD(vcov_type=bad) @pytest.mark.parametrize("bad", ["classical", "hc2", "hc2_bm", "conley"]) - def test_fit_revalidates_after_set_params(self, bad): - data = generate_test_data(n_units=60, seed=4) + def test_set_params_rejects_bad_vcov_eagerly(self, bad): + # set_params validates via constructor probe (transactional per the + # locked v4 rule): the bad value raises at set_params and the + # estimator is unchanged. est = TwoStageDiD() - est.set_params(vcov_type=bad) # sklearn mutate-then-validate-at-use - assert est.vcov_type == bad with pytest.raises(ValueError, match=bad): - est.fit(data, outcome="outcome", unit="unit", time="time", first_treat="first_treat") + est.set_params(vcov_type=bad) + assert est.vcov_type == "hc1" def test_rejection_messages_are_methodology_specific(self): with pytest.raises(ValueError, match="GMM"):