Address PR #113 Round 7 feedback: fix Python-only LOOCV subsampling

Fix ValueError in joint method when control_obs exceeds max_loocv_samples
without Rust backend. np.random.choice cannot directly sample from a list
of tuples - now samples indices first, then indexes into the list (matching
the pattern already used in the twostep method).

Add test to verify Python-only joint LOOCV subsampling works correctly.

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

Author: igerber

diff_diff/trop.py

@@ -1412,12 +1412,12 @@ def _fit_joint(
                 if control_mask[t, i] and not np.isnan(Y[t, i])
             ]
 
-            # Subsample if needed
-            if len(control_obs) > self.max_loocv_samples:
-                rng = np.random.default_rng(self.seed)
-                control_obs = list(
-                    rng.choice(control_obs, size=self.max_loocv_samples, replace=False)
-                )
+            # Subsample if needed (sample indices to avoid ValueError on list of tuples)
+            rng = np.random.default_rng(self.seed)
+            max_loocv = min(self.max_loocv_samples, len(control_obs))
+            if len(control_obs) > max_loocv:
+                indices = rng.choice(len(control_obs), size=max_loocv, replace=False)
+                control_obs = [control_obs[idx] for idx in indices]
 
             # Grid search with true LOOCV
             for lambda_time_val in self.lambda_time_grid:

tests/test_rust_backend.py

@@ -1554,31 +1554,18 @@ def test_trop_joint_nan_exclusion_rust_python_parity(self):
         trop_rust = TROP(**trop_params)
         results_rust = trop_rust.fit(df.copy(), 'outcome', 'treated', 'unit', 'time')
 
-        # Run with Python-only backend
-        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 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_Python(**trop_params)
+        # Run with Python-only backend using mock.patch to avoid module reload issues
+        # (Module reload breaks isinstance() checks in other tests due to class identity)
+        from unittest.mock import patch
+        import sys
+        trop_module = sys.modules['diff_diff.trop']
+
+        with patch.object(trop_module, 'HAS_RUST_BACKEND', False), \
+             patch.object(trop_module, '_rust_loocv_grid_search_joint', None), \
+             patch.object(trop_module, '_rust_bootstrap_trop_variance_joint', None):
+
+            trop_python = TROP(**trop_params)
             results_python = trop_python.fit(df.copy(), 'outcome', 'treated', 'unit', 'time')
-        finally:
-            # Restore original backend setting
-            if old_backend is None:
-                os.environ.pop('DIFF_DIFF_BACKEND', None)
-            else:
-                os.environ['DIFF_DIFF_BACKEND'] = old_backend
-            import importlib
-            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"
@@ -1657,31 +1644,18 @@ def test_trop_joint_treated_pre_nan_rust_python_parity(self):
         trop_rust = TROP(**trop_params)
         results_rust = trop_rust.fit(df.copy(), 'outcome', 'treated', 'unit', 'time')
 
-        # Run with Python-only backend
-        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 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_Python(**trop_params)
+        # Run with Python-only backend using mock.patch to avoid module reload issues
+        # (Module reload breaks isinstance() checks in other tests due to class identity)
+        from unittest.mock import patch
+        import sys
+        trop_module = sys.modules['diff_diff.trop']
+
+        with patch.object(trop_module, 'HAS_RUST_BACKEND', False), \
+             patch.object(trop_module, '_rust_loocv_grid_search_joint', None), \
+             patch.object(trop_module, '_rust_bootstrap_trop_variance_joint', None):
+
+            trop_python = TROP(**trop_params)
             results_python = trop_python.fit(df.copy(), 'outcome', 'treated', 'unit', 'time')
-        finally:
-            # Restore original backend setting
-            if old_backend is None:
-                os.environ.pop('DIFF_DIFF_BACKEND', None)
-            else:
-                os.environ['DIFF_DIFF_BACKEND'] = old_backend
-            import importlib
-            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"

tests/test_trop.py

@@ -3236,3 +3236,60 @@ def test_joint_rejects_staggered_adoption(self):
         trop = TROP(method="joint")
         with pytest.raises(ValueError, match="staggered adoption"):
             trop.fit(df, 'outcome', 'treated', 'unit', 'time')
+
+    def test_joint_python_loocv_subsampling(self):
+        """Test that joint method works with Python-only LOOCV when control_obs > max_loocv_samples.
+
+        This tests the fix for PR #113 Round 7 feedback (P1): Python fallback
+        LOOCV sampling could raise ValueError when control_obs is a list of tuples.
+        """
+        from unittest.mock import patch
+        import sys
+
+        np.random.seed(42)
+        # Create data with many control observations (> default max_loocv_samples=500)
+        n_units, n_periods = 30, 25  # 30*25 = 750 observations, most are control
+        n_treated = 3
+        n_post = 3
+
+        data = []
+        for i in range(n_units):
+            is_treated = i < n_treated
+            for t in range(n_periods):
+                post = t >= (n_periods - n_post)
+                y = 10.0 + i * 0.1 + t * 0.1 + np.random.randn() * 0.5
+                treatment_indicator = 1 if (is_treated and post) else 0
+                if treatment_indicator:
+                    y += 2.0
+                data.append({
+                    'unit': i,
+                    'time': t,
+                    'outcome': y,
+                    'treated': treatment_indicator,
+                })
+
+        df = pd.DataFrame(data)
+
+        # Patch to force Python backend and set small max_loocv_samples
+        trop_module = sys.modules['diff_diff.trop']
+
+        with patch.object(trop_module, 'HAS_RUST_BACKEND', False), \
+             patch.object(trop_module, '_rust_loocv_grid_search_joint', None), \
+             patch.object(trop_module, '_rust_bootstrap_trop_variance_joint', None):
+
+            # Use small max_loocv_samples to trigger subsampling
+            trop_est = TROP(
+                method="joint",
+                lambda_time_grid=[1.0],
+                lambda_unit_grid=[1.0],
+                lambda_nn_grid=[0.0],
+                max_loocv_samples=100,  # Force subsampling (control_obs > 100)
+                n_bootstrap=0,
+                seed=42
+            )
+
+            # This should not raise ValueError
+            results = trop_est.fit(df, 'outcome', 'treated', 'unit', 'time')
+
+            assert isinstance(results, TROPResults)
+            assert np.isfinite(results.att)