Address AI review findings for TROP split

P1: Restore absorbing-state remediation guidance in global path error message
P2: Extract _soft_threshold_svd to module-level function, eliminating MRO
    dependency between TROPGlobalMixin and TROPLocalMixin
P2: Add 5 regression tests for module split (dispatch, finite lambda_nn,
    error message consistency)
P3: Remove unused unit_to_idx/period_to_idx in _fit_with_fixed_lambda

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

Author: igerber

diff_diff/trop_global.py

@@ -24,6 +24,7 @@
     _rust_bootstrap_trop_variance_global,
     _rust_loocv_grid_search_global,
 )
+from diff_diff.trop_local import _soft_threshold_svd
 from diff_diff.trop_results import TROPResults
 from diff_diff.utils import safe_inference
 
@@ -39,8 +40,6 @@ class TROPGlobalMixin:
     - Inference params: ``alpha``, ``n_bootstrap``, ``seed``
     - State: ``results_``, ``is_fitted_``
 
-    The ``_solve_global_with_lowrank`` method calls ``self._soft_threshold_svd``
-    which is defined in ``TROPLocalMixin`` and resolved via Python MRO.
     """
 
     # Type hints for attributes accessed from the main TROP class
@@ -395,9 +394,6 @@ def _solve_global_with_lowrank(
         The (1-W) masking is already applied to delta by _compute_global_weights,
         so treated observations have zero weight and do not affect the fit.
 
-        Note: calls ``self._soft_threshold_svd`` which is defined in
-        ``TROPLocalMixin`` and resolved via Python MRO on the ``TROP`` class.
-
         Parameters
         ----------
         Y : np.ndarray
@@ -468,7 +464,7 @@ def _solve_global_with_lowrank(
 
                 # Proximal step: soft-threshold singular values
                 L_inner_prev = L_inner
-                L_inner = self._soft_threshold_svd(gradient_step, threshold)
+                L_inner = _soft_threshold_svd(gradient_step, threshold)
                 t_fista = t_fista_new
 
                 # Convergence check (L_inner_prev holds the pre-SVD value)
@@ -564,7 +560,9 @@ def _fit_global(
         if violating_units:
             raise ValueError(
                 f"Treatment indicator is not an absorbing state for units: {violating_units}. "
-                f"D[t, unit] must be monotonic non-decreasing."
+                f"D[t, unit] must be monotonic non-decreasing (once treated, always treated). "
+                f"If this is event-study style data, convert to absorbing state: "
+                f"D[t, i] = 1 for all t >= first treatment period."
             )
 
         # Identify treated observations

diff_diff/trop_local.py

@@ -27,6 +27,76 @@
 from diff_diff.trop_results import _PrecomputedStructures
 
 
+# Module-level convergence tolerance for SVD singular value truncation.
+# Singular values below this threshold after soft-thresholding are treated
+# as zero to improve numerical stability.
+_CONVERGENCE_TOL_SVD: float = 1e-10
+
+
+def _soft_threshold_svd(
+    M: np.ndarray,
+    threshold: float,
+    convergence_tol: float = _CONVERGENCE_TOL_SVD,
+) -> np.ndarray:
+    """
+    Apply soft-thresholding to singular values (proximal operator for nuclear norm).
+
+    Parameters
+    ----------
+    M : np.ndarray
+        Input matrix.
+    threshold : float
+        Soft-thresholding parameter.
+    convergence_tol : float, default=1e-10
+        Singular values below this after thresholding are treated as zero.
+
+    Returns
+    -------
+    np.ndarray
+        Matrix with soft-thresholded singular values.
+    """
+    if threshold <= 0:
+        return M
+
+    # Handle NaN/Inf values in input
+    if not np.isfinite(M).all():
+        M = np.nan_to_num(M, nan=0.0, posinf=0.0, neginf=0.0)
+
+    try:
+        U, s, Vt = np.linalg.svd(M, full_matrices=False)
+    except np.linalg.LinAlgError:
+        # SVD failed, return zero matrix
+        return np.zeros_like(M)
+
+    # Check for numerical issues in SVD output
+    if not (np.isfinite(U).all() and np.isfinite(s).all() and np.isfinite(Vt).all()):
+        # SVD produced non-finite values, return zero matrix
+        return np.zeros_like(M)
+
+    s_thresh = np.maximum(s - threshold, 0)
+
+    # Use truncated reconstruction with only non-zero singular values
+    nonzero_mask = s_thresh > convergence_tol
+    if not np.any(nonzero_mask):
+        return np.zeros_like(M)
+
+    # Truncate to non-zero components for numerical stability
+    U_trunc = U[:, nonzero_mask]
+    s_trunc = s_thresh[nonzero_mask]
+    Vt_trunc = Vt[nonzero_mask, :]
+
+    # Compute result, suppressing expected numerical warnings from
+    # ill-conditioned matrices during alternating minimization
+    with np.errstate(divide='ignore', over='ignore', invalid='ignore'):
+        result = (U_trunc * s_trunc) @ Vt_trunc
+
+    # Replace any NaN/Inf in result with zeros
+    if not np.isfinite(result).all():
+        result = np.nan_to_num(result, nan=0.0, posinf=0.0, neginf=0.0)
+
+    return result
+
+
 class TROPLocalMixin:
     """Mixin providing local (observation-specific) estimation for TROP.
 
