diff --git a/autoarray/inversion/mesh/interpolator/rectangular_kernel.py b/autoarray/inversion/mesh/interpolator/rectangular_kernel.py new file mode 100644 index 00000000..ac897728 --- /dev/null +++ b/autoarray/inversion/mesh/interpolator/rectangular_kernel.py @@ -0,0 +1,347 @@ +""" +Kernel-density CDF variant of the rectangular adaptive interpolator. + +The adaptive rectangular mesh transforms source-plane coordinates through a +per-axis CDF so that the uniform mesh pixels adapt to the density (or weight) +of the traced points. The "linear" variant in ``rectangular.py`` uses a +piecewise-linear empirical CDF whose knots are the traced points themselves: +whenever interp queries coincide with knots (imaging at pixelization +over-sampling 1, the interferometer sparse path) the likelihood becomes +exactly piecewise-constant in the mass/shear parameters, and every rank swap +between traced points leaves a kink even when gradients do flow. + +This module replaces the empirical CDF with a smooth kernel-density CDF +(the RTU formulation of Enzi et al., arXiv:2606.30620): per axis + + F(x) = sum_i w_i * Phi((x - x_i) / h) + +where ``Phi`` is the standard normal CDF, ``x_i`` the traced points, ``w_i`` +uniform weights (density adaption) or the normalized adapt-image weights +(image adaption), and ``h`` a bandwidth tied to the mesh resolution. + +Properties, by construction: + +- strictly monotone (no jitter hack, no ``_enforce_strict_monotone``); +- C-infinity in the queries AND the traced-point positions — no ranks, no + sorts, no ``argsort`` anywhere, so there is nothing to swap; +- duplicate-safe: coincident traced points simply stack their weights (no + ``1 / Δknot`` terms). + +The forward transform is evaluated exactly at the query points (keeping the +C-infinity guarantee) and rescaled so the data bounding box maps onto the +unit square exactly — matching the linear variant's convention, including +clamping to [0, 1] outside the data range. The inverse — only needed at +fixed unit-square grid values (mesh pixel centres/edges) — is a linear-interp +lookup on a small fixed table of ``n_knots`` knots spanning the data range; +because the inverse queries are constants, gradients flow smoothly through +the table values. + +The linear meshes in ``rectangular.py`` and the spline meshes in +``rectangular_spline.py`` are untouched; per the pattern established there, +the bilinear-weights steps are copied rather than shared so the variants stay +independently auditable. +""" +import numpy as np +from functools import partial +from typing import Optional + +from autoconf import cached_property + +from autoarray.inversion.mesh.interpolator.rectangular import ( + InterpolatorRectangular, + reverse_interp, + reverse_interp_np, +) + + +KERNEL_CDF_DEFAULT_BANDWIDTH: float = 1.0 +KERNEL_CDF_DEFAULT_KNOTS: int = 64 + +_SQRT2 = np.sqrt(2.0) + + +def _norm_cdf(t, xp): + """Standard normal CDF, xp-aware (scipy erf on numpy, jax.scipy under jax).""" + if xp is np: + from scipy.special import erf + else: + from jax.scipy.special import erf + + return 0.5 * (1.0 + erf(t / _SQRT2)) + + +def create_transforms_kernel( + traced_points, + mesh_pixels: int, + mesh_weight_map=None, + bandwidth: float = KERNEL_CDF_DEFAULT_BANDWIDTH, + n_knots: int = KERNEL_CDF_DEFAULT_KNOTS, + xp=np, +): + """ + Build the per-axis kernel-density CDF transform pair. + + Drop-in for ``rectangular.create_transforms``: the returned ``transform`` + maps (scaled) source-plane coordinates into the unit square and + ``inv_transform`` maps unit-square coordinates back to the (scaled) + source plane. + + Parameters + ---------- + traced_points + The (N, 2) scaled source-plane coordinates the CDF adapts to. + mesh_pixels + Mesh pixels per axis; sets the default bandwidth scale + ``h = bandwidth * data_span / mesh_pixels`` per axis. + mesh_weight_map + Optional (N,) weights from the adapt image (image adaption). ``None`` + gives uniform weights (density adaption). + bandwidth + Bandwidth in units of the mesh pixel scale (data span / mesh_pixels). + Smaller values track the point density more sharply (approaching the + empirical CDF and its staircase); larger values smooth the mesh + geometry towards uniform. + n_knots + Size of the fixed knot table used to invert the CDF. + xp + The array library to use (numpy or jax.numpy). + """ + points = traced_points + N = points.shape[0] + + if mesh_weight_map is None: + w = xp.full((N,), 1.0 / N) + else: + w = mesh_weight_map / xp.sum(mesh_weight_map) + + lo = xp.min(points, axis=0) + hi = xp.max(points, axis=0) + span = hi - lo + h = bandwidth * span / mesh_pixels + + def F_raw(q): + t = (q[:, None, :] - points[None, :, :]) / h[None, None, :] + return xp.sum(w[None, :, None] * _norm_cdf(t, xp), axis=1) + + # The unit square maps onto the data bounding box exactly, matching the + # linear variant's convention (its empirical-CDF knots end at the extreme + # points); the kernel tails outside [lo, hi] are absorbed by the rescale. + u = xp.linspace(0.0, 1.0, n_knots) + knots = lo[None, :] + u[:, None] * span[None, :] + + F_knots_raw = F_raw(knots) + F_lo = F_knots_raw[0] + F_hi = F_knots_raw[-1] + denom = F_hi - F_lo + + F_knots = (F_knots_raw - F_lo[None, :]) / denom[None, :] + + def transform(q): + F_q = (F_raw(q) - F_lo[None, :]) / denom[None, :] + return xp.clip(F_q, 0.0, 1.0) + + if xp.__name__.startswith("jax"): + inv_transform = partial(reverse_interp, F_knots, knots) + return transform, inv_transform + + inv_transform = partial(reverse_interp_np, F_knots, knots) + return transform, inv_transform + + +def adaptive_rectangular_transformed_grid_from_kernel( + data_grid, + grid, + mesh_pixels: int, + mesh_weight_map=None, + bandwidth: float = KERNEL_CDF_DEFAULT_BANDWIDTH, + n_knots: int = KERNEL_CDF_DEFAULT_KNOTS, + xp=np, +): + """Kernel-CDF version of ``rectangular.adaptive_rectangular_transformed_grid_from``.""" + mu = data_grid.mean(axis=0) + scale = data_grid.std(axis=0).min() + source_grid_scaled = (data_grid - mu) / scale + + _, inv_transform = create_transforms_kernel( + source_grid_scaled, + mesh_pixels=mesh_pixels, + mesh_weight_map=mesh_weight_map, + bandwidth=bandwidth, + n_knots=n_knots, + xp=xp, + ) + + def inv_full(U): + return inv_transform(U) * scale + mu + + return inv_full(grid) + + +def adaptive_rectangular_areas_from_kernel( + source_grid_shape, + data_grid, + mesh_weight_map=None, + bandwidth: float = KERNEL_CDF_DEFAULT_BANDWIDTH, + n_knots: int = KERNEL_CDF_DEFAULT_KNOTS, + xp=np, +): + """Kernel-CDF version of ``mesh_geometry.rectangular.adaptive_rectangular_areas_from``.""" + edges_y = xp.linspace(1, 0, source_grid_shape[0] + 1) + edges_x = xp.linspace(0, 1, source_grid_shape[1] + 1) + + mu = data_grid.mean(axis=0) + scale = data_grid.std(axis=0).min() + source_grid_scaled = (data_grid - mu) / scale + + _, inv_transform = create_transforms_kernel( + source_grid_scaled, + mesh_pixels=source_grid_shape[0], + mesh_weight_map=mesh_weight_map, + bandwidth=bandwidth, + n_knots=n_knots, + xp=xp, + ) + + def inv_full(U): + return inv_transform(U) * scale + mu + + pixel_edges = inv_full(xp.stack([edges_y, edges_x]).T) + pixel_lengths = xp.diff(pixel_edges, axis=0).squeeze() + + dy = pixel_lengths[:, 0] + dx = pixel_lengths[:, 1] + + return xp.abs(xp.outer(dy, dx).flatten()) + + +def adaptive_rectangular_mappings_weights_via_interpolation_from_kernel( + source_grid_size: int, + data_grid, + data_grid_over_sampled, + mesh_weight_map=None, + bandwidth: float = KERNEL_CDF_DEFAULT_BANDWIDTH, + n_knots: int = KERNEL_CDF_DEFAULT_KNOTS, + xp=np, +): + """Kernel-CDF version of the linear helper in ``rectangular.py``. + + Steps 1–2 build the kernel transforms. Steps 3–7 (floor/ceil, flatten, + bilinear weights) are identical to the linear path — copied here to keep + the variants independently auditable. + """ + mu = data_grid.mean(axis=0) + scale = data_grid.std(axis=0).min() + source_grid_scaled = (data_grid - mu) / scale + + transform, _ = create_transforms_kernel( + source_grid_scaled, + mesh_pixels=source_grid_size, + mesh_weight_map=mesh_weight_map, + bandwidth=bandwidth, + n_knots=n_knots, + xp=xp, + ) + + grid_over_sampled_scaled = (data_grid_over_sampled - mu) / scale + grid_over_sampled_transformed = transform(grid_over_sampled_scaled) + grid_over_index = (source_grid_size - 3) * grid_over_sampled_transformed + 1 + + 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]) + + idx_tl = xp.stack([ix_up, iy_down], axis=1) + idx_tr = xp.stack([ix_up, iy_up], axis=1) + idx_br = xp.stack([ix_down, iy_up], axis=1) + idx_bl = xp.stack([ix_down, iy_down], axis=1) + + def flatten(idx, n): + return (n - idx[:, 0]) * n + idx[:, 1] + + flat_tl = flatten(idx_tl, source_grid_size) + flat_tr = flatten(idx_tr, source_grid_size) + flat_bl = flatten(idx_bl, source_grid_size) + flat_br = flatten(idx_br, source_grid_size) + + flat_indices = xp.stack([flat_tl, flat_tr, flat_bl, flat_br], axis=1).astype( + "int64" + ) + + 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 + weights = xp.stack([w_tl, w_tr, w_bl, w_br], axis=1) + + return flat_indices, weights + + +class InterpolatorRectangularKernel(InterpolatorRectangular): + """Kernel-density-CDF adaptive rectangular interpolator. + + Subclasses :class:`InterpolatorRectangular` so that existing + ``isinstance(..., InterpolatorRectangular)`` dispatch sites (e.g. + :mod:`autoarray.plot.inversion`) treat it as an adaptive rectangular + interpolator, and the source-plane mesh reconstruction renders through + the same ``pcolormesh`` path. + """ + + def __init__( + self, + mesh, + mesh_grid, + data_grid, + mesh_weight_map, + adapt_data: Optional[np.ndarray] = None, + bandwidth: float = KERNEL_CDF_DEFAULT_BANDWIDTH, + n_knots: int = KERNEL_CDF_DEFAULT_KNOTS, + xp=np, + ): + super().__init__( + mesh=mesh, + mesh_grid=mesh_grid, + data_grid=data_grid, + mesh_weight_map=mesh_weight_map, + adapt_data=adapt_data, + xp=xp, + ) + self.bandwidth = bandwidth + self.n_knots = n_knots + + @cached_property + def mesh_geometry(self): + from autoarray.inversion.mesh.mesh_geometry.rectangular import ( + MeshGeometryRectangular, + ) + + return MeshGeometryRectangular( + mesh=self.mesh, + mesh_grid=self.mesh_grid, + data_grid=self.data_grid, + mesh_weight_map=self.mesh_weight_map, + kernel_bandwidth=self.bandwidth, + kernel_knots=self.n_knots, + xp=self._xp, + ) + + @cached_property + def _mappings_sizes_weights(self): + mappings, weights = ( + adaptive_rectangular_mappings_weights_via_interpolation_from_kernel( + source_grid_size=self.mesh.shape[0], + data_grid=self.data_grid.array, + data_grid_over_sampled=self.data_grid.over_sampled.array, + mesh_weight_map=self.mesh_weight_map, + bandwidth=self.bandwidth, + n_knots=self.n_knots, + xp=self._xp, + ) + ) + + sizes = 4 * self._xp.ones(len(mappings), dtype="int") + + return mappings, sizes, weights diff --git a/autoarray/inversion/mesh/mesh/__init__.py b/autoarray/inversion/mesh/mesh/__init__.py index ec50dd5e..c37b637c 100644 --- a/autoarray/inversion/mesh/mesh/__init__.py +++ b/autoarray/inversion/mesh/mesh/__init__.py @@ -1,6 +1,8 @@ from .abstract import AbstractMesh as Mesh from .rectangular_adapt_density import RectangularAdaptDensity from .rectangular_adapt_image import RectangularAdaptImage +from .rectangular_kernel_adapt_density import RectangularKernelAdaptDensity +from .rectangular_kernel_adapt_image import RectangularKernelAdaptImage from .rectangular_spline_adapt_density import RectangularSplineAdaptDensity from .rectangular_spline_adapt_image import RectangularSplineAdaptImage from .rectangular_rotated_adapt_image import RectangularRotatedAdaptImage diff --git a/autoarray/inversion/mesh/mesh/rectangular_kernel_adapt_density.py b/autoarray/inversion/mesh/mesh/rectangular_kernel_adapt_density.py new file mode 100644 index 00000000..0247b52b --- /dev/null +++ b/autoarray/inversion/mesh/mesh/rectangular_kernel_adapt_density.py @@ -0,0 +1,95 @@ +import numpy as np +from typing import Optional, Tuple + +from autoarray.inversion.mesh.border_relocator import BorderRelocator +from autoarray.inversion.mesh.mesh.rectangular_adapt_density import ( + RectangularAdaptDensity, + overlay_grid_from, +) +from autoarray.inversion.mesh.interpolator.rectangular_kernel import ( + KERNEL_CDF_DEFAULT_BANDWIDTH, + KERNEL_CDF_DEFAULT_KNOTS, +) +from autoarray.structures.grids.irregular_2d import Grid2DIrregular +from autoarray.structures.grids.uniform_2d import Grid2D + + +class RectangularKernelAdaptDensity(RectangularAdaptDensity): + def __init__( + self, + shape: Tuple[int, int] = (3, 3), + bandwidth: float = KERNEL_CDF_DEFAULT_BANDWIDTH, + n_knots: int = KERNEL_CDF_DEFAULT_KNOTS, + ): + """ + Density-adaptive rectangular mesh using a smooth kernel-density CDF + transform in place of the empirical-CDF linear-interp transform used + by `RectangularAdaptDensity`. + + The kernel CDF is strictly monotone by construction, C-infinity in + both the interp queries and the traced-point positions (no ranks, no + sorts), and duplicate-safe. It therefore carries correct smooth + mass/shear gradients in every configuration — including imaging at + pixelization over-sampling 1 and the interferometer sparse path, + where the empirical CDF makes the likelihood exactly + piecewise-constant. + + Parameters + ---------- + shape + The 2D dimensions of the rectangular pixel grid + ``(total_y_pixels, total_x_pixels)``. + bandwidth + Kernel bandwidth in units of the mesh pixel scale + (data span / mesh pixels per axis). Smaller values track the + point density more sharply (approaching the empirical CDF); + larger values smooth the mesh geometry towards uniform. + n_knots + Size of the fixed knot table used to invert the CDF. + """ + super().__init__(shape=shape) + self.bandwidth = float(bandwidth) + self.n_knots = int(n_knots) + + @property + def interpolator_cls(self): + from autoarray.inversion.mesh.interpolator.rectangular_kernel import ( + InterpolatorRectangularKernel, + ) + + return InterpolatorRectangularKernel + + def interpolator_from( + self, + source_plane_data_grid: Grid2D, + source_plane_mesh_grid: Grid2DIrregular, + border_relocator: Optional[BorderRelocator] = None, + adapt_data: np.ndarray = None, + xp=np, + ): + """See ``RectangularAdaptDensity.interpolator_from``; forwards + ``bandwidth`` / ``n_knots`` to the kernel interpolator.""" + relocated_grid = self.relocated_grid_from( + border_relocator=border_relocator, + source_plane_data_grid=source_plane_data_grid, + xp=xp, + ) + + mesh_grid = overlay_grid_from( + shape_native=self.shape, + grid=relocated_grid.over_sampled, + xp=xp, + ) + + mesh_weight_map = self.mesh_weight_map_from(adapt_data=adapt_data, xp=xp) + + return self.interpolator_cls( + mesh=self, + data_grid=relocated_grid, + mesh_grid=Grid2DIrregular(mesh_grid), + mesh_weight_map=mesh_weight_map, + adapt_data=adapt_data, + bandwidth=self.bandwidth, + n_knots=self.n_knots, + xp=xp, + ) diff --git a/autoarray/inversion/mesh/mesh/rectangular_kernel_adapt_image.py b/autoarray/inversion/mesh/mesh/rectangular_kernel_adapt_image.py new file mode 100644 index 00000000..fda3a842 --- /dev/null +++ b/autoarray/inversion/mesh/mesh/rectangular_kernel_adapt_image.py @@ -0,0 +1,98 @@ +import numpy as np +from typing import Optional, Tuple + +from autoarray.inversion.mesh.border_relocator import BorderRelocator +from autoarray.inversion.mesh.mesh.rectangular_adapt_density import overlay_grid_from +from autoarray.inversion.mesh.mesh.rectangular_adapt_image import ( + RectangularAdaptImage, +) +from autoarray.inversion.mesh.interpolator.rectangular_kernel import ( + KERNEL_CDF_DEFAULT_BANDWIDTH, + KERNEL_CDF_DEFAULT_KNOTS, +) +from autoarray.structures.grids.irregular_2d import Grid2DIrregular +from autoarray.structures.grids.uniform_2d import Grid2D + + +class RectangularKernelAdaptImage(RectangularAdaptImage): + def __init__( + self, + shape: Tuple[int, int] = (3, 3), + weight_power: float = 1.0, + weight_floor: float = 0.0, + bandwidth: float = KERNEL_CDF_DEFAULT_BANDWIDTH, + n_knots: int = KERNEL_CDF_DEFAULT_KNOTS, + ): + """ + Image-adaptive rectangular mesh using a smooth kernel-density CDF + transform in place of the empirical-CDF linear-interp transform used + by `RectangularAdaptImage`. + + Inherits the adapt-image-driven ``mesh_weight_map_from`` from its + parent (``weight_power`` / ``weight_floor`` control); the adapt-image + weights become the kernel weights ``w_i``. The kernel CDF is strictly + monotone by construction, C-infinity in queries and point positions + (no ranks, no sorts), and duplicate-safe, so mass/shear gradients are + smooth in every configuration. + + Parameters + ---------- + shape + The 2D dimensions of the rectangular pixel grid. + weight_power + Exponent applied to the adapt-image weights. + weight_floor + Minimum weight applied to prevent low-intensity pixels from + becoming unconstrained. + bandwidth + Kernel bandwidth in units of the mesh pixel scale + (data span / mesh pixels per axis). + n_knots + Size of the fixed knot table used to invert the CDF. + """ + super().__init__( + shape=shape, weight_power=weight_power, weight_floor=weight_floor + ) + self.bandwidth = float(bandwidth) + self.n_knots = int(n_knots) + + @property + def interpolator_cls(self): + from autoarray.inversion.mesh.interpolator.rectangular_kernel import ( + InterpolatorRectangularKernel, + ) + + return InterpolatorRectangularKernel + + def interpolator_from( + self, + source_plane_data_grid: Grid2D, + source_plane_mesh_grid: Grid2DIrregular, + border_relocator: Optional[BorderRelocator] = None, + adapt_data: np.ndarray = None, + xp=np, + ): + relocated_grid = self.relocated_grid_from( + border_relocator=border_relocator, + source_plane_data_grid=source_plane_data_grid, + xp=xp, + ) + + mesh_grid = overlay_grid_from( + shape_native=self.shape, + grid=relocated_grid.over_sampled, + xp=xp, + ) + + mesh_weight_map = self.mesh_weight_map_from(adapt_data=adapt_data, xp=xp) + + return self.interpolator_cls( + mesh=self, + data_grid=relocated_grid, + mesh_grid=Grid2DIrregular(mesh_grid), + mesh_weight_map=mesh_weight_map, + adapt_data=adapt_data, + bandwidth=self.bandwidth, + n_knots=self.n_knots, + xp=xp, + ) diff --git a/autoarray/inversion/mesh/mesh_geometry/abstract.py b/autoarray/inversion/mesh/mesh_geometry/abstract.py index 8ac7f1dc..d0addc69 100644 --- a/autoarray/inversion/mesh/mesh_geometry/abstract.py +++ b/autoarray/inversion/mesh/mesh_geometry/abstract.py @@ -10,6 +10,8 @@ def __init__( data_grid, mesh_weight_map=None, spline_deg=None, + kernel_bandwidth=None, + kernel_knots=None, xp=np, ): @@ -20,6 +22,10 @@ def __init__( # When non-None, rectangular geometry uses the spline-CDF helpers # instead of the linear-interp CDF (areas / edges transforms only). self.spline_deg = spline_deg + # When non-None, rectangular geometry uses the kernel-density-CDF + # helpers instead of the linear-interp CDF (areas / edges transforms). + self.kernel_bandwidth = kernel_bandwidth + self.kernel_knots = kernel_knots self._use_jax = xp is not np @property diff --git a/autoarray/inversion/mesh/mesh_geometry/rectangular.py b/autoarray/inversion/mesh/mesh_geometry/rectangular.py index f5c6fb25..64e61a8e 100644 --- a/autoarray/inversion/mesh/mesh_geometry/rectangular.py +++ b/autoarray/inversion/mesh/mesh_geometry/rectangular.py @@ -473,7 +473,9 @@ def areas_transformed(self): When ``spline_deg`` is set (spline-CDF meshes), routes through the polynomial+Hermite-spline variant to keep the areas consistent with - the mapper's CDF choice. Otherwise uses the linear-interp CDF. + the mapper's CDF choice; when ``kernel_bandwidth`` is set + (kernel-CDF meshes), through the kernel-density variant. Otherwise + uses the linear-interp CDF. """ if self.spline_deg is not None: from autoarray.inversion.mesh.interpolator.rectangular_spline import ( @@ -488,6 +490,20 @@ def areas_transformed(self): xp=self._xp, ) + if self.kernel_bandwidth is not None: + from autoarray.inversion.mesh.interpolator.rectangular_kernel import ( + adaptive_rectangular_areas_from_kernel, + ) + + return adaptive_rectangular_areas_from_kernel( + source_grid_shape=self.shape_native, + data_grid=self.data_grid.over_sampled, + mesh_weight_map=self.mesh_weight_map, + bandwidth=self.kernel_bandwidth, + n_knots=self.kernel_knots, + xp=self._xp, + ) + return adaptive_rectangular_areas_from( source_grid_shape=self.shape_native, data_grid=self.data_grid.over_sampled, @@ -510,16 +526,29 @@ def areas_for_magnification(self): @property def edges_transformed(self): """ - A class packing the ndarrays describing the neighbors of every pixel in the rectangular pixelization (see - `Neighbors` for a complete description of the neighboring scheme). - - The neighbors of a rectangular pixelization are computed by exploiting the uniform and symmetric nature of the - rectangular grid, as described in the method `rectangular_neighbors_from`. + The source-plane cell edges of every mesh pixel, transformed through + the adaptive CDF — one more edge than pixels per axis, shape + ``(n + 1, 2)`` packing (y_edges, x_edges). + + The interpolation mapper is node-based: reconstruction value + ``(row r, col c)`` lives at node ``U_y = (n_y - r - 1)/(n_y - 3)``, + ``U_x = (c - 1)/(n_x - 3)`` in unit CDF space (a guard ring occupies + the border; rows are flipped by the mapper's ``row = n - i``). The + cell edges are therefore the node midpoints — a uniform ``[0, 1]`` + partition drew every cell up to ~1.5 mesh pixels away from where the + mapper scatters its flux (issue #372). Guard-node edges fall outside + the CDF domain and clamp to the data span in the inverse interp, + squashing the border guard cells to the region they actually cover. """ - # edges defined in 0 -> 1 space, there is one more edge than pixel centers on each side - edges_y = self._xp.linspace(1, 0, self.shape_native[0] + 1) - edges_x = self._xp.linspace(0, 1, self.shape_native[1] + 1) + # Node-midpoint edges in unit CDF space (see docstring); the guard + # denominator keeps the degenerate n <= 3 meshes finite (their + # interpolator maps every point to the single interior node anyway). + n_y, n_x = self.shape_native + rows = self._xp.arange(n_y + 1) + cols = self._xp.arange(n_x + 1) + edges_y = (n_y - rows - 0.5) / max(n_y - 3, 1) + edges_x = (cols - 1.5) / max(n_x - 3, 1) edges_reshaped = self._xp.stack([edges_y, edges_x]).T @@ -536,6 +565,21 @@ def edges_transformed(self): xp=self._xp, ) + if self.kernel_bandwidth is not None: + from autoarray.inversion.mesh.interpolator.rectangular_kernel import ( + adaptive_rectangular_transformed_grid_from_kernel, + ) + + return adaptive_rectangular_transformed_grid_from_kernel( + data_grid=self.data_grid.array, + grid=edges_reshaped, + mesh_pixels=self.shape_native[0], + mesh_weight_map=self.mesh_weight_map, + bandwidth=self.kernel_bandwidth, + n_knots=self.kernel_knots, + xp=self._xp, + ) + from autoarray.inversion.mesh.interpolator.rectangular import ( adaptive_rectangular_transformed_grid_from, ) diff --git a/test_autoarray/inversion/pixelization/mesh/test_rectangular_kernel.py b/test_autoarray/inversion/pixelization/mesh/test_rectangular_kernel.py new file mode 100644 index 00000000..30dbedd4 --- /dev/null +++ b/test_autoarray/inversion/pixelization/mesh/test_rectangular_kernel.py @@ -0,0 +1,340 @@ +"""Unit tests for the kernel-density-CDF rectangular meshes. + +Pure numpy — no JAX imports here. Cross-xp / gradient certification lives in +autolens_workspace_test/scripts/jax_grad per the project's "no JAX in unit +tests" rule. +""" +import numpy as np +import pytest + +import autoarray as aa +from autoarray.inversion.mesh.interpolator.rectangular_kernel import ( + KERNEL_CDF_DEFAULT_BANDWIDTH, + KERNEL_CDF_DEFAULT_KNOTS, + InterpolatorRectangularKernel, + adaptive_rectangular_areas_from_kernel, + adaptive_rectangular_mappings_weights_via_interpolation_from_kernel, + create_transforms_kernel, +) +from autoarray.inversion.mesh.interpolator.rectangular import ( + adaptive_rectangular_mappings_weights_via_interpolation_from, +) +from autoarray.inversion.mesh.mesh_geometry.rectangular import ( + adaptive_rectangular_areas_from, +) + + +# --------------------------------------------------------------------------- +# Construction +# --------------------------------------------------------------------------- + + +def test__construction__shape_and_kernel_kwargs__default(): + density = aa.mesh.RectangularKernelAdaptDensity(shape=(5, 7)) + image = aa.mesh.RectangularKernelAdaptImage(shape=(5, 7)) + + assert density.shape == (5, 7) + assert density.bandwidth == KERNEL_CDF_DEFAULT_BANDWIDTH + assert density.n_knots == KERNEL_CDF_DEFAULT_KNOTS + assert image.shape == (5, 7) + assert image.bandwidth == KERNEL_CDF_DEFAULT_BANDWIDTH + assert image.n_knots == KERNEL_CDF_DEFAULT_KNOTS + # inherited + assert image.weight_power == 1.0 + assert image.weight_floor == 0.0 + + +def test__construction__kernel_kwargs__overridden(): + density = aa.mesh.RectangularKernelAdaptDensity( + shape=(3, 3), bandwidth=0.5, n_knots=128 + ) + image = aa.mesh.RectangularKernelAdaptImage( + shape=(3, 3), weight_power=2.0, weight_floor=0.1, bandwidth=2.0, n_knots=32 + ) + + assert density.bandwidth == 0.5 + assert density.n_knots == 128 + assert image.bandwidth == 2.0 + assert image.n_knots == 32 + assert image.weight_power == 2.0 + assert image.weight_floor == 0.1 + + +def test__construction__minimum_shape_raises_same_as_parent(): + with pytest.raises(aa.exc.MeshException): + aa.mesh.RectangularKernelAdaptDensity(shape=(2, 3)) + with pytest.raises(aa.exc.MeshException): + aa.mesh.RectangularKernelAdaptImage(shape=(3, 2)) + + +# --------------------------------------------------------------------------- +# Interpolator dispatch +# --------------------------------------------------------------------------- + + +def test__interpolator_cls__is_InterpolatorRectangularKernel(): + density = aa.mesh.RectangularKernelAdaptDensity(shape=(3, 3)) + image = aa.mesh.RectangularKernelAdaptImage(shape=(3, 3)) + + assert density.interpolator_cls is InterpolatorRectangularKernel + assert image.interpolator_cls is InterpolatorRectangularKernel + + +# --------------------------------------------------------------------------- +# mesh_weight_map_from inheritance +# --------------------------------------------------------------------------- + + +def test__mesh_weight_map_from__density__returns_none(): + density = aa.mesh.RectangularKernelAdaptDensity(shape=(3, 3)) + assert density.mesh_weight_map_from(adapt_data=None) is None + + +def test__mesh_weight_map_from__image__returns_weighted_normalized(): + image = aa.mesh.RectangularKernelAdaptImage( + shape=(3, 3), weight_power=2.0, weight_floor=0.0 + ) + + class _Stub: + def __init__(self, arr): + self.array = arr + + adapt = _Stub(np.array([1.0, 2.0, 4.0, 0.0, 8.0])) + w = image.mesh_weight_map_from(adapt_data=adapt) + + expected = np.array([1.0, 4.0, 16.0, 1e-24, 64.0]) + expected = expected / expected.sum() + assert w == pytest.approx(expected, rel=1e-6) + + +# --------------------------------------------------------------------------- +# CDF transform properties +# --------------------------------------------------------------------------- + + +def _seeded_inputs(M=128, K=400, seed=0): + rng = np.random.default_rng(seed) + data_grid = rng.standard_normal((M, 2)) + data_grid_over = rng.standard_normal((K, 2)) * 0.8 + weights = rng.uniform(0.1, 1.0, size=M) + weights = weights / weights.sum() + return data_grid, data_grid_over, weights + + +def test__create_transforms_kernel__strictly_monotone_in_queries(): + data_grid, _, weights = _seeded_inputs(seed=1) + + fwd, _ = create_transforms_kernel( + data_grid, mesh_pixels=16, mesh_weight_map=weights, xp=np + ) + + # Dense interior queries per axis, strictly increasing. + q = np.linspace(data_grid.min(axis=0), data_grid.max(axis=0), 500) + F = fwd(q) + assert np.all(np.diff(F, axis=0) > 0.0) + + +def test__create_transforms_kernel__fwd_maps_data_range_to_unit_square(): + data_grid, _, weights = _seeded_inputs(seed=3) + + fwd, _ = create_transforms_kernel( + data_grid, mesh_pixels=16, mesh_weight_map=weights, xp=np + ) + + F = fwd(data_grid) + assert F.min() >= 0.0 + assert F.max() <= 1.0 + + # The data bounding box corners map exactly onto the unit square corners. + lo = data_grid.min(axis=0) + hi = data_grid.max(axis=0) + corners = fwd(np.stack([lo, hi])) + assert corners[0] == pytest.approx([0.0, 0.0], abs=1e-12) + assert corners[1] == pytest.approx([1.0, 1.0], abs=1e-12) + + +def test__create_transforms_kernel__roundtrip_matches_identity(): + data_grid, _, weights = _seeded_inputs(seed=1) + + fwd, rev = create_transforms_kernel( + data_grid, mesh_pixels=16, mesh_weight_map=weights, n_knots=512, xp=np + ) + + probe = np.array([[0.1, 0.1], [0.5, 0.5], [0.9, 0.9]]) + roundtrip = fwd(rev(probe)) + assert roundtrip == pytest.approx(probe, abs=1e-4) + + +def test__create_transforms_kernel__unweighted_roundtrip_default_knots(): + data_grid, _, _ = _seeded_inputs(seed=2) + + fwd, rev = create_transforms_kernel(data_grid, mesh_pixels=16, xp=np) + + probe = np.array([[0.25, 0.3], [0.5, 0.5], [0.7, 0.75]]) + roundtrip = fwd(rev(probe)) + assert roundtrip == pytest.approx(probe, abs=2e-3) + + +def test__create_transforms_kernel__duplicate_points_are_safe(): + """Coincident traced points must not produce NaN/inf or break + monotonicity — the empirical CDF's 1/Δknot failure mode.""" + rng = np.random.default_rng(4) + base = rng.standard_normal((32, 2)) + data_grid = np.concatenate([base, base[:8], base[:1].repeat(4, axis=0)]) + + fwd, rev = create_transforms_kernel(data_grid, mesh_pixels=8, xp=np) + + q = np.linspace(data_grid.min(axis=0), data_grid.max(axis=0), 200) + F = fwd(q) + assert np.all(np.isfinite(F)) + assert np.all(np.diff(F, axis=0) > 0.0) + + probe = np.array([[0.2, 0.4], [0.6, 0.8]]) + assert np.all(np.isfinite(rev(probe))) + + +def test__create_transforms_kernel__large_bandwidth_tends_to_uniform(): + """As h → ∞ the kernel CDF linearises, so the transform approaches the + plain affine map of the data bounding box onto the unit square.""" + data_grid, _, _ = _seeded_inputs(seed=5) + + fwd, _ = create_transforms_kernel( + data_grid, mesh_pixels=16, bandwidth=1000.0, xp=np + ) + + lo = data_grid.min(axis=0) + hi = data_grid.max(axis=0) + q = np.linspace(lo, hi, 50) + expected = (q - lo) / (hi - lo) + assert fwd(q) == pytest.approx(expected, abs=1e-4) + + +# --------------------------------------------------------------------------- +# Mapper output shape + content consistency with linear variant +# --------------------------------------------------------------------------- + + +def test__mappings_sizes_weights__shapes_match_linear(): + data_grid, over, weights = _seeded_inputs() + + idx_k, w_k = adaptive_rectangular_mappings_weights_via_interpolation_from_kernel( + source_grid_size=16, + data_grid=data_grid, + data_grid_over_sampled=over, + mesh_weight_map=weights, + xp=np, + ) + idx_l, w_l = adaptive_rectangular_mappings_weights_via_interpolation_from( + source_grid_size=16, + data_grid=data_grid, + data_grid_over_sampled=over, + mesh_weight_map=weights, + xp=np, + ) + + assert idx_k.shape == idx_l.shape == (400, 4) + assert w_k.shape == w_l.shape == (400, 4) + assert np.allclose(w_k.sum(axis=1), 1.0, atol=1e-10) + assert np.allclose(w_l.sum(axis=1), 1.0, atol=1e-10) + + +def test__mappings__flat_indices_mostly_match_linear(): + """Kernel and linear CDF route most points to the same cell — only points + near cell boundaries can differ. Assert at least 70% exact match.""" + data_grid, over, weights = _seeded_inputs() + + idx_k, _ = adaptive_rectangular_mappings_weights_via_interpolation_from_kernel( + source_grid_size=16, + data_grid=data_grid, + data_grid_over_sampled=over, + mesh_weight_map=weights, + xp=np, + ) + idx_l, _ = adaptive_rectangular_mappings_weights_via_interpolation_from( + source_grid_size=16, + data_grid=data_grid, + data_grid_over_sampled=over, + mesh_weight_map=weights, + xp=np, + ) + + exact_match_frac = float((idx_k == idx_l).mean()) + assert exact_match_frac >= 0.70, ( + f"kernel vs linear flat_index exact-match fraction = {exact_match_frac:.3f} " + "(expected ≥ 0.70)" + ) + + +# --------------------------------------------------------------------------- +# Areas +# --------------------------------------------------------------------------- + + +def test__areas_from_kernel__positive_finite__total_matches_linear(): + data_grid, _, weights = _seeded_inputs(seed=6) + + areas_k = adaptive_rectangular_areas_from_kernel( + source_grid_shape=(12, 12), + data_grid=data_grid, + mesh_weight_map=weights, + xp=np, + ) + areas_l = adaptive_rectangular_areas_from( + source_grid_shape=(12, 12), + data_grid=data_grid, + mesh_weight_map=weights, + xp=np, + ) + + assert areas_k.shape == (144,) + assert np.all(np.isfinite(areas_k)) + assert np.all(areas_k > 0.0) + # Both variants map the unit square onto the data bounding box, so the + # per-axis edge differences telescope to the same span: totals agree. + assert areas_k.sum() == pytest.approx(areas_l.sum(), rel=1e-8) + + +# --------------------------------------------------------------------------- +# Interpolator construction — exercise the class directly without the full +# BorderRelocator / Grid2D pipeline (covered by the jax_grad certification). +# --------------------------------------------------------------------------- + + +def test__InterpolatorRectangularKernel__mappings_sizes_weights_via_property(): + class _StubGrid: + def __init__(self, arr): + self.array = arr + self.over_sampled = self + self._array = arr + + def __getattr__(self, item): + return getattr(self._array, item) + + rng = np.random.default_rng(5) + data_grid = _StubGrid(rng.standard_normal((64, 2))) + + mesh = aa.mesh.RectangularKernelAdaptDensity(shape=(6, 6), bandwidth=0.8) + interpolator = InterpolatorRectangularKernel( + mesh=mesh, + mesh_grid=_StubGrid(rng.standard_normal((36, 2))), + data_grid=data_grid, + mesh_weight_map=None, + bandwidth=0.8, + xp=np, + ) + + mappings, sizes, weights = interpolator._mappings_sizes_weights + assert mappings.shape == (64, 4) + assert sizes.shape == (64,) + assert weights.shape == (64, 4) + assert np.all(sizes == 4) + assert np.allclose(weights.sum(axis=1), 1.0, atol=1e-10) + + geometry = interpolator.mesh_geometry + assert geometry.kernel_bandwidth == 0.8 + assert geometry.kernel_knots == KERNEL_CDF_DEFAULT_KNOTS + + areas = geometry.areas_transformed + assert areas.shape == (36,) + assert np.all(np.isfinite(areas)) + assert np.all(areas > 0.0) diff --git a/test_autoarray/inversion/pixelization/mesh_geometry/test_rectangular.py b/test_autoarray/inversion/pixelization/mesh_geometry/test_rectangular.py index eadeb26e..d04fb74b 100644 --- a/test_autoarray/inversion/pixelization/mesh_geometry/test_rectangular.py +++ b/test_autoarray/inversion/pixelization/mesh_geometry/test_rectangular.py @@ -200,3 +200,53 @@ def test__edges_transformed(mask_2d_7x7): ), abs=1e-8, ) + + +def test__edges_transformed__aligned_with_interpolation_node_convention(): + """ + Regression test for issue #372: the pcolormesh plotting path draws value + (row r, col c) inside the cell bounded by edges_transformed — that cell + must be centred on the interpolation mapper's node for that value, not on + a uniform [0, 1] partition (which shifted plots by ~1.5 mesh pixels). + + A delta function scattered through the mapper must land in plotted cells + whose weighted centroid matches the input point to sub-cell precision. + """ + from autoarray.inversion.mesh.interpolator.rectangular import ( + adaptive_rectangular_mappings_weights_via_interpolation_from, + adaptive_rectangular_transformed_grid_from, + ) + + n = 10 + rng = np.random.default_rng(0) + data_grid = rng.uniform(-1.0, 1.0, (5000, 2)) + test_point = np.array([[0.3, -0.2]]) + + flat_indices, weights = ( + adaptive_rectangular_mappings_weights_via_interpolation_from( + source_grid_size=n, + data_grid=data_grid, + data_grid_over_sampled=test_point, + ) + ) + + # The node-midpoint unit-space edges (what edges_transformed now builds), + # pushed through the same CDF transform. + rows = np.arange(n + 1) + edges_y = (n - rows - 0.5) / (n - 3) + edges_x = (rows - 1.5) / (n - 3) + edges = np.stack([edges_y, edges_x]).T + edges_t = adaptive_rectangular_transformed_grid_from(data_grid, edges) + y_edges, x_edges = edges_t.T + + centroid_y = 0.0 + centroid_x = 0.0 + for flat, weight in zip(flat_indices[0], weights[0]): + r, c = flat // n, flat % n + centroid_y += weight * 0.5 * (y_edges[r] + y_edges[r + 1]) + centroid_x += weight * 0.5 * (x_edges[c] + x_edges[c + 1]) + + # Half a mesh cell in these units is ~0.15; the pre-fix uniform edges + # missed by ~0.4 in y. + assert centroid_y == pytest.approx(test_point[0, 0], abs=0.1) + assert centroid_x == pytest.approx(test_point[0, 1], abs=0.1)