From 4ccd920221cf7157438c4fca7607850254817aa8 Mon Sep 17 00:00:00 2001
From: Trent Nelson <trentn@nvidia.com>
Date: Thu, 1 Oct 2026 01:37:33 -0700
Subject: [PATCH] Avoid redundant compaction in xFit LM iterations

Filter normal-equation rows only when a fallback was needed.
Skip the second residual-budget filter for analytic derivatives;
keep it for finite differences, with a boundary regression test.

Signed-off-by: Trent Nelson <trentn@nvidia.com>
---
 src/cuphoton/xfit/solver.py | 36 +++++++++++++++++++--------------
 tests/xfit/test_solver.py   | 40 +++++++++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+), 15 deletions(-)

diff --git a/src/cuphoton/xfit/solver.py b/src/cuphoton/xfit/solver.py
index c73c4e8..a3644f8 100644
--- a/src/cuphoton/xfit/solver.py
+++ b/src/cuphoton/xfit/solver.py
@@ -523,13 +523,13 @@ 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]
-            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]
+                active_indices = active_indices[valid_normal_equations]
+                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]
             diagnostics_current[active_indices] = True
         else:
             if (
@@ -605,14 +605,20 @@ def finite_rows(value: BackendArray) -> BackendArray:
         residual_active = residual_active[valid_step]
         step = step[valid_step]
 
-        can_evaluate = evaluations[active_indices] < settings.max_evaluations
-        status[active_indices[~can_evaluate]] = int(LMStatus.MAX_EVALUATIONS)
-        active_indices = active_indices[can_evaluate]
-        if active_indices.shape[0] == 0:
-            continue
-        x_active = x_active[can_evaluate]
-        residual_active = residual_active[can_evaluate]
-        step = step[can_evaluate]
+        # Analytic derivatives leave the iteration-entry budget unchanged.
+        if problem.jacobian is None or settings.use_finite_difference:
+            can_evaluate = (
+                evaluations[active_indices] < settings.max_evaluations
+            )
+            status[active_indices[~can_evaluate]] = int(
+                LMStatus.MAX_EVALUATIONS
+            )
+            active_indices = active_indices[can_evaluate]
+            if active_indices.shape[0] == 0:
+                continue
+            x_active = x_active[can_evaluate]
+            residual_active = residual_active[can_evaluate]
+            step = step[can_evaluate]
 
         trial_x = x_active + step
         trial_residual = _call_residual(problem, trial_x, active_indices)
diff --git a/tests/xfit/test_solver.py b/tests/xfit/test_solver.py
index 079bf47..bde828e 100644
--- a/tests/xfit/test_solver.py
+++ b/tests/xfit/test_solver.py
@@ -130,6 +130,46 @@ def residual(x, *, indices, out=None):
     assert residual_rows_evaluated == int(result.evaluations.sum())
 
 
+@pytest.mark.parametrize(
+    "with_analytic_jacobian",
+    [False, True],
+    ids=["residual-only", "forced-finite-difference"],
+)
+def test_finite_difference_budget_can_end_before_the_trial(
+    with_analytic_jacobian: bool,
+) -> None:
+    residual_calls = 0
+
+    def residual(x, *, indices):
+        nonlocal residual_calls
+        del indices
+        residual_calls += 1
+        return x.copy()
+
+    def jacobian(x, *, indices):
+        del x, indices
+        raise AssertionError("analytic Jacobian must not be called")
+
+    initial = np.ones((2, 2))
+    result = batched_levenberg_marquardt(
+        BatchedLeastSquaresProblem(
+            residual, jacobian if with_analytic_jacobian else None
+        ),
+        initial,
+        config=LMConfig(
+            max_evaluations=3,
+            use_finite_difference=with_analytic_jacobian,
+        ),
+    )
+
+    # The initial residual and two perturbations leave no trial evaluation.
+    assert residual_calls == 3
+    assert result.evaluations.tolist() == [3, 3]
+    assert result.status.tolist() == [LMStatus.MAX_EVALUATIONS] * 2
+    np.testing.assert_array_equal(result.parameters, initial)
+    assert result.rank.tolist() == [2, 2]
+
+
 def test_callbacks_do_not_need_to_accept_an_out_argument() -> None:
     target = np.asarray([[1.0], [-2.0]])