From 273e152eaa64aebcb042d0684302f31f560ca0e0 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 23 Sep 2026 16:44:35 +0200 Subject: [PATCH 1/2] Drop the throwaway jnp.unique from the JAX triangle vertex table (#568) CoordinateArrayTriangles._vertices_and_indices now returns the flat (3N, 2) vertex table and an arange(3N).reshape(N, 3) index map. Under jit, jnp.unique needs a static size, so the "deduplicated" table was 3N rows anyway (no deflection evaluations saved) while costing a lexicographic sort of 3N fp64 rows twice per solver refinement step. NaN padding rows from for_indexes trace to NaN triangles, which every Shape.mask rejects, so containment is unchanged. The NumPy sibling CoordinateArrayTrianglesNp still deduplicates (dynamic shapes). See autolens_profiling#297. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../structures/triangles/coordinate_array.py | 43 +++++++++++-------- 1 file changed, 26 insertions(+), 17 deletions(-) diff --git a/autoarray/structures/triangles/coordinate_array.py b/autoarray/structures/triangles/coordinate_array.py index 1fd70fbd0..eab8e3497 100644 --- a/autoarray/structures/triangles/coordinate_array.py +++ b/autoarray/structures/triangles/coordinate_array.py @@ -283,23 +283,23 @@ def neighborhood(self) -> "CoordinateArrayTriangles": @property def _vertices_and_indices(self): + """ + The flat ``(3N, 2)`` vertex table and the ``(N, 3)`` index map into it. + + On this JAX path the table is deliberately *not* deduplicated: vertex ``3 * i + k`` is + vertex ``k`` of triangle ``i`` and ``indices`` is simply ``arange(3N).reshape(N, 3)``. + Under ``jit``, ``jnp.unique`` needs a static ``size``, so a deduplicated table was padded + back to 3N rows anyway (no deflection evaluations saved) while costing a lexicographic + sort of 3N fp64 rows, twice per solver refinement step (PyAutoArray#568, + autolens_profiling#297). NaN padding rows (from `for_indexes`) trace to NaN triangles, + which every `Shape.mask` rejects, so containment is unchanged. The NumPy sibling + `CoordinateArrayTrianglesNp` still deduplicates, because its shapes are dynamic. + """ import jax.numpy as jnp - flat_triangles = self.triangles.reshape(-1, 2) - vertices, inverse_indices = jnp.unique( - flat_triangles, - axis=0, - return_inverse=True, - size=3 * self.coordinates.shape[0], - equal_nan=True, - fill_value=jnp.nan, - ) - - nan_mask = jnp.isnan(vertices).any(axis=1) - inverse_indices = jnp.where(nan_mask[inverse_indices], -1, inverse_indices) - - indices = inverse_indices.reshape(-1, 3) - return vertices, indices + flat = self.triangles.reshape(-1, 2) + indices = jnp.arange(flat.shape[0]).reshape(-1, 3) + return flat, indices def with_vertices(self, vertices: np.ndarray) -> ArrayTriangles: """ @@ -350,14 +350,23 @@ def for_indexes(self, indexes: np.ndarray) -> "CoordinateArrayTriangles": @property def vertices(self) -> np.ndarray: """ - The unique vertices of the triangles. + The vertices of the triangles as a flat ``(3N, 2)`` table, row ``3 * i + k`` being vertex + ``k`` of triangle ``i``. + + Not deduplicated on this JAX path (see `_vertices_and_indices`): a static-shape + ``jnp.unique`` returned 3N rows regardless and only added a sort. Rows of NaN padding + triangles are NaN. `CoordinateArrayTrianglesNp.vertices` is deduplicated. """ return self._vertices_and_indices[0] @property def indices(self) -> np.ndarray: """ - The indices of the vertices of the triangles. + The indices of the vertices of the triangles, an ``(N, 3)`` map into `vertices`. + + On this JAX path it is ``arange(3N).reshape(N, 3)`` (no deduplication, see + `_vertices_and_indices`); padding triangles keep valid indices and are carried as NaN + vertices instead of ``-1`` entries. """ return self._vertices_and_indices[1] From 25894d10a7497d1efa6bc140032dba2f49d1795f Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 23 Sep 2026 16:44:35 +0200 Subject: [PATCH 2/2] Add JAX CoordinateArrayTriangles vertex-table tests (#568) Round trip of the non-deduplicated vertex table (eager and jit), NaN padding from for_indexes never contained, kept-triangle parity with the deduplicating NumPy sibling on a small lattice (interior, vertex, shared edge and centroid source points), NumPy deduplication pinned, and an HLO guard that the traced containment carries no sort op (fails on main 11b93476 with 39 sort matches). Co-Authored-By: Claude Opus 5.5 (1M context) --- .../triangles/test_coordinate_jax.py | 177 ++++++++++++++++++ 1 file changed, 177 insertions(+) create mode 100644 test_autoarray/structures/triangles/test_coordinate_jax.py diff --git a/test_autoarray/structures/triangles/test_coordinate_jax.py b/test_autoarray/structures/triangles/test_coordinate_jax.py new file mode 100644 index 000000000..10085f60f --- /dev/null +++ b/test_autoarray/structures/triangles/test_coordinate_jax.py @@ -0,0 +1,177 @@ +""" +Tests of the JAX `CoordinateArrayTriangles` vertex table and its containment path. + +`CoordinateArrayTriangles._vertices_and_indices` returns the flat, non-deduplicated ``(3N, 2)`` +vertex table with an ``arange`` index map (PyAutoArray#568, autolens_profiling#297): under ``jit`` a +static-size ``jnp.unique`` produced 3N rows anyway and only added a lexicographic sort. These tests +pin that the table round-trips to the triangles exactly, that NaN padding stays NaN and is never +returned by containment, that the kept triangles match the deduplicating NumPy sibling, and that +the traced containment carries no sort. +""" + +import importlib +import re + +import numpy as np +import pytest + + +# jax is an `[optional]` extra and is absent on the NumPy-only matrix env: every test in this module +# skips there. +if importlib.util.find_spec("jax") is None: + pytestmark = pytest.mark.skip(reason="requires jax (the [optional] extras)") + + def test__placeholder_requires_jax(): # pragma: no cover + pass + +else: + import jax + import jax.numpy as jnp + + jax.config.update("jax_enable_x64", True) + + from autoarray.structures.triangles.array import MAX_CONTAINING_SIZE + from autoarray.structures.triangles.coordinate_array import ( + CoordinateArrayTriangles, + ) + from autoarray.structures.triangles.coordinate_array_np import ( + CoordinateArrayTrianglesNp, + ) + from autoarray.structures.triangles.shape import Point + + LIMITS = dict(y_min=-1.0, y_max=1.0, x_min=-1.0, x_max=1.0, scale=0.5) + + def _lattice(): + return CoordinateArrayTriangles.for_limits_and_scale(**LIMITS) + + def _lattice_np(): + return CoordinateArrayTrianglesNp.for_limits_and_scale(**LIMITS) + + def _round_trip(coordinates): + triangles = CoordinateArrayTriangles(coordinates=coordinates, side_length=0.5) + return ( + triangles.triangles, + triangles.vertices.reshape(-1, 3, 2), + triangles.with_vertices(triangles.vertices).triangles, + ) + + def _containing(coordinates): + triangles = CoordinateArrayTriangles(coordinates=coordinates, side_length=0.5) + return triangles.with_vertices(triangles.vertices).containing_indices( + Point(0.1, 0.2) + ) + + def _source_points(triangles_np): + """ + Five source points on the lattice: two interior points, one on a lattice vertex, one on the + midpoint of an edge shared by two triangles, and one centroid. The edge is the horizontal one + (constant element 1), whose midpoint the barycentric test keeps on both sides; the midpoints + of the slanted edges round outside both triangles on NumPy and JAX alike. + """ + tri = np.asarray(triangles_np.triangles) + interior = int(np.argmin(np.linalg.norm(tri.mean(axis=1), axis=1))) + vertex = tri[interior, 0] + edge_midpoint = 0.5 * (tri[interior, 1] + tri[interior, 2]) + centroid = tri[interior].mean(axis=0) + return [ + (0.1, 0.2), + (-0.37, 0.41), + tuple(vertex), + tuple(edge_midpoint), + tuple(centroid), + ] + + def _kept_rows(triangles, point): + kept = triangles.for_indexes( + triangles.with_vertices(triangles.vertices).containing_indices( + Point(*point) + ) + ) + rows = np.asarray(kept.triangles).reshape(-1, 6) + rows = rows[np.all(np.isfinite(rows), axis=1)] + return {tuple(np.round(row, 8)) for row in rows} + + +def test__vertices_round_trip_to_triangles(): + triangles = _lattice() + + assert triangles.vertices.shape == (3 * triangles.coordinates.shape[0], 2) + assert np.array_equal( + np.asarray(triangles.indices), + np.arange(3 * triangles.coordinates.shape[0]).reshape(-1, 3), + ) + assert np.array_equal( + np.asarray(triangles.vertices.reshape(-1, 3, 2)), + np.asarray(triangles.triangles), + ) + assert np.array_equal( + np.asarray(triangles.with_vertices(triangles.vertices).triangles), + np.asarray(triangles.triangles), + ) + + +def test__vertices_round_trip_to_triangles__jit(): + """ + Compared inside one jitted computation: XLA may fuse the eager and jitted triangle arithmetic + differently (the two differ by up to 1 ulp), so the exact round trip is asserted between the + outputs of the same compiled program. + """ + triangles, vertices, with_vertices = jax.jit(_round_trip)(_lattice().coordinates) + + assert np.array_equal(np.asarray(vertices), np.asarray(triangles)) + assert np.array_equal(np.asarray(with_vertices), np.asarray(triangles)) + + +def test__nan_padding_stays_nan_and_is_never_contained(): + lattice = _lattice() + point = tuple(np.asarray(lattice.triangles[0]).mean(axis=0)) + + padded = lattice.for_indexes(jnp.array([0, 2, -1, -1])) + + coordinates = np.asarray(padded.coordinates) + assert np.all(np.isfinite(coordinates[:2])) + assert np.all(np.isnan(coordinates[2:])) + + triangles = np.asarray(padded.with_vertices(padded.vertices).triangles) + assert np.all(np.isfinite(triangles[:2])) + assert np.all(np.isnan(triangles[2:])) + + containing = np.asarray( + padded.with_vertices(padded.vertices).containing_indices(Point(*point)) + ) + assert containing[0] == 0 + assert np.all(containing[1:] == -1) + + +@pytest.mark.parametrize("point_index", range(5)) +def test__kept_triangles_match_numpy(point_index): + triangles_np = _lattice_np() + triangles = _lattice() + point = _source_points(triangles_np)[point_index] + + kept_np = _kept_rows(triangles_np, point) + kept = _kept_rows(triangles, point) + + assert 0 < len(kept_np) < MAX_CONTAINING_SIZE + assert kept == kept_np + + +def test__numpy_vertices_are_deduplicated(): + triangles_np = _lattice_np() + n = triangles_np.coordinates.shape[0] + + assert triangles_np.vertices.shape[0] < 3 * n + assert np.unique(triangles_np.vertices, axis=0).shape[0] == ( + triangles_np.vertices.shape[0] + ) + assert np.array_equal( + triangles_np.vertices[triangles_np.indices], triangles_np.triangles + ) + + +def test_no_sort(): + coordinates = _lattice().coordinates + compiled = jax.jit(_containing).lower(coordinates).compile().as_text() + + sorts = re.findall(r"\bsort\b", compiled, flags=re.IGNORECASE) + assert not sorts, f"traced containment carries {len(sorts)} sort op(s)"