Reuse compacted row indices in the xFit solver

src/cuphoton/xfit/solver.py

@@ -523,13 +523,16 @@ def finite_rows(value: BackendArray) -> BackendArray:
                     valid_normal_equations[recovered_positions] = True
                     jacobians[recovered_indices] = fallback_jacobian
                     jacobian_current[recovered_indices] = True
-            active_indices = active_indices[valid_normal_equations]
+            # Reuse row indices so each array does not repeat Boolean-mask
+            # compaction and its device-to-host count synchronization.
+            valid_rows = ap.flatnonzero(valid_normal_equations)
+            active_indices = active_indices[valid_rows]
             if active_indices.shape[0] == 0:
                 continue
-            x_active = x_active[valid_normal_equations]
-            residual_active = residual_active[valid_normal_equations]
-            gradient = gradient[valid_normal_equations]
-            hessian = hessian[valid_normal_equations]
+            x_active = x_active[valid_rows]
+            residual_active = residual_active[valid_rows]
+            gradient = gradient[valid_rows]
+            hessian = hessian[valid_rows]
             diagnostics_current[active_indices] = True
         else:
             if (
@@ -558,12 +561,13 @@ def finite_rows(value: BackendArray) -> BackendArray:
             status[active_indices[~valid_jacobian]] = int(
                 LMStatus.INVALID_RESIDUAL
             )
-            active_indices = active_indices[valid_jacobian]
+            valid_rows = ap.flatnonzero(valid_jacobian)
+            active_indices = active_indices[valid_rows]
             if active_indices.shape[0] == 0:
                 continue
-            x_active = x_active[valid_jacobian]
-            residual_active = residual_active[valid_jacobian]
-            jacobian = jacobian[valid_jacobian]
+            x_active = x_active[valid_rows]
+            residual_active = residual_active[valid_rows]
+            jacobian = jacobian[valid_rows]
             jacobians[active_indices] = jacobian
             jacobian_current[active_indices] = True
             gradient = ap.einsum("knm,km->kn", jacobian, residual_active)
@@ -577,14 +581,15 @@ def finite_rows(value: BackendArray) -> BackendArray:
         status[active_indices[gradient_converged]] = int(
             LMStatus.CONVERGED_G_TOL
         )
-        active_indices = active_indices[~gradient_converged]
+        unconverged_rows = ap.flatnonzero(~gradient_converged)
+        active_indices = active_indices[unconverged_rows]
         if active_indices.shape[0] == 0:
             continue
-        x_active = x_active[~gradient_converged]
-        residual_active = residual_active[~gradient_converged]
-        gradient = gradient[~gradient_converged]
-        hessian = hessian[~gradient_converged]
-        diagonal = diagonal[~gradient_converged]
+        x_active = x_active[unconverged_rows]
+        residual_active = residual_active[unconverged_rows]
+        gradient = gradient[unconverged_rows]
+        hessian = hessian[unconverged_rows]
+        diagonal = diagonal[unconverged_rows]
 
         system = hessian + (
             damping[active_indices, None, None]
@@ -598,21 +603,23 @@ def finite_rows(value: BackendArray) -> BackendArray:
             continue
         valid_step = finite_rows(step)
         status[active_indices[~valid_step]] = int(LMStatus.SINGULAR)
-        active_indices = active_indices[valid_step]
+        valid_rows = ap.flatnonzero(valid_step)
+        active_indices = active_indices[valid_rows]
         if active_indices.shape[0] == 0:
             continue
-        x_active = x_active[valid_step]
-        residual_active = residual_active[valid_step]
-        step = step[valid_step]
+        x_active = x_active[valid_rows]
+        residual_active = residual_active[valid_rows]
+        step = step[valid_rows]
 
         can_evaluate = evaluations[active_indices] < settings.max_evaluations
         status[active_indices[~can_evaluate]] = int(LMStatus.MAX_EVALUATIONS)
-        active_indices = active_indices[can_evaluate]
+        evaluable_rows = ap.flatnonzero(can_evaluate)
+        active_indices = active_indices[evaluable_rows]
         if active_indices.shape[0] == 0:
             continue
-        x_active = x_active[can_evaluate]
-        residual_active = residual_active[can_evaluate]
-        step = step[can_evaluate]
+        x_active = x_active[evaluable_rows]
+        residual_active = residual_active[evaluable_rows]
+        step = step[evaluable_rows]
 
         trial_x = x_active + step
         trial_residual = _call_residual(problem, trial_x, active_indices)
@@ -625,11 +632,12 @@ def finite_rows(value: BackendArray) -> BackendArray:
         reduction = cost - trial_cost
         accepted = trial_valid & (reduction > 0)
 
-        accepted_indices = active_indices[accepted]
+        accepted_rows = ap.flatnonzero(accepted)
+        accepted_indices = active_indices[accepted_rows]
         if accepted_indices.shape[0] > 0:
-            accepted_step = step[accepted]
-            x[accepted_indices] = trial_x[accepted]
-            residuals[accepted_indices] = trial_residual[accepted]
+            accepted_step = step[accepted_rows]
+            x[accepted_indices] = trial_x[accepted_rows]
+            residuals[accepted_indices] = trial_residual[accepted_rows]
             jacobians[accepted_indices] = ap.nan
             jacobian_current[accepted_indices] = False
             diagnostics_current[accepted_indices] = False
@@ -638,8 +646,8 @@ def finite_rows(value: BackendArray) -> BackendArray:
                 damping[accepted_indices] * settings.damping_decrease,
             )
             rejected_steps[accepted_indices] = 0
-            f_converged = reduction[accepted] <= (
-                settings.f_tol * ap.maximum(1.0, cost[accepted])
+            f_converged = reduction[accepted_rows] <= (
+                settings.f_tol * ap.maximum(1.0, cost[accepted_rows])
             )
             x_converged = ap.linalg.norm(accepted_step, axis=1) <= (
                 settings.x_tol * (settings.x_tol + 1.0)

tests/xfit/test_solver.py

@@ -4,6 +4,8 @@
 
 from __future__ import annotations
 
+from dataclasses import fields
+
 import numpy as np
 import pytest
 from scipy.optimize import least_squares
@@ -103,6 +105,129 @@ def jacobian(x, *, indices, out=None):
     assert result.evaluations.tolist() == [1, 3, 3]
 
 
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+@pytest.mark.parametrize("use_normal_equations", [False, True])
+def test_mixed_batch_matches_independent_rows(
+    dtype, use_normal_equations: bool
+) -> None:
+    # Interleave successful fits, initially invalid rows, an already solved
+    # row, an invalid derivative, a wrong-signed derivative, and a residual
+    # that becomes invalid at every trial. Each callback uses original row
+    # identities even after the active batch shrinks.
+    target = np.asarray([3.0, 1.0, 4.0, 2.0, -1.0, -2.0, 1.0], dtype=dtype)
+    initial = np.asarray([0.0, 0.0, 4.0, 0.0, 0.0, 5.0, 0.0], dtype=dtype)
+    slopes = np.arange(1, 8, dtype=dtype)
+
+    def solve(rows):
+        evaluated = np.zeros(rows.size, dtype=np.int32)
+
+        def residual(x, *, indices):
+            assert np.all(np.diff(indices) > 0)
+            original = rows[indices]
+            evaluated[indices] += 1
+            value = (x - target[original, None]) * slopes[original, None]
+            value[original == 1] = np.nan
+            value[(original == 6) & (x[:, 0] != 0.0)] = np.nan
+            return value
+
+        def jacobian(x, *, indices):
+            del x
+            assert np.all(np.diff(indices) > 0)
+            original = rows[indices]
+            value = slopes[original, None, None].copy()
+            value[original == 3] = np.nan
+            value[original == 4] *= -1.0
+            return value
+
+        def normal_equations(x, residuals, *, indices):
+            jac = jacobian(x, indices=indices)
+            return (
+                np.einsum("knm,km->kn", jac, residuals),
+                np.einsum("knm,kpm->knp", jac, jac),
+            )
+
+        result = batched_levenberg_marquardt(
+            BatchedLeastSquaresProblem(
+                residual,
+                jacobian,
+                normal_equations=(
+                    normal_equations if use_normal_equations else None
+                ),
+            ),
+            initial[rows, None],
+            config=LMConfig(max_evaluations=6),
+        )
+        np.testing.assert_array_equal(result.evaluations, evaluated)
+        return result
+
+    batched = solve(np.arange(initial.size))
+    independent = [solve(np.asarray([row])) for row in range(initial.size)]
+    for field in fields(batched):
+        expected = np.concatenate(
+            [getattr(result, field.name) for result in independent]
+        )
+        np.testing.assert_array_equal(getattr(batched, field.name), expected)
+
+    assert batched.converged.tolist() == [
+        True,
+        False,
+        True,
+        False,
+        False,
+        True,
+        False,
+    ]
+    assert batched.status[[1, 2, 3, 4, 6]].tolist() == [
+        LMStatus.INVALID_RESIDUAL,
+        LMStatus.CONVERGED_G_TOL,
+        LMStatus.INVALID_RESIDUAL,
+        LMStatus.MAX_EVALUATIONS,
+        LMStatus.MAX_EVALUATIONS,
+    ]
+    assert batched.evaluations[[1, 2, 3, 4, 6]].tolist() == [1, 1, 1, 6, 6]
+
+
+def test_nonfinite_step_does_not_displace_other_active_rows(
+    monkeypatch,
+) -> None:
+    original_solve = np.linalg.solve
+
+    def solve_with_one_failed_row(system, rhs):
+        step = original_solve(system, rhs)
+        step[rhs[:, 0, 0] == -2.0] = np.nan
+        return step
+
+    monkeypatch.setattr(np.linalg, "solve", solve_with_one_failed_row)
+
+    def residual(x, *, indices):
+        del indices
+        return x.copy()
+
+    def jacobian(x, *, indices):
+        del indices
+        return np.ones((x.shape[0], 1, 1))
+
+    result = batched_levenberg_marquardt(
+        BatchedLeastSquaresProblem(residual, jacobian),
+        np.asarray([[1.0], [2.0], [3.0]]),
+        config=LMConfig(f_tol=1.0, max_evaluations=2),
+    )
+
+    assert result.status.tolist() == [
+        LMStatus.CONVERGED_F_TOL,
+        LMStatus.SINGULAR,
+        LMStatus.CONVERGED_F_TOL,
+    ]
+    assert result.evaluations.tolist() == [2, 1, 2]
+    np.testing.assert_array_equal(
+        result.parameters[:, 0],
+        np.asarray([1.0, 2.0, 3.0]) - np.asarray([1.0, 0.0, 3.0]) / 1.001,
+    )
+    assert result.rank.tolist() == [1, 1, 1]
+    np.testing.assert_array_equal(result.residuals, result.parameters)
+    np.testing.assert_array_equal(result.gradient, result.parameters)
+
+
 def test_finite_difference_jacobian_converges_and_counts_evaluations() -> (
     None
 ):