From 3dfd869625e9d0fd9d13d855fc4740e7ff7327a6 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Thu, 17 Sep 2026 12:50:15 +0000 Subject: [PATCH 1/2] fix: make the rank-CDF mesh transform tie-safe; NumPy backend fp64 under mixed precision (#552) The JAX-vs-NumPy log-likelihood gap of the RectangularBilinearAdaptImage + Adapt smoke script (15.5 nats on the 316-pixel dataset, 5.8 nats on the old 716-pixel one) is not a mixed-precision loss: mixed vs fp64 within either backend differs by ~1e-4 nats, the fnnls/jaxnnls solvers agree to 1e-12 on identical inputs, and the FFT and real-space convolutions agree to 1e-16. The whole gap sits in `create_transforms_rank`'s weighted branch, which accumulates the adapt-image weights in `argsort` order: `np.argsort` (quicksort, unstable) and `jnp.argsort` (stable) order the heavily tied coordinates of an unlensed image grid differently (553 of 632 permutation entries, 80 sub-pixels landing in a different mesh cell). - interpolator/rectangular.py: every point of a tie block takes the block's final cumulative weight (the right-continuous empirical CDF), so the transform depends only on the point multiset, not on the backend's tie-break or the input order. Untied points are unchanged bit-for-bit. - mapper_util.mapping_matrix_from / Convolver.mapping_matrix_native_from: the fp32 allocation under `use_mixed_precision` is now JAX-only, so the NumPy backend is the fp64 reference the Settings docstring promises. - inversion_util.curvature_matrix_via_mapping_matrix_from: one fp64 expression on every backend; the former fp32 `1/noise_map` on the JAX branch never produced an fp32 accumulation (the blurred mapping matrix is already fp64) and only weighted F inconsistently with D. - Settings docstring updated to match. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_018bJQFXsFc4iCqpUywydex2 --- .../inversion/inversion/inversion_util.py | 23 ++++++--------- autoarray/inversion/mappers/mapper_util.py | 11 +++++-- .../mesh/interpolator/rectangular.py | 29 +++++++++++++++++++ autoarray/operators/convolver.py | 8 ++++- autoarray/settings.py | 23 ++++++++++----- 5 files changed, 70 insertions(+), 24 deletions(-) diff --git a/autoarray/inversion/inversion/inversion_util.py b/autoarray/inversion/inversion/inversion_util.py index 3dd67fc7b..9485bb240 100644 --- a/autoarray/inversion/inversion/inversion_util.py +++ b/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 diff --git a/autoarray/inversion/mappers/mapper_util.py b/autoarray/inversion/mappers/mapper_util.py index 629d8dfac..ae0394f85 100644 --- a/autoarray/inversion/mappers/mapper_util.py +++ b/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,) diff --git a/autoarray/inversion/mesh/interpolator/rectangular.py b/autoarray/inversion/mesh/interpolator/rectangular.py index 1c0ebd6ed..20118f373 100644 --- a/autoarray/inversion/mesh/interpolator/rectangular.py +++ 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) diff --git a/autoarray/operators/convolver.py b/autoarray/operators/convolver.py index c3b5bd72f..7611a72fc 100644 --- a/autoarray/operators/convolver.py +++ b/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] diff --git a/autoarray/settings.py b/autoarray/settings.py index e72c9e07b..226a21128 100644 --- a/autoarray/settings.py +++ b/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 From 09da4c40a4fabc7c7a6f52a348a2ee81c1802d97 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Thu, 17 Sep 2026 12:54:57 +0000 Subject: [PATCH 2/2] test: first use_mixed_precision unit tests and tie-safe rank-CDF tests (#552) - interpolator/test_rectangular.py: the weighted rank CDF on tied points agrees between NumPy and JAX, is invariant to the input ordering, keeps the plain cumsum/interp value on untied points, and takes the whole block's weight at a tie. - mappers/test_mapper_util.py: mapping_matrix_from stays float64 under numpy with use_mixed_precision=True, is float32 under jax.numpy, float64 on both without the flag. - inversion/test_inversion_util.py: the curvature matrix under mixed precision matches NumPy on the JAX backend and is weighted with the fp64 1/noise_map (a noise value whose reciprocal is not float32-representable). - operators/test_convolver.py: the native cube dtype per backend, and the mixed-precision JAX FFT blurred matrix against the NumPy fp64 result. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_018bJQFXsFc4iCqpUywydex2 --- .../inversion/test_inversion_util.py | 103 ++++++++++++ .../interpolator/test_rectangular.py | 151 +++++++++++++++++- .../pixelization/mappers/test_mapper_util.py | 89 +++++++++++ test_autoarray/operators/test_convolver.py | 131 +++++++++++++++ 4 files changed, 472 insertions(+), 2 deletions(-) diff --git a/test_autoarray/inversion/inversion/test_inversion_util.py b/test_autoarray/inversion/inversion/test_inversion_util.py index 7aa5944fc..4bbbb7fcb 100644 --- a/test_autoarray/inversion/inversion/test_inversion_util.py +++ b/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) diff --git a/test_autoarray/inversion/pixelization/interpolator/test_rectangular.py b/test_autoarray/inversion/pixelization/interpolator/test_rectangular.py index 4cd5c1f28..485262aef 100644 --- a/test_autoarray/inversion/pixelization/interpolator/test_rectangular.py +++ b/test_autoarray/inversion/pixelization/interpolator/test_rectangular.py @@ -1,11 +1,16 @@ """Unit tests for the adaptive rectangular interpolator (rank and kernel-density CDF transforms). -Pure numpy — no JAX imports here. Cross-xp / gradient certification lives in +Numpy by default — cross-xp / gradient certification lives in autolens_workspace_test/scripts/jax_grad per the project's "no JAX in unit -tests" rule. +tests" rule. The one exception is the tie-block section at the bottom of this +module: the defect it pins (PyAutoArray#552) *was* a NumPy-vs-JAX divergence, +so a NumPy-only test cannot express it. Those tests import jax inside the test +body, behind `requires_jax`, exactly as test_mapper_util.py does. """ +import importlib.util + import numpy as np import pytest @@ -559,3 +564,145 @@ def test__forward_transform__windowed_numba_matches_dense_reference(weighted): ) np.testing.assert_allclose(fwd(q), reference, rtol=0.0, atol=1e-12) + + +# --------------------------------------------------------------------------- +# Tied coordinates in the weighted rank-CDF transform (PyAutoArray#552) +# +# On an unlensed image grid every column and every row of the traced points is +# a tie block (316 points over 38 distinct coordinates per axis, up to 18-fold +# ties). The weighted branch used to accumulate the adapt weights in `argsort` +# order, and `np.argsort` (quicksort, unstable) and `jnp.argsort` (stable) +# order tied entries differently — moving sub-pixels between mesh cells and +# opening a 15.47-nat NumPy-vs-JAX log-likelihood gap. Giving every point of a +# tie block the block's final cumulative weight (the right-continuous +# empirical CDF) makes the transform a function of the point multiset alone. +# --------------------------------------------------------------------------- + +# jax is an `[optional]` extra and is absent on the NumPy-only matrix env, so +# the JAX parity tests skip rather than fail there (same convention as +# test_mapper_util.py). +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 _tied_lattice(n=7, seed=0): + """ + An n x n lattice — n distinct coordinates per axis, n-fold ties on every + one of them — with random positive weights normalised to 1, and a set of + off-lattice queries. + """ + rng = np.random.default_rng(seed) + + axis = np.linspace(-1.0, 1.0, n) + points = np.array([[y, x] for y in axis for x in axis]) + + weights = rng.uniform(0.1, 1.0, size=points.shape[0]) + weights = weights / weights.sum() + + queries = rng.uniform(-1.2, 1.2, size=(23, 2)) + + return points, weights, queries + + +@requires_jax +def test__create_transforms_rank__weighted__tied_points__numpy_matches_jax(): + import jax + + jax.config.update("jax_enable_x64", True) + import jax.numpy as jnp + + points, weights, queries = _tied_lattice() + + # The fixture must actually tie, or the test proves nothing. + assert np.unique(points[:, 0]).size == 7 + assert points.shape[0] == 49 + + fwd_np, rev_np = create_transforms_rank(points, mesh_weight_map=weights, xp=np) + fwd_jax, rev_jax = create_transforms_rank( + jnp.asarray(points), mesh_weight_map=jnp.asarray(weights), xp=jnp + ) + + forward_np = fwd_np(queries) + forward_jax = np.asarray(fwd_jax(jnp.asarray(queries))) + + assert forward_jax == pytest.approx(forward_np, abs=1.0e-12) + + # Interior unit-square probes for the inverse. + probes = np.random.default_rng(5).uniform(0.05, 0.95, size=(23, 2)) + + inverse_np = rev_np(probes) + inverse_jax = np.asarray(rev_jax(jnp.asarray(probes))) + + assert inverse_jax == pytest.approx(inverse_np, abs=1.0e-12) + + +def test__create_transforms_rank__weighted__tied_points__permutation_invariant(): + # The transform is a function of the (point, weight) multiset: permuting + # the inputs is the same perturbation the backends' argsort tie-breaks + # apply, without needing a second backend to express it. + points, weights, queries = _tied_lattice(seed=1) + + fwd, rev = create_transforms_rank(points, mesh_weight_map=weights, xp=np) + + order = np.random.default_rng(7).permutation(points.shape[0]) + fwd_permuted, rev_permuted = create_transforms_rank( + points[order], mesh_weight_map=weights[order], xp=np + ) + + probes = np.random.default_rng(9).uniform(0.05, 0.95, size=(23, 2)) + + assert fwd_permuted(queries) == pytest.approx(fwd(queries), abs=1.0e-12) + assert rev_permuted(probes) == pytest.approx(rev(probes), abs=1.0e-12) + + +def test__create_transforms_rank__weighted__untied_points__matches_plain_cumsum_interp(): + # Regression guard on the other side of the fix: with no ties the + # right-continuous CDF must reduce, bit-for-bit, to the pre-fix formula + # (interpolate cumsum(w[argsort]) at the sorted coordinates). + points, weights, queries = _tied_lattice(seed=2) + + jitter = np.random.default_rng(11).uniform(-0.01, 0.01, size=points.shape) + points = points + jitter + + assert np.unique(points[:, 0]).size == points.shape[0] + assert np.unique(points[:, 1]).size == points.shape[0] + + fwd, _ = create_transforms_rank(points, mesh_weight_map=weights, xp=np) + + expected = np.empty((queries.shape[0], 2)) + for d in range(2): + order = np.argsort(points[:, d], kind="stable") + expected[:, d] = np.interp( + queries[:, d], + points[order, d], + np.cumsum(weights[order]), + left=0, + right=1, + ) + + assert fwd(queries) == pytest.approx(expected, abs=0.0) + + +def test__create_transforms_rank__weighted__cdf_at_a_tie_is_the_whole_block(): + # Right-continuity: F(v) is the total weight of every point <= v, for + # every tied coordinate v — not a partial sum through the block. + points, weights, _ = _tied_lattice(seed=3) + + fwd, _ = create_transforms_rank(points, mesh_weight_map=weights, xp=np) + + for d in range(2): + for value in np.unique(points[:, d]): + query = np.zeros((1, 2)) + query[0, d] = value + + expected = weights[points[:, d] <= value].sum() + + assert fwd(query)[0, d] == pytest.approx(expected, abs=1.0e-12) + + # ... and the transform stays monotone non-decreasing along each axis. + dense = np.linspace(points.min(axis=0) - 0.3, points.max(axis=0) + 0.3, 500) + + assert np.all(np.diff(fwd(dense), axis=0) >= 0.0) diff --git a/test_autoarray/inversion/pixelization/mappers/test_mapper_util.py b/test_autoarray/inversion/pixelization/mappers/test_mapper_util.py index a3c550312..9ce921a90 100644 --- a/test_autoarray/inversion/pixelization/mappers/test_mapper_util.py +++ b/test_autoarray/inversion/pixelization/mappers/test_mapper_util.py @@ -567,3 +567,92 @@ def total_signal(data): gradient = jax.grad(total_signal)(jnp.asarray(adapt_data)) assert np.isfinite(np.asarray(gradient)).all() + + +# ---------------------------------------------------------------------------- +# Mixed precision is a JAX-only allocation choice (PyAutoArray#552). +# +# `use_mixed_precision` buys GPU throughput on the JAX path. Applying it on the +# NumPy path as well silently removed the fp64 reference a NumPy fit is +# supposed to be: the workspace smoke scripts compare a NumPy fit against a +# mixed-precision JAX fit and call the difference the mixed-precision error. +# ---------------------------------------------------------------------------- + + +def _mixed_precision_mapping_kwargs(): + """ + Three image pixels, each mapped by two weighted sub-pixels, with weights + that are not representable in fp32 so a downcast is visible in the values + as well as in the dtype. + """ + return dict( + pix_indexes_for_sub_slim_index=np.array([[0, 1], [1, 2], [2, 0]]), + pix_size_for_sub_slim_index=np.array([2, 2, 2]), + pix_weights_for_sub_slim_index=np.array( + [[1.0 / 3.0, 2.0 / 3.0], [0.7, 0.3], [0.1, 0.9]] + ), + pixels=3, + total_mask_pixels=3, + slim_index_for_sub_slim_index=np.array([0, 1, 2]), + sub_fraction=np.array([1.0, 1.0, 1.0]), + ) + + +def test__mapping_matrix_from__mixed_precision__numpy_stays_float64(): + kwargs = _mixed_precision_mapping_kwargs() + + mapping_matrix = aa.util.mapper.mapping_matrix_from( + **kwargs, use_mixed_precision=True + ) + + assert mapping_matrix.dtype == np.float64 + + mapping_matrix = aa.util.mapper.mapping_matrix_from( + **kwargs, use_mixed_precision=False + ) + + assert mapping_matrix.dtype == np.float64 + + +@requires_jax +def test__mapping_matrix_from__mixed_precision__jax_is_float32_and_matches_numpy(): + import jax + + jax.config.update("jax_enable_x64", True) + import jax.numpy as jnp + + kwargs = _mixed_precision_mapping_kwargs() + + numpy_matrix = aa.util.mapper.mapping_matrix_from( + **kwargs, use_mixed_precision=True + ) + jax_matrix = aa.util.mapper.mapping_matrix_from( + **kwargs, use_mixed_precision=True, xp=jnp + ) + + assert numpy_matrix.dtype == np.float64 + assert jax_matrix.dtype == jnp.float32 + + assert np.asarray(jax_matrix) == pytest.approx(numpy_matrix, abs=1.0e-6) + + +@requires_jax +def test__mapping_matrix_from__no_mixed_precision__float64_on_both_backends(): + import jax + + jax.config.update("jax_enable_x64", True) + import jax.numpy as jnp + + kwargs = _mixed_precision_mapping_kwargs() + + numpy_matrix = aa.util.mapper.mapping_matrix_from( + **kwargs, use_mixed_precision=False + ) + jax_matrix = aa.util.mapper.mapping_matrix_from( + **kwargs, use_mixed_precision=False, xp=jnp + ) + + assert numpy_matrix.dtype == np.float64 + assert jax_matrix.dtype == jnp.float64 + + assert np.asarray(jax_matrix) == pytest.approx(numpy_matrix, abs=1.0e-14) diff --git a/test_autoarray/operators/test_convolver.py b/test_autoarray/operators/test_convolver.py index 2aee39511..f0011323e 100644 --- a/test_autoarray/operators/test_convolver.py +++ b/test_autoarray/operators/test_convolver.py @@ -1,3 +1,5 @@ +import importlib.util + from astropy import units from astropy.modeling import functional_models from astropy.coordinates import Angle @@ -680,3 +682,132 @@ def test__convolved_mapping_matrix_via_real_space_np__matches_image_convolution_ assert convolved[:, index] == pytest.approx( np.array(convolved_image), abs=1.0e-12 ) + + +# ---------------------------------------------------------------------------- +# Mixed precision is a JAX-only allocation choice (PyAutoArray#552). +# +# The native mapping-matrix cube is the big allocation the fp32 path exists +# for, and it only buys anything on a GPU. Downcasting it under `xp=np` as +# well removed the fp64 reference the workspace smoke scripts compare a +# mixed-precision JAX fit against. +# ---------------------------------------------------------------------------- + +# 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 _mixed_precision_convolver_inputs(): + """ + The single-source-pixel example of + `test__convolved_mapping_matrix_from__single_source_pixel__...`: a 4x4 + unmasked region inside a 6x6 mask, an asymmetric 3x3 kernel, and three + source pixels each reached by one image pixel. + """ + mask = aa.Mask2D( + mask=np.array( + [ + [True, True, True, True, True, True], + [True, False, False, False, False, True], + [True, False, False, False, False, True], + [True, False, False, False, False, True], + [True, False, False, False, False, True], + [True, True, True, True, True, True], + ] + ), + pixel_scales=1.0, + ) + + kernel = aa.Array2D.no_mask( + values=[[0, 0.0, 0], [0.4, 0.2, 0.3], [0, 0.1, 0]], + pixel_scales=mask.pixel_scales, + ) + + mapping_matrix = np.zeros((16, 3)) + mapping_matrix[7, 1] = 1.0 + mapping_matrix[9, 0] = 1.0 + mapping_matrix[10, 2] = 1.0 + + return mask, kernel, mapping_matrix + + +def test__mapping_matrix_native_from__mixed_precision__numpy_stays_float64(): + mask, kernel, mapping_matrix = _mixed_precision_convolver_inputs() + + convolver = aa.Convolver(kernel=kernel) + + for use_mixed_precision in (True, False): + native = convolver.mapping_matrix_native_from( + mapping_matrix=mapping_matrix, + mask=mask, + use_mixed_precision=use_mixed_precision, + ) + + assert native.dtype == np.float64 + + +@requires_jax +def test__mapping_matrix_native_from__mixed_precision__jax_is_float32(): + import jax + + jax.config.update("jax_enable_x64", True) + import jax.numpy as jnp + + mask, kernel, mapping_matrix = _mixed_precision_convolver_inputs() + + convolver = aa.Convolver(kernel=kernel) + + native_mixed = convolver.mapping_matrix_native_from( + mapping_matrix=jnp.asarray(mapping_matrix), + mask=mask, + use_mixed_precision=True, + xp=jnp, + ) + native_fp64 = convolver.mapping_matrix_native_from( + mapping_matrix=jnp.asarray(mapping_matrix), + mask=mask, + use_mixed_precision=False, + xp=jnp, + ) + + assert native_mixed.dtype == jnp.float32 + assert native_fp64.dtype == jnp.float64 + + assert np.asarray(native_mixed) == pytest.approx( + np.asarray(native_fp64), abs=1.0e-6 + ) + + +@requires_jax +def test__convolved_mapping_matrix_from__mixed_precision__jax_fft_matches_numpy_fp64(): + # The fp32 cube feeds a complex64 forward FFT, but the kernel multiply is + # deliberately complex128, so the convolved matrix comes back fp64 and + # within fp32 round-off of the fp64 NumPy reference. + import jax + + jax.config.update("jax_enable_x64", True) + import jax.numpy as jnp + + mask, kernel, mapping_matrix = _mixed_precision_convolver_inputs() + + reference = aa.Convolver(kernel=kernel).convolved_mapping_matrix_from( + mapping_matrix, mask + ) + + assert reference.dtype == np.float64 + + convolver_fft = aa.Convolver(kernel=kernel, use_fft=True) + + blurred = convolver_fft.convolved_mapping_matrix_from( + jnp.asarray(mapping_matrix), + mask, + use_mixed_precision=True, + xp=jnp, + ) + + assert blurred.dtype == jnp.float64 + assert np.asarray(blurred) == pytest.approx(reference, abs=1.0e-5)