@@ -378,61 +448,8 @@ def _soft_threshold_svd(
         M: np.ndarray,
         threshold: float,
     ) -> np.ndarray:
-        """
-        Apply soft-thresholding to singular values (proximal operator for nuclear norm).
-
-        Parameters
-        ----------
-        M : np.ndarray
-            Input matrix.
-        threshold : float
-            Soft-thresholding parameter.
-
-        Returns
-        -------
-        np.ndarray
-            Matrix with soft-thresholded singular values.
-        """
-        if threshold <= 0:
-            return M
-
-        # Handle NaN/Inf values in input
-        if not np.isfinite(M).all():
-            M = np.nan_to_num(M, nan=0.0, posinf=0.0, neginf=0.0)
-
-        try:
-            U, s, Vt = np.linalg.svd(M, full_matrices=False)
-        except np.linalg.LinAlgError:
-            # SVD failed, return zero matrix
-            return np.zeros_like(M)
-
-        # Check for numerical issues in SVD output
-        if not (np.isfinite(U).all() and np.isfinite(s).all() and np.isfinite(Vt).all()):
-            # SVD produced non-finite values, return zero matrix
-            return np.zeros_like(M)
-
-        s_thresh = np.maximum(s - threshold, 0)
-
-        # Use truncated reconstruction with only non-zero singular values
-        nonzero_mask = s_thresh > self.CONVERGENCE_TOL_SVD
-        if not np.any(nonzero_mask):
-            return np.zeros_like(M)
-
-        # Truncate to non-zero components for numerical stability
-        U_trunc = U[:, nonzero_mask]
-        s_trunc = s_thresh[nonzero_mask]
-        Vt_trunc = Vt[nonzero_mask, :]
-
-        # Compute result, suppressing expected numerical warnings from
-        # ill-conditioned matrices during alternating minimization
-        with np.errstate(divide='ignore', over='ignore', invalid='ignore'):
-            result = (U_trunc * s_trunc) @ Vt_trunc
-
-        # Replace any NaN/Inf in result with zeros
-        if not np.isfinite(result).all():
-            result = np.nan_to_num(result, nan=0.0, posinf=0.0, neginf=0.0)
-
-        return result
+        """Delegate to module-level ``_soft_threshold_svd``."""
+        return _soft_threshold_svd(M, threshold, self.CONVERGENCE_TOL_SVD)
 
     def _weighted_nuclear_norm_solve(
         self,
@@ -948,9 +965,6 @@ def _fit_with_fixed_lambda(
         n_units = len(all_units)
         n_periods = len(all_periods)
 
-        unit_to_idx = {u: i for i, u in enumerate(all_units)}
-        period_to_idx = {p: i for i, p in enumerate(all_periods)}
-
         # Vectorized: use pivot for O(1) reshaping instead of O(n) iterrows loop
         Y = (
             data.pivot(index=time, columns=unit, values=outcome)

tests/test_trop.py

@@ -3662,3 +3662,97 @@ def test_local_bootstrap_zero_draws_returns_nan_se(self):
         assert np.isnan(se), f"SE should be NaN when 0 draws succeed, got {se}"
         assert len(dist) == 0
 
+
+class TestTROPModuleSplit:
+    """Regression tests for the trop.py -> trop_global.py / trop_local.py split."""
+
+    @staticmethod
+    def _make_panel():
+        """Create a simple balanced panel for split regression tests."""
+        rng = np.random.default_rng(42)
+        n_units, n_periods = 8, 6
+        rows = []
+        for i in range(n_units):
+            treated = i < 3  # 3 treated, 5 control
+            for t in range(n_periods):
+                y = rng.normal(0, 1)
+                if treated and t >= 4:
+                    y += 2.0  # treatment effect
+                rows.append({
+                    "unit": i, "time": t, "outcome": y,
+                    "treated": 1 if treated and t >= 4 else 0,
+                })
+        return pd.DataFrame(rows)
+
+    def test_global_absorbing_state_error_has_remediation_guidance(self):
+        """Global path ValueError for non-absorbing D includes remediation text."""
+        df = self._make_panel()
+        # Break absorbing state: unit 0 goes 0->1->0
+        df.loc[(df["unit"] == 0) & (df["time"] == 5), "treated"] = 0
+
+        with pytest.raises(ValueError, match="once treated, always treated"):
+            TROP(method="global").fit(df, "outcome", "treated", "unit", "time")
+
+        with pytest.raises(ValueError, match="convert to absorbing state"):
+            TROP(method="global").fit(df, "outcome", "treated", "unit", "time")
+
+    def test_global_finite_lambda_nn_exercises_lowrank_path(self):
+        """method='global' with finite lambda_nn successfully fits the low-rank solver."""
+        df = self._make_panel()
+        trop_est = TROP(
+            method="global",
+            lambda_time_grid=[0.0],
+            lambda_unit_grid=[0.0],
+            lambda_nn_grid=[0.1],  # finite -> exercises _solve_global_with_lowrank
+            n_bootstrap=5,
+            seed=42,
+        )
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            result = trop_est.fit(df, "outcome", "treated", "unit", "time")
+        assert np.isfinite(result.att)
+
+    def test_local_finite_lambda_nn_exercises_nuclear_norm(self):
+        """method='local' with finite lambda_nn exercises weighted nuclear norm solver."""
+        df = self._make_panel()
+        trop_est = TROP(
+            method="local",
+            lambda_time_grid=[0.0],
+            lambda_unit_grid=[0.0],
+            lambda_nn_grid=[0.1],  # finite -> exercises _weighted_nuclear_norm_solve
+            n_bootstrap=5,
+            seed=42,
+        )
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            result = trop_est.fit(df, "outcome", "treated", "unit", "time")
+        assert np.isfinite(result.att)
+
+    def test_method_dispatch_global_uses_fit_global(self):
+        """method='global' dispatches to _fit_global from TROPGlobalMixin."""
+        from unittest.mock import patch
+
+        df = self._make_panel()
+        trop_est = TROP(method="global", n_bootstrap=2, seed=42)
+
+        with patch.object(TROP, '_fit_global', wraps=trop_est._fit_global) as mock_fg:
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore")
+                trop_est.fit(df, "outcome", "treated", "unit", "time")
+            mock_fg.assert_called_once()
+
+    def test_method_dispatch_local_does_not_use_fit_global(self):
+        """method='local' does NOT call _fit_global."""
+        from unittest.mock import patch
+
+        df = self._make_panel()
+        trop_est = TROP(method="local", n_bootstrap=2, seed=42,
+                        lambda_time_grid=[0.0], lambda_unit_grid=[0.0],
+                        lambda_nn_grid=[np.inf])
+
+        with patch.object(TROP, '_fit_global') as mock_fg:
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore")
+                trop_est.fit(df, "outcome", "treated", "unit", "time")
+            mock_fg.assert_not_called()
+