Address PR #113 Round 6 feedback: fix staggered adoption check and test reloads

- Fix simultaneous-adoption check to use observed periods only, avoiding
  false positives on unbalanced panels where missing entries were filled as 0
- Add zero-weight guard in Python joint solver matching Rust's behavior
- Fix backend parity tests to properly reload trop module using sys.modules

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

Author: igerber

diff_diff/trop.py

@@ -1079,6 +1079,11 @@ def _solve_joint_no_lowrank(
 
         sqrt_weights = np.sqrt(np.maximum(weights, 0))
 
+        # Check for all-zero weights (matches Rust's sum_w < 1e-10 check)
+        sum_w = np.sum(weights)
+        if sum_w < 1e-10:
+            raise ValueError("All weights are zero - cannot estimate")
+
         # Build design matrix: [intercept, unit_dummies, time_dummies, treatment]
         # Total columns: 1 + n_units + n_periods + 1
         # But we need to drop one unit and one time dummy for identification
@@ -1321,11 +1326,17 @@ def _fit_joint(
             raise ValueError("Need at least 2 pre-treatment periods")
 
         # Check for staggered adoption (joint method requires simultaneous treatment)
+        # Use only observed periods (skip missing) to avoid false positives on unbalanced panels
         first_treat_by_unit = []
         for i in treated_unit_idx:
-            treated_periods_i = np.where(D[:, i] == 1)[0]
-            if len(treated_periods_i) > 0:
-                first_treat_by_unit.append(treated_periods_i[0])
+            observed_mask = ~missing_mask[:, i]
+            # Get D values for observed periods only
+            observed_d = D[observed_mask, i]
+            observed_periods = np.where(observed_mask)[0]
+            # Find first treatment among observed periods
+            treated_idx = np.where(observed_d == 1)[0]
+            if len(treated_idx) > 0:
+                first_treat_by_unit.append(observed_periods[treated_idx[0]])
 
         unique_starts = sorted(set(first_treat_by_unit))
         if len(unique_starts) > 1:

tests/test_rust_backend.py

@@ -1559,11 +1559,15 @@ def test_trop_joint_nan_exclusion_rust_python_parity(self):
         try:
             os.environ['DIFF_DIFF_BACKEND'] = 'python'
             # Need to reimport to pick up new backend setting
+            # Must reload both _backend AND trop modules since trop imports
+            # HAS_RUST_BACKEND and Rust functions at module load time
             import importlib
-            import diff_diff._backend
-            importlib.reload(diff_diff._backend)
+            import sys
+            importlib.reload(sys.modules['diff_diff._backend'])
+            importlib.reload(sys.modules['diff_diff.trop'])
+            from diff_diff.trop import TROP as TROP_Python
 
-            trop_python = TROP(**trop_params)
+            trop_python = TROP_Python(**trop_params)
             results_python = trop_python.fit(df.copy(), 'outcome', 'treated', 'unit', 'time')
         finally:
             # Restore original backend setting
@@ -1572,8 +1576,9 @@ def test_trop_joint_nan_exclusion_rust_python_parity(self):
             else:
                 os.environ['DIFF_DIFF_BACKEND'] = old_backend
             import importlib
-            import diff_diff._backend
-            importlib.reload(diff_diff._backend)
+            import sys
+            importlib.reload(sys.modules['diff_diff._backend'])
+            importlib.reload(sys.modules['diff_diff.trop'])
 
         # Both should produce finite results
         assert np.isfinite(results_rust.att), f"Rust ATT {results_rust.att} should be finite"
@@ -1656,11 +1661,16 @@ def test_trop_joint_treated_pre_nan_rust_python_parity(self):
         old_backend = os.environ.get('DIFF_DIFF_BACKEND')
         try:
             os.environ['DIFF_DIFF_BACKEND'] = 'python'
+            # Need to reimport to pick up new backend setting
+            # Must reload both _backend AND trop modules since trop imports
+            # HAS_RUST_BACKEND and Rust functions at module load time
             import importlib
-            import diff_diff._backend
-            importlib.reload(diff_diff._backend)
+            import sys
+            importlib.reload(sys.modules['diff_diff._backend'])
+            importlib.reload(sys.modules['diff_diff.trop'])
+            from diff_diff.trop import TROP as TROP_Python
 
-            trop_python = TROP(**trop_params)
+            trop_python = TROP_Python(**trop_params)
             results_python = trop_python.fit(df.copy(), 'outcome', 'treated', 'unit', 'time')
         finally:
             # Restore original backend setting
@@ -1669,8 +1679,9 @@ def test_trop_joint_treated_pre_nan_rust_python_parity(self):
             else:
                 os.environ['DIFF_DIFF_BACKEND'] = old_backend
             import importlib
-            import diff_diff._backend
-            importlib.reload(diff_diff._backend)
+            import sys
+            importlib.reload(sys.modules['diff_diff._backend'])
+            importlib.reload(sys.modules['diff_diff.trop'])
 
         # Both should produce finite results
         assert np.isfinite(results_rust.att), f"Rust ATT {results_rust.att} should be finite"