Thread survey_design through bacon_decompose() and triple_difference() convenience functions from PR #226 review (round 15)

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

Author: igerber

diff_diff/bacon.py

@@ -1043,6 +1043,7 @@ def bacon_decompose(
     time: str,
     first_treat: str,
     weights: str = "approximate",
+    survey_design: object = None,
 ) -> BaconDecompositionResults:
     """
     Convenience function for Goodman-Bacon decomposition.
@@ -1119,4 +1120,4 @@ def bacon_decompose(
     CallawaySantAnna : Robust estimator that avoids forbidden comparisons
     """
     decomp = BaconDecomposition(weights=weights)
-    return decomp.fit(data, outcome, unit, time, first_treat)
+    return decomp.fit(data, outcome, unit, time, first_treat, survey_design=survey_design)

diff_diff/triple_diff.py

@@ -1702,6 +1702,7 @@ def triple_difference(
     cluster: Optional[str] = None,
     alpha: float = 0.05,
     rank_deficient_action: str = "warn",
+    survey_design: object = None,
 ) -> TripleDifferenceResults:
     """
     Estimate Triple Difference (DDD) treatment effect.
@@ -1776,4 +1777,5 @@ def triple_difference(
         partition=partition,
         time=time,
         covariates=covariates,
+        survey_design=survey_design,
     )

tests/test_survey_phase3.py

@@ -425,6 +425,21 @@ def test_exact_weights_survey_weighted(self, staggered_survey_data):
         diffs = [abs(exact_weights[k] - approx_weights[k]) for k in common_keys]
         assert max(diffs) > 1e-10, "Exact and approximate weights should differ"
 
+    def test_convenience_function(self, staggered_survey_data):
+        """bacon_decompose() convenience function threads survey_design."""
+        from diff_diff.bacon import bacon_decompose
+
+        sd = SurveyDesign(weights="weight")
+        result = bacon_decompose(
+            staggered_survey_data,
+            "outcome",
+            "unit",
+            "time",
+            "first_treat",
+            survey_design=sd,
+        )
+        assert result.survey_metadata is not None
+
 
 # =============================================================================
 # TripleDifference
@@ -555,6 +570,22 @@ def test_reg_with_covariates_survey(self, ddd_survey_data):
         # Survey df should be used for inference
         assert result.survey_metadata.df_survey is not None
 
+    def test_convenience_function(self, ddd_survey_data):
+        """triple_difference() convenience function threads survey_design."""
+        from diff_diff.triple_diff import triple_difference
+
+        sd = SurveyDesign(weights="weight")
+        result = triple_difference(
+            ddd_survey_data,
+            "outcome",
+            "group",
+            "partition",
+            "time",
+            estimation_method="reg",
+            survey_design=sd,
+        )
+        assert result.survey_metadata is not None
+
 
 # =============================================================================
 # ContinuousDiD