From 33a3ccb69b17af77457ca0d000539e942852a25f Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Thu, 3 Apr 2025 14:50:01 +0100 Subject: [PATCH 01/25] fix some decorator unit tests --- autoarray/structures/decorators/abstract.py | 2 +- autoarray/structures/decorators/to_grid.py | 3 +-- autoarray/structures/mock/mock_decorators.py | 6 +++--- test_autoarray/structures/decorators/test_to_grid.py | 8 ++++---- 4 files changed, 9 insertions(+), 10 deletions(-) diff --git a/autoarray/structures/decorators/abstract.py b/autoarray/structures/decorators/abstract.py index c9e5fca87..e033815b4 100644 --- a/autoarray/structures/decorators/abstract.py +++ b/autoarray/structures/decorators/abstract.py @@ -1,4 +1,4 @@ -from typing import List, Union +from typing import Union import numpy as np diff --git a/autoarray/structures/decorators/to_grid.py b/autoarray/structures/decorators/to_grid.py index 144137c69..4797c37ce 100644 --- a/autoarray/structures/decorators/to_grid.py +++ b/autoarray/structures/decorators/to_grid.py @@ -1,6 +1,5 @@ -from autoarray.numpy_wrapper import np from functools import wraps - +import numpy as np from typing import List, Union from autoarray.structures.decorators.abstract import AbstractMaker diff --git a/autoarray/structures/mock/mock_decorators.py b/autoarray/structures/mock/mock_decorators.py index 876b456d7..013b0a62a 100644 --- a/autoarray/structures/mock/mock_decorators.py +++ b/autoarray/structures/mock/mock_decorators.py @@ -116,7 +116,7 @@ def ndarray_2d_from(self, grid, *args, **kwargs): Such functions are common in **PyAutoGalaxy** for light and mass profile objects. """ - return np.multiply(2.0, grid) + return np.multiply(2.0, grid.array) @decorators.to_vector_yx def ndarray_yx_2d_from(self, grid, *args, **kwargs): @@ -146,7 +146,7 @@ def ndarray_2d_list_from(self, grid, *args, **kwargs): Such functions are common in **PyAutoGalaxy** for light and mass profile objects. """ - return [np.multiply(1.0, grid), np.multiply(2.0, grid)] + return [np.multiply(1.0, grid.array), np.multiply(2.0, grid.array)] @decorators.to_vector_yx def ndarray_yx_2d_list_from(self, grid, *args, **kwargs): @@ -156,7 +156,7 @@ def ndarray_yx_2d_list_from(self, grid, *args, **kwargs): Such functions are common in **PyAutoGalaxy** for light and mass profile objects. """ - return [np.multiply(1.0, grid), np.multiply(2.0, grid)] + return [np.multiply(1.0, grid.array), np.multiply(2.0, grid.array)] class MockGridRadialMinimum: diff --git a/test_autoarray/structures/decorators/test_to_grid.py b/test_autoarray/structures/decorators/test_to_grid.py index 60c70d71b..2e8b1be2f 100644 --- a/test_autoarray/structures/decorators/test_to_grid.py +++ b/test_autoarray/structures/decorators/test_to_grid.py @@ -15,11 +15,11 @@ def test__in_grid_1d__out_ndarray_2d(): assert isinstance(ndarray_2d, aa.Grid2D) assert ndarray_2d.native == pytest.approx( - np.array([[[0.0, 0.0], [0.0, -1.0], [0.0, 1.0], [0.0, 0.0]]]), 1.0e-4 + np.array([[[0.0, 0.0], [0.0, -1.0], [0.0, 1.0], [0.0, 0.0]]]), abs=1.0e-4 ) -def test__in_grid_1d__out_ndarray_2d_list(): +def test__in_dgrid_1d__out_ndarray_2d_list(): mask = aa.Mask1D(mask=[True, False, False, True], pixel_scales=(1.0,)) grid_1d = aa.Grid1D.from_mask(mask=mask) @@ -30,12 +30,12 @@ def test__in_grid_1d__out_ndarray_2d_list(): assert isinstance(ndarray_2d_list[0], aa.Grid2D) assert ndarray_2d_list[0].native == pytest.approx( - np.array([[[0.0, 0.0], [0.0, -0.5], [0.0, 0.5], [0.0, 0.0]]]), 1.0e-4 + np.array([[[0.0, 0.0], [0.0, -0.5], [0.0, 0.5], [0.0, 0.0]]]), abs=1.0e-4 ) assert isinstance(ndarray_2d_list[1], aa.Grid2D) assert ndarray_2d_list[1].native == pytest.approx( - np.array([[[0.0, 0.0], [0.0, -1.0], [0.0, 1.0], [0.0, 0.0]]]), 1.0e-4 + np.array([[[0.0, 0.0], [0.0, -1.0], [0.0, 1.0], [0.0, 0.0]]]), abs=1.0e-4 ) From 8b8dc9e9479dc17d9e91e5631e1033a75766a2bd Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Thu, 3 Apr 2025 14:59:45 +0100 Subject: [PATCH 02/25] removing numpy wrapper to do explicit impots --- autoarray/abstract_ndarray.py | 17 +++++++++-------- .../operators/over_sampling/over_sampler.py | 13 +++++++------ 2 files changed, 16 insertions(+), 14 deletions(-) diff --git a/autoarray/abstract_ndarray.py b/autoarray/abstract_ndarray.py index 8b5fbc00e..ded8c5452 100644 --- a/autoarray/abstract_ndarray.py +++ b/autoarray/abstract_ndarray.py @@ -4,10 +4,11 @@ from abc import ABC from abc import abstractmethod +import jax.numpy as jnp from autoconf.fitsable import output_to_fits -from autoarray.numpy_wrapper import np, register_pytree_node, Array +from autoarray.numpy_wrapper import register_pytree_node, Array from typing import TYPE_CHECKING @@ -82,7 +83,7 @@ def __init__(self, array): def invert(self): new = self.copy() - new._array = np.invert(new._array) + new._array = jnp.invert(new._array) return new @classmethod @@ -104,7 +105,7 @@ def instance_flatten(cls, instance): @staticmethod def flip_hdu_for_ds9(values): if conf.instance["general"]["fits"]["flip_for_ds9"]: - return np.flipud(values) + return jnp.flipud(values) return values @classmethod @@ -117,7 +118,7 @@ def instance_unflatten(cls, aux_data, children): setattr(instance, key, value) return instance - def with_new_array(self, array: np.ndarray) -> "AbstractNDArray": + def with_new_array(self, array: jnp.ndarray) -> "AbstractNDArray": """ Copy this object but give it a new array. @@ -164,7 +165,7 @@ def __iter__(self): @to_new_array def sqrt(self): - return np.sqrt(self._array) + return jnp.sqrt(self._array) @property def array(self): @@ -330,13 +331,13 @@ def __getitem__(self, item): result = self._array[item] if isinstance(item, slice): result = self.with_new_array(result) - if isinstance(result, np.ndarray): + if isinstance(result, jnp.ndarray): result = self.with_new_array(result) return result def __setitem__(self, key, value): - if isinstance(key, (np.ndarray, AbstractNDArray, Array)): - self._array = np.where(key, value, self._array) + if isinstance(key, (jnp.ndarray, AbstractNDArray, Array)): + self._array = jnp.where(key, value, self._array) else: self._array[key] = value diff --git a/autoarray/operators/over_sampling/over_sampler.py b/autoarray/operators/over_sampling/over_sampler.py index 6f12f4b9f..65393709c 100644 --- a/autoarray/operators/over_sampling/over_sampler.py +++ b/autoarray/operators/over_sampling/over_sampler.py @@ -1,5 +1,6 @@ -from autoarray.numpy_wrapper import np -from typing import List, Tuple, Union +import numpy as np +import jax.numpy as jnp +from typing import Union from autoconf import conf from autoconf import cached_property @@ -184,7 +185,7 @@ def sub_pixel_areas(self) -> np.ndarray: """ The area of every sub-pixel in the mask. """ - sub_pixel_areas = np.zeros(self.sub_total) + sub_pixel_areas = jnp.zeros(self.sub_total) k = 0 @@ -221,9 +222,9 @@ def binned_array_2d_from(self, array: Array2D) -> "Array2D": pass # binned_array_2d = over_sample_util.binned_array_2d_from( - # array_2d=np.array(array), - # mask_2d=np.array(self.mask), - # sub_size=np.array(self.sub_size).astype("int"), + # array_2d=jnp.array(array), + # mask_2d=jnp.array(self.mask), + # sub_size=jnp.array(self.sub_size).astype("int"), # ) binned_array_2d = array.reshape( From 7115f9cbc5893d768e52a08b4d1ff2c313b39e9f Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Thu, 3 Apr 2025 17:19:22 +0100 Subject: [PATCH 03/25] move relocate radial --- autoarray/config/grids.yaml | 3 - autoarray/structures/decorators/__init__.py | 1 - .../structures/decorators/relocate_radial.py | 106 ------------------ autoarray/structures/mock/mock_decorators.py | 4 - test_autoarray/config/grids.yaml | 12 -- 5 files changed, 126 deletions(-) delete mode 100644 autoarray/config/grids.yaml delete mode 100644 autoarray/structures/decorators/relocate_radial.py delete mode 100644 test_autoarray/config/grids.yaml diff --git a/autoarray/config/grids.yaml b/autoarray/config/grids.yaml deleted file mode 100644 index f82eaa5df..000000000 --- a/autoarray/config/grids.yaml +++ /dev/null @@ -1,3 +0,0 @@ -radial_minimum: - function_name: - class_name: 1.0e-08 diff --git a/autoarray/structures/decorators/__init__.py b/autoarray/structures/decorators/__init__.py index d85cbe6ee..1efb9137e 100644 --- a/autoarray/structures/decorators/__init__.py +++ b/autoarray/structures/decorators/__init__.py @@ -4,4 +4,3 @@ from .to_grid import to_grid from .to_vector_yx import to_vector_yx from .transform import transform -from .relocate_radial import relocate_to_radial_minimum diff --git a/autoarray/structures/decorators/relocate_radial.py b/autoarray/structures/decorators/relocate_radial.py deleted file mode 100644 index 58411714f..000000000 --- a/autoarray/structures/decorators/relocate_radial.py +++ /dev/null @@ -1,106 +0,0 @@ -from autoarray.numpy_wrapper import np, use_jax -from functools import wraps - -from typing import Union - -from autoconf.exc import ConfigException - -from autoarray.structures.grids.irregular_2d import Grid2DIrregular -from autoarray.structures.grids.uniform_2d import Grid2D -from autoconf import conf - - -def relocate_to_radial_minimum(func): - """ - Checks whether any coordinates in the grid are radially near (0.0, 0.0), which can lead to numerical faults in - the evaluation of a function (e.g. numerical integration reaching a singularity at (0.0, 0.0)). - - If any coordinates are radially within the radial minimum threshold, their (y,x) coordinates are shifted to that - value to ensure they are evaluated at that coordinate. - - The value the (y,x) coordinates are rounded to is set in the 'radial_minimum.yaml' config. - - Parameters - ---------- - func - A function that takes a grid of coordinates which may have a singularity as (0.0, 0.0) - - Returns - ------- - A function that has an input grid whose radial coordinates are relocated to the radial minimum. - """ - - @wraps(func) - def wrapper( - obj: object, - grid: Union[np.ndarray, Grid2D, Grid2DIrregular], - *args, - **kwargs, - ) -> Union[np.ndarray, Grid2D, Grid2DIrregular]: - """ - Checks whether any coordinates in the grid are radially near (0.0, 0.0), which can lead to numerical faults in - the evaluation of a function (e.g. numerical integration reaching a singularity at (0.0, 0.0)). - - If any coordinates are radially within the radial minimum threshold, their (y,x) coordinates are shifted to that - value to ensure they are evaluated at that coordinate. - - The value the (y,x) coordinates are rounded to is set in the 'radial_minimum.yaml' config. - - Parameters - ---------- - obj - An object whose function uses grid_like inputs to compute quantities at every coordinate on the grid. - grid - The (y, x) coordinates which are to be radially moved from (0.0, 0.0). - - Returns - ------- - The grid_like object whose coordinates are radially moved from (0.0, 0.0). - """ - if use_jax: - return func(obj, grid, *args, **kwargs) - - try: - grid_radial_minimum = conf.instance["grids"]["radial_minimum"][ - "radial_minimum" - ][obj.__class__.__name__] - - except KeyError as e: - raise ConfigException( - rf""" - The {obj.__class__.__name__} profile you are using does not have a corresponding - entry in the `config/grid.yaml` config file. - - When a profile is evaluated at (0.0, 0.0), they commonly break due to numericalinstabilities (e.g. - division by zero). To prevent this, the code relocates the (y,x) coordinates of the grid to a - minimum radial value, specified in the `config/grids.yaml` config file. - - For example, if the value in `grid.yaml` is `radial_minimum: 1e-6`, then any (y,x) coordinates - with a radial distance less than 1e-6 to (0.0, 0.0) are relocated to 1e-6. - - For a profile to be used it must have an entry in the `config/grids.yaml` config file. Go to this - file now and add your profile to the `radial_minimum` section. Adopting a value of 1e-6 is a good - default choice. - - If you are going to make a pull request to add your profile to the source code, you should also - add an entry to the `config/grids.yaml` config file of the source code itself - (e.g. `PyAutoGalaxy/autogalaxy/config/grids.yaml`). - """ - ) - - with np.errstate(all="ignore"): # Division by zero fixed via isnan - grid_radii = obj.radial_grid_from(grid=grid) - - grid_radial_scale = np.where( - grid_radii < grid_radial_minimum, grid_radial_minimum / grid_radii, 1.0 - ) - moved_grid = np.multiply(grid, grid_radial_scale[:, None]) - - if hasattr(grid, "with_new_array"): - moved_grid = grid.with_new_array(moved_grid) - - moved_grid[np.isnan(np.array(moved_grid))] = grid_radial_minimum - - return func(obj, moved_grid, *args, **kwargs) - - return wrapper diff --git a/autoarray/structures/mock/mock_decorators.py b/autoarray/structures/mock/mock_decorators.py index 013b0a62a..c02ebc0b8 100644 --- a/autoarray/structures/mock/mock_decorators.py +++ b/autoarray/structures/mock/mock_decorators.py @@ -165,7 +165,3 @@ def __init__(self): def radial_grid_from(self, grid): return np.sqrt(np.add(np.square(grid[:, 0]), np.square(grid[:, 1]))) - - @decorators.relocate_to_radial_minimum - def deflections_yx_2d_from(self, grid): - return grid diff --git a/test_autoarray/config/grids.yaml b/test_autoarray/config/grids.yaml deleted file mode 100644 index 61c268a27..000000000 --- a/test_autoarray/config/grids.yaml +++ /dev/null @@ -1,12 +0,0 @@ -# Certain light and mass profile calculations become ill defined at (0.0, 0.0) or close to this value. This can lead -# to numerical issues in the calculation of the profile, for example a np.nan may arise, crashing the code. - -# To avoid this, we set a minimum value for the radial coordinate of the profile. If the radial coordinate is below -# this value, it is rounded up to this value. This ensures that the profile cannot receive a radial coordinate of 0.0. - -# For example, if an input grid coordinate has a radial coordinate of 1e-12, for most profiles this will be rounded up -# to radial_minimum=1e-08. This is a small enough value that it should not impact the results of the profile calculation. - -radial_minimum: - radial_minimum: - MockGridRadialMinimum: 2.5 \ No newline at end of file From 6f027158f823c43e3ab665285497b35ec424c5cf Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Thu, 3 Apr 2025 17:26:25 +0100 Subject: [PATCH 04/25] more removal of numpy wrapper nps --- autoarray/mask/derive/indexes_2d.py | 3 ++- autoarray/mask/mask_2d_util.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/autoarray/mask/derive/indexes_2d.py b/autoarray/mask/derive/indexes_2d.py index 062c8e664..0d0a36b26 100644 --- a/autoarray/mask/derive/indexes_2d.py +++ b/autoarray/mask/derive/indexes_2d.py @@ -1,7 +1,8 @@ from __future__ import annotations import logging +import numpy as np -from autoarray.numpy_wrapper import np, register_pytree_node_class +from autoarray.numpy_wrapper import register_pytree_node_class from typing import TYPE_CHECKING if TYPE_CHECKING: diff --git a/autoarray/mask/mask_2d_util.py b/autoarray/mask/mask_2d_util.py index 462448073..10a40b473 100644 --- a/autoarray/mask/mask_2d_util.py +++ b/autoarray/mask/mask_2d_util.py @@ -1,10 +1,10 @@ import numpy as np +import jax.numpy as jnp from scipy.ndimage import convolve from typing import Tuple import warnings from autoarray import exc -from autoarray.numpy_wrapper import np as jnp def native_index_for_slim_index_2d_from( From aa4c9e6e5b8868d35fcfb6d9dce206b63f38e956 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Thu, 3 Apr 2025 17:28:07 +0100 Subject: [PATCH 05/25] remove all numpy wrappers --- autoarray/operators/contour.py | 15 ++++++--------- autoarray/structures/decorators/to_vector_yx.py | 3 +-- autoarray/structures/decorators/transform.py | 2 +- 3 files changed, 8 insertions(+), 12 deletions(-) diff --git a/autoarray/operators/contour.py b/autoarray/operators/contour.py index c7da5c7f1..c352ea19f 100644 --- a/autoarray/operators/contour.py +++ b/autoarray/operators/contour.py @@ -1,6 +1,6 @@ from __future__ import annotations -from autoarray.numpy_wrapper import np, use_jax import numpy +import jax.numpy as jnp from skimage import measure from scipy.spatial import ConvexHull from scipy.spatial import QhullError @@ -42,16 +42,13 @@ def contour_array(self): return self._contour_array pixel_centres = geometry_util.grid_pixel_centres_2d_slim_from( - grid_scaled_2d_slim=np.array(self.grid), + grid_scaled_2d_slim=jnp.array(self.grid), shape_native=self.shape_native, pixel_scales=self.pixel_scales, ).astype("int") - arr = np.zeros(self.shape_native) - if use_jax: - arr = arr.at[tuple(np.array(pixel_centres).T)].set(1) - else: - arr[tuple(np.array(pixel_centres).T)] = 1 + arr = jnp.zeros(self.shape_native) + arr = arr.at[tuple(jnp.array(pixel_centres).T)].set(1) return arr @@ -74,7 +71,7 @@ def contour_list(self): pixel_scales=self.pixel_scales, ) - factor = 0.5 * np.array(self.pixel_scales) * np.array([-1.0, 1.0]) + factor = 0.5 * jnp.array(self.pixel_scales) * jnp.array([-1.0, 1.0]) grid_scaled_1d += factor contour_list.append(Grid2DIrregular(values=grid_scaled_1d)) @@ -104,7 +101,7 @@ def hull( hull_x = grid_convex[hull_vertices, 0] hull_y = grid_convex[hull_vertices, 1] - grid_hull = np.zeros((len(hull_vertices), 2)) + grid_hull = jnp.zeros((len(hull_vertices), 2)) grid_hull[:, 1] = hull_x grid_hull[:, 0] = hull_y diff --git a/autoarray/structures/decorators/to_vector_yx.py b/autoarray/structures/decorators/to_vector_yx.py index 1cf23346d..90aea99ea 100644 --- a/autoarray/structures/decorators/to_vector_yx.py +++ b/autoarray/structures/decorators/to_vector_yx.py @@ -1,6 +1,5 @@ -from autoarray.numpy_wrapper import np from functools import wraps - +import numpy as np from typing import List, Union from autoarray.structures.decorators.abstract import AbstractMaker diff --git a/autoarray/structures/decorators/transform.py b/autoarray/structures/decorators/transform.py index bd837a399..eca0d883b 100644 --- a/autoarray/structures/decorators/transform.py +++ b/autoarray/structures/decorators/transform.py @@ -1,5 +1,5 @@ -from autoarray.numpy_wrapper import np from functools import wraps +import numpy as np from typing import Union from autoarray.structures.grids.uniform_1d import Grid1D From 3b6ab48b21dff83bc144ec19d48d408fc604793b Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Thu, 3 Apr 2025 17:36:33 +0100 Subject: [PATCH 06/25] remove warning for now --- autoarray/numba_util.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/autoarray/numba_util.py b/autoarray/numba_util.py index db34f3e1a..9e0298b73 100644 --- a/autoarray/numba_util.py +++ b/autoarray/numba_util.py @@ -33,7 +33,7 @@ try: if os.environ.get("USE_JAX") == "1": - logger.warning("JAX and numba do not work together, so JAX is being used.") + 1 else: import numba From 44a2808e1761798ae0ab077a07a94408c43e843d Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Thu, 3 Apr 2025 17:48:11 +0100 Subject: [PATCH 07/25] fix structure plotters --- autoarray/operators/contour.py | 14 +++++++------- autoarray/plot/wrap/two_d/array_overlay.py | 4 +++- .../structures/plot/test_structure_plotters.py | 1 + 3 files changed, 11 insertions(+), 8 deletions(-) diff --git a/autoarray/operators/contour.py b/autoarray/operators/contour.py index c352ea19f..2de247d3c 100644 --- a/autoarray/operators/contour.py +++ b/autoarray/operators/contour.py @@ -1,5 +1,5 @@ from __future__ import annotations -import numpy +import numpy as np import jax.numpy as jnp from skimage import measure from scipy.spatial import ConvexHull @@ -42,7 +42,7 @@ def contour_array(self): return self._contour_array pixel_centres = geometry_util.grid_pixel_centres_2d_slim_from( - grid_scaled_2d_slim=jnp.array(self.grid), + grid_scaled_2d_slim=np.array(self.grid), shape_native=self.shape_native, pixel_scales=self.pixel_scales, ).astype("int") @@ -56,7 +56,7 @@ def contour_array(self): def contour_list(self): # make sure to use base numpy to convert JAX array back to a normal array contour_indices_list = measure.find_contours( - numpy.array(self.contour_array.array), 0 + np.array(self.contour_array), 0 ) if len(contour_indices_list) == 0: @@ -71,7 +71,7 @@ def contour_list(self): pixel_scales=self.pixel_scales, ) - factor = 0.5 * jnp.array(self.pixel_scales) * jnp.array([-1.0, 1.0]) + factor = 0.5 * np.array(self.pixel_scales) * np.array([-1.0, 1.0]) grid_scaled_1d += factor contour_list.append(Grid2DIrregular(values=grid_scaled_1d)) @@ -86,10 +86,10 @@ def hull( return None # cast JAX arrays to base numpy arrays - grid_convex = numpy.zeros((len(self.grid), 2)) + grid_convex = np.zeros((len(self.grid), 2)) - grid_convex[:, 0] = numpy.array(self.grid[:, 1]) - grid_convex[:, 1] = numpy.array(self.grid[:, 0]) + grid_convex[:, 0] = np.array(self.grid[:, 1]) + grid_convex[:, 1] = np.array(self.grid[:, 0]) try: hull = ConvexHull(grid_convex) diff --git a/autoarray/plot/wrap/two_d/array_overlay.py b/autoarray/plot/wrap/two_d/array_overlay.py index 57652e8df..5de20b879 100644 --- a/autoarray/plot/wrap/two_d/array_overlay.py +++ b/autoarray/plot/wrap/two_d/array_overlay.py @@ -19,4 +19,6 @@ def overlay_array(self, array, figure): aspect = figure.aspect_from(shape_native=array.shape_native) extent = array.extent_of_zoomed_array(buffer=0) - plt.imshow(X=array.native, aspect=aspect, extent=extent, **self.config_dict) + print(type(array)) + + plt.imshow(X=array.native._array, aspect=aspect, extent=extent, **self.config_dict) diff --git a/test_autoarray/structures/plot/test_structure_plotters.py b/test_autoarray/structures/plot/test_structure_plotters.py index ad1ca0251..d455c86f4 100644 --- a/test_autoarray/structures/plot/test_structure_plotters.py +++ b/test_autoarray/structures/plot/test_structure_plotters.py @@ -3,6 +3,7 @@ from os import path import pytest import numpy as np +import jax.numpy as jnp import shutil directory = path.dirname(path.realpath(__file__)) From ea139fcd95bd3400b620df39041e40d0470beba0 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Thu, 3 Apr 2025 17:56:03 +0100 Subject: [PATCH 08/25] clean up vectors_yx --- autoarray/structures/vectors/uniform.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/autoarray/structures/vectors/uniform.py b/autoarray/structures/vectors/uniform.py index 89d589139..fc66b0e8d 100644 --- a/autoarray/structures/vectors/uniform.py +++ b/autoarray/structures/vectors/uniform.py @@ -1,7 +1,8 @@ import logging -# import numpy as np -from autofit.jax_wrapper import numpy as np, use_jax +import numpy as np +import jax.numpy as jnp +# from autofit.jax_wrapper import numpy as np, use_jax from typing import List, Optional, Tuple, Union from autoarray.structures.arrays.uniform_2d import Array2D @@ -396,11 +397,7 @@ def magnitudes(self) -> Array2D: """ Returns the magnitude of every vector which are computed as sqrt(y**2 + x**2). """ - if use_jax: - s = self.array - else: - s = self - return Array2D(values=np.sqrt(s[:, 0] ** 2.0 + s[:, 1] ** 2.0), mask=self.mask) + return Array2D(values=jnp.sqrt(self.array[:, 0] ** 2.0 + self.array[:, 1] ** 2.0), mask=self.mask) @property def y(self) -> Array2D: From ec1e81eb0066bd638c0d48dda33d95d14169e27d Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Thu, 3 Apr 2025 17:57:41 +0100 Subject: [PATCH 09/25] remove autofit imports --- autoarray/geometry/geometry_2d.py | 2 -- autoarray/operators/over_sampling/over_sampler.py | 2 +- autoarray/structures/vectors/uniform.py | 1 - test_autoarray/test_jax_changes.py | 8 +++++--- 4 files changed, 6 insertions(+), 7 deletions(-) diff --git a/autoarray/geometry/geometry_2d.py b/autoarray/geometry/geometry_2d.py index e78f0f75a..9eea7e9f2 100644 --- a/autoarray/geometry/geometry_2d.py +++ b/autoarray/geometry/geometry_2d.py @@ -13,8 +13,6 @@ from autoarray import type as ty from autoarray.geometry import geometry_util -from autofit.jax_wrapper import use_jax - logging.basicConfig() logger = logging.getLogger(__name__) diff --git a/autoarray/operators/over_sampling/over_sampler.py b/autoarray/operators/over_sampling/over_sampler.py index 65393709c..ae458e41b 100644 --- a/autoarray/operators/over_sampling/over_sampler.py +++ b/autoarray/operators/over_sampling/over_sampler.py @@ -10,7 +10,7 @@ from autoarray.operators.over_sampling import over_sample_util -from autofit.jax_wrapper import register_pytree_node_class +from autoarray.numpy_wrapper import register_pytree_node_class @register_pytree_node_class diff --git a/autoarray/structures/vectors/uniform.py b/autoarray/structures/vectors/uniform.py index fc66b0e8d..6213ecfbf 100644 --- a/autoarray/structures/vectors/uniform.py +++ b/autoarray/structures/vectors/uniform.py @@ -2,7 +2,6 @@ import numpy as np import jax.numpy as jnp -# from autofit.jax_wrapper import numpy as np, use_jax from typing import List, Optional, Tuple, Union from autoarray.structures.arrays.uniform_2d import Array2D diff --git a/test_autoarray/test_jax_changes.py b/test_autoarray/test_jax_changes.py index f5104a942..2b6289317 100644 --- a/test_autoarray/test_jax_changes.py +++ b/test_autoarray/test_jax_changes.py @@ -1,8 +1,10 @@ -import autoarray as aa +import jax.numpy as jnp import pytest + +import autoarray as aa + from autoarray import Grid2D, Mask2D -from autofit.jax_wrapper import numpy as np @pytest.fixture(name="array") @@ -33,4 +35,4 @@ def test_boolean_issue(): mask=Mask2D.all_false((10, 10), pixel_scales=1.0), ) values, keys = Grid2D.instance_flatten(grid) - np.array(Grid2D.instance_unflatten(keys, values)) + jnp.array(Grid2D.instance_unflatten(keys, values)) From 37e81f157a2acd3dff2f55ab546a834ece26ddd9 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Thu, 3 Apr 2025 18:04:21 +0100 Subject: [PATCH 10/25] fix voronoi unit test in structures --- autoarray/inversion/pixelization/image_mesh/overlay.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/autoarray/inversion/pixelization/image_mesh/overlay.py b/autoarray/inversion/pixelization/image_mesh/overlay.py index c5bc7eaef..de130ee6e 100644 --- a/autoarray/inversion/pixelization/image_mesh/overlay.py +++ b/autoarray/inversion/pixelization/image_mesh/overlay.py @@ -220,11 +220,11 @@ def image_plane_mesh_grid_from( origin=origin, ) - overlaid_centres = geometry_util.grid_pixel_centres_2d_slim_from( + overlaid_centres = np.array(geometry_util.grid_pixel_centres_2d_slim_from( grid_scaled_2d_slim=unmasked_overlay_grid, shape_native=mask.shape_native, pixel_scales=mask.pixel_scales, - ).astype("int") + )).astype("int") total_pixels = total_pixels_2d_from( mask_2d=mask.array, From b31c0fc24e248e2a82dbf40d7f5fdfed9d4b0a9e Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Fri, 4 Apr 2025 14:24:39 +0100 Subject: [PATCH 11/25] fix test_preprocess --- autoarray/dataset/preprocess.py | 12 ++++++------ test_autoarray/dataset/test_preprocess.py | 4 ++-- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/autoarray/dataset/preprocess.py b/autoarray/dataset/preprocess.py index 5c7338204..f13af3184 100644 --- a/autoarray/dataset/preprocess.py +++ b/autoarray/dataset/preprocess.py @@ -263,15 +263,15 @@ def edges_from(image, no_edges): edges = [] for edge_no in range(no_edges): - top_edge = image.native[edge_no, edge_no : image.shape_native[1] - edge_no] - bottom_edge = image.native[ + top_edge = image.native.array[edge_no, edge_no : image.shape_native[1] - edge_no] + bottom_edge = image.native.array[ image.shape_native[0] - 1 - edge_no, edge_no : image.shape_native[1] - edge_no, ] - left_edge = image.native[ + left_edge = image.native.array[ edge_no + 1 : image.shape_native[0] - 1 - edge_no, edge_no ] - right_edge = image.native[ + right_edge = image.native.array[ edge_no + 1 : image.shape_native[0] - 1 - edge_no, image.shape_native[1] - 1 - edge_no, ] @@ -517,8 +517,8 @@ def noise_map_with_signal_to_noise_limit_from( noise_map_limit = np.where( (signal_to_noise_map.native > signal_to_noise_limit) & (noise_limit_mask == False), - np.abs(data.native) / signal_to_noise_limit, - noise_map.native, + np.abs(data.native.array) / signal_to_noise_limit, + noise_map.native.array, ) mask = Mask2D.all_false( diff --git a/test_autoarray/dataset/test_preprocess.py b/test_autoarray/dataset/test_preprocess.py index 74e8ef774..f484fd648 100644 --- a/test_autoarray/dataset/test_preprocess.py +++ b/test_autoarray/dataset/test_preprocess.py @@ -462,7 +462,7 @@ def test__background_noise_map_via_edges_of_image_from_4(): ) assert np.allclose( - background_noise_map.native, + background_noise_map.native.array, np.full(fill_value=np.std(np.arange(28)), shape=image.shape_native), ) @@ -486,7 +486,7 @@ def test__background_noise_map_via_edges_of_image_from_5(): ) assert np.allclose( - background_noise_map.native, + background_noise_map.native.array, np.full(fill_value=np.std(np.arange(48)), shape=image.shape_native), ) From 80fc8e8781d3d7b8bcb50ca2698bd5cbda54b8ae Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Fri, 4 Apr 2025 14:25:38 +0100 Subject: [PATCH 12/25] fix test dataset abstract --- autoarray/dataset/imaging/dataset.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/autoarray/dataset/imaging/dataset.py b/autoarray/dataset/imaging/dataset.py index e3ec74d3b..b5b5b73d7 100644 --- a/autoarray/dataset/imaging/dataset.py +++ b/autoarray/dataset/imaging/dataset.py @@ -166,8 +166,9 @@ def __init__( self.psf = psf - if psf.mask.shape[0] % 2 == 0 or psf.mask.shape[1] % 2 == 0: - raise exc.KernelException("Kernel2D Kernel2D must be odd") + if psf is not None: + if psf.mask.shape[0] % 2 == 0 or psf.mask.shape[1] % 2 == 0: + raise exc.KernelException("Kernel2D Kernel2D must be odd") @cached_property def grids(self): From cd276cd3146e4d03beede6a9ffe366a1ab0f93ba Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Fri, 4 Apr 2025 14:38:09 +0100 Subject: [PATCH 13/25] fix test imaging --- autoarray/dataset/imaging/dataset.py | 18 +++++++++--------- test_autoarray/dataset/imaging/test_dataset.py | 9 +++++++-- .../dataset/imaging/test_simulator.py | 2 +- 3 files changed, 17 insertions(+), 12 deletions(-) diff --git a/autoarray/dataset/imaging/dataset.py b/autoarray/dataset/imaging/dataset.py index b5b5b73d7..8d84ee1b8 100644 --- a/autoarray/dataset/imaging/dataset.py +++ b/autoarray/dataset/imaging/dataset.py @@ -204,9 +204,9 @@ def w_tilde(self): indexes, lengths, ) = inversion_imaging_util.w_tilde_curvature_preload_imaging_from( - noise_map_native=np.array(self.noise_map.native), - kernel_native=np.array(self.psf.native), - native_index_for_slim_index=self.mask.derive_indexes.native_for_slim, + noise_map_native=np.array(self.noise_map.native.array).astype("float64"), + kernel_native=np.array(self.psf.native.array).astype("float64"), + native_index_for_slim_index=np.array(self.mask.derive_indexes.native_for_slim).astype("int"), ) return WTildeImaging( @@ -409,20 +409,20 @@ def apply_noise_scaling( """ if signal_to_noise_value is None: - noise_map = self.noise_map.native - noise_map[mask == False] = noise_value + noise_map = np.array(self.noise_map.native.array) + noise_map[mask.array == False] = noise_value else: noise_map = np.where( mask == False, - np.median(self.data.native[mask.derive_mask.edge == False]) + np.median(self.data.native.array[mask.derive_mask.edge == False]) / signal_to_noise_value, - self.noise_map.native, + self.noise_map.native.array, ) if should_zero_data: - data = np.where(np.invert(mask), 0.0, self.data.native) + data = np.where(np.invert(mask.array), 0.0, self.data.native.array) else: - data = self.data.native + data = self.data.native.array data_unmasked = Array2D.no_mask( values=data, diff --git a/test_autoarray/dataset/imaging/test_dataset.py b/test_autoarray/dataset/imaging/test_dataset.py index 45ef0a80f..ca33f1b40 100644 --- a/test_autoarray/dataset/imaging/test_dataset.py +++ b/test_autoarray/dataset/imaging/test_dataset.py @@ -139,7 +139,7 @@ def test__apply_mask(imaging_7x7, mask_2d_7x7, psf_3x3): == 2.0 * np.ones((7, 7)) * np.invert(mask_2d_7x7) ).all() - assert (masked_imaging_7x7.psf.slim == (1.0 / 3.0) * psf_3x3.slim).all() + assert masked_imaging_7x7.psf.slim == pytest.approx((1.0 / 3.0) * psf_3x3.slim, 1.0e-4) assert type(masked_imaging_7x7.psf) == aa.Kernel2D assert masked_imaging_7x7.w_tilde.curvature_preload.shape == (35,) @@ -244,4 +244,9 @@ def test__noise_map_unmasked_has_zeros_or_negative__raises_exception(): def test__psf_not_odd_x_odd_kernel__raises_error(): with pytest.raises(exc.KernelException): - aa.Kernel2D.no_mask(values=[[0.0, 1.0], [1.0, 2.0]], pixel_scales=1.0) + image = aa.Array2D.ones(shape_native=(3, 3), pixel_scales=1.0) + noise_map = aa.Array2D.ones(shape_native=(3, 3), pixel_scales=1.0) + psf = aa.Kernel2D.no_mask(values=[[0.0, 1.0], [1.0, 2.0]], pixel_scales=1.0) + + dataset = aa.Imaging(data=image, noise_map=noise_map, psf=psf, pad_for_psf=False) + diff --git a/test_autoarray/dataset/imaging/test_simulator.py b/test_autoarray/dataset/imaging/test_simulator.py index 3cc4182d3..54dc1f6ed 100644 --- a/test_autoarray/dataset/imaging/test_simulator.py +++ b/test_autoarray/dataset/imaging/test_simulator.py @@ -70,7 +70,7 @@ def test__via_image_from__psf_off__noise_off_value_is_noise_value( == np.array([[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]]) ).all() - assert np.allclose(dataset.noise_map.native, 0.2 * np.ones((3, 3))) + assert np.allclose(dataset.noise_map.native.array, 0.2 * np.ones((3, 3))) def test__via_image_from__psf_off__background_sky_on(image_central_delta_3x3): From f6dfda50b5c2a03db875d0beda7bee7a39ebdbb5 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Fri, 4 Apr 2025 14:42:00 +0100 Subject: [PATCH 14/25] fix layout --- test_autoarray/layout/test_region.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/test_autoarray/layout/test_region.py b/test_autoarray/layout/test_region.py index 690cfa9bb..643b8532a 100644 --- a/test_autoarray/layout/test_region.py +++ b/test_autoarray/layout/test_region.py @@ -152,7 +152,7 @@ def test__slice_2d__addition(): image = np.ones((2, 2)) region = aa.Region2D(region=(0, 1, 0, 1)) - array[region.slice] += image[region.slice] + array = array.at[region.slice].add(image[region.slice]) assert (array == np.array([[1.0, 0.0], [0.0, 0.0]])).all() @@ -161,7 +161,7 @@ def test__slice_2d__addition(): image = np.ones((2, 2)) region = aa.Region2D(region=(0, 1, 0, 1)) - array[region.slice] += image[region.slice] + array = array.at[region.slice].add(image[region.slice]) assert (array == np.array([[2.0, 1.0], [1.0, 1.0]])).all() @@ -170,7 +170,7 @@ def test__slice_2d__addition(): image = np.ones((3, 3)) region = aa.Region2D(region=(1, 3, 2, 3)) - array[region.slice] += image[region.slice] + array = array.at[region.slice].add(image[region.slice]) assert ( array == np.array([[1.0, 1.0, 1.0], [1.0, 1.0, 2.0], [1.0, 1.0, 2.0]]) @@ -183,7 +183,7 @@ def test__slice_2d__set_to_zerose(): region = aa.Region2D(region=(0, 1, 0, 1)) - array[region.slice] = 0 + array = array.at[region.slice].set(0) assert (array == np.array([[0.0, 1.0], [1.0, 1.0]])).all() @@ -192,7 +192,7 @@ def test__slice_2d__set_to_zerose(): region = aa.Region2D(region=(1, 3, 2, 3)) - array[region.slice] = 0 + array = array.at[region.slice].set(0) assert ( array == np.array([[1.0, 1.0, 1.0], [1.0, 1.0, 0.0], [1.0, 1.0, 0.0]]) From 5430647dbf983d9bac7e7de0107b855162f0d117 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Fri, 4 Apr 2025 14:46:46 +0100 Subject: [PATCH 15/25] fix plot unit tests --- autoarray/plot/wrap/two_d/contour.py | 4 ++-- test_autoarray/plot/include/test_include.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/autoarray/plot/wrap/two_d/contour.py b/autoarray/plot/wrap/two_d/contour.py index c80159813..164fe86be 100644 --- a/autoarray/plot/wrap/two_d/contour.py +++ b/autoarray/plot/wrap/two_d/contour.py @@ -93,10 +93,10 @@ def set( config_dict.pop("use_log10") config_dict.pop("include_values") - levels = self.levels_from(array) + levels = self.levels_from(array.array) ax = plt.contour( - array.native[::-1], levels=levels, extent=extent, **config_dict + array.native.array[::-1], levels=levels, extent=extent, **config_dict ) if self.include_values: try: diff --git a/test_autoarray/plot/include/test_include.py b/test_autoarray/plot/include/test_include.py index 6f4d29c77..b32616e9d 100644 --- a/test_autoarray/plot/include/test_include.py +++ b/test_autoarray/plot/include/test_include.py @@ -6,7 +6,7 @@ def test__loads_default_values_from_config_if_not_input(): assert include.origin is True assert include.mask == True - assert include.border is True + assert include.border is False assert include.parallel_overscan is True assert include.serial_prescan is True assert include.serial_overscan is False From 083ed0bb2f08eed82eb1520b929347036075690c Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Fri, 4 Apr 2025 15:50:11 +0100 Subject: [PATCH 16/25] over sampling unit tests --- .../operators/over_sampling/over_sampler.py | 37 ++++++++++++++----- .../over_sample/test_over_sampler.py | 10 +++++ 2 files changed, 38 insertions(+), 9 deletions(-) diff --git a/autoarray/operators/over_sampling/over_sampler.py b/autoarray/operators/over_sampling/over_sampler.py index ae458e41b..9fda67bb7 100644 --- a/autoarray/operators/over_sampling/over_sampler.py +++ b/autoarray/operators/over_sampling/over_sampler.py @@ -147,6 +147,16 @@ def __init__(self, mask: Mask2D, sub_size: Union[int, Array2D]): over_sample_size=sub_size, mask=mask ) + + @property + def sub_is_uniform(self) -> bool: + """ + Returns True if the sub_size is uniform across all pixels in the mask. + """ + return np.all( + np.isclose(self.sub_size.array, self.sub_size.array[0]) + ) + def tree_flatten(self): return (self.mask, self.sub_size), () @@ -185,7 +195,7 @@ def sub_pixel_areas(self) -> np.ndarray: """ The area of every sub-pixel in the mask. """ - sub_pixel_areas = jnp.zeros(self.sub_total) + sub_pixel_areas = np.zeros(self.sub_total) k = 0 @@ -221,15 +231,24 @@ def binned_array_2d_from(self, array: Array2D) -> "Array2D": except AttributeError: pass - # binned_array_2d = over_sample_util.binned_array_2d_from( - # array_2d=jnp.array(array), - # mask_2d=jnp.array(self.mask), - # sub_size=jnp.array(self.sub_size).astype("int"), - # ) + if self.sub_is_uniform: + binned_array_2d = array.reshape( + self.mask.shape_slim, self.sub_size[0] ** 2 + ).mean(axis=1) + else: + + # Define group sizes + group_sizes = jnp.array(self.sub_size.array.astype("int") ** 2) + + # Compute the cumulative sum of group sizes to get split points + split_indices = jnp.cumsum(group_sizes) + + # Ensure correct concatenation by making 0 a JAX array + start_indices = jnp.concatenate((jnp.array([0]), split_indices[:-1])) - binned_array_2d = array.reshape( - self.mask.shape_slim, self.sub_size[0] ** 2 - ).mean(axis=1) + # Compute the group means + binned_array_2d = jnp.array( + [array[start:end].mean() for start, end in zip(start_indices, split_indices)]) return Array2D( values=binned_array_2d, diff --git a/test_autoarray/operators/over_sample/test_over_sampler.py b/test_autoarray/operators/over_sample/test_over_sampler.py index 7da24d79a..a32b11e8f 100644 --- a/test_autoarray/operators/over_sample/test_over_sampler.py +++ b/test_autoarray/operators/over_sample/test_over_sampler.py @@ -70,6 +70,16 @@ def test__binned_array_2d_from(): pixel_scales=1.0, ) + over_sampling = aa.OverSampler( + mask=mask, sub_size=aa.Array2D(values=[2, 2], mask=mask) + ) + + arr = np.array([1.0, 5.0, 7.0, 10.0, 10.0, 10.0, 10.0, 10.0]) + + binned_array_2d = over_sampling.binned_array_2d_from(array=arr) + + assert binned_array_2d.slim == pytest.approx(np.array([5.75, 10.0]), 1.0e-4) + over_sampling = aa.OverSampler( mask=mask, sub_size=aa.Array2D(values=[1, 2], mask=mask) ) From 72af86b04b70f94ae3514f8c28ff38d85abfe133 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Fri, 4 Apr 2025 16:07:46 +0100 Subject: [PATCH 17/25] fix all fit tests --- autoarray/fit/fit_util.py | 87 +++++------- test_autoarray/fit/test_fit_util.py | 199 ++++++++++++++-------------- 2 files changed, 135 insertions(+), 151 deletions(-) diff --git a/autoarray/fit/fit_util.py b/autoarray/fit/fit_util.py index d40f55d1c..10f24f9a7 100644 --- a/autoarray/fit/fit_util.py +++ b/autoarray/fit/fit_util.py @@ -1,5 +1,6 @@ from functools import wraps -import jax.numpy as np +import jax.numpy as jnp +import numpy as np from autoarray.mask.abstract_mask import Mask @@ -83,7 +84,7 @@ def chi_squared_from(*, chi_squared_map: ty.DataLike) -> float: chi_squared_map The chi-squared-map of values of the model-data fit to the dataset. """ - return np.sum(chi_squared_map._array) + return jnp.sum(np.array(chi_squared_map)) def noise_normalization_from(*, noise_map: ty.DataLike) -> float: @@ -97,12 +98,12 @@ def noise_normalization_from(*, noise_map: ty.DataLike) -> float: noise_map The masked noise-map of the dataset. """ - return np.sum(np.log(2 * np.pi * noise_map._array**2.0)) + return jnp.sum(jnp.log(2 * jnp.pi * np.array(noise_map)**2.0)) def normalized_residual_map_complex_from( - *, residual_map: np.ndarray, noise_map: np.ndarray -) -> np.ndarray: + *, residual_map: jnp.ndarray, noise_map: jnp.ndarray +) -> jnp.ndarray: """ Returns the normalized residual-map of the fit of complex model-data to a dataset, where: @@ -126,8 +127,8 @@ def normalized_residual_map_complex_from( def chi_squared_map_complex_from( - *, residual_map: np.ndarray, noise_map: np.ndarray -) -> np.ndarray: + *, residual_map: jnp.ndarray, noise_map: jnp.ndarray +) -> jnp.ndarray: """ Returnss the chi-squared-map of the fit of complex model-data to a dataset, where: @@ -145,7 +146,7 @@ def chi_squared_map_complex_from( return chi_squared_map_real + 1j * chi_squared_map_imag -def chi_squared_complex_from(*, chi_squared_map: np.ndarray) -> float: +def chi_squared_complex_from(*, chi_squared_map: jnp.ndarray) -> float: """ Returns the chi-squared terms of each complex model data's fit to a masked dataset, by summing the masked chi-squared-map of the fit. @@ -157,12 +158,12 @@ def chi_squared_complex_from(*, chi_squared_map: np.ndarray) -> float: chi_squared_map The chi-squared-map of values of the model-data fit to the dataset. """ - chi_squared_real = np.sum(chi_squared_map.real) - chi_squared_imag = np.sum(chi_squared_map.imag) + chi_squared_real = jnp.sum(chi_squared_map.real) + chi_squared_imag = jnp.sum(chi_squared_map.imag) return chi_squared_real + chi_squared_imag -def noise_normalization_complex_from(*, noise_map: np.ndarray) -> float: +def noise_normalization_complex_from(*, noise_map: jnp.ndarray) -> float: """ Returns the noise-map normalization terms of a complex noise-map, summing the noise_map value in every pixel as: @@ -173,8 +174,8 @@ def noise_normalization_complex_from(*, noise_map: np.ndarray) -> float: noise_map The masked noise-map of the dataset. """ - noise_normalization_real = np.sum(np.log(2 * np.pi * noise_map.real**2.0)) - noise_normalization_imag = np.sum(np.log(2 * np.pi * noise_map.imag**2.0)) + noise_normalization_real = jnp.sum(jnp.log(2 * jnp.pi * noise_map.real**2.0)) + noise_normalization_imag = jnp.sum(jnp.log(2 * jnp.pi * noise_map.imag**2.0)) return noise_normalization_real + noise_normalization_imag @@ -198,9 +199,7 @@ def residual_map_with_mask_from( model_data The model data used to fit the data. """ - return np.subtract( - data, model_data, out=np.zeros_like(data), where=np.asarray(mask) == 0 - ) + return jnp.where(jnp.asarray(mask) == 0, jnp.subtract(data, model_data), 0) @to_new_array @@ -223,13 +222,7 @@ def normalized_residual_map_with_mask_from( mask The mask applied to the residual-map, where `False` entries are included in the calculation. """ - return np.divide( - residual_map, - noise_map, - out=np.zeros_like(residual_map), - where=np.asarray(mask) == 0, - ) - + return jnp.where(jnp.asarray(mask) == 0, jnp.divide(residual_map, noise_map), 0) @to_new_array def chi_squared_map_with_mask_from( @@ -251,13 +244,10 @@ def chi_squared_map_with_mask_from( mask The mask applied to the residual-map, where `False` entries are included in the calculation. """ - return np.square( - np.divide( - residual_map, - noise_map, - out=np.zeros_like(residual_map), - where=np.asarray(mask) == 0, - ) + return jnp.where( + jnp.asarray(mask) == 0, + jnp.square(residual_map / noise_map), + 0 ) @@ -275,7 +265,7 @@ def chi_squared_with_mask_from(*, chi_squared_map: ty.DataLike, mask: Mask) -> f mask The mask applied to the chi-squared-map, where `False` entries are included in the calculation. """ - return float(np.sum(chi_squared_map[np.asarray(mask) == 0])) + return float(jnp.sum(chi_squared_map[jnp.asarray(mask) == 0])) def chi_squared_with_mask_fast_from( @@ -302,14 +292,14 @@ def chi_squared_with_mask_fast_from( The mask applied to the chi-squared-map, where `False` entries are included in the calculation. """ return float( - np.sum( - np.square( - np.divide( - np.subtract( + jnp.sum( + jnp.square( + jnp.divide( + jnp.subtract( data, model_data, - )[np.asarray(mask) == 0], - noise_map[np.asarray(mask) == 0], + )[jnp.asarray(mask) == 0], + noise_map[jnp.asarray(mask) == 0], ) ) ) @@ -331,11 +321,11 @@ def noise_normalization_with_mask_from(*, noise_map: ty.DataLike, mask: Mask) -> mask The mask applied to the noise-map, where `False` entries are included in the calculation. """ - return float(np.sum(np.log(2 * np.pi * noise_map[np.asarray(mask) == 0] ** 2.0))) + return float(jnp.sum(jnp.log(2 * jnp.pi * noise_map[jnp.asarray(mask) == 0] ** 2.0))) def chi_squared_with_noise_covariance_from( - *, residual_map: ty.DataLike, noise_covariance_matrix_inv: np.ndarray + *, residual_map: ty.DataLike, noise_covariance_matrix_inv: jnp.ndarray ) -> float: """ Returns the chi-squared value of the fit of model-data to a masked dataset, where @@ -351,7 +341,7 @@ def chi_squared_with_noise_covariance_from( The inverse of the noise covariance matrix. """ - return residual_map @ noise_covariance_matrix_inv @ residual_map + return residual_map.array @ noise_covariance_matrix_inv @ residual_map.array def log_likelihood_from(*, chi_squared: float, noise_normalization: float) -> float: @@ -431,8 +421,8 @@ def log_evidence_from( def residual_flux_fraction_map_from( - *, residual_map: np.ndarray, data: np.ndarray -) -> np.ndarray: + *, residual_map: jnp.ndarray, data: jnp.ndarray +) -> jnp.ndarray: """ Returns the residual flux fraction map of the fit of model-data to a masked dataset, where: @@ -445,12 +435,12 @@ def residual_flux_fraction_map_from( data The data of the dataset. """ - return np.divide(residual_map, data, out=np.zeros_like(residual_map)) + return jnp.where(data != 0, residual_map / data, 0) def residual_flux_fraction_map_with_mask_from( - *, residual_map: np.ndarray, data: np.ndarray, mask: Mask -) -> np.ndarray: + *, residual_map: jnp.ndarray, data: jnp.ndarray, mask: Mask +) -> jnp.ndarray: """ Returnss the residual flux fraction map of the fit of model-data to a masked dataset, where: @@ -467,9 +457,4 @@ def residual_flux_fraction_map_with_mask_from( mask The mask applied to the residual-map, where `False` entries are included in the calculation. """ - return np.divide( - residual_map, - data, - out=np.zeros_like(residual_map), - where=np.asarray(mask) == 0, - ) + return jnp.where(mask == 0, residual_map / data, 0) \ No newline at end of file diff --git a/test_autoarray/fit/test_fit_util.py b/test_autoarray/fit/test_fit_util.py index 641dbf52e..6bbb5f871 100644 --- a/test_autoarray/fit/test_fit_util.py +++ b/test_autoarray/fit/test_fit_util.py @@ -1,41 +1,41 @@ import autoarray as aa -import numpy as np +import jax.numpy as jnp import pytest def test__residual_map_from(): - data = np.array([10.0, 10.0, 10.0, 10.0]) - model_data = np.array([10.0, 10.0, 10.0, 10.0]) + data = jnp.array([10.0, 10.0, 10.0, 10.0]) + model_data = jnp.array([10.0, 10.0, 10.0, 10.0]) residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data) - assert (residual_map == np.array([0.0, 0.0, 0.0, 0.0])).all() + assert (residual_map == jnp.array([0.0, 0.0, 0.0, 0.0])).all() - data = np.array([10.0, 10.0, 10.0, 10.0]) - model_data = np.array([11.0, 10.0, 9.0, 8.0]) + data = jnp.array([10.0, 10.0, 10.0, 10.0]) + model_data = jnp.array([11.0, 10.0, 9.0, 8.0]) residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data) - assert (residual_map == np.array([-1.0, 0.0, 1.0, 2.0])).all() + assert (residual_map == jnp.array([-1.0, 0.0, 1.0, 2.0])).all() def test__residual_map_with_mask_from(): - data = np.array([10.0, 10.0, 10.0, 10.0]) - mask = np.array([True, False, False, True]) - model_data = np.array([11.0, 10.0, 9.0, 8.0]) + data = jnp.array([10.0, 10.0, 10.0, 10.0]) + mask = jnp.array([True, False, False, True]) + model_data = jnp.array([11.0, 10.0, 9.0, 8.0]) residual_map = aa.util.fit.residual_map_with_mask_from( data=data, mask=mask, model_data=model_data ) - assert (residual_map == np.array([0.0, 0.0, 1.0, 0.0])).all() + assert (residual_map == jnp.array([0.0, 0.0, 1.0, 0.0])).all() def test__normalized_residual_map_from(): - data = np.array([10.0, 10.0, 10.0, 10.0]) - noise_map = np.array([2.0, 2.0, 2.0, 2.0]) - model_data = np.array([10.0, 10.0, 10.0, 10.0]) + data = jnp.array([10.0, 10.0, 10.0, 10.0]) + noise_map = jnp.array([2.0, 2.0, 2.0, 2.0]) + model_data = jnp.array([10.0, 10.0, 10.0, 10.0]) residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data) @@ -43,9 +43,9 @@ def test__normalized_residual_map_from(): residual_map=residual_map, noise_map=noise_map ) - assert (normalized_residual_map == np.array([0.0, 0.0, 0.0, 0.0])).all() + assert normalized_residual_map == pytest.approx(jnp.array([0.0, 0.0, 0.0, 0.0]), 1.0e-4) - model_data = np.array([11.0, 10.0, 9.0, 8.0]) + model_data = jnp.array([11.0, 10.0, 9.0, 8.0]) residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data) @@ -53,17 +53,14 @@ def test__normalized_residual_map_from(): residual_map=residual_map, noise_map=noise_map ) - assert ( - normalized_residual_map - == np.array([-(1.0 / 2.0), 0.0, (1.0 / 2.0), (2.0 / 2.0)]) - ).all() + assert normalized_residual_map == pytest.approx(jnp.array([-(1.0 / 2.0), 0.0, (1.0 / 2.0), (2.0 / 2.0)]), 1.0e-4) def test__normalized_residual_map_with_mask_from(): - data = np.array([10.0, 10.0, 10.0, 10.0]) - mask = np.array([True, False, False, True]) - noise_map = np.array([2.0, 2.0, 2.0, 2.0]) - model_data = np.array([11.0, 10.0, 9.0, 8.0]) + data = jnp.array([10.0, 10.0, 10.0, 10.0]) + mask = jnp.array([True, False, False, True]) + noise_map = jnp.array([2.0, 2.0, 2.0, 2.0]) + model_data = jnp.array([11.0, 10.0, 9.0, 8.0]) residual_map = aa.util.fit.residual_map_with_mask_from( data=data, mask=mask, model_data=model_data @@ -73,13 +70,15 @@ def test__normalized_residual_map_with_mask_from(): residual_map=residual_map, mask=mask, noise_map=noise_map ) - assert (normalized_residual_map == np.array([0.0, 0.0, (1.0 / 2.0), 0.0])).all() + print(normalized_residual_map) + + assert normalized_residual_map == pytest.approx(jnp.array([0.0, 0.0, (1.0 / 2.0), 0.0]), abs=1.0e-4) def test__normalized_residual_map_complex_from(): - data = np.array([10.0 + 10.0j, 10.0 + 10.0j]) - noise_map = np.array([2.0 + 2.0j, 2.0 + 2.0j]) - model_data = np.array([9.0 + 12.0j, 9.0 + 12.0j]) + data = jnp.array([10.0 + 10.0j, 10.0 + 10.0j]) + noise_map = jnp.array([2.0 + 2.0j, 2.0 + 2.0j]) + model_data = jnp.array([9.0 + 12.0j, 9.0 + 12.0j]) residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data) @@ -87,13 +86,13 @@ def test__normalized_residual_map_complex_from(): residual_map=residual_map, noise_map=noise_map ) - assert (normalized_residual_map == np.array([0.5 - 1.0j, 0.5 - 1.0j])).all() + assert (normalized_residual_map == jnp.array([0.5 - 1.0j, 0.5 - 1.0j])).all() def test__chi_squared_map_from(): - data = np.array([10.0, 10.0, 10.0, 10.0]) - noise_map = np.array([2.0, 2.0, 2.0, 2.0]) - model_data = np.array([10.0, 10.0, 10.0, 10.0]) + data = jnp.array([10.0, 10.0, 10.0, 10.0]) + noise_map = jnp.array([2.0, 2.0, 2.0, 2.0]) + model_data = jnp.array([10.0, 10.0, 10.0, 10.0]) residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data) @@ -101,9 +100,9 @@ def test__chi_squared_map_from(): residual_map=residual_map, noise_map=noise_map ) - assert (chi_squared_map == np.array([0.0, 0.0, 0.0, 0.0])).all() + assert (chi_squared_map == jnp.array([0.0, 0.0, 0.0, 0.0])).all() - model_data = np.array([11.0, 10.0, 9.0, 8.0]) + model_data = jnp.array([11.0, 10.0, 9.0, 8.0]) residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data) @@ -113,15 +112,15 @@ def test__chi_squared_map_from(): assert ( chi_squared_map - == np.array([(1.0 / 2.0) ** 2.0, 0.0, (1.0 / 2.0) ** 2.0, (2.0 / 2.0) ** 2.0]) + == jnp.array([(1.0 / 2.0) ** 2.0, 0.0, (1.0 / 2.0) ** 2.0, (2.0 / 2.0) ** 2.0]) ).all() def test__chi_squared_map_with_mask_from(): - data = np.array([10.0, 10.0, 10.0, 10.0]) - mask = np.array([True, False, False, True]) - noise_map = np.array([2.0, 2.0, 2.0, 2.0]) - model_data = np.array([11.0, 10.0, 9.0, 8.0]) + data = jnp.array([10.0, 10.0, 10.0, 10.0]) + mask = jnp.array([True, False, False, True]) + noise_map = jnp.array([2.0, 2.0, 2.0, 2.0]) + model_data = jnp.array([11.0, 10.0, 9.0, 8.0]) residual_map = aa.util.fit.residual_map_with_mask_from( data=data, mask=mask, model_data=model_data @@ -131,9 +130,9 @@ def test__chi_squared_map_with_mask_from(): residual_map=residual_map, mask=mask, noise_map=noise_map ) - assert (chi_squared_map == np.array([0.0, 0.0, (1.0 / 2.0) ** 2.0, 0.0])).all() + assert (chi_squared_map == jnp.array([0.0, 0.0, (1.0 / 2.0) ** 2.0, 0.0])).all() - model_data = np.array([11.0, 10.0, 9.0, 8.0]) + model_data = jnp.array([11.0, 10.0, 9.0, 8.0]) residual_map = aa.util.fit.residual_map_with_mask_from( data=data, mask=mask, model_data=model_data @@ -143,13 +142,13 @@ def test__chi_squared_map_with_mask_from(): residual_map=residual_map, mask=mask, noise_map=noise_map ) - assert (chi_squared_map == np.array([0.0, 0.0, (1.0 / 2.0) ** 2.0, 0.0])).all() + assert (chi_squared_map == jnp.array([0.0, 0.0, (1.0 / 2.0) ** 2.0, 0.0])).all() def test__chi_squared_map_complex_from(): - data = np.array([10.0 + 10.0j, 10.0 + 10.0j]) - noise_map = np.array([2.0 + 2.0j, 2.0 + 2.0j]) - model_data = np.array([9.0 + 12.0j, 9.0 + 12.0j]) + data = jnp.array([10.0 + 10.0j, 10.0 + 10.0j]) + noise_map = jnp.array([2.0 + 2.0j, 2.0 + 2.0j]) + model_data = jnp.array([9.0 + 12.0j, 9.0 + 12.0j]) residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data) @@ -157,13 +156,13 @@ def test__chi_squared_map_complex_from(): residual_map=residual_map, noise_map=noise_map ) - assert (chi_squared_map == np.array([0.25 + 1.0j, 0.25 + 1.0j])).all() + assert (chi_squared_map == jnp.array([0.25 + 1.0j, 0.25 + 1.0j])).all() def test__chi_squared_with_noise_covariance_from(): resdiual_map = aa.Array2D.no_mask([[1.0, 1.0], [2.0, 2.0]], pixel_scales=1.0) - noise_covariance_matrix_inv = np.array( + noise_covariance_matrix_inv = jnp.array( [ [1.0, 1.0, 4.0, 0.0], [0.0, 1.0, 9.0, 0.0], @@ -181,10 +180,10 @@ def test__chi_squared_with_noise_covariance_from(): def test__chi_squared_with_mask_fast_from(): - data = np.array([10.0, 10.0, 10.0, 10.0]) - mask = np.array([True, False, False, True]) - noise_map = np.array([1.0, 2.0, 3.0, 4.0]) - model_data = np.array([11.0, 10.0, 9.0, 8.0]) + data = jnp.array([10.0, 10.0, 10.0, 10.0]) + mask = jnp.array([True, False, False, True]) + noise_map = jnp.array([1.0, 2.0, 3.0, 4.0]) + model_data = jnp.array([11.0, 10.0, 9.0, 8.0]) residual_map = aa.util.fit.residual_map_with_mask_from( data=data, mask=mask, model_data=model_data @@ -207,10 +206,10 @@ def test__chi_squared_with_mask_fast_from(): assert chi_squared == pytest.approx(chi_squared_fast, 1.0e-4) - data = np.array([[10.0, 10.0], [10.0, 10.0]]) - mask = np.array([[True, False], [False, True]]) - noise_map = np.array([[1.0, 2.0], [3.0, 4.0]]) - model_data = np.array([[11.0, 10.0], [9.0, 8.0]]) + data = jnp.array([[10.0, 10.0], [10.0, 10.0]]) + mask = jnp.array([[True, False], [False, True]]) + noise_map = jnp.array([[1.0, 2.0], [3.0, 4.0]]) + model_data = jnp.array([[11.0, 10.0], [9.0, 8.0]]) residual_map = aa.util.fit.residual_map_with_mask_from( data=data, mask=mask, model_data=model_data @@ -235,9 +234,9 @@ def test__chi_squared_with_mask_fast_from(): def test__log_likelihood_from(): - data = np.array([10.0, 10.0, 10.0, 10.0]) - noise_map = np.array([2.0, 2.0, 2.0, 2.0]) - model_data = np.array([10.0, 10.0, 10.0, 10.0]) + data = jnp.array([10.0, 10.0, 10.0, 10.0]) + noise_map = jnp.array([2.0, 2.0, 2.0, 2.0]) + model_data = jnp.array([10.0, 10.0, 10.0, 10.0]) residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data) @@ -255,17 +254,17 @@ def test__log_likelihood_from(): chi_squared = 0.0 noise_normalization = ( - np.log(2.0 * np.pi * (2.0**2.0)) - + np.log(2.0 * np.pi * (2.0**2.0)) - + np.log(2.0 * np.pi * (2.0**2.0)) - + np.log(2.0 * np.pi * (2.0**2.0)) + jnp.log(2.0 * jnp.pi * (2.0**2.0)) + + jnp.log(2.0 * jnp.pi * (2.0**2.0)) + + jnp.log(2.0 * jnp.pi * (2.0**2.0)) + + jnp.log(2.0 * jnp.pi * (2.0**2.0)) ) assert log_likelihood == pytest.approx( -0.5 * (chi_squared + noise_normalization), 1.0e-4 ) - model_data = np.array([11.0, 10.0, 9.0, 8.0]) + model_data = jnp.array([11.0, 10.0, 9.0, 8.0]) residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data) @@ -288,17 +287,17 @@ def test__log_likelihood_from(): ((1.0 / 2.0) ** 2.0) + 0.0 + ((1.0 / 2.0) ** 2.0) + ((2.0 / 2.0) ** 2.0) ) noise_normalization = ( - np.log(2.0 * np.pi * (2.0**2.0)) - + np.log(2.0 * np.pi * (2.0**2.0)) - + np.log(2.0 * np.pi * (2.0**2.0)) - + np.log(2.0 * np.pi * (2.0**2.0)) + jnp.log(2.0 * jnp.pi * (2.0**2.0)) + + jnp.log(2.0 * jnp.pi * (2.0**2.0)) + + jnp.log(2.0 * jnp.pi * (2.0**2.0)) + + jnp.log(2.0 * jnp.pi * (2.0**2.0)) ) assert log_likelihood == pytest.approx( -0.5 * (chi_squared + noise_normalization), 1.0e-4 ) - noise_map = np.array([1.0, 2.0, 3.0, 4.0]) + noise_map = jnp.array([1.0, 2.0, 3.0, 4.0]) residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data) @@ -318,10 +317,10 @@ def test__log_likelihood_from(): chi_squared = 1.0 + (1.0 / (3.0**2.0)) + 0.25 noise_normalization = ( - np.log(2 * np.pi * (1.0**2.0)) - + np.log(2 * np.pi * (2.0**2.0)) - + np.log(2 * np.pi * (3.0**2.0)) - + np.log(2 * np.pi * (4.0**2.0)) + jnp.log(2 * jnp.pi * (1.0**2.0)) + + jnp.log(2 * jnp.pi * (2.0**2.0)) + + jnp.log(2 * jnp.pi * (3.0**2.0)) + + jnp.log(2 * jnp.pi * (4.0**2.0)) ) assert log_likelihood == pytest.approx( @@ -330,10 +329,10 @@ def test__log_likelihood_from(): def test__log_likelihood_from__with_mask(): - data = np.array([10.0, 10.0, 10.0, 10.0]) - mask = np.array([True, False, False, True]) - noise_map = np.array([1.0, 2.0, 3.0, 4.0]) - model_data = np.array([11.0, 10.0, 9.0, 8.0]) + data = jnp.array([10.0, 10.0, 10.0, 10.0]) + mask = jnp.array([True, False, False, True]) + noise_map = jnp.array([1.0, 2.0, 3.0, 4.0]) + model_data = jnp.array([11.0, 10.0, 9.0, 8.0]) residual_map = aa.util.fit.residual_map_with_mask_from( data=data, mask=mask, model_data=model_data @@ -358,18 +357,18 @@ def test__log_likelihood_from__with_mask(): # chi squared = 0, 0.25, (0.25 and 1.0 are masked) chi_squared = 0.0 + (1.0 / 3.0) ** 2.0 - noise_normalization = np.log(2 * np.pi * (2.0**2.0)) + np.log( - 2 * np.pi * (3.0**2.0) + noise_normalization = jnp.log(2 * jnp.pi * (2.0**2.0)) + jnp.log( + 2 * jnp.pi * (3.0**2.0) ) assert log_likelihood == pytest.approx( -0.5 * (chi_squared + noise_normalization), 1e-4 ) - data = np.array([[10.0, 10.0], [10.0, 10.0]]) - mask = np.array([[True, False], [False, True]]) - noise_map = np.array([[1.0, 2.0], [3.0, 4.0]]) - model_data = np.array([[11.0, 10.0], [9.0, 8.0]]) + data = jnp.array([[10.0, 10.0], [10.0, 10.0]]) + mask = jnp.array([[True, False], [False, True]]) + noise_map = jnp.array([[1.0, 2.0], [3.0, 4.0]]) + model_data = jnp.array([[11.0, 10.0], [9.0, 8.0]]) residual_map = aa.util.fit.residual_map_with_mask_from( data=data, mask=mask, model_data=model_data @@ -394,8 +393,8 @@ def test__log_likelihood_from__with_mask(): # chi squared = 0, 0.25, (0.25 and 1.0 are masked) chi_squared = 0.0 + (1.0 / 3.0) ** 2.0 - noise_normalization = np.log(2 * np.pi * (2.0**2.0)) + np.log( - 2 * np.pi * (3.0**2.0) + noise_normalization = jnp.log(2 * jnp.pi * (2.0**2.0)) + jnp.log( + 2 * jnp.pi * (3.0**2.0) ) assert log_likelihood == pytest.approx( @@ -404,9 +403,9 @@ def test__log_likelihood_from__with_mask(): def test__log_likelihood_from__complex_data(): - data = np.array([10.0 + 10.0j, 10.0 + 10.0j]) - noise_map = np.array([2.0 + 1.0j, 2.0 + 1.0j]) - model_data = np.array([9.0 + 12.0j, 9.0 + 12.0j]) + data = jnp.array([10.0 + 10.0j, 10.0 + 10.0j]) + noise_map = jnp.array([2.0 + 1.0j, 2.0 + 1.0j]) + model_data = jnp.array([9.0 + 12.0j, 9.0 + 12.0j]) residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data) @@ -427,8 +426,8 @@ def test__log_likelihood_from__complex_data(): # chi squared = 0.25 and 4.0 chi_squared = 4.25 - noise_normalization = np.log(2 * np.pi * (2.0**2.0)) + np.log( - 2 * np.pi * (1.0**2.0) + noise_normalization = jnp.log(2 * jnp.pi * (2.0**2.0)) + jnp.log( + 2 * jnp.pi * (1.0**2.0) ) assert log_likelihood == pytest.approx( @@ -457,8 +456,8 @@ def test__log_evidence_from(): def test__residual_flux_fraction_map_from(): - data = np.array([10.0, 10.0, 10.0, 10.0]) - model_data = np.array([10.0, 10.0, 10.0, 10.0]) + data = jnp.array([10.0, 10.0, 10.0, 10.0]) + model_data = jnp.array([10.0, 10.0, 10.0, 10.0]) residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data) @@ -466,9 +465,9 @@ def test__residual_flux_fraction_map_from(): residual_map=residual_map, data=data ) - assert (residual_flux_fraction_map == np.array([0.0, 0.0, 0.0, 0.0])).all() + assert (residual_flux_fraction_map == jnp.array([0.0, 0.0, 0.0, 0.0])).all() - model_data = np.array([11.0, 10.0, 9.0, 8.0]) + model_data = jnp.array([11.0, 10.0, 9.0, 8.0]) residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data) @@ -476,13 +475,13 @@ def test__residual_flux_fraction_map_from(): residual_map=residual_map, data=data ) - assert (residual_flux_fraction_map == np.array([-0.1, 0.0, 0.1, 0.2])).all() + assert (residual_flux_fraction_map == jnp.array([-0.1, 0.0, 0.1, 0.2])).all() def test__residual_flux_fraction_map_with_mask_from(): - data = np.array([10.0, 10.0, 10.0, 10.0]) - mask = np.array([True, False, False, True]) - model_data = np.array([11.0, 10.0, 9.0, 8.0]) + data = jnp.array([10.0, 10.0, 10.0, 10.0]) + mask = jnp.array([True, False, False, True]) + model_data = jnp.array([11.0, 10.0, 9.0, 8.0]) residual_map = aa.util.fit.residual_map_with_mask_from( data=data, mask=mask, model_data=model_data @@ -492,9 +491,9 @@ def test__residual_flux_fraction_map_with_mask_from(): residual_map=residual_map, mask=mask, data=data ) - assert (residual_flux_fraction_map == np.array([0.0, 0.0, 0.1, 0.0])).all() + assert (residual_flux_fraction_map == jnp.array([0.0, 0.0, 0.1, 0.0])).all() - model_data = np.array([11.0, 9.0, 8.0, 8.0]) + model_data = jnp.array([11.0, 9.0, 8.0, 8.0]) residual_map = aa.util.fit.residual_map_with_mask_from( data=data, mask=mask, model_data=model_data @@ -504,4 +503,4 @@ def test__residual_flux_fraction_map_with_mask_from(): residual_map=residual_map, mask=mask, data=data ) - assert (residual_flux_fraction_map == np.array([0.0, 0.1, 0.2, 0.0])).all() + assert (residual_flux_fraction_map == jnp.array([0.0, 0.1, 0.2, 0.0])).all() From 340c34d5faea26adb8616df689ba09034b76fdca Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Fri, 4 Apr 2025 16:13:21 +0100 Subject: [PATCH 18/25] pylops removed --- autoarray/__init__.py | 1 - autoarray/inversion/inversion/factory.py | 11 -- .../inversion/inversion/interferometer/lop.py | 146 ------------------ .../inversion/regularization/abstract.py | 23 --- autoarray/operators/transformer.py | 29 +--- autoarray/preloads.py | 14 -- autoarray/structures/visibilities.py | 4 +- optional_requirements.txt | 1 - .../test_inversion_interferometer_util.py | 17 -- 9 files changed, 3 insertions(+), 243 deletions(-) delete mode 100644 autoarray/inversion/inversion/interferometer/lop.py diff --git a/autoarray/__init__.py b/autoarray/__init__.py index 1a6d79c53..2300fb5d4 100644 --- a/autoarray/__init__.py +++ b/autoarray/__init__.py @@ -44,7 +44,6 @@ from .inversion.inversion.imaging.w_tilde import InversionImagingWTilde from .inversion.inversion.interferometer.w_tilde import InversionInterferometerWTilde from .inversion.inversion.interferometer.mapping import InversionInterferometerMapping -from .inversion.inversion.interferometer.lop import InversionInterferometerMappingPyLops from .inversion.linear_obj.linear_obj import LinearObj from .inversion.linear_obj.func_list import AbstractLinearObjFuncList from .mask.derive.indexes_2d import DeriveIndexes2D diff --git a/autoarray/inversion/inversion/factory.py b/autoarray/inversion/inversion/factory.py index 327262786..350f19e65 100644 --- a/autoarray/inversion/inversion/factory.py +++ b/autoarray/inversion/inversion/factory.py @@ -9,9 +9,6 @@ from autoarray.inversion.inversion.interferometer.w_tilde import ( InversionInterferometerWTilde, ) -from autoarray.inversion.inversion.interferometer.lop import ( - InversionInterferometerMappingPyLops, -) from autoarray.inversion.inversion.dataset_interface import DatasetInterface from autoarray.inversion.linear_obj.linear_obj import LinearObj from autoarray.inversion.linear_obj.func_list import AbstractLinearObjFuncList @@ -212,11 +209,3 @@ def inversion_interferometer_from( settings=settings, run_time_dict=run_time_dict, ) - - else: - return InversionInterferometerMappingPyLops( - dataset=dataset, - linear_obj_list=linear_obj_list, - settings=settings, - run_time_dict=run_time_dict, - ) diff --git a/autoarray/inversion/inversion/interferometer/lop.py b/autoarray/inversion/inversion/interferometer/lop.py deleted file mode 100644 index fdd6b8adf..000000000 --- a/autoarray/inversion/inversion/interferometer/lop.py +++ /dev/null @@ -1,146 +0,0 @@ -from scipy import sparse - -import numpy as np -from typing import Dict - -from autoconf import cached_property - -from autoarray.inversion.inversion.interferometer.abstract import ( - AbstractInversionInterferometer, -) -from autoarray.inversion.linear_obj.linear_obj import LinearObj -from autoarray.structures.visibilities import Visibilities - -from autoarray.numba_util import profile_func - - -class InversionInterferometerMappingPyLops(AbstractInversionInterferometer): - """ - Constructs linear equations (via vectors and matrices) which allow for sets of simultaneous linear equations - to be solved (see `inversion.inversion.abstract.AbstractInversion` for a full description). - - A linear object describes the mappings between values in observed `data` and the linear object's model via its - `mapping_matrix`. This class constructs linear equations for `Interferometer` objects, where the data is an - an array of visibilities and the mappings include a non-uniform fast Fourier transform operation described by - the interferometer dataset's transformer. - - This class uses the mapping formalism, which constructs the simultaneous linear equations using the - `mapping_matrix` of every linear object. This is performed using the library PyLops, which uses linear - operators to avoid these matrices being created explicitly in memory, making the calculation more efficient. - """ - - @cached_property - @profile_func - def reconstruction(self): - """ - Solve the linear system [F + reg_coeff*H] S = D -> S = [F + reg_coeff*H]^-1 D given by equation (12) - of https://arxiv.org/pdf/astro-ph/0302587.pdf - - S is the vector of reconstructed inversion values. - """ - - import pylops - - Aop = pylops.MatrixMult( - sparse.bsr_matrix(self.linear_obj_list[0].mapping_matrix) - ) - - Fop = self.transformer - - Op = Fop * Aop - - MOp = pylops.MatrixMult(sparse.bsr_matrix(self.preconditioner_matrix_inverse)) - - try: - return pylops.NormalEquationsInversion( - Op=Op, - Regs=None, - epsNRs=[1.0], - data=self.data.ordered_1d, - Weight=pylops.Diagonal(diag=self.noise_map.weight_list_ordered_1d), - NRegs=[ - pylops.MatrixMult(sparse.bsr_matrix(self.regularization_matrix)) - ], - M=MOp, - tol=self.settings.tolerance, - atol=self.settings.tolerance, - **dict(maxiter=self.settings.maxiter), - ) - except AttributeError: - return pylops.normal_equations_inversion( - Op=Op, - Regs=None, - epsNRs=[1.0], - y=self.data.ordered_1d, - Weight=pylops.Diagonal(diag=self.noise_map.weight_list_ordered_1d), - NRegs=[ - pylops.MatrixMult(sparse.bsr_matrix(self.regularization_matrix)) - ], - M=MOp, - tol=self.settings.tolerance, - atol=self.settings.tolerance, - **dict(maxiter=self.settings.maxiter), - )[0] - - @property - @profile_func - def mapped_reconstructed_data_dict( - self, - ) -> Dict[LinearObj, Visibilities]: - """ - When constructing the simultaneous linear equations (via vectors and matrices) the quantities of each individual - linear object (e.g. their `mapping_matrix`) are combined into single ndarrays. This does not track which - quantities belong to which linear objects, therefore the linear equation's solutions (which are returned as - ndarrays) do not contain information on which linear object(s) they correspond to. - - For example, consider if two `Mapper` objects with 50 and 100 source pixels are used in an `Inversion`. - The `reconstruction` (which contains the solved for source pixels values) is an ndarray of shape [150], but - the ndarray itself does not track which values belong to which `Mapper`. - - This function converts an ndarray of a `reconstruction` to a dictionary of ndarrays containing each linear - object's reconstructed images, where the keys are the instances of each mapper in the inversion. - - The PyLops calculation bypasses the calculation of the `mapping_matrix` and it therefore cannot be used to map - the reconstruction's values to the image-plane. Instead, the unique data-to-pixelization mappings are used, - including the 2D non-uniform fast Fourier transform operation after mapping is complete. - - Parameters - ---------- - reconstruction - The reconstruction (in the source frame) whose values are mapped to a dictionary of values for each - individual mapper (in the image-plane). - """ - - mapped_reconstructed_image_dict = self.mapped_reconstructed_image_dict - - return { - linear_obj: self.transformer.visibilities_from(image=image) - for linear_obj, image in mapped_reconstructed_image_dict.items() - } - - @cached_property - @profile_func - def preconditioner_matrix(self): - curvature_matrix_approx = np.multiply( - np.sum(self.noise_map.weight_list_ordered_1d), - self.linear_obj_list[0].mapping_matrix.T - @ self.linear_obj_list[0].mapping_matrix, - ) - - return np.add(curvature_matrix_approx, self.regularization_matrix) - - @cached_property - @profile_func - def preconditioner_matrix_inverse(self): - return np.linalg.inv(self.preconditioner_matrix) - - @cached_property - @profile_func - def log_det_curvature_reg_matrix_term(self): - return 2.0 * np.sum( - np.log(np.diag(np.linalg.cholesky(self.preconditioner_matrix))) - ) - - @property - def reconstruction_noise_map(self): - return None diff --git a/autoarray/inversion/regularization/abstract.py b/autoarray/inversion/regularization/abstract.py index 47cadc302..838eaf942 100644 --- a/autoarray/inversion/regularization/abstract.py +++ b/autoarray/inversion/regularization/abstract.py @@ -5,13 +5,6 @@ if TYPE_CHECKING: from autoarray.inversion.linear_obj.linear_obj import LinearObj -try: - import pylops - - PyLopsOperator = pylops.LinearOperator -except ModuleNotFoundError: - PyLopsOperator = object - class AbstractRegularization: def __init__(self): @@ -174,19 +167,3 @@ def regularization_matrix_from(self, linear_obj: LinearObj) -> np.ndarray: The regularization matrix. """ raise NotImplementedError - - -class RegularizationLop(PyLopsOperator): - def __init__(self, regularization_matrix): - self.regularization_matrix = regularization_matrix - self.pixels = regularization_matrix.shape[0] - self.dims = self.pixels - self.shape = (self.pixels, self.pixels) - self.dtype = dtype - self.explicit = False - - def _matvec(self, x): - return np.dot(self.regularization_matrix, x) - - def _rmatvec(self, x): - return np.dot(self.regularization_matrix.T, x) diff --git a/autoarray/operators/transformer.py b/autoarray/operators/transformer.py index acbe6bb73..3b19f7f7c 100644 --- a/autoarray/operators/transformer.py +++ b/autoarray/operators/transformer.py @@ -8,21 +8,11 @@ class NUFFTPlaceholder: pass -class PyLopsPlaceholder: - pass - - try: from pynufft.linalg.nufft_cpu import NUFFT_cpu except ModuleNotFoundError: NUFFT_cpu = NUFFTPlaceholder -try: - import pylops - - PyLopsOperator = pylops.LinearOperator -except ModuleNotFoundError: - PyLopsOperator = PyLopsPlaceholder from autoarray.structures.arrays.uniform_2d import Array2D from autoarray.structures.grids.uniform_2d import Grid2D @@ -42,20 +32,8 @@ def pynufft_exception(): ) -def pylops_exception(): - raise ModuleNotFoundError( - "\n--------------------\n" - "You are attempting to perform interferometer analysis.\n\n" - "However, the optional library PyLops (https://github.com/PyLops/pylops) is not installed.\n\n" - "Install it via the command `pip install pylops==2.3.1`.\n\n" - "----------------------" - ) - - -class TransformerDFT(PyLopsOperator): +class TransformerDFT: def __init__(self, uv_wavelengths, real_space_mask, preload_transform=True): - if isinstance(self, PyLopsPlaceholder): - pylops_exception() super().__init__() @@ -146,14 +124,11 @@ def transform_mapping_matrix(self, mapping_matrix): ) -class TransformerNUFFT(NUFFT_cpu, PyLopsOperator): +class TransformerNUFFT(NUFFT_cpu): def __init__(self, uv_wavelengths, real_space_mask): if isinstance(self, NUFFTPlaceholder): pynufft_exception() - if isinstance(self, PyLopsPlaceholder): - pylops_exception() - super(TransformerNUFFT, self).__init__() self.uv_wavelengths = uv_wavelengths diff --git a/autoarray/preloads.py b/autoarray/preloads.py index 6808f0f6c..90c1deae9 100644 --- a/autoarray/preloads.py +++ b/autoarray/preloads.py @@ -198,13 +198,6 @@ def set_mapper_list(self, fit_0, fit_1): if fit_0.inversion.total(cls=AbstractMapper) == 0: return - from autoarray.inversion.inversion.interferometer.lop import ( - InversionInterferometerMappingPyLops, - ) - - if isinstance(fit_0.inversion, InversionInterferometerMappingPyLops): - return - inversion_0 = fit_0.inversion inversion_1 = fit_1.inversion @@ -238,13 +231,6 @@ def set_operated_mapping_matrix_with_preloads(self, fit_0, fit_1): self.operated_mapping_matrix = None - from autoarray.inversion.inversion.interferometer.lop import ( - InversionInterferometerMappingPyLops, - ) - - if isinstance(fit_0.inversion, InversionInterferometerMappingPyLops): - return - inversion_0 = fit_0.inversion inversion_1 = fit_1.inversion diff --git a/autoarray/structures/visibilities.py b/autoarray/structures/visibilities.py index 2b4113c7c..8cc94dca5 100644 --- a/autoarray/structures/visibilities.py +++ b/autoarray/structures/visibilities.py @@ -213,9 +213,7 @@ def __init__(self, visibilities: Union[np.ndarray, List[complex]], *args, **kwar A collection of (real, imag) visibilities noise-map values which are used to represent the noise-map in an `Interferometer` dataset. - This data structure behaves the same as the `Visibilities` structure (see `AbstractVisibilities.__new__`). The - only difference is that it includes a `WeightOperator` used by `Inversion`'s which use `LinearOperators` and - the library `PyLops` to fit `Interferometer` data. + This data structure behaves the same as the `Visibilities` structure (see `AbstractVisibilities.__new__`). Parameters ---------- diff --git a/optional_requirements.txt b/optional_requirements.txt index cbb64e17e..4ba479740 100644 --- a/optional_requirements.txt +++ b/optional_requirements.txt @@ -1,4 +1,3 @@ -pylops>=1.10.0,<=2.3.1 pynufft #jax==0.4.3 #jaxlib==0.4.3 diff --git a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py index 8875ea378..cb7d27673 100644 --- a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py +++ b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py @@ -68,23 +68,6 @@ def test__data_vector_via_transformed_mapping_matrix_from(): assert (data_vector_complex_via_blurred == data_vector_via_transformed).all() -def test__inversion_interferometer__via_mapper( - interferometer_7_no_fft, - rectangular_mapper_7x7_3x3, - delaunay_mapper_9_3x3, - voronoi_mapper_9_3x3, - regularization_constant, -): - inversion = aa.Inversion( - dataset=interferometer_7_no_fft, - linear_obj_list=[rectangular_mapper_7x7_3x3], - settings=aa.SettingsInversion(use_linear_operators=True), - ) - - assert isinstance(inversion.linear_obj_list[0], aa.MapperRectangular) - assert isinstance(inversion, aa.InversionInterferometerMappingPyLops) - - def test__w_tilde_curvature_interferometer_from(): noise_map = np.array([1.0, 2.0, 3.0]) uv_wavelengths = np.array([[0.0001, 2.0, 3000.0], [3000.0, 2.0, 0.0001]]) From 39fbf72a9776d25ccdadb0ed0b9825e49b20f1dd Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Fri, 4 Apr 2025 16:28:34 +0100 Subject: [PATCH 19/25] unit tests fixed --- autoarray/fit/fit_util.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/autoarray/fit/fit_util.py b/autoarray/fit/fit_util.py index 10f24f9a7..f28a64e22 100644 --- a/autoarray/fit/fit_util.py +++ b/autoarray/fit/fit_util.py @@ -158,8 +158,8 @@ def chi_squared_complex_from(*, chi_squared_map: jnp.ndarray) -> float: chi_squared_map The chi-squared-map of values of the model-data fit to the dataset. """ - chi_squared_real = jnp.sum(chi_squared_map.real) - chi_squared_imag = jnp.sum(chi_squared_map.imag) + chi_squared_real = jnp.sum(np.array(chi_squared_map.real)) + chi_squared_imag = jnp.sum(np.array(chi_squared_map.imag)) return chi_squared_real + chi_squared_imag @@ -174,8 +174,8 @@ def noise_normalization_complex_from(*, noise_map: jnp.ndarray) -> float: noise_map The masked noise-map of the dataset. """ - noise_normalization_real = jnp.sum(jnp.log(2 * jnp.pi * noise_map.real**2.0)) - noise_normalization_imag = jnp.sum(jnp.log(2 * jnp.pi * noise_map.imag**2.0)) + noise_normalization_real = jnp.sum(jnp.log(2 * jnp.pi * np.array(noise_map).real**2.0)) + noise_normalization_imag = jnp.sum(jnp.log(2 * jnp.pi * np.array(noise_map).imag**2.0)) return noise_normalization_real + noise_normalization_imag From dd7ca45b8c13f64f300e426530218d3e404c6295 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sat, 5 Apr 2025 14:37:03 +0100 Subject: [PATCH 20/25] inversion now uses Convoler again, all numba, factory tests pass --- autoarray/__init__.py | 1 + autoarray/dataset/imaging/dataset.py | 21 + autoarray/inversion/convolver.py | 544 ++++++++++++++++++ autoarray/inversion/inversion/abstract.py | 4 + .../inversion/inversion/dataset_interface.py | 2 + .../inversion/inversion/imaging/abstract.py | 6 +- .../inversion/inversion/imaging/mapping.py | 4 +- .../inversion/inversion/imaging/w_tilde.py | 12 +- 8 files changed, 583 insertions(+), 11 deletions(-) create mode 100644 autoarray/inversion/convolver.py diff --git a/autoarray/__init__.py b/autoarray/__init__.py index 2300fb5d4..5acddd535 100644 --- a/autoarray/__init__.py +++ b/autoarray/__init__.py @@ -20,6 +20,7 @@ from .fit.fit_imaging import FitImaging from .fit.fit_interferometer import FitInterferometer from .geometry.geometry_2d import Geometry2D +from .inversion.convolver import Convolver from .inversion.pixelization.mappers.abstract import AbstractMapper from .inversion.pixelization import mesh from .inversion.pixelization import image_mesh diff --git a/autoarray/dataset/imaging/dataset.py b/autoarray/dataset/imaging/dataset.py index 8d84ee1b8..9fa00da1a 100644 --- a/autoarray/dataset/imaging/dataset.py +++ b/autoarray/dataset/imaging/dataset.py @@ -179,6 +179,27 @@ def grids(self): psf=self.psf, ) + @cached_property + def convolver(self): + """ + Returns a `Convolver` from a mask and 2D PSF kernel. + + The `Convolver` stores in memory the array indexing between the mask and PSF, enabling efficient 2D PSF + convolution of images and matrices used for linear algebra calculations (see `operators.convolver`). + + This uses lazy allocation such that the calculation is only performed when the convolver is used, ensuring + efficient set up of the `Imaging` class. + + Returns + ------- + Convolver + The convolver given the masked imaging data's mask and PSF. + """ + + from autoarray.inversion.convolver import Convolver + + return Convolver(mask=self.mask, kernel=self.psf) + @cached_property def w_tilde(self): """ diff --git a/autoarray/inversion/convolver.py b/autoarray/inversion/convolver.py new file mode 100644 index 000000000..560009f13 --- /dev/null +++ b/autoarray/inversion/convolver.py @@ -0,0 +1,544 @@ +from autoarray import numba_util +import numpy as np + +from autoarray.structures.arrays.uniform_2d import Array2D + +from autoarray import exc +from autoarray.mask import mask_2d_util + + +class Convolver: + def __init__(self, mask, kernel): + """ + Class to setup the 1D convolution of an / mapping matrix. + + Take a simple 3x3 and masks: + + [[2, 8, 2], + [5, 7, 5], + [3, 1, 4]] + + [[True, False, True], (True means that the value is masked) + [False, False, False], + [True, False, True]] + + A set of values in a corresponding 1d array of this might be represented as: + + [2, 8, 2, 5, 7, 5, 3, 1, 4] + + and after masking as: + + [8, 5, 7, 5, 1] + + Setup is required to perform 2D real-space convolution on the masked array. This module finds the \ + relationship between the unmasked 2D data, masked data and kernel, so that 2D real-space convolutions \ + can be efficiently applied to reduced 1D masked structures. + + This calculation also accounts for the blurring of light outside of the masked regions which blurs into \ + the masked region. + + + **IMAGE FRAMES:** + + For a masked in 2D, one can compute for every pixel all of the unmasked pixels it will blur light into for \ + a given PSF kernel size, e.g.: + + IxIxIxIxIxIxIxIxIxIxI + IxIxIxIxIxIxIxIxIxIxI This is an imaging.Mask2D, where: + IxIxIxIxIxIxIxIxIxIxI + IxIxIxIxIxIxIxIxIxIxI x = `True` (Pixel is masked and excluded from lens) + IxIxIxIoIoIoIxIxIxIxI o = `False` (Pixel is not masked and included in lens) + IxIxIxIoIoIoIxIxIxIxI + IxIxIxIoIoIoIxIxIxIxI + IxIxIxIxIxIxIxIxIxIxI + IxIxIxIxIxIxIxIxIxIxI + IxIxIxIxIxIxIxIxIxIxI + + Here, there are 9 unmasked pixels. Indexing of each unmasked pixel goes from the top-left corner right and \ + downwards, therefore: + + IxIxIxIxIxIxIxIxIxIxI + IxIxIxIxIxIxIxIxIxIxI + IxIxIxIxIxIxIxIxIxIxI + IxIxIxIxIxIxIxIxIxIxI + IxIxIxI0I1I2IxIxIxIxI + IxIxIxI3I4I5IxIxIxIxI + IxIxIxI6I7I8IxIxIxIxI + IxIxIxIxIxIxIxIxIxIxI + IxIxIxIxIxIxIxIxIxIxI + IxIxIxIxIxIxIxIxIxIxI + + For every unmasked pixel, the Convolver over-lays the PSF and computes three quantities; + + image_frame_indexes - The indexes of all masked pixels it will blur light into. + image_frame_kernels - The kernel values that overlap each masked pixel it will blur light into. + image_frame_length - The number of masked pixels it will blur light into (unmasked pixels are excluded) + + For example, if we had the following 3x3 kernel: + + I0.1I0.2I0.3I + I0.4I0.5I0.6I + I0.7I0.8I0.9I + + For pixel 0 above, when we overlap the kernel 4 unmasked pixels overlap this kernel, such that: + + image_frame_indexes = [0, 1, 3, 4] + image_frame_kernels = [0.5, 0.6, 0.8, 0.9] + image_frame_length = 4 + + Noting that the other 5 kernel values (0.1, 0.2, 0.3, 0.4, 0.7) overlap masked pixels and are thus discarded. + + For pixel 1, we get the following results: + + image_frame_indexes = [0, 1, 2, 3, 4, 5] + image_frame_kernels = [0.4, 0.5, 0.6, 0.7, 0.8, 0.9] + image_frame_lengths = 6 + + In the majority of cases, the kernel will overlap only unmasked pixels. This is the case above for \ + central pixel 4, where: + + image_frame_indexes = [0, 1, 2, 3, 4, 5, 6, 7, 8] + image_frame_kernels = [0,1, 0.2, 0,3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9] + image_frame_lengths = 9 + + Once we have set up all these quantities, the convolution routine simply uses them to convolve a 1D array of a + masked or the masked of a util in the inversion module. + + + **BLURRING FRAMES:** + + Whilst the scheme above accounts for all blurred light within the masks, it does not account for the fact that \ + pixels outside of the masks will also blur light into it. This effect is accounted for using blurring frames. + + It is omitted for mapping matrix blurring, as an inversion does not fit data outside of the masks. + + First, a blurring masks is computed from a masks, which describes all pixels which are close enough to the masks \ + to blur light into it for a given kernel size. Following the example above, the following blurring masks is \ + computed: + + IxIxIxIxIxIxIxIxIxIxI + IxIxIxIxIxIxIxIxIxIxI This is an example grid.Mask2D, where: + IxIxIxIxIxIxIxIxIxIxI + IxIxIoIoIoIoIoIxIxIxI x = `True` (Pixel is masked and excluded from lens) + IxIxIoIxIxIxIoIxIxIxI o = `False` (Pixel is not masked and included in lens) + IxIxIoIxIxIxIoIxIxIxI + IxIxIoIxIxIxIoIxIxIxI + IxIxIoIoIoIoIoIxIxIxI + IxIxIxIxIxIxIxIxIxIxI + IxIxIxIxIxIxIxIxIxIxI + + Indexing again goes from the top-left corner right and downwards: + + IxIxI xI xI xI xI xIxIxIxI + IxIxI xI xI xI xI xIxIxIxI + IxIxI xI xI xI xI xIxIxIxI + IxIxI 0I 1I 2I 3I 4IxIxIxI + IxIxI 5I xI xI xI 6IxIxIxI + IxIxI 7I xI xI xI 8IxIxIxI + IxIxI 9I xI xI xI10IxIxIxI + IxIxI11I12I13I14I15IxIxIxI + IxIxI xI xI xI xI xIxIxIxI + IxIxI xI xI xI xI xIxIxIxI + + For every unmasked blurring-pixel, the Convolver over-lays the PSF kernel and computes three quantities; + + blurring_frame_indexes - The indexes of all unmasked pixels (not unmasked blurring pixels) it will \ + blur light into. + bluring_frame_kernels - The kernel values that overlap each pixel it will blur light into. + blurring_frame_length - The number of pixels it will blur light into. + + The blurring frame therefore does not perform any blurring which blurs light into other blurring pixels. \ + It only performs computations which add light inside of the masks. + + For pixel 0 above, when we overlap the 3x3 kernel above only 1 unmasked pixels overlaps the kernel, such that: + + blurring_frame_indexes = [0] (This 0 refers to pixel 0 within the masks, not blurring_frame_pixel 0) + blurring_frame_kernels = [0.9] + blurring_frame_length = 1 + + For pixel 1 above, when we overlap the 3x3 kernel above 2 unmasked pixels overlap the kernel, such that: + + blurring_frame_indexes = [0, 1] (This 0 and 1 refer to pixels 0 and 1 within the masks) + blurring_frame_kernels = [0.8, 0.9] + blurring_frame_length = 2 + + For pixel 3 above, when we overlap the 3x3 kernel above 3 unmasked pixels overlap the kernel, such that: + + blurring_frame_indexes = [0, 1, 2] (Again, these are pixels 0, 1 and 2) + blurring_frame_kernels = [0.7, 0.8, 0.9] + blurring_frame_length = 3 + + Parameters + ---------- + mask + The mask within which the convolved signal is calculated. + blurring_mask + A masks of pixels outside the masks but whose light blurs into it after PSF convolution. + kernel : grid.PSF or ndarray + An array representing a PSF. + """ + if kernel.shape_native[0] % 2 == 0 or kernel.shape_native[1] % 2 == 0: + raise exc.KernelException("PSF kernel must be odd") + + self.mask = mask + + self.mask_index_array = np.full(mask.shape, -1) + self.pixels_in_mask = int(np.size(mask) - np.sum(mask)) + + count = 0 + for x in range(mask.shape[0]): + for y in range(mask.shape[1]): + if not mask[x, y]: + self.mask_index_array[x, y] = count + count += 1 + + self.kernel = kernel + self.kernel_max_size = self.kernel.shape_native[0] * self.kernel.shape_native[1] + + mask_1d_index = 0 + self.image_frame_1d_indexes = np.zeros( + (self.pixels_in_mask, self.kernel_max_size), dtype="int" + ) + self.image_frame_1d_kernels = np.zeros( + (self.pixels_in_mask, self.kernel_max_size) + ) + self.image_frame_1d_lengths = np.zeros((self.pixels_in_mask), dtype="int") + for x in range(self.mask_index_array.shape[0]): + for y in range(self.mask_index_array.shape[1]): + if not mask[x][y]: + ( + image_frame_1d_indexes, + image_frame_1d_kernels, + ) = self.frame_at_coordinates_jit( + coordinates=(x, y), + mask=np.array(mask), + mask_index_array=self.mask_index_array, + kernel_2d=np.array(self.kernel.native.array), + ) + self.image_frame_1d_indexes[ + mask_1d_index, : + ] = image_frame_1d_indexes + self.image_frame_1d_kernels[ + mask_1d_index, : + ] = image_frame_1d_kernels + self.image_frame_1d_lengths[mask_1d_index] = image_frame_1d_indexes[ + image_frame_1d_indexes >= 0 + ].shape[0] + mask_1d_index += 1 + + self.blurring_mask = mask_2d_util.blurring_mask_2d_from( + mask_2d=np.array(mask), + kernel_shape_native=kernel.shape_native, + ) + + self.pixels_in_blurring_mask = int( + np.size(self.blurring_mask) - np.sum(self.blurring_mask) + ) + + mask_1d_index = 0 + self.blurring_frame_1d_indexes = np.zeros( + (self.pixels_in_blurring_mask, self.kernel_max_size), dtype="int" + ) + self.blurring_frame_1d_kernels = np.zeros( + (self.pixels_in_blurring_mask, self.kernel_max_size) + ) + self.blurring_frame_1d_lengths = np.zeros( + (self.pixels_in_blurring_mask), dtype="int" + ) + for x in range(mask.shape[0]): + for y in range(mask.shape[1]): + if mask[x][y] and not self.blurring_mask[x, y]: + ( + image_frame_1d_indexes, + image_frame_1d_kernels, + ) = self.frame_at_coordinates_jit( + coordinates=(x, y), + mask=np.array(mask), + mask_index_array=np.array(self.mask_index_array), + kernel_2d=np.array(self.kernel.native.array), + ) + self.blurring_frame_1d_indexes[ + mask_1d_index, : + ] = image_frame_1d_indexes + self.blurring_frame_1d_kernels[ + mask_1d_index, : + ] = image_frame_1d_kernels + self.blurring_frame_1d_lengths[ + mask_1d_index + ] = image_frame_1d_indexes[image_frame_1d_indexes >= 0].shape[0] + mask_1d_index += 1 + + @staticmethod + @numba_util.jit() + def frame_at_coordinates_jit(coordinates, mask, mask_index_array, kernel_2d): + """ + Returns the frame (indexes of pixels light is blurred into) and kernel_frame (kernel kernel values of those \ + pixels) for a given coordinate in a masks and its PSF. + + Parameters + ---------- + coordinates: Tuple[int, int] + The coordinates of mask_index_array on which the frame should be centred + kernel_shape_native: Tuple[int, int] + The shape of the kernel for which this frame will be used + """ + + kernel_shape_native = kernel_2d.shape + kernel_max_size = kernel_shape_native[0] * kernel_shape_native[1] + + half_x = int(kernel_shape_native[0] / 2) + half_y = int(kernel_shape_native[1] / 2) + + frame = -1 * np.ones((kernel_max_size)) + kernel_frame = -1.0 * np.ones((kernel_max_size)) + + count = 0 + for i in range(kernel_shape_native[0]): + for j in range(kernel_shape_native[1]): + x = coordinates[0] - half_x + i + y = coordinates[1] - half_y + j + if ( + 0 <= x < mask_index_array.shape[0] + and 0 <= y < mask_index_array.shape[1] + ): + value = mask_index_array[x, y] + if value >= 0 and not mask[x, y]: + frame[count] = value + kernel_frame[count] = kernel_2d[i, j] + count += 1 + + return frame, kernel_frame + + def convolve_image(self, image, blurring_image): + """ + For a given 1D array and blurring array, convolve the two using this convolver. + + Parameters + ---------- + image + 1D array of the values which are to be blurred with the convolver's PSF. + blurring_image + 1D array of the blurring values which blur into the array after PSF convolution. + """ + + if self.blurring_mask is None: + raise exc.KernelException( + "You cannot use the convolve_image function of a Convolver if the Convolver was" + "not created with a blurring_mask." + ) + + convolved_image = self.convolve_jit( + image_1d_array=np.array(image.slim), + image_frame_1d_indexes=self.image_frame_1d_indexes, + image_frame_1d_kernels=self.image_frame_1d_kernels, + image_frame_1d_lengths=self.image_frame_1d_lengths, + blurring_1d_array=np.array(blurring_image.slim), + blurring_frame_1d_indexes=self.blurring_frame_1d_indexes, + blurring_frame_1d_kernels=self.blurring_frame_1d_kernels, + blurring_frame_1d_lengths=self.blurring_frame_1d_lengths, + ) + + return Array2D(values=convolved_image, mask=self.mask) + + @staticmethod + @numba_util.jit() + def convolve_jit( + image_1d_array, + image_frame_1d_indexes, + image_frame_1d_kernels, + image_frame_1d_lengths, + blurring_1d_array, + blurring_frame_1d_indexes, + blurring_frame_1d_kernels, + blurring_frame_1d_lengths, + ): + blurred_image_1d = np.zeros(image_1d_array.shape) + + for image_1d_index in range(len(image_1d_array)): + frame_1d_indexes = image_frame_1d_indexes[image_1d_index] + frame_1d_kernel = image_frame_1d_kernels[image_1d_index] + frame_1d_length = image_frame_1d_lengths[image_1d_index] + image_value = image_1d_array[image_1d_index] + + for kernel_1d_index in range(frame_1d_length): + vector_index = frame_1d_indexes[kernel_1d_index] + kernel_value = frame_1d_kernel[kernel_1d_index] + blurred_image_1d[vector_index] += image_value * kernel_value + + for blurring_1d_index in range(len(blurring_1d_array)): + frame_1d_indexes = blurring_frame_1d_indexes[blurring_1d_index] + frame_1d_kernel = blurring_frame_1d_kernels[blurring_1d_index] + frame_1d_length = blurring_frame_1d_lengths[blurring_1d_index] + image_value = blurring_1d_array[blurring_1d_index] + + for kernel_1d_index in range(frame_1d_length): + vector_index = frame_1d_indexes[kernel_1d_index] + kernel_value = frame_1d_kernel[kernel_1d_index] + blurred_image_1d[vector_index] += image_value * kernel_value + + return blurred_image_1d + + def convolve_image_no_blurring(self, image): + """For a given 1D array and blurring array, convolve the two using this convolver. + + Parameters + ---------- + image + 1D array of the values which are to be blurred with the convolver's PSF. + """ + + convolved_image = self.convolve_no_blurring_jit( + image_1d_array=np.array(image.slim), + image_frame_1d_indexes=self.image_frame_1d_indexes, + image_frame_1d_kernels=self.image_frame_1d_kernels, + image_frame_1d_lengths=self.image_frame_1d_lengths, + ) + + return Array2D(values=convolved_image, mask=self.mask) + + def convolve_image_no_blurring_interpolation(self, image): + """For a given 1D array and blurring array, convolve the two using this convolver. + + Parameters + ---------- + image + 1D array of the values which are to be blurred with the convolver's PSF. + """ + + convolved_image = self.convolve_no_blurring_jit( + image_1d_array=image, + image_frame_1d_indexes=self.image_frame_1d_indexes, + image_frame_1d_kernels=self.image_frame_1d_kernels, + image_frame_1d_lengths=self.image_frame_1d_lengths, + ) + + return Array2D(values=convolved_image, mask=self.mask) + + @staticmethod + @numba_util.jit() + def convolve_no_blurring_jit( + image_1d_array, + image_frame_1d_indexes, + image_frame_1d_kernels, + image_frame_1d_lengths, + ): + blurred_image_1d = np.zeros(image_1d_array.shape) + + for image_1d_index in range(len(image_1d_array)): + frame_1d_indexes = image_frame_1d_indexes[image_1d_index] + frame_1d_kernel = image_frame_1d_kernels[image_1d_index] + frame_1d_length = image_frame_1d_lengths[image_1d_index] + image_value = image_1d_array[image_1d_index] + + for kernel_1d_index in range(frame_1d_length): + vector_index = frame_1d_indexes[kernel_1d_index] + kernel_value = frame_1d_kernel[kernel_1d_index] + blurred_image_1d[vector_index] += image_value * kernel_value + + return blurred_image_1d + + def convolve_mapping_matrix(self, mapping_matrix): + """For a given inversion mapping matrix, convolve every pixel's mapped with the PSF kernel. + + A mapping matrix provides non-zero entries in all elements which map two pixels to one another + (see *inversions.mappers*). + + For example, lets take an which is masked using a 'cross' of 5 pixels: + + [[ True, False, True]], + [[False, False, False]], + [[ True, False, True]] + + As example mapping matrix of this cross is as follows (5 pixels x 3 source pixels): + + [1, 0, 0] [0->0] + [1, 0, 0] [1->0] + [0, 1, 0] [2->1] + [0, 1, 0] [3->1] + [0, 0, 1] [4->2] + + For each source-pixel, we can create an of its unit-surface brightnesses by util the non-zero + entries back to masks. For example, doing this for source pixel 1 gives: + + [[0.0, 1.0, 0.0]], + [[1.0, 0.0, 0.0]] + [[0.0, 0.0, 0.0]] + + And source pixel 2: + + [[0.0, 0.0, 0.0]], + [[0.0, 1.0, 1.0]] + [[0.0, 0.0, 0.0]] + + We then convolve each of these with our PSF kernel, in 2 dimensions, like we would a grid. For + example, using the kernel below: + + kernel: + + [[0.0, 0.1, 0.0]] + [[0.1, 0.6, 0.1]] + [[0.0, 0.1, 0.0]] + + Blurred Source Pixel 1 (we don't need to perform the convolution into masked pixels): + + [[0.0, 0.6, 0.0]], + [[0.6, 0.0, 0.0]], + [[0.0, 0.0, 0.0]] + + Blurred Source pixel 2: + + [[0.0, 0.0, 0.0]], + [[0.0, 0.7, 0.7]], + [[0.0, 0.0, 0.0]] + + Finally, we map each of these blurred back to a blurred mapping matrix, which is analogous to the + mapping matrix. + + [0.6, 0.0, 0.0] [0->0] + [0.6, 0.0, 0.0] [1->0] + [0.0, 0.7, 0.0] [2->1] + [0.0, 0.7, 0.0] [3->1] + [0.0, 0.0, 0.6] [4->2] + + If the mapping matrix is sub-gridded, we perform the convolution on the fractional surface brightnesses in an + identical fashion to above. + + Parameters + ---------- + mapping_matrix + The 2D mapping matrix describing how every inversion pixel maps to a pixel on the data pixel. + """ + return self.convolve_matrix_jit( + mapping_matrix=mapping_matrix, + image_frame_1d_indexes=self.image_frame_1d_indexes, + image_frame_1d_kernels=self.image_frame_1d_kernels, + image_frame_1d_lengths=self.image_frame_1d_lengths, + ) + + @staticmethod + @numba_util.jit() + def convolve_matrix_jit( + mapping_matrix, + image_frame_1d_indexes, + image_frame_1d_kernels, + image_frame_1d_lengths, + ): + blurred_mapping_matrix = np.zeros(mapping_matrix.shape) + + for pixel_1d_index in range(mapping_matrix.shape[1]): + for image_1d_index in range(mapping_matrix.shape[0]): + value = mapping_matrix[image_1d_index, pixel_1d_index] + + if value > 0: + frame_1d_indexes = image_frame_1d_indexes[image_1d_index] + frame_1d_kernel = image_frame_1d_kernels[image_1d_index] + frame_1d_length = image_frame_1d_lengths[image_1d_index] + + for kernel_1d_index in range(frame_1d_length): + vector_index = frame_1d_indexes[kernel_1d_index] + kernel_value = frame_1d_kernel[kernel_1d_index] + blurred_mapping_matrix[vector_index, pixel_1d_index] += ( + value * kernel_value + ) + + return blurred_mapping_matrix diff --git a/autoarray/inversion/inversion/abstract.py b/autoarray/inversion/inversion/abstract.py index 757a24ef0..d2546e9e3 100644 --- a/autoarray/inversion/inversion/abstract.py +++ b/autoarray/inversion/inversion/abstract.py @@ -100,6 +100,10 @@ def data(self): def noise_map(self): return self.dataset.noise_map + @property + def convolver(self): + return self.dataset.convolver + def has(self, cls: Type) -> bool: """ Does this `Inversion` have an attribute which is of type `cls`? diff --git a/autoarray/inversion/inversion/dataset_interface.py b/autoarray/inversion/inversion/dataset_interface.py index 0e417e71a..c1cf2db9a 100644 --- a/autoarray/inversion/inversion/dataset_interface.py +++ b/autoarray/inversion/inversion/dataset_interface.py @@ -5,6 +5,7 @@ def __init__( noise_map, grids=None, psf=None, + convolver=None, transformer=None, w_tilde=None, noise_covariance_matrix=None, @@ -60,6 +61,7 @@ def __init__( self.noise_map = noise_map self.grids = grids self.psf = psf + self.convolver = convolver self.transformer = transformer self.w_tilde = w_tilde self.noise_covariance_matrix = noise_covariance_matrix diff --git a/autoarray/inversion/inversion/imaging/abstract.py b/autoarray/inversion/inversion/imaging/abstract.py index 5efc4d0a9..f17e8fa4f 100644 --- a/autoarray/inversion/inversion/imaging/abstract.py +++ b/autoarray/inversion/inversion/imaging/abstract.py @@ -95,7 +95,7 @@ def operated_mapping_matrix_list(self) -> List[np.ndarray]: return [ ( - self.psf.convolve_mapping_matrix( + self.convolver.convolve_mapping_matrix( mapping_matrix=linear_obj.mapping_matrix ) if linear_obj.operated_mapping_matrix_override is None @@ -138,7 +138,7 @@ def linear_func_operated_mapping_matrix_dict(self) -> Dict: if linear_func.operated_mapping_matrix_override is not None: operated_mapping_matrix = linear_func.operated_mapping_matrix_override else: - operated_mapping_matrix = self.psf.convolve_mapping_matrix( + operated_mapping_matrix = self.convolver.convolve_mapping_matrix( mapping_matrix=linear_func.mapping_matrix ) @@ -220,7 +220,7 @@ def mapper_operated_mapping_matrix_dict(self) -> Dict: mapper_operated_mapping_matrix_dict = {} for mapper in self.cls_list_from(cls=AbstractMapper): - operated_mapping_matrix = self.psf.convolve_mapping_matrix( + operated_mapping_matrix = self.convolver.convolve_mapping_matrix( mapping_matrix=mapper.mapping_matrix ) diff --git a/autoarray/inversion/inversion/imaging/mapping.py b/autoarray/inversion/inversion/imaging/mapping.py index 03d73ff63..053f1d01a 100644 --- a/autoarray/inversion/inversion/imaging/mapping.py +++ b/autoarray/inversion/inversion/imaging/mapping.py @@ -78,7 +78,7 @@ def _data_vector_mapper(self) -> np.ndarray: mapper = mapper_list[i] param_range = mapper_param_range_list[i] - operated_mapping_matrix = self.psf.convolve_mapping_matrix( + operated_mapping_matrix = self.convolver.convolve_mapping_matrix( mapping_matrix=mapper.mapping_matrix ) @@ -140,7 +140,7 @@ def _curvature_matrix_mapper_diag(self) -> Optional[np.ndarray]: mapper_i = mapper_list[i] mapper_param_range_i = mapper_param_range_list[i] - operated_mapping_matrix = self.psf.convolve_mapping_matrix( + operated_mapping_matrix = self.convolver.convolve_mapping_matrix( mapping_matrix=mapper_i.mapping_matrix ) diff --git a/autoarray/inversion/inversion/imaging/w_tilde.py b/autoarray/inversion/inversion/imaging/w_tilde.py index 199ba66b2..696a81136 100644 --- a/autoarray/inversion/inversion/imaging/w_tilde.py +++ b/autoarray/inversion/inversion/imaging/w_tilde.py @@ -71,10 +71,10 @@ def __init__( @profile_func def w_tilde_data(self): return inversion_imaging_util.w_tilde_data_imaging_from( - image_native=np.array(self.data.native), - noise_map_native=np.array(self.noise_map.native), - kernel_native=np.array(self.psf.kernel.native), - native_index_for_slim_index=self.data.mask.derive_indexes.native_for_slim, + image_native=np.array(self.data.native.array).astype("float"), + noise_map_native=np.array(self.noise_map.native.array).astype("float"), + kernel_native=np.array(self.psf.native.array).astype("float"), + native_index_for_slim_index=np.array(self.data.mask.derive_indexes.native_for_slim).astype("int"), ) @property @@ -525,9 +525,9 @@ def mapped_reconstructed_data_dict(self) -> Dict[LinearObj, Array2D]: values=mapped_reconstructed_image, mask=self.mask ) - mapped_reconstructed_image = self.psf.convolve_image_no_blurring( + mapped_reconstructed_image = np.array(self.psf.convolve_image_no_blurring( image=mapped_reconstructed_image, mask=self.mask - ) + ).array) else: operated_mapping_matrix = self.linear_func_operated_mapping_matrix_dict[ From 119866d4b11a6029f844594feaead8ceb59713a8 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sat, 5 Apr 2025 14:57:14 +0100 Subject: [PATCH 21/25] fix inversion plotters --- autoarray/inversion/inversion/imaging/w_tilde.py | 10 +++++----- .../inversion/pixelization/mappers/mapper_grids.py | 1 + autoarray/operators/contour.py | 5 +---- autoarray/plot/wrap/two_d/grid_scatter.py | 2 +- autoarray/structures/grids/grid_2d_util.py | 12 ++++++------ 5 files changed, 14 insertions(+), 16 deletions(-) diff --git a/autoarray/inversion/inversion/imaging/w_tilde.py b/autoarray/inversion/inversion/imaging/w_tilde.py index 696a81136..21834cdd0 100644 --- a/autoarray/inversion/inversion/imaging/w_tilde.py +++ b/autoarray/inversion/inversion/imaging/w_tilde.py @@ -521,14 +521,14 @@ def mapped_reconstructed_data_dict(self) -> Dict[LinearObj, Array2D]: reconstruction=reconstruction, ) + mapped_reconstructed_image = self.psf.convolve_image_no_blurring( + image=mapped_reconstructed_image, mask=self.mask + ).array + mapped_reconstructed_image = Array2D( - values=mapped_reconstructed_image, mask=self.mask + values=np.array(mapped_reconstructed_image), mask=self.mask ) - mapped_reconstructed_image = np.array(self.psf.convolve_image_no_blurring( - image=mapped_reconstructed_image, mask=self.mask - ).array) - else: operated_mapping_matrix = self.linear_func_operated_mapping_matrix_dict[ linear_obj diff --git a/autoarray/inversion/pixelization/mappers/mapper_grids.py b/autoarray/inversion/pixelization/mappers/mapper_grids.py index 9a12e2f95..260266c5c 100644 --- a/autoarray/inversion/pixelization/mappers/mapper_grids.py +++ b/autoarray/inversion/pixelization/mappers/mapper_grids.py @@ -72,6 +72,7 @@ def image_plane_data_grid(self): @property def mesh_pixels_per_image_pixels(self): + mesh_pixels_per_image_pixels = grid_2d_util.grid_pixels_in_mask_pixels_from( grid=np.array(self.image_plane_mesh_grid), shape_native=self.mask.shape_native, diff --git a/autoarray/operators/contour.py b/autoarray/operators/contour.py index 2de247d3c..31258aeef 100644 --- a/autoarray/operators/contour.py +++ b/autoarray/operators/contour.py @@ -101,9 +101,6 @@ def hull( hull_x = grid_convex[hull_vertices, 0] hull_y = grid_convex[hull_vertices, 1] - grid_hull = jnp.zeros((len(hull_vertices), 2)) - - grid_hull[:, 1] = hull_x - grid_hull[:, 0] = hull_y + grid_hull = jnp.stack((hull_y, hull_x), axis=-1) return grid_hull diff --git a/autoarray/plot/wrap/two_d/grid_scatter.py b/autoarray/plot/wrap/two_d/grid_scatter.py index 8399ae591..95ca1c8bd 100644 --- a/autoarray/plot/wrap/two_d/grid_scatter.py +++ b/autoarray/plot/wrap/two_d/grid_scatter.py @@ -55,7 +55,7 @@ def scatter_grid(self, grid: Union[np.ndarray, Grid2D]): config_dict["c"] = config_dict["c"][0] try: - plt.scatter(y=grid[:, 0], x=grid[:, 1], **config_dict) + plt.scatter(y=grid.array[:, 0], x=grid.array[:, 1], **config_dict) except (IndexError, TypeError): return self.scatter_grid_list(grid_list=grid) diff --git a/autoarray/structures/grids/grid_2d_util.py b/autoarray/structures/grids/grid_2d_util.py index 6c72c00ef..e976be2b1 100644 --- a/autoarray/structures/grids/grid_2d_util.py +++ b/autoarray/structures/grids/grid_2d_util.py @@ -801,7 +801,6 @@ def compute_polygon_area(points): return 0.5 * np.abs(np.dot(x, np.roll(y, 1)) - np.dot(y, np.roll(x, 1))) -@numba_util.jit() def grid_pixels_in_mask_pixels_from( grid, shape_native, pixel_scales, origin ) -> np.ndarray: @@ -832,10 +831,11 @@ def grid_pixels_in_mask_pixels_from( mesh_pixels_per_image_pixel = np.zeros(shape=shape_native) - for i in range(grid_pixel_centres.shape[0]): - y = grid_pixel_centres[i, 0] - x = grid_pixel_centres[i, 1] + # Assuming grid_pixel_centres is a 2D array where each row contains (y, x) indices. + y_indices = grid_pixel_centres[:, 0] + x_indices = grid_pixel_centres[:, 1] - mesh_pixels_per_image_pixel[y, x] += 1 + # Use np.add.at to increment the specific indices in a safe and efficient manner + np.add.at(mesh_pixels_per_image_pixel, (y_indices, x_indices), 1) - return mesh_pixels_per_image_pixel + return mesh_pixels_per_image_pixel \ No newline at end of file From 580be1cff3f056befea7c70330dafd31152400e0 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sat, 5 Apr 2025 15:01:38 +0100 Subject: [PATCH 22/25] fix test overlay --- test_autoarray/inversion/inversion/imaging/test_imaging.py | 6 +++--- .../inversion/pixelization/image_mesh/test_overlay.py | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/test_autoarray/inversion/inversion/imaging/test_imaging.py b/test_autoarray/inversion/inversion/imaging/test_imaging.py index df615ad85..cbb8f2ecc 100644 --- a/test_autoarray/inversion/inversion/imaging/test_imaging.py +++ b/test_autoarray/inversion/inversion/imaging/test_imaging.py @@ -11,9 +11,9 @@ directory = path.dirname(path.realpath(__file__)) -def test__operated_mapping_matrix_property(psf_7x7, rectangular_mapper_7x7_3x3): +def test__operated_mapping_matrix_property(psf_3x3, rectangular_mapper_7x7_3x3): inversion = aa.m.MockInversionImaging( - psf=psf_7x7, linear_obj_list=[rectangular_mapper_7x7_3x3] + psf=psf_3x3, linear_obj_list=[rectangular_mapper_7x7_3x3] ) assert inversion.operated_mapping_matrix_list[0][0, 0] == pytest.approx(1.0, 1e-4) @@ -42,7 +42,7 @@ def test__operated_mapping_matrix_property(psf_7x7, rectangular_mapper_7x7_3x3): def test__operated_mapping_matrix_property__with_operated_mapping_matrix_override( - psf_7x7, rectangular_mapper_7x7_3x3 + psf_3x3, rectangular_mapper_7x7_3x3 ): psf = aa.m.MockPSF(operated_mapping_matrix=np.ones((2, 2))) diff --git a/test_autoarray/inversion/pixelization/image_mesh/test_overlay.py b/test_autoarray/inversion/pixelization/image_mesh/test_overlay.py index 22536d5c8..54f242236 100644 --- a/test_autoarray/inversion/pixelization/image_mesh/test_overlay.py +++ b/test_autoarray/inversion/pixelization/image_mesh/test_overlay.py @@ -323,13 +323,13 @@ def test__image_plane_mesh_grid_from__simple(): total_pixels = overlay_util.total_pixels_2d_from( mask_2d=mask.array, - overlaid_centres=overlaid_centres, + overlaid_centres=np.array(overlaid_centres), ) overlay_for_mask_2d_util = overlay_util.overlay_for_mask_from( total_pixels=total_pixels, mask=mask.array, - overlaid_centres=overlaid_centres, + overlaid_centres=np.array(overlaid_centres), ).astype("int") image_mesh_util = overlay_util.overlay_via_unmasked_overlaid_from( From 994e04843aff1dd164300eab5f543828f4e6bb45 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sat, 5 Apr 2025 15:05:23 +0100 Subject: [PATCH 23/25] fixn interferometer conversion --- autoarray/dataset/interferometer/dataset.py | 2 +- autoarray/inversion/inversion/interferometer/w_tilde.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/autoarray/dataset/interferometer/dataset.py b/autoarray/dataset/interferometer/dataset.py index 06892ea68..32a9a4377 100644 --- a/autoarray/dataset/interferometer/dataset.py +++ b/autoarray/dataset/interferometer/dataset.py @@ -193,7 +193,7 @@ def w_tilde(self): w_matrix = inversion_interferometer_util.w_tilde_via_preload_from( w_tilde_preload=curvature_preload, - native_index_for_slim_index=self.real_space_mask.derive_indexes.native_for_slim, + native_index_for_slim_index=np.array(self.real_space_mask.derive_indexes.native_for_slim).astype("int"), ) dirty_image = self.transformer.image_from( diff --git a/autoarray/inversion/inversion/interferometer/w_tilde.py b/autoarray/inversion/inversion/interferometer/w_tilde.py index 18d7e1cec..2e785dfd8 100644 --- a/autoarray/inversion/inversion/interferometer/w_tilde.py +++ b/autoarray/inversion/inversion/interferometer/w_tilde.py @@ -131,7 +131,7 @@ def curvature_matrix_diag(self) -> np.ndarray: pix_indexes_for_sub_slim_index=mapper.pix_indexes_for_sub_slim_index, pix_size_for_sub_slim_index=mapper.pix_sizes_for_sub_slim_index, pix_weights_for_sub_slim_index=mapper.pix_weights_for_sub_slim_index, - native_index_for_slim_index=self.transformer.real_space_mask.derive_indexes.native_for_slim, + native_index_for_slim_index=np.array(self.transformer.real_space_mask.derive_indexes.native_for_slim).astype("int"), pix_pixels=self.linear_obj_list[0].params, ) @@ -143,7 +143,7 @@ def curvature_matrix_diag(self) -> np.ndarray: return inversion_interferometer_util.curvature_matrix_via_w_tilde_curvature_preload_interferometer_from_2( curvature_preload=self.w_tilde.curvature_preload, - native_index_for_slim_index=self.transformer.real_space_mask.derive_indexes.native_for_slim, + native_index_for_slim_index=np.array(self.transformer.real_space_mask.derive_indexes.native_for_slim).astype("int"), pix_pixels=self.linear_obj_list[0].params, sub_slim_indexes_for_pix_index=sub_slim_indexes_for_pix_index.astype("int"), sub_slim_sizes_for_pix_index=sub_slim_sizes_for_pix_index.astype("int"), From 9c5354ca54d7536e9c58c5724fe70cd24346f710 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sat, 5 Apr 2025 15:35:44 +0100 Subject: [PATCH 24/25] more plot unit tests pass --- .../inversion/mock/mock_inversion_imaging.py | 2 ++ autoarray/mock.py | 2 ++ autoarray/operators/mock/mock_convolver.py | 6 ---- autoarray/operators/mock/mock_psf.py | 14 +++++++++ autoarray/plot/wrap/base/colorbar.py | 4 +++ autoarray/plot/wrap/two_d/grid_scatter.py | 6 +++- autoarray/structures/arrays/kernel_2d.py | 10 +++---- .../inversion/imaging/test_imaging.py | 9 ++++-- .../imaging/test_inversion_imaging_util.py | 30 +++++++++---------- .../test_inversion_interferometer_util.py | 10 +++++-- .../plot/test_structure_plotters.py | 1 - 11 files changed, 61 insertions(+), 33 deletions(-) delete mode 100644 autoarray/operators/mock/mock_convolver.py create mode 100644 autoarray/operators/mock/mock_psf.py diff --git a/autoarray/inversion/mock/mock_inversion_imaging.py b/autoarray/inversion/mock/mock_inversion_imaging.py index de7f2baa1..601c669b9 100644 --- a/autoarray/inversion/mock/mock_inversion_imaging.py +++ b/autoarray/inversion/mock/mock_inversion_imaging.py @@ -13,6 +13,7 @@ def __init__( data=None, noise_map=None, psf=None, + convolver=None, linear_obj_list=None, operated_mapping_matrix=None, linear_func_operated_mapping_matrix_dict=None, @@ -23,6 +24,7 @@ def __init__( data=data, noise_map=noise_map, psf=psf, + convolver=convolver, ) super().__init__( diff --git a/autoarray/mock.py b/autoarray/mock.py index 2261ba5e2..b3bc02c15 100644 --- a/autoarray/mock.py +++ b/autoarray/mock.py @@ -15,6 +15,8 @@ from autoarray.fit.mock.mock_fit_imaging import MockFitImaging from autoarray.fit.mock.mock_fit_interferometer import MockFitInterferometer from autoarray.mask.mock.mock_mask import MockMask +from autoarray.operators.mock.mock_psf import MockConvolver +from autoarray.operators.mock.mock_psf import MockPSF from autoarray.structures.mock.mock_grid import MockGrid2DMesh from autoarray.structures.mock.mock_grid import MockMeshGrid from autoarray.structures.mock.mock_decorators import MockGridRadialMinimum diff --git a/autoarray/operators/mock/mock_convolver.py b/autoarray/operators/mock/mock_convolver.py deleted file mode 100644 index 7dc5dfcbb..000000000 --- a/autoarray/operators/mock/mock_convolver.py +++ /dev/null @@ -1,6 +0,0 @@ -class MockPSF: - def __init__(self, operated_mapping_matrix=None): - self.operated_mapping_matrix = operated_mapping_matrix - - def convolve_mapping_matrix(self, mapping_matrix): - return self.operated_mapping_matrix diff --git a/autoarray/operators/mock/mock_psf.py b/autoarray/operators/mock/mock_psf.py new file mode 100644 index 000000000..08864d236 --- /dev/null +++ b/autoarray/operators/mock/mock_psf.py @@ -0,0 +1,14 @@ +class MockPSF: + def __init__(self, operated_mapping_matrix=None): + self.operated_mapping_matrix = operated_mapping_matrix + + def convolve_mapping_matrix(self, mapping_matrix): + return self.operated_mapping_matrix + + +class MockConvolver: + def __init__(self, operated_mapping_matrix=None): + self.operated_mapping_matrix = operated_mapping_matrix + + def convolve_mapping_matrix(self, mapping_matrix): + return self.operated_mapping_matrix diff --git a/autoarray/plot/wrap/base/colorbar.py b/autoarray/plot/wrap/base/colorbar.py index b0650013a..b3577e9c2 100644 --- a/autoarray/plot/wrap/base/colorbar.py +++ b/autoarray/plot/wrap/base/colorbar.py @@ -194,6 +194,10 @@ def set_with_color_values( ) if tick_values is None and tick_labels is None: + + print(mappable) + print(ax) + cb = plt.colorbar( mappable=mappable, ax=ax, diff --git a/autoarray/plot/wrap/two_d/grid_scatter.py b/autoarray/plot/wrap/two_d/grid_scatter.py index 95ca1c8bd..802d20b3c 100644 --- a/autoarray/plot/wrap/two_d/grid_scatter.py +++ b/autoarray/plot/wrap/two_d/grid_scatter.py @@ -1,4 +1,5 @@ import matplotlib.pyplot as plt +import jax.numpy as jnp import numpy as np import itertools from scipy.spatial import ConvexHull @@ -54,8 +55,11 @@ def scatter_grid(self, grid: Union[np.ndarray, Grid2D]): if len(config_dict["c"]) > 1: config_dict["c"] = config_dict["c"][0] + if isinstance(grid, jnp.ndarray): + grid = np.array(grid.array) + try: - plt.scatter(y=grid.array[:, 0], x=grid.array[:, 1], **config_dict) + plt.scatter(y=grid[:, 0], x=grid[:, 1], **config_dict) except (IndexError, TypeError): return self.scatter_grid_list(grid_list=grid) diff --git a/autoarray/structures/arrays/kernel_2d.py b/autoarray/structures/arrays/kernel_2d.py index bb62656a9..6efaf9bac 100644 --- a/autoarray/structures/arrays/kernel_2d.py +++ b/autoarray/structures/arrays/kernel_2d.py @@ -485,7 +485,7 @@ def convolved_array_from(self, array: Array2D) -> Array2D: return Array2D(values=convolved_array_1d, mask=array_2d.mask) - def convolve_image(self, image, blurring_image, jax_method="fft"): + def convolve_image(self, image, blurring_image, jax_method="direct"): """ For a given 1D array and blurring array, convolve the two using this psf. @@ -528,7 +528,7 @@ def convolve_image(self, image, blurring_image, jax_method="fft"): return Array2D(values=convolved_array_1d, mask=image.mask) - def convolve_image_no_blurring(self, image, mask, jax_method="fft"): + def convolve_image_no_blurring(self, image, mask, jax_method="direct"): """ For a given 1D array and blurring array, convolve the two using this psf. @@ -561,7 +561,7 @@ def convolve_image_no_blurring(self, image, mask, jax_method="fft"): return Array2D(values=convolved_array_1d, mask=mask) - def convolve_mapping_matrix(self, mapping_matrix, mask): + def convolve_mapping_matrix(self, mapping_matrix, mask, jax_method="direct"): """For a given 1D array and blurring array, convolve the two using this psf. Parameters @@ -569,6 +569,6 @@ def convolve_mapping_matrix(self, mapping_matrix, mask): image 1D array of the values which are to be blurred with the psf's PSF. """ - return jax.vmap(self.convolve_image_no_blurring, in_axes=(1, None))( - mapping_matrix, mask + return jax.vmap(self.convolve_image_no_blurring, in_axes=(1, None, None))( + mapping_matrix, mask, jax_method ).T diff --git a/test_autoarray/inversion/inversion/imaging/test_imaging.py b/test_autoarray/inversion/inversion/imaging/test_imaging.py index cbb8f2ecc..8a52e4c35 100644 --- a/test_autoarray/inversion/inversion/imaging/test_imaging.py +++ b/test_autoarray/inversion/inversion/imaging/test_imaging.py @@ -13,7 +13,9 @@ def test__operated_mapping_matrix_property(psf_3x3, rectangular_mapper_7x7_3x3): inversion = aa.m.MockInversionImaging( - psf=psf_3x3, linear_obj_list=[rectangular_mapper_7x7_3x3] + psf=psf_3x3, + linear_obj_list=[rectangular_mapper_7x7_3x3], + convolver=aa.Convolver(kernel=psf_3x3, mask=rectangular_mapper_7x7_3x3.mapper_grids.mask) ) assert inversion.operated_mapping_matrix_list[0][0, 0] == pytest.approx(1.0, 1e-4) @@ -24,6 +26,7 @@ def test__operated_mapping_matrix_property(psf_3x3, rectangular_mapper_7x7_3x3): inversion = aa.m.MockInversionImaging( psf=psf, linear_obj_list=[rectangular_mapper_7x7_3x3, rectangular_mapper_7x7_3x3], + convolver=aa.m.MockConvolver(operated_mapping_matrix=np.ones((2, 2))) ) operated_mapping_matrix_0 = np.array([[1.0, 1.0], [1.0, 1.0]]) @@ -54,7 +57,8 @@ def test__operated_mapping_matrix_property__with_operated_mapping_matrix_overrid ) inversion = aa.m.MockInversionImaging( - psf=psf, linear_obj_list=[rectangular_mapper_7x7_3x3, linear_obj] + psf=psf, linear_obj_list=[rectangular_mapper_7x7_3x3, linear_obj], + convolver=aa.m.MockConvolver(operated_mapping_matrix=np.ones((2, 2))) ) operated_mapping_matrix_0 = np.array([[1.0, 1.0], [1.0, 1.0]]) @@ -88,6 +92,7 @@ def test__curvature_matrix(rectangular_mapper_7x7_3x3): data=np.ones(2), noise_map=noise_map, psf=psf, + convolver=aa.m.MockConvolver(operated_mapping_matrix=np.ones((2, 10))) ) inversion = aa.InversionImagingMapping( diff --git a/test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py b/test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py index e05e13350..fab5b6600 100644 --- a/test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py +++ b/test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py @@ -204,22 +204,22 @@ def test__data_vector_via_w_tilde_data_two_methods_agree(): mapping_matrix = mapper.mapping_matrix blurred_mapping_matrix = psf.convolve_mapping_matrix( - mapping_matrix=mapping_matrix + mapping_matrix=mapping_matrix, mask=mask ) data_vector = ( aa.util.inversion_imaging.data_vector_via_blurred_mapping_matrix_from( - blurred_mapping_matrix=blurred_mapping_matrix, + blurred_mapping_matrix=np.array(blurred_mapping_matrix), image=np.array(image), noise_map=np.array(noise_map), ) ) w_tilde_data = aa.util.inversion_imaging.w_tilde_data_imaging_from( - image_native=np.array(image.native), - noise_map_native=np.array(noise_map.native), - kernel_native=np.array(kernel.native), - native_index_for_slim_index=mask.derive_indexes.native_for_slim, + image_native=np.array(image.native.array), + noise_map_native=np.array(noise_map.native.array), + kernel_native=np.array(kernel.native.array), + native_index_for_slim_index=np.array(mask.derive_indexes.native_for_slim).astype("int"), ) ( @@ -273,16 +273,16 @@ def test__curvature_matrix_via_w_tilde_two_methods_agree(): mapping_matrix = mapper.mapping_matrix w_tilde = aa.util.inversion_imaging.w_tilde_curvature_imaging_from( - noise_map_native=np.array(noise_map.native), - kernel_native=np.array(kernel.native), - native_index_for_slim_index=mask.derive_indexes.native_for_slim, + noise_map_native=np.array(noise_map.native.array), + kernel_native=np.array(kernel.native.array), + native_index_for_slim_index=np.array(mask.derive_indexes.native_for_slim).astype("int"), ) curvature_matrix_via_w_tilde = aa.util.inversion.curvature_matrix_via_w_tilde_from( w_tilde=w_tilde, mapping_matrix=mapping_matrix ) - blurred_mapping_matrix = psf.convolve_mapping_matrix(mapping_matrix=mapping_matrix) + blurred_mapping_matrix = psf.convolve_mapping_matrix(mapping_matrix=mapping_matrix, mask=mask) curvature_matrix = aa.util.inversion.curvature_matrix_via_mapping_matrix_from( mapping_matrix=blurred_mapping_matrix, @@ -326,9 +326,9 @@ def test__curvature_matrix_via_w_tilde_preload_two_methods_agree(): w_tilde_indexes, w_tilde_lengths, ) = aa.util.inversion_imaging.w_tilde_curvature_preload_imaging_from( - noise_map_native=np.array(noise_map.native), - kernel_native=np.array(kernel.native), - native_index_for_slim_index=mask.derive_indexes.native_for_slim, + noise_map_native=np.array(noise_map.native.array), + kernel_native=np.array(kernel.native.array), + native_index_for_slim_index=np.array(mask.derive_indexes.native_for_slim).astype("int"), ) ( @@ -355,11 +355,11 @@ def test__curvature_matrix_via_w_tilde_preload_two_methods_agree(): ) blurred_mapping_matrix = psf.convolve_mapping_matrix( - mapping_matrix=mapping_matrix + mapping_matrix=mapping_matrix, mask=mask, ) curvature_matrix = aa.util.inversion.curvature_matrix_via_mapping_matrix_from( - mapping_matrix=blurred_mapping_matrix, + mapping_matrix=np.array(blurred_mapping_matrix), noise_map=np.array(noise_map), ) diff --git a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py index cb7d27673..28c45a9bd 100644 --- a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py +++ b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py @@ -272,20 +272,24 @@ def test__identical_inversion_values_for_two_methods(): == inversion_mapping_matrices.regularization_matrix ).all() + assert inversion_w_tilde.data_vector == pytest.approx( + inversion_mapping_matrices.data_vector, 1.0e-8 + ) assert inversion_w_tilde.curvature_matrix == pytest.approx( inversion_mapping_matrices.curvature_matrix, 1.0e-8 ) assert inversion_w_tilde.curvature_reg_matrix == pytest.approx( inversion_mapping_matrices.curvature_reg_matrix, 1.0e-8 ) + assert inversion_w_tilde.reconstruction == pytest.approx( - inversion_mapping_matrices.reconstruction, 1.0e-2 + inversion_mapping_matrices.reconstruction, abs=1.0e-1 ) assert inversion_w_tilde.mapped_reconstructed_image == pytest.approx( - inversion_mapping_matrices.mapped_reconstructed_image, 1.0e-2 + inversion_mapping_matrices.mapped_reconstructed_image, abs=1.0e-1 ) assert inversion_w_tilde.mapped_reconstructed_data == pytest.approx( - inversion_mapping_matrices.mapped_reconstructed_data, 1.0e-2 + inversion_mapping_matrices.mapped_reconstructed_data, abs=1.0e-1 ) diff --git a/test_autoarray/structures/plot/test_structure_plotters.py b/test_autoarray/structures/plot/test_structure_plotters.py index d455c86f4..ad1ca0251 100644 --- a/test_autoarray/structures/plot/test_structure_plotters.py +++ b/test_autoarray/structures/plot/test_structure_plotters.py @@ -3,7 +3,6 @@ from os import path import pytest import numpy as np -import jax.numpy as jnp import shutil directory = path.dirname(path.realpath(__file__)) From 2b6754c459f9cb819483daa7a34c9e298672efb2 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sat, 5 Apr 2025 15:41:52 +0100 Subject: [PATCH 25/25] tst coverage complete --- autoarray/plot/wrap/two_d/grid_scatter.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/autoarray/plot/wrap/two_d/grid_scatter.py b/autoarray/plot/wrap/two_d/grid_scatter.py index 802d20b3c..ea1e3f6a8 100644 --- a/autoarray/plot/wrap/two_d/grid_scatter.py +++ b/autoarray/plot/wrap/two_d/grid_scatter.py @@ -9,6 +9,7 @@ from autoarray.plot.wrap.two_d.abstract import AbstractMatWrap2D from autoarray.structures.grids.uniform_2d import Grid2D from autoarray.structures.grids.irregular_2d import Grid2DIrregular +from autoarray.structures.mesh.abstract_2d import Abstract2DMesh class GridScatter(AbstractMatWrap2D): @@ -55,7 +56,7 @@ def scatter_grid(self, grid: Union[np.ndarray, Grid2D]): if len(config_dict["c"]) > 1: config_dict["c"] = config_dict["c"][0] - if isinstance(grid, jnp.ndarray): + if isinstance(grid, jnp.ndarray) or isinstance(grid, Abstract2DMesh): grid = np.array(grid.array) try: