feat(imputation): leave-one-out conservative variance (BJS 2024 App. A.9)

CHANGELOG.md

@@ -38,6 +38,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
     change; the machine-precision hetero/cluster lock is deferred — needs an unbalanced-DGP golden).
 
 ### Added
+- **`ImputationDiD` leave-one-out conservative variance** (`leave_one_out`, default `False`) — the
+  Borusyak-Jaravel-Spiess (2024) Supplementary Appendix A.9 finite-sample refinement. The non-LOO
+  auxiliary aggregate `tau_tilde_g` (eq. 8) is built from the fitted `tau_hat_it` and so partially
+  overfits to the noise `epsilon_it`, biasing the conservative variance downward; `leave_one_out=True`
+  recomputes each unit's group aggregate excluding that unit (implemented efficiently by rescaling each
+  treated auxiliary residual by `1/(1 - v_ig**2/sum_j v_jg**2)`, exactly equivalent to the direct
+  leave-one-out at the per-unit cluster sum), yielding a larger, less-downward-biased SE (Prop. A8:
+  unbiased for an upper bound at the default unit clustering). Point estimates are unchanged. A group
+  with a single positive-weight unit (LOO undefined, App. A.9 fn. 51) keeps the non-LOO residual with a
+  `UserWarning`. `leave_one_out` is recorded on `ImputationDiDResults` (and `to_dict()` / `summary()`);
+  replicate-weight survey designs raise `NotImplementedError` (their variance bypasses the
+  influence-function path). Default `False` preserves R `didimputation` parity.
 - **`ContinuousDiD` lowest-dose-as-control** (`control_group="lowest_dose"`, CGBS 2024 Remark 3.1) for
   settings with no untreated group (`P(D=0) = 0`): the lowest-dose group `d_L` becomes the comparison
   and the estimand is `ATT(d) − ATT(d_L)` (with `ATT(d_L) = 0` the omitted reference). It is a

METHODOLOGY_REVIEW.md

@@ -561,7 +561,7 @@ and covariate-adjusted specifications.)
 - Multiplier bootstrap on the Theorem-3 influence function (library extension, not in the paper).
 - Survey-design TSL variance on the influence function (library extension).
 - NaN inference for undefined statistics; Proposition-5 refuse-to-estimate (NaN + warning).
-- Leave-one-out variance refinement (Supplementary Appendix A.9) not implemented (finite-sample refinement; tracked as future work).
+- Leave-one-out variance refinement (Supplementary Appendix A.9) implemented as the opt-in `leave_one_out` parameter (default `False`, preserving R `didimputation` parity); efficient residual rescale `1/(1 - v_ig^2/sum_j v_jg^2)` exactly equivalent to direct leave-one-out (see REGISTRY ImputationDiD Note).
 
 **R Comparison Results** (`didimputation` v0.5.0, fixed-seed panel, `benchmarks/data/didimputation_golden.json`):
 

TODO.md

