fix(bayes): v0.9.16 self-review — cohort dropna alignment + poly_u sentinel

Two round-2 fixes caught in self-review after v0.9.16 landed:

1. `bayes_did(cohort=...)` used two independent `.dropna()` passes
   (one in `_prepare_did_frame`, one in the cohort factoriser) which
   silently misaligned `cohort_codes` with `DID` whenever the cohort
   column had NaN rows that y/treat/post/unit/time did not. Synthetic
   test data never triggered this; real panels with missing cohort
   info would have either shape-errored in PyMC or, worse, fit with
   shuffled cohort labels. Fix threads `cohort` through
   `_prepare_did_frame` so a single `.dropna()` governs every array.

2. `poly_u: int = 2` default forced bivariate_normal callers to eat
   a spurious "overriding poly_u=2 -> poly_u=1" UserWarning on every
   default-arg call. Switch to `poly_u: Optional[int] = None` and
   resolve method-specifically: 2 for polynomial/hv_latent, 1 for
   bivariate_normal. Explicit non-1 values with bivariate_normal
   still warn before overriding.

New regression test `test_cohort_nan_rows_dropped_consistently`
poisons 15 rows' cohort values with NaN and asserts the ordering of
per-cohort τ posteriors is still recovered.

All 118 bayes_mte/did/iv tests green; full suite 2411/2411 green.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: brycewang-stanford

src/statspai/bayes/did.py

@@ -29,8 +29,19 @@ def _prepare_did_frame(
     unit: Optional[str],
     time: Optional[str],
     covariates: Optional[List[str]],
+    cohort: Optional[str] = None,
 ) -> dict:
-    """Validate, drop NA, and extract arrays for DID."""
+    """Validate, drop NA, and extract arrays for DID.
+
+    ``cohort`` is threaded through here (rather than being handled
+    separately in ``bayes_did``) so its NaN rows get dropped in the
+    same pass as ``y / treat / post / unit / time / covariates``.
+    Using two independent ``dropna`` calls — which is what v0.9.16's
+    first draft did — caused a silent row-alignment bug whenever
+    ``cohort`` had NaN rows that the other columns did not: the
+    cohort-codes array and ``DID`` array ended up with different
+    lengths (or same length but different rows).
+    """
     required = [y, treat, post]
     for c in required:
         if c not in data.columns:
