Fix not-yet-treated control mask to respect anticipation parameter

The not-yet-treated control group in both ContinuousDiD and
CallawaySantAnna used `G > t` instead of `G > t + anticipation`,
incorrectly including cohorts in the anticipation window as controls.
This matches R's `did::compute.att_gt()` logic.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: igerber

diff_diff/continuous_did.py

@@ -690,7 +690,9 @@ def _compute_dose_response_gt(
             control_mask = never_treated_mask
         else:
             # Not-yet-treated: never-treated + first_treat > t
-            control_mask = never_treated_mask | ((unit_cohorts > t) & (unit_cohorts != g))
+            control_mask = never_treated_mask | (
+                (unit_cohorts > t + self.anticipation) & (unit_cohorts != g)
+            )
         n_control = int(np.sum(control_mask))
         if n_control == 0:
             warnings.warn(

diff_diff/staggered.py

@@ -491,7 +491,9 @@ def _compute_att_gt_fast(
         else:  # not_yet_treated
             # Not yet treated at time t: never-treated OR (first_treat > t AND not cohort g)
             # Must exclude cohort g since they are the treated group for this ATT(g,t)
-            control_mask = never_treated_mask | ((unit_cohorts > t) & (unit_cohorts != g))
+            control_mask = never_treated_mask | (
+                (unit_cohorts > t + self.anticipation) & (unit_cohorts != g)
+            )
 
         # Extract outcomes for base and post periods
         y_base = outcome_matrix[:, base_col]

docs/methodology/REGISTRY.md

@@ -341,6 +341,8 @@ The multiplier bootstrap uses random weights w_i with E[w]=0 and Var(w)=1:
 - Anticipation: `anticipation` parameter shifts reference period
   - Group aggregation includes periods t >= g - anticipation (not just t >= g)
   - Both analytical SE and bootstrap SE aggregation respect anticipation
+  - Not-yet-treated + anticipation: control mask uses `G > t + anticipation`
+    (not just `G > t`) to exclude cohorts in the anticipation window
 - Rank-deficient design matrix (covariate collinearity):
   - Detection: Pivoted QR decomposition with tolerance `1e-07` (R's `qr()` default)
   - Handling: Warns and drops linearly dependent columns, sets NA for dropped coefficients (R-style, matches `lm()`)
@@ -378,7 +380,7 @@ The multiplier bootstrap uses random weights w_i with E[w]=0 and Var(w)=1:
   - Always excludes cohort g from controls when computing ATT(g,t)
   - This applies to both pre-treatment (t < g) and post-treatment (t >= g) periods
   - For pre-treatment periods: even though cohort g hasn't been treated yet at time t, they are the treated group for this ATT(g,t) and cannot serve as their own controls
-  - Control mask: `never_treated OR (first_treat > t AND first_treat != g)`
+  - Control mask: `never_treated OR (first_treat > t + anticipation AND first_treat != g)`
 
 **Reference implementation(s):**
 - R: `did::att_gt()` (Callaway & Sant'Anna's official package)
@@ -427,6 +429,10 @@ This is stronger than standard PT because it conditions on specific dose values.
 - **All-same dose**: B-spline basis collapses; ACRT(d) = 0 everywhere.
 - **Rank deficiency**: When n_treated <= n_basis, cell is skipped.
 - **Balanced panel required**: Matches R `contdid` v0.1.0.
+- **Anticipation + not-yet-treated**: Control mask uses `G > t + anticipation`
+  (not just `G > t`) to exclude cohorts in the anticipation window from
+  not-yet-treated controls. When `anticipation=0` (default), behavior is
+  unchanged.
 - **Boundary knots**: Knots are built once from all treated doses (global, not per-cell) to ensure a common basis across (g,t) cells for aggregation. Evaluation grid is clamped to training-dose boundary knots (`range(dose)`). R's `contdid` v0.1.0 has an inconsistency where `splines2::bSpline(dvals)` uses `range(dvals)` instead of `range(dose)`, which can produce extrapolation artifacts at dose grid extremes. Our approach avoids extrapolation and is methodologically sound.
 
 ### Implementation Checklist

tests/test_continuous_did.py

@@ -801,6 +801,66 @@ def test_anticipation_event_study(self):
         )
         assert np.isfinite(results.event_study_effects[-1]["effect"])
 
+    def test_anticipation_not_yet_treated_excludes_anticipation_window(self):
+        """Not-yet-treated controls must exclude cohorts in the anticipation window.
+
+        With anticipation=1 and cohort g=3, computing ATT(g=3, t=4) should use
+        threshold t + anticipation = 5, so cohort g=5 (unit_cohorts == 5) fails
+        > 5 and is correctly excluded. Without the fix, threshold is t=4 and
+        cohort g=5 passes > 4, contaminating controls with treated units.
+        """
+        rng = np.random.default_rng(42)
+        n_per_group = 20
+        periods = [1, 2, 3, 4, 5, 6]
+
+        rows = []
+        # Never-treated group
+        for i in range(n_per_group):
+            uid = i
+            for t in periods:
+                rows.append({
+                    "unit": uid, "period": t, "first_treat": 0,
+                    "dose": 0.0, "outcome": rng.normal(0, 0.5),
+                })
+
+        # Early cohort: g=3, treatment effect = +5*dose at t>=3
+        for i in range(n_per_group):
+            uid = n_per_group + i
+            d = rng.uniform(1, 3)
+            for t in periods:
+                y = rng.normal(0, 0.5) + (5.0 * d if t >= 3 else 0)
+                rows.append({
+                    "unit": uid, "period": t, "first_treat": 3,
+                    "dose": d, "outcome": y,
+                })
+
+        # Late cohort: g=5, treatment effect = +5*dose at t>=5
+        for i in range(n_per_group):
+            uid = 2 * n_per_group + i
+            d = rng.uniform(1, 3)
+            for t in periods:
+                y = rng.normal(0, 0.5) + (5.0 * d if t >= 5 else 0)
+                rows.append({
+                    "unit": uid, "period": t, "first_treat": 5,
+                    "dose": d, "outcome": y,
+                })
+
+        data = pd.DataFrame(rows)
+
+        est = ContinuousDiD(
+            anticipation=1, control_group="not_yet_treated", n_bootstrap=0,
+        )
+        results = est.fit(
+            data, "outcome", "unit", "period", "first_treat", "dose",
+        )
+
+        assert np.isfinite(results.overall_att), (
+            "overall_att should be finite with anticipation + not_yet_treated"
+        )
+        assert results.dose_response_att is not None, (
+            "dose-response curve should exist"
+        )
+
 
 class TestEmptyPostTreatment:
     """Test guard for empty post-treatment cells."""

tests/test_staggered.py

@@ -2692,6 +2692,40 @@ def test_group_effects_anticipation_boundary(self):
             post_treatment = [t for (gg, t) in gt_for_group if t >= g]
             assert len(post_treatment) > 0, "Should have post-treatment effects"
 
+    def test_not_yet_treated_with_anticipation_excludes_anticipation_window(self):
+        """Not-yet-treated controls must exclude cohorts in the anticipation window.
+
+        With anticipation=1, the control mask should use G > t + anticipation
+        (not just G > t). Without the fix, cohorts about to be treated are
+        incorrectly included as controls, biasing pre-treatment ATTs toward
+        the treatment effect (~3.0) instead of near zero.
+        """
+        data = generate_staggered_data(
+            n_units=100, n_periods=10, n_cohorts=2,
+            treatment_effect=3.0, seed=42,
+        )
+
+        cs = CallawaySantAnna(anticipation=1, control_group="not_yet_treated")
+        result = cs.fit(
+            data, outcome="outcome", unit="unit",
+            time="time", first_treat="first_treat",
+        )
+
+        groups = sorted(
+            g for g in data[data["first_treat"] > 0]["first_treat"].unique()
+        )
+
+        for g in groups:
+            for (gg, t), eff in result.group_time_effects.items():
+                if gg != g:
+                    continue
+                # Pre-treatment: t < g - anticipation
+                if t < g - 1:
+                    assert abs(eff["effect"]) < 1.5, (
+                        f"Pre-treatment ATT(g={g}, t={t}) = {eff['effect']:.3f} "
+                        f"should be near zero (< 1.5); contaminated controls?"
+                    )
+
 
 class TestCallawaySantAnnaTStatNaN:
     """Tests for NaN t_stat when SE is invalid."""