Add SunAbraham survey validation + TwoStage NotImplementedError propagation

P1: SunAbraham now calls _validate_unit_constant_survey() before
    Rao-Wu collapse, rejecting within-unit-varying survey columns
P2: TwoStageDiD re-raises NotImplementedError (e.g. lonely_psu=adjust)
    instead of swallowing it in the generic except Exception handler

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

Author: igerber

diff_diff/sun_abraham.py

@@ -502,12 +502,21 @@ def fit(
             raise ValueError(f"Missing columns: {missing}")
 
         # Resolve survey design if provided
-        from diff_diff.survey import _resolve_effective_cluster, _resolve_survey_for_fit
+        from diff_diff.survey import (
+            _resolve_effective_cluster,
+            _resolve_survey_for_fit,
+            _validate_unit_constant_survey,
+        )
 
         resolved_survey, survey_weights, survey_weight_type, survey_metadata = (
             _resolve_survey_for_fit(survey_design, data, "analytical")
         )
 
+        # Validate survey columns are constant within units (required for
+        # unit-level collapse in Rao-Wu bootstrap)
+        if resolved_survey is not None:
+            _validate_unit_constant_survey(data, unit, survey_design)
+
         # Bootstrap + survey supported via Rao-Wu rescaled bootstrap.
         # Determine Rao-Wu eligibility from the *original* survey_design
         # (before cluster-as-PSU injection which adds PSU to weights-only designs).

diff_diff/two_stage.py

@@ -585,6 +585,8 @@ def fit(
                     aggregate=aggregate,
                     resolved_survey=resolved_survey,
                 )
+            except NotImplementedError:
+                raise  # Don't swallow explicit rejections (e.g. lonely_psu="adjust")
             except Exception as e:
                 warnings.warn(
                     f"Bootstrap failed: {e}. Skipping bootstrap inference.",