@@ -30,7 +30,6 @@ The `Origin` column (Actionable tables) and the `PR` column (Deferred tables) bo
 |-------|----------|--------|--------|----------|
 | `SyntheticControl` conformal (CWZ 2021) extensions: (a) one-sided / signed-`t` variants (§7); (b) covariates in the conformal proxy (`X_jt`, eqs 4/6 — current proxy is outcomes-only); (c) AR / innovation-permutation path (Lemmas 5-7) for time-series proxies. The joint test, pointwise CIs, and average-effect CI have landed. | `conformal.py`, `synthetic_control_results.py` | CWZ-2021 | Heavy | Low |
 | `ContinuousDiD` CGBS-2024 extensions. (a) `covariates=` kwarg — **DONE (reg/dr)**; (b) discrete-treatment saturated regression (`treatment_type="discrete"`) — **DONE**; (c) lowest-dose-as-control per Remark 3.1 when `P(D=0)=0` (`control_group="lowest_dose"`) — **DONE** (discrete + continuous mass-point, single-cohort; estimand `ATT(d)−ATT(d_L)`; see REGISTRY Note #7). Remaining (all deferred `NotImplementedError`, documented): `estimation_method="ipw"` on the dose curve (scalar-adjustment / degenerate); `covariates=` × `survey_design=` (weighted OR + weighted nuisance IF); multi-cohort **heterogeneous-support** discrete aggregation (support-aware: average each dose only over the cohorts that observe it); **multi-cohort `lowest_dose`** (within-cohort `d_L` reference + support-aware cross-cohort aggregation); and **`covariates=` × `lowest_dose`** (conditional-PT-relative-to-`d_L` estimand). Single-cohort / 2-period / shared-support multi-cohort are supported. | `continuous_did.py` | CGBS-2024 | Heavy | Low |
-| `ImputationDiD` LOO conservative-variance refinement (BJS 2024 Supp. Appendix A.9) — a finite-sample improvement to the auxiliary-model residuals reducing overfit of `tau_tilde_g` to `epsilon`. Asymptotic Theorem-3 variance is implemented and matches R `didimputation` (which also omits LOO by default). | `imputation.py` | imputation-validation | Mid | Low |
 | `TwoWayFixedEffects(vcov_type in {hc2, hc2_bm})` with replicate-weight designs raises `NotImplementedError` (`twfe.py:~233`). The replicate path re-demeans per replicate, which doesn't compose with the full-dummy HC2/HC2-BM build — a correct impl needs per-replicate full-dummy refit. Workaround: `hc1` for replicate-weight CR1. | `twfe.py::fit` | follow-up | Heavy | Low |
 | `CallawaySantAnna` no-covariate ipw treats the propensity as unconditional (no correction term); R's `did` fits an intercept-only logit whose estimation effect is identically zero in the IF, so values match — decide whether to document-only (current REGISTRY note) or mirror R's code path for structural parity. | `staggered.py::_ipw_estimation` (no-cov branch) | CS-scaling | Quick | Low |
 | Fold the R `did` 2.5.1 ipw aggregation yardsticks (hardcoded with provenance in `test_golden_ipw_aggregation_se_vs_r_did_251`) into `csdid_golden_values.json` on the next fixture regeneration — the generator already emits the ipw `aggte` blocks; switch the test to read the JSON. | `benchmarks/R/generate_csdid_test_values.R`, `tests/test_csdid_ported.py` | CS-scaling | Quick | Low |

diff_diff/guides/llms-full.txt

@@ -392,6 +392,7 @@ ImputationDiD(
     horizon_max: int | None = None,              # Max event-study horizon
     aux_partition: str = "cohort_horizon",        # "cohort_horizon", "cohort", or "horizon"
     pretrends: bool = False,                     # Include pre-treatment horizons in event study
+    leave_one_out: bool = False,                 # BJS 2024 App. A.9 leave-one-out finite-sample variance (larger, less-downward-biased SE)
 )
 ```
 

diff_diff/imputation.py

@@ -127,6 +127,27 @@ class ImputationDiD(ImputationDiDBootstrapMixin):
         pre-trends assessment. Pre-period effects should be ~0 under
         parallel trends. Only affects event_study aggregation; overall
         ATT and group aggregation are unchanged.
+    leave_one_out : bool, default=False
+        If True, apply the Borusyak-Jaravel-Spiess (2024) Supplementary
+        Appendix A.9 leave-one-out finite-sample refinement to the
+        conservative variance. The non-LOO auxiliary aggregate ``tau_tilde_g``
+        is built from the fitted ``tau_hat_it`` and thus partially overfits to
+        the noise ``epsilon_it``, biasing the variance downward. LOO recomputes
+        each unit's group aggregate excluding that unit -- implemented
+        efficiently by rescaling each treated auxiliary residual by
+        ``1 / (1 - v_ig**2 / sum_j v_jg**2)`` (App. A.9), which is exactly
+        equivalent to the direct leave-one-out at the per-unit cluster sum.
+        Yields a larger, less-downward-biased SE (Prop. A8: unbiased for an
+        upper bound). Default False preserves R ``didimputation`` parity; the
+        refinement is an option in the authors' Stata ``did_imputation``. LOO
+        is undefined for a group with a single positive-weight unit (App. A.9
+        footnote 51): such groups fall back to the non-LOO residual with a
+        UserWarning. The Prop. A8 direction (LOO >= non-LOO) is guaranteed at
+        the default unit clustering; coarser ``cluster=`` / analytical
+        ``survey_design=`` / ``n_bootstrap`` compositions apply the same rescale
+        but are a library extension beyond the paper's derivation.
+        Replicate-weight survey designs raise ``NotImplementedError`` (their
+        variance bypasses the influence-function path where the rescale lives).
 
     Attributes
     ----------
@@ -182,6 +203,7 @@ def __init__(
         horizon_max: Optional[int] = None,
         aux_partition: str = "cohort_horizon",
         pretrends: bool = False,
+        leave_one_out: bool = False,
     ):
         if rank_deficient_action not in ("warn", "error", "silent"):
             raise ValueError(
@@ -199,6 +221,7 @@ def __init__(
                 f"got '{aux_partition}'"
             )
         self._validate_vcov_type(vcov_type)
+        self._validate_leave_one_out(leave_one_out)
 
         self.anticipation = anticipation
         self.alpha = alpha
@@ -211,6 +234,7 @@ def __init__(
         self.horizon_max = horizon_max
         self.aux_partition = aux_partition
         self.pretrends = pretrends
+        self.leave_one_out = leave_one_out
 
         self.is_fitted_ = False
         self.results_: Optional[ImputationDiDResults] = None
@@ -275,6 +299,7 @@ def fit(
         # mutations (e.g. set_params(vcov_type="classical")) are re-checked
         # at use rather than silently accepted by the parameter setter.
         self._validate_vcov_type(self.vcov_type)
+        self._validate_leave_one_out(self.leave_one_out)
 
         # Validate inputs
         required_cols = [outcome, unit, time, first_treat]
@@ -345,6 +370,22 @@ def fit(
                 "the design structure implicitly) or use a non-replicate "
                 "survey design (with explicit strata/psu/fpc)."
             )
+        # Reject replicate-weight + leave_one_out=: the BJS 2024 App. A.9
+        # refinement rescales the conservative influence-function auxiliary
+        # residuals, but replicate-weight variance is computed by per-replicate
+        # point-estimate refits (not the IF path), so leave_one_out would
+        # silently have no effect. Fail-closed (no-silent-failures).
+        if _uses_replicate_imp and self.leave_one_out:
+            raise NotImplementedError(
+                "ImputationDiD(leave_one_out=True) is not supported with "
+                "replicate-weight survey designs. The leave-one-out refinement "
+                "(Borusyak, Jaravel & Spiess 2024, Supp. App. A.9) rescales the "
+                "conservative influence-function residuals, but replicate-weight "
+                "variance is computed by per-replicate refits and does not use "
+                "that path — leave_one_out would silently have no effect. Use a "
+                "non-replicate (Taylor-linearization) survey design, or "
+                "leave_one_out=False."
+            )
         # Validate within-unit constancy for panel survey designs
         if resolved_survey is not None:
             _validate_unit_constant_survey(data, unit, survey_design)
@@ -968,6 +1009,7 @@ def _refit_imp(w_r):
             vcov_type=self.vcov_type,
             cluster_name=_cluster_name_for_results,
             n_clusters=_n_clusters_for_results,
+            leave_one_out=self.leave_one_out,
         )
 
         self.is_fitted_ = True
@@ -1683,6 +1725,8 @@ def _compute_auxiliary_residuals_treated(
         # _compute_target_weights) poison its whole group via 0 * NaN = NaN. The
         # previous observation-level pandas sum relied on skipna to drop them.
         contrib = (v_treated != 0.0) & np.isfinite(tau_hat)
+        loo_factor: Optional[pd.Series] = None
+        n_single_loo = 0
         if contrib.any():
             inner = pd.DataFrame(
                 {
@@ -1702,6 +1746,12 @@ def _compute_auxiliary_residuals_treated(
             )
             den_ok = per_group["den"].abs() >= 1e-15
             tau_tilde_map = (per_group["num"] / per_group["den"]).where(den_ok)
+            # BJS 2024 App. A.9 leave-one-out refinement: rescale each treated
+            # residual by 1/(1 - v_ig^2 / sum_j v_jg^2) (== the direct-LOO tau_tilde
+            # exactly, at the per-unit cluster sum). Reuses a_{i,g} = per_unit['a']
+            # and sum_j v_jg^2 = per_group['den']; applied to epsilon_treated below.
+            if self.leave_one_out:
+                loo_factor, n_single_loo = self._leave_one_out_factor(per_unit, per_group)
         else:
             tau_tilde_map = pd.Series(dtype=float)
 
@@ -1720,6 +1770,28 @@ def _compute_auxiliary_residuals_treated(
         # Auxiliary residuals
         epsilon_treated = tau_hat - tau_tilde
 
+        # Leave-one-out rescale (BJS 2024 App. A.9): map each treated obs to its
+        # (group, unit) factor and inflate the residual. Non-contributing rows
+        # (v_it == 0, psi == 0 anyway) and single-positive-weight-unit groups
+        # (LOO undefined, fn. 51) keep factor 1.0.
+        if self.leave_one_out and loo_factor is not None:
+            obs_index = pd.MultiIndex.from_arrays(
+                [group_codes, df_1[unit].values], names=["g", "u"]
+            )
+            factor_per_obs = loo_factor.reindex(obs_index).to_numpy(dtype=float)
+            factor_per_obs = np.where(np.isfinite(factor_per_obs), factor_per_obs, 1.0)
+            epsilon_treated = epsilon_treated * factor_per_obs
+            if n_single_loo > 0:
+                warnings.warn(
+                    f"leave_one_out=True: {n_single_loo} auxiliary group(s) have a single "
+                    f"positive-weight unit, where the leave-one-out variance is undefined "
+                    f"(Borusyak, Jaravel & Spiess 2024, Supp. App. A.9 fn. 51); those groups "
+                    f"keep the non-leave-out residual. A coarser aux_partition reduces "
+                    f"singleton groups.",
+                    UserWarning,
+                    stacklevel=2,
+                )
+
         return epsilon_treated
 
     def _compute_residuals_untreated(
@@ -2496,6 +2568,7 @@ def get_params(self) -> Dict[str, Any]:
             "horizon_max": self.horizon_max,
             "aux_partition": self.aux_partition,
             "pretrends": self.pretrends,
+            "leave_one_out": self.leave_one_out,
         }
 
     def set_params(self, **params) -> "ImputationDiD":
@@ -2507,6 +2580,75 @@ def set_params(self, **params) -> "ImputationDiD":
                 raise ValueError(f"Unknown parameter: {key}")
         return self
 
+    @staticmethod
+    def _validate_leave_one_out(leave_one_out: Any) -> None:
+        """Validate ``leave_one_out`` is a strict bool.
+
+        Called from ``__init__`` AND ``fit()`` so sklearn-style
+        ``set_params(leave_one_out=...)`` mutations are re-checked at use
+        time -- the naive ``set_params`` setter would otherwise accept a
+        truthy string (e.g. "yes") and silently run the LOO refinement.
+        """
+        if not isinstance(leave_one_out, bool):
+            raise TypeError(f"leave_one_out must be a bool, got {type(leave_one_out).__name__}")
+
+    @staticmethod
+    def _leave_one_out_factor(
+        per_unit: pd.DataFrame, per_group: pd.DataFrame
+    ) -> Tuple[pd.Series, int]:
+        """Per-(group, unit) leave-one-out residual-rescale factor (BJS 2024 A.9).
+
+        ``factor_{g,i} = 1 / (1 - v_ig**2 / sum_j v_jg**2)`` with
+        ``v_ig = per_unit['a']`` and ``sum_j v_jg**2 = per_group['den']``. This
+        rescale of ``epsilon_tilde_it`` reproduces the direct leave-one-out
+        aggregate ``tau_tilde_it^LO`` exactly at the per-unit cluster sum
+        ``psi_i = sum_t v_it * epsilon_tilde_it`` (App. A.9). A group with a
+        single positive-weight unit has ``v_ig**2 == sum_j v_jg**2`` so the
+        factor diverges (LOO undefined, App. A.9 fn. 51); those groups fall back
+        to ``1.0`` (non-LOO). A genuinely unit-dominated but >=2-unit group keeps
+        its large finite factor -- that is the paper's intended inflation.
+
+        Returns
+        -------
+        (factor : pd.Series indexed like ``per_unit`` (g, u), n_single_unit_groups : int)
+        """
+        a = per_unit["a"].to_numpy(dtype=float)
+        sq = a**2
+        g_level = per_unit.index.get_level_values("g")
+        u_level = per_unit.index.get_level_values("u")
+        den = per_group["den"].reindex(g_level).to_numpy(dtype=float)  # D_g per (g,u)
+
+        # A group is "singleton" for LOO (App. A.9 fn. 51) when fewer than two
+        # units carry positive squared weight -- covers a true 1-unit group AND
+        # the effective-singleton case (>=2 rows, only one with a_ig != 0).
+        pos_per_group = pd.Series(sq > 0.0, index=per_unit.index).groupby(level="g").sum()
+        single_groups = pos_per_group.index[pos_per_group < 2]
+        is_single = pd.Index(g_level).isin(single_groups)
+
+        den_ok = np.abs(den) >= 1e-15
+        # factor = D_g / (D_g - v_ig^2) = D_g / sum_{j!=i} v_jg^2. Compute the
+        # leave-one-out denominator as the sum of the OTHER units' squared
+        # weights -- NOT as D_g - v_ig^2 after forming the ratio: for a genuinely
+        # dominated (but >=2-unit) group the subtraction loses precision (and can
+        # cancel to 0/negative) in float64 -- a finite-but-wrong or silently
+        # non-LOO factor. The fast subtraction is accurate away from the
+        # near-cancellation boundary; wherever the leave-one-out mass is a tiny
+        # fraction of D (relative loss of >~1e-6), recompute it exactly as the
+        # drop-then-sum of the OTHER units' squared weights. At most one unit per
+        # group can be that dominant, so the recompute stays O(units).
+        other_mass = den - sq
+        suspect = (~is_single) & den_ok & (other_mass <= 1e-6 * den)
+        if suspect.any():
+            sq_series = pd.Series(sq, index=per_unit.index)
+            for pos in np.nonzero(suspect)[0]:
+                grp = sq_series.xs(g_level[pos], level="g")
+                other_mass[pos] = float(grp.drop(u_level[pos]).sum())
+        # Fall back to non-LOO (factor 1.0) only where LOO is genuinely undefined:
+        # a singleton group (fn. 51), a degenerate den, or no other positive mass.
+        fallback = is_single | ~den_ok | (other_mass <= 0.0)
+        factor = np.where(fallback, 1.0, den / np.where(fallback, 1.0, other_mass))
+        return pd.Series(factor, index=per_unit.index), int(len(single_groups))
+
     @staticmethod
     def _validate_vcov_type(vcov_type: str) -> None:
         """Validate ``vcov_type`` membership against ImputationDiD's

diff_diff/imputation_results.py

@@ -150,6 +150,9 @@ class ImputationDiDResults:
     vcov_type: str = field(default="hc1")
     cluster_name: Optional[str] = field(default=None)
     n_clusters: Optional[int] = field(default=None)
+    # BJS 2024 Supp. App. A.9 leave-one-out finite-sample variance refinement
+    # (opt-in). Recorded here so reported SEs are self-describing.
+    leave_one_out: bool = field(default=False)
 
     # --- Inference-field aliases (balance/external-adapter compatibility) ---
     @property
@@ -260,6 +263,8 @@ def summary(self, alpha: Optional[float] = None) -> str:
                 lines.append(f"{'Variance estimator:':<30} {vcov_label:>15}")
         if self.n_clusters is not None and self.bootstrap_results is None:
             lines.append(f"{'Number of clusters:':<30} {self.n_clusters:>15}")
+        if self.leave_one_out:
+            lines.append(f"{'Leave-one-out variance:':<30} {'A.9 (BJS 2024)':>15}")
 
         lines.append("")
 
@@ -504,6 +509,7 @@ def to_dict(self) -> Dict[str, Any]:
             "alpha": self.alpha,
             "anticipation": self.anticipation,
             "vcov_type": self.vcov_type,
+            "leave_one_out": self.leave_one_out,
         }
         if self.cluster_name is not None:
             result["cluster_name"] = self.cluster_name