From 72fb01d14db82bd0b946c50c30dfe7d446ba2671 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 26 Aug 2026 00:53:49 +0000 Subject: [PATCH] fix: correct mirrored row weights + round-off-dependent cells in rectangular mapper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `adaptive_rectangular_mappings_weights_via_interpolation_from` — the bilinear mapper shared by every adaptive rectangular mesh — had two coupled defects. 1. Mirrored row weights. `t_row` is the fractional distance measured FROM `ix_down`, so `ix_up` must carry `t_row` and `ix_down` `1 - t_row`. They were the other way round, mirroring the interpolation in the row direction. The column weights were, and remain, correctly paired. 2. Round-off-dependent cell assignment. `ix_up = ceil(g)` collapses onto `ix_down` wherever `g` is exactly integral. `transform()` ends in `clip(F_q, 0.0, 1.0)`, so saturated points land on exactly integer `g` systematically, not by chance. There the cell degenerated and `t_row` was forced to 0 by the `+ 1e-12` guard; a 1-ULP change in the traced grid moved a point off the plateau and jumped its weight a whole mesh row. (2) is why the eager and jitted likelihoods disagreed by ~1.6e-3 on an otherwise smooth surface, tripping `assert_eager_jit_consistent` in autolens_workspace_test (issue #279). (1) is why (2) could not be fixed without a behaviour change. Both are regressions, not design choices. The mapper was correct when introduced in fd11b178 (2025-06-24); 8f007957 (2025-09-15) forked it for the adaptive meshes and mirrored the columns, and 9b1c91cf (2025-09-23) "fixed mappings and weights" fixed the columns and broke the rows. The correct formulation never left the package — it survives verbatim for the uniform mesh at interpolator/rectangular_uniform.py:72-99, which is what this restores. Affected meshes (all inherit `InterpolatorRectangular` from `RectangularRTUAdaptDensity`): RectangularRTUAdaptDensity, RectangularRTUAdaptImage, RectangularBilinearAdaptDensity, RectangularBilinearAdaptImage. Not affected: RectangularUniform, Delaunay, DelaunayNN, KNearestNeighbor, KNNBarycentric. Adds the two regression tests whose absence let this survive eleven months and a "fix" commit: partition of unity plus linear reproduction (the property any consistent mis-pairing satisfies the first but not the second), and continuity of the cell assignment across integer boundaries. Both fail on the previous implementation — linear reproduction by ~1.0 in the row axis, continuity by a jump of 5.999 (a two-row flip) — and pass here. Behaviour change: reconstructions using an adaptive rectangular mesh shift. Measured on autolens_workspace_test/scripts/imaging/jax_likelihood/rectangular.py, log likelihood moves -651692.997799 -> -650470.379097 (+1222.6, a better fit). Downstream `EXPECTED_LOG_*` constants need regenerating; see #279. Validated: pytest test_autoarray/ 1222 passed on Python 3.12 and 3.13. autolens_workspace_test/scripts/interferometer/jax_grad/gradient.py green on both legs, all four variants, eager/jit gap exactly 0.0 with the guard left at rtol=1e-10 (3.13 83s, 3.12 86s against the 300s cap). Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_016azuS2UbS3mFvfxkDFGsKj --- .../mesh/interpolator/rectangular.py | 43 +++++-- .../interpolator/test_rectangular.py | 105 ++++++++++++++++++ 2 files changed, 136 insertions(+), 12 deletions(-) diff --git a/autoarray/inversion/mesh/interpolator/rectangular.py b/autoarray/inversion/mesh/interpolator/rectangular.py index ab16ce294..1c0ebd6ed 100644 --- a/autoarray/inversion/mesh/interpolator/rectangular.py +++ b/autoarray/inversion/mesh/interpolator/rectangular.py @@ -544,11 +544,23 @@ def adaptive_rectangular_mappings_weights_via_interpolation_from( grid_over_sampled_transformed = transform_func(grid_over_sampled_scaled) grid_over_index = (source_grid_size - 3) * grid_over_sampled_transformed + 1 - # --- Step 4. Floor/ceil indices --- - ix_down = xp.floor(grid_over_index[:, 0]) - ix_up = xp.ceil(grid_over_index[:, 0]) - iy_down = xp.floor(grid_over_index[:, 1]) - iy_up = xp.ceil(grid_over_index[:, 1]) + # --- Step 4. Bracketing indices --- + # + # `ix_up` is ALWAYS `ix_down + 1`, never `ceil(grid_over_index)`. + # `transform()` ends in `xp.clip(F_q, 0.0, 1.0)`, so saturated points land + # on EXACTLY integer `grid_over_index` values. Under `ceil` the bracketing + # cell collapses there (`ix_up == ix_down`) and the interpolation weight is + # forced onto a single row; a 1-ULP change in the traced grid then moves + # such a point off the plateau and jumps its weight a whole mesh row. That + # is a discontinuity in the discretisation, and it is what made the eager + # and jitted likelihoods disagree by ~1.6e-3 on an otherwise smooth + # surface (autolens_workspace_test#279). Bracketing with `ix_down + 1` is + # continuous at integer coordinates, so the cell assignment no longer + # depends on round-off. Clamp so the `+ 1` cannot leave the mesh. + ix_down = xp.clip(xp.floor(grid_over_index[:, 0]), 0, source_grid_size - 2) + iy_down = xp.clip(xp.floor(grid_over_index[:, 1]), 0, source_grid_size - 2) + ix_up = ix_down + 1 + iy_up = iy_down + 1 # --- Step 5. Four corners --- idx_tl = xp.stack([ix_up, iy_down], axis=1) @@ -570,13 +582,20 @@ def flatten(idx, n): ) # --- Step 7. Bilinear interpolation weights --- - t_row = (grid_over_index[:, 0] - ix_down) / (ix_up - ix_down + 1e-12) - t_col = (grid_over_index[:, 1] - iy_down) / (iy_up - iy_down + 1e-12) - - w_tl = (1 - t_row) * (1 - t_col) - w_tr = (1 - t_row) * t_col - w_bl = t_row * (1 - t_col) - w_br = t_row * t_col + # + # `t_row` / `t_col` are fractional distances measured FROM the `down` node, + # so the `up` node carries `t` and the `down` node carries `1 - t`. The row + # weights were previously the other way round (`ix_up` carried `1 - t_row`, + # `ix_down` carried `t_row`), mirroring the interpolation in the row + # direction; the column weights were, and remain, correctly paired. No + # `+ 1e-12` guard is needed now the bracket is always exactly one cell wide. + t_row = grid_over_index[:, 0] - ix_down + t_col = grid_over_index[:, 1] - iy_down + + w_tl = t_row * (1 - t_col) + w_tr = t_row * t_col + w_bl = (1 - t_row) * (1 - t_col) + w_br = (1 - t_row) * t_col weights = xp.stack([w_tl, w_tr, w_bl, w_br], axis=1) return flat_indices, weights diff --git a/test_autoarray/inversion/pixelization/interpolator/test_rectangular.py b/test_autoarray/inversion/pixelization/interpolator/test_rectangular.py index fee895705..4cd5c1f28 100644 --- a/test_autoarray/inversion/pixelization/interpolator/test_rectangular.py +++ b/test_autoarray/inversion/pixelization/interpolator/test_rectangular.py @@ -293,6 +293,111 @@ def test__mappings_sizes_weights__shapes_and_weight_normalization(): assert np.allclose(w.sum(axis=1), 1.0, atol=1e-10) +def _index_space_nodes(idx, n): + """ + Recover each mapped pixel's (row, col) position in index space. + + Inverse of the module's ``flatten(ix, iy) = (n - ix) * n + iy``. + """ + return n - idx // n, idx % n + + +def test__mappings_sizes_weights__reproduces_the_query_position(): + """ + Bilinear interpolation must be exact for linear functions, which means the + four weights have to reconstruct the query itself: + + sum_i w_i * node_i == grid_over_index + + Partition of unity alone does NOT imply this — it is satisfied by any + consistent mis-pairing of corners to weights, which is exactly how the row + weights came to be mirrored (`ix_up` carrying `1 - t_row` instead of + `t_row`) and stayed that way from 2025-09-23 to 2026-08-26. The mirroring + was smooth, so gradient checks passed; only this property catches it. + See autolens_workspace_test#279. + """ + n = 16 + data_grid, over, weights = _seeded_inputs(seed=11) + + idx, w = adaptive_rectangular_mappings_weights_via_interpolation_from( + source_grid_size=n, + data_grid=data_grid, + data_grid_over_sampled=over, + mesh_weight_map=weights, + xp=np, + ) + + # Rebuild the index-space query the function discretises internally. + mu, scale = data_grid.mean(axis=0), data_grid.std(axis=0) + transform_func, _ = create_transforms( + (data_grid - mu) / scale, mesh_pixels=n, mesh_weight_map=weights, xp=np + ) + grid_over_index = (n - 3) * transform_func((over - mu) / scale) + 1 + + node_row, node_col = _index_space_nodes(idx, n) + + assert np.allclose(w.sum(axis=1), 1.0, atol=1e-10) + assert np.allclose((w * node_row).sum(axis=1), grid_over_index[:, 0], atol=1e-10) + assert np.allclose((w * node_col).sum(axis=1), grid_over_index[:, 1], atol=1e-10) + + +def test__mappings_sizes_weights__cell_assignment_is_continuous_at_integers(): + """ + ``transform()`` ends in ``clip(F_q, 0.0, 1.0)``, so saturated queries land on + EXACTLY integer ``grid_over_index`` values — systematically, not by chance. + Bracketing those with ``ceil`` collapsed the cell (``ix_up == ix_down``), so + a 1-ULP move off the plateau jumped a point's weight a whole mesh row. That + made the likelihood depend on floating-point association, which is how the + eager and jitted evaluations came to disagree by ~1.6e-3. + + Here the property is asserted directly on the interpolated value of a linear + ramp: approaching an integer row coordinate from either side must converge + to the value AT that coordinate. + """ + n = 16 + data_grid, _, weights = _seeded_inputs(seed=12) + mu, scale = data_grid.mean(axis=0), data_grid.std(axis=0) + + transform_func, inv = create_transforms( + (data_grid - mu) / scale, mesh_pixels=n, mesh_weight_map=weights, xp=np + ) + + def interpolated_ramp(over): + idx, w = adaptive_rectangular_mappings_weights_via_interpolation_from( + source_grid_size=n, + data_grid=data_grid, + data_grid_over_sampled=over, + mesh_weight_map=weights, + xp=np, + ) + node_row, node_col = _index_space_nodes(idx, n) + # A linear function of position; bilinear interpolation reproduces it + # exactly, so any discontinuity here is a cell-assignment jump. + return (w * (3.0 * node_row - 2.0 * node_col)).sum(axis=1) + + # Sweep the row coordinate densely across the whole data range. In index + # space that spans [1, n - 2], so the sweep crosses every interior integer + # boundary; a jump at any crossing is a cell-assignment discontinuity. + # This is deliberately placement-free — the inverse transform is a knot + # lookup and cannot land a query on an integer precisely enough to probe + # one boundary directly. + lo, hi = data_grid[:, 0].min(), data_grid[:, 0].max() + sweep = np.stack( + [np.linspace(lo, hi, 20001), np.full(20001, np.median(data_grid[:, 1]))], + axis=1, + ) + + values = interpolated_ramp(sweep) + steps = np.abs(np.diff(values)) + + # The ramp must actually vary, or continuity is vacuous. + assert values.max() - values.min() > 1.0 + + # A whole-row flip moves the ramp by ~3.0 (its row coefficient); a + # continuous scheme moves by ~the sweep resolution. + assert steps.max() < 0.05 + + # --------------------------------------------------------------------------- # Areas # ---------------------------------------------------------------------------