Fix TROP joint method Rust/Python parity issues (PR #113 feedback)

Address P1 review feedback:
- P1-2: Align nuclear-norm threshold scaling by using eta * lambda_nn
  for soft-threshold SVD step in Python (matching Rust implementation)
- P1-1: Add comprehensive NaN handling in _compute_joint_weights,
  _solve_joint_no_lowrank, and _solve_joint_with_lowrank

Add tests for NaN handling parity between backends.

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

Author: igerber

diff_diff/trop.py

@@ -910,23 +910,32 @@ def _compute_joint_weights(
         delta_time = np.exp(-lambda_time * dist_time)
 
         # Unit weights: RMSE to average treated trajectory over pre-periods
-        # Compute average treated trajectory
-        average_treated = np.mean(Y[:, treated_unit_idx], axis=1)
+        # Compute average treated trajectory (use nanmean to handle NaN)
+        average_treated = np.nanmean(Y[:, treated_unit_idx], axis=1)
 
         # Pre-period mask: 1 in pre, 0 in post
         pre_mask = np.ones(n_periods, dtype=float)
         pre_mask[-treated_periods:] = 0.0
 
         # Compute RMS distance for each unit
         # dist_unit[i] = sqrt(sum_pre(avg_tr - Y_i)^2 / n_pre)
-        diff_sq = ((average_treated[:, np.newaxis] - Y) ** 2) * pre_mask[:, np.newaxis]
+        # Use NaN-safe operations: treat NaN differences as 0 (excluded)
+        diff = average_treated[:, np.newaxis] - Y
+        diff_sq = np.where(np.isfinite(diff), diff ** 2, 0.0) * pre_mask[:, np.newaxis]
+
+        # Count valid observations per unit in pre-period
+        valid_count = np.sum(
+            np.isfinite(Y) * pre_mask[:, np.newaxis], axis=0
+        )
         sum_sq = np.sum(diff_sq, axis=0)
         n_pre = np.sum(pre_mask)
 
         if n_pre == 0:
             raise ValueError("No pre-treatment periods")
 
-        dist_unit = np.sqrt(sum_sq / n_pre)
+        # Use valid count per unit (avoid division by zero)
+        valid_count = np.maximum(valid_count, 1)
+        dist_unit = np.sqrt(sum_sq / valid_count)
         delta_unit = np.exp(-lambda_unit * dist_unit)
 
         # Outer product: (n_periods x n_units)
@@ -1050,6 +1059,15 @@ def _solve_joint_no_lowrank(
         y = Y.flatten()  # length n_periods * n_units
         w = D.flatten()
         weights = delta.flatten()
+
+        # Handle NaN values: zero weight for NaN outcomes/weights, impute with 0
+        # This ensures NaN observations don't contribute to estimation
+        valid_y = np.isfinite(y)
+        valid_w = np.isfinite(weights)
+        valid_mask = valid_y & valid_w
+        weights = np.where(valid_mask, weights, 0.0)
+        y = np.where(valid_mask, y, 0.0)
+
         sqrt_weights = np.sqrt(np.maximum(weights, 0))
 
         # Build design matrix: [intercept, unit_dummies, time_dummies, treatment]
@@ -1132,20 +1150,24 @@ def _solve_joint_with_lowrank(
         """
         n_periods, n_units = Y.shape
 
+        # Handle NaN values: impute with 0 for computations
+        # The solver will also zero weights for NaN observations
+        Y_safe = np.where(np.isfinite(Y), Y, 0.0)
+
         # Initialize L = 0
         L = np.zeros((n_periods, n_units))
 
         for iteration in range(max_iter):
             L_old = L.copy()
 
             # Step 1: Fix L, solve for (mu, alpha, beta, tau)
-            # Adjusted outcome: Y - L
-            Y_adj = Y - L
+            # Adjusted outcome: Y - L (using NaN-safe Y)
+            Y_adj = Y_safe - L
             mu, alpha, beta, tau = self._solve_joint_no_lowrank(Y_adj, D, delta)
 
             # Step 2: Fix (mu, alpha, beta, tau), update L
-            # Residual: R = Y - mu - alpha - beta - tau*D
-            R = Y - mu - alpha[np.newaxis, :] - beta[:, np.newaxis] - tau * D
+            # Residual: R = Y - mu - alpha - beta - tau*D (using NaN-safe Y)
+            R = Y_safe - mu - alpha[np.newaxis, :] - beta[:, np.newaxis] - tau * D
 
             # Weighted proximal step for L (soft-threshold SVD)
             # Normalize weights
@@ -1160,7 +1182,9 @@ def _solve_joint_with_lowrank(
             gradient_step = L + delta_norm * (R - L)
 
             # Soft-threshold singular values
-            L = self._soft_threshold_svd(gradient_step, lambda_nn)
+            # Use eta * lambda_nn for proper proximal step size (matches Rust)
+            eta = 1.0 / delta_max if delta_max > 0 else 1.0
+            L = self._soft_threshold_svd(gradient_step, eta * lambda_nn)
 
             # Check convergence
             if np.max(np.abs(L - L_old)) < tol:

tests/test_rust_backend.py

@@ -1364,6 +1364,65 @@ def test_trop_joint_and_twostep_agree_in_direction(self):
         # Both should have same sign (both positive for true_effect=2.0)
         assert np.sign(results_joint.att) == np.sign(results_twostep.att)
 
+    def test_trop_joint_handles_nan_outcomes(self):
+        """Test TROP joint method handles NaN outcome values gracefully."""
+        import pandas as pd
+        from diff_diff import TROP
+
+        np.random.seed(42)
+        n_units, n_periods = 20, 10
+        n_treated = 5
+        n_post = 3
+        true_effect = 2.0
+
+        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.2 + t * 0.3 + np.random.randn() * 0.5
+                treatment_indicator = 1 if (is_treated and post) else 0
+                if treatment_indicator:
+                    y += true_effect
+                data.append({
+                    'unit': i,
+                    'time': t,
+                    'outcome': y,
+                    'treated': treatment_indicator,
+                })
+
+        df = pd.DataFrame(data)
+
+        # Introduce NaN values in control observations (pre-treatment periods)
+        # Set 5% of control pre-treatment observations to NaN
+        nan_indices = []
+        for idx, row in df.iterrows():
+            if row['treated'] == 0 and row['time'] < (n_periods - n_post):
+                if np.random.rand() < 0.05:
+                    nan_indices.append(idx)
+        df.loc[nan_indices, 'outcome'] = np.nan
+
+        n_nan = len(nan_indices)
+        assert n_nan > 0, "Should have introduced some NaN values"
+
+        trop = TROP(
+            method="joint",
+            lambda_time_grid=[0.0, 1.0],
+            lambda_unit_grid=[0.0, 1.0],
+            lambda_nn_grid=[0.0, 0.1],
+            n_bootstrap=20,
+            seed=42
+        )
+        results = trop.fit(df, 'outcome', 'treated', 'unit', 'time')
+
+        # Results should be finite (NaN observations are excluded)
+        assert np.isfinite(results.att), f"ATT {results.att} should be finite with NaN data"
+        assert np.isfinite(results.se), f"SE {results.se} should be finite with NaN data"
+        assert results.se >= 0, "SE should be non-negative"
+
+        # ATT should still be positive (true effect is positive)
+        assert results.att > 0, f"ATT {results.att:.2f} should be positive"
+
 
 class TestFallbackWhenNoRust:
     """Test that pure Python fallback works when Rust is unavailable."""

tests/test_trop.py

@@ -2948,3 +2948,69 @@ def test_joint_loocv_score_internal(self, simple_panel_data):
             treated_periods, n_units, n_periods
         )
         assert np.isfinite(score2) or np.isinf(score2), "Score should be finite or inf"
+
+    def test_joint_handles_nan_outcomes(self, simple_panel_data):
+        """Joint method handles NaN outcome values gracefully."""
+        # Introduce NaN in some control observations
+        data = simple_panel_data.copy()
+        control_mask = data['treated'] == 0
+        control_indices = data[control_mask].index.tolist()
+
+        # Set 5 random control observations to NaN
+        np.random.seed(42)
+        nan_indices = np.random.choice(control_indices, size=5, replace=False)
+        data.loc[nan_indices, 'outcome'] = np.nan
+
+        trop_est = TROP(
+            method="joint",
+            lambda_time_grid=[0.0, 1.0],
+            lambda_unit_grid=[0.0, 1.0],
+            lambda_nn_grid=[0.0, 0.1],
+            n_bootstrap=10,
+            seed=42,
+        )
+        results = trop_est.fit(
+            data,
+            outcome="outcome",
+            treatment="treated",
+            unit="unit",
+            time="period",
+        )
+
+        # Results should be finite (NaN observations excluded)
+        assert np.isfinite(results.att), "ATT should be finite with NaN data"
+        assert np.isfinite(results.se), "SE should be finite with NaN data"
+        # ATT should be positive (true effect is 3.0)
+        assert results.att > 0, "ATT should be positive"
+
+    def test_joint_with_lowrank_handles_nan(self, simple_panel_data):
+        """Joint method with low-rank handles NaN values correctly."""
+        # Introduce NaN in some control observations
+        data = simple_panel_data.copy()
+        control_mask = data['treated'] == 0
+        control_indices = data[control_mask].index.tolist()
+
+        # Set 3 random control observations to NaN
+        np.random.seed(123)
+        nan_indices = np.random.choice(control_indices, size=3, replace=False)
+        data.loc[nan_indices, 'outcome'] = np.nan
+
+        trop_est = TROP(
+            method="joint",
+            lambda_time_grid=[0.0],
+            lambda_unit_grid=[0.0],
+            lambda_nn_grid=[0.1],  # Finite lambda_nn enables low-rank
+            n_bootstrap=10,
+            seed=42,
+        )
+        results = trop_est.fit(
+            data,
+            outcome="outcome",
+            treatment="treated",
+            unit="unit",
+            time="period",
+        )
+
+        # Results should be finite
+        assert np.isfinite(results.att), "ATT should be finite with NaN data"
+        assert np.isfinite(results.se), "SE should be finite with NaN data"