Clean up unused parameters from TROP Rust API

Address code review feedback to remove unused API parameters:

Rust backend (trop.rs):
- Remove control_unit_idx and unit_dist_matrix from loocv_grid_search
- Remove control_unit_idx, treated_obs_t/i, unit_dist_matrix from
  bootstrap_trop_variance
- Remove unit_dist_boot computation in bootstrap (no longer needed)
- Remove control_units and unit_dist from internal functions

Python (trop.py):
- Update _rust_loocv_grid_search call to use new signature
- Update _rust_bootstrap_trop_variance call to use new signature
- Remove unused variable preparation for removed parameters

Tests (test_rust_backend.py):
- Update test calls to use new API signatures
- Remove unused variable assignments

The precomputed unit_dist_matrix is no longer needed by the Rust backend
since per-observation distances are computed dynamically to properly
exclude the target period per Equation 3 of the paper.

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

Author: igerber

diff_diff/trop.py

@@ -784,16 +784,14 @@ def fit(
                 # Prepare inputs for Rust function
                 control_mask_u8 = control_mask.astype(np.uint8)
                 time_dist_matrix = self._precomputed["time_dist_matrix"].astype(np.int64)
-                unit_dist_matrix = self._precomputed["unit_dist_matrix"]
-                control_unit_idx_i64 = control_unit_idx.astype(np.int64)
 
                 lambda_time_arr = np.array(self.lambda_time_grid, dtype=np.float64)
                 lambda_unit_arr = np.array(self.lambda_unit_grid, dtype=np.float64)
                 lambda_nn_arr = np.array(self.lambda_nn_grid, dtype=np.float64)
 
                 best_lt, best_lu, best_ln, best_score = _rust_loocv_grid_search(
-                    Y, D.astype(np.float64), control_mask_u8, control_unit_idx_i64,
-                    unit_dist_matrix, time_dist_matrix,
+                    Y, D.astype(np.float64), control_mask_u8,
+                    time_dist_matrix,
                     lambda_time_arr, lambda_unit_arr, lambda_nn_arr,
                     self.max_loocv_samples, self.max_iter, self.tol,
                     self.seed if self.seed is not None else 0
@@ -1510,21 +1508,15 @@ def _bootstrap_variance(
         # Try Rust backend for parallel bootstrap (5-15x speedup)
         if (HAS_RUST_BACKEND and _rust_bootstrap_trop_variance is not None
                 and self._precomputed is not None and Y is not None
-                and D is not None and control_unit_idx is not None):
+                and D is not None):
             try:
-                # Prepare inputs
-                treated_observations = self._precomputed["treated_observations"]
-                treated_t = np.array([t for t, i in treated_observations], dtype=np.int64)
-                treated_i = np.array([i for t, i in treated_observations], dtype=np.int64)
                 control_mask = self._precomputed["control_mask"]
+                time_dist_matrix = self._precomputed["time_dist_matrix"].astype(np.int64)
 
                 bootstrap_estimates, se = _rust_bootstrap_trop_variance(
                     Y, D.astype(np.float64),
                     control_mask.astype(np.uint8),
-                    control_unit_idx.astype(np.int64),
-                    treated_t, treated_i,
-                    self._precomputed["unit_dist_matrix"],
-                    self._precomputed["time_dist_matrix"].astype(np.int64),
+                    time_dist_matrix,
                     lambda_time, lambda_unit, lambda_nn,
                     self.n_bootstrap, self.max_iter, self.tol,
                     self.seed if self.seed is not None else 0

rust/src/trop.rs

@@ -172,15 +172,13 @@ fn compute_pair_distance(
 /// # Returns
 /// (best_lambda_time, best_lambda_unit, best_lambda_nn, best_score)
 #[pyfunction]
-#[pyo3(signature = (y, d, control_mask, control_unit_idx, unit_dist_matrix, time_dist_matrix, lambda_time_grid, lambda_unit_grid, lambda_nn_grid, max_loocv_samples, max_iter, tol, seed))]
+#[pyo3(signature = (y, d, control_mask, time_dist_matrix, lambda_time_grid, lambda_unit_grid, lambda_nn_grid, max_loocv_samples, max_iter, tol, seed))]
 #[allow(clippy::too_many_arguments)]
 pub fn loocv_grid_search<'py>(
     _py: Python<'py>,
     y: PyReadonlyArray2<'py, f64>,
     d: PyReadonlyArray2<'py, f64>,
     control_mask: PyReadonlyArray2<'py, u8>,
-    control_unit_idx: PyReadonlyArray1<'py, i64>,
-    unit_dist_matrix: PyReadonlyArray2<'py, f64>,
     time_dist_matrix: PyReadonlyArray2<'py, i64>,
     lambda_time_grid: PyReadonlyArray1<'py, f64>,
     lambda_unit_grid: PyReadonlyArray1<'py, f64>,
@@ -193,19 +191,11 @@ pub fn loocv_grid_search<'py>(
     let y_arr = y.as_array();
     let d_arr = d.as_array();
     let control_mask_arr = control_mask.as_array();
-    let control_unit_idx_arr = control_unit_idx.as_array();
-    let unit_dist_arr = unit_dist_matrix.as_array();
     let time_dist_arr = time_dist_matrix.as_array();
     let lambda_time_vec: Vec<f64> = lambda_time_grid.as_array().to_vec();
     let lambda_unit_vec: Vec<f64> = lambda_unit_grid.as_array().to_vec();
     let lambda_nn_vec: Vec<f64> = lambda_nn_grid.as_array().to_vec();
 
-    // Convert control_unit_idx to Vec<usize>
-    let control_units: Vec<usize> = control_unit_idx_arr
-        .iter()
-        .map(|&idx| idx as usize)
-        .collect();
-
     // Get control observations for LOOCV
     let control_obs = get_control_observations(
         &y_arr,
@@ -232,8 +222,6 @@ pub fn loocv_grid_search<'py>(
                 &y_arr,
                 &d_arr,
                 &control_mask_arr,
-                &control_units,
-                &unit_dist_arr,
                 &time_dist_arr,
                 &control_obs,
                 lambda_time,
@@ -293,8 +281,6 @@ fn loocv_score_for_params(
     y: &ArrayView2<f64>,
     d: &ArrayView2<f64>,
     control_mask: &ArrayView2<u8>,
-    control_units: &[usize],
-    unit_dist: &ArrayView2<f64>,
     time_dist: &ArrayView2<i64>,
     control_obs: &[(usize, usize)],
     lambda_time: f64,
@@ -311,7 +297,6 @@ fn loocv_score_for_params(
 
     for &(t, i) in control_obs {
         // Compute observation-specific weight matrix
-        // Issue A+B fix: pass y and d for dynamic control sets and per-obs distances
         let weight_matrix = compute_weight_matrix(
             y,
             d,
@@ -321,8 +306,6 @@ fn loocv_score_for_params(
             t,
             lambda_time,
             lambda_unit,
-            control_units,
-            unit_dist,
             time_dist,
         );
 
@@ -410,8 +393,6 @@ fn compute_weight_matrix(
     target_period: usize,
     lambda_time: f64,
     lambda_unit: f64,
-    _control_units: &[usize],  // Kept for API compatibility but not used
-    _unit_dist: &ArrayView2<f64>,  // Not used - we compute per-observation distances
     time_dist: &ArrayView2<i64>,
 ) -> Array2<f64> {
     // Time weights for this target period: θ_s = exp(-λ_time × |t - s|)
@@ -707,17 +688,13 @@ fn max_abs_diff_2d(a: &Array2<f64>, b: &Array2<f64>) -> f64 {
 /// # Returns
 /// (bootstrap_estimates, standard_error)
 #[pyfunction]
-#[pyo3(signature = (y, d, control_mask, control_unit_idx, treated_obs_t, treated_obs_i, unit_dist_matrix, time_dist_matrix, lambda_time, lambda_unit, lambda_nn, n_bootstrap, max_iter, tol, seed))]
+#[pyo3(signature = (y, d, control_mask, time_dist_matrix, lambda_time, lambda_unit, lambda_nn, n_bootstrap, max_iter, tol, seed))]
 #[allow(clippy::too_many_arguments)]
 pub fn bootstrap_trop_variance<'py>(
     py: Python<'py>,
     y: PyReadonlyArray2<'py, f64>,
     d: PyReadonlyArray2<'py, f64>,
     control_mask: PyReadonlyArray2<'py, u8>,
-    control_unit_idx: PyReadonlyArray1<'py, i64>,
-    treated_obs_t: PyReadonlyArray1<'py, i64>,
-    treated_obs_i: PyReadonlyArray1<'py, i64>,
-    unit_dist_matrix: PyReadonlyArray2<'py, f64>,
     time_dist_matrix: PyReadonlyArray2<'py, i64>,
     lambda_time: f64,
     lambda_unit: f64,
@@ -730,17 +707,11 @@ pub fn bootstrap_trop_variance<'py>(
     let y_arr = y.as_array().to_owned();
     let d_arr = d.as_array().to_owned();
     let control_mask_arr = control_mask.as_array().to_owned();
-    let unit_dist_arr = unit_dist_matrix.as_array().to_owned();
     let time_dist_arr = time_dist_matrix.as_array().to_owned();
 
     let n_units = y_arr.ncols();
     let n_periods = y_arr.nrows();
 
-    // Note: control_unit_idx, treated_obs_t, treated_obs_i are passed for API
-    // compatibility but not used directly - each bootstrap iteration recomputes
-    // control units and treated observations from the resampled data.
-    let _ = (control_unit_idx, treated_obs_t, treated_obs_i);
-
     // Issue D fix: Identify treated and control units for stratified sampling
     // Following paper's Algorithm 3 (page 27): sample N_0 control and N_1 treated separately
     let mut original_treated_units: Vec<usize> = Vec::new();
@@ -784,18 +755,13 @@ pub fn bootstrap_trop_variance<'py>(
             let mut y_boot = Array2::<f64>::zeros((n_periods, n_units));
             let mut d_boot = Array2::<f64>::zeros((n_periods, n_units));
             let mut control_mask_boot = Array2::<u8>::zeros((n_periods, n_units));
-            let mut unit_dist_boot = Array2::<f64>::zeros((n_units, n_units));
 
             for (new_idx, &old_idx) in sampled_units.iter().enumerate() {
                 for t in 0..n_periods {
                     y_boot[[t, new_idx]] = y_arr[[t, old_idx]];
                     d_boot[[t, new_idx]] = d_arr[[t, old_idx]];
                     control_mask_boot[[t, new_idx]] = control_mask_arr[[t, old_idx]];
                 }
-
-                for (new_j, &old_j) in sampled_units.iter().enumerate() {
-                    unit_dist_boot[[new_idx, new_j]] = unit_dist_arr[[old_idx, old_j]];
-                }
             }
 
             // Get treated observations in bootstrap sample
@@ -829,7 +795,6 @@ pub fn bootstrap_trop_variance<'py>(
             let mut tau_values = Vec::with_capacity(boot_treated.len());
 
             for (t, i) in boot_treated {
-                // Issue A+B fix: pass y and d for dynamic control sets and per-obs distances
                 let weight_matrix = compute_weight_matrix(
                     &y_boot.view(),
                     &d_boot.view(),
@@ -839,8 +804,6 @@ pub fn bootstrap_trop_variance<'py>(
                     t,
                     lambda_time,
                     lambda_unit,
-                    &boot_control_units,
-                    &unit_dist_boot.view(),
                     &time_dist_arr.view(),
                 );
 

tests/test_rust_backend.py

@@ -657,11 +657,8 @@ def test_loocv_grid_search_returns_valid_params(self):
         D[6:, 0] = 1.0
 
         control_mask = (D == 0).astype(np.uint8)
-        control_unit_idx = np.array([1, 2, 3, 4, 5], dtype=np.int64)
 
-        # Compute distance matrices
-        from diff_diff._rust_backend import compute_unit_distance_matrix
-        unit_dist = compute_unit_distance_matrix(Y, D)
+        # Compute time distance matrix
         time_dist = np.abs(
             np.arange(n_periods)[:, np.newaxis] - np.arange(n_periods)[np.newaxis, :]
         ).astype(np.int64)
@@ -671,8 +668,7 @@ def test_loocv_grid_search_returns_valid_params(self):
         lambda_nn = np.array([0.0, 0.1], dtype=np.float64)
 
         best_lt, best_lu, best_ln, score = loocv_grid_search(
-            Y, D, control_mask, control_unit_idx,
-            unit_dist, time_dist,
+            Y, D, control_mask, time_dist,
             lambda_time, lambda_unit, lambda_nn,
             50, 100, 1e-6, 42
         )
@@ -685,7 +681,7 @@ def test_loocv_grid_search_returns_valid_params(self):
 
     def test_bootstrap_variance_shape(self):
         """Test bootstrap returns correct shapes."""
-        from diff_diff._rust_backend import bootstrap_trop_variance, compute_unit_distance_matrix
+        from diff_diff._rust_backend import bootstrap_trop_variance
 
         np.random.seed(42)
         n_periods, n_units = 8, 6
@@ -694,20 +690,15 @@ def test_bootstrap_variance_shape(self):
         D[6:, 0] = 1.0  # Treat unit 0 in last 2 periods
 
         control_mask = (D == 0).astype(np.uint8)
-        control_unit_idx = np.array([1, 2, 3, 4, 5], dtype=np.int64)
-        treated_t = np.array([6, 7], dtype=np.int64)
-        treated_i = np.array([0, 0], dtype=np.int64)
 
-        unit_dist = compute_unit_distance_matrix(Y, D)
+        # Compute time distance matrix
         time_dist = np.abs(
             np.arange(n_periods)[:, np.newaxis] - np.arange(n_periods)[np.newaxis, :]
         ).astype(np.int64)
 
         n_bootstrap = 20
         estimates, se = bootstrap_trop_variance(
-            Y, D, control_mask, control_unit_idx,
-            treated_t, treated_i,
-            unit_dist, time_dist,
+            Y, D, control_mask, time_dist,
             1.0, 1.0, 0.1,  # lambda values
             n_bootstrap, 100, 1e-6, 42
         )
@@ -718,7 +709,7 @@ def test_bootstrap_variance_shape(self):
 
     def test_bootstrap_reproducibility(self):
         """Test bootstrap is reproducible with same seed."""
-        from diff_diff._rust_backend import bootstrap_trop_variance, compute_unit_distance_matrix
+        from diff_diff._rust_backend import bootstrap_trop_variance
 
         np.random.seed(42)
         n_periods, n_units = 8, 6
@@ -727,26 +718,19 @@ def test_bootstrap_reproducibility(self):
         D[6:, 0] = 1.0
 
         control_mask = (D == 0).astype(np.uint8)
-        control_unit_idx = np.array([1, 2, 3, 4, 5], dtype=np.int64)
-        treated_t = np.array([6, 7], dtype=np.int64)
-        treated_i = np.array([0, 0], dtype=np.int64)
 
-        unit_dist = compute_unit_distance_matrix(Y, D)
+        # Compute time distance matrix
         time_dist = np.abs(
             np.arange(n_periods)[:, np.newaxis] - np.arange(n_periods)[np.newaxis, :]
         ).astype(np.int64)
 
         # Run twice with same seed
         est1, se1 = bootstrap_trop_variance(
-            Y, D, control_mask, control_unit_idx,
-            treated_t, treated_i,
-            unit_dist, time_dist,
+            Y, D, control_mask, time_dist,
             1.0, 1.0, 0.1, 20, 100, 1e-6, 42
         )
         est2, se2 = bootstrap_trop_variance(
-            Y, D, control_mask, control_unit_idx,
-            treated_t, treated_i,
-            unit_dist, time_dist,
+            Y, D, control_mask, time_dist,
             1.0, 1.0, 0.1, 20, 100, 1e-6, 42
         )