fix: tie-safe rank-CDF mesh transform; NumPy stays fp64 under mixed precision (#552)

autoarray/inversion/inversion/inversion_util.py

@@ -123,20 +123,15 @@ def curvature_matrix_via_mapping_matrix_from(
     xp
         The array module to use (`numpy` by default; pass `jax.numpy` for JAX support).
     """
-    # NumPy path: keep it simple + stable
-    if xp is np:
-        A = mapping_matrix / noise_map[:, None]
-        curvature_matrix = xp.dot(A.T, A)
-    else:
-        # Choose compute dtype
-
-        compute_dtype = xp.float32 if settings.use_mixed_precision else xp.float64
-        out_dtype = xp.float64  # always return float64 for downstream stability
-
-        A = mapping_matrix
-        w = (1.0 / noise_map).astype(compute_dtype)
-        A = A * w[:, None]
-        curvature_matrix = xp.dot(A.T, A).astype(out_dtype)
+    # The noise weighting is fp64 on every backend, whatever
+    # `settings.use_mixed_precision` says. The mapping matrix arrives fp64 from
+    # every operated path (the mixed-precision FFT upcasts at the kernel
+    # multiply), so the former fp32 `1 / noise_map` on the JAX branch never
+    # bought an fp32 accumulation -- it only rounded the weights of F
+    # inconsistently with the fp64 `1 / noise_map**2` of the data vector
+    # (PyAutoArray#552).
+    A = xp.asarray(mapping_matrix, dtype=xp.float64) / noise_map[:, None]
+    curvature_matrix = xp.dot(A.T, A)
 
     if (
         add_to_curvature_diag

autoarray/inversion/mappers/mapper_util.py

@@ -307,8 +307,15 @@ def mapping_matrix_from(
     w64 = xp.asarray(pix_weights_for_sub_slim_index, dtype=xp.float64)
     frac64 = xp.asarray(sub_fraction, dtype=xp.float64)
 
-    # Output dtype only (big allocation)
-    out_dtype = xp.float32 if use_mixed_precision else xp.float64
+    # Output dtype only (big allocation). fp32 is a JAX-only (GPU throughput)
+    # choice: the NumPy backend is the fp64 parity oracle of the workspace smoke
+    # scripts and always allocates float64, as the `Settings` docstring promises
+    # (PyAutoArray#552).
+    out_dtype = (
+        xp.float32
+        if (use_mixed_precision and xp.__name__.startswith("jax"))
+        else xp.float64
+    )
 
     # 1) Flatten
     flat_pixidx = pix_idx.reshape(-1)  # (M_sub*B,)

autoarray/inversion/mesh/interpolator/rectangular.py

@@ -191,6 +191,14 @@ def create_transforms_rank(traced_points, mesh_weight_map=None, xp=np):
     The CDF is the empirical (weighted) rank CDF of the traced points — a
     sort plus a cumulative sum, linearly interpolated between points — so a
     likelihood evaluation costs O(N log N) with no kernel hyperparameters.
+
+    Tied coordinates (an unlensed image grid ties every column and every row)
+    are handled in the weighted branch by giving all points of a tie block
+    the block's final cumulative weight, i.e. the right-continuous empirical
+    CDF, which makes the transform independent of the backend's ``argsort``
+    tie-break and of the input ordering (PyAutoArray#552). The unweighted
+    branch keeps its historical rank staircase (ranks are the same whichever
+    tied point is picked, so it was never backend-dependent).
     Recovered from the pre-consolidation implementation (PR #402 deleted it
     when the kernel CDF took the plain class names); see the module
     docstring for the gradient trade-off versus the kernel CDF.
@@ -222,6 +230,27 @@ def create_transforms_rank(traced_points, mesh_weight_map=None, xp=np):
         t = xp.take_along_axis(t, sdx, axis=0)
         t = xp.cumsum(t, axis=0)
 
+        # Tied coordinates: every point of a tie block takes the block's final
+        # cumulative weight -- the right-continuous empirical CDF -- so the
+        # transform is a function of the point multiset only. Without this the
+        # partial sums inside a block follow whatever order ``argsort`` gave the
+        # tied entries, and ``np.argsort`` (quicksort, unstable) and
+        # ``jnp.argsort`` (stable) give different ones: on an unlensed image grid
+        # (up to 18 points tied per coordinate) that moved sub-pixels between
+        # mesh cells and opened a 15-nat NumPy-vs-JAX log-likelihood gap
+        # (PyAutoArray#552). Untied points are unchanged bit-for-bit.
+        t = xp.stack(
+            [
+                t[
+                    xp.searchsorted(sort_points[:, d], sort_points[:, d], side="right")
+                    - 1,
+                    d,
+                ]
+                for d in range(2)
+            ],
+            axis=1,
+        )
+
     if xp.__name__.startswith("jax"):
         transform = partial(forward_interp, sort_points, t)
         inv_transform = partial(reverse_interp, t, sort_points)

autoarray/operators/convolver.py

@@ -900,7 +900,13 @@ def mapping_matrix_native_from(
             Contains contributions from both the main mapping matrix and, if provided,
             the blurring mapping matrix.
         """
-        dtype_native = xp.float32 if use_mixed_precision else xp.float64
+        # fp32 is a JAX-only choice; the NumPy backend always builds the cube in
+        # fp64 (see `Settings.use_mixed_precision`, PyAutoArray#552).
+        dtype_native = (
+            xp.float32
+            if (use_mixed_precision and xp.__name__.startswith("jax"))
+            else xp.float64
+        )
 
         n_src = mapping_matrix.shape[1]
 

autoarray/settings.py

@@ -47,11 +47,11 @@ def __init__(
               algebra stay fp64. Pixelization meshes with K ≫ 40 source pixels accumulate enough fp32
               round-off through NNLS / log-determinant to shift ``figure_of_merit`` by O(1) units; the upcast
               preserves precision while the cheaper fp32 scatter and forward FFT are kept.
-            - The mapping matrix native cube allocation in
-              :func:`autoarray.inversion.mappers.mapper_util.mapping_matrix_from` — output dtype becomes fp32.
-            - The internal compute dtype of the curvature matrix accumulation in
-              :func:`autoarray.inversion.inversion.inversion_util.curvature_matrix_via_mapping_matrix_from` —
-              the noise-weighted ``A.T @ A`` is formed in fp32 then cast to fp64 for downstream stability.
+            - The mapping matrix allocation in
+              :func:`autoarray.inversion.mappers.mapper_util.mapping_matrix_from` and the native cube in
+              :meth:`Convolver.mapping_matrix_native_from` — output dtype becomes fp32 on the JAX backend
+              only; under ``xp=np`` both stay fp64, so a NumPy fit is a true fp64 reference for a
+              mixed-precision JAX fit (PyAutoArray#552).
 
             Empirical platform notes:
 
@@ -64,6 +64,12 @@ def __init__(
 
             Paths that intentionally stay in fp64:
 
+            - The noise weighting of the curvature matrix in
+              :func:`autoarray.inversion.inversion.inversion_util.curvature_matrix_via_mapping_matrix_from`:
+              ``F`` and the data vector ``D`` are both formed with fp64 ``1 / noise_map`` on every backend,
+              so the linear system is weighted consistently (an earlier fp32 ``1 / noise_map`` on the JAX
+              branch never produced an fp32 accumulation, because the blurred mapping matrix is already
+              fp64 — it only biased ``F`` against ``D``; PyAutoArray#552).
             - The NNLS reconstruction (jaxnnls / Cholesky factor + cho_solve) in
               :func:`autoarray.inversion.inversion.inversion_util.reconstruction_positive_only_from`. Active-set
               and PDIP solvers are sensitive to fp32 noise on ill-conditioned source meshes.
@@ -73,8 +79,11 @@ def __init__(
 
             Empirical numerical impact on the MGE imaging regression (HST-shaped, 15k masked pixels, 40 linear
             Gaussians): Δlog-likelihood ≈ 1e-4 absolute at log-likelihood ≈ 27,400. Well below the natural χ²
-            sampling noise floor (σ ≈ √(2N) ≈ 175). Pixelization paths with K ≫ 40 source pixels are more
-            sensitive — verify on representative integration tests before turning on for production fits.
+            sampling noise floor (σ ≈ √(2N) ≈ 175). On a 316-pixel ``RectangularBilinearAdaptImage`` +
+            ``Adapt`` inversion (17×17 mesh) the whole mixed-precision effect on the JAX path is
+            Δlog-likelihood ≈ 1e-4 nats against the fp64 path (PyAutoArray#552); pixelization paths with
+            K ≫ 40 source pixels are more sensitive — verify on representative integration tests before
+            turning on for production fits.
 
             If `False` (default), all paths run in fp64.
         use_positive_only_solver

test_autoarray/inversion/inversion/test_inversion_util.py

@@ -1,3 +1,5 @@
+import importlib.util
+
 import autoarray as aa
 import numpy as np
 import pytest
@@ -281,3 +283,104 @@ def test__reconstruction_positive_only_from__singular_matrix_test_mode_returns_d
     assert reconstruction.shape == data_vector.shape
     assert np.all(np.isfinite(reconstruction))
     assert (reconstruction == 1.0).all()
+
+
+# ----------------------------------------------------------------------------
+# The curvature matrix's noise weighting is fp64 on every backend
+# (PyAutoArray#552).
+#
+# The JAX branch used to round `1 / noise_map` to fp32 when mixed precision was
+# on. That never bought an fp32 accumulation — the blurred mapping matrix
+# arrives fp64 from every operated path — it only weighted `F` with fp32
+# reciprocals while the data vector `D` kept fp64 `1 / noise_map ** 2`, so the
+# linear system was inconsistently weighted.
+# ----------------------------------------------------------------------------
+
+# jax is an `[optional]` extra and is absent on the NumPy-only matrix env, so
+# the JAX parity tests skip rather than fail there.
+requires_jax = pytest.mark.skipif(
+    importlib.util.find_spec("jax") is None,
+    reason="requires jax (installed via the [optional] extras; absent on the NumPy-only matrix env)",
+)
+
+
+def _curvature_mixed_precision_inputs():
+    """
+    A mapping matrix with fp32-visible entries over a noise-map whose
+    reciprocal is *not* representable in fp32 (1 / 3 and 1 / 7), so an fp32
+    rounding of the weights shows up in `F` at the 1e-8 relative level.
+    """
+    mapping_matrix = np.array(
+        [
+            [1.0, 1.0, 0.0],
+            [1.0, 0.0, 0.0],
+            [0.0, 1.0, 0.0],
+            [0.0, 1.0, 1.0],
+            [0.3, 0.0, 0.7],
+            [0.0, 0.2, 0.8],
+        ]
+    )
+
+    noise_map = np.array([3.0, 3.0, 7.0, 3.0, 7.0, 3.0])
+
+    return mapping_matrix, noise_map
+
+
+@requires_jax
+def test__curvature_matrix_via_mapping_matrix_from__mixed_precision__jax_matches_numpy():
+    import jax
+
+    jax.config.update("jax_enable_x64", True)
+    import jax.numpy as jnp
+
+    mapping_matrix, noise_map = _curvature_mixed_precision_inputs()
+
+    settings = aa.Settings(use_mixed_precision=True)
+
+    numpy_curvature = aa.util.inversion.curvature_matrix_via_mapping_matrix_from(
+        mapping_matrix=mapping_matrix, noise_map=noise_map, settings=settings
+    )
+    jax_curvature = aa.util.inversion.curvature_matrix_via_mapping_matrix_from(
+        mapping_matrix=jnp.asarray(mapping_matrix),
+        noise_map=jnp.asarray(noise_map),
+        settings=settings,
+        xp=jnp,
+    )
+
+    assert numpy_curvature.dtype == np.float64
+    assert jax_curvature.dtype == jnp.float64
+
+    assert np.asarray(jax_curvature) == pytest.approx(numpy_curvature, rel=1.0e-12)
+
+
+@requires_jax
+def test__curvature_matrix_via_mapping_matrix_from__mixed_precision__weights_are_fp64():
+    # The parity test above would also pass if *both* backends rounded the
+    # weights to fp32, so pin the absolute value: `F` must be the fp64
+    # `(A / sigma).T @ (A / sigma)`, and must NOT be the fp32-weighted one.
+    import jax
+
+    jax.config.update("jax_enable_x64", True)
+    import jax.numpy as jnp
+
+    mapping_matrix, noise_map = _curvature_mixed_precision_inputs()
+
+    weighted = mapping_matrix / noise_map[:, None]
+    expected = np.dot(weighted.T, weighted)
+
+    weighted_fp32 = mapping_matrix * (1.0 / noise_map).astype(np.float32)[:, None]
+    expected_fp32_weights = np.dot(weighted_fp32.T, weighted_fp32)
+
+    # The fixture has to be able to tell the two apart in the first place.
+    assert not np.allclose(expected, expected_fp32_weights, rtol=1.0e-14, atol=0.0)
+
+    jax_curvature = np.asarray(
+        aa.util.inversion.curvature_matrix_via_mapping_matrix_from(
+            mapping_matrix=jnp.asarray(mapping_matrix),
+            noise_map=jnp.asarray(noise_map),
+            settings=aa.Settings(use_mixed_precision=True),
+            xp=jnp,
+        )
+    )
+
+    assert jax_curvature == pytest.approx(expected, rel=1.0e-14)