diff --git a/autoarray/exc.py b/autoarray/exc.py index a27366149..8be02eebd 100644 --- a/autoarray/exc.py +++ b/autoarray/exc.py @@ -72,7 +72,7 @@ class MeshException(Exception): """ Raises exceptions associated with the `inversion/mesh` modules and `Mesh` classes. - For example if a `RectangularAdaptDensity` mesh has dimensions below 3x3. + For example if a `RectangularRTUAdaptDensity` mesh has dimensions below 3x3. """ pass diff --git a/autoarray/fixtures.py b/autoarray/fixtures.py index be4d64205..905036075 100644 --- a/autoarray/fixtures.py +++ b/autoarray/fixtures.py @@ -368,7 +368,7 @@ def make_border_relocator_2d_7x7(): def make_rectangular_mapper_7x7_3x3(): - from autoarray.inversion.mesh.mesh.rectangular_adapt_density import ( + from autoarray.inversion.mesh.mesh.rectangular_rtu_adapt_density import ( overlay_grid_from, ) diff --git a/autoarray/inversion/inversion/abstract.py b/autoarray/inversion/inversion/abstract.py index eb1291ce0..8bc50d2bd 100644 --- a/autoarray/inversion/inversion/abstract.py +++ b/autoarray/inversion/inversion/abstract.py @@ -174,8 +174,8 @@ def cls_list_from(self, cls: Type, cls_filtered: Optional[Type] = None) -> List: - If the input is `cls=aa.mesh.Mesh`, a list containing all pixelizations in the class are returned. - - If `cls=aa.mesh.Mesh` and `cls_filtered=aa.mesh.RectangularAdaptDensity`, a list of all pixelizations - excluding those which are `RectangularAdaptDensity` pixelizations will be returned. + - If `cls=aa.mesh.Mesh` and `cls_filtered=aa.mesh.RectangularRTUAdaptDensity`, a list of all pixelizations + excluding those which are `RectangularRTUAdaptDensity` pixelizations will be returned. Parameters ---------- diff --git a/autoarray/inversion/linear_obj/neighbors.py b/autoarray/inversion/linear_obj/neighbors.py index 6136d4538..0e312e839 100644 --- a/autoarray/inversion/linear_obj/neighbors.py +++ b/autoarray/inversion/linear_obj/neighbors.py @@ -4,14 +4,14 @@ class Neighbors(np.ndarray): def __new__(cls, arr: np.ndarray, sizes: np.ndarray): """ - Class packaging ndarrays describing the neighbors of every pixel in a mesh (e.g. `RectangularAdaptDensity`). + Class packaging ndarrays describing the neighbors of every pixel in a mesh (e.g. `RectangularRTUAdaptDensity`). The array `arr` contains the pixel indexes of the neighbors of every pixel. Its has shape [total_pixels, max_neighbors_in_single_pixel]. The array `sizes` contains the number of neighbors of every pixel in the pixelixzation. - For example, for a 3x3 `RectangularAdaptDensity` grid: + For example, for a 3x3 `RectangularRTUAdaptDensity` grid: - `total_pixels=9` and `max_neighbors_in_single_pixel=4` (because the central pixel has 4 neighbors whereas edge / corner pixels have 3 and 2). diff --git a/autoarray/inversion/mesh/interpolator/delaunay.py b/autoarray/inversion/mesh/interpolator/delaunay.py index 14621c38c..6c751b7d4 100644 --- a/autoarray/inversion/mesh/interpolator/delaunay.py +++ b/autoarray/inversion/mesh/interpolator/delaunay.py @@ -786,10 +786,10 @@ def _mappings_sizes_weights(self): The `sub_slim_index` refers to the masked data sub-pixels and `pix_indexes` the pixelization pixel indexes, for example: - - `pix_indexes_for_sub_slim_index[0, 0] = 2`: The data's first (index 0) sub-pixel maps to the RectangularAdaptDensity + - `pix_indexes_for_sub_slim_index[0, 0] = 2`: The data's first (index 0) sub-pixel maps to the RectangularRTUAdaptDensity pixelization's third (index 2) pixel. - - `pix_indexes_for_sub_slim_index[2, 0] = 4`: The data's third (index 2) sub-pixel maps to the RectangularAdaptDensity + - `pix_indexes_for_sub_slim_index[2, 0] = 4`: The data's third (index 2) sub-pixel maps to the RectangularRTUAdaptDensity pixelization's fifth (index 4) pixel. The second dimension of the array `pix_indexes_for_sub_slim_index`, which is 0 in both examples above, is used diff --git a/autoarray/inversion/mesh/interpolator/rectangular.py b/autoarray/inversion/mesh/interpolator/rectangular.py index 47a01a461..ab16ce294 100644 --- a/autoarray/inversion/mesh/interpolator/rectangular.py +++ b/autoarray/inversion/mesh/interpolator/rectangular.py @@ -1,10 +1,23 @@ """ -The adaptive rectangular interpolator, built on a kernel-density CDF transform. - -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 CDF is a smooth kernel-density CDF (the RTU -formulation of Enzi et al., arXiv:2606.30620): per axis +The adaptive rectangular interpolator, built on a per-axis CDF transform of +the traced source-plane coordinates so that the uniform mesh pixels adapt to +the density (or weight) of the traced points. Two transforms are provided, +selected by the ``transform`` argument threaded through this module: + +**Empirical rank CDF** (``transform="rank"`` — the ``RectangularBilinear*`` +meshes). The per-axis CDF is the empirical (weighted) rank CDF of the traced +points: a sort plus a cumulative sum, linearly interpolated between points. +O(N log N) per likelihood evaluation with no kernel hyperparameters — the +fast, conceptually simple CPU default. The trade-off is smoothness: the +transform is piecewise-linear in the queries and its dependence on the +traced-point *positions* passes through ranks, so with pixelization +over-sampling 1 (and on the interferometer sparse path, which has no +over-sampling) the likelihood is exactly piecewise-constant in mass/shear — +gradients are identically zero. Gradient (JAX sampler) users need +``over_sample_size_pixelization >= 4`` on imaging, or the kernel-CDF meshes. + +**Kernel-density CDF** (``transform="kernel"`` — the ``RectangularRTU*`` +meshes, the RTU formulation of Enzi et al., arXiv:2606.30620): per axis F(x) = sum_i w_i * Phi((x - x_i) / h) @@ -12,23 +25,28 @@ uniform weights (density adaption) or the normalized adapt-image weights (image adaption), and ``h`` a bandwidth tied to the mesh resolution. -Properties, by construction: +Properties of the kernel CDF, by construction: - strictly monotone; - C-infinity in the queries AND the traced-point positions — no ranks, no sorts, no ``argsort`` anywhere, so there is nothing to swap and the likelihood carries smooth mass/shear gradients in every configuration (including pixelization over-sampling 1 and the interferometer sparse - path, where an empirical point-rank CDF is exactly piecewise-constant); + path, where the empirical rank CDF is exactly piecewise-constant); - duplicate-safe: coincident traced points simply stack their weights. -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, 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 +The kernel 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, 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. +constants, gradients flow smoothly through the table values. The exact +kernel sum is the O(M_sub x N_data) erf evaluation that dominates the CPU +likelihood at production scale — the reason the rank CDF is the CPU default. + +Both transforms warp the same lattice; the 4-pixel bilinear interpolation of +``adaptive_rectangular_mappings_weights_via_interpolation_from`` is shared. """ import math @@ -44,6 +62,35 @@ from autoarray.inversion.mesh.interpolator.abstract import AbstractInterpolator +def forward_interp(xp, yp, x): + import jax + import jax.numpy as jnp + + return jax.vmap(jnp.interp, in_axes=(1, 1, 1, None, None), out_axes=(1))( + x, xp, yp, 0, 1 + ) + + +def forward_interp_np(xp, yp, x): + """ + xp: (N, M) + yp: (N, M) + x : (K, M) ← query points per column + """ + + if yp.ndim == 1 and xp.ndim == 2: + yp = np.broadcast_to(yp[:, None], xp.shape) + + K, M = x.shape + + out = np.empty((K, 2), dtype=xp.dtype) + + for j in range(2): + out[:, j] = np.interp(x[:, j], xp[:, j], yp[:, j], left=0, right=1) + + return out + + def reverse_interp(xp, yp, x): import jax import jax.numpy as jnp @@ -136,6 +183,55 @@ def _norm_cdf(t, xp): return 0.5 * (1.0 + erf(t / _SQRT2)) +def create_transforms_rank(traced_points, mesh_weight_map=None, xp=np): + """ + Build the per-axis empirical rank-CDF transform pair (the ``Bilinear`` + meshes' lattice transform). + + 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. + 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. + + 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_weight_map + Optional (N,) weights from the adapt image (image adaption). ``None`` + gives uniform weights (density adaption). + xp + The array library to use (numpy or jax.numpy). + """ + N = traced_points.shape[0] + + if mesh_weight_map is None: + t = xp.arange(1, N + 1) / (N + 1) + t = xp.stack([t, t], axis=1) + sort_points = xp.sort(traced_points, axis=0) + else: + sdx = xp.argsort(traced_points, axis=0) + sort_points = xp.take_along_axis(traced_points, sdx, axis=0) + t = xp.stack([mesh_weight_map, mesh_weight_map], axis=1) + t = xp.take_along_axis(t, sdx, axis=0) + t = xp.cumsum(t, axis=0) + + if xp.__name__.startswith("jax"): + transform = partial(forward_interp, sort_points, t) + inv_transform = partial(reverse_interp, t, sort_points) + return transform, inv_transform + + transform = partial(forward_interp_np, sort_points, t) + inv_transform = partial(reverse_interp_np, t, sort_points) + return transform, inv_transform + + def create_transforms( traced_points, mesh_pixels: int, @@ -253,6 +349,36 @@ def transform(q): return transform, inv_transform +def _transforms_from( + source_grid_scaled, + mesh_pixels: int, + mesh_weight_map, + bandwidth: float, + n_knots: int, + transform: str, + xp, +): + """ + Dispatch to the requested lattice-transform pair: the empirical rank CDF + (``"rank"``, the Bilinear meshes — ``bandwidth`` / ``n_knots`` unused) or + the kernel-density CDF (``"kernel"``, the RTU meshes). + """ + if transform == "rank": + return create_transforms_rank( + source_grid_scaled, mesh_weight_map=mesh_weight_map, xp=xp + ) + if transform != "kernel": + raise ValueError(f"transform must be 'rank' or 'kernel', got {transform!r}.") + return create_transforms( + source_grid_scaled, + mesh_pixels=mesh_pixels, + mesh_weight_map=mesh_weight_map, + bandwidth=bandwidth, + n_knots=n_knots, + xp=xp, + ) + + def adaptive_rectangular_transformed_grid_from( data_grid, grid, @@ -260,19 +386,21 @@ def adaptive_rectangular_transformed_grid_from( mesh_weight_map=None, bandwidth: float = KERNEL_CDF_DEFAULT_BANDWIDTH, n_knots: int = KERNEL_CDF_DEFAULT_KNOTS, + transform: str = "kernel", xp=np, ): - """Map unit-square coordinates back to the source plane via the kernel CDF.""" + """Map unit-square coordinates back to the source plane via the adaptive CDF.""" mu = data_grid.mean(axis=0) scale = data_grid.std(axis=0).min() source_grid_scaled = (data_grid - mu) / scale - _, inv_transform = create_transforms( + _, inv_transform = _transforms_from( source_grid_scaled, mesh_pixels=mesh_pixels, mesh_weight_map=mesh_weight_map, bandwidth=bandwidth, n_knots=n_knots, + transform=transform, xp=xp, ) @@ -288,9 +416,10 @@ def adaptive_rectangular_areas_from( mesh_weight_map=None, bandwidth: float = KERNEL_CDF_DEFAULT_BANDWIDTH, n_knots: int = KERNEL_CDF_DEFAULT_KNOTS, + transform: str = "kernel", xp=np, ): - """The source-plane area of every mesh pixel under the kernel-CDF transform.""" + """The source-plane area of every mesh pixel under the adaptive CDF transform.""" edges_y = xp.linspace(1, 0, source_grid_shape[0] + 1) edges_x = xp.linspace(0, 1, source_grid_shape[1] + 1) @@ -298,12 +427,13 @@ def adaptive_rectangular_areas_from( scale = data_grid.std(axis=0).min() source_grid_scaled = (data_grid - mu) / scale - _, inv_transform = create_transforms( + _, inv_transform = _transforms_from( source_grid_scaled, mesh_pixels=source_grid_shape[0], mesh_weight_map=mesh_weight_map, bandwidth=bandwidth, n_knots=n_knots, + transform=transform, xp=xp, ) @@ -326,6 +456,7 @@ def adaptive_rectangular_mappings_weights_via_interpolation_from( mesh_weight_map=None, bandwidth: float = KERNEL_CDF_DEFAULT_BANDWIDTH, n_knots: int = KERNEL_CDF_DEFAULT_KNOTS, + transform: str = "kernel", xp=np, ): """ @@ -347,8 +478,9 @@ def adaptive_rectangular_mappings_weights_via_interpolation_from( Steps performed: 1. Normalize the source-plane grid by subtracting its mean and dividing by the minimum axis standard deviation (to balance scaling). - 2. Construct the kernel-CDF forward/inverse transforms which map the grid - into the unit square [0,1]^2. + 2. Construct the adaptive-CDF forward/inverse transforms (rank or + kernel, per ``transform``) which map the grid into the unit square + [0,1]^2. 3. Transform the oversampled source-plane grid into [0,1]^2, then scale it to index space `[0, source_grid_size)`. 4. Compute floor/ceil along x and y axes to find the enclosing rectangular cell. @@ -373,9 +505,14 @@ def adaptive_rectangular_mappings_weights_via_interpolation_from( The weight map used to weight the creation of the rectangular mesh grid, which adapts the size of the mesh pixels to where the source is reconstructed. bandwidth - Bandwidth in units of the mesh pixel scale (see ``create_transforms``). + Bandwidth in units of the mesh pixel scale (see ``create_transforms``; + kernel transform only). n_knots - Size of the fixed knot table used to invert the CDF. + Size of the fixed knot table used to invert the CDF (kernel transform + only). + transform + The lattice transform: ``"rank"`` (empirical rank CDF, the Bilinear + meshes) or ``"kernel"`` (kernel-density CDF, the RTU meshes). Returns ------- @@ -392,18 +529,19 @@ def adaptive_rectangular_mappings_weights_via_interpolation_from( source_grid_scaled = (data_grid - mu) / scale # --- Step 2. Build transforms --- - transform, _ = create_transforms( + transform_func, _ = _transforms_from( source_grid_scaled, mesh_pixels=source_grid_size, mesh_weight_map=mesh_weight_map, bandwidth=bandwidth, n_knots=n_knots, + transform=transform, xp=xp, ) # --- Step 3. Transform oversampled grid into index space --- grid_over_sampled_scaled = (data_grid_over_sampled - mu) / scale - grid_over_sampled_transformed = transform(grid_over_sampled_scaled) + 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 --- @@ -445,13 +583,15 @@ def flatten(idx, n): class InterpolatorRectangular(AbstractInterpolator): - """Adaptive rectangular interpolator (kernel-density CDF). + """Adaptive rectangular interpolator (rank or kernel-density CDF). A grid of (y,x) coordinates which represent an adaptive rectangular pixelization, ordered such that pixels begin from the top-row and go rightwards and then downwards. The mesh pixels adapt to the density (or - adapt-image weights) of the traced points through the smooth kernel-CDF - transform of ``create_transforms``. + adapt-image weights) of the traced points through the per-axis CDF + transform selected by ``transform``: the empirical rank CDF of + ``create_transforms_rank`` (the Bilinear meshes) or the smooth kernel + CDF of ``create_transforms`` (the RTU meshes). This class is used in conjunction with the `inversion/pixelizations` package to create rectangular pixelizations and mappers that perform an @@ -467,6 +607,7 @@ def __init__( adapt_data: Optional[np.ndarray] = None, bandwidth: float = KERNEL_CDF_DEFAULT_BANDWIDTH, n_knots: int = KERNEL_CDF_DEFAULT_KNOTS, + transform: str = "kernel", xp=np, ): super().__init__( @@ -479,6 +620,7 @@ def __init__( self.mesh_weight_map = mesh_weight_map self.bandwidth = bandwidth self.n_knots = n_knots + self.transform = transform @cached_property def mesh_geometry(self): @@ -493,6 +635,7 @@ def mesh_geometry(self): mesh_weight_map=self.mesh_weight_map, kernel_bandwidth=self.bandwidth, kernel_knots=self.n_knots, + transform=self.transform, xp=self._xp, ) @@ -506,6 +649,7 @@ def _mappings_sizes_weights(self): mesh_weight_map=self.mesh_weight_map, bandwidth=self.bandwidth, n_knots=self.n_knots, + transform=self.transform, xp=self._xp, ) ) diff --git a/autoarray/inversion/mesh/mesh/__init__.py b/autoarray/inversion/mesh/mesh/__init__.py index 808df01d2..492587cb6 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_bilinear_adapt_density import RectangularBilinearAdaptDensity +from .rectangular_bilinear_adapt_image import RectangularBilinearAdaptImage +from .rectangular_rtu_adapt_density import RectangularRTUAdaptDensity +from .rectangular_rtu_adapt_image import RectangularRTUAdaptImage from .rectangular_uniform import RectangularUniform from .delaunay import Delaunay from .delaunay_nn import DelaunayNN diff --git a/autoarray/inversion/mesh/mesh/rectangular_bilinear_adapt_density.py b/autoarray/inversion/mesh/mesh/rectangular_bilinear_adapt_density.py new file mode 100644 index 000000000..5087afa44 --- /dev/null +++ b/autoarray/inversion/mesh/mesh/rectangular_bilinear_adapt_density.py @@ -0,0 +1,83 @@ +from typing import Tuple + +from autoarray.inversion.mesh.mesh.rectangular_rtu_adapt_density import ( + RectangularRTUAdaptDensity, +) + + +class RectangularBilinearAdaptDensity(RectangularRTUAdaptDensity): + + def __init__( + self, + shape: Tuple[int, int] = (3, 3), + ): + """ + A rectangular mesh of pixels used to reconstruct a source on a regular + grid, whose pixels adapt to the density of the traced source-plane + coordinates through the empirical rank-CDF transform. + + The mesh is defined by a 2D shape `(total_y_pixels, total_x_pixels)` and + is indexed in row-major order: + + - Index 0 corresponds to the top-left pixel. + - Indices increase from left to right across each row, + and from top to bottom across rows. + + Adaptive behaviour + ------------------ + The mesh adapts to the spatial density of the traced points through the + per-axis empirical (rank) CDF of the traced coordinates — a sort plus a + cumulative sum, linearly interpolated between points: mesh pixels + shrink where many coordinates land (e.g. regions of high magnification + in gravitational lensing) and grow where sampling is sparse. The + inversion therefore achieves higher effective resolution in these + regions without changing the fixed rectangular topology. + + The transform is conceptually simple (no kernel hyperparameters) and + costs O(N log N) per likelihood evaluation, making this the fast + default rectangular mesh for CPU fitting. + + When to use Bilinear vs RTU + --------------------------- + The rank CDF depends on the traced-point positions through their + ranks, so with pixelization over-sampling 1 the likelihood is exactly + piecewise-constant in mass/shear — gradients are identically zero. + Gradient-based (JAX) samplers therefore need + `over_sample_size_pixelization >= 4` on imaging data, or the + `RectangularRTUAdaptDensity` / `RectangularRTUAdaptImage` meshes + (smooth kernel-CDF transform, correct gradients in every + configuration). The interferometer sparse path has no over-sampling, + so gradient work there must use the RTU meshes (or + `RectangularUniform`). The RTU meshes are also the recommended option + on GPU, where their kernel-CDF cost is not the bottleneck. + + Edge handling + ------------- + Boundary (edge) pixels are automatically identified through the mesh + neighbour structure. These edge pixels may be internally excluded + (zeroed) during inversion to improve numerical stability and reduce + edge artefacts. This zeroing is determined by the mesh connectivity + and does not require manual specification of boundary indices. + + Parameters + ---------- + shape : Tuple[int, int] + The 2D dimensions of the rectangular pixel grid + `(total_y_pixels, total_x_pixels)`. + + Raises + ------ + MeshException + If either dimension is less than 3, as a minimum of 3×3 pixels + is required to define interior and boundary structure. + """ + super().__init__(shape=shape) + + @property + def interpolator_kwargs(self) -> dict: + """ + Extra keyword arguments `interpolator_from` forwards to + `interpolator_cls` — the rank-CDF transform selector; the empirical + rank CDF has no kernel hyperparameters. + """ + return {"transform": "rank"} diff --git a/autoarray/inversion/mesh/mesh/rectangular_bilinear_adapt_image.py b/autoarray/inversion/mesh/mesh/rectangular_bilinear_adapt_image.py new file mode 100644 index 000000000..48b6a925b --- /dev/null +++ b/autoarray/inversion/mesh/mesh/rectangular_bilinear_adapt_image.py @@ -0,0 +1,81 @@ +from typing import Tuple + +from autoarray.inversion.mesh.mesh.rectangular_rtu_adapt_image import ( + RectangularRTUAdaptImage, +) + + +class RectangularBilinearAdaptImage(RectangularRTUAdaptImage): + + def __init__( + self, + shape: Tuple[int, int] = (3, 3), + weight_power: float = 1.0, + weight_floor: float = 0.0, + ): + """ + A rectangular mesh of pixels used to reconstruct a source on a regular + grid, with adaptive weighting driven by an external adapt image and + the empirical rank-CDF transform. + + The mesh geometry is fixed and defined by a 2D shape + `(total_y_pixels, total_x_pixels)`. Pixels are indexed in row-major order: + + - Index 0 corresponds to the top-left pixel. + - Indices increase left-to-right across rows and top-to-bottom + between rows. + + Adaptive behaviour (adapt image) + -------------------------------- + Like `RectangularRTUAdaptImage`, this mesh adapts using an *adapt + image*: weights that emphasise specific regions of the source plane, + typically bright regions of a previously estimated reconstruction. + Pixels corresponding to higher adapt-image intensity receive increased + weighting, controlled by `weight_power` and `weight_floor` (see + `RectangularRTUAdaptImage` for the full description). + + The weighted lattice transform is the per-axis empirical rank CDF of + the traced coordinates (a sort plus a weighted cumulative sum) — no + kernel hyperparameters and O(N log N) per likelihood evaluation, + making this the fast default adaptive rectangular mesh for CPU + fitting. + + When to use Bilinear vs RTU + --------------------------- + See `RectangularBilinearAdaptDensity`: gradient-based (JAX) samplers + need `over_sample_size_pixelization >= 4` on imaging data, or the RTU + meshes; interferometer gradient work must use the RTU meshes, which + are also the recommended option on GPU. + + Edge handling + ------------- + Boundary (edge) pixels are automatically identified via the mesh + neighbour structure and may be internally excluded (zeroed) during + inversion to improve numerical stability and reduce edge artefacts. + + Parameters + ---------- + shape : Tuple[int, int] + The 2D dimensions of the rectangular pixel grid + `(total_y_pixels, total_x_pixels)`. + weight_power : float, optional + Exponent applied to the adapt-image weights to control the strength + of adaptivity. + weight_floor : float, optional + Minimum weight applied to ensure numerical stability in low-intensity + regions. + """ + super().__init__( + shape=shape, + weight_power=weight_power, + weight_floor=weight_floor, + ) + + @property + def interpolator_kwargs(self) -> dict: + """ + Extra keyword arguments `interpolator_from` forwards to + `interpolator_cls` — the rank-CDF transform selector; the empirical + rank CDF has no kernel hyperparameters. + """ + return {"transform": "rank"} diff --git a/autoarray/inversion/mesh/mesh/rectangular_adapt_density.py b/autoarray/inversion/mesh/mesh/rectangular_rtu_adapt_density.py similarity index 92% rename from autoarray/inversion/mesh/mesh/rectangular_adapt_density.py rename to autoarray/inversion/mesh/mesh/rectangular_rtu_adapt_density.py index 1e45b3db9..9d1469499 100644 --- a/autoarray/inversion/mesh/mesh/rectangular_adapt_density.py +++ b/autoarray/inversion/mesh/mesh/rectangular_rtu_adapt_density.py @@ -61,9 +61,9 @@ def overlay_grid_from( return grid_slim -class RectangularAdaptDensity(AbstractMesh): +class RectangularRTUAdaptDensity(AbstractMesh): # Rectangular meshes do not support split regularization -- their interpolators provide no - # split-cross mappings. Inherited by `RectangularUniform` and `RectangularAdaptImage`. + # split-cross mappings. Inherited by `RectangularUniform` and `RectangularRTUAdaptImage`. supports_split_regularization = False def __init__( @@ -75,7 +75,8 @@ def __init__( """ A rectangular mesh of pixels used to reconstruct a source on a regular grid, whose pixels adapt to the density of the traced source-plane - coordinates. + coordinates through the smooth RTU kernel-density CDF transform + (Enzi et al., arXiv:2606.30620). The mesh is defined by a 2D shape `(total_y_pixels, total_x_pixels)` and is indexed in row-major order: @@ -84,6 +85,17 @@ def __init__( - Indices increase from left to right across each row, and from top to bottom across rows. + When to use RTU vs Bilinear + --------------------------- + The RTU meshes are the recommended advanced option: on GPU, for JAX + gradient-based samplers, and on the interferometer sparse path (where + they are the only adaptive rectangular meshes with usable gradients). + Their kernel-CDF evaluation is an O(M_sub x N_data) erf sum that + dominates the CPU likelihood at production scale, so for CPU-only + fitting the `RectangularBilinearAdaptDensity` / + `RectangularBilinearAdaptImage` meshes (empirical rank-CDF transform, + no hyperparameters) are the faster default. + Adaptive behaviour ------------------ The mesh adapts to the spatial density of the traced points through a diff --git a/autoarray/inversion/mesh/mesh/rectangular_adapt_image.py b/autoarray/inversion/mesh/mesh/rectangular_rtu_adapt_image.py similarity index 94% rename from autoarray/inversion/mesh/mesh/rectangular_adapt_image.py rename to autoarray/inversion/mesh/mesh/rectangular_rtu_adapt_image.py index 8e8db7a16..3b1c8537d 100644 --- a/autoarray/inversion/mesh/mesh/rectangular_adapt_image.py +++ b/autoarray/inversion/mesh/mesh/rectangular_rtu_adapt_image.py @@ -1,12 +1,12 @@ import numpy as np from typing import Optional, Tuple -from autoarray.inversion.mesh.mesh.rectangular_adapt_density import ( - RectangularAdaptDensity, +from autoarray.inversion.mesh.mesh.rectangular_rtu_adapt_density import ( + RectangularRTUAdaptDensity, ) -class RectangularAdaptImage(RectangularAdaptDensity): +class RectangularRTUAdaptImage(RectangularRTUAdaptDensity): def __init__( self, @@ -76,7 +76,7 @@ def __init__( regions. bandwidth : float, optional Kernel bandwidth in units of the mesh pixel scale (see - `RectangularAdaptDensity`). Defaults to the kernel default. + `RectangularRTUAdaptDensity`). Defaults to the kernel default. n_knots : int, optional Size of the fixed knot table used to invert the CDF. Defaults to the kernel default. diff --git a/autoarray/inversion/mesh/mesh/rectangular_uniform.py b/autoarray/inversion/mesh/mesh/rectangular_uniform.py index 22851684a..fad20eeda 100644 --- a/autoarray/inversion/mesh/mesh/rectangular_uniform.py +++ b/autoarray/inversion/mesh/mesh/rectangular_uniform.py @@ -1,9 +1,9 @@ -from autoarray.inversion.mesh.mesh.rectangular_adapt_density import ( - RectangularAdaptDensity, +from autoarray.inversion.mesh.mesh.rectangular_rtu_adapt_density import ( + RectangularRTUAdaptDensity, ) -class RectangularUniform(RectangularAdaptDensity): +class RectangularUniform(RectangularRTUAdaptDensity): """ A uniform rectangular mesh of pixels used to reconstruct a source on a regular grid, with no adaptive weighting. @@ -21,7 +21,7 @@ class RectangularUniform(RectangularAdaptDensity): Uniform behaviour ----------------- - Unlike `RectangularAdaptDensity` and `RectangularAdaptImage`, this mesh + Unlike the adaptive rectangular meshes (`RectangularBilinearAdapt*`, `RectangularRTUAdapt*`), this mesh applies no adaptive weighting based on data density or an adapt image. All pixels are treated equally in the reconstruction, and the effective resolution is determined solely by the fixed mesh geometry and the diff --git a/autoarray/inversion/mesh/mesh_geometry/abstract.py b/autoarray/inversion/mesh/mesh_geometry/abstract.py index 0e0a43d1a..958f27a02 100644 --- a/autoarray/inversion/mesh/mesh_geometry/abstract.py +++ b/autoarray/inversion/mesh/mesh_geometry/abstract.py @@ -11,6 +11,7 @@ def __init__( mesh_weight_map=None, kernel_bandwidth=None, kernel_knots=None, + transform="kernel", xp=np, ): @@ -22,6 +23,11 @@ def __init__( # areas / edges transforms; None falls back to the kernel defaults. self.kernel_bandwidth = kernel_bandwidth self.kernel_knots = kernel_knots + # The rectangular lattice transform the areas / edges route through: + # "kernel" (the RTU meshes and the uniform mesh's fallback) or "rank" + # (the Bilinear meshes) — must match the mapper's transform so the + # geometry is consistent with where the mapper scatters its flux. + self.transform = transform 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 f30c1ce3b..d11be49ad 100644 --- a/autoarray/inversion/mesh/mesh_geometry/rectangular.py +++ b/autoarray/inversion/mesh/mesh_geometry/rectangular.py @@ -464,6 +464,7 @@ def areas_transformed(self): mesh_weight_map=self.mesh_weight_map, bandwidth=bandwidth, n_knots=n_knots, + transform=self.transform, xp=self._xp, ) @@ -532,6 +533,7 @@ def edges_transformed(self): mesh_weight_map=self.mesh_weight_map, bandwidth=bandwidth, n_knots=n_knots, + transform=self.transform, xp=self._xp, ) diff --git a/autoarray/inversion/pixelization.py b/autoarray/inversion/pixelization.py index aa9431f01..477bf059e 100644 --- a/autoarray/inversion/pixelization.py +++ b/autoarray/inversion/pixelization.py @@ -36,7 +36,7 @@ def __init__( - **Left:** Observed image of a galaxy. - **Centre:** The (y,x) grid of coordinates corresponding to the centre of each pixel in the observed image. The centre of each pixel is shown by a magenta point. - - **Right:** An overlaid ``RectangularAdaptDensity`` ``mesh``, where the square pixel boundaries of this mesh are shown by + - **Right:** An overlaid ``RectangularBilinearAdaptDensity`` ``mesh``, where the square pixel boundaries of this mesh are shown by dashed black lines. Red points highlight a subset of points. They are used below to illustrate additional behaviour @@ -44,7 +44,7 @@ def __init__( The **centre** and **right** panels show the core functionality of a ``Pixelization``. It represents the mappings between an image's (y,x) grid of coordinates and a ``Mesh``'s (y,x) grid of coordinates (in the - example above the centres of ``RectangularAdaptDensity`` pixels, shown by the dashed black lines). + example above the centres of ``RectangularBilinearAdaptDensity`` pixels, shown by the dashed black lines). **Image-Plane Example (Masked)** @@ -63,7 +63,7 @@ def __init__( :width: 240 The behaviour is analogous to the non-masked case, however only unmasked pixel's in the image's (y,x) grid - of coordinates are included when pairing with the ``RectangularAdaptDensity`` mesh. + of coordinates are included when pairing with the ``RectangularBilinearAdaptDensity`` mesh. **Source-Plane Example (Masked)** @@ -94,7 +94,7 @@ def __init__( - ``Inversion``: Use the ``Pixelization`` to reconstruct the data on the mesh via linear algebra. - ``Regularization``: Apply smoothing to the solutions computed using an ``Inversion``. - In the example above, a ``RectangularAdaptDensity`` ``Mesh`` object is used. Other meshes are available (e.g. + In the example above, a ``RectangularBilinearAdaptDensity`` ``Mesh`` object is used. Other meshes are available (e.g. ``Delaunay``). **Source Code API** @@ -112,7 +112,7 @@ def __init__( Parameters ---------- mesh - The mesh object (e.g. RectangularAdaptDensity grid of pixels, Delaunay triangulation) describing the + The mesh object (e.g. RectangularBilinearAdaptDensity grid of pixels, Delaunay triangulation) describing the pixels of the `Pixelization`. regularization The regularization object that can smooth ``Pixelization`` pixels with one another when it is used to @@ -127,7 +127,7 @@ def __init__( grid_2d = al.Grid2D.uniform(shape_native=(50, 50), pixel_scales=0.1) - mesh = al.mesh.RectangularAdaptDensity(shape=(10, 10)) + mesh = al.mesh.RectangularBilinearAdaptDensity(shape=(10, 10)) pixelization = al.Pixelization(mesh=mesh) @@ -137,7 +137,7 @@ def __init__( import autogalaxy as ag - mesh = af.Model(ag.mesh.RectangularAdaptDensity) + mesh = af.Model(ag.mesh.RectangularBilinearAdaptDensity) mesh.shape_0 = af.UniformPrior(lower_limit=10, upper_limit=20) mesh.shape_1 = af.UniformPrior(lower_limit=10, upper_limit=20) @@ -164,7 +164,8 @@ def __init__( Split regularization regularizes using a cross of four points around each pixel centre, which requires the mesh to provide split-cross mappings. The rectangular meshes - (`RectangularUniform`, `RectangularAdaptDensity`, `RectangularAdaptImage`) do not + (`RectangularUniform`, `RectangularBilinearAdaptDensity`, `RectangularBilinearAdaptImage`, + `RectangularRTUAdaptDensity`, `RectangularRTUAdaptImage`) do not provide them. Use either: diff --git a/autoarray/util/misc_util.py b/autoarray/util/misc_util.py index f9ae56cc8..73a0188f0 100644 --- a/autoarray/util/misc_util.py +++ b/autoarray/util/misc_util.py @@ -35,8 +35,8 @@ def cls_list_from( - If the input is `cls=aa.mesh.Mesh`, a list containing all pixelizations in the class are returned. - - If `cls=aa.mesh.Mesh` and `cls_filtered=aa.mesh.RectangularAdaptDensity`, a list of all pixelizations - excluding those which are `RectangularAdaptDensity` pixelizations will be returned. + - If `cls=aa.mesh.Mesh` and `cls_filtered=aa.mesh.RectangularRTUAdaptDensity`, a list of all pixelizations + excluding those which are `RectangularRTUAdaptDensity` pixelizations will be returned. Parameters ---------- diff --git a/files/ghost_peak_experiment.py b/files/ghost_peak_experiment.py index ba23189fc..c7533ebb5 100644 --- a/files/ghost_peak_experiment.py +++ b/files/ghost_peak_experiment.py @@ -16,7 +16,7 @@ for the geometric diagnostic). 2. Build a synthetic adapt image as a sum of K isotropic Gaussian peaks of matched amplitude. -3. Compute ``mesh_weight_map`` exactly as ``RectangularAdaptImage`` does: +3. Compute ``mesh_weight_map`` exactly as ``RectangularRTUAdaptImage`` does: ``clip(b, eps, None) ** power`` then normalise. 4. Build the spline CDF transforms via ``create_transforms_spline`` (the same function ``RectangularSplineAdaptImage`` uses). @@ -146,7 +146,7 @@ def run_experiment(peaks_label, peaks, real_zones, ghost_zones, probe_radius=0.15): traced = sample_traced_points(n_traced, seed=0) brightness = brightness_at(traced, peaks=peaks, sigma=sigma) - # Mimic RectangularAdaptImage.mesh_weight_map_from with power=1, floor=0 + # Mimic RectangularRTUAdaptImage.mesh_weight_map_from with power=1, floor=0 weight_map = np.clip(brightness, 1e-12, None) weight_map = weight_map / weight_map.sum() diff --git a/test_autoarray/inversion/pixelization/interpolator/test_rectangular.py b/test_autoarray/inversion/pixelization/interpolator/test_rectangular.py index ba8bd160c..fee895705 100644 --- a/test_autoarray/inversion/pixelization/interpolator/test_rectangular.py +++ b/test_autoarray/inversion/pixelization/interpolator/test_rectangular.py @@ -1,4 +1,5 @@ -"""Unit tests for the adaptive rectangular interpolator (kernel-density CDF). +"""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 autolens_workspace_test/scripts/jax_grad per the project's "no JAX in unit @@ -16,6 +17,7 @@ adaptive_rectangular_areas_from, adaptive_rectangular_mappings_weights_via_interpolation_from, create_transforms, + create_transforms_rank, ) @@ -137,6 +139,139 @@ def test__create_transforms__chunked_forward_is_block_size_invariant(): assert batched == pytest.approx(row_wise, abs=0.0) +# --------------------------------------------------------------------------- +# Empirical rank-CDF transform (the Bilinear meshes) +# --------------------------------------------------------------------------- + + +def test__create_transforms_rank__unweighted_cdf_is_ranks_at_points(): + data_grid, _, _ = _seeded_inputs(seed=1) + N = data_grid.shape[0] + + fwd, _ = create_transforms_rank(data_grid, xp=np) + + # At the sorted points themselves the empirical CDF is exactly the rank + # values (i + 1) / (N + 1), per axis. + sort_points = np.sort(data_grid, axis=0) + expected = np.arange(1, N + 1) / (N + 1) + F = fwd(sort_points) + assert F[:, 0] == pytest.approx(expected, abs=1e-12) + assert F[:, 1] == pytest.approx(expected, abs=1e-12) + + +def test__create_transforms_rank__weighted_cdf_is_cumsum_at_points(): + data_grid, _, weights = _seeded_inputs(seed=2) + + fwd, _ = create_transforms_rank(data_grid, mesh_weight_map=weights, xp=np) + + for d in range(2): + order = np.argsort(data_grid[:, d]) + expected = np.cumsum(weights[order]) + F = fwd(np.sort(data_grid, axis=0)) + assert F[:, d] == pytest.approx(expected, abs=1e-12) + + +def test__create_transforms_rank__monotone_and_bounded(): + data_grid, _, weights = _seeded_inputs(seed=3) + + fwd, _ = create_transforms_rank(data_grid, mesh_weight_map=weights, xp=np) + + q = np.linspace(data_grid.min(axis=0) - 0.5, data_grid.max(axis=0) + 0.5, 500) + F = fwd(q) + assert np.all(np.diff(F, axis=0) >= 0.0) + assert F.min() >= 0.0 + assert F.max() <= 1.0 + + +def test__create_transforms_rank__roundtrip_matches_identity(): + data_grid, _, _ = _seeded_inputs(seed=4) + N = data_grid.shape[0] + + fwd, rev = create_transforms_rank(data_grid, xp=np) + + # Interior unit-square probes (inside [1/(N+1), N/(N+1)] where the + # piecewise-linear CDF is invertible) round-trip exactly. + probe = np.array([[0.1, 0.1], [0.5, 0.5], [0.9, 0.9]]) + assert 0.1 > 1.0 / (N + 1) and 0.9 < N / (N + 1) + roundtrip = fwd(rev(probe)) + assert roundtrip == pytest.approx(probe, abs=1e-12) + + +def test__rank__mappings_sizes_weights__shapes_and_weight_normalization(): + data_grid, over, weights = _seeded_inputs() + + idx, w = adaptive_rectangular_mappings_weights_via_interpolation_from( + source_grid_size=16, + data_grid=data_grid, + data_grid_over_sampled=over, + mesh_weight_map=weights, + transform="rank", + xp=np, + ) + + assert idx.shape == (400, 4) + assert w.shape == (400, 4) + assert np.allclose(w.sum(axis=1), 1.0, atol=1e-10) + + +def test__rank__areas__positive_finite__total_is_bounding_box_area(): + data_grid, _, weights = _seeded_inputs(seed=6) + + areas = adaptive_rectangular_areas_from( + source_grid_shape=(12, 12), + data_grid=data_grid, + mesh_weight_map=weights, + transform="rank", + xp=np, + ) + + assert areas.shape == (144,) + assert np.all(np.isfinite(areas)) + assert np.all(areas > 0.0) + span = data_grid.max(axis=0) - data_grid.min(axis=0) + assert areas.sum() == pytest.approx(span[0] * span[1], rel=1e-8) + + +def test__rank__areas__adapt_to_point_density(): + """A dense cluster of traced points must shrink the mesh pixels covering + it relative to a sparse region — the adaptive property the rank CDF + exists to provide.""" + rng = np.random.default_rng(11) + cluster = rng.normal(loc=-1.0, scale=0.05, size=(400, 2)) + sparse = rng.uniform(low=-2.0, high=2.0, size=(100, 2)) + data_grid = np.concatenate([cluster, sparse]) + + areas = adaptive_rectangular_areas_from( + source_grid_shape=(10, 10), + data_grid=data_grid, + transform="rank", + xp=np, + ) + + # Strong adaptivity: the cluster holds 80% of the rank mass, so the cells + # covering it shrink by orders of magnitude relative to the sparse + # outskirts. A uniform lattice would have every cell equal. + assert areas.min() < areas.max() / 100.0 + + # The smallest cells cover the cluster: every cell of side <~ 4 sigma + # contains cluster points, so its area is far below the uniform cell area. + uniform_cell = areas.sum() / areas.size + assert areas.min() < uniform_cell / 100.0 + + +def test__transforms_from__invalid_transform_raises(): + data_grid, over, _ = _seeded_inputs() + + with pytest.raises(ValueError): + adaptive_rectangular_mappings_weights_via_interpolation_from( + source_grid_size=16, + data_grid=data_grid, + data_grid_over_sampled=over, + transform="spline", + xp=np, + ) + + # --------------------------------------------------------------------------- # Mapper output shapes # --------------------------------------------------------------------------- @@ -201,7 +336,7 @@ def __getattr__(self, item): rng = np.random.default_rng(5) data_grid = _StubGrid(rng.standard_normal((64, 2))) - mesh = aa.mesh.RectangularAdaptDensity(shape=(6, 6), bandwidth=0.8) + mesh = aa.mesh.RectangularRTUAdaptDensity(shape=(6, 6), bandwidth=0.8) interpolator = InterpolatorRectangular( mesh=mesh, mesh_grid=_StubGrid(rng.standard_normal((36, 2))), @@ -221,6 +356,50 @@ def __getattr__(self, item): geometry = interpolator.mesh_geometry assert geometry.kernel_bandwidth == 0.8 assert geometry.kernel_knots == KERNEL_CDF_DEFAULT_KNOTS + assert geometry.transform == "kernel" + + areas = geometry.areas_transformed + assert areas.shape == (36,) + assert np.all(np.isfinite(areas)) + assert np.all(areas > 0.0) + + +def test__InterpolatorRectangular__rank_transform_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.RectangularBilinearAdaptDensity(shape=(6, 6)) + interpolator = InterpolatorRectangular( + mesh=mesh, + mesh_grid=_StubGrid(rng.standard_normal((36, 2))), + data_grid=data_grid, + mesh_weight_map=None, + **mesh.interpolator_kwargs, + xp=np, + ) + + assert interpolator.transform == "rank" + + 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) + + # The geometry routes areas / edges through the same rank transform as + # the mapper, so the two stay consistent. + geometry = interpolator.mesh_geometry + assert geometry.transform == "rank" areas = geometry.areas_transformed assert areas.shape == (36,) @@ -258,8 +437,13 @@ def test__forward_transform__windowed_numba_matches_dense_reference(weighted): data_grid, data_grid_over, weights = _seeded_inputs(M=300, K=500, seed=7) # queries beyond the data bounding box exercise the saturated tails - q = np.concatenate([data_grid_over, data_grid.min(axis=0) - 1.0 + np.zeros((1, 2)), - data_grid.max(axis=0) + 1.0 + np.zeros((1, 2))]) + q = np.concatenate( + [ + data_grid_over, + data_grid.min(axis=0) - 1.0 + np.zeros((1, 2)), + data_grid.max(axis=0) + 1.0 + np.zeros((1, 2)), + ] + ) fwd, _ = create_transforms( data_grid, mesh_pixels=16, mesh_weight_map=weights if weighted else None, xp=np diff --git a/test_autoarray/inversion/pixelization/mappers/test_rectangular.py b/test_autoarray/inversion/pixelization/mappers/test_rectangular.py index 2594571c7..d895d3c52 100644 --- a/test_autoarray/inversion/pixelization/mappers/test_rectangular.py +++ b/test_autoarray/inversion/pixelization/mappers/test_rectangular.py @@ -2,7 +2,7 @@ import autoarray as aa -from autoarray.inversion.mesh.mesh.rectangular_adapt_density import ( +from autoarray.inversion.mesh.mesh.rectangular_rtu_adapt_density import ( overlay_grid_from, ) from autoarray.inversion.mesh.interpolator.rectangular_uniform import ( @@ -61,7 +61,7 @@ def test__pixel_signals_from__rectangular_adapt_density_mesh__matches_util( shape_native=(3, 3), grid=grid_2d_sub_1_7x7.over_sampled, buffer=1e-8 ) - mesh = aa.mesh.RectangularAdaptDensity(shape=(3, 3)) + mesh = aa.mesh.RectangularRTUAdaptDensity(shape=(3, 3)) # The adapt image is defined on the data's own mask, as it is for a real fit. diff --git a/test_autoarray/inversion/pixelization/mesh/test_rectangular.py b/test_autoarray/inversion/pixelization/mesh/test_rectangular.py index f0baf1d7d..3a4a9f127 100644 --- a/test_autoarray/inversion/pixelization/mesh/test_rectangular.py +++ b/test_autoarray/inversion/pixelization/mesh/test_rectangular.py @@ -11,7 +11,7 @@ from autoarray.inversion.mesh.interpolator.rectangular_uniform import ( InterpolatorRectangularUniform, ) -from autoarray.inversion.mesh.mesh.rectangular_adapt_density import ( +from autoarray.inversion.mesh.mesh.rectangular_rtu_adapt_density import ( overlay_grid_from, ) @@ -150,8 +150,8 @@ def test__overlay_grid_from__pixel_centres__3x3_grid__pixel_centres(): def test__construction__shape_and_kernel_kwargs__default(): - density = aa.mesh.RectangularAdaptDensity(shape=(5, 7)) - image = aa.mesh.RectangularAdaptImage(shape=(5, 7)) + density = aa.mesh.RectangularRTUAdaptDensity(shape=(5, 7)) + image = aa.mesh.RectangularRTUAdaptImage(shape=(5, 7)) assert density.shape == (5, 7) assert density.bandwidth == KERNEL_CDF_DEFAULT_BANDWIDTH @@ -165,8 +165,10 @@ def test__construction__shape_and_kernel_kwargs__default(): def test__construction__kernel_kwargs__overridden(): - density = aa.mesh.RectangularAdaptDensity(shape=(3, 3), bandwidth=0.5, n_knots=128) - image = aa.mesh.RectangularAdaptImage( + density = aa.mesh.RectangularRTUAdaptDensity( + shape=(3, 3), bandwidth=0.5, n_knots=128 + ) + image = aa.mesh.RectangularRTUAdaptImage( shape=(3, 3), weight_power=2.0, weight_floor=0.1, bandwidth=2.0, n_knots=32 ) @@ -180,9 +182,9 @@ def test__construction__kernel_kwargs__overridden(): def test__construction__minimum_shape_raises(): with pytest.raises(aa.exc.MeshException): - aa.mesh.RectangularAdaptDensity(shape=(2, 3)) + aa.mesh.RectangularRTUAdaptDensity(shape=(2, 3)) with pytest.raises(aa.exc.MeshException): - aa.mesh.RectangularAdaptImage(shape=(3, 2)) + aa.mesh.RectangularRTUAdaptImage(shape=(3, 2)) with pytest.raises(aa.exc.MeshException): aa.mesh.RectangularUniform(shape=(2, 2)) @@ -193,8 +195,10 @@ def test__construction__minimum_shape_raises(): def test__interpolator_cls_and_kwargs(): - density = aa.mesh.RectangularAdaptDensity(shape=(3, 3), bandwidth=0.5, n_knots=128) - image = aa.mesh.RectangularAdaptImage(shape=(3, 3)) + density = aa.mesh.RectangularRTUAdaptDensity( + shape=(3, 3), bandwidth=0.5, n_knots=128 + ) + image = aa.mesh.RectangularRTUAdaptImage(shape=(3, 3)) uniform = aa.mesh.RectangularUniform(shape=(3, 3)) assert density.interpolator_cls is InterpolatorRectangular @@ -216,12 +220,12 @@ def test__interpolator_cls_and_kwargs(): def test__mesh_weight_map_from__density__returns_none(): - density = aa.mesh.RectangularAdaptDensity(shape=(3, 3)) + density = aa.mesh.RectangularRTUAdaptDensity(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.RectangularAdaptImage( + image = aa.mesh.RectangularRTUAdaptImage( shape=(3, 3), weight_power=2.0, weight_floor=0.0 ) @@ -235,3 +239,78 @@ def __init__(self, arr): 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) + + +# --------------------------------------------------------------------------- +# Bilinear (rank-CDF) meshes +# --------------------------------------------------------------------------- + + +def test__bilinear__construction_and_interpolator_dispatch(): + density = aa.mesh.RectangularBilinearAdaptDensity(shape=(5, 7)) + image = aa.mesh.RectangularBilinearAdaptImage( + shape=(5, 7), weight_power=2.0, weight_floor=0.1 + ) + + assert density.shape == (5, 7) + assert image.shape == (5, 7) + assert image.weight_power == 2.0 + assert image.weight_floor == 0.1 + + assert density.interpolator_cls is InterpolatorRectangular + assert image.interpolator_cls is InterpolatorRectangular + + # The rank CDF has no kernel hyperparameters — only the transform selector. + assert density.interpolator_kwargs == {"transform": "rank"} + assert image.interpolator_kwargs == {"transform": "rank"} + + +def test__bilinear__no_kernel_hyperparameters_in_signature(): + with pytest.raises(TypeError): + aa.mesh.RectangularBilinearAdaptDensity(shape=(3, 3), bandwidth=0.5) + with pytest.raises(TypeError): + aa.mesh.RectangularBilinearAdaptImage(shape=(3, 3), n_knots=128) + + +def test__bilinear__minimum_shape_raises(): + with pytest.raises(aa.exc.MeshException): + aa.mesh.RectangularBilinearAdaptDensity(shape=(2, 3)) + with pytest.raises(aa.exc.MeshException): + aa.mesh.RectangularBilinearAdaptImage(shape=(3, 2)) + + +def test__bilinear__mesh_weight_map_matches_rtu(): + class _Stub: + def __init__(self, arr): + self.array = arr + + adapt = _Stub(np.array([1.0, 2.0, 4.0, 0.0, 8.0])) + + bilinear = aa.mesh.RectangularBilinearAdaptImage( + shape=(3, 3), weight_power=2.0, weight_floor=0.0 + ) + rtu = aa.mesh.RectangularRTUAdaptImage( + shape=(3, 3), weight_power=2.0, weight_floor=0.0 + ) + + assert bilinear.mesh_weight_map_from(adapt_data=adapt) == pytest.approx( + rtu.mesh_weight_map_from(adapt_data=adapt), rel=1e-12 + ) + + density = aa.mesh.RectangularBilinearAdaptDensity(shape=(3, 3)) + assert density.mesh_weight_map_from(adapt_data=None) is None + + +def test__bilinear__split_regularization_not_supported(): + assert ( + aa.mesh.RectangularBilinearAdaptDensity( + shape=(15, 15) + ).supports_split_regularization + is False + ) + assert ( + aa.mesh.RectangularBilinearAdaptImage( + shape=(15, 15) + ).supports_split_regularization + is False + ) diff --git a/test_autoarray/inversion/pixelization/mesh_geometry/test_rectangular.py b/test_autoarray/inversion/pixelization/mesh_geometry/test_rectangular.py index 63de1444a..370df69dd 100644 --- a/test_autoarray/inversion/pixelization/mesh_geometry/test_rectangular.py +++ b/test_autoarray/inversion/pixelization/mesh_geometry/test_rectangular.py @@ -3,7 +3,7 @@ import autoarray as aa -from autoarray.inversion.mesh.mesh.rectangular_adapt_density import ( +from autoarray.inversion.mesh.mesh.rectangular_rtu_adapt_density import ( overlay_grid_from, ) @@ -18,7 +18,7 @@ def test__rectangular_neighbors_from(): # I3I4I5I # I6I7I8I - (neighbors, neighbors_sizes) = rectangular_neighbors_from(shape_native=(3, 3)) + neighbors, neighbors_sizes = rectangular_neighbors_from(shape_native=(3, 3)) # TODO : Use pytest.parameterize @@ -38,7 +38,7 @@ def test__rectangular_neighbors_from(): # I4I5I 6I 7I # I8I9I10I11I - (neighbors, neighbors_sizes) = rectangular_neighbors_from(shape_native=(3, 4)) + neighbors, neighbors_sizes = rectangular_neighbors_from(shape_native=(3, 4)) assert (neighbors[0] == [1, 4, -1, -1]).all() assert (neighbors[1] == [0, 2, 5, -1]).all() @@ -60,7 +60,7 @@ def test__rectangular_neighbors_from(): # I6I 7I 8I # I9I10I11I - (neighbors, neighbors_sizes) = rectangular_neighbors_from(shape_native=(4, 3)) + neighbors, neighbors_sizes = rectangular_neighbors_from(shape_native=(4, 3)) assert (neighbors[0] == [1, 3, -1, -1]).all() assert (neighbors[1] == [0, 2, 4, -1]).all() @@ -82,7 +82,7 @@ def test__rectangular_neighbors_from(): # I8 I 9I10I11I # I12I13I14I15I - (neighbors, neighbors_sizes) = rectangular_neighbors_from(shape_native=(4, 4)) + neighbors, neighbors_sizes = rectangular_neighbors_from(shape_native=(4, 4)) assert (neighbors[0] == [1, 4, -1, -1]).all() assert (neighbors[1] == [0, 2, 5, -1]).all() @@ -122,7 +122,7 @@ def test__neighbors__compare_to_mesh_util(): mesh=mesh, mesh_grid=mesh_grid, data_grid=None ) - (neighbors_util, neighbors_sizes_util) = rectangular_neighbors_from( + neighbors_util, neighbors_sizes_util = rectangular_neighbors_from( shape_native=(7, 5) ) @@ -151,7 +151,7 @@ def test__areas_transformed(mask_2d_7x7): mesh_grid = overlay_grid_from(shape_native=(3, 3), grid=grid, buffer=1e-8) - mesh = aa.mesh.RectangularAdaptDensity(shape=(3, 3)) + mesh = aa.mesh.RectangularRTUAdaptDensity(shape=(3, 3)) interpolator = mesh.interpolator_from( source_plane_data_grid=grid, @@ -198,7 +198,7 @@ def test__edges_transformed(mask_2d_7x7): mesh_grid = overlay_grid_from(shape_native=(3, 3), grid=grid, buffer=1e-8) - mesh = aa.mesh.RectangularAdaptDensity(shape=(3, 3)) + mesh = aa.mesh.RectangularRTUAdaptDensity(shape=(3, 3)) interpolator = mesh.interpolator_from( source_plane_data_grid=grid, diff --git a/test_autoarray/inversion/pixelization/test_split_regularization_support.py b/test_autoarray/inversion/pixelization/test_split_regularization_support.py index 28192f1d2..f632923b5 100644 --- a/test_autoarray/inversion/pixelization/test_split_regularization_support.py +++ b/test_autoarray/inversion/pixelization/test_split_regularization_support.py @@ -7,7 +7,7 @@ - ``RectangularUniform`` -> ``AttributeError: 'InterpolatorRectangularUniform' object has no attribute '_mappings_sizes_weights_split'`` -- ``RectangularAdaptDensity`` / ``RectangularAdaptImage`` +- ``RectangularRTUAdaptDensity`` / ``RectangularRTUAdaptImage`` -> ``IndexError: index 4 is out of bounds for axis 0 with size 4`` (``InterpolatorRectangular`` returned the plain 4-corner mappings from a pass-through that claimed split "reuses the same mappings") @@ -21,11 +21,12 @@ import autoarray as aa from autoarray import exc - RECTANGULAR_MESHES = [ aa.mesh.RectangularUniform, - aa.mesh.RectangularAdaptDensity, - aa.mesh.RectangularAdaptImage, + aa.mesh.RectangularBilinearAdaptDensity, + aa.mesh.RectangularBilinearAdaptImage, + aa.mesh.RectangularRTUAdaptDensity, + aa.mesh.RectangularRTUAdaptImage, ] SPLIT_REGULARIZATIONS = [ @@ -43,8 +44,10 @@ @pytest.mark.parametrize("mesh_cls", RECTANGULAR_MESHES) @pytest.mark.parametrize("regularization_cls", SPLIT_REGULARIZATIONS) -def test__rectangular_mesh_with_split_regularization__raises(mesh_cls, regularization_cls): - """All 9 rectangular-mesh x split-regularization combinations are rejected.""" +def test__rectangular_mesh_with_split_regularization__raises( + mesh_cls, regularization_cls +): + """All 15 rectangular-mesh x split-regularization combinations are rejected.""" with pytest.raises(exc.PixelizationException) as error: aa.Pixelization( @@ -73,7 +76,9 @@ def test__rectangular_mesh_with_non_split_regularization__is_allowed(mesh_cls): @pytest.mark.parametrize("mesh_cls", ADAPTIVE_MESHES) @pytest.mark.parametrize("regularization_cls", SPLIT_REGULARIZATIONS) -def test__adaptive_mesh_with_split_regularization__is_allowed(mesh_cls, regularization_cls): +def test__adaptive_mesh_with_split_regularization__is_allowed( + mesh_cls, regularization_cls +): """Split regularization remains supported on the meshes that implement it.""" pixelization = aa.Pixelization( @@ -96,9 +101,18 @@ def test__rectangular_mesh_without_regularization__is_allowed(mesh_cls): def test__capability_flags(): """The flags the guard reads, asserted directly so a future mesh can't silently regress.""" - assert aa.mesh.RectangularUniform(shape=(15, 15)).supports_split_regularization is False - assert aa.mesh.RectangularAdaptDensity(shape=(15, 15)).supports_split_regularization is False - assert aa.mesh.RectangularAdaptImage(shape=(15, 15)).supports_split_regularization is False + assert ( + aa.mesh.RectangularUniform(shape=(15, 15)).supports_split_regularization + is False + ) + assert ( + aa.mesh.RectangularRTUAdaptDensity(shape=(15, 15)).supports_split_regularization + is False + ) + assert ( + aa.mesh.RectangularRTUAdaptImage(shape=(15, 15)).supports_split_regularization + is False + ) assert aa.mesh.Delaunay(pixels=100).supports_split_regularization is True assert aa.mesh.DelaunayNN(pixels=100).supports_split_regularization is True assert aa.mesh.KNNBarycentric(pixels=100).supports_split_regularization is True @@ -116,7 +130,9 @@ def test__interpolator_rectangular_has_no_split_mappings(): would reintroduce a silent-looking API that fails one frame later. """ - from autoarray.inversion.mesh.interpolator.rectangular import InterpolatorRectangular + from autoarray.inversion.mesh.interpolator.rectangular import ( + InterpolatorRectangular, + ) from autoarray.inversion.mesh.interpolator.rectangular_uniform import ( InterpolatorRectangularUniform, )