@@ -50,6 +61,10 @@ def _prepare_did_frame(
             if c not in data.columns:
                 raise ValueError(f"Covariate '{c}' not found in data")
         cols.extend(covariates)
+    if cohort is not None:
+        if cohort not in data.columns:
+            raise ValueError(f"Column '{cohort}' (cohort) not found in data")
+        cols.append(cohort)
 
     clean = data[cols].dropna().reset_index(drop=True)
     n = len(clean)
@@ -94,6 +109,18 @@ def _prepare_did_frame(
                 "time random effect."
             )
 
+    cohort_codes: Optional[np.ndarray] = None
+    cohort_labels: list = []
+    if cohort is not None:
+        codes, uniques = pd.factorize(clean[cohort], sort=True)
+        cohort_codes = np.asarray(codes, dtype=np.int64)
+        cohort_labels = list(uniques)
+        if len(cohort_labels) < 2:
+            raise ValueError(
+                f"cohort column '{cohort}' has < 2 distinct values; "
+                "per-cohort ATT requires at least 2 cohorts."
+            )
+
     X: Optional[np.ndarray] = None
     if covariates:
         X = clean[covariates].to_numpy(dtype=float)
@@ -108,6 +135,8 @@ def _prepare_did_frame(
         'n_units': n_units,
         'time_idx': time_idx,
         'n_times': n_times,
+        'cohort_idx': cohort_codes,
+        'cohort_labels': cohort_labels,
         'X': X,
         'covariates': list(covariates) if covariates else [],
     }
@@ -207,7 +236,15 @@ def bayes_did(
     """
     pm, _ = _require_pymc()
 
-    prep = _prepare_did_frame(data, y, treat, post, unit, time, covariates)
+    # Thread `cohort` into `_prepare_did_frame` so every NaN gets
+    # dropped in a single pass and `cohort_codes` is guaranteed to
+    # align row-for-row with `DID`. A prior implementation did two
+    # independent dropna() calls and silently mis-aligned the two
+    # arrays whenever `cohort` had NaN rows the other columns did
+    # not (not caught by synthetic-data tests that have no NaN).
+    prep = _prepare_did_frame(
+        data, y, treat, post, unit, time, covariates, cohort=cohort,
+    )
     n = prep['n']
     Y = prep['Y']
     T = prep['T']
@@ -218,34 +255,8 @@ def bayes_did(
     use_unit_re = prep['unit_idx'] is not None
     use_time_re = prep['time_idx'] is not None
 
-    # Cohort wiring: factorise once so the model gets integer codes and
-    # the result carries back the user's original labels for tidy(). We
-    # factorise on the dropna'd frame to keep index alignment with DID.
-    cohort_codes: Optional[np.ndarray] = None
-    cohort_labels: list = []
-    if cohort is not None:
-        if cohort not in data.columns:
-            raise ValueError(f"Column '{cohort}' (cohort) not found in data")
-        cols_for_factor = [y, treat, post]
-        if unit is not None:
-            cols_for_factor.append(unit)
-        if time is not None:
-            cols_for_factor.append(time)
-        if covariates:
-            cols_for_factor.extend(covariates)
-        cols_for_factor.append(cohort)
-        clean_for_cohort = (
-            data[cols_for_factor].dropna().reset_index(drop=True)
-        )
-        codes, uniques = pd.factorize(clean_for_cohort[cohort], sort=True)
-        cohort_codes = np.asarray(codes, dtype=np.int64)
-        cohort_labels = list(uniques)
-        n_cohorts = len(cohort_labels)
-        if n_cohorts < 2:
-            raise ValueError(
-                f"cohort column '{cohort}' has < 2 distinct values; "
-                "per-cohort ATT requires at least 2 cohorts."
-            )
+    cohort_codes = prep['cohort_idx']
+    cohort_labels = prep['cohort_labels']
 
     mu_ate, sigma_ate = prior_ate
 

src/statspai/bayes/mte.py

@@ -115,7 +115,7 @@ def bayes_mte(
     mte_method: str = 'polynomial',
     selection: str = 'uniform',
     u_grid: Optional[np.ndarray] = None,
-    poly_u: int = 2,
+    poly_u: Optional[int] = None,
     prior_coef_sigma: float = 10.0,
     prior_mte_sigma: float = 5.0,
     prior_noise: float = 5.0,
@@ -188,10 +188,15 @@ def bayes_mte(
     u_grid : np.ndarray, optional
         Grid of propensity-to-be-treated values on which to evaluate
         the MTE posterior. Default: ``np.linspace(0.05, 0.95, 19)``.
-    poly_u : int, default 2
+    poly_u : int, optional
         Polynomial order of the MTE in ``U_D``. ``poly_u=0`` reduces
         to a constant treatment effect; ``poly_u=2`` captures
-        U-shaped or inverted-U selection on gains.
+        U-shaped or inverted-U selection on gains. Default resolves
+        per ``mte_method``: 2 for ``'polynomial'`` / ``'hv_latent'``,
+        and 1 for ``'bivariate_normal'`` (the model is inherently
+        linear in ``V`` there and ignores any other value). Passing
+        an explicit non-1 value with ``mte_method='bivariate_normal'``
+        emits a UserWarning before the override.
     prior_mte_sigma : float, default 5.0
         SD on each MTE polynomial coefficient (Normal prior).
     prior_coef_sigma, prior_noise : float
@@ -258,6 +263,13 @@ def bayes_mte(
             f"got {selection!r}"
         )
 
+    # Resolve `poly_u=None` to the method-appropriate default so that
+    # the signature default no longer requires a fixed integer and
+    # 'bivariate_normal' stops emitting a spurious "poly_u=2 -> 1"
+    # warning on every default-arg call.
+    if poly_u is None:
+        poly_u = 1 if mte_method == 'bivariate_normal' else 2
+
     # Bivariate-normal HV is inherently a V-scale joint-first-stage
     # model; raise loudly instead of silently mis-specifying. The
     # model also fixes poly_u=1 because MTE(v) is closed-form linear

tests/test_bayes_did_cohort.py

@@ -178,3 +178,43 @@ def test_cohort_weight_recorded_in_model_info(staggered_did_data):
     # treated cohorts; control cohort gets weight 0).
     assert all(wi >= 0 for wi in w)
     assert abs(sum(w) - 1.0) < 1e-9
+
+
+def test_cohort_nan_rows_dropped_consistently(staggered_did_data):
+    """Regression test for a row-alignment bug caught in self-review:
+    if ``cohort`` has NaN rows that ``y/treat/post`` do not, the
+    cohort-codes array and DID array must still be length-aligned.
+    Both are produced by a single ``dropna()`` pass inside
+    ``_prepare_did_frame``, so any extra NaN in the cohort column
+    removes the row for ALL arrays uniformly."""
+    df = staggered_did_data.copy()
+    # Poison 15 rows' cohort values with NaN while leaving y/treat/post
+    # intact. If dropna were handled separately for cohort vs DID, this
+    # would desynchronise the arrays and either (a) error with a
+    # length-mismatch from PyMC or (b) silently fit with shuffled cohorts.
+    rng = np.random.default_rng(77)
+    poison_idx = rng.choice(df.index, size=15, replace=False)
+    df.loc[poison_idx, 'cohort'] = np.nan
+
+    r = bayes_did(df, y='y', treat='treat', post='post',
+                  cohort='cohort',
+                  draws=150, tune=150, chains=2, progressbar=False,
+                  random_state=78)
+    # Sanity: the model fits without shape errors (the bug manifested
+    # as a PyMC shape-mismatch exception).
+    assert len(r.cohort_labels) == 3
+    # Setting a cell to NaN promotes the whole column from int64 to
+    # float64 in pandas, so cohort labels become floats (-1.0, 2019.0,
+    # 2020.0). The tidy infrastructure stringifies them consistently —
+    # look up posteriors via whatever stringification the result uses.
+    summaries = r.cohort_summaries
+    label_2019 = next(
+        k for k in summaries if str(k) in ('2019', '2019.0')
+    )
+    label_2020 = next(
+        k for k in summaries if str(k) in ('2020', '2020.0')
+    )
+    m_2019 = summaries[label_2019]['posterior_mean']
+    m_2020 = summaries[label_2020]['posterior_mean']
+    # Ordering (2.0 > 0.5) is preserved after dropping 15 rows.
+    assert m_2019 > m_2020