diff --git a/autoarray/config/general.yaml b/autoarray/config/general.yaml index e4dc1b6f5..b06ae2707 100644 --- a/autoarray/config/general.yaml +++ b/autoarray/config/general.yaml @@ -20,6 +20,7 @@ inversion: regularization_term_method: matmul # How the Bayesian-evidence regularization term s^T H s is computed. "matmul" (default) is the historical s @ (H @ s) against the explicitly formed regularization matrix; "cho_solve" evaluates coefficient * s^T C^-1 s for the kernel schemes (Matern/Gaussian/Exponential/MaternAdapt) via one Cholesky solve of their covariance C, avoiding the explicit inverse whose round-off is amplified by cond(C) (~1e9 on clustered traced mesh vertices). Opt-in, non-default; does not change the default evidence. Schemes with no such factorization fall back to the formed matrix. interferometer_numba_nnz_per_source_max: null # How a NumPy (xp=np) single-mapper interferometer inversion chooses between the numba `direct_conv` curvature path (InversionInterferometerSparseNumba) and the FFT one (InversionInterferometerSparse). null (default) = cost model: predicts both routes' single-thread time from source pixels S, total mapping non-zeros and the extent rectangle (FFT ~ S x rfft2+irfft2 cost of the padded (2ny,2nx) extent, direct_conv ~ nnz x ny x nx) and picks the faster; calibrated on autolens_profiling#226 and the Aris 512x512 mask (PyAutoArray#628). A positive value is a hard gate in mean non-zeros per source column (mapper.pix_sizes_for_sub_slim_index.sum() / mapper.params): numba at or below it, FFT above (the pre-#628 behaviour, default was 60.0). Set 0 to disable the numba path entirely. interferometer_route: null # Force the interferometer curvature route for benchmarking: null (default) routes by interferometer_numba_nnz_per_source_max; "fft" always uses the FFT route; "numba" uses the numba direct_conv route whenever it can represent the inversion (NumPy, one mapper, no linear functions, no over-sampling), overriding the gate and the 0 kill switch. + interferometer_direct_conv_symmetric_max_work_fraction: 0.7 # When the serial (numba parallel: false) direct_conv curvature route runs, use the symmetric row-restricted kernel (curvature_direct_conv_symmetric: source columns ranked by their top extent row, each convolving only the extent rows later-ranked columns span and only each row's mask x-range, then a mirror pass) when its geometry work fraction (its convolution work / the full kernel's) is at or below this value. Same F to round-off (computed entries bit-identical). Synthetic Aris-512 geometry (compact columns): f ~0.41, 1.81-1.90x single-thread on RAL; Aris's real 512 Delaunay mesh (multiply imaged, tall columns): f ~0.51, ~1.46x measured (PyAutoArray#635). The route cost model scales t_dc by min(1, 1.32 x work fraction), so the kernel wins below f ~0.76; 0.7 keeps a predicted gain >= 1.08x. Set 0 to always use the full kernel. The parallel kernel is unaffected. numba: use_numba: true cache: true diff --git a/autoarray/inversion/inversion/factory.py b/autoarray/inversion/inversion/factory.py index 3b908f034..f47c22cc5 100644 --- a/autoarray/inversion/inversion/factory.py +++ b/autoarray/inversion/inversion/factory.py @@ -11,9 +11,15 @@ from autoarray.inversion.inversion.interferometer.sparse import ( InversionInterferometerSparse, ) +from autoarray.inversion.inversion.interferometer.abstract import ( + real_space_mask_from, +) from autoarray.inversion.inversion.interferometer_numba.sparse import ( InversionInterferometerSparseNumba, ) +from autoarray.inversion.inversion.interferometer_numba import ( + sparse as interferometer_numba_sparse, +) from autoarray.inversion.inversion.interferometer_numba import ( inversion_interferometer_numba_util, ) @@ -218,6 +224,10 @@ def inversion_interferometer_from( dataset.sparse_operator.y_shape, dataset.sparse_operator.x_shape, ), + # Only the NumPy cost model reads the mask, so the JAX path never resolves it. + real_space_mask=( + real_space_mask_from(dataset=dataset) if xp is np else None + ), ): return InversionInterferometerSparseNumba( dataset=dataset, @@ -248,6 +258,7 @@ def _use_interferometer_numba( settings: Settings = None, xp=np, extent_shape: Optional[Tuple[int, int]] = None, + real_space_mask=None, ) -> bool: """ Whether an interferometer inversion is routed to the numba `direct_conv` curvature @@ -281,7 +292,14 @@ def _use_interferometer_numba( which predicts both routes' time from the source pixels, the mapping operator's non-zeros and the extent rectangle `extent_shape`, and routes to numba when `direct_conv` is predicted no slower. Without an `extent_shape` the model cannot - be evaluated and the FFT route is used. + be evaluated and the FFT route is used. When numba is serial and + `interferometer_direct_conv_symmetric_max_work_fraction` is positive, the + symmetric kernel's work fraction is computed from the mapper's triplets and + `real_space_mask` (`symmetric_work_fraction_from`), and if it is at or below that + threshold -- i.e. `InversionInterferometerSparseNumba` will run the cheaper + symmetric kernel -- it is passed to the cost model, which scales the predicted + `direct_conv` time by it. Without a `real_space_mask` whose extent matches + `extent_shape`, the full kernel's cost is predicted. - `import numba` succeeds. NaN mapping weights (an overflowing or degenerate `DelaunayNN` Sibson row) are @@ -304,6 +322,9 @@ def _use_interferometer_numba( The `(ny, nx)` extent rectangle the `W~` operator lives on, i.e. the dataset's `sparse_operator.y_shape, sparse_operator.x_shape` (`= mask.shape_native_masked_pixels`). Needed only by the default cost model. + real_space_mask + The real-space mask the inversion reconstructs on (`real_space_mask_from`), used + only by the default cost model for the symmetric kernel's work fraction. """ if xp is not np: return False @@ -360,6 +381,12 @@ def _use_interferometer_numba( source_pixels=int(mapper.params), nnz=int(np.asarray(mapper.pix_sizes_for_sub_slim_index).sum()), extent_shape=extent_shape, + work_fraction=_symmetric_work_fraction_for_route( + mapper=mapper, + settings=settings, + extent_shape=extent_shape, + real_space_mask=real_space_mask, + ), ) ) @@ -372,3 +399,56 @@ def _use_interferometer_numba( return False return True + + +def _symmetric_work_fraction_for_route( + mapper, settings: Settings, extent_shape, real_space_mask +) -> Optional[float]: + """ + The work fraction to give the route cost model: the symmetric kernel's + `work_fraction` when `InversionInterferometerSparseNumba` would run it (serial numba, + the gate positive, finite weights and the fraction at or below it), else `None` + (predict the full kernel). + + Computed from the mapper's triplets and the mask's extent indices, with the same + statistics the inversion's `symmetric_restriction_inputs` uses, so the gate here and + the kernel choice there cannot disagree. + """ + threshold = settings.interferometer_direct_conv_symmetric_max_work_fraction + + if ( + threshold <= 0.0 + or real_space_mask is None + or interferometer_numba_sparse._numba_parallel() + ): + return None + + if tuple(int(n) for n in real_space_mask.shape_native_masked_pixels) != tuple( + int(n) for n in extent_shape + ): + return None + + # Non-finite weights make the inversion fall back to the full kernel. + sizes = np.asarray(mapper.pix_sizes_for_sub_slim_index) + weights = np.asarray(mapper.pix_weights_for_sub_slim_index) + valid = np.arange(weights.shape[1])[None, :] < sizes[:, None] + + if not np.all(np.isfinite(weights[valid])): + return None + + extent_index_for_masked_pixel = np.asarray( + real_space_mask.extent_index_for_masked_pixel + )[np.asarray(mapper.slim_index_for_sub_slim_index)] + + work_fraction = inversion_interferometer_numba_util.symmetric_work_fraction_from( + pix_indexes_for_sub_slim_index=mapper.pix_indexes_for_sub_slim_index, + pix_sizes_for_sub_slim_index=mapper.pix_sizes_for_sub_slim_index, + extent_index_for_masked_pixel=extent_index_for_masked_pixel, + extent_shape=extent_shape, + pix_pixels=int(mapper.params), + ) + + if work_fraction > threshold: + return None + + return work_fraction diff --git a/autoarray/inversion/inversion/interferometer/abstract.py b/autoarray/inversion/inversion/interferometer/abstract.py index 9db642643..edca3928c 100644 --- a/autoarray/inversion/inversion/interferometer/abstract.py +++ b/autoarray/inversion/inversion/interferometer/abstract.py @@ -15,6 +15,28 @@ from autoarray.inversion.inversion import inversion_util +def real_space_mask_from(dataset) -> Mask2D: + """ + The real-space mask an interferometer inversion of `dataset` reconstructs on: an + `Interferometer`'s `real_space_mask`, else a transformer's `real_space_mask`, else the + dataset's own `mask` (a `DatasetInterface`'s `grids.lp.mask`). + + A module function so the inversion factory can read the same mask before it builds an + inversion (`AbstractInversionInterferometer.mask` returns this). + """ + real_space_mask = getattr(dataset, "real_space_mask", None) + + if real_space_mask is not None: + return real_space_mask + + transformer = getattr(dataset, "transformer", None) + + if transformer is not None: + return transformer.real_space_mask + + return dataset.mask + + class AbstractInversionInterferometer(AbstractInversion): def __init__( self, @@ -70,17 +92,7 @@ def mask(self) -> Mask2D: `real_space_mask`, else the dataset's own `mask` (a `DatasetInterface`'s `grids.lp.mask`). """ - real_space_mask = getattr(self.dataset, "real_space_mask", None) - - if real_space_mask is not None: - return real_space_mask - - transformer = getattr(self.dataset, "transformer", None) - - if transformer is not None: - return transformer.real_space_mask - - return self.dataset.mask + return real_space_mask_from(dataset=self.dataset) @cached_property def operated_mapping_matrix_list(self) -> List[np.ndarray]: diff --git a/autoarray/inversion/inversion/interferometer_numba/inversion_interferometer_numba_util.py b/autoarray/inversion/inversion/interferometer_numba/inversion_interferometer_numba_util.py index 215b2ab3b..065723bdb 100644 --- a/autoarray/inversion/inversion/interferometer_numba/inversion_interferometer_numba_util.py +++ b/autoarray/inversion/inversion/interferometer_numba/inversion_interferometer_numba_util.py @@ -103,6 +103,180 @@ def curvature_direct_conv( return curvature_matrix +@numba_util.jit() +def curvature_direct_conv_symmetric( + preload, + iy, + ix, + flat, + indptr, + col, + val, + cscptr, + csc_row, + csc_val, + ny, + nx, + pix_pixels, + order, + rank, + ylo_suffix, + xlo_row, + xhi_row, +): + """ + `F = Aᵀ W~ A` by the :func:`curvature_direct_conv` convolution, skipping the work + that `F`'s symmetry and the mask's footprint make redundant. + + :func:`curvature_direct_conv` convolves every source column over the whole `(ny, nx)` + extent rectangle and gathers every entry of its row of `F`. Two cuts remove work + that cannot change the result: + + 1. **Symmetry, via y-ranked columns.** The source columns are processed in the order + `order` (by the top extent row `ymin` of their non-zeros, see + :func:`symmetric_restriction_inputs_from`); `rank` is the inverse permutation. + Column `sp` of rank `r` only computes the entries `F[sp, c]` with + `rank[c] >= r`, and the others are filled by a mirror pass at the end + (`F[c, sp] = F[sp, c]`). Every pixel `A[:, c]` touches lies in an extent row + `>= ymin(c)`, so column `sp` only needs `u = W~ A[:, sp]` on the rows + `>= ylo_suffix[r] = min_{rank(c) >= r} ymin(c)` -- the rows above are neither + convolved nor gathered. + 2. **The mask's x-range per row.** `u` is only ever read at masked pixels, so each + extent row `jy` is convolved only over `[xlo_row[jy], xhi_row[jy])`, the x-range + of the masked pixels in that row (rows with no masked pixel are skipped). + + Both cuts are exact for *any* column ordering: whatever `order` is, each computed + entry `F[sp, c]` receives exactly the contributions :func:`curvature_direct_conv` + gives it (the restricted rows and columns hold every masked pixel `A[:, c]` touches), + and every entry is either computed or mirrored. The ordering only decides how much + work is skipped. Ranking by each column's top row `ymin` (stable; empty columns last) + makes `ylo_suffix[r] = ymin(order[r])`, the largest the suffix minimum can be for + every rank, so column `r` convolves only from its own top row down: about half the + rows on average for compact columns, more for tall ones (`work_fraction` measures it + exactly). On Aris's real 512x512 Delaunay mesh, where multiple imaging makes columns + tall (median span 90 of 341 rows), `ymin` order gives f = 0.51 against 0.59 for + centroid order (PyAutoArray#635). + + The mirror relies on `F` being symmetric, i.e. on `W~` being even in the pixel offset + (`preload[-dy, -dx] = preload[dy, dx]` with wrapped indices), which the + interferometer `nufft_precision_operator` is by construction. + + The gather is column-major and branch-free: for rank `r` it computes, for each later- + ranked column `c = order[k]`, `k >= r`, the dot product + `F[sp, c] = Σ_{t in column c} val_t · u[flat[row_t]]` over that column's own non-zeros + -- `Σ_{k >= r} nnz_k` multiply-adds, about half of the full kernel's `nnz` per column, + with no per-entry rank test. The CSC non-zeros are first copied into rank order + (`ranked_flat` / `ranked_val`, `O(nnz)`), so the sweep for rank `r` is one contiguous + pass over the tail of those arrays. The pixels column `c` touches lie in the rows + `>= ylo_suffix[r]` and inside the mask x-range by construction, where `u` has just + been computed. + + Each computed entry is accumulated in the same order as in + :func:`curvature_direct_conv`: the same `t` loop per extent cell for `u`, and, within + a column, the CSC non-zeros in increasing pixel index (:func:`kernel_inputs_from` + sorts them stably), which is the order that kernel's CSR sweep adds them in. So every + computed entry is bit-identical to that kernel's entry, whatever the pixel ordering. + A mirrored entry is the transposed computed one, so it equals the full kernel's entry + to round-off: the full kernel's own `F` is symmetric only to round-off. + + The kernel keeps a single `u` buffer and is serial; the parallel path uses + :func:`direct_conv_parallel_kernel`, whose columns cost the same whatever their order. + + Parameters + ---------- + preload, iy, ix, flat, indptr, col, val, cscptr, csc_row, csc_val, ny, nx, pix_pixels + As :func:`curvature_direct_conv`. + order + The source columns in processing order (`[pix_pixels]`). + rank + The inverse permutation of `order`: `rank[order[r]] = r`. + ylo_suffix + The first extent row column `order[r]` must convolve (`[pix_pixels]`, + non-decreasing; `ny` for none). + xlo_row, xhi_row + Per extent row, the half-open x-range `[xlo, xhi)` of its masked pixels + (`xhi <= xlo` for a row with none). + + Returns + ------- + The curvature matrix `F`, of shape `[pix_pixels, pix_pixels]`. + """ + nx2 = 2 * nx + + curvature_matrix = np.zeros((pix_pixels, pix_pixels)) + u = np.zeros(ny * nx) + + # The CSC non-zeros in rank order, with the pixel's extent-flat index resolved. + ranked_ptr = np.zeros(pix_pixels + 1, dtype=np.int64) + + for k in range(pix_pixels): + c = order[k] + ranked_ptr[k + 1] = ranked_ptr[k] + cscptr[c + 1] - cscptr[c] + + ranked_flat = np.empty(ranked_ptr[pix_pixels], dtype=np.int64) + ranked_val = np.empty(ranked_ptr[pix_pixels]) + + for k in range(pix_pixels): + c = order[k] + q = ranked_ptr[k] + + for t in range(cscptr[c], cscptr[c + 1]): + ranked_flat[q] = flat[csc_row[t]] + ranked_val[q] = csc_val[t] + q += 1 + + for r in range(pix_pixels): + sp = order[r] + y0 = ylo_suffix[r] + + u[y0 * nx :] = 0.0 + + for t in range(cscptr[sp], cscptr[sp + 1]): + i0 = csc_row[t] + wi = csc_val[t] + i_y = iy[i0] + i_x = ix[i0] + off = nx2 - i_x + + for jy in range(y0, ny): + x0 = xlo_row[jy] + x1 = xhi_row[jy] + + if x1 <= x0: + continue + + dy = jy - i_y + base = jy * nx + + for jx in range(x0, min(i_x, x1)): + u[base + jx] += wi * preload[dy, off + jx] + + for jx in range(max(i_x, x0), x1): + u[base + jx] += wi * preload[dy, jx - i_x] + + row = curvature_matrix[sp] + + for k in range(r, pix_pixels): + value = 0.0 + + for t in range(ranked_ptr[k], ranked_ptr[k + 1]): + value += ranked_val[t] * u[ranked_flat[t]] + + row[order[k]] = value + + for a in range(pix_pixels): + for b in range(pix_pixels): + if rank[b] > rank[a]: + curvature_matrix[b, a] = curvature_matrix[a, b] + + return curvature_matrix + + +# The `symmetric_restriction_inputs_from` keys `curvature_direct_conv_symmetric` takes after +# the `curvature_direct_conv` arguments, in argument order. +SYMMETRIC_KERNEL_KEYS = ("order", "rank", "ylo_suffix", "xlo_row", "xhi_row") + + _PARALLEL_CACHE: dict = {} @@ -280,6 +454,215 @@ def kernel_inputs_from( } +def _symmetric_restriction_from( + column_counts, column_ymin, iy, ix, ny, nx, pix_pixels +) -> dict: + """ + The column ordering, row restriction and per-row mask x-range of + :func:`curvature_direct_conv_symmetric`, plus its `work_fraction`, from per-source- + column statistics of the mapping operator and the masked pixels' extent positions. + + Shared by :func:`symmetric_restriction_inputs_from` (from the kernel's CSC arrays) and + :func:`symmetric_work_fraction_from` (from a mapper's triplets, for the factory's cost + model), so the gate and the kernel see the same work fraction: `column_ymin` is an + integer minimum, identical however it was computed, so both get the same order. + """ + pix_pixels = int(pix_pixels) + ny, nx = int(ny), int(nx) + + column_counts = np.asarray(column_counts, dtype=np.int64) + column_ymin = np.asarray(column_ymin, dtype=np.int64) + non_empty = column_counts > 0 + + # Rank by each column's top extent row, stably. An empty column has `ymin = ny`, which + # exceeds every real row, so it is ranked last and restricted to nothing. + column_ymin = np.where(non_empty, column_ymin, ny) + + order = np.argsort(column_ymin, kind="stable").astype(np.int64) + rank = np.empty(pix_pixels, dtype=np.int64) + rank[order] = np.arange(pix_pixels, dtype=np.int64) + + # Column of rank r needs the rows spanned by every column of rank >= r. In `ymin` order + # this is just ymin(order[r]); the suffix minimum keeps it right for any order. + ylo_suffix = np.ascontiguousarray( + np.minimum.accumulate(column_ymin[order][::-1])[::-1], dtype=np.int64 + ) + + iy = np.asarray(iy, dtype=np.int64) + ix = np.asarray(ix, dtype=np.int64) + + pix_by_row = np.argsort(iy, kind="stable").astype(np.int64) + pix_by_row_ptr = np.searchsorted( + iy[pix_by_row], np.arange(ny + 1, dtype=np.int64) + ).astype(np.int64) + + # Half-open masked x-range of each extent row; a row with no masked pixel keeps the + # empty range [nx, 0), which the kernel skips. + xlo_row = np.full(ny, nx, dtype=np.int64) + xhi_row = np.zeros(ny, dtype=np.int64) + + row_non_empty = np.diff(pix_by_row_ptr) > 0 + + if np.any(row_non_empty): + x_by_row = ix[pix_by_row] + starts = pix_by_row_ptr[:-1][row_non_empty] + xlo_row[row_non_empty] = np.minimum.reduceat(x_by_row, starts) + xhi_row[row_non_empty] = np.maximum.reduceat(x_by_row, starts) + 1 + + width = np.maximum(xhi_row - xlo_row, 0) + width_from_row = np.zeros(ny + 1, dtype=np.int64) + width_from_row[:ny] = np.cumsum(width[::-1])[::-1] + + nnz = int(column_counts.sum()) + + if nnz == 0 or ny * nx == 0: + work_fraction = 0.0 + else: + work = np.sum( + column_counts[order].astype(np.float64) + * width_from_row[ylo_suffix].astype(np.float64) + ) + work_fraction = float(work / (float(nnz) * float(ny) * float(nx))) + + return { + "order": order, + "rank": rank, + "ylo_suffix": ylo_suffix, + "xlo_row": xlo_row, + "xhi_row": xhi_row, + "work_fraction": work_fraction, + } + + +def symmetric_restriction_inputs_from( + iy, ix, cscptr, csc_row, ny, nx, pix_pixels +) -> dict: + """ + The extra inputs of :func:`curvature_direct_conv_symmetric` and the fraction of + :func:`curvature_direct_conv`'s convolution work it performs. + + - `order` / `rank`: the source columns sorted (stably) by their top extent row + `ymin` (the lowest row of their non-zeros), and its inverse permutation. Empty + columns are ranked last. + - `ylo_suffix[r]`: the suffix minimum, over the columns of rank `>= r`, of each + column's lowest extent row (`ny` when they are all empty). Column `order[r]` only + convolves the rows `>= ylo_suffix[r]`. + - `xlo_row` / `xhi_row`: per extent row the half-open x-range of its masked pixels + (`[nx, 0)`, i.e. empty, for a row with none). + - `work_fraction`: the convolution work of the symmetric kernel relative to the full + one, + + Σ_cols nnz_col · Σ_{y >= ylo_suffix(col)} width(y) / (nnz · ny · nx), + + with `width(y) = xhi_row[y] - xlo_row[y]`. It is a property of the geometry and the + mesh only (PyAutoArray#635): f ~ 0.41 on the synthetic Aris-512 geometry (compact + 3-NN columns), f ~ 0.51 on Aris's real 512x512 Delaunay mesh (multiply imaged, + tall columns). + + Everything is vectorised NumPy (`reduceat` over the CSC and row-sorted segments), as + it runs once per likelihood evaluation. The keys the kernel takes, in its argument + order, are :data:`SYMMETRIC_KERNEL_KEYS`. + + Parameters + ---------- + iy, ix + The extent-grid row / column of every masked pixel (:func:`kernel_inputs_from`). + cscptr, csc_row + The mapping operator's CSC column pointers and row (masked pixel) indices. + ny, nx + The extent rectangle's shape. + pix_pixels + The number of source pixels. + + Returns + ------- + A dict keyed by the kernel's extra argument names (:data:`SYMMETRIC_KERNEL_KEYS`), + plus `work_fraction`. + """ + iy = np.asarray(iy, dtype=np.int64) + cscptr = np.asarray(cscptr, dtype=np.int64) + csc_row = np.asarray(csc_row, dtype=np.int64) + + pix_pixels = int(pix_pixels) + ny = int(ny) + + column_counts = np.diff(cscptr) + non_empty = column_counts > 0 + + column_ymin = np.full(pix_pixels, ny, dtype=np.int64) + + if np.any(non_empty): + column_ymin[non_empty] = np.minimum.reduceat( + iy[csc_row], cscptr[:-1][non_empty] + ) + + return _symmetric_restriction_from( + column_counts=column_counts, + column_ymin=column_ymin, + iy=iy, + ix=ix, + ny=ny, + nx=nx, + pix_pixels=pix_pixels, + ) + + +def symmetric_work_fraction_from( + pix_indexes_for_sub_slim_index: np.ndarray, + pix_sizes_for_sub_slim_index: np.ndarray, + extent_index_for_masked_pixel: np.ndarray, + extent_shape, + pix_pixels: int, +) -> float: + """ + The `work_fraction` of :func:`symmetric_restriction_inputs_from`, straight from a + mapper's triplets, for the factory's route decision. + + It skips the CSC sort :func:`kernel_inputs_from` performs (`bincount` and + `np.minimum.at` over the valid triplets instead), and returns exactly the value the + inversion's own restriction inputs will give. + + Parameters + ---------- + pix_indexes_for_sub_slim_index, pix_sizes_for_sub_slim_index + The mapper's dense triplet arrays (the weights do not enter). + extent_index_for_masked_pixel + The mask's flat extent-grid index of every masked pixel, ordered to match the + triplet rows. + extent_shape + The `(ny, nx)` extent rectangle. + pix_pixels + The number of source pixels, i.e. `mapper.params`. + """ + flat = np.asarray(extent_index_for_masked_pixel, dtype=np.int64) + sizes = np.asarray(pix_sizes_for_sub_slim_index, dtype=np.int64) + indexes = np.asarray(pix_indexes_for_sub_slim_index, dtype=np.int64) + + ny, nx = int(extent_shape[0]), int(extent_shape[1]) + pix_pixels = int(pix_pixels) + + iy = flat // nx + ix = flat % nx + + valid = np.arange(indexes.shape[1], dtype=np.int64)[None, :] < sizes[:, None] + col = indexes[valid] + y_of_nnz = np.repeat(iy, sizes) + + column_counts = np.bincount(col, minlength=pix_pixels).astype(np.int64) + column_ymin = np.full(pix_pixels, ny, dtype=np.int64) + np.minimum.at(column_ymin, col, y_of_nnz) + + return _symmetric_restriction_from( + column_counts=column_counts, + column_ymin=column_ymin, + iy=iy, + ix=ix, + ny=ny, + nx=nx, + pix_pixels=pix_pixels, + )["work_fraction"] + + def nnz_per_source_column_from(mapper) -> float: """ The mean number of non-zero mapping weights per source column of `A`. @@ -469,8 +852,54 @@ def direct_conv_work_from(nnz: int, extent_shape) -> float: FFT_ROUTE_SECONDS_PER_UNIT = 7.62e-10 DIRECT_CONV_SECONDS_PER_UNIT = 1.77e-09 +# Time of the serial symmetric kernel (`curvature_direct_conv_symmetric`) relative to the full +# `curvature_direct_conv`, per unit of its `work_fraction` f: `t_dc` is multiplied by +# `min(1, kappa * f)` when the symmetric kernel runs. Fitted by least squares on +# log(t_sym / t_full) to quiet-core single-thread timings of this library kernel on RAL (job +# 411835, euclid-ral-compute-1, load < 1, A-C-C-A order, spreads <= 1.004; PyAutoPulse#62; +# PyAutoArray#635), synthetic Aris-512 geometry (extent 341x341, circular mask, 91k masked +# pixels, 3-NN weights, nnz 274,083), f from `symmetric_restriction_inputs_from` on that +# geometry. These rows were measured with the columns in centroid-row order, the kernel's +# order at the time; on this geometry of compact columns the `ymin` order it now uses gives +# f = 0.414 / 0.408 / 0.403, so the ordering barely matters here: +# +# S f t_full t_sym measured ratio predicted (kappa f) +# 1,428 0.416 22.04 s 12.19 s 0.553 0.548 +# 2,540 0.411 22.52 s 12.17 s 0.540 0.540 +# 5,712 0.404 23.88 s 12.59 s 0.527 0.532 +# +# Measured / predicted lies in 0.99-1.01 (log-rms 0.008). kappa > 1 because the x-range splits +# and the `S^2` mirror pass cost a little beyond the convolution work f counts; the ranked CSC +# gather (half the non-zeros per column, branch-free) keeps the ratio flat in S. With it the +# Aris 512 S = 1,459 point (measured FFT 28.3 s vs direct_conv 38.4 s on the full kernel) is +# predicted direct_conv-faster (t_fft / t_dc 0.89 -> ~1.6 at f = 0.42). +# +# Check, not part of the fit: Aris's real 512x512 data with its step-8 Delaunay mesh (S = 1,459, +# nnz 273,753, extent 341^2), laptop single-thread timings under falling load. Multiple imaging +# feeds one source vertex from 2-4 lensed images, so the columns are tall (median span 90 of +# 341 rows, against 15 on the synthetic geometry) and the ordering matters: +# +# order f t_full t_sym measured speedup predicted 1 / (kappa f) +# centroid row 0.587 40.4 s 30.5 s 1.32x 1.29x +# ymin 0.511 40.4 s 27.6 s 1.46x 1.48x +# +# Re-timed in ymin order on RAL (jobs 412490 / 412492, euclid-ral-compute-1, exclusive, load +# <= 1, compile excluded), same data, step 8 / 6 / 4 meshes: +# +# vertices f t_full t_sym measured speedup predicted 1 / (kappa f) +# 1,459 0.511 22.43 s 15.08 s 1.49x 1.48x +# 2,577 0.504 22.83 s 14.95 s 1.53x 1.50x +# 5,746 0.495 24.22 s 15.35 s 1.58x 1.53x +# +# The kernel wins while kappa f < 1, i.e. f < 0.76; the packaged gate +# (`interferometer_direct_conv_symmetric_max_work_fraction: 0.7`) keeps a predicted gain of at +# least 1.08x. +DIRECT_CONV_SYMMETRIC_WORK_FACTOR = 1.32 -def curvature_route_seconds_from(source_pixels: int, nnz: int, extent_shape): + +def curvature_route_seconds_from( + source_pixels: int, nnz: int, extent_shape, work_fraction=None +): """ Predicted single-thread seconds of the FFT route and of the `direct_conv` route. @@ -495,6 +924,13 @@ def curvature_route_seconds_from(source_pixels: int, nnz: int, extent_shape): autolens_profiling#226 extents (70^2-280^2, crossover ~60) and wrong on the Aris 512x512 mask (model crossover ~166, measured between 106 and 188). + When the serial symmetric kernel (:func:`curvature_direct_conv_symmetric`) will run, + the caller passes its `work_fraction` and + + t_dc *= min(1, DIRECT_CONV_SYMMETRIC_WORK_FACTOR · work_fraction) + + `work_fraction=None` (the default) predicts the full kernel, unchanged. + Returns ------- `(t_fft, t_direct_conv)` in seconds on the calibration machine. @@ -506,4 +942,9 @@ def curvature_route_seconds_from(source_pixels: int, nnz: int, extent_shape): nnz=nnz, extent_shape=extent_shape ) + if work_fraction is not None: + t_direct_conv *= min( + 1.0, DIRECT_CONV_SYMMETRIC_WORK_FACTOR * float(work_fraction) + ) + return t_fft, t_direct_conv diff --git a/autoarray/inversion/inversion/interferometer_numba/sparse.py b/autoarray/inversion/inversion/interferometer_numba/sparse.py index cc5364a12..d7919f0ce 100644 --- a/autoarray/inversion/inversion/interferometer_numba/sparse.py +++ b/autoarray/inversion/inversion/interferometer_numba/sparse.py @@ -184,15 +184,67 @@ def kernel_index_arrays(self) -> dict: return inputs + @cached_property + def symmetric_restriction_inputs(self) -> dict: + """ + The column ordering, row restriction and per-row mask x-range of the serial + symmetric kernel, plus its `work_fraction` + (`inversion_interferometer_numba_util.symmetric_restriction_inputs_from`). + """ + inputs = self.kernel_index_arrays + + return inversion_interferometer_numba_util.symmetric_restriction_inputs_from( + iy=inputs["iy"], + ix=inputs["ix"], + cscptr=inputs["cscptr"], + csc_row=inputs["csc_row"], + ny=inputs["ny"], + nx=inputs["nx"], + pix_pixels=inputs["pix_pixels"], + ) + + @property + def use_symmetric_kernel(self) -> bool: + """ + Whether `curvature_matrix_diag` uses the serial symmetric, row-restricted kernel + (`curvature_direct_conv_symmetric`) rather than the full `curvature_direct_conv`. + + Only on the serial path, and only when its work fraction is at or below + `Settings.interferometer_direct_conv_symmetric_max_work_fraction` (`0` = never). + The parallel path always uses the full parallel kernel. + + A mapping operator with non-finite weights (an overflowing or degenerate + `DelaunayNN` Sibson row) also uses the full kernel. The two kernels agree on + finite weights, but a NaN weight spreads through `F` along the full kernel's + column x row loops -- including into the rows of empty source columns, which the + symmetric kernel never visits -- and the evidence must come out NaN exactly as on + the FFT and mapping routes (pinned in `test_interferometer_numba.py`). The check + is one pass over the non-zeros. + """ + if _numba_parallel(): + return False + + threshold = self.settings.interferometer_direct_conv_symmetric_max_work_fraction + + if threshold <= 0.0: + return False + + if not np.all(np.isfinite(self.kernel_index_arrays["val"])): + return False + + return self.symmetric_restriction_inputs["work_fraction"] <= threshold + @cached_property def curvature_matrix_diag(self) -> np.ndarray: """ `F = Aᵀ W~ A` for the inversion's single mapper, from the `direct_conv` numba kernel. - The returned matrix is complete and symmetric (the kernel loops the full source - column x masked pixel space rather than halving on symmetry), so it drops straight - into the parent's single-mapper branch with no mirroring pass. + The returned matrix is complete and symmetric, so it drops straight into the + parent's single-mapper branch with no mirroring pass: the full kernels loop the + whole source column x masked pixel space, and the serial symmetric kernel (used + when `use_symmetric_kernel`) computes one triangle in its ranked order and mirrors + it itself. """ inputs = self.kernel_index_arrays @@ -202,12 +254,7 @@ def curvature_matrix_diag(self) -> np.ndarray: ) ) - if _numba_parallel(): - kernel = inversion_interferometer_numba_util.direct_conv_parallel_kernel() - else: - kernel = inversion_interferometer_numba_util.curvature_direct_conv - - return kernel( + args = ( preload, inputs["iy"], inputs["ix"], @@ -223,6 +270,24 @@ def curvature_matrix_diag(self) -> np.ndarray: inputs["pix_pixels"], ) + if _numba_parallel(): + return inversion_interferometer_numba_util.direct_conv_parallel_kernel()( + *args + ) + + if self.use_symmetric_kernel: + restriction = self.symmetric_restriction_inputs + + return inversion_interferometer_numba_util.curvature_direct_conv_symmetric( + *args, + *[ + restriction[key] + for key in inversion_interferometer_numba_util.SYMMETRIC_KERNEL_KEYS + ], + ) + + return inversion_interferometer_numba_util.curvature_direct_conv(*args) + def _numba_parallel() -> bool: """ diff --git a/autoarray/settings.py b/autoarray/settings.py index a69832582..58b0124a7 100644 --- a/autoarray/settings.py +++ b/autoarray/settings.py @@ -23,6 +23,7 @@ def __init__( regularization_term_method: Optional[str] = None, interferometer_numba_nnz_per_source_max: Optional[float] = None, interferometer_route: Optional[str] = None, + interferometer_direct_conv_symmetric_max_work_fraction: Optional[float] = None, positive_only_solver: Optional[str] = None, certified_pass_budget: Optional[int] = None, certified_fallback: Optional[str] = None, @@ -208,6 +209,13 @@ def __init__( (`InversionInterferometerSparseNumba`) whenever it can represent the inversion (NumPy backend, one mapper, no linear functions, no over-sampling, numba installed), overriding the geometry gate and the `0` kill switch. + interferometer_direct_conv_symmetric_max_work_fraction + When the serial numba `direct_conv` route runs, the largest geometry work + fraction at which it uses the symmetric, row-restricted kernel + (`curvature_direct_conv_symmetric`) instead of the full one. `None` (default) + reads the packaged value (`0.7`); `0` always uses the full kernel. The result + is the same either way (to round-off); only the time differs. See the property + of the same name. positive_only_solver Which solver the JAX (`xp=jnp`) positive-only reconstruction uses. `None` (default) reads the packaged value ``"pdip"``. @@ -275,6 +283,9 @@ def __init__( interferometer_numba_nnz_per_source_max ) self._interferometer_route = interferometer_route + self._interferometer_direct_conv_symmetric_max_work_fraction = ( + interferometer_direct_conv_symmetric_max_work_fraction + ) self._positive_only_solver = positive_only_solver self._certified_pass_budget = certified_pass_budget self._certified_fallback = certified_fallback @@ -290,6 +301,8 @@ def __init__( self.nnls_preconditioning_no_mapper if interferometer_route is not None: self.interferometer_route + if interferometer_direct_conv_symmetric_max_work_fraction is not None: + self.interferometer_direct_conv_symmetric_max_work_fraction @property def use_positive_only_solver(self): @@ -452,6 +465,51 @@ def interferometer_route(self) -> Optional[str]: return value + @property + def interferometer_direct_conv_symmetric_max_work_fraction(self) -> float: + """ + The work-fraction gate of the serial symmetric `direct_conv` kernel: `0.7` + (packaged default), `0` for the kill switch (always the full kernel). + + `curvature_direct_conv_symmetric` exploits `F`'s symmetry (source columns ranked by + their top extent row, each convolving only the rows later-ranked columns span) and + each row's mask x-range. Its `work_fraction` + (`inversion_interferometer_numba_util.symmetric_restriction_inputs_from`) is the + fraction of the full kernel's convolution work it still does (PyAutoArray#635): + + - the synthetic Aris-512 geometry (compact 3-NN columns): f ~ 0.41, 1.81-1.90x + faster single-threaded on RAL; + - Aris's real 512x512 Delaunay mesh (multiply imaged, tall columns): f ~ 0.51, + measured ~1.46x, as the cost model's `1 / (kappa f)` predicts. + + The cost model (`DIRECT_CONV_SYMMETRIC_WORK_FACTOR`, kappa = 1.32) says the kernel + wins while `kappa f < 1`, i.e. f < 0.76; the default 0.7 keeps a predicted gain of + at least 1.08x. With serial numba and `work_fraction <= this value` (value > 0) the + symmetric kernel is used, and the route cost model is told so; the parallel path + always uses the full kernel. Must be non-negative. + """ + value = self._interferometer_direct_conv_symmetric_max_work_fraction + + if value is None: + value = self._inversion_config_value( + "interferometer_direct_conv_symmetric_max_work_fraction", 0.7 + ) + + # A `null` config entry (the style of the neighbouring keys) means the default. + if value is None: + value = 0.7 + + value = float(value) + + if not value >= 0.0: + raise ValueError( + "interferometer_direct_conv_symmetric_max_work_fraction=" + f"{value!r} is invalid; expected a non-negative number (0 disables the " + "symmetric kernel)." + ) + + return value + def _inversion_config_value(self, key, default): # A workspace `general.yaml` normally omits the newer keys, so autoconf's config-path # list falls through to autoarray's packaged value. The fallback fires only when the diff --git a/test_autoarray/inversion/inversion/interferometer_numba/test_curvature_direct_conv_symmetric.py b/test_autoarray/inversion/inversion/interferometer_numba/test_curvature_direct_conv_symmetric.py new file mode 100644 index 000000000..8bd8779e8 --- /dev/null +++ b/test_autoarray/inversion/inversion/interferometer_numba/test_curvature_direct_conv_symmetric.py @@ -0,0 +1,693 @@ +import numpy as np +import pytest + +import autoarray as aa + +from autoarray.inversion.inversion.factory import _use_interferometer_numba +from autoarray.inversion.inversion.interferometer_numba import ( + inversion_interferometer_numba_util as numba_util, +) +from autoarray.inversion.inversion.interferometer_numba import sparse as numba_sparse +from autoarray.inversion.inversion.interferometer_numba.sparse import ( + InversionInterferometerSparseNumba, +) + +pytest.importorskip("numba") + +RESTRICTION_KEYS = list(numba_util.SYMMETRIC_KERNEL_KEYS) + + +def _prep_reference(iy, ix, cscptr, csc_row, ny, nx, pix_pixels): + """ + The straightforward (Python-loop) set-up of the PyAutoArray#635 prototype, ranking + columns by their top row `ymin` (stable, empty columns last), kept as the oracle for + the vectorised `symmetric_restriction_inputs_from`. + """ + S = pix_pixels + ymin_by_column = np.array( + [ + ( + iy[csc_row[cscptr[s] : cscptr[s + 1]]].min() + if cscptr[s + 1] > cscptr[s] + else ny + ) + for s in range(S) + ] + ) + order = np.argsort(ymin_by_column, kind="stable").astype(np.int64) + rank = np.empty(S, np.int64) + rank[order] = np.arange(S) + ymin = ymin_by_column[order] + ylo_suffix = np.minimum.accumulate(ymin[::-1])[::-1].astype(np.int64) + + # In ymin order the suffix minimum is each column's own top row. + assert np.array_equal(ylo_suffix, ymin) + xlo = np.full(ny, nx, np.int64) + xhi = np.zeros(ny, np.int64) + for y, x in zip(iy, ix): + xlo[y] = min(xlo[y], x) + xhi[y] = max(xhi[y], x + 1) + + width = np.maximum(xhi - xlo, 0) + cum = np.concatenate([np.cumsum(width[::-1])[::-1], [0]]) + nnz_col = np.diff(cscptr)[order] + work_fraction = float((nnz_col * cum[ylo_suffix]).sum() / (nnz_col.sum() * ny * nx)) + + return [order, rank, ylo_suffix, xlo, xhi], work_fraction + + +def _even_preload_from(ny, nx, seed=1): + """ + A random `(2ny, 2nx)` operator that is even in the wrapped offset, as `W~` is, so + `F = Aᵀ W~ A` is symmetric (to round-off) and the mirror pass is valid. + """ + P = np.random.default_rng(seed).standard_normal((2 * ny, 2 * nx)) + return np.ascontiguousarray(0.5 * (P + np.roll(P[::-1, ::-1], 1, axis=(0, 1)))) + + +def _mask_from(kind): + if kind == "square": + ny = nx = 9 + m = np.ones((ny, nx), dtype=bool) + elif kind == "circular": + ny = nx = 15 + yy, xx = np.mgrid[:ny, :nx] + m = (yy - 7) ** 2 + (xx - 7) ** 2 <= 7.5**2 + elif kind == "irregular": + # Two offset blobs with ragged rows and two empty rows between them. + ny, nx = 14, 12 + rng = np.random.default_rng(5) + m = np.zeros((ny, nx), dtype=bool) + m[0:6, 0:7] = rng.random((6, 7)) < 0.7 + m[8:14, 4:12] = rng.random((6, 8)) < 0.7 + m[0, 0] = m[13, 11] = True + else: + raise ValueError(kind) + + yy, xx = np.mgrid[:ny, :nx] + iy = yy[m].astype(np.int64) + ix = xx[m].astype(np.int64) + + return iy, ix, ny, nx + + +def _triplets_from(kind, iy, ix, ny, nx, seed=0): + """ + Mapper-like `[N_pix, P_max]` triplets: 3 nearest of a random vertex set (Delaunay-like) + or the 4 corners of a coarse grid cell (rectangular-like), with a few extra source + pixels no masked pixel maps to (empty columns), shuffled so the column index order is + not the ranked (`ymin`) order. + """ + rng = np.random.default_rng(seed) + n_pix = iy.shape[0] + points = np.c_[iy, ix].astype(float) + + if kind == "delaunay": + S_used = 17 + vertices = np.c_[rng.uniform(0, ny, S_used), rng.uniform(0, nx, S_used)] + d = ((points[:, None, :] - vertices[None, :, :]) ** 2).sum(-1) + indexes = np.argsort(d, axis=1)[:, :3] + sizes = np.full(n_pix, 3) + # Ragged rows: some pixels map to 1 or 2 vertices only (padded with junk). + sizes[::5] = 2 + sizes[::7] = 1 + elif kind == "rectangular": + cell = 3 + gy, gx = ny // cell + 2, nx // cell + 2 + S_used = gy * gx + cy, cx = iy // cell, ix // cell + indexes = np.c_[ + cy * gx + cx, + cy * gx + cx + 1, + (cy + 1) * gx + cx, + (cy + 1) * gx + cx + 1, + ] + sizes = np.full(n_pix, 4) + else: + raise ValueError(kind) + + weights = rng.random(indexes.shape) + weights /= weights.sum(axis=1, keepdims=True) + + S = S_used + 3 + relabel = rng.permutation(S) + indexes = relabel[indexes] + # Padding entries beyond `sizes` hold junk, as real mapper rows may. + pad = np.arange(indexes.shape[1])[None, :] >= sizes[:, None] + indexes = np.where(pad, S - 1, indexes) + + return indexes, sizes, weights, S + + +def _kernel_args_from(mask_kind, triplet_kind, seed=0): + iy, ix, ny, nx = _mask_from(mask_kind) + indexes, sizes, weights, S = _triplets_from(triplet_kind, iy, ix, ny, nx, seed) + + inputs = numba_util.kernel_inputs_from( + pix_indexes_for_sub_slim_index=indexes, + pix_sizes_for_sub_slim_index=sizes, + pix_weights_for_sub_slim_index=weights, + extent_index_for_masked_pixel=iy * nx + ix, + extent_shape=(ny, nx), + pix_pixels=S, + ) + + args = [_even_preload_from(ny, nx)] + [ + inputs[k] + for k in [ + "iy", + "ix", + "flat", + "indptr", + "col", + "val", + "cscptr", + "csc_row", + "csc_val", + "ny", + "nx", + "pix_pixels", + ] + ] + + restriction = numba_util.symmetric_restriction_inputs_from( + iy=inputs["iy"], + ix=inputs["ix"], + cscptr=inputs["cscptr"], + csc_row=inputs["csc_row"], + ny=ny, + nx=nx, + pix_pixels=S, + ) + + return args, inputs, restriction, (indexes, sizes, iy * nx + ix, (ny, nx), S) + + +GEOMETRIES = [ + ("square", "delaunay"), + ("square", "rectangular"), + ("circular", "delaunay"), + ("circular", "rectangular"), + ("irregular", "delaunay"), + ("irregular", "rectangular"), +] + + +@pytest.mark.parametrize("mask_kind, triplet_kind", GEOMETRIES) +def test__symmetric_kernel__matches_full_kernel(mask_kind, triplet_kind): + args, inputs, restriction, _ = _kernel_args_from(mask_kind, triplet_kind) + + assert np.any(np.diff(inputs["cscptr"]) == 0) # some empty source columns + + F_full = numba_util.curvature_direct_conv(*args) + F_sym = numba_util.curvature_direct_conv_symmetric( + *args, *[restriction[k] for k in RESTRICTION_KEYS] + ) + + # The entries the kernel computes (rank(col) >= rank(row)) are bit-identical to the + # full kernel's; the mirrored ones equal it to the full kernel's own asymmetry. + rank = restriction["rank"] + computed = rank[None, :] >= rank[:, None] + + assert np.array_equal(F_sym[computed], F_full[computed]) + + scale = np.abs(F_full).max() + + assert np.abs(F_sym - F_full).max() <= 1.0e-14 * scale + assert np.array_equal(F_sym, F_sym.T) + + +def test__symmetric_kernel__computed_entries_bit_identical_for_shuffled_pixel_order(): + """ + The column-major gather adds a column's non-zeros in increasing pixel index, as the + full kernel's CSR sweep does, so exactness does not need row-major pixels. + """ + iy, ix, ny, nx = _mask_from("irregular") + perm = np.random.default_rng(7).permutation(iy.shape[0]) + iy, ix = iy[perm], ix[perm] + indexes, sizes, weights, S = _triplets_from("delaunay", iy, ix, ny, nx) + + inputs = numba_util.kernel_inputs_from( + pix_indexes_for_sub_slim_index=indexes, + pix_sizes_for_sub_slim_index=sizes, + pix_weights_for_sub_slim_index=weights, + extent_index_for_masked_pixel=iy * nx + ix, + extent_shape=(ny, nx), + pix_pixels=S, + ) + args = [_even_preload_from(ny, nx)] + [ + inputs[k] + for k in [ + "iy", + "ix", + "flat", + "indptr", + "col", + "val", + "cscptr", + "csc_row", + "csc_val", + "ny", + "nx", + "pix_pixels", + ] + ] + restriction = numba_util.symmetric_restriction_inputs_from( + inputs["iy"], + inputs["ix"], + inputs["cscptr"], + inputs["csc_row"], + ny, + nx, + S, + ) + + F_full = numba_util.curvature_direct_conv(*args) + F_sym = numba_util.curvature_direct_conv_symmetric( + *args, *[restriction[k] for k in RESTRICTION_KEYS] + ) + + rank = restriction["rank"] + computed = rank[None, :] >= rank[:, None] + + assert np.array_equal(F_sym[computed], F_full[computed]) + assert np.abs(F_sym - F_full).max() <= 1.0e-14 * np.abs(F_full).max() + + +def test__symmetric_kernel__upper_triangle_bit_identical_when_columns_are_in_ymin_order(): + """ + With the source columns already in `ymin` order the ranked triangle is the index + upper triangle, so `np.triu` of the two kernels is bit-identical. + """ + args, inputs, restriction, triplets = _kernel_args_from("circular", "delaunay") + + indexes, sizes, flat, extent_shape, S = triplets + relabel = np.empty(S, dtype=np.int64) + relabel[restriction["order"]] = np.arange(S) + + rng = np.random.default_rng(0) + weights = rng.random(indexes.shape) + + inputs = numba_util.kernel_inputs_from( + pix_indexes_for_sub_slim_index=relabel[indexes], + pix_sizes_for_sub_slim_index=sizes, + pix_weights_for_sub_slim_index=weights, + extent_index_for_masked_pixel=flat, + extent_shape=extent_shape, + pix_pixels=S, + ) + args = [args[0]] + [ + inputs[k] + for k in [ + "iy", + "ix", + "flat", + "indptr", + "col", + "val", + "cscptr", + "csc_row", + "csc_val", + "ny", + "nx", + "pix_pixels", + ] + ] + restriction = numba_util.symmetric_restriction_inputs_from( + inputs["iy"], + inputs["ix"], + inputs["cscptr"], + inputs["csc_row"], + inputs["ny"], + inputs["nx"], + S, + ) + + assert np.array_equal(restriction["order"], np.arange(S)) + + F_full = numba_util.curvature_direct_conv(*args) + F_sym = numba_util.curvature_direct_conv_symmetric( + *args, *[restriction[k] for k in RESTRICTION_KEYS] + ) + + assert np.array_equal(np.triu(F_sym), np.triu(F_full)) + + +def test__symmetric_kernel__control__wrong_restriction_fails(): + """ + Control: dropping the first extent row from the restriction of the first-ranked + column must change `F`, so the equality tests above can fail. + """ + args, _, restriction, _ = _kernel_args_from("circular", "delaunay") + + ylo_suffix = restriction["ylo_suffix"].copy() + ylo_suffix[0] += 1 + extra = [restriction[k] for k in RESTRICTION_KEYS] + extra[2] = ylo_suffix + + F_full = numba_util.curvature_direct_conv(*args) + F_bad = numba_util.curvature_direct_conv_symmetric(*args, *extra) + + assert np.abs(F_bad - F_full).max() > 1.0e-6 * np.abs(F_full).max() + + +@pytest.mark.parametrize("mask_kind, triplet_kind", GEOMETRIES) +def test__symmetric_restriction_inputs_from__matches_reference(mask_kind, triplet_kind): + _, inputs, restriction, triplets = _kernel_args_from(mask_kind, triplet_kind) + + reference, work_fraction = _prep_reference( + inputs["iy"], + inputs["ix"], + inputs["cscptr"], + inputs["csc_row"], + inputs["ny"], + inputs["nx"], + inputs["pix_pixels"], + ) + + assert len(reference) == len(RESTRICTION_KEYS) + + for key, expected in zip(RESTRICTION_KEYS, reference): + assert restriction[key].dtype == np.int64 + assert np.array_equal(restriction[key], expected), key + + assert restriction["work_fraction"] == pytest.approx(work_fraction, rel=1.0e-14) + + indexes, sizes, flat, extent_shape, S = triplets + + assert ( + numba_util.symmetric_work_fraction_from( + pix_indexes_for_sub_slim_index=indexes, + pix_sizes_for_sub_slim_index=sizes, + extent_index_for_masked_pixel=flat, + extent_shape=extent_shape, + pix_pixels=S, + ) + == restriction["work_fraction"] + ) + + +def test__work_fraction__hand_checked_geometry(): + """ + Extent 4x4. Masked pixels: row 0 x = 0, 1; row 1 none; row 2 x = 1, 2, 3; row 3 x = 2. + Column 0 maps pixel (3, 2); column 1 maps (0, 0) and (2, 1); column 2 is empty. + + Top rows (ymin) 3 / 0 / empty -> order [1, 0, 2]; ylo_suffix = [0, 3, 4]. Row widths + 2 / 0 / 3 / 1, so the rows from 0 sum to 6 and from 3 to 1: + work_fraction = (2 * 6 + 1 * 1 + 0) / (3 * 16) = 13 / 48. + """ + iy = np.array([0, 0, 2, 2, 2, 3]) + ix = np.array([0, 1, 1, 2, 3, 2]) + # Pixel 0 -> column 1, pixel 2 -> column 1, pixel 5 -> column 0. + indexes = np.array([[1], [0], [1], [0], [0], [0]]) + sizes = np.array([1, 0, 1, 0, 0, 1]) + weights = np.ones((6, 1)) + + inputs = numba_util.kernel_inputs_from( + pix_indexes_for_sub_slim_index=indexes, + pix_sizes_for_sub_slim_index=sizes, + pix_weights_for_sub_slim_index=weights, + extent_index_for_masked_pixel=iy * 4 + ix, + extent_shape=(4, 4), + pix_pixels=3, + ) + + restriction = numba_util.symmetric_restriction_inputs_from( + inputs["iy"], inputs["ix"], inputs["cscptr"], inputs["csc_row"], 4, 4, 3 + ) + + assert restriction["order"].tolist() == [1, 0, 2] + assert restriction["rank"].tolist() == [1, 0, 2] + assert restriction["ylo_suffix"].tolist() == [0, 3, 4] + assert restriction["xlo_row"].tolist() == [0, 4, 1, 2] + assert restriction["xhi_row"].tolist() == [2, 0, 4, 3] + assert restriction["work_fraction"] == pytest.approx(13.0 / 48.0, rel=1.0e-15) + + args = [_even_preload_from(4, 4)] + [ + inputs[k] + for k in [ + "iy", + "ix", + "flat", + "indptr", + "col", + "val", + "cscptr", + "csc_row", + "csc_val", + "ny", + "nx", + "pix_pixels", + ] + ] + + F_full = numba_util.curvature_direct_conv(*args) + F_sym = numba_util.curvature_direct_conv_symmetric( + *args, *[restriction[k] for k in RESTRICTION_KEYS] + ) + + assert np.abs(F_sym - F_full).max() <= 1.0e-14 * np.abs(F_full).max() + + +def test__curvature_route_seconds_from__work_fraction(): + extent_shape = (341, 341) + + t_fft, t_dc = numba_util.curvature_route_seconds_from( + source_pixels=1459, nnz=273661, extent_shape=extent_shape + ) + + assert numba_util.curvature_route_seconds_from( + source_pixels=1459, nnz=273661, extent_shape=extent_shape, work_fraction=None + ) == (t_fft, t_dc) + + # Aris 512 S = 1,459 measured FFT 28.3 s vs full direct_conv 38.4 s: the full kernel is + # predicted slower ... + assert t_dc > t_fft + + # ... and the symmetric kernel, at the work fraction of Aris's real 512 Delaunay mesh + # (0.511, multiply imaged tall columns; the synthetic geometry gives 0.414), faster. + t_fft_sym, t_dc_sym = numba_util.curvature_route_seconds_from( + source_pixels=1459, + nnz=273661, + extent_shape=extent_shape, + work_fraction=0.511, + ) + + assert t_fft_sym == t_fft + assert t_dc_sym == pytest.approx( + t_dc * numba_util.DIRECT_CONV_SYMMETRIC_WORK_FACTOR * 0.511, rel=1.0e-12 + ) + assert t_dc_sym < t_fft + + # The factor never makes the prediction slower than the full kernel. + assert numba_util.curvature_route_seconds_from( + source_pixels=1459, nnz=273661, extent_shape=extent_shape, work_fraction=0.9 + ) == (t_fft, t_dc) + + +def test__settings__symmetric_max_work_fraction(): + assert aa.Settings().interferometer_direct_conv_symmetric_max_work_fraction == 0.7 + assert ( + aa.Settings( + interferometer_direct_conv_symmetric_max_work_fraction=0 + ).interferometer_direct_conv_symmetric_max_work_fraction + == 0.0 + ) + assert ( + aa.Settings( + interferometer_direct_conv_symmetric_max_work_fraction=0.7 + ).interferometer_direct_conv_symmetric_max_work_fraction + == 0.7 + ) + + with pytest.raises(ValueError): + aa.Settings(interferometer_direct_conv_symmetric_max_work_fraction=-0.1) + + +def _dataset_from(seed=0, n_visibilities=7): + """ + A 13x13 interferometer dataset with a circular mask (extent 9x9) and its + sparse-operator form, large enough for a Delaunay mesh to give a work fraction below + one. Built here rather than imported, as in `test_interferometer_numba.py`. + """ + yy, xx = np.mgrid[:13, :13] + mask = aa.Mask2D( + mask=(yy - 6) ** 2 + (xx - 6) ** 2 > 4.5**2, + pixel_scales=1.0, + ) + + rng = np.random.default_rng(seed=seed) + + dataset = aa.Interferometer( + data=aa.Visibilities(visibilities=rng.normal(size=(n_visibilities, 2))), + noise_map=aa.VisibilitiesNoiseMap(visibilities=np.ones((n_visibilities, 2))), + uv_wavelengths=rng.normal(size=(n_visibilities, 2)), + real_space_mask=mask, + transformer_class=aa.TransformerDFT, + ) + + return mask, dataset.apply_sparse_operator(use_jax=False) + + +def _delaunay_mapper_from(mask, pixels=16, seed=3): + grid = aa.Grid2D.from_mask(mask=mask, over_sample_size=1) + + mesh = aa.mesh.Delaunay(pixels=pixels) + + rng = np.random.default_rng(seed=seed) + + interpolator = mesh.interpolator_from( + source_plane_data_grid=grid, + source_plane_mesh_grid=aa.Grid2DIrregular( + rng.uniform(-4.5, 4.5, size=(pixels, 2)) + ), + ) + + return aa.Mapper( + interpolator=interpolator, + regularization=aa.reg.Constant(coefficient=1.0), + ) + + +def _log_evidence_terms_from(inversion): + return ( + float(inversion.fast_chi_squared) + + float(inversion.regularization_term) + + float(inversion.log_det_curvature_reg_matrix_term) + - float(inversion.log_det_regularization_matrix_term) + ) + + +def test__inversion__gate_selects_the_kernel_and_leaves_the_likelihood_unchanged( + monkeypatch, +): + monkeypatch.setattr(numba_sparse, "_numba_parallel", lambda: False) + + calls = {"symmetric": 0} + symmetric = numba_util.curvature_direct_conv_symmetric + + def spy(*args): + calls["symmetric"] += 1 + return symmetric(*args) + + monkeypatch.setattr(numba_util, "curvature_direct_conv_symmetric", spy) + + mask, dataset_sparse = _dataset_from() + mapper = _delaunay_mapper_from(mask=mask) + + def inversion_from(threshold): + return InversionInterferometerSparseNumba( + dataset=dataset_sparse, + linear_obj_list=[mapper], + settings=aa.Settings( + interferometer_direct_conv_symmetric_max_work_fraction=threshold + ), + xp=np, + ) + + inversion_off = inversion_from(threshold=0.0) + work_fraction = inversion_off.symmetric_restriction_inputs["work_fraction"] + + assert 0.0 < work_fraction < 1.0 + assert not inversion_off.use_symmetric_kernel + + curvature_off = np.asarray(inversion_off.curvature_matrix) + evidence_off = _log_evidence_terms_from(inversion_off) + + assert calls["symmetric"] == 0 + + inversion_below = inversion_from(threshold=0.999 * work_fraction) + assert not inversion_below.use_symmetric_kernel + np.asarray(inversion_below.curvature_matrix) + assert calls["symmetric"] == 0 + + inversion_on = inversion_from(threshold=work_fraction) + assert inversion_on.use_symmetric_kernel + + curvature_on = np.asarray(inversion_on.curvature_matrix) + evidence_on = _log_evidence_terms_from(inversion_on) + + assert calls["symmetric"] == 1 + + np.testing.assert_allclose( + curvature_on, + curvature_off, + rtol=0.0, + atol=1.0e-13 * np.abs(curvature_off).max(), + ) + assert evidence_on == pytest.approx(evidence_off, rel=1.0e-12) + + # The parallel path always uses the full parallel kernel. + monkeypatch.setattr(numba_sparse, "_numba_parallel", lambda: True) + assert not inversion_from(threshold=1.0).use_symmetric_kernel + + +def test__factory__passes_the_work_fraction_to_the_cost_model(monkeypatch): + mask, dataset_sparse = _dataset_from() + mapper = _delaunay_mapper_from(mask=mask) + + seen = [] + route_seconds = numba_util.curvature_route_seconds_from + + def spy(**kwargs): + seen.append(kwargs.get("work_fraction")) + return route_seconds(**kwargs) + + monkeypatch.setattr(numba_util, "curvature_route_seconds_from", spy) + monkeypatch.setattr(numba_sparse, "_numba_parallel", lambda: False) + + work_fraction = InversionInterferometerSparseNumba( + dataset=dataset_sparse, linear_obj_list=[mapper], xp=np + ).symmetric_restriction_inputs["work_fraction"] + + extent_shape = ( + dataset_sparse.sparse_operator.y_shape, + dataset_sparse.sparse_operator.x_shape, + ) + + def route(threshold, real_space_mask=mask): + _use_interferometer_numba( + linear_obj_list=[mapper], + settings=aa.Settings( + interferometer_direct_conv_symmetric_max_work_fraction=threshold + ), + xp=np, + extent_shape=extent_shape, + real_space_mask=real_space_mask, + ) + return seen[-1] + + assert route(threshold=1.0) == work_fraction + assert route(threshold=work_fraction) == work_fraction + assert route(threshold=0.999 * work_fraction) is None + assert route(threshold=0.0) is None + assert route(threshold=1.0, real_space_mask=None) is None + + monkeypatch.setattr(numba_sparse, "_numba_parallel", lambda: True) + assert route(threshold=1.0) is None + + +def test__inversion__non_finite_weights_use_the_full_kernel(monkeypatch): + """ + NaN mapping weights must give the full kernel's NaN pattern (and so a NaN evidence), + which the symmetric kernel does not reproduce on empty source columns' rows. + """ + monkeypatch.setattr(numba_sparse, "_numba_parallel", lambda: False) + + mask, dataset_sparse = _dataset_from() + mapper = _delaunay_mapper_from(mask=mask) + + inversion = InversionInterferometerSparseNumba( + dataset=dataset_sparse, + linear_obj_list=[mapper], + settings=aa.Settings( + interferometer_direct_conv_symmetric_max_work_fraction=1.0 + ), + xp=np, + ) + + assert inversion.use_symmetric_kernel + + inversion.kernel_index_arrays["val"][0] = np.nan + + assert not inversion.use_symmetric_kernel diff --git a/test_autoarray/inversion/inversion/interferometer_numba/test_interferometer_numba.py b/test_autoarray/inversion/inversion/interferometer_numba/test_interferometer_numba.py index b6fa0cca7..2ce79da08 100644 --- a/test_autoarray/inversion/inversion/interferometer_numba/test_interferometer_numba.py +++ b/test_autoarray/inversion/inversion/interferometer_numba/test_interferometer_numba.py @@ -709,6 +709,13 @@ def wrapper(*args, **kwargs): monkeypatch.setattr( numba_util, "curvature_direct_conv", counted(numba_util.curvature_direct_conv) ) + # The serial path may run the symmetric kernel instead (gate + # `interferometer_direct_conv_symmetric_max_work_fraction`); either counts once. + monkeypatch.setattr( + numba_util, + "curvature_direct_conv_symmetric", + counted(numba_util.curvature_direct_conv_symmetric), + ) parallel_kernel = numba_util.direct_conv_parallel_kernel monkeypatch.setattr( numba_util,