Address code review feedback for CallawaySantAnna covariates

diff_diff/staggered.py

@@ -593,9 +593,6 @@ def fit(
                 "Use n_bootstrap=0 for analytical standard errors."
             )
 
-        # Store covariates for use in estimation
-        self._covariates = covariates
-
         # Create working copy
         df = data.copy()
 
@@ -970,7 +967,10 @@ def _ipw_estimation(
             # Adjusted variance for IPW
             se = np.sqrt(var_t / n_t + var_c * (1 - p_treat) / (n_c * p_treat)) if (n_t > 0 and n_c > 0 and p_treat > 0) else 0.0
 
-            inf_func = np.array([])  # Placeholder
+            # Influence function (for aggregation)
+            inf_treated = (treated_change - np.mean(treated_change)) / n_t
+            inf_control = (control_change - np.mean(control_change)) / n_c
+            inf_func = np.concatenate([inf_treated, -inf_control])
 
         return att, se, inf_func
 

docs/reviews/pr22_covariate_adjustment_review.md

@@ -0,0 +1,65 @@
+# Code Review: PR #22 - CallawaySantAnna Covariate Adjustment
+
+**Reviewer:** Claude
+**Date:** 2026-01-03
+**Status:** Approved with suggestions
+
+## Summary
+
+This PR implements covariate adjustment for the `CallawaySantAnna` estimator, enabling conditional parallel trends assumptions. This addresses a key 1.0 blocker from the roadmap.
+
+## Changes Made During Review
+
+The following issues were fixed directly in this review:
+
+1. **Removed unused instance variable** (`staggered.py:596-597`)
+   - `self._covariates = covariates` was stored but never used
+   - Covariates are passed through the method chain instead
+
+2. **Fixed empty influence function in IPW** (`staggered.py:883-886`)
+   - The unconditional IPW case returned `np.array([])` as a placeholder
+   - Now properly computes the influence function for consistency with other methods
+
+3. **Added test for extreme propensity scores** (`test_staggered.py:656-707`)
+   - Tests that propensity score clipping handles near-perfect separation gracefully
+
+## Suggestions for Future Work
+
+### Standard Error Approximation
+
+The SE calculation in `_outcome_regression` (line 776) ignores estimation error in the regression coefficients:
+
+```python
+# Approximate SE (ignoring estimation error in beta for simplicity)
+se = np.sqrt(var_t / n_t + var_c / n_c)
+```
+
+For a full sandwich variance estimator, see Sant'Anna & Zhao (2020). This is a reasonable approximation for now but could be improved in a future release.
+
+### Propensity Score Model Caching
+
+In `_doubly_robust`, the propensity score is estimated independently (line 932). For efficiency with large datasets, consider refactoring to share the propensity model between IPW and DR estimators when both are called.
+
+### Additional Test Coverage
+
+Consider adding tests for:
+- Near-collinear covariates (testing `_linear_regression` fallback to pseudo-inverse)
+- Missing values in covariates (testing the fallback warning path)
+- Very small sample sizes per group-time cell
+
+### Documentation
+
+The class docstring could include an example showing covariate-adjusted usage:
+
+```python
+# When parallel trends only holds conditional on covariates
+cs = CallawaySantAnna(estimation_method='dr')  # doubly robust
+results = cs.fit(data, outcome='outcome', unit='unit',
+                 time='time', first_treat='first_treat',
+                 covariates=['age', 'income'])
+```
+
+## References
+
+- Callaway, B., & Sant'Anna, P. H. (2021). Difference-in-Differences with multiple time periods. Journal of Econometrics, 225(2), 200-230.
+- Sant'Anna, P. H., & Zhao, J. (2020). Doubly robust difference-in-differences estimators. Journal of Econometrics, 219(1), 101-122.

tests/test_staggered.py

@@ -652,3 +652,56 @@ def test_treatment_effect_recovery_with_covariates(self):
         # Note: we use a generous bound due to finite sample variance
         assert results.overall_att > 0, "ATT should be positive"
         assert abs(results.overall_att - 3.0) < 2.0, f"ATT={results.overall_att} too far from 3.0"
+
+    def test_extreme_propensity_scores(self):
+        """Test handling of covariates that strongly predict treatment.
+
+        When covariates nearly perfectly separate treated/control units,
+        propensity scores approach 0 or 1. The estimator should handle
+        this gracefully via propensity score clipping.
+        """
+        np.random.seed(42)
+        n_units = 100
+        n_periods = 8
+
+        # Generate unit and time identifiers
+        units = np.repeat(np.arange(n_units), n_periods)
+        times = np.tile(np.arange(n_periods), n_units)
+
+        # Create a covariate that strongly predicts treatment
+        # High values -> treated, low values -> never-treated
+        x_strong = np.random.randn(n_units)
+        x_strong_expanded = np.repeat(x_strong, n_periods)
+
+        # Assign treatment based on covariate (top 50% treated at period 4)
+        first_treat = np.zeros(n_units)
+        first_treat[x_strong > np.median(x_strong)] = 4
+        first_treat_expanded = np.repeat(first_treat, n_periods)
+
+        # Generate outcomes
+        post = (times >= first_treat_expanded) & (first_treat_expanded > 0)
+        outcomes = 1.0 + 0.5 * x_strong_expanded + 2.0 * post + np.random.randn(len(units)) * 0.3
+
+        data = pd.DataFrame({
+            'unit': units,
+            'time': times,
+            'outcome': outcomes,
+            'first_treat': first_treat_expanded.astype(int),
+            'x_strong': x_strong_expanded,
+        })
+
+        # IPW should handle extreme propensity scores via clipping
+        cs = CallawaySantAnna(estimation_method='ipw')
+        results = cs.fit(
+            data,
+            outcome='outcome',
+            unit='unit',
+            time='time',
+            first_treat='first_treat',
+            covariates=['x_strong']
+        )
+
+        # Should produce valid results (not NaN or inf)
+        assert np.isfinite(results.overall_att), "ATT should be finite"
+        assert np.isfinite(results.overall_se), "SE should be finite"
+        assert results.overall_se > 0, "SE should be positive"