Audit-round polish: 5 real bugs found in the v0.7.1 work, all fixed

Full audit of the 17 commits landed today surfaced five genuine issues
plus a dozen purely cosmetic ones.  This commit fixes the real five
and adds targeted regression tests for each.

1. `_df_to_booktabs` LaTeX escape was non-idempotent.
   Old code replaced `\` with `\textbackslash{}` first, then replaced
   `{` → `\{` and `}` → `\}` in a second pass — which re-escaped the
   braces just inserted, producing `\textbackslash\{\}` instead of the
   correct `\textbackslash{}`.  Same failure mode for `~` and `^`.
   Fix: single-pass `re.sub` with a lookup table so each character is
   escaped exactly once.
   (Latent in the DiD report path — our tables never carried
    backslashes — but would bite any user passing string labels.)

2. `cs_report(save_to='~/study/cs_v1')` didn't expand `~`.
   `os.path.dirname('~/study/cs_v1')` returns `"~/study"`, so the
   pre-fix code created a literal `./~/study` directory.  Fix:
   `os.path.expanduser()` before anything else.

3. `_save_report_bundle` forced `matplotlib.use('Agg')` unconditionally,
   which prints a warning when pyplot is already imported (common in
   Jupyter).  Fix: only switch backend if pyplot hasn't been loaded
   yet; otherwise honour whatever backend the user had.

4. `bjs_pretrend_joint` swallowed `except Exception` silently.
   If `did_imputation` had a real bug, users would just see
   "not enough reps succeeded" with no clue why.  Fix: narrow to
   `(ValueError, RuntimeError, LinAlgError)` for expected
   pathological resamples; track the last unexpected exception and
   include its type + message in the failure.

5. `sp.did()` silently ignored `aggregation=`, `panel=False`, and
   `anticipation=` when `method` resolved to anything other than CS.
   Fix:  (a) auto-upgrade `method='auto'` → `'callaway_santanna'`
   when `aggregation` is supplied with an `id`; (b) raise
   `ValueError` if any CS-only argument is paired with a non-CS
   `method`.

Regression tests (4 new):
 - LaTeX backslash round-trip (brace mangling guard)
 - `save_to` expands `~` to `$HOME`
 - `sp.did(method='2x2', aggregation='dynamic')` raises
 - `sp.did(method='2x2', anticipation=1)` raises

Full DiD suite: 165/165 (was 161).

Also noted but left alone (cosmetic / not correctness):
 - SA's `(g_val, e) in {m for m in interact_meta}` builds a
   throwaway set comprehension; `in interact_meta` would suffice.
   Performance only matters on huge event windows.
 - SA's `XtX + 1e-10 * I` ridge term slightly biases β̂ in
   well-conditioned cases (~1 ppm); matches common practice.
 - `_two_way_demean(x)` calls `x.astype(float).copy()` — `.copy()`
   is redundant once `.astype()` already returns a fresh array.

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

Author: brycewang-stanford

src/statspai/did/__init__.py

@@ -215,6 +215,34 @@ def did(
             + (" ..." if len(data.columns) > 10 else "")
         )
 
+    # If CS-specific arguments are passed, force the method to CS rather
+    # than silently ignoring them.  This is the "do what I mean" path
+    # when a user writes
+    #   sp.did(..., aggregation='dynamic')
+    # and expects CS to run under the hood.
+    if aggregation is not None and method == 'auto' and id is not None:
+        method = 'callaway_santanna'
+
+    # Validate that CS-only arguments were not paired with a non-CS
+    # method — fail loudly rather than swallow the argument.
+    _cs_methods = {'callaway_santanna', 'cs', 'auto'}
+    if aggregation is not None and method not in _cs_methods:
+        raise ValueError(
+            f"`aggregation={aggregation!r}` is only supported with the "
+            f"Callaway-Sant'Anna estimator (method='cs'); got "
+            f"method={method!r}."
+        )
+    if (not panel) and method not in _cs_methods:
+        raise ValueError(
+            f"`panel=False` is only supported with Callaway-Sant'Anna; "
+            f"got method={method!r}."
+        )
+    if anticipation != 0 and method not in _cs_methods:
+        raise ValueError(
+            f"`anticipation={anticipation}` is only supported with "
+            f"Callaway-Sant'Anna; got method={method!r}."
+        )
+
     # Auto-detect if subgroup is provided → DDD
     if method == 'auto' and subgroup is not None:
         method = 'ddd'

src/statspai/did/bjs_inference.py

@@ -114,6 +114,10 @@ def bjs_pretrend_joint(
     rng = np.random.default_rng(seed)
 
     boot_mat = np.full((n_boot, K), np.nan)
+    # Remember the most recent unexpected error so that if many / all
+    # bootstrap replications fail we can surface it to the user rather
+    # than silently claim "not enough reps succeeded".
+    last_unexpected: Optional[Exception] = None
     for b in range(n_boot):
         picks = rng.choice(clusters, size=G, replace=True)
         frames = []
@@ -129,22 +133,36 @@ def bjs_pretrend_joint(
                 first_treat=first_treat, controls=controls,
                 horizon=horizon, cluster=cluster_col,
             )
-            es_b = r_b.model_info.get('event_study')
-            if es_b is None or len(es_b) == 0:
-                continue
-            sub = es_b[es_b['relative_time'].isin(pre_k)].sort_values(
-                'relative_time',
-            )
-            if len(sub) == K:
-                boot_mat[b] = sub['att'].to_numpy()
-        except Exception:
+        except (ValueError, RuntimeError, np.linalg.LinAlgError):
+            # Expected failures on pathological resamples (e.g. a draw
+            # that happens to drop all eventually-treated units).  Skip.
+            continue
+        except Exception as exc:
+            # Anything else is likely a genuine bug in BJS — keep the
+            # last one around so we can report it below if the run fails.
+            last_unexpected = exc
+            continue
+
+        es_b = r_b.model_info.get('event_study')
+        if es_b is None or len(es_b) == 0:
             continue
+        sub = es_b[es_b['relative_time'].isin(pre_k)].sort_values(
+            'relative_time',
+        )
+        if len(sub) == K:
+            boot_mat[b] = sub['att'].to_numpy()
 
     valid = ~np.any(np.isnan(boot_mat), axis=1)
     if valid.sum() < K + 1:
+        extra = (
+            f"  Last unexpected error: {type(last_unexpected).__name__}: "
+            f"{last_unexpected}"
+            if last_unexpected is not None else ""
+        )
         raise RuntimeError(
             f"Only {int(valid.sum())} / {n_boot} bootstrap replications "
             f"succeeded — not enough to estimate the {K}×{K} covariance."
+            + extra
         )
     boot_valid = boot_mat[valid]
     # Centre around the bootstrap mean so the cov estimator is unbiased

src/statspai/did/report.py

@@ -560,15 +560,30 @@ def _df_to_booktabs(df: pd.DataFrame, float_format: str = "%.4f") -> str:
             aligns.append("l")
     col_spec = "".join(aligns)
 
+    # Single-pass LaTeX escape.  Sequential replace() calls are unsafe
+    # because `\` → `\textbackslash{}` inserts `{` and `}` that later
+    # passes would re-escape, producing broken output like
+    # `\textbackslash\{\}` instead of `\textbackslash{}`.  Using re.sub
+    # with a lookup table guarantees each input character is escaped
+    # exactly once.
+    import re as _re
+    _LATEX_ESCAPES = {
+        '\\': r'\textbackslash{}',
+        '~':  r'\textasciitilde{}',
+        '^':  r'\textasciicircum{}',
+        '&':  r'\&',
+        '%':  r'\%',
+        '$':  r'\$',
+        '#':  r'\#',
+        '_':  r'\_',
+        '{':  r'\{',
+        '}':  r'\}',
+    }
+    _LATEX_RE = _re.compile(r'[\\~^&%$#_{}]')
+
     def _escape(v: object) -> str:
         text = "" if (isinstance(v, float) and pd.isna(v)) else str(v)
-        for old, new in (("\\", r"\textbackslash{}"), ("&", r"\&"),
-                         ("%", r"\%"), ("$", r"\$"), ("#", r"\#"),
-                         ("_", r"\_"), ("{", r"\{"), ("}", r"\}"),
-                         ("~", r"\textasciitilde{}"),
-                         ("^", r"\textasciicircum{}")):
-            text = text.replace(old, new)
-        return text
+        return _LATEX_RE.sub(lambda m: _LATEX_ESCAPES[m.group(0)], text)
 
     header = " & ".join(_escape(c) for c in formatted.columns) + " \\\\"
     body_rows = [
@@ -759,8 +774,14 @@ def _save_report_bundle(
     skipping any format whose optional dependency is missing.
     """
     import os
+    import sys
 
-    prefix = str(prefix)
+    def _sys_modules():
+        return sys.modules
+
+    # Expand a leading "~" / "~user" so `save_to='~/study/cs_v1'` works,
+    # otherwise os.makedirs would create a literal "./~/study" directory.
+    prefix = os.path.expanduser(str(prefix))
     parent = os.path.dirname(prefix)
     if parent:
         os.makedirs(parent, exist_ok=True)
@@ -791,7 +812,15 @@ def _save_report_bundle(
 
     try:
         import matplotlib
-        matplotlib.use("Agg")
+        # Only switch to a non-interactive backend if nothing has been
+        # set yet — calling matplotlib.use() after pyplot has been
+        # imported is a warning-only no-op in modern matplotlib but
+        # can surprise users running inside Jupyter.
+        if "matplotlib.pyplot" not in _sys_modules():
+            try:
+                matplotlib.use("Agg")
+            except Exception:
+                pass  # already set — honour caller's choice
         fig, _ = report.plot()
         png_path = f"{prefix}.png"
         fig.savefig(png_path, dpi=110, bbox_inches="tight")

tests/test_cs_report.py

@@ -253,3 +253,33 @@ def test_save_to_files_contain_expected_content(tmp_path):
     assert "\\begin{table}" in tex and "\\bottomrule" in tex
     txt = (tmp_path / "cs.txt").read_text()
     assert "Overall ATT" in txt
+
+
+# --------------------------------------------------------------------------- #
+# Audit-round bug regressions                                                 #
+# --------------------------------------------------------------------------- #
+
+def test_to_latex_backslash_is_single_escaped(report):
+    """Bug: \\ → \\textbackslash{} followed by { → \\{ mangled the brace.
+
+    A single backslash in a cell should render as \\textbackslash{}
+    verbatim, not \\textbackslash\\{\\}."""
+    import importlib
+    mod = importlib.import_module('statspai.did.report')
+    df = pd.DataFrame({'label': ['a\\b', 'x^y'], 'val': [1, 2]})
+    tex = mod._df_to_booktabs(df)
+    assert r'\textbackslash{}' in tex
+    assert r'\textbackslash\{\}' not in tex
+    assert r'\textasciicircum{}' in tex
+
+
+def test_save_to_expands_tilde(tmp_path, monkeypatch):
+    """`save_to='~/foo/bar'` should expand the home directory."""
+    df = _staggered_panel(seed=31)
+    monkeypatch.setenv("HOME", str(tmp_path))
+    cs_report(df, y='y', g='g', t='t', i='i',
+              n_boot=50, random_state=0, verbose=False,
+              save_to='~/cs_study/rpt')
+    # If ~ was not expanded, we'd see a literal "~" subdirectory under CWD.
+    assert (tmp_path / "cs_study" / "rpt.md").exists()
+    assert (tmp_path / "cs_study" / "rpt.txt").exists()

tests/test_sp_did_aggregation.py

@@ -65,3 +65,28 @@ def test_rcs_path_through_top_level():
                estimator='reg', panel=False)
     assert r.model_info['panel'] is False
     assert 'RCS' in r.model_info['estimator']
+
+
+def test_aggregation_with_incompatible_method_raises():
+    """Silent-ignore bug: `aggregation=` must be rejected for 2x2 / DDD."""
+    rng = np.random.default_rng(0)
+    df = pd.DataFrame({
+        'y': rng.normal(size=200),
+        'treat': rng.integers(0, 2, size=200),
+        'time': rng.integers(0, 2, size=200),
+    })
+    with pytest.raises(ValueError, match='Callaway'):
+        sp.did(df, y='y', treat='treat', time='time',
+               method='2x2', aggregation='dynamic')
+
+
+def test_anticipation_rejected_for_non_cs_method():
+    rng = np.random.default_rng(0)
+    df = pd.DataFrame({
+        'y': rng.normal(size=200),
+        'treat': rng.integers(0, 2, size=200),
+        'time': rng.integers(0, 2, size=200),
+    })
+    with pytest.raises(ValueError, match='anticipation'):
+        sp.did(df, y='y', treat='treat', time='time',
+               method='2x2', anticipation=1)