sp.did(aggregation=...) — one-call CS estimation + aggte in the top-level API

Teach the top-level `sp.did()` dispatcher to pipe a Callaway-Sant'Anna
result straight through `aggte()` when the caller passes a new
`aggregation=` argument (one of 'simple', 'dynamic', 'group',
'calendar'), so users no longer have to hand-compose:

    cs = callaway_santanna(...)
    es = aggte(cs, type='dynamic', n_boot=..., random_state=...)

Instead:

    es = sp.did(df, y=..., treat=..., time=..., id=...,
                aggregation='dynamic',
                n_boot=500, random_state=0)

Also surface `panel=` and `anticipation=` as first-class arguments
on `sp.did()` so the top-level dispatcher can drive the RCS and
anticipation paths without `**kwargs` acrobatics.

Tests (tests/test_sp_did_aggregation.py, 7 cases):
 - no-aggregation path still returns the raw (g, t) grid
 - four aggregation values produce the right number of rows and
   attach the uniform-band columns on non-simple aggregations
 - invalid aggregation raises ValueError
 - panel=False RCS branch is reachable through sp.did()

Full DiD suite still green.

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

Author: brycewang-stanford

src/statspai/did/__init__.py

@@ -80,6 +80,12 @@ def did(
     treat_unit=None,
     treat_time=None,
     se_method: str = 'placebo',
+    # CS aggte-dispatch (v0.7.1+)
+    aggregation: Optional[str] = None,
+    n_boot: int = 1000,
+    random_state: Optional[int] = None,
+    panel: bool = True,
+    anticipation: int = 0,
     **kwargs,
 ) -> CausalResult:
     """
@@ -135,6 +141,22 @@ def did(
         For SDID: treatment time.
     se_method : str, default 'placebo'
         For SDID: 'placebo', 'bootstrap', or 'jackknife'.
+    aggregation : str, optional
+        When set and ``method`` is Callaway–Sant'Anna, the raw ATT(g,t)
+        result is passed through :func:`aggte` with ``type=aggregation``
+        (``'simple'``, ``'dynamic'``, ``'group'``, or ``'calendar'``),
+        delivering the aggregated ATT with Mammen multiplier-bootstrap
+        uniform confidence bands in a single call.
+    n_boot : int, default 1000
+        Bootstrap replications for the multiplier bootstrap when
+        ``aggregation`` is set.
+    random_state : int, optional
+        Seed for the multiplier bootstrap.
+    panel : bool, default True
+        Forwarded to :func:`callaway_santanna`; set ``panel=False`` for
+        repeated cross-sections.
+    anticipation : int, default 0
+        Forwarded to :func:`callaway_santanna`.
 
     Returns
     -------
@@ -253,12 +275,25 @@ def did(
             raise ValueError(
                 "'id' (unit identifier) is required for staggered DID."
             )
-        return callaway_santanna(
+        cs_result = callaway_santanna(
             data, y=y, g=treat, t=time, i=id,
             x=covariates, estimator=estimator,
             control_group=control_group,
             base_period=base_period, alpha=alpha,
+            anticipation=anticipation, panel=panel,
         )
+        if aggregation is not None:
+            if aggregation not in ('simple', 'dynamic', 'group', 'calendar'):
+                raise ValueError(
+                    f"aggregation must be one of "
+                    f"'simple'/'dynamic'/'group'/'calendar', "
+                    f"got {aggregation!r}"
+                )
+            return aggte(
+                cs_result, type=aggregation, alpha=alpha,
+                n_boot=n_boot, random_state=random_state,
+            )
+        return cs_result
 
     if method in ('sun_abraham', 'sa', 'sunab'):
         if id is None:

tests/test_sp_did_aggregation.py

@@ -0,0 +1,67 @@
+"""Tests for the `sp.did(..., aggregation='dynamic'|...)` dispatch path."""
+import numpy as np
+import pandas as pd
+import pytest
+
+import statspai as sp
+
+
+def _panel(seed=0):
+    rng = np.random.default_rng(seed)
+    rows = []
+    for u in range(80):
+        g = [3, 5, 7, 0][u // 20]
+        ui = rng.normal(scale=0.3)
+        for t in range(1, 9):
+            te = max(0, t - g + 1) * 0.5 if g > 0 else 0
+            rows.append({'i': u, 't': t, 'g': g,
+                         'y': ui + 0.2 * t + te + rng.normal()})
+    return pd.DataFrame(rows)
+
+
+def test_no_aggregation_returns_cs_result():
+    df = _panel()
+    r = sp.did(df, y='y', treat='g', time='t', id='i')
+    # Without aggregation the result should carry the full (g,t) grid.
+    assert len(r.detail) == 21
+    assert 'aggregation' not in r.model_info
+
+
+@pytest.mark.parametrize("agg,expected_rows", [
+    ('simple', 1),
+    ('dynamic', 11),  # relative_time -6..-1 (excl -1 base) + 0..5
+    ('group', 3),
+    ('calendar', 6),
+])
+def test_aggregation_dispatch(agg, expected_rows):
+    df = _panel()
+    r = sp.did(df, y='y', treat='g', time='t', id='i',
+               aggregation=agg, n_boot=100, random_state=0)
+    assert r.model_info['aggregation'] == agg
+    assert len(r.detail) == expected_rows
+    # Aggregated results carry uniform-band columns except 'simple'.
+    if agg != 'simple':
+        assert {'cband_lower', 'cband_upper'} <= set(r.detail.columns)
+
+
+def test_invalid_aggregation_raises():
+    df = _panel()
+    with pytest.raises(ValueError, match='aggregation'):
+        sp.did(df, y='y', treat='g', time='t', id='i', aggregation='bogus')
+
+
+def test_rcs_path_through_top_level():
+    """sp.did(..., panel=False) should route through the RCS branch."""
+    rng = np.random.default_rng(1)
+    rows = []
+    for obs in range(400):
+        g = int(rng.choice([3, 5, 0], p=[0.3, 0.3, 0.4]))
+        t = int(rng.integers(1, 7))
+        te = max(0, t - g + 1) * 0.5 if g > 0 else 0
+        y = 0.2 * t + te + rng.normal()
+        rows.append({'obs': obs, 't': t, 'g': g, 'y': y})
+    df = pd.DataFrame(rows)
+    r = sp.did(df, y='y', treat='g', time='t', id='obs',
+               estimator='reg', panel=False)
+    assert r.model_info['panel'] is False
+    assert 'RCS' in r.model_info['estimator']