diff --git a/autoarray/__init__.py b/autoarray/__init__.py index 4944c874..4a81a3b2 100644 --- a/autoarray/__init__.py +++ b/autoarray/__init__.py @@ -8,7 +8,7 @@ from . import type from . import util from . import fixtures -from . import mock as m +from . import mock as m from .dataset import preprocess from .dataset.abstract.dataset import AbstractDataset diff --git a/autoarray/dataset/grids.py b/autoarray/dataset/grids.py index 666e47fe..60cf8bae 100644 --- a/autoarray/dataset/grids.py +++ b/autoarray/dataset/grids.py @@ -17,7 +17,6 @@ def __init__( over_sample_size_lp: Union[int, Array2D], over_sample_size_pixelization: Union[int, Array2D], psf: Optional[Kernel2D] = None, - use_sparse_operator: bool = False, ): """ Contains grids of (y,x) Cartesian coordinates at the centre of every pixel in the dataset's image and @@ -66,8 +65,6 @@ def __init__( self._blurring = None self._border_relocator = None - self.use_sparse_operator = use_sparse_operator - @property def lp(self): @@ -120,7 +117,6 @@ def border_relocator(self) -> BorderRelocator: self._border_relocator = BorderRelocator( mask=self.mask, sub_size=self.over_sample_size_pixelization, - use_sparse_operator=self.use_sparse_operator, ) return self._border_relocator diff --git a/autoarray/dataset/imaging/dataset.py b/autoarray/dataset/imaging/dataset.py index 8ef91049..cc278b67 100644 --- a/autoarray/dataset/imaging/dataset.py +++ b/autoarray/dataset/imaging/dataset.py @@ -192,14 +192,11 @@ def __init__( if psf.mask.shape[0] % 2 == 0 or psf.mask.shape[1] % 2 == 0: raise exc.KernelException("Kernel2D Kernel2D must be odd") - use_sparse_operator = True if sparse_operator is not None else False - self.grids = GridsDataset( mask=self.data.mask, over_sample_size_lp=self.over_sample_size_lp, over_sample_size_pixelization=self.over_sample_size_pixelization, psf=self.psf, - use_sparse_operator=use_sparse_operator, ) self.sparse_operator = sparse_operator diff --git a/autoarray/dataset/interferometer/dataset.py b/autoarray/dataset/interferometer/dataset.py index 7e72ddcd..7c52585e 100644 --- a/autoarray/dataset/interferometer/dataset.py +++ b/autoarray/dataset/interferometer/dataset.py @@ -96,13 +96,10 @@ def __init__( real_space_mask=real_space_mask, ) - use_sparse_operator = True if sparse_operator is not None else False - self.grids = GridsDataset( mask=self.real_space_mask, over_sample_size_lp=self.over_sample_size_lp, over_sample_size_pixelization=self.over_sample_size_pixelization, - use_sparse_operator=use_sparse_operator, ) self.sparse_operator = sparse_operator diff --git a/autoarray/inversion/inversion/imaging_numba/inversion_imaging_numba_util.py b/autoarray/inversion/inversion/imaging_numba/inversion_imaging_numba_util.py index 7b2ba971..7baa7645 100644 --- a/autoarray/inversion/inversion/imaging_numba/inversion_imaging_numba_util.py +++ b/autoarray/inversion/inversion/imaging_numba/inversion_imaging_numba_util.py @@ -803,74 +803,6 @@ def mapped_reconstructed_data_via_image_to_pix_unique_from( return mapped_reconstructed_data -@numba_util.jit() -def relocated_grid_via_jit_from(grid, border_grid): - """ - Relocate the coordinates of a grid to its border if they are outside the border, where the border is - defined as all pixels at the edge of the grid's mask (see *mask._border_1d_indexes*). - - This is performed as follows: - - 1: Use the mean value of the grid's y and x coordinates to determine the origin of the grid. - 2: Compute the radial distance of every grid coordinate from the origin. - 3: For every coordinate, find its nearest pixel in the border. - 4: Determine if it is outside the border, by comparing its radial distance from the origin to its paired - border pixel's radial distance. - 5: If its radial distance is larger, use the ratio of radial distances to move the coordinate to the - border (if its inside the border, do nothing). - - The method can be used on uniform or irregular grids, however for irregular grids the border of the - 'image-plane' mask is used to define border pixels. - - Parameters - ---------- - grid - The grid (uniform or irregular) whose pixels are to be relocated to the border edge if outside it. - border_grid : Grid2D - The grid of border (y,x) coordinates. - """ - - grid_relocated = np.zeros(grid.shape) - grid_relocated[:, :] = grid[:, :] - - border_origin = np.zeros(2) - border_origin[0] = np.mean(border_grid[:, 0]) - border_origin[1] = np.mean(border_grid[:, 1]) - border_grid_radii = np.sqrt( - np.add( - np.square(np.subtract(border_grid[:, 0], border_origin[0])), - np.square(np.subtract(border_grid[:, 1], border_origin[1])), - ) - ) - border_min_radii = np.min(border_grid_radii) - - grid_radii = np.sqrt( - np.add( - np.square(np.subtract(grid[:, 0], border_origin[0])), - np.square(np.subtract(grid[:, 1], border_origin[1])), - ) - ) - - for pixel_index in range(grid.shape[0]): - if grid_radii[pixel_index] > border_min_radii: - closest_pixel_index = np.argmin( - np.square(grid[pixel_index, 0] - border_grid[:, 0]) - + np.square(grid[pixel_index, 1] - border_grid[:, 1]) - ) - - move_factor = ( - border_grid_radii[closest_pixel_index] / grid_radii[pixel_index] - ) - - if move_factor < 1.0: - grid_relocated[pixel_index, :] = ( - move_factor * (grid[pixel_index, :] - border_origin[:]) - + border_origin[:] - ) - - return grid_relocated - - class SparseLinAlgImagingNumba: def __init__( self, diff --git a/autoarray/inversion/pixelization/border_relocator.py b/autoarray/inversion/pixelization/border_relocator.py index b5d9fe87..4ec44def 100644 --- a/autoarray/inversion/pixelization/border_relocator.py +++ b/autoarray/inversion/pixelization/border_relocator.py @@ -203,68 +203,115 @@ def sub_border_slim_from(mask, sub_size): ).astype("int") -def relocated_grid_from(grid, border_grid, xp=np): +def ellipse_params_via_border_pca_from(border_grid, xp=np, eps=1e-12): """ - Relocate the coordinates of a grid to its border if they are outside the border, where the border is - defined as all pixels at the edge of the grid's mask (see *mask._border_1d_indexes*). + Estimate origin, semi-axes (a,b), and rotation phi from border points using PCA. + Works well for circle/ellipse-like borders. - This is performed as follows: + Parameters + ---------- + border_grid : (M,2) + Border coordinates in (y, x) order. + xp : module + numpy-like module (np, jnp, cupy, etc.). + eps : float + Numerical safety epsilon. + + Returns + ------- + origin : (2,) + Estimated center (y0, x0). + a : float + Semi-major axis. + b : float + Semi-minor axis. + phi : float + Rotation angle in radians. + """ + + origin = xp.mean(border_grid, axis=0) + + dy = border_grid[:, 0] - origin[0] + dx = border_grid[:, 1] - origin[1] + + # Build data matrix in (x, y) order for PCA + X = xp.stack([dx, dy], axis=1) # (M,2) + + # Covariance matrix + denom = xp.maximum(X.shape[0] - 1, 1) + C = (X.T @ X) / denom # (2,2) + + # Eigen-decomposition (ascending eigenvalues) + evals, evecs = xp.linalg.eigh(C) + + # Major axis = eigenvector with largest eigenvalue + v_major = evecs[:, -1] # (2,) in (x,y) + + phi = xp.arctan2(v_major[1], v_major[0]) + + # Rotate border points into ellipse-aligned frame + c = xp.cos(phi) + s = xp.sin(phi) + + xprime = c * dx + s * dy + yprime = -s * dx + c * dy + + # Semi-axes from maximal extent + a = xp.max(xp.abs(xprime)) + eps + b = xp.max(xp.abs(yprime)) + eps - 1: Use the mean value of the grid's y and x coordinates to determine the origin of the grid. - 2: Compute the radial distance of every grid coordinate from the origin. - 3: For every coordinate, find its nearest pixel in the border. - 4: Determine if it is outside the border, by comparing its radial distance from the origin to its paired - border pixel's radial distance. - 5: If its radial distance is larger, use the ratio of radial distances to move the coordinate to the - border (if its inside the border, do nothing). + return origin, a, b, phi - The method can be used on uniform or irregular grids, however for irregular grids the border of the - 'image-plane' mask is used to define border pixels. + +def relocated_grid_via_ellipse_border_from(grid, origin, a, b, phi, xp=np, eps=1e-12): + """ + Rotated ellipse centered at origin with semi-axes a (major, x'), b (minor, y'), + rotated by phi radians (counterclockwise). Parameters ---------- - grid - The grid (uniform or irregular) whose pixels are to be relocated to the border edge if outside it. - border_grid : Grid2D - The grid of border (y,x) coordinates. + grid : (N,2) + Coordinates in (y, x) order. + origin : (2,) + Ellipse center (y0, x0). + a, b : float + Semi-major and semi-minor axes. + phi : float + Rotation angle in radians. + xp : module + numpy-like module (np, jnp, cupy, etc.). + eps : float + Numerical safety epsilon. """ - # Compute origin (center) of the border grid - border_origin = xp.mean(border_grid, axis=0) + # shift to origin + dy = grid[:, 0] - origin[0] + dx = grid[:, 1] - origin[1] - # Radii from origin - grid_radii = xp.linalg.norm(grid - border_origin, axis=1) # (N,) - border_radii = xp.linalg.norm(border_grid - border_origin, axis=1) # (M,) - border_min_radius = xp.min(border_radii) + c = xp.cos(phi) + s = xp.sin(phi) - # Determine which points are outside - outside_mask = grid_radii > border_min_radius # (N,) + # rotate into ellipse-aligned frame + xprime = c * dx + s * dy + yprime = -s * dx + c * dy - # To compute nearest border point for each grid point, we must do it for all and then mask later - # Compute all distances: (N, M) - diffs = grid[:, None, :] - border_grid[None, :, :] # (N, M, 2) - dists_squared = xp.sum(diffs**2, axis=2) # (N, M) - closest_indices = xp.argmin(dists_squared, axis=1) # (N,) + # ellipse radius in normalized coords + q = (xprime / a) ** 2 + (yprime / b) ** 2 - # Get border radius for closest border point to each grid point - matched_border_radii = border_radii[closest_indices] # (N,) + outside = q > 1.0 + scale = 1.0 / xp.sqrt(xp.maximum(q, 1.0 + eps)) - # Ratio of border to grid radius - move_factors = matched_border_radii / grid_radii # (N,) + # scale back to boundary + xprime2 = xprime * scale + yprime2 = yprime * scale - # Only move if: - # - the point is outside the border - # - the matched border point is closer to the origin (i.e. move_factor < 1) - apply_move = xp.logical_and(outside_mask, move_factors < 1.0) # (N,) + # rotate back to original frame + dx2 = c * xprime2 - s * yprime2 + dy2 = s * xprime2 + c * yprime2 - # Compute moved positions (for all points, but will select with mask) - direction_vectors = grid - border_origin # (N, 2) - moved_grid = move_factors[:, None] * direction_vectors + border_origin # (N, 2) + moved = xp.stack([origin[0] + dy2, origin[1] + dx2], axis=1) - # Select which grid points to move - relocated_grid = xp.where(apply_move[:, None], moved_grid, grid) # (N, 2) - - return relocated_grid + return xp.where(outside[:, None], moved, grid) class BorderRelocator: @@ -272,7 +319,6 @@ def __init__( self, mask: Mask2D, sub_size: Union[int, Array2D], - use_sparse_operator: bool = False, ): """ Relocates source plane coordinates that trace outside the mask’s border in the source-plane back onto the @@ -330,8 +376,6 @@ def __init__( self.sub_border_grid = sub_grid[self.sub_border_slim] - self.use_sparse_operator = use_sparse_operator - def relocated_grid_from(self, grid: Grid2D, xp=np) -> Grid2D: """ Relocate the coordinates of a grid to the border of this grid if they are outside the border, where the @@ -359,32 +403,17 @@ def relocated_grid_from(self, grid: Grid2D, xp=np) -> Grid2D: if len(self.sub_border_grid) == 0: return grid - if self.use_sparse_operator is False or xp.__name__.startswith("jax"): - - values = relocated_grid_from( - grid=grid.array, border_grid=grid.array[self.border_slim], xp=xp - ) - - over_sampled = relocated_grid_from( - grid=grid.over_sampled.array, - border_grid=grid.over_sampled.array[self.sub_border_slim], - xp=xp, - ) - - else: - - from autoarray.inversion.inversion.imaging_numba import ( - inversion_imaging_numba_util, - ) + origin, a, b, phi = ellipse_params_via_border_pca_from( + grid.array[self.border_slim], xp=xp + ) - values = inversion_imaging_numba_util.relocated_grid_via_jit_from( - grid=grid.array, border_grid=grid[self.border_slim] - ) + values = relocated_grid_via_ellipse_border_from( + grid=grid.array, origin=origin, a=a, b=b, phi=phi, xp=xp + ) - over_sampled = inversion_imaging_numba_util.relocated_grid_via_jit_from( - grid=grid.over_sampled.array, - border_grid=grid.over_sampled[self.sub_border_slim], - ) + over_sampled = relocated_grid_via_ellipse_border_from( + grid=grid.over_sampled.array, origin=origin, a=a, b=b, phi=phi, xp=xp + ) return Grid2D( values=values, @@ -411,24 +440,13 @@ def relocated_mesh_grid_from( if len(self.sub_border_grid) == 0: return mesh_grid - if self.use_sparse_operator is False or xp.__name__.startswith("jax"): - - relocated_grid = relocated_grid_from( - grid=mesh_grid.array, - border_grid=grid[self.sub_border_slim], - xp=xp, - ) - - else: - - from autoarray.inversion.inversion.imaging import ( - inversion_imaging_numba_util, - ) + origin, a, b, phi = ellipse_params_via_border_pca_from( + grid.array[self.border_slim], xp=xp + ) - relocated_grid = inversion_imaging_numba_util.relocated_grid_via_jit_from( - grid=mesh_grid.array, - border_grid=grid[self.sub_border_slim], - ) + relocated_grid = relocated_grid_via_ellipse_border_from( + grid=mesh_grid.array, origin=origin, a=a, b=b, phi=phi, xp=xp + ) return Grid2DIrregular( values=relocated_grid, diff --git a/autoarray/inversion/pixelization/mesh/delaunay.py b/autoarray/inversion/pixelization/mesh/delaunay.py index 603f3742..ba7ca379 100644 --- a/autoarray/inversion/pixelization/mesh/delaunay.py +++ b/autoarray/inversion/pixelization/mesh/delaunay.py @@ -126,20 +126,17 @@ def mapper_grids_from( relocated_mesh_grid = self.relocated_mesh_grid_from( border_relocator=border_relocator, - source_plane_data_grid=relocated_grid.over_sampled, + source_plane_data_grid=Grid2DIrregular(relocated_grid.over_sampled), source_plane_mesh_grid=source_plane_mesh_grid, xp=xp, ) - try: - source_plane_mesh_grid = self.mesh_grid_from( - source_plane_data_grid=relocated_grid.over_sampled, - source_plane_mesh_grid=relocated_mesh_grid, - preloads=preloads, - xp=xp, - ) - except ValueError as e: - raise e + source_plane_mesh_grid = self.mesh_grid_from( + source_plane_data_grid=relocated_grid.over_sampled, + source_plane_mesh_grid=relocated_mesh_grid, + preloads=preloads, + xp=xp, + ) return MapperGrids( mask=mask, diff --git a/test_autoarray/inversion/pixelization/mesh/test_abstract.py b/test_autoarray/inversion/pixelization/mesh/test_abstract.py index c00b6817..c408cbff 100644 --- a/test_autoarray/inversion/pixelization/mesh/test_abstract.py +++ b/test_autoarray/inversion/pixelization/mesh/test_abstract.py @@ -18,6 +18,8 @@ def test__grid_is_relocated_via_border(grid_2d_7x7): ) image_mesh = image_mesh.image_plane_mesh_grid_from(mask=mask, adapt_data=None) + image_mesh = aa.Grid2DIrregular(image_mesh) + grid[8, 0] = 100.0 border_relocator = aa.BorderRelocator(mask=mask, sub_size=1) diff --git a/test_autoarray/inversion/pixelization/test_border_relocator.py b/test_autoarray/inversion/pixelization/test_border_relocator.py index db190269..2b6c3cd6 100644 --- a/test_autoarray/inversion/pixelization/test_border_relocator.py +++ b/test_autoarray/inversion/pixelization/test_border_relocator.py @@ -356,7 +356,7 @@ def test__relocated_grid_from__outside_border_includes_relocations(): relocated_grid = border_relocator.relocated_grid_from(grid=grid) assert relocated_grid.over_sampled[1] == pytest.approx( - [0.97783243, 0.00968151], 1e-4 + [0.949953439, 0.009405479597], 1e-4 ) @@ -375,4 +375,4 @@ def test__relocated_grid_from__positive_origin_included_in_relocate(): relocated_grid = border_relocator.relocated_grid_from(grid=grid) - assert relocated_grid.over_sampled[1] == pytest.approx([1.97783243, 1.0], 1e-4) + assert relocated_grid.over_sampled[1] == pytest.approx([1.95, 1.0], 1e-4)