diff --git a/autolens/point/solver/implicit_diff.py b/autolens/point/solver/implicit_diff.py index 263f01694..048a236de 100644 --- a/autolens/point/solver/implicit_diff.py +++ b/autolens/point/solver/implicit_diff.py @@ -34,7 +34,13 @@ Padded rows (the ``inf`` sentinels of the fixed ``MAX_CONTAINING_SIZE`` output) are constants of the output shape; their tangent is forced to zero so they cannot inject -NaNs into the batch. +NaNs into the batch. Because reverse mode transposes the rule into row-summed +cotangents, masking the *output* alone is not enough: the padded rows' solve inputs +are also sanitized (Jacobian evaluated at a real solved image, ``a_mat`` replaced by +the identity, ``rhs`` zeroed) so no NaN ever enters the linear algebra. Evaluating +the Jacobian at the padded rows' former ``(0, 0)`` placeholder was itself a NaN +source — profile centres typically sit at the origin, where deflection Jacobians +are singular (cluster host halos; PyAutoLens#678 phase B). Known limitation — free cosmology parameters: ``Tracer`` is registered with ``cosmology`` as ``no_flatten`` aux, so a cosmology carrying traced parameters (a free @@ -83,7 +89,15 @@ def implicit_tangents_from(jac_alpha, dalpha, dbeta, finite, xp): """ identity = xp.eye(2) a_mat = identity[None, :, :] - jac_alpha + # Sanitize padded rows before the solve: a non-finite padded ``a_mat`` row + # survives the output masking in forward mode but not in reverse mode, + # where the transpose solves against ``a_mat`` row-by-row and sums the + # cotangents — one NaN padded row contaminates every parameter's gradient. + # Real (finite) rows are untouched, so legitimate near-critical divergence + # is still surfaced. + a_mat = xp.where(finite[:, None, None], a_mat, identity[None, :, :]) rhs = dalpha + dbeta[None, :] + rhs = xp.where(finite[:, None], rhs, 0.0) dtheta = xp.linalg.solve(a_mat, rhs[..., None])[..., 0] return xp.where(finite[:, None], dtheta, 0.0) @@ -171,7 +185,17 @@ def solve_padded_jvp(primals, tangents): theta = solve_padded(tracer, beta) finite = xp.isfinite(theta).all(axis=1) - theta_safe = xp.where(finite[:, None], theta, 0.0) + # Padded rows are anchored at the first real solved image rather than + # (0, 0): profile centres typically sit at the origin, where deflection + # Jacobians are singular (an NFW's jacfwd at its own centre is NaN), + # and although the padded rows' tangents are masked below, reverse mode + # transposes the rule into row-summed cotangents — one NaN row poisons + # the gradient of every parameter. At a real image the Jacobian is + # finite and the padded rows' contributions are exactly zeroed. With + # zero solved images the anchor row is itself non-finite and the + # gradient is NaN — the likelihood is already invalid there. + anchor = theta[xp.argmax(finite)] + theta_safe = xp.where(finite[:, None], theta, anchor[None, :]) def deflections_single(position, tracer_): return deflections_from(position[None, :], tracer_)[0] diff --git a/test_autolens/point/triangles/test_implicit_diff.py b/test_autolens/point/triangles/test_implicit_diff.py index c29718e74..f25d85ebb 100644 --- a/test_autolens/point/triangles/test_implicit_diff.py +++ b/test_autolens/point/triangles/test_implicit_diff.py @@ -37,6 +37,34 @@ def test_implicit_tangents_solve_the_linear_system(): np.testing.assert_array_equal(dtheta[i], 0.0) +def test_implicit_tangents_padded_nan_rows_never_reach_the_solve(): + """ + Padded rows carry whatever the Jacobian evaluated at their placeholder position + produced — NaN when that position sits on a profile centre. The rule must + sanitize those rows before the solve (identity ``a_mat``, zero ``rhs``): in + reverse mode the transpose solves against ``a_mat`` row-by-row and sums the + cotangents, so a NaN padded row would contaminate every parameter's gradient + even though the forward output masks it (#678 phase B, cluster cells). + """ + jac_alpha = np.array( + [ + [[0.5, 0.0], [0.0, 0.5]], + [[np.nan, np.nan], [np.nan, np.nan]], + ] + ) + dalpha = np.array([[1.0, 2.0], [np.nan, np.nan]]) + dbeta = np.array([0.1, -0.2]) + finite = np.array([True, False]) + + dtheta = implicit_diff.implicit_tangents_from( + jac_alpha=jac_alpha, dalpha=dalpha, dbeta=dbeta, finite=finite, xp=np + ) + + assert np.isfinite(dtheta).all() + np.testing.assert_allclose(dtheta[0], (dalpha[0] + dbeta) / 0.5, rtol=1e-12) + np.testing.assert_array_equal(dtheta[1], 0.0) + + def test_implicit_tangents_near_critical_diverge_unclamped(): # det(I - J) -> 0: the tangent must diverge with the true solve, never be clamped. eps = 1e-12