diff --git a/src/cuphoton/xfit/solver.py b/src/cuphoton/xfit/solver.py index c73c4e81..ca6574b7 100644 --- a/src/cuphoton/xfit/solver.py +++ b/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) diff --git a/tests/xfit/test_solver.py b/tests/xfit/test_solver.py index 079bf479..7f825ccb 100644 --- a/tests/xfit/test_solver.py +++ b/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 ):