From 33a3ccb69b17af77457ca0d000539e942852a25f Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Thu, 3 Apr 2025 14:50:01 +0100 Subject: [PATCH 01/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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: From 204904281e53dc1e8d87d7abcc9e5689f6c03342 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sat, 5 Apr 2025 16:55:50 +0100 Subject: [PATCH 26/62] small changes to jax array handling --- autoarray/geometry/geometry_util.py | 27 ++++++++++++++--------- autoarray/structures/arrays/uniform_1d.py | 6 +++++ autoarray/structures/grids/uniform_2d.py | 4 +++- 3 files changed, 25 insertions(+), 12 deletions(-) diff --git a/autoarray/geometry/geometry_util.py b/autoarray/geometry/geometry_util.py index e2c6c8898..9ca5fc48b 100644 --- a/autoarray/geometry/geometry_util.py +++ b/autoarray/geometry/geometry_util.py @@ -445,7 +445,12 @@ def transform_grid_2d_to_reference_frame( grid The 2d grid of (y, x) coordinates which are transformed to a new reference frame. """ - shifted_grid_2d = np.array(grid_2d) - jnp.array(centre) + try: + grid_2d = grid_2d.array + except AttributeError: + pass + + shifted_grid_2d = grid_2d - jnp.array(centre) radius = jnp.sqrt(jnp.sum(shifted_grid_2d**2.0, axis=1)) theta_coordinate_to_profile = jnp.arctan2( @@ -477,23 +482,23 @@ def transform_grid_2d_from_reference_frame( The 2d grid of (y, x) coordinates which are transformed to a new reference frame. """ - cos_angle = np.cos(np.radians(angle)) - sin_angle = np.sin(np.radians(angle)) + cos_angle = jnp.cos(jnp.radians(angle)) + sin_angle = jnp.sin(jnp.radians(angle)) - y = np.add( - np.add( - np.multiply(grid_2d[:, 1], sin_angle), np.multiply(grid_2d[:, 0], cos_angle) + y = jnp.add( + jnp.add( + jnp.multiply(grid_2d[:, 1], sin_angle), jnp.multiply(grid_2d[:, 0], cos_angle) ), centre[0], ) - x = np.add( - np.add( - np.multiply(grid_2d[:, 1], cos_angle), - -np.multiply(grid_2d[:, 0], sin_angle), + x = jnp.add( + jnp.add( + jnp.multiply(grid_2d[:, 1], cos_angle), + -jnp.multiply(grid_2d[:, 0], sin_angle), ), centre[1], ) - return np.vstack((y, x)).T + return jnp.vstack((y, x)).T def grid_pixels_2d_slim_from( diff --git a/autoarray/structures/arrays/uniform_1d.py b/autoarray/structures/arrays/uniform_1d.py index 4dbf1692f..234de540c 100644 --- a/autoarray/structures/arrays/uniform_1d.py +++ b/autoarray/structures/arrays/uniform_1d.py @@ -24,6 +24,12 @@ def __init__( header: Optional[Header] = None, store_native: bool = False, ): + + try: + values = values._array + except AttributeError: + values = values + values = array_1d_util.convert_array_1d( array_1d=values, mask_1d=mask, diff --git a/autoarray/structures/grids/uniform_2d.py b/autoarray/structures/grids/uniform_2d.py index 670cbfcbe..9a837d5b1 100644 --- a/autoarray/structures/grids/uniform_2d.py +++ b/autoarray/structures/grids/uniform_2d.py @@ -180,7 +180,7 @@ def __init__( self.over_sampler = OverSampler(sub_size=over_sample_size, mask=mask) if over_sampled is None: - self.over_sampled = ( + over_sampled = ( over_sample_util.grid_2d_slim_over_sampled_via_mask_from( mask_2d=np.array(self.mask), pixel_scales=self.mask.pixel_scales, @@ -189,6 +189,8 @@ def __init__( ) ) + self.over_sampled = Grid2DIrregular(values=over_sampled) + else: self.over_sampled = over_sampled From 0da7719fe41ed1f599efb69a129d3a3e5fee977d Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sat, 5 Apr 2025 20:43:14 +0100 Subject: [PATCH 27/62] clean up grid contou --- autoarray/operators/contour.py | 8 +++++++- autoarray/structures/arrays/kernel_2d.py | 5 ++++- autoarray/structures/grids/uniform_2d.py | 1 + autoarray/structures/vectors/uniform.py | 5 +++++ 4 files changed, 17 insertions(+), 2 deletions(-) diff --git a/autoarray/operators/contour.py b/autoarray/operators/contour.py index 31258aeef..b9e0a244c 100644 --- a/autoarray/operators/contour.py +++ b/autoarray/operators/contour.py @@ -55,8 +55,14 @@ def contour_array(self): @property def contour_list(self): # make sure to use base numpy to convert JAX array back to a normal array + + if isinstance(self.contour_array, jnp.ndarray): + contour_array = np.array(self.contour_array) + else: + contour_array = np.array(self.contour_array.array) + contour_indices_list = measure.find_contours( - np.array(self.contour_array), 0 + contour_array, 0 ) if len(contour_indices_list) == 0: diff --git a/autoarray/structures/arrays/kernel_2d.py b/autoarray/structures/arrays/kernel_2d.py index 6efaf9bac..4c31a471f 100644 --- a/autoarray/structures/arrays/kernel_2d.py +++ b/autoarray/structures/arrays/kernel_2d.py @@ -549,7 +549,10 @@ def convolve_image_no_blurring(self, image, mask, jax_method="direct"): expanded_array_native = jnp.zeros(mask.shape) - expanded_array_native = expanded_array_native.at[slim_to_native].set(image) + if isinstance(image, np.ndarray) or isinstance(image, jnp.ndarray): + expanded_array_native = expanded_array_native.at[slim_to_native].set(image) + else: + expanded_array_native = expanded_array_native.at[slim_to_native].set(image.array) kernel = np.array(self.native.array) diff --git a/autoarray/structures/grids/uniform_2d.py b/autoarray/structures/grids/uniform_2d.py index 9a837d5b1..879aba63c 100644 --- a/autoarray/structures/grids/uniform_2d.py +++ b/autoarray/structures/grids/uniform_2d.py @@ -159,6 +159,7 @@ def __init__( not uniform (e.g. due to gravitational lensing) is cannot be computed internally in this function. If the over sampled grid is not passed in it is computed assuming uniformity. """ + values = grid_2d_util.convert_grid_2d( grid_2d=np.array(values), mask_2d=mask, diff --git a/autoarray/structures/vectors/uniform.py b/autoarray/structures/vectors/uniform.py index 6213ecfbf..62dd7592d 100644 --- a/autoarray/structures/vectors/uniform.py +++ b/autoarray/structures/vectors/uniform.py @@ -140,6 +140,11 @@ def __init__( if type(values) is list: values = np.asarray(values) + try: + values = values._array + except AttributeError: + values = values + values = grid_2d_util.convert_grid_2d( grid_2d=values, mask_2d=mask, store_native=store_native ) From 5e31ce02a832744f11c7970ef47deb59268e142e Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 13:17:56 +0100 Subject: [PATCH 28/62] arrays now store maks in jax --- autoarray/operators/contour.py | 2 ++ autoarray/structures/arrays/array_2d_util.py | 11 ++++++----- autoarray/structures/arrays/kernel_2d.py | 2 +- autoarray/structures/arrays/uniform_2d.py | 2 +- autoarray/structures/grids/uniform_2d.py | 4 ++-- .../arrays/files/array/output_test/array.fits | Bin 5760 -> 5760 bytes test_autoarray/structures/arrays/test_repr.py | 2 +- 7 files changed, 13 insertions(+), 10 deletions(-) diff --git a/autoarray/operators/contour.py b/autoarray/operators/contour.py index b9e0a244c..804c86a49 100644 --- a/autoarray/operators/contour.py +++ b/autoarray/operators/contour.py @@ -65,6 +65,8 @@ def contour_list(self): contour_array, 0 ) + print(contour_indices_list) + if len(contour_indices_list) == 0: return [] diff --git a/autoarray/structures/arrays/array_2d_util.py b/autoarray/structures/arrays/array_2d_util.py index d8a0f7e10..63ed6ca06 100644 --- a/autoarray/structures/arrays/array_2d_util.py +++ b/autoarray/structures/arrays/array_2d_util.py @@ -122,21 +122,23 @@ def convert_array_2d( check_array_2d_and_mask_2d(array_2d=array_2d, mask_2d=mask_2d) + mask_2d = jnp.array(mask_2d.array) + is_native = len(array_2d.shape) == 2 if is_native and not skip_mask: - array_2d *= np.invert(mask_2d) + array_2d *= jnp.invert(mask_2d) if is_native == store_native: return array_2d elif not store_native: return array_2d_slim_from( - array_2d_native=np.array(array_2d), - mask_2d=np.array(mask_2d), + array_2d_native=array_2d, + mask_2d=mask_2d, ) array_2d = array_2d_native_from( array_2d_slim=array_2d, - mask_2d=np.array(mask_2d), + mask_2d=mask_2d, ) return array_2d @@ -587,7 +589,6 @@ def array_2d_native_from( ) -@partial(jax.jit, static_argnums=(1,)) def array_2d_via_indexes_from( array_2d_slim: np.ndarray, shape: Tuple[int, int], diff --git a/autoarray/structures/arrays/kernel_2d.py b/autoarray/structures/arrays/kernel_2d.py index 4c31a471f..b436a4077 100644 --- a/autoarray/structures/arrays/kernel_2d.py +++ b/autoarray/structures/arrays/kernel_2d.py @@ -255,7 +255,7 @@ def from_gaussian( """ grid = Grid2D.uniform(shape_native=shape_native, pixel_scales=pixel_scales) - grid_shifted = np.subtract(grid, centre) + grid_shifted = np.subtract(grid.array, centre) grid_radius = np.sqrt(np.sum(grid_shifted**2.0, 1)) theta_coordinate_to_profile = np.arctan2( grid_shifted[:, 0], grid_shifted[:, 1] diff --git a/autoarray/structures/arrays/uniform_2d.py b/autoarray/structures/arrays/uniform_2d.py index 11c478ad5..c527104e0 100644 --- a/autoarray/structures/arrays/uniform_2d.py +++ b/autoarray/structures/arrays/uniform_2d.py @@ -962,7 +962,7 @@ def from_yx_and_values( ) grid_pixels = geometry_util.grid_pixel_indexes_2d_slim_from( - grid_scaled_2d_slim=np.array(grid.slim), + grid_scaled_2d_slim=grid.slim, shape_native=shape_native, pixel_scales=pixel_scales, ) diff --git a/autoarray/structures/grids/uniform_2d.py b/autoarray/structures/grids/uniform_2d.py index 879aba63c..3e1152e4e 100644 --- a/autoarray/structures/grids/uniform_2d.py +++ b/autoarray/structures/grids/uniform_2d.py @@ -161,7 +161,7 @@ def __init__( """ values = grid_2d_util.convert_grid_2d( - grid_2d=np.array(values), + grid_2d=values, mask_2d=mask, store_native=store_native, ) @@ -320,7 +320,7 @@ def from_yx_1d( x = np.asarray(x) return cls.no_mask( - values=np.stack((y, x), axis=-1), + values=jnp.stack((y, x), axis=-1), shape_native=shape_native, pixel_scales=pixel_scales, origin=origin, diff --git a/test_autoarray/structures/arrays/files/array/output_test/array.fits b/test_autoarray/structures/arrays/files/array/output_test/array.fits index 0cff0a8f7db90d3ded28c644cd66e3791627feeb..dab9da2fb676efc023e53b73faa54fe78fc4000c 100644 GIT binary patch delta 56 kcmZqBZP1;N!(?o$Vgmz%Jzl(dBKLc)$qTsk0j$6e2mk;8 delta 84 hcmZqBZP1;N!(?W%G4C>$;|B&XuqT_|T*&>83jh#~5uX46 diff --git a/test_autoarray/structures/arrays/test_repr.py b/test_autoarray/structures/arrays/test_repr.py index 8ea8c6b9d..e70c98056 100644 --- a/test_autoarray/structures/arrays/test_repr.py +++ b/test_autoarray/structures/arrays/test_repr.py @@ -3,4 +3,4 @@ def test_repr(): array = aa.Array2D.no_mask([[1, 2], [3, 4]], pixel_scales=1) - assert repr(array) == "Array2D([1, 2, 3, 4])" + assert repr(array) == "Array([1, 2, 3, 4], dtype=int32)" From 3357cbb539e09abcac4c09bf70f81d7540f1e743 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 13:27:50 +0100 Subject: [PATCH 29/62] array 2d always stored as jax --- autoarray/structures/arrays/array_2d_util.py | 9 +++++++-- autoarray/structures/arrays/uniform_2d.py | 5 ----- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/autoarray/structures/arrays/array_2d_util.py b/autoarray/structures/arrays/array_2d_util.py index 63ed6ca06..a286924c5 100644 --- a/autoarray/structures/arrays/array_2d_util.py +++ b/autoarray/structures/arrays/array_2d_util.py @@ -11,7 +11,6 @@ from autoarray.mask import mask_2d_util from autoarray import exc -from functools import partial def convert_array(array: Union[np.ndarray, List]) -> np.ndarray: @@ -23,8 +22,14 @@ def convert_array(array: Union[np.ndarray, List]) -> np.ndarray: array : list or ndarray The array which may be converted to an ndarray """ + + try: + array = array.array + except AttributeError: + pass + if isinstance(array, np.ndarray) or isinstance(array, list): - array = np.asarray(array) + array = jnp.asarray(array) elif isinstance(array, jnp.ndarray): array = jax.lax.cond( type(array) is list, lambda _: jnp.asarray(array), lambda _: array, None diff --git a/autoarray/structures/arrays/uniform_2d.py b/autoarray/structures/arrays/uniform_2d.py index c527104e0..1ab5db723 100644 --- a/autoarray/structures/arrays/uniform_2d.py +++ b/autoarray/structures/arrays/uniform_2d.py @@ -232,11 +232,6 @@ def __init__( if conf.instance["general"]["structures"]["native_binned_only"]: store_native = True - try: - values = values._array - except AttributeError: - values = values - values = array_2d_util.convert_array_2d( array_2d=values, mask_2d=mask, From b9e4cb5bd53db54888fd2d82e01c63b84c0e7bc2 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 13:32:14 +0100 Subject: [PATCH 30/62] update array_1d_slim_from --- autoarray/structures/arrays/array_1d_util.py | 26 ++++++-------------- autoarray/structures/arrays/array_2d_util.py | 9 ++----- autoarray/structures/arrays/uniform_1d.py | 5 ---- 3 files changed, 10 insertions(+), 30 deletions(-) diff --git a/autoarray/structures/arrays/array_1d_util.py b/autoarray/structures/arrays/array_1d_util.py index 5e6565326..7e126b0a5 100644 --- a/autoarray/structures/arrays/array_1d_util.py +++ b/autoarray/structures/arrays/array_1d_util.py @@ -1,4 +1,5 @@ from __future__ import annotations +import jax.numpy as jnp import numpy as np from typing import TYPE_CHECKING, List, Union @@ -38,26 +39,26 @@ def convert_array_1d( If True, the ndarray is stored in its native format [total_y_pixels, total_x_pixels]. This avoids mapping large data arrays to and from the slim / native formats, which can be a computational bottleneck. """ - array_1d = array_2d_util.convert_array(array=array_1d) is_native = array_1d.shape[0] == mask_1d.shape_native[0] + mask_1d = jnp.array(mask_1d.array) + if is_native == store_native: return array_1d elif not store_native: return array_1d_slim_from( - array_1d_native=np.array(array_1d), - mask_1d=np.array(mask_1d), + array_1d_native=array_1d, + mask_1d=mask_1d, ) return array_1d_native_from( array_1d_slim=array_1d, - mask_1d=np.array(mask_1d), + mask_1d=mask_1d, ) -@numba_util.jit() def array_1d_slim_from( array_1d_native: np.ndarray, mask_1d: np.ndarray, @@ -105,19 +106,8 @@ def array_1d_slim_from( array_1d_slim = array_1d_slim_from(array_1d_native, array_2d=array_2d) """ - - total_pixels = mask_1d_util.total_pixels_1d_from( - mask_1d=mask_1d, - ) - - line_1d_slim = np.zeros(shape=total_pixels) - index = 0 - - for x in range(mask_1d.shape[0]): - if not mask_1d[x]: - line_1d_slim[index] = array_1d_native[x] - index += 1 - + unmasked_indices = ~mask_1d + line_1d_slim = array_1d_native[unmasked_indices] return line_1d_slim diff --git a/autoarray/structures/arrays/array_2d_util.py b/autoarray/structures/arrays/array_2d_util.py index a286924c5..4b74fa07e 100644 --- a/autoarray/structures/arrays/array_2d_util.py +++ b/autoarray/structures/arrays/array_2d_util.py @@ -28,13 +28,8 @@ def convert_array(array: Union[np.ndarray, List]) -> np.ndarray: except AttributeError: pass - if isinstance(array, np.ndarray) or isinstance(array, list): - array = jnp.asarray(array) - elif isinstance(array, jnp.ndarray): - array = jax.lax.cond( - type(array) is list, lambda _: jnp.asarray(array), lambda _: array, None - ) - return array + return jnp.asarray(array) + def check_array_2d(array_2d: np.ndarray): diff --git a/autoarray/structures/arrays/uniform_1d.py b/autoarray/structures/arrays/uniform_1d.py index 234de540c..9b0af1af2 100644 --- a/autoarray/structures/arrays/uniform_1d.py +++ b/autoarray/structures/arrays/uniform_1d.py @@ -25,11 +25,6 @@ def __init__( store_native: bool = False, ): - try: - values = values._array - except AttributeError: - values = values - values = array_1d_util.convert_array_1d( array_1d=values, mask_1d=mask, From d132365b35c8c12eb5473cdc82e52ec1450be9de Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 13:36:58 +0100 Subject: [PATCH 31/62] array 1D stuff all updated to support JAX --- autoarray/mask/mask_1d_util.py | 49 +------------------- autoarray/structures/arrays/array_1d_util.py | 13 +----- test_autoarray/mask/test_mask_1d_util.py | 8 ---- 3 files changed, 4 insertions(+), 66 deletions(-) diff --git a/autoarray/mask/mask_1d_util.py b/autoarray/mask/mask_1d_util.py index 3d9943c19..cfc9e7dff 100644 --- a/autoarray/mask/mask_1d_util.py +++ b/autoarray/mask/mask_1d_util.py @@ -1,43 +1,7 @@ +import jax.numpy as jnp import numpy as np -from autoarray import numba_util - -@numba_util.jit() -def total_pixels_1d_from(mask_1d: np.ndarray) -> int: - """ - Returns the total number of unmasked pixels in a mask. - - Parameters - ---------- - mask_1d - A 2D array of bools, where `False` values are unmasked and included when counting pixels. - - Returns - ------- - int - The total number of pixels that are unmasked. - - Examples - -------- - - mask = np.array([[True, False, True], - [False, False, False] - [True, False, True]]) - - total_regular_pixels = total_regular_pixels_from(mask=mask) - """ - - total_regular_pixels = 0 - - for x in range(mask_1d.shape[0]): - if not mask_1d[x]: - total_regular_pixels += 1 - - return total_regular_pixels - - -@numba_util.jit() def native_index_for_slim_index_1d_from( mask_1d: np.ndarray, ) -> np.ndarray: @@ -70,14 +34,5 @@ def native_index_for_slim_index_1d_from( """ - total_pixels = total_pixels_1d_from(mask_1d=mask_1d) - native_index_for_slim_index_1d = np.zeros(shape=total_pixels) - - slim_index = 0 - - for x in range(mask_1d.shape[0]): - if not mask_1d[x]: - native_index_for_slim_index_1d[slim_index] = x - slim_index += 1 - + native_index_for_slim_index_1d = jnp.flatnonzero(~mask_1d) return native_index_for_slim_index_1d diff --git a/autoarray/structures/arrays/array_1d_util.py b/autoarray/structures/arrays/array_1d_util.py index 7e126b0a5..c8f519ae2 100644 --- a/autoarray/structures/arrays/array_1d_util.py +++ b/autoarray/structures/arrays/array_1d_util.py @@ -6,7 +6,6 @@ if TYPE_CHECKING: from autoarray.mask.mask_1d import Mask1D -from autoarray import numba_util from autoarray.mask import mask_1d_util from autoarray.structures.arrays import array_2d_util @@ -127,8 +126,6 @@ def array_1d_native_from( native_index_for_slim_index_1d=native_index_for_slim_index_1d, ) - -@numba_util.jit() def array_1d_via_indexes_1d_from( array_1d_slim: np.ndarray, shape: int, @@ -167,11 +164,5 @@ def array_1d_via_indexes_1d_from( ndarray The native 1D array of values mapped from the slimmed array with dimensions (total_x_pixels). """ - array_1d_native = np.zeros(shape) - - for slim_index in range(len(native_index_for_slim_index_1d)): - array_1d_native[native_index_for_slim_index_1d[slim_index]] = array_1d_slim[ - slim_index - ] - - return array_1d_native + array_1d_native = jnp.zeros(shape) + return array_1d_native.at[native_index_for_slim_index_1d].set(array_1d_slim) diff --git a/test_autoarray/mask/test_mask_1d_util.py b/test_autoarray/mask/test_mask_1d_util.py index c8251471f..1b032c37f 100644 --- a/test_autoarray/mask/test_mask_1d_util.py +++ b/test_autoarray/mask/test_mask_1d_util.py @@ -1,14 +1,6 @@ -from autoarray import exc from autoarray import util import numpy as np -import pytest - - -def test__total_image_pixels_1d_from(): - mask_1d = np.array([False, True, False, False, False, True]) - - assert util.mask_1d.total_pixels_1d_from(mask_1d=mask_1d) == 4 def test__native_index_for_slim_index_1d_from(): From bd88fee9fce571f43e1a641fa40798864f9d814c Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 13:43:51 +0100 Subject: [PATCH 32/62] grid_1d_slim_via_mask_from to JAX --- autoarray/geometry/geometry_util.py | 2 -- autoarray/structures/arrays/array_2d_util.py | 1 - autoarray/structures/grids/grid_1d_util.py | 26 +++++++------------- autoarray/structures/grids/grid_2d_util.py | 19 +++++++------- 4 files changed, 19 insertions(+), 29 deletions(-) diff --git a/autoarray/geometry/geometry_util.py b/autoarray/geometry/geometry_util.py index 9ca5fc48b..e92eea131 100644 --- a/autoarray/geometry/geometry_util.py +++ b/autoarray/geometry/geometry_util.py @@ -57,7 +57,6 @@ def convert_pixel_scales_1d(pixel_scales: ty.PixelScales) -> Tuple[float]: return pixel_scales -@numba_util.jit() def central_pixel_coordinates_1d_from( shape_slim: Tuple[int], ) -> Union[Tuple[float], Tuple[float]]: @@ -85,7 +84,6 @@ def central_pixel_coordinates_1d_from( return (float(shape_slim[0] - 1) / 2,) -@numba_util.jit() def central_scaled_coordinate_1d_from( shape_slim: Tuple[float], pixel_scales: ty.PixelScales, diff --git a/autoarray/structures/arrays/array_2d_util.py b/autoarray/structures/arrays/array_2d_util.py index 4b74fa07e..84248c2f4 100644 --- a/autoarray/structures/arrays/array_2d_util.py +++ b/autoarray/structures/arrays/array_2d_util.py @@ -1,5 +1,4 @@ from __future__ import annotations -import jax import jax.numpy as jnp import numpy as np from typing import TYPE_CHECKING, List, Tuple, Union diff --git a/autoarray/structures/grids/grid_1d_util.py b/autoarray/structures/grids/grid_1d_util.py index aa0592c47..bb99ccf27 100644 --- a/autoarray/structures/grids/grid_1d_util.py +++ b/autoarray/structures/grids/grid_1d_util.py @@ -1,5 +1,6 @@ from __future__ import annotations import numpy as np +import jax.numpy as jnp from typing import TYPE_CHECKING, List, Union, Tuple if TYPE_CHECKING: @@ -43,16 +44,18 @@ def convert_grid_1d( is_native = grid_1d.shape[0] == mask_1d.shape_native[0] + mask_1d = jnp.array(mask_1d.array) + if is_native == store_native: return grid_1d elif not store_native: return grid_1d_slim_from( grid_1d_native=grid_1d, - mask_1d=np.array(mask_1d), + mask_1d=mask_1d, ) return grid_1d_native_from( grid_1d_slim=grid_1d, - mask_1d=np.array(mask_1d), + mask_1d=mask_1d, ) @@ -95,7 +98,6 @@ def grid_1d_slim_via_shape_slim_from( ) -@numba_util.jit() def grid_1d_slim_via_mask_from( mask_1d: np.ndarray, pixel_scales: ty.PixelScales, @@ -131,23 +133,13 @@ def grid_1d_slim_via_mask_from( mask = np.array([True, False, True, False, False, False]) grid_slim = grid_1d_via_mask_from(mask_1d=mask_1d, pixel_scales=(0.5, 0.5), origin=(0.0, 0.0)) """ - - total_pixels = mask_1d_util.total_pixels_1d_from(mask_1d) - - grid_1d = np.zeros(shape=(total_pixels,)) - centres_scaled = geometry_util.central_scaled_coordinate_1d_from( shape_slim=mask_1d.shape, pixel_scales=pixel_scales, origin=origin ) - - index = 0 - - for x in range(mask_1d.shape[0]): - if not mask_1d[x]: - grid_1d[index] = (x - centres_scaled[0]) * pixel_scales[0] - index += 1 - - return grid_1d + indices = jnp.arange(mask_1d.shape[0]) + unmasked = jnp.logical_not(mask_1d) + coords = (indices - centres_scaled[0]) * pixel_scales[0] + return coords[unmasked] def grid_1d_slim_from( diff --git a/autoarray/structures/grids/grid_2d_util.py b/autoarray/structures/grids/grid_2d_util.py index e976be2b1..8c934d0c0 100644 --- a/autoarray/structures/grids/grid_2d_util.py +++ b/autoarray/structures/grids/grid_2d_util.py @@ -1,7 +1,6 @@ from __future__ import annotations import numpy as np import jax.numpy as jnp -import jax from typing import TYPE_CHECKING, List, Optional, Tuple, Union @@ -15,6 +14,16 @@ from autoarray import type as ty +def convert_grid(grid: Union[np.ndarray, List]) -> np.ndarray: + + try: + grid = grid.array + except AttributeError: + pass + + return jnp.asarray(grid) + + def check_grid_slim(grid, shape_native): if shape_native is None: raise exc.GridException( @@ -35,14 +44,6 @@ def check_grid_slim(grid, shape_native): """ ) - -def convert_grid(grid: Union[np.ndarray, List]) -> np.ndarray: - if type(grid) is list: - grid = np.asarray(grid) - - return grid - - def check_grid_2d(grid_2d: np.ndarray): if grid_2d.shape[-1] != 2: raise exc.GridException( From 5f8877bf3f3086c4f9e80665ccf83307c387db0c Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 13:50:10 +0100 Subject: [PATCH 33/62] grids now use JAX --- autoarray/structures/grids/grid_1d_util.py | 2 -- autoarray/structures/grids/grid_2d_util.py | 16 +++++++++------- autoarray/structures/grids/uniform_2d.py | 12 ++++-------- .../structures/grids/test_uniform_2d.py | 6 +++--- 4 files changed, 16 insertions(+), 20 deletions(-) diff --git a/autoarray/structures/grids/grid_1d_util.py b/autoarray/structures/grids/grid_1d_util.py index bb99ccf27..de1c0c725 100644 --- a/autoarray/structures/grids/grid_1d_util.py +++ b/autoarray/structures/grids/grid_1d_util.py @@ -7,10 +7,8 @@ from autoarray.mask.mask_1d import Mask1D from autoarray.structures.arrays import array_1d_util -from autoarray import numba_util from autoarray.geometry import geometry_util from autoarray.structures.grids import grid_2d_util -from autoarray.mask import mask_1d_util from autoarray import type as ty diff --git a/autoarray/structures/grids/grid_2d_util.py b/autoarray/structures/grids/grid_2d_util.py index 8c934d0c0..4dfe97ed7 100644 --- a/autoarray/structures/grids/grid_2d_util.py +++ b/autoarray/structures/grids/grid_2d_util.py @@ -113,20 +113,22 @@ def convert_grid_2d( is_native = len(grid_2d.shape) == 3 + mask_2d = jnp.array(mask_2d.array) + if is_native: - grid_2d[:, :, 0] *= np.invert(mask_2d) - grid_2d[:, :, 1] *= np.invert(mask_2d) + grid_2d = grid_2d.at[:, :, 0].multiply(jnp.invert(mask_2d)) + grid_2d = grid_2d.at[:, :, 1].multiply(jnp.invert(mask_2d)) if is_native == store_native: return grid_2d elif not store_native: return grid_2d_slim_from( - grid_2d_native=np.array(grid_2d), - mask=np.array(mask_2d), + grid_2d_native=grid_2d, + mask=mask_2d, ) return grid_2d_native_from( - grid_2d_slim=np.array(grid_2d), - mask_2d=np.array(mask_2d), + grid_2d_slim=grid_2d, + mask_2d=mask_2d, ) @@ -724,7 +726,7 @@ def grid_2d_native_from( mask_2d=mask_2d, ) - return np.stack((grid_2d_native_y, grid_2d_native_x), axis=-1) + return jnp.stack((grid_2d_native_y, grid_2d_native_x), axis=-1) @numba_util.jit() diff --git a/autoarray/structures/grids/uniform_2d.py b/autoarray/structures/grids/uniform_2d.py index 3e1152e4e..bea896b01 100644 --- a/autoarray/structures/grids/uniform_2d.py +++ b/autoarray/structures/grids/uniform_2d.py @@ -842,14 +842,10 @@ def squared_distances_to_coordinate_from( coordinate The (y,x) coordinate from which the squared distance of every grid (y,x) coordinate is computed. """ - if isinstance(self, jnp.ndarray): - squared_distances = jnp.square( - self.array[:, 0] - coordinate[0] - ) + jnp.square(self.array[:, 1] - coordinate[1]) - else: - squared_distances = np.square(self[:, 0] - coordinate[0]) + np.square( - self[:, 1] - coordinate[1] - ) + squared_distances = jnp.square( + self.array[:, 0] - coordinate[0] + ) + jnp.square(self.array[:, 1] - coordinate[1]) + return Array2D(values=squared_distances, mask=self.mask) def distances_to_coordinate_from( diff --git a/test_autoarray/structures/grids/test_uniform_2d.py b/test_autoarray/structures/grids/test_uniform_2d.py index 78813329e..ae5d53b10 100644 --- a/test_autoarray/structures/grids/test_uniform_2d.py +++ b/test_autoarray/structures/grids/test_uniform_2d.py @@ -441,7 +441,7 @@ def test__from_mask(): mask = aa.Mask2D(mask=mask, pixel_scales=(2.0, 2.0)) grid_via_util = aa.util.grid_2d.grid_2d_slim_via_mask_from( - mask_2d=np.array(mask), pixel_scales=(2.0, 2.0) + mask_2d=mask, pixel_scales=(2.0, 2.0) ) grid_2d = aa.Grid2D.from_mask(mask=mask) @@ -451,8 +451,8 @@ def test__from_mask(): assert grid_2d.pixel_scales == (2.0, 2.0) grid_2d_native = aa.util.grid_2d.grid_2d_native_from( - grid_2d_slim=np.array(grid_2d), - mask_2d=np.array(mask), + grid_2d_slim=grid_2d.array, + mask_2d=mask, ) assert (grid_2d_native == grid_2d.native).all() From 6b7d43341d342e39507e900798791a075581eed5 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 13:55:49 +0100 Subject: [PATCH 34/62] fix vectors --- autoarray/structures/vectors/irregular.py | 9 +++------ autoarray/structures/vectors/uniform.py | 8 -------- 2 files changed, 3 insertions(+), 14 deletions(-) diff --git a/autoarray/structures/vectors/irregular.py b/autoarray/structures/vectors/irregular.py index 5895149fb..e6501e46a 100644 --- a/autoarray/structures/vectors/irregular.py +++ b/autoarray/structures/vectors/irregular.py @@ -1,5 +1,6 @@ import logging import numpy as np +import jax.numpy as jnp from typing import List, Tuple, Union from autoarray.structures.vectors.abstract import AbstractVectorYX2D @@ -43,10 +44,6 @@ def __init__( grid The irregular grid of (y,x) coordinates where each vector is located. """ - - if type(values) is list: - values = np.asarray(values) - self.grid = Grid2DIrregular(values=grid) super().__init__(values) @@ -120,13 +117,13 @@ def vectors_within_radius( squared_distances = self.grid.distances_to_coordinate_from(coordinate=centre) mask = squared_distances < radius - if np.all(mask == False): + if jnp.all(mask == False): raise exc.VectorYXException( "The input radius removed all vectors / points on the grid." ) return VectorYX2DIrregular( - values=self[mask], grid=Grid2DIrregular(self.grid[mask]) + values=jnp.array(self.array)[mask], grid=Grid2DIrregular(self.grid[mask]) ) def vectors_within_annulus( diff --git a/autoarray/structures/vectors/uniform.py b/autoarray/structures/vectors/uniform.py index 62dd7592d..4da4836b2 100644 --- a/autoarray/structures/vectors/uniform.py +++ b/autoarray/structures/vectors/uniform.py @@ -137,14 +137,6 @@ def __init__( mapping large data arrays to and from the slim / native formats, which can be a computational bottleneck. """ - if type(values) is list: - values = np.asarray(values) - - try: - values = values._array - except AttributeError: - values = values - values = grid_2d_util.convert_grid_2d( grid_2d=values, mask_2d=mask, store_native=store_native ) From a6480a9089d3e0f3c9e260166eb5e47fde2951c3 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 14:06:09 +0100 Subject: [PATCH 35/62] fix some visualizatoin --- autoarray/plot/mat_plot/two_d.py | 6 +++--- autoarray/plot/visuals/two_d.py | 7 ++++--- autoarray/plot/wrap/two_d/grid_scatter.py | 3 --- .../arrays/files/array/output_test/array.fits | Bin 5760 -> 5760 bytes 4 files changed, 7 insertions(+), 9 deletions(-) diff --git a/autoarray/plot/mat_plot/two_d.py b/autoarray/plot/mat_plot/two_d.py index 9fcec3657..e28e81f6a 100644 --- a/autoarray/plot/mat_plot/two_d.py +++ b/autoarray/plot/mat_plot/two_d.py @@ -438,11 +438,11 @@ def plot_grid( if y_errors is None and x_errors is None: self.grid_scatter.scatter_grid_colored( - grid=grid, color_array=color_array, cmap=cmap + grid=grid.array, color_array=color_array, cmap=cmap ) else: self.grid_errorbar.errorbar_grid_colored( - grid=grid, + grid=grid.array, cmap=cmap, color_array=color_array, y_errors=y_errors, @@ -496,7 +496,7 @@ def plot_grid( self.contour.set(array=color_array, extent=extent, use_log10=self.use_log10) visuals_2d.plot_via_plotter( - plotter=self, grid_indexes=grid, geometry=grid.geometry + plotter=self, grid_indexes=grid.array, geometry=grid.geometry ) if not self.is_for_subplot: diff --git a/autoarray/plot/visuals/two_d.py b/autoarray/plot/visuals/two_d.py index d1da534ff..80e67fbab 100644 --- a/autoarray/plot/visuals/two_d.py +++ b/autoarray/plot/visuals/two_d.py @@ -1,3 +1,4 @@ +import numpy as np from matplotlib import patches as ptch from typing import List, Optional, Union @@ -54,13 +55,13 @@ def plot_via_plotter(self, plotter, grid_indexes=None, mapper=None, geometry=Non ) if self.mask is not None: - plotter.mask_scatter.scatter_grid(grid=self.mask.derive_grid.edge) + plotter.mask_scatter.scatter_grid(grid=np.array(self.mask.derive_grid.edge.array)) if self.border is not None: - plotter.border_scatter.scatter_grid(grid=self.border) + plotter.border_scatter.scatter_grid(grid=self.border.array) if self.grid is not None: - plotter.grid_scatter.scatter_grid(grid=self.grid) + plotter.grid_scatter.scatter_grid(grid=self.grid.array) if self.mesh_grid is not None: plotter.mesh_grid_scatter.scatter_grid(grid=self.mesh_grid) diff --git a/autoarray/plot/wrap/two_d/grid_scatter.py b/autoarray/plot/wrap/two_d/grid_scatter.py index ea1e3f6a8..526aeffb9 100644 --- a/autoarray/plot/wrap/two_d/grid_scatter.py +++ b/autoarray/plot/wrap/two_d/grid_scatter.py @@ -56,9 +56,6 @@ 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) or isinstance(grid, Abstract2DMesh): - grid = np.array(grid.array) - try: plt.scatter(y=grid[:, 0], x=grid[:, 1], **config_dict) except (IndexError, TypeError): diff --git a/test_autoarray/structures/arrays/files/array/output_test/array.fits b/test_autoarray/structures/arrays/files/array/output_test/array.fits index dab9da2fb676efc023e53b73faa54fe78fc4000c..0cff0a8f7db90d3ded28c644cd66e3791627feeb 100644 GIT binary patch delta 84 hcmZqBZP1;N!(?W%G4C>$;|B&XuqT_|T*&>83jh#~5uX46 delta 56 kcmZqBZP1;N!(?o$Vgmz%Jzl(dBKLc)$qTsk0j$6e2mk;8 From 1fde8b18c7223df8cf4aea20b408b6db04fd9a3c Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 14:14:11 +0100 Subject: [PATCH 36/62] risky change to figure open --- autoarray/plot/mat_plot/two_d.py | 2 ++ autoarray/plot/wrap/base/colorbar.py | 3 --- autoarray/plot/wrap/base/figure.py | 10 ++++------ 3 files changed, 6 insertions(+), 9 deletions(-) diff --git a/autoarray/plot/mat_plot/two_d.py b/autoarray/plot/mat_plot/two_d.py index e28e81f6a..fca62f0cf 100644 --- a/autoarray/plot/mat_plot/two_d.py +++ b/autoarray/plot/mat_plot/two_d.py @@ -417,6 +417,7 @@ def plot_grid( if not self.is_for_subplot: fig, ax = self.figure.open() + print(ax) else: ax = self.setup_subplot() @@ -450,6 +451,7 @@ def plot_grid( ) if self.colorbar is not None: + colorbar = self.colorbar.set_with_color_values( units=self.units, cmap=self.cmap.cmap, diff --git a/autoarray/plot/wrap/base/colorbar.py b/autoarray/plot/wrap/base/colorbar.py index b3577e9c2..1c9a8c5d8 100644 --- a/autoarray/plot/wrap/base/colorbar.py +++ b/autoarray/plot/wrap/base/colorbar.py @@ -195,9 +195,6 @@ 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/base/figure.py b/autoarray/plot/wrap/base/figure.py index 238264fd9..edea67465 100644 --- a/autoarray/plot/wrap/base/figure.py +++ b/autoarray/plot/wrap/base/figure.py @@ -82,12 +82,10 @@ def open(self): """ Wraps the Matplotlib method 'plt.figure' for opening a figure. """ - if not plt.fignum_exists(num=1): - config_dict = self.config_dict - config_dict.pop("aspect") - fig = plt.figure(**config_dict) - return fig, plt.gca() - return None, None + config_dict = self.config_dict + config_dict.pop("aspect") + fig = plt.figure(**config_dict) + return fig, plt.gca() def close(self): """ From 875483cf5f3d0f0889d6d133094de7aaef21fc80 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 14:17:15 +0100 Subject: [PATCH 37/62] fix structurre plotters now JAX arrays used --- autoarray/plot/mat_plot/one_d.py | 22 +++++++++++----------- autoarray/plot/mat_plot/two_d.py | 2 +- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/autoarray/plot/mat_plot/one_d.py b/autoarray/plot/mat_plot/one_d.py index 733f2ff83..aa5eb084b 100644 --- a/autoarray/plot/mat_plot/one_d.py +++ b/autoarray/plot/mat_plot/one_d.py @@ -166,7 +166,7 @@ def plot_yx( text_manual_dict_y=None, bypass: bool = False, ): - if (y is None) or np.count_nonzero(y) == 0 or np.isnan(y).all(): + if (y.array is None) or np.count_nonzero(y.array) == 0 or np.isnan(y.array).all(): return ax = None @@ -198,8 +198,8 @@ def plot_yx( label = self.legend.label or auto_labels.legend self.yx_plot.plot_y_vs_x( - y=y, - x=x, + y=y.array, + x=x.array, label=label, plot_axis_type=plot_axis_type, y_errors=y_errors, @@ -229,18 +229,18 @@ def plot_yx( plt.yscale("symlog") if x_errors is not None: - min_value_x = np.nanmin(x - x_errors) - max_value_x = np.nanmax(x + x_errors) + min_value_x = np.nanmin(x.array - x_errors) + max_value_x = np.nanmax(x.array + x_errors) else: - min_value_x = np.nanmin(x) - max_value_x = np.nanmax(x) + min_value_x = np.nanmin(x.array) + max_value_x = np.nanmax(x.array) if y_errors is not None: - min_value_y = np.nanmin(y - y_errors) - max_value_y = np.nanmax(y + y_errors) + min_value_y = np.nanmin(y.array - y_errors) + max_value_y = np.nanmax(y.array + y_errors) else: - min_value_y = np.nanmin(y) - max_value_y = np.nanmax(y) + min_value_y = np.nanmin(y.array) + max_value_y = np.nanmax(y.array) if should_plot_zero: if min_value_y > 0: diff --git a/autoarray/plot/mat_plot/two_d.py b/autoarray/plot/mat_plot/two_d.py index fca62f0cf..64811a830 100644 --- a/autoarray/plot/mat_plot/two_d.py +++ b/autoarray/plot/mat_plot/two_d.py @@ -313,7 +313,7 @@ def plot_array( aspect = self.figure.aspect_from(shape_native=array.shape_native) - norm = self.cmap.norm_from(array=array, use_log10=self.use_log10) + norm = self.cmap.norm_from(array=array.array, use_log10=self.use_log10) origin = conf.instance["visualize"]["general"]["general"]["imshow_origin"] From 1088baf5caec3b156ed7437e001942bd38bf1ff9 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 14:20:30 +0100 Subject: [PATCH 38/62] over sampler fix --- autoarray/operators/over_sampling/over_sampler.py | 4 ++-- autoarray/structures/grids/uniform_1d.py | 5 +++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/autoarray/operators/over_sampling/over_sampler.py b/autoarray/operators/over_sampling/over_sampler.py index 9fda67bb7..ccc09f04a 100644 --- a/autoarray/operators/over_sampling/over_sampler.py +++ b/autoarray/operators/over_sampling/over_sampler.py @@ -227,13 +227,13 @@ def binned_array_2d_from(self, array: Array2D) -> "Array2D": return self try: - array = array.slim + array = array.slim.array except AttributeError: pass if self.sub_is_uniform: binned_array_2d = array.reshape( - self.mask.shape_slim, self.sub_size[0] ** 2 + self.mask.shape_slim, self.sub_size.array[0] ** 2 ).mean(axis=1) else: diff --git a/autoarray/structures/grids/uniform_1d.py b/autoarray/structures/grids/uniform_1d.py index 53a9ec756..0a736e324 100644 --- a/autoarray/structures/grids/uniform_1d.py +++ b/autoarray/structures/grids/uniform_1d.py @@ -1,5 +1,6 @@ from __future__ import annotations import numpy as np +import jax.numpy as jnp from typing import List, Union, Tuple from autoarray.structures.abstract_structure import Structure @@ -311,8 +312,8 @@ def grid_2d_radial_projected_from(self, angle: float = 0.0) -> Grid2DIrregular: The projected and rotated 2D grid of (y,x) coordinates. """ - grid = np.zeros((self.mask.pixels_in_mask, 2)) - grid[:, 1] = self.slim + grid = jnp.zeros((self.mask.pixels_in_mask, 2)) + grid = grid.at[:, 1].set(self.slim.array) grid = geometry_util.transform_grid_2d_to_reference_frame( grid_2d=grid, centre=(0.0, 0.0), angle=angle From f72b15e86fee6a040313d1c1cb8043fc099a7fcc Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 14:25:21 +0100 Subject: [PATCH 39/62] fix kmesh --- autoarray/structures/mesh/delaunay_2d.py | 2 +- autoarray/structures/mesh/rectangular_2d.py | 2 +- autoarray/structures/mesh/triangulation_2d.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/autoarray/structures/mesh/delaunay_2d.py b/autoarray/structures/mesh/delaunay_2d.py index a6b7f9012..e12d4898c 100644 --- a/autoarray/structures/mesh/delaunay_2d.py +++ b/autoarray/structures/mesh/delaunay_2d.py @@ -71,7 +71,7 @@ def interpolated_array_from( interpolated_array = mesh_util.delaunay_interpolated_array_from( shape_native=shape_native, - interpolation_grid_slim=interpolation_grid.slim, + interpolation_grid_slim=np.array(interpolation_grid.slim.array), delaunay=self.delaunay, pixel_values=values, ) diff --git a/autoarray/structures/mesh/rectangular_2d.py b/autoarray/structures/mesh/rectangular_2d.py index 0e4eab108..7541d592a 100644 --- a/autoarray/structures/mesh/rectangular_2d.py +++ b/autoarray/structures/mesh/rectangular_2d.py @@ -178,7 +178,7 @@ def interpolated_array_from( shape_native=shape_native, extent=extent ) - interpolated_array = griddata(points=self, values=values, xi=interpolation_grid) + interpolated_array = griddata(points=self.array, values=values, xi=interpolation_grid) interpolated_array = interpolated_array.reshape(shape_native) diff --git a/autoarray/structures/mesh/triangulation_2d.py b/autoarray/structures/mesh/triangulation_2d.py index b5c9fedc9..7d96415b4 100644 --- a/autoarray/structures/mesh/triangulation_2d.py +++ b/autoarray/structures/mesh/triangulation_2d.py @@ -96,7 +96,7 @@ def delaunay(self) -> scipy.spatial.Delaunay: `MeshException`, which helps exception handling in the `inversion` package. """ try: - return scipy.spatial.Delaunay(np.asarray([self[:, 0], self[:, 1]]).T) + return scipy.spatial.Delaunay(np.asarray([self.array[:, 0], self.array[:, 1]]).T) except (ValueError, OverflowError, scipy.spatial.qhull.QhullError) as e: raise exc.MeshException() from e From ac68a08b07726ca985d4d91e670eec5f805e4b0b Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 14:26:58 +0100 Subject: [PATCH 40/62] fix test repr --- test_autoarray/structures/arrays/test_repr.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test_autoarray/structures/arrays/test_repr.py b/test_autoarray/structures/arrays/test_repr.py index e70c98056..6e806a10d 100644 --- a/test_autoarray/structures/arrays/test_repr.py +++ b/test_autoarray/structures/arrays/test_repr.py @@ -3,4 +3,4 @@ def test_repr(): array = aa.Array2D.no_mask([[1, 2], [3, 4]], pixel_scales=1) - assert repr(array) == "Array([1, 2, 3, 4], dtype=int32)" + assert repr(array) == "Array([1, 2, 3, 4], dtype=int64)" From 678673e02de0af777c08709bb2e29d6cb12f5546 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 15:09:48 +0100 Subject: [PATCH 41/62] fix test preprocess --- autoarray/dataset/preprocess.py | 9 +++-- test_autoarray/dataset/test_preprocess.py | 42 +++++++---------------- 2 files changed, 17 insertions(+), 34 deletions(-) diff --git a/autoarray/dataset/preprocess.py b/autoarray/dataset/preprocess.py index f13af3184..f411b8302 100644 --- a/autoarray/dataset/preprocess.py +++ b/autoarray/dataset/preprocess.py @@ -149,7 +149,7 @@ def noise_map_via_data_eps_and_exposure_time_map_from(data_eps, exposure_time_ma The exposure time at every data-point of the data. """ return data_eps.with_new_array( - np.abs(data_eps * exposure_time_map) ** 0.5 / exposure_time_map + np.abs(data_eps.array * exposure_time_map.array) ** 0.5 / exposure_time_map.array ) @@ -406,9 +406,10 @@ def poisson_noise_via_data_eps_from(data_eps, exposure_time_map, seed=-1): An array describing simulated poisson noise_maps """ setup_random_seed(seed) - image_counts = np.multiply(data_eps, exposure_time_map) + + image_counts = np.multiply(data_eps.array, exposure_time_map.array) return data_eps - np.divide( - np.random.poisson(image_counts, data_eps.shape), exposure_time_map + np.random.poisson(image_counts, data_eps.shape), exposure_time_map.array ) @@ -506,8 +507,6 @@ def noise_map_with_signal_to_noise_limit_from( from autoarray.structures.arrays.uniform_1d import Array1D from autoarray.structures.arrays.uniform_2d import Array2D - # TODO : Refacotr into a util - signal_to_noise_map = data / noise_map signal_to_noise_map[signal_to_noise_map < 0] = 0 diff --git a/test_autoarray/dataset/test_preprocess.py b/test_autoarray/dataset/test_preprocess.py index f484fd648..8ea4e4bd9 100644 --- a/test_autoarray/dataset/test_preprocess.py +++ b/test_autoarray/dataset/test_preprocess.py @@ -493,11 +493,11 @@ def test__background_noise_map_via_edges_of_image_from_5(): def test__exposure_time_map_from_exposure_time_and_inverse_noise_map(): exposure_time = 6.0 - background_noise_map = aa.Array2D.full( - fill_value=0.25, shape_native=(3, 3), pixel_scales=1.0 - ) - background_noise_map[0] = 0.5 + background_noise_map = aa.Array2D.no_mask( + [[0.5, 0.25, 0.25], [0.25, 0.25, 0.25], [0.25, 0.25, 0.25]], pixel_scales=1.0 + ) + exposure_time_map = ( aa.preprocess.exposure_time_map_via_exposure_time_and_background_noise_map_from( exposure_time=exposure_time, background_noise_map=background_noise_map @@ -556,6 +556,7 @@ def test__poisson_noise_from_data(): def test__data_with_poisson_noised_added(): data = aa.Array2D.zeros(shape_native=(2, 2), pixel_scales=1.0) exposure_time_map = aa.Array2D.ones(shape_native=(2, 2), pixel_scales=1.0) + data_with_poisson_noise = aa.preprocess.data_eps_with_poisson_noise_added( data_eps=data, exposure_time_map=exposure_time_map, seed=1 ) @@ -661,49 +662,32 @@ def test__data_with_complex_gaussian_noise_added(): def test__noise_map_with_signal_to_noise_limit_from(): - image = aa.Array2D.full(fill_value=20.0, shape_native=(2, 2), pixel_scales=1.0) - image[3] = 5.0 - noise_map_array = aa.Array2D.full( - fill_value=5.0, shape_native=(2, 2), pixel_scales=1.0 - ) - noise_map_array[3] = 2.0 + image = aa.Array2D.no_mask(values=[[20, 20], [20, 5]], pixel_scales=1.0) + noise_map = aa.Array2D.no_mask(values=[[5, 5], [5, 2]], pixel_scales=1.0) noise_map = aa.preprocess.noise_map_with_signal_to_noise_limit_from( - data=image, noise_map=noise_map_array, signal_to_noise_limit=100.0 + data=image, noise_map=noise_map, signal_to_noise_limit=100.0 ) assert (noise_map.slim == np.array([5.0, 5.0, 5.0, 2.0])).all() - image = aa.Array2D.full(fill_value=20.0, shape_native=(2, 2), pixel_scales=1.0) - image[3] = 5.0 - - noise_map_array = aa.Array2D.full( - fill_value=5.0, shape_native=(2, 2), pixel_scales=1.0 - ) - noise_map_array[3] = 2.0 + noise_map = aa.Array2D.no_mask(values=[[5, 5], [5, 2]], pixel_scales=1.0) noise_map = aa.preprocess.noise_map_with_signal_to_noise_limit_from( - data=image, noise_map=noise_map_array, signal_to_noise_limit=2.0 + data=image, noise_map=noise_map, signal_to_noise_limit=2.0 ) assert (noise_map.native == np.array([[10.0, 10.0], [10.0, 2.5]])).all() - image = aa.Array2D.full(fill_value=20.0, shape_native=(2, 2), pixel_scales=1.0) - image[2] = 5.0 - image[3] = 5.0 - - noise_map_array = aa.Array2D.full( - fill_value=5.0, shape_native=(2, 2), pixel_scales=1.0 - ) - noise_map_array[2] = 2.0 - noise_map_array[3] = 2.0 + image = aa.Array2D.no_mask(values=[[20, 20], [5, 5]], pixel_scales=1.0) + noise_map = aa.Array2D.no_mask(values=[[5, 5], [2, 2]], pixel_scales=1.0) mask = aa.Mask2D(mask=[[True, False], [False, True]], pixel_scales=1.0) noise_map = aa.preprocess.noise_map_with_signal_to_noise_limit_from( data=image, - noise_map=noise_map_array, + noise_map=noise_map, signal_to_noise_limit=2.0, noise_limit_mask=mask, ) From c611bb62b2695fdaf4f8dc6eb803361b776bc279 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 15:13:37 +0100 Subject: [PATCH 42/62] fix abstract dataset tests --- .../dataset/abstract/test_dataset.py | 19 +++++++++++-------- test_autoarray/dataset/test_preprocess.py | 2 +- 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/test_autoarray/dataset/abstract/test_dataset.py b/test_autoarray/dataset/abstract/test_dataset.py index dd99c1201..1bd36a2fe 100644 --- a/test_autoarray/dataset/abstract/test_dataset.py +++ b/test_autoarray/dataset/abstract/test_dataset.py @@ -115,15 +115,18 @@ def test__grid_settings__sub_size(image_7x7, noise_map_7x7): def test__new_imaging_with_arrays_trimmed_via_kernel_shape(): - data = aa.Array2D.full(fill_value=20.0, shape_native=(3, 3), pixel_scales=1.0) - data[4] = 5.0 - noise_map_array = aa.Array2D.full( - fill_value=5.0, shape_native=(3, 3), pixel_scales=1.0 + data = aa.Array2D.no_mask( + values=[[20.0, 20.0, 20.0], [20.0, 5.0, 20.0], [20.0, 20.0, 20.0]], + pixel_scales=1.0, ) - noise_map_array[4] = 2.0 - dataset = ds.AbstractDataset(data=data, noise_map=noise_map_array) + noise_map = aa.Array2D.no_mask( + values=[[20.0, 20.0, 20.0], [20.0, 2.0, 20.0], [20.0, 20.0, 20.0]], + pixel_scales=1.0, + ) + + dataset = ds.AbstractDataset(data=data, noise_map=noise_map) dataset_trimmed = dataset.trimmed_after_convolution_from(kernel_shape=(3, 3)) @@ -146,8 +149,8 @@ def test__apply_over_sampling(image_7x7, noise_map_7x7): grid_sub_2 = dataset_7x7.grids.lp grids_pixelization_sub_2 = dataset_7x7.grids.pixelization - dataset_7x7.grids.__dict__["lp"][0][0] = 100.0 - dataset_7x7.grids.__dict__["pixelization"][0][0] = 100.0 + dataset_7x7.grids.lp = dataset_7x7.grids.lp.at[0, 0].set(100.0) + dataset_7x7.grids.pixelization = dataset_7x7.grids.pixelization.at[0, 0].set(100.0) assert dataset_7x7.grids.lp[0][0] == pytest.approx(100.0, 1.0e-4) assert dataset_7x7.grids.pixelization[0][0] == pytest.approx(100.0, 1.0e-4) diff --git a/test_autoarray/dataset/test_preprocess.py b/test_autoarray/dataset/test_preprocess.py index 8ea4e4bd9..108253a08 100644 --- a/test_autoarray/dataset/test_preprocess.py +++ b/test_autoarray/dataset/test_preprocess.py @@ -497,7 +497,7 @@ def test__exposure_time_map_from_exposure_time_and_inverse_noise_map(): background_noise_map = aa.Array2D.no_mask( [[0.5, 0.25, 0.25], [0.25, 0.25, 0.25], [0.25, 0.25, 0.25]], pixel_scales=1.0 ) - + exposure_time_map = ( aa.preprocess.exposure_time_map_via_exposure_time_and_background_noise_map_from( exposure_time=exposure_time, background_noise_map=background_noise_map From e1899a1b9b2127db53338d53cdd0c5b6aa266372 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 15:20:09 +0100 Subject: [PATCH 43/62] fix test imaigng data --- test_autoarray/dataset/imaging/test_dataset.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/test_autoarray/dataset/imaging/test_dataset.py b/test_autoarray/dataset/imaging/test_dataset.py index ca33f1b40..417e1bc18 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 == pytest.approx((1.0 / 3.0) * psf_3x3.slim, 1.0e-4) + assert masked_imaging_7x7.psf.slim == pytest.approx((1.0 / 3.0) * psf_3x3.slim.array, 1.0e-4) assert type(masked_imaging_7x7.psf) == aa.Kernel2D assert masked_imaging_7x7.w_tilde.curvature_preload.shape == (35,) @@ -156,10 +156,19 @@ def test__apply_noise_scaling(imaging_7x7, mask_2d_7x7): assert masked_imaging_7x7.noise_map.native[4, 4] == 1e5 -def test__apply_noise_scaling__use_signal_to_noise_value(imaging_7x7, mask_2d_7x7): - imaging_7x7 = copy.copy(imaging_7x7) +def test__apply_noise_scaling__use_signal_to_noise_value(image_7x7, psf_3x3, noise_map_7x7, mask_2d_7x7): - imaging_7x7.data[24] = 2.0 + image_7x7 = np.array(image_7x7.native.array) + image_7x7[3,3] = 2.0 + + image_7x7 = aa.Array2D(values=image_7x7, mask=mask_2d_7x7) + + imaging_7x7 = aa.Imaging( + data=image_7x7, + psf=psf_3x3, + noise_map=noise_map_7x7, + over_sample_size_lp=1, + ) masked_imaging_7x7 = imaging_7x7.apply_noise_scaling( mask=mask_2d_7x7, signal_to_noise_value=0.1, should_zero_data=False From 7d5ba268336107e078cdd91b2ced65fe002acfbf Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 15:22:51 +0100 Subject: [PATCH 44/62] fix simulator --- autoarray/dataset/imaging/simulator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/autoarray/dataset/imaging/simulator.py b/autoarray/dataset/imaging/simulator.py index 449a3d4b6..b1bb23198 100644 --- a/autoarray/dataset/imaging/simulator.py +++ b/autoarray/dataset/imaging/simulator.py @@ -151,7 +151,7 @@ def via_image_from( pixel_scales=image.pixel_scales, ) - if np.isnan(noise_map).any(): + if np.isnan(noise_map.array).any(): raise exc.DatasetException( "The noise-map has NaN values in it. This suggests your exposure time and / or" "background sky levels are too low, creating signal counts at or close to 0.0." From 1feabff569f2d00680e68880c2a79b37ed918a1f Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 15:24:49 +0100 Subject: [PATCH 45/62] fix interferometrer --- autoarray/operators/transformer.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/autoarray/operators/transformer.py b/autoarray/operators/transformer.py index 3b19f7f7c..5be47539f 100644 --- a/autoarray/operators/transformer.py +++ b/autoarray/operators/transformer.py @@ -48,12 +48,12 @@ def __init__(self, uv_wavelengths, real_space_mask, preload_transform=True): if preload_transform: self.preload_real_transforms = transformer_util.preload_real_transforms( - grid_radians=np.array(self.grid), + grid_radians=np.array(self.grid.array), uv_wavelengths=self.uv_wavelengths, ) self.preload_imag_transforms = transformer_util.preload_imag_transforms( - grid_radians=np.array(self.grid), + grid_radians=np.array(self.grid.array), uv_wavelengths=self.uv_wavelengths, ) @@ -79,14 +79,14 @@ def __init__(self, uv_wavelengths, real_space_mask, preload_transform=True): def visibilities_from(self, image): if self.preload_transform: visibilities = transformer_util.visibilities_via_preload_jit_from( - image_1d=np.array(image), + image_1d=np.array(image.array), preloaded_reals=self.preload_real_transforms, preloaded_imags=self.preload_imag_transforms, ) else: visibilities = transformer_util.visibilities_jit( - image_1d=np.array(image.slim), + image_1d=np.array(image.slim.array), grid_radians=np.array(self.grid), uv_wavelengths=self.uv_wavelengths, ) From 2c02e39c088ca2afee7d8b99b71389cee09e6b14 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 15:33:34 +0100 Subject: [PATCH 46/62] dataset tests pass --- autoarray/operators/transformer.py | 2 +- autoarray/plot/mat_plot/one_d.py | 35 ++++++++++++++++++++---------- 2 files changed, 24 insertions(+), 13 deletions(-) diff --git a/autoarray/operators/transformer.py b/autoarray/operators/transformer.py index 5be47539f..c5b1d3a50 100644 --- a/autoarray/operators/transformer.py +++ b/autoarray/operators/transformer.py @@ -96,7 +96,7 @@ def visibilities_from(self, image): def image_from(self, visibilities, use_adjoint_scaling: bool = False): image_slim = transformer_util.image_via_jit_from( n_pixels=self.grid.shape[0], - grid_radians=np.array(self.grid), + grid_radians=np.array(self.grid.array), uv_wavelengths=self.uv_wavelengths, visibilities=visibilities.in_array, ) diff --git a/autoarray/plot/mat_plot/one_d.py b/autoarray/plot/mat_plot/one_d.py index aa5eb084b..79b721bba 100644 --- a/autoarray/plot/mat_plot/one_d.py +++ b/autoarray/plot/mat_plot/one_d.py @@ -166,7 +166,18 @@ def plot_yx( text_manual_dict_y=None, bypass: bool = False, ): - if (y.array is None) or np.count_nonzero(y.array) == 0 or np.isnan(y.array).all(): + + try: + y = y.array + except AttributeError: + pass + + try: + x = x.array + except AttributeError: + pass + + if (y is None) or np.count_nonzero(y) == 0 or np.isnan(y).all(): return ax = None @@ -185,7 +196,7 @@ def plot_yx( x = np.arange(len(y)) use_integers = True pixel_scales = (x[1] - x[0],) - x = Array1D.no_mask(values=x, pixel_scales=pixel_scales) + x = Array1D.no_mask(values=x, pixel_scales=pixel_scales).array if self.yx_plot.plot_axis_type is None: plot_axis_type = "linear" @@ -198,8 +209,8 @@ def plot_yx( label = self.legend.label or auto_labels.legend self.yx_plot.plot_y_vs_x( - y=y.array, - x=x.array, + y=y, + x=x, label=label, plot_axis_type=plot_axis_type, y_errors=y_errors, @@ -229,18 +240,18 @@ def plot_yx( plt.yscale("symlog") if x_errors is not None: - min_value_x = np.nanmin(x.array - x_errors) - max_value_x = np.nanmax(x.array + x_errors) + min_value_x = np.nanmin(x - x_errors) + max_value_x = np.nanmax(x + x_errors) else: - min_value_x = np.nanmin(x.array) - max_value_x = np.nanmax(x.array) + min_value_x = np.nanmin(x) + max_value_x = np.nanmax(x) if y_errors is not None: - min_value_y = np.nanmin(y.array - y_errors) - max_value_y = np.nanmax(y.array + y_errors) + min_value_y = np.nanmin(y - y_errors) + max_value_y = np.nanmax(y + y_errors) else: - min_value_y = np.nanmin(y.array) - max_value_y = np.nanmax(y.array) + min_value_y = np.nanmin(y) + max_value_y = np.nanmax(y) if should_plot_zero: if min_value_y > 0: From 87931b6975e6034dd4a8819160f21dbd968f422f Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 15:38:46 +0100 Subject: [PATCH 47/62] fix fit tests --- autoarray/fit/fit_dataset.py | 4 ++-- autoarray/fit/fit_util.py | 4 ++-- test_autoarray/fit/test_fit_imaging.py | 6 +++--- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/autoarray/fit/fit_dataset.py b/autoarray/fit/fit_dataset.py index c2f498949..a21b98885 100644 --- a/autoarray/fit/fit_dataset.py +++ b/autoarray/fit/fit_dataset.py @@ -86,7 +86,7 @@ def chi_squared(self) -> float: """ Returns the chi-squared terms of the model data's fit to an dataset, by summing the chi-squared-map. """ - return fit_util.chi_squared_from(chi_squared_map=self.chi_squared_map) + return fit_util.chi_squared_from(chi_squared_map=self.chi_squared_map.array) @property def noise_normalization(self) -> float: @@ -95,7 +95,7 @@ def noise_normalization(self) -> float: [Noise_Term] = sum(log(2*pi*[Noise]**2.0)) """ - return fit_util.noise_normalization_from(noise_map=self.noise_map) + return fit_util.noise_normalization_from(noise_map=self.noise_map.array) @property def log_likelihood(self) -> float: diff --git a/autoarray/fit/fit_util.py b/autoarray/fit/fit_util.py index f28a64e22..2ad7d0806 100644 --- a/autoarray/fit/fit_util.py +++ b/autoarray/fit/fit_util.py @@ -84,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 jnp.sum(np.array(chi_squared_map)) + return jnp.sum(chi_squared_map) def noise_normalization_from(*, noise_map: ty.DataLike) -> float: @@ -98,7 +98,7 @@ def noise_normalization_from(*, noise_map: ty.DataLike) -> float: noise_map The masked noise-map of the dataset. """ - return jnp.sum(jnp.log(2 * jnp.pi * np.array(noise_map)**2.0)) + return jnp.sum(jnp.log(2 * jnp.pi * noise_map**2.0)) def normalized_residual_map_complex_from( diff --git a/test_autoarray/fit/test_fit_imaging.py b/test_autoarray/fit/test_fit_imaging.py index c0c525e71..5b7e49446 100644 --- a/test_autoarray/fit/test_fit_imaging.py +++ b/test_autoarray/fit/test_fit_imaging.py @@ -30,7 +30,7 @@ def test__data_and_model_are_identical__no_masking__check_values_are_correct(): assert fit.chi_squared == 0.0 assert fit.reduced_chi_squared == 0.0 - assert fit.noise_normalization == np.sum(np.log(2 * np.pi * noise_map**2.0)) + assert fit.noise_normalization == np.sum(np.log(2 * np.pi * noise_map.array**2.0)) assert fit.log_likelihood == -0.5 * (fit.chi_squared + fit.noise_normalization) @@ -59,7 +59,7 @@ def test__data_and_model_are_different__include_masking__check_values_are_correc assert fit.chi_squared == 0.25 assert fit.reduced_chi_squared == 0.25 / 3.0 - assert fit.noise_normalization == np.sum(np.log(2 * np.pi * noise_map**2.0)) + assert fit.noise_normalization == np.sum(np.log(2 * np.pi * noise_map.array**2.0)) assert fit.log_likelihood == -0.5 * (fit.chi_squared + fit.noise_normalization) @@ -92,7 +92,7 @@ def test__data_and_model_are_identical__inversion_included__changes_certain_prop assert fit.chi_squared == 0.0 assert fit.reduced_chi_squared == 0.0 - assert fit.noise_normalization == np.sum(np.log(2 * np.pi * noise_map**2.0)) + assert fit.noise_normalization == np.sum(np.log(2 * np.pi * noise_map.array**2.0)) assert fit.log_likelihood == -0.5 * (fit.chi_squared + fit.noise_normalization) assert fit.log_likelihood_with_regularization == -0.5 * ( From 0fa9646a9736a0a3f02eb6c448490514e349533d Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 16:14:46 +0100 Subject: [PATCH 48/62] mask array 2d choose type based on if input is jax or numpy --- autoarray/dataset/imaging/dataset.py | 2 ++ autoarray/structures/arrays/array_2d_util.py | 21 ++++++++++----- autoarray/structures/arrays/uniform_2d.py | 10 ++++---- .../dataset/imaging/test_dataset.py | 1 + test_autoarray/layout/test_region.py | 4 +-- .../arrays/files/array/output_test/array.fits | Bin 5760 -> 5760 bytes .../structures/arrays/test_kernel_2d.py | 24 +++++++++--------- test_autoarray/structures/arrays/test_repr.py | 2 +- 8 files changed, 37 insertions(+), 27 deletions(-) diff --git a/autoarray/dataset/imaging/dataset.py b/autoarray/dataset/imaging/dataset.py index 9fa00da1a..7380cdcba 100644 --- a/autoarray/dataset/imaging/dataset.py +++ b/autoarray/dataset/imaging/dataset.py @@ -135,6 +135,8 @@ def __init__( f"the image and noise-map yourself." ) + print(type(data.array)) + super().__init__( data=data, noise_map=noise_map, diff --git a/autoarray/structures/arrays/array_2d_util.py b/autoarray/structures/arrays/array_2d_util.py index 84248c2f4..eadbd1b6f 100644 --- a/autoarray/structures/arrays/array_2d_util.py +++ b/autoarray/structures/arrays/array_2d_util.py @@ -27,7 +27,10 @@ def convert_array(array: Union[np.ndarray, List]) -> np.ndarray: except AttributeError: pass - return jnp.asarray(array) + if isinstance(array, list): + array = np.asarray(array) + + return array @@ -119,27 +122,31 @@ def convert_array_2d( """ array_2d = convert_array(array=array_2d).copy() - check_array_2d_and_mask_2d(array_2d=array_2d, mask_2d=mask_2d) + if isinstance(array_2d, np.ndarray): + is_numpy = True + else: + is_numpy = False - mask_2d = jnp.array(mask_2d.array) + check_array_2d_and_mask_2d(array_2d=array_2d, mask_2d=mask_2d) is_native = len(array_2d.shape) == 2 if is_native and not skip_mask: - array_2d *= jnp.invert(mask_2d) + array_2d *= np.invert(mask_2d) if is_native == store_native: - return array_2d + return np.array(array_2d) if is_numpy else jnp.array(array_2d) elif not store_native: - return array_2d_slim_from( + array_2d = array_2d_slim_from( array_2d_native=array_2d, mask_2d=mask_2d, ) + return np.array(array_2d) if is_numpy else jnp.array(array_2d) array_2d = array_2d_native_from( array_2d_slim=array_2d, mask_2d=mask_2d, ) - return array_2d + return np.array(array_2d) if is_numpy else jnp.array(array_2d) def convert_array_2d_to_slim(array_2d: np.ndarray, mask_2d: Mask2D) -> np.ndarray: diff --git a/autoarray/structures/arrays/uniform_2d.py b/autoarray/structures/arrays/uniform_2d.py index 1ab5db723..574d2fa5d 100644 --- a/autoarray/structures/arrays/uniform_2d.py +++ b/autoarray/structures/arrays/uniform_2d.py @@ -459,7 +459,7 @@ def zoomed_around_mask(self, buffer: int = 1) -> "Array2D": """ extracted_array_2d = array_2d_util.extracted_array_2d_from( - array_2d=np.array(self.native._array), + array_2d=np.array(self.native.array), y0=self.mask.zoom_region[0] - buffer, y1=self.mask.zoom_region[1] + buffer, x0=self.mask.zoom_region[2] - buffer, @@ -493,7 +493,7 @@ def extent_of_zoomed_array(self, buffer: int = 1) -> np.ndarray: The number pixels around the extracted array used as a buffer. """ extracted_array_2d = array_2d_util.extracted_array_2d_from( - array_2d=np.array(self.native._array), + array_2d=np.array(self.native.array), y0=self.mask.zoom_region[0] - buffer, y1=self.mask.zoom_region[1] + buffer, x0=self.mask.zoom_region[2] - buffer, @@ -527,7 +527,7 @@ def resized_from( """ resized_array_2d = array_2d_util.resized_array_2d_from( - array_2d=np.array(self.native._array), resized_shape=new_shape + array_2d=np.array(self.native.array), resized_shape=new_shape ) resized_mask = self.mask.resized_from( @@ -587,14 +587,14 @@ def trimmed_after_convolution_from( psf_cut_x = int(np.ceil(kernel_shape[1] / 2)) - 1 array_y = int(self.mask.shape[0]) array_x = int(self.mask.shape[1]) - trimmed_array_2d = self.native[ + trimmed_array_2d = self.native.array[ psf_cut_y : array_y - psf_cut_y, psf_cut_x : array_x - psf_cut_x ] resized_mask = self.mask.resized_from(new_shape=trimmed_array_2d.shape) array = array_2d_util.convert_array_2d( - array_2d=trimmed_array_2d._array, mask_2d=resized_mask + array_2d=trimmed_array_2d, mask_2d=resized_mask ) return Array2D( diff --git a/test_autoarray/dataset/imaging/test_dataset.py b/test_autoarray/dataset/imaging/test_dataset.py index 417e1bc18..f2f2b6d27 100644 --- a/test_autoarray/dataset/imaging/test_dataset.py +++ b/test_autoarray/dataset/imaging/test_dataset.py @@ -105,6 +105,7 @@ def test__from_fits(): def test__output_to_fits(imaging_7x7, test_data_path): + imaging_7x7.output_to_fits( data_path=path.join(test_data_path, "data.fits"), psf_path=path.join(test_data_path, "psf.fits"), diff --git a/test_autoarray/layout/test_region.py b/test_autoarray/layout/test_region.py index 643b8532a..20809cc7d 100644 --- a/test_autoarray/layout/test_region.py +++ b/test_autoarray/layout/test_region.py @@ -31,7 +31,7 @@ def test__slice_1d__addition(): region = aa.Region1D(region=(0, 1)) - arr_1d[region.slice] += image[region.slice] + arr_1d = arr_1d.array.at[region.slice].add(image[region.slice]) assert (arr_1d == np.array([2.0, 2.0, 3.0, 4.0])).all() @@ -39,7 +39,7 @@ def test__slice_1d__addition(): region = aa.Region1D(region=(2, 4)) - arr_1d[region.slice] += image[region.slice] + arr_1d = arr_1d.array.at[region.slice].add(image[region.slice]) assert (arr_1d == np.array([1.0, 2.0, 4.0, 5.0])).all() diff --git a/test_autoarray/structures/arrays/files/array/output_test/array.fits b/test_autoarray/structures/arrays/files/array/output_test/array.fits index 0cff0a8f7db90d3ded28c644cd66e3791627feeb..dab9da2fb676efc023e53b73faa54fe78fc4000c 100644 GIT binary patch delta 56 kcmZqBZP1;N!(?o$Vgmz%Jzl(dBKLc)$qTsk0j$6e2mk;8 delta 84 hcmZqBZP1;N!(?W%G4C>$;|B&XuqT_|T*&>83jh#~5uX46 diff --git a/test_autoarray/structures/arrays/test_kernel_2d.py b/test_autoarray/structures/arrays/test_kernel_2d.py index 4cf8b92d7..39ba25842 100644 --- a/test_autoarray/structures/arrays/test_kernel_2d.py +++ b/test_autoarray/structures/arrays/test_kernel_2d.py @@ -150,7 +150,7 @@ def test__rescaled_with_odd_dimensions_from__evens_to_odds(): rescale_factor=0.5, normalize=True ) assert kernel_2d.pixel_scales == (2.0, 2.0) - assert (kernel_2d.native == (1.0 / 9.0) * np.ones((3, 3))).all() + assert kernel_2d.native == pytest.approx((1.0 / 9.0) * np.ones((3, 3)), 1.0e-4) array_2d = np.ones((9, 9)) kernel_2d = aa.Kernel2D.no_mask(values=array_2d, pixel_scales=1.0, normalize=False) @@ -158,7 +158,7 @@ def test__rescaled_with_odd_dimensions_from__evens_to_odds(): rescale_factor=0.333333333333333, normalize=True ) assert kernel_2d.pixel_scales == (3.0, 3.0) - assert (kernel_2d.native == (1.0 / 9.0) * np.ones((3, 3))).all() + assert kernel_2d.native == pytest.approx((1.0 / 9.0) * np.ones((3, 3)), 1.0e-4) array_2d = np.ones((18, 6)) kernel_2d = aa.Kernel2D.no_mask(values=array_2d, pixel_scales=1.0, normalize=False) @@ -166,7 +166,7 @@ def test__rescaled_with_odd_dimensions_from__evens_to_odds(): rescale_factor=0.5, normalize=True ) assert kernel_2d.pixel_scales == (2.0, 2.0) - assert (kernel_2d.native == (1.0 / 27.0) * np.ones((9, 3))).all() + assert kernel_2d.native == pytest.approx((1.0 / 27.0) * np.ones((9, 3)), 1.0e-4) array_2d = np.ones((6, 18)) kernel_2d = aa.Kernel2D.no_mask(values=array_2d, pixel_scales=1.0, normalize=False) @@ -174,7 +174,7 @@ def test__rescaled_with_odd_dimensions_from__evens_to_odds(): rescale_factor=0.5, normalize=True ) assert kernel_2d.pixel_scales == (2.0, 2.0) - assert (kernel_2d.native == (1.0 / 27.0) * np.ones((3, 9))).all() + assert kernel_2d.native == pytest.approx((1.0 / 27.0) * np.ones((3, 9)), 1.0e-4) def test__rescaled_with_odd_dimensions_from__different_scalings(): @@ -183,7 +183,7 @@ def test__rescaled_with_odd_dimensions_from__different_scalings(): rescale_factor=2.0, normalize=True ) assert kernel_2d.pixel_scales == (0.4, 0.4) - assert (kernel_2d.native == (1.0 / 25.0) * np.ones((5, 5))).all() + assert kernel_2d.native == pytest.approx((1.0 / 25.0) * np.ones((5, 5)), 1.0e-4) kernel_2d = aa.Kernel2D.ones( shape_native=(40, 40), pixel_scales=1.0, normalize=False @@ -192,7 +192,7 @@ def test__rescaled_with_odd_dimensions_from__different_scalings(): rescale_factor=0.1, normalize=True ) assert kernel_2d.pixel_scales == (8.0, 8.0) - assert (kernel_2d.native == (1.0 / 25.0) * np.ones((5, 5))).all() + assert kernel_2d.native == pytest.approx((1.0 / 25.0) * np.ones((5, 5)), 1.0e-4) kernel_2d = aa.Kernel2D.ones(shape_native=(2, 4), pixel_scales=1.0, normalize=False) kernel_2d = kernel_2d.rescaled_with_odd_dimensions_from( @@ -201,7 +201,7 @@ def test__rescaled_with_odd_dimensions_from__different_scalings(): assert kernel_2d.pixel_scales[0] == pytest.approx(0.4, 1.0e-4) assert kernel_2d.pixel_scales[1] == pytest.approx(0.4444444, 1.0e-4) - assert (kernel_2d.native == (1.0 / 45.0) * np.ones((5, 9))).all() + assert kernel_2d.native == pytest.approx((1.0 / 45.0) * np.ones((5, 9)), 1.0e-4) kernel_2d = aa.Kernel2D.ones(shape_native=(4, 2), pixel_scales=1.0, normalize=False) kernel_2d = kernel_2d.rescaled_with_odd_dimensions_from( @@ -209,7 +209,7 @@ def test__rescaled_with_odd_dimensions_from__different_scalings(): ) assert kernel_2d.pixel_scales[0] == pytest.approx(0.4444444, 1.0e-4) assert kernel_2d.pixel_scales[1] == pytest.approx(0.4, 1.0e-4) - assert (kernel_2d.native == (1.0 / 45.0) * np.ones((9, 5))).all() + assert kernel_2d.native == pytest.approx((1.0 / 45.0) * np.ones((9, 5)), 1.0e-4) kernel_2d = aa.Kernel2D.ones(shape_native=(6, 4), pixel_scales=1.0, normalize=False) kernel_2d = kernel_2d.rescaled_with_odd_dimensions_from( @@ -217,7 +217,7 @@ def test__rescaled_with_odd_dimensions_from__different_scalings(): ) assert kernel_2d.pixel_scales == pytest.approx((2.0, 1.3333333333), 1.0e-4) - assert (kernel_2d.native == (1.0 / 9.0) * np.ones((3, 3))).all() + assert kernel_2d.native == pytest.approx((1.0 / 9.0) * np.ones((3, 3)), 1.0e-4) kernel_2d = aa.Kernel2D.ones( shape_native=(9, 12), pixel_scales=1.0, normalize=False @@ -227,7 +227,7 @@ def test__rescaled_with_odd_dimensions_from__different_scalings(): ) assert kernel_2d.pixel_scales == pytest.approx((3.0, 2.4), 1.0e-4) - assert (kernel_2d.native == (1.0 / 15.0) * np.ones((3, 5))).all() + assert kernel_2d.native == pytest.approx((1.0 / 15.0) * np.ones((3, 5)), 1.0e-4) kernel_2d = aa.Kernel2D.ones(shape_native=(4, 6), pixel_scales=1.0, normalize=False) kernel_2d = kernel_2d.rescaled_with_odd_dimensions_from( @@ -235,7 +235,7 @@ def test__rescaled_with_odd_dimensions_from__different_scalings(): ) assert kernel_2d.pixel_scales == pytest.approx((1.33333333333, 2.0), 1.0e-4) - assert (kernel_2d.native == (1.0 / 9.0) * np.ones((3, 3))).all() + assert kernel_2d.native == pytest.approx((1.0 / 9.0) * np.ones((3, 3)), 1.0e-4) kernel_2d = aa.Kernel2D.ones( shape_native=(12, 9), pixel_scales=1.0, normalize=False @@ -244,7 +244,7 @@ def test__rescaled_with_odd_dimensions_from__different_scalings(): rescale_factor=0.33333333333, normalize=True ) assert kernel_2d.pixel_scales == pytest.approx((2.4, 3.0), 1.0e-4) - assert (kernel_2d.native == (1.0 / 15.0) * np.ones((5, 3))).all() + assert kernel_2d.native == pytest.approx((1.0 / 15.0) * np.ones((5, 3)), 1.0e-4) def test__from_as_gaussian_via_alma_fits_header_parameters__identical_to_astropy_gaussian_model(): diff --git a/test_autoarray/structures/arrays/test_repr.py b/test_autoarray/structures/arrays/test_repr.py index 6e806a10d..8ea8c6b9d 100644 --- a/test_autoarray/structures/arrays/test_repr.py +++ b/test_autoarray/structures/arrays/test_repr.py @@ -3,4 +3,4 @@ def test_repr(): array = aa.Array2D.no_mask([[1, 2], [3, 4]], pixel_scales=1) - assert repr(array) == "Array([1, 2, 3, 4], dtype=int64)" + assert repr(array) == "Array2D([1, 2, 3, 4])" From 578bc4c7122a80d2867008843da84db29a99bdde Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 16:18:02 +0100 Subject: [PATCH 49/62] array 1d now use convert rules --- autoarray/structures/arrays/array_1d_util.py | 17 ++++++++++------- autoarray/structures/arrays/array_2d_util.py | 17 +++++++---------- 2 files changed, 17 insertions(+), 17 deletions(-) diff --git a/autoarray/structures/arrays/array_1d_util.py b/autoarray/structures/arrays/array_1d_util.py index c8f519ae2..bc3b8819a 100644 --- a/autoarray/structures/arrays/array_1d_util.py +++ b/autoarray/structures/arrays/array_1d_util.py @@ -40,22 +40,25 @@ def convert_array_1d( """ array_1d = array_2d_util.convert_array(array=array_1d) + is_numpy = True if isinstance(array_1d, np.ndarray) else False + is_native = array_1d.shape[0] == mask_1d.shape_native[0] mask_1d = jnp.array(mask_1d.array) if is_native == store_native: - return array_1d + array_1d = array_1d elif not store_native: - return array_1d_slim_from( + array_1d = array_1d_slim_from( array_1d_native=array_1d, mask_1d=mask_1d, ) - - return array_1d_native_from( - array_1d_slim=array_1d, - mask_1d=mask_1d, - ) + else: + array_1d = array_1d_native_from( + array_1d_slim=array_1d, + mask_1d=mask_1d, + ) + return np.array(array_1d) if is_numpy else jnp.array(array_1d) def array_1d_slim_from( diff --git a/autoarray/structures/arrays/array_2d_util.py b/autoarray/structures/arrays/array_2d_util.py index eadbd1b6f..880421e72 100644 --- a/autoarray/structures/arrays/array_2d_util.py +++ b/autoarray/structures/arrays/array_2d_util.py @@ -122,10 +122,7 @@ def convert_array_2d( """ array_2d = convert_array(array=array_2d).copy() - if isinstance(array_2d, np.ndarray): - is_numpy = True - else: - is_numpy = False + is_numpy = True if isinstance(array_2d, np.ndarray) else False check_array_2d_and_mask_2d(array_2d=array_2d, mask_2d=mask_2d) @@ -135,17 +132,17 @@ def convert_array_2d( array_2d *= np.invert(mask_2d) if is_native == store_native: - return np.array(array_2d) if is_numpy else jnp.array(array_2d) + array_2d = array_2d elif not store_native: array_2d = array_2d_slim_from( array_2d_native=array_2d, mask_2d=mask_2d, ) - return np.array(array_2d) if is_numpy else jnp.array(array_2d) - array_2d = array_2d_native_from( - array_2d_slim=array_2d, - mask_2d=mask_2d, - ) + else: + array_2d = array_2d_native_from( + array_2d_slim=array_2d, + mask_2d=mask_2d, + ) return np.array(array_2d) if is_numpy else jnp.array(array_2d) From 32f55f7bee0b7fe79b9377eb65c3b3769a1b3389 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 16:22:35 +0100 Subject: [PATCH 50/62] grid 2d casting --- autoarray/structures/grids/grid_2d_util.py | 32 ++++++++++++++-------- 1 file changed, 20 insertions(+), 12 deletions(-) diff --git a/autoarray/structures/grids/grid_2d_util.py b/autoarray/structures/grids/grid_2d_util.py index 4dfe97ed7..511495b60 100644 --- a/autoarray/structures/grids/grid_2d_util.py +++ b/autoarray/structures/grids/grid_2d_util.py @@ -21,7 +21,10 @@ def convert_grid(grid: Union[np.ndarray, List]) -> np.ndarray: except AttributeError: pass - return jnp.asarray(grid) + if isinstance(grid, list): + grid = np.asarray(grid) + + return grid def check_grid_slim(grid, shape_native): @@ -109,28 +112,33 @@ def convert_grid_2d( grid_2d = convert_grid(grid=grid_2d) + is_numpy = True if isinstance(grid_2d, np.ndarray) else False + check_grid_2d_and_mask_2d(grid_2d=grid_2d, mask_2d=mask_2d) is_native = len(grid_2d.shape) == 3 - mask_2d = jnp.array(mask_2d.array) - if is_native: - grid_2d = grid_2d.at[:, :, 0].multiply(jnp.invert(mask_2d)) - grid_2d = grid_2d.at[:, :, 1].multiply(jnp.invert(mask_2d)) + if not is_numpy: + grid_2d = grid_2d.at[:, :, 0].multiply(jnp.invert(mask_2d)) + grid_2d = grid_2d.at[:, :, 1].multiply(jnp.invert(mask_2d)) + else: + grid_2d[:, :, 0] *= np.invert(mask_2d) + grid_2d[:, :, 1] *= np.invert(mask_2d) if is_native == store_native: - return grid_2d + grid_2d = grid_2d elif not store_native: - return grid_2d_slim_from( + grid_2d = grid_2d_slim_from( grid_2d_native=grid_2d, mask=mask_2d, ) - return grid_2d_native_from( - grid_2d_slim=grid_2d, - mask_2d=mask_2d, - ) - + else: + grid_2d = grid_2d_native_from( + grid_2d_slim=grid_2d, + mask_2d=mask_2d, + ) + return np.array(grid_2d) if is_numpy else jnp.array(grid_2d) def convert_grid_2d_to_slim( grid_2d: Union[np.ndarray, List], mask_2d: Mask2D From b5601e82d0c13e03da7b08175b0a351850149c7b Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 16:24:41 +0100 Subject: [PATCH 51/62] grid 2d follows rules --- autoarray/structures/grids/uniform_2d.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/autoarray/structures/grids/uniform_2d.py b/autoarray/structures/grids/uniform_2d.py index bea896b01..52e01e557 100644 --- a/autoarray/structures/grids/uniform_2d.py +++ b/autoarray/structures/grids/uniform_2d.py @@ -862,7 +862,7 @@ def distances_to_coordinate_from( squared_distance = self.squared_distances_to_coordinate_from( coordinate=coordinate ) - distances = np.sqrt(squared_distance.array) + distances = jnp.sqrt(squared_distance.array) return Array2D(values=distances, mask=self.mask) def grid_2d_radial_projected_shape_slim_from( From 49b8dc50472128e29a58982972195ab8f7d0c120 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 16:25:51 +0100 Subject: [PATCH 52/62] grid 1d no conversion --- autoarray/structures/grids/grid_1d_util.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/autoarray/structures/grids/grid_1d_util.py b/autoarray/structures/grids/grid_1d_util.py index de1c0c725..ad6b093ce 100644 --- a/autoarray/structures/grids/grid_1d_util.py +++ b/autoarray/structures/grids/grid_1d_util.py @@ -40,21 +40,23 @@ def convert_grid_1d( grid_1d = grid_2d_util.convert_grid(grid=grid_1d) - is_native = grid_1d.shape[0] == mask_1d.shape_native[0] + is_numpy = True if isinstance(grid_1d, np.ndarray) else False - mask_1d = jnp.array(mask_1d.array) + is_native = grid_1d.shape[0] == mask_1d.shape_native[0] if is_native == store_native: - return grid_1d + grid_1d = grid_1d elif not store_native: - return grid_1d_slim_from( + grid_1d = grid_1d_slim_from( grid_1d_native=grid_1d, mask_1d=mask_1d, ) - return grid_1d_native_from( - grid_1d_slim=grid_1d, - mask_1d=mask_1d, - ) + else: + grid_1d = grid_1d_native_from( + grid_1d_slim=grid_1d, + mask_1d=mask_1d, + ) + return np.array(grid_1d) if is_numpy else jnp.array(grid_1d) def grid_1d_slim_via_shape_slim_from( From d3c5bbfec376225a13de0088c1b520bff58f2071 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 16:30:03 +0100 Subject: [PATCH 53/62] same rules for grids --- autoarray/structures/arrays/array_1d_util.py | 4 +--- autoarray/structures/arrays/array_2d_util.py | 2 +- autoarray/structures/grids/grid_1d_util.py | 2 +- autoarray/structures/grids/grid_2d_util.py | 18 +++++++++--------- 4 files changed, 12 insertions(+), 14 deletions(-) diff --git a/autoarray/structures/arrays/array_1d_util.py b/autoarray/structures/arrays/array_1d_util.py index bc3b8819a..f7c943724 100644 --- a/autoarray/structures/arrays/array_1d_util.py +++ b/autoarray/structures/arrays/array_1d_util.py @@ -44,8 +44,6 @@ def convert_array_1d( is_native = array_1d.shape[0] == mask_1d.shape_native[0] - mask_1d = jnp.array(mask_1d.array) - if is_native == store_native: array_1d = array_1d elif not store_native: @@ -124,7 +122,7 @@ def array_1d_native_from( ).astype("int") return array_1d_via_indexes_1d_from( - array_1d_slim=np.array(array_1d_slim), + array_1d_slim=array_1d_slim, shape=shape, native_index_for_slim_index_1d=native_index_for_slim_index_1d, ) diff --git a/autoarray/structures/arrays/array_2d_util.py b/autoarray/structures/arrays/array_2d_util.py index 880421e72..17c27a29f 100644 --- a/autoarray/structures/arrays/array_2d_util.py +++ b/autoarray/structures/arrays/array_2d_util.py @@ -129,7 +129,7 @@ def convert_array_2d( is_native = len(array_2d.shape) == 2 if is_native and not skip_mask: - array_2d *= np.invert(mask_2d) + array_2d *= ~mask_2d if is_native == store_native: array_2d = array_2d diff --git a/autoarray/structures/grids/grid_1d_util.py b/autoarray/structures/grids/grid_1d_util.py index ad6b093ce..82aa4514e 100644 --- a/autoarray/structures/grids/grid_1d_util.py +++ b/autoarray/structures/grids/grid_1d_util.py @@ -171,7 +171,7 @@ def grid_1d_slim_from( """ return array_1d_util.array_1d_slim_from( - array_1d_native=np.array(grid_1d_native), + array_1d_native=grid_1d_native, mask_1d=mask_1d, ) diff --git a/autoarray/structures/grids/grid_2d_util.py b/autoarray/structures/grids/grid_2d_util.py index 511495b60..33ad88eb6 100644 --- a/autoarray/structures/grids/grid_2d_util.py +++ b/autoarray/structures/grids/grid_2d_util.py @@ -120,11 +120,11 @@ def convert_grid_2d( if is_native: if not is_numpy: - grid_2d = grid_2d.at[:, :, 0].multiply(jnp.invert(mask_2d)) - grid_2d = grid_2d.at[:, :, 1].multiply(jnp.invert(mask_2d)) + grid_2d = grid_2d.at[:, :, 0].multiply(~mask_2d) + grid_2d = grid_2d.at[:, :, 1].multiply(~mask_2d) else: - grid_2d[:, :, 0] *= np.invert(mask_2d) - grid_2d[:, :, 1] *= np.invert(mask_2d) + grid_2d[:, :, 0] *= ~mask_2d + grid_2d[:, :, 1] *= ~mask_2d if is_native == store_native: grid_2d = grid_2d @@ -682,16 +682,16 @@ def grid_2d_slim_from( """ grid_1d_slim_y = array_2d_util.array_2d_slim_from( - array_2d_native=np.array(grid_2d_native[:, :, 0]), - mask_2d=np.array(mask), + array_2d_native=grid_2d_native[:, :, 0], + mask_2d=mask, ) grid_1d_slim_x = array_2d_util.array_2d_slim_from( - array_2d_native=np.array(grid_2d_native[:, :, 1]), - mask_2d=np.array(mask), + array_2d_native=grid_2d_native[:, :, 1], + mask_2d=mask, ) - return np.stack((grid_1d_slim_y, grid_1d_slim_x), axis=-1) + return jnp.stack((grid_1d_slim_y, grid_1d_slim_x), axis=-1) def grid_2d_native_from( From 61a72b8ee976f042d65c5fea64ef84d38a24d550 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 16:43:23 +0100 Subject: [PATCH 54/62] fix dataset unit tests --- autoarray/dataset/imaging/dataset.py | 2 -- autoarray/structures/grids/uniform_2d.py | 8 ++++---- .../dataset/abstract/test_dataset.py | 12 ++++-------- .../dataset/imaging/test_dataset.py | 4 ++-- test_autoarray/dataset/test_preprocess.py | 7 ++----- .../arrays/files/array/output_test/array.fits | Bin 5760 -> 5760 bytes 6 files changed, 12 insertions(+), 21 deletions(-) diff --git a/autoarray/dataset/imaging/dataset.py b/autoarray/dataset/imaging/dataset.py index 7380cdcba..9fa00da1a 100644 --- a/autoarray/dataset/imaging/dataset.py +++ b/autoarray/dataset/imaging/dataset.py @@ -135,8 +135,6 @@ def __init__( f"the image and noise-map yourself." ) - print(type(data.array)) - super().__init__( data=data, noise_map=noise_map, diff --git a/autoarray/structures/grids/uniform_2d.py b/autoarray/structures/grids/uniform_2d.py index 52e01e557..3a2e4a068 100644 --- a/autoarray/structures/grids/uniform_2d.py +++ b/autoarray/structures/grids/uniform_2d.py @@ -320,7 +320,7 @@ def from_yx_1d( x = np.asarray(x) return cls.no_mask( - values=jnp.stack((y, x), axis=-1), + values=np.stack((y, x), axis=-1), shape_native=shape_native, pixel_scales=pixel_scales, origin=origin, @@ -547,14 +547,14 @@ def from_mask( The mask whose masked pixels are used to setup the grid. """ - grid_1d = grid_2d_util.grid_2d_slim_via_mask_from( - mask_2d=mask._array, + grid_2d = grid_2d_util.grid_2d_slim_via_mask_from( + mask_2d=mask, pixel_scales=mask.pixel_scales, origin=mask.origin, ) return Grid2D( - values=grid_1d, + values=np.array(grid_2d), mask=mask, over_sample_size=over_sample_size, ) diff --git a/test_autoarray/dataset/abstract/test_dataset.py b/test_autoarray/dataset/abstract/test_dataset.py index 1bd36a2fe..770d9442b 100644 --- a/test_autoarray/dataset/abstract/test_dataset.py +++ b/test_autoarray/dataset/abstract/test_dataset.py @@ -32,9 +32,7 @@ def test__signal_to_noise_map(): dataset = ds.AbstractDataset(data=array, noise_map=noise_map) - assert ( - dataset.signal_to_noise_map.native == np.array([[0.1, 0.2], [0.1, 1.0]]) - ).all() + assert dataset.signal_to_noise_map.native == pytest.approx(np.array([[0.1, 0.2], [0.1, 1.0]]), 1.0e-4) assert dataset.signal_to_noise_max == 1.0 array = aa.Array2D.no_mask([[-1.0, 2.0], [3.0, -4.0]], pixel_scales=1.0) @@ -43,9 +41,7 @@ def test__signal_to_noise_map(): dataset = ds.AbstractDataset(data=array, noise_map=noise_map) - assert ( - dataset.signal_to_noise_map.native == np.array([[0.0, 0.2], [0.1, 0.0]]) - ).all() + assert dataset.signal_to_noise_map.native == pytest.approx(np.array([[0.0, 0.2], [0.1, 0.0]]), 1.0e-4) assert dataset.signal_to_noise_max == 0.2 @@ -149,8 +145,8 @@ def test__apply_over_sampling(image_7x7, noise_map_7x7): grid_sub_2 = dataset_7x7.grids.lp grids_pixelization_sub_2 = dataset_7x7.grids.pixelization - dataset_7x7.grids.lp = dataset_7x7.grids.lp.at[0, 0].set(100.0) - dataset_7x7.grids.pixelization = dataset_7x7.grids.pixelization.at[0, 0].set(100.0) + dataset_7x7.grids.__dict__["lp"][0][0] = 100.0 + dataset_7x7.grids.__dict__["pixelization"][0][0] = 100.0 assert dataset_7x7.grids.lp[0][0] == pytest.approx(100.0, 1.0e-4) assert dataset_7x7.grids.pixelization[0][0] == pytest.approx(100.0, 1.0e-4) diff --git a/test_autoarray/dataset/imaging/test_dataset.py b/test_autoarray/dataset/imaging/test_dataset.py index f2f2b6d27..1665865f2 100644 --- a/test_autoarray/dataset/imaging/test_dataset.py +++ b/test_autoarray/dataset/imaging/test_dataset.py @@ -78,7 +78,7 @@ def test__from_fits(): ) assert (dataset.data.native == np.ones((3, 3))).all() - assert (dataset.psf.native == (1.0 / 9.0) * np.ones((3, 3))).all() + assert dataset.psf.native == pytest.approx((1.0 / 9.0) * np.ones((3, 3)), 1.0e-4) assert (dataset.noise_map.native == 3.0 * np.ones((3, 3))).all() assert dataset.pixel_scales == (0.1, 0.1) @@ -96,7 +96,7 @@ def test__from_fits(): ) assert (dataset.data.native == np.ones((3, 3))).all() - assert (dataset.psf.native == (1.0 / 9.0) * np.ones((3, 3))).all() + assert dataset.psf.native == pytest.approx((1.0 / 9.0) * np.ones((3, 3)), 1.0e-4) assert (dataset.noise_map.native == 3.0 * np.ones((3, 3))).all() assert dataset.pixel_scales == (0.1, 0.1) diff --git a/test_autoarray/dataset/test_preprocess.py b/test_autoarray/dataset/test_preprocess.py index 108253a08..67107be4a 100644 --- a/test_autoarray/dataset/test_preprocess.py +++ b/test_autoarray/dataset/test_preprocess.py @@ -133,15 +133,12 @@ def test__noise_map_from_image_exposure_time_map(): data_eps=image, exposure_time_map=exposure_time_map ) - assert ( - poisson_noise_map.native - == np.array( + assert poisson_noise_map.native == pytest.approx(np.array( [ [np.sqrt(5.0), np.sqrt(6.0) / 2.0], [np.sqrt(30.0) / 3.0, np.sqrt(80.0) / 4.0], ] - ) - ).all() + ), 1.0e-4) def test__noise_map_from_image_exposure_time_map_and_background_noise_map(): diff --git a/test_autoarray/structures/arrays/files/array/output_test/array.fits b/test_autoarray/structures/arrays/files/array/output_test/array.fits index dab9da2fb676efc023e53b73faa54fe78fc4000c..0cff0a8f7db90d3ded28c644cd66e3791627feeb 100644 GIT binary patch delta 84 hcmZqBZP1;N!(?W%G4C>$;|B&XuqT_|T*&>83jh#~5uX46 delta 56 kcmZqBZP1;N!(?o$Vgmz%Jzl(dBKLc)$qTsk0j$6e2mk;8 From 15cd0b5a3ef7b266d7f7369a84bb69329b9855c3 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 16:59:04 +0100 Subject: [PATCH 55/62] fix some unit tests --- autoarray/mask/derive/grid_2d.py | 8 ++++---- autoarray/operators/over_sampling/over_sample_util.py | 5 ++--- autoarray/structures/arrays/uniform_1d.py | 2 +- autoarray/structures/arrays/uniform_2d.py | 2 +- autoarray/structures/grids/uniform_1d.py | 6 +++--- autoarray/structures/grids/uniform_2d.py | 4 ++-- autoarray/structures/vectors/irregular.py | 3 +++ 7 files changed, 16 insertions(+), 14 deletions(-) diff --git a/autoarray/mask/derive/grid_2d.py b/autoarray/mask/derive/grid_2d.py index 225c9a43b..068e40a1b 100644 --- a/autoarray/mask/derive/grid_2d.py +++ b/autoarray/mask/derive/grid_2d.py @@ -112,7 +112,7 @@ def all_false(self) -> Grid2D: origin=self.mask.origin, ) - return Grid2D(values=grid_slim, mask=self.mask.derive_mask.all_false) + return Grid2D(values=np.array(grid_slim), mask=self.mask.derive_mask.all_false) @property def unmasked(self) -> Grid2D: @@ -158,12 +158,12 @@ def unmasked(self) -> Grid2D: """ from autoarray.structures.grids.uniform_2d import Grid2D - grid_1d = grid_2d_util.grid_2d_slim_via_mask_from( - mask_2d=np.array(self.mask), + grid_2d = grid_2d_util.grid_2d_slim_via_mask_from( + mask_2d=self.mask, pixel_scales=self.mask.pixel_scales, origin=self.mask.origin, ) - return Grid2D(values=grid_1d, mask=self.mask) + return Grid2D(values=np.array(grid_2d), mask=self.mask) @property def edge(self) -> Grid2D: diff --git a/autoarray/operators/over_sampling/over_sample_util.py b/autoarray/operators/over_sampling/over_sample_util.py index 084d0bd0d..2a1a247c6 100644 --- a/autoarray/operators/over_sampling/over_sample_util.py +++ b/autoarray/operators/over_sampling/over_sample_util.py @@ -1,6 +1,6 @@ from __future__ import annotations import numpy as np -from typing import TYPE_CHECKING, Union, List, Tuple +from typing import TYPE_CHECKING, Union from typing import List, Tuple from autoarray.structures.arrays.uniform_2d import Array2D @@ -12,7 +12,6 @@ from autoarray.mask.mask_2d import Mask2D from autoarray import numba_util -from autoarray.mask import mask_2d_util from autoarray import type as ty @@ -601,7 +600,7 @@ def over_sample_size_via_radial_bins_from( radial_grid = grid.distances_to_coordinate_from(coordinate=centre) sub_size_of_centre = sub_size_radial_bins_from( - radial_grid=np.array(radial_grid), + radial_grid=np.array(radial_grid.array), sub_size_list=np.array(sub_size_list), radial_list=np.array(radial_list), ) diff --git a/autoarray/structures/arrays/uniform_1d.py b/autoarray/structures/arrays/uniform_1d.py index 9b0af1af2..d708a68ad 100644 --- a/autoarray/structures/arrays/uniform_1d.py +++ b/autoarray/structures/arrays/uniform_1d.py @@ -88,7 +88,7 @@ def no_mask( origin=origin, ) - return Array1D(values=values, mask=mask, header=header) + return Array1D(values=np.array(values), mask=mask, header=header) @classmethod def full( diff --git a/autoarray/structures/arrays/uniform_2d.py b/autoarray/structures/arrays/uniform_2d.py index 574d2fa5d..66c881ff3 100644 --- a/autoarray/structures/arrays/uniform_2d.py +++ b/autoarray/structures/arrays/uniform_2d.py @@ -692,7 +692,7 @@ def no_mask( origin=origin, ) - return Array2D(values=values, mask=mask, header=header) + return Array2D(values=np.array(values), mask=mask, header=header) @classmethod def full( diff --git a/autoarray/structures/grids/uniform_1d.py b/autoarray/structures/grids/uniform_1d.py index 0a736e324..49b748b2e 100644 --- a/autoarray/structures/grids/uniform_1d.py +++ b/autoarray/structures/grids/uniform_1d.py @@ -179,7 +179,7 @@ def no_mask( origin=origin, ) - return Grid1D(values=values, mask=mask) + return Grid1D(values=np.array(values), mask=mask) @classmethod def from_mask(cls, mask: Mask1D) -> "Grid1D": @@ -196,12 +196,12 @@ def from_mask(cls, mask: Mask1D) -> "Grid1D": """ grid_1d = grid_1d_util.grid_1d_slim_via_mask_from( - mask_1d=np.array(mask), + mask_1d=mask.array, pixel_scales=mask.pixel_scales, origin=mask.origin, ) - return Grid1D(values=grid_1d, mask=mask) + return Grid1D(values=np.array(grid_1d), mask=mask) @classmethod def uniform( diff --git a/autoarray/structures/grids/uniform_2d.py b/autoarray/structures/grids/uniform_2d.py index 3a2e4a068..10262dd4a 100644 --- a/autoarray/structures/grids/uniform_2d.py +++ b/autoarray/structures/grids/uniform_2d.py @@ -247,7 +247,7 @@ def no_mask( ) return Grid2D( - values=values, + values=np.array(values), mask=mask, over_sample_size=over_sample_size, ) @@ -548,7 +548,7 @@ def from_mask( """ grid_2d = grid_2d_util.grid_2d_slim_via_mask_from( - mask_2d=mask, + mask_2d=mask.array, pixel_scales=mask.pixel_scales, origin=mask.origin, ) diff --git a/autoarray/structures/vectors/irregular.py b/autoarray/structures/vectors/irregular.py index e6501e46a..b49ec3868 100644 --- a/autoarray/structures/vectors/irregular.py +++ b/autoarray/structures/vectors/irregular.py @@ -44,6 +44,9 @@ def __init__( grid The irregular grid of (y,x) coordinates where each vector is located. """ + if type(values) is list: + values = np.asarray(values) + self.grid = Grid2DIrregular(values=grid) super().__init__(values) From d33ab4454ff187063308b221057b486aee58bda5 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 17:14:49 +0100 Subject: [PATCH 56/62] mask now also follows strictire typing rules --- autoarray/dataset/interferometer/dataset.py | 2 +- autoarray/geometry/geometry_util.py | 2 -- autoarray/mask/mask_1d.py | 4 ++-- autoarray/mask/mask_2d.py | 21 +++++++++------------ autoarray/plot/visuals/two_d.py | 2 +- autoarray/plot/wrap/base/figure.py | 10 ++++++---- test_autoarray/layout/test_region.py | 16 ++++++++-------- 7 files changed, 27 insertions(+), 30 deletions(-) diff --git a/autoarray/dataset/interferometer/dataset.py b/autoarray/dataset/interferometer/dataset.py index 32a9a4377..4cf7220c2 100644 --- a/autoarray/dataset/interferometer/dataset.py +++ b/autoarray/dataset/interferometer/dataset.py @@ -205,7 +205,7 @@ def w_tilde(self): return WTildeInterferometer( w_matrix=w_matrix, curvature_preload=curvature_preload, - dirty_image=dirty_image, + dirty_image=np.array(dirty_image.array), real_space_mask=self.real_space_mask, noise_map_value=self.noise_map[0], ) diff --git a/autoarray/geometry/geometry_util.py b/autoarray/geometry/geometry_util.py index e92eea131..7179fb879 100644 --- a/autoarray/geometry/geometry_util.py +++ b/autoarray/geometry/geometry_util.py @@ -119,7 +119,6 @@ def central_scaled_coordinate_1d_from( return (x_pixel,) -@numba_util.jit() def pixel_coordinates_1d_from( scaled_coordinates_1d: Tuple[float], shape_slim: Tuple[int], @@ -137,7 +136,6 @@ def pixel_coordinates_1d_from( return (x_pixel,) -@numba_util.jit() def scaled_coordinates_1d_from( pixel_coordinates_1d: Tuple[float], shape_slim: Tuple[int], diff --git a/autoarray/mask/mask_1d.py b/autoarray/mask/mask_1d.py index 8c36d8866..6ceacae7f 100644 --- a/autoarray/mask/mask_1d.py +++ b/autoarray/mask/mask_1d.py @@ -59,7 +59,7 @@ def __init__( mask = np.asarray(mask).astype("bool") if invert: - mask = np.invert(mask) + mask = ~mask if type(pixel_scales) is float: pixel_scales = (pixel_scales,) @@ -153,7 +153,7 @@ def from_fits( """ return cls( - array_1d_util.numpy_array_1d_via_fits_from(file_path=file_path, hdu=hdu), + mask=np.array(array_1d_util.numpy_array_1d_via_fits_from(file_path=file_path, hdu=hdu)), pixel_scales=pixel_scales, origin=origin, ) diff --git a/autoarray/mask/mask_2d.py b/autoarray/mask/mask_2d.py index 05bd77852..5c65f8972 100644 --- a/autoarray/mask/mask_2d.py +++ b/autoarray/mask/mask_2d.py @@ -200,11 +200,8 @@ def __init__( if type(mask) is list: mask = np.asarray(mask).astype("bool") - if not isinstance(mask, np.ndarray): - mask = mask._array - if invert: - mask = np.invert(mask) + mask = ~mask pixel_scales = geometry_util.convert_pixel_scales_2d(pixel_scales=pixel_scales) @@ -327,7 +324,7 @@ def circular( ) return cls( - mask=mask, + mask=np.array(mask), pixel_scales=pixel_scales, origin=origin, invert=invert, @@ -381,7 +378,7 @@ def circular_annular( ) return cls( - mask=mask, + mask=np.array(mask), pixel_scales=pixel_scales, origin=origin, invert=invert, @@ -438,7 +435,7 @@ def elliptical( ) return cls( - mask=mask, + mask=np.array(mask), pixel_scales=pixel_scales, origin=origin, invert=invert, @@ -508,7 +505,7 @@ def elliptical_annular( ) return cls( - mask=mask, + mask=np.array(mask), pixel_scales=pixel_scales, origin=origin, invert=invert, @@ -556,7 +553,7 @@ def from_pixel_coordinates( ) return cls( - mask=mask, + mask=np.array(mask), pixel_scales=pixel_scales, origin=origin, invert=invert, @@ -594,7 +591,7 @@ def from_fits( mask = np.invert(mask.astype("bool")) mask = Mask2D( - mask=mask, + mask=np.array(mask), pixel_scales=pixel_scales, origin=origin, ) @@ -743,7 +740,7 @@ def rescaled_from(self, rescale_factor) -> Mask2D: from autoarray.mask.mask_2d import Mask2D rescaled_mask = mask_2d_util.rescaled_mask_2d_from( - mask_2d=np.array(self), + mask_2d=self.array, rescale_factor=rescale_factor, ) @@ -802,7 +799,7 @@ def resized_from(self, new_shape, pad_value: int = 0.0) -> Mask2D: """ resized_mask = array_2d_util.resized_array_2d_from( - array_2d=np.array(self._array), + array_2d=self.array, resized_shape=new_shape, pad_value=pad_value, ).astype("bool") diff --git a/autoarray/plot/visuals/two_d.py b/autoarray/plot/visuals/two_d.py index 80e67fbab..eacca5edc 100644 --- a/autoarray/plot/visuals/two_d.py +++ b/autoarray/plot/visuals/two_d.py @@ -64,7 +64,7 @@ def plot_via_plotter(self, plotter, grid_indexes=None, mapper=None, geometry=Non plotter.grid_scatter.scatter_grid(grid=self.grid.array) if self.mesh_grid is not None: - plotter.mesh_grid_scatter.scatter_grid(grid=self.mesh_grid) + plotter.mesh_grid_scatter.scatter_grid(grid=self.mesh_grid.array) if self.positions is not None: plotter.positions_scatter.scatter_grid(grid=self.positions) diff --git a/autoarray/plot/wrap/base/figure.py b/autoarray/plot/wrap/base/figure.py index edea67465..238264fd9 100644 --- a/autoarray/plot/wrap/base/figure.py +++ b/autoarray/plot/wrap/base/figure.py @@ -82,10 +82,12 @@ def open(self): """ Wraps the Matplotlib method 'plt.figure' for opening a figure. """ - config_dict = self.config_dict - config_dict.pop("aspect") - fig = plt.figure(**config_dict) - return fig, plt.gca() + if not plt.fignum_exists(num=1): + config_dict = self.config_dict + config_dict.pop("aspect") + fig = plt.figure(**config_dict) + return fig, plt.gca() + return None, None def close(self): """ diff --git a/test_autoarray/layout/test_region.py b/test_autoarray/layout/test_region.py index 20809cc7d..78bfba73f 100644 --- a/test_autoarray/layout/test_region.py +++ b/test_autoarray/layout/test_region.py @@ -31,7 +31,7 @@ def test__slice_1d__addition(): region = aa.Region1D(region=(0, 1)) - arr_1d = arr_1d.array.at[region.slice].add(image[region.slice]) + arr_1d[region.slice] += image[region.slice] assert (arr_1d == np.array([2.0, 2.0, 3.0, 4.0])).all() @@ -39,7 +39,7 @@ def test__slice_1d__addition(): region = aa.Region1D(region=(2, 4)) - arr_1d = arr_1d.array.at[region.slice].add(image[region.slice]) + arr_1d[region.slice] += image[region.slice] assert (arr_1d == np.array([1.0, 2.0, 4.0, 5.0])).all() @@ -152,7 +152,7 @@ def test__slice_2d__addition(): image = np.ones((2, 2)) region = aa.Region2D(region=(0, 1, 0, 1)) - array = array.at[region.slice].add(image[region.slice]) + array[region.slice] += 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 = array.at[region.slice].add(image[region.slice]) + array[region.slice] += 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 = array.at[region.slice].add(image[region.slice]) + array[region.slice] += 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 = array.at[region.slice].set(0) + array[region.slice] = 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 = array.at[region.slice].set(0) + array[region.slice] = 0 assert ( array == np.array([[1.0, 1.0, 1.0], [1.0, 1.0, 0.0], [1.0, 1.0, 0.0]]) @@ -362,4 +362,4 @@ def test__exception__negative_coordinates_2d(): aa.Region2D(region=(0, 0, -1, 2)) with pytest.raises(exc.RegionException): - aa.Region2D(region=(0, 1, 2, -1)) + aa.Region2D(region=(0, 1, 2, -1)) \ No newline at end of file From 3aef08ac68a58652e403091461f24e1098a550e9 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 17:57:28 +0100 Subject: [PATCH 57/62] but of typing simplication --- .../over_sampling/over_sample_util.py | 21 ++----------------- .../operators/over_sampling/over_sampler.py | 12 +++++------ autoarray/structures/grids/uniform_2d.py | 2 +- 3 files changed, 9 insertions(+), 26 deletions(-) diff --git a/autoarray/operators/over_sampling/over_sample_util.py b/autoarray/operators/over_sampling/over_sample_util.py index 2a1a247c6..d3488a457 100644 --- a/autoarray/operators/over_sampling/over_sample_util.py +++ b/autoarray/operators/over_sampling/over_sample_util.py @@ -44,9 +44,9 @@ def over_sample_size_convert_to_array_2d_from( if isinstance(over_sample_size, int): over_sample_size = np.full( fill_value=over_sample_size, shape=mask.pixels_in_mask - ).astype("int") + ) - return Array2D(values=over_sample_size, mask=mask) + return Array2D(values=np.array(over_sample_size).astype("int"), mask=mask) @numba_util.jit() @@ -444,23 +444,6 @@ def grid_2d_slim_over_sampled_via_mask_from( for y1 in range(sub): for x1 in range(sub): - # if use_jax: - # # while this makes it run, it is very, very slow - # grid_slim = grid_slim.at[sub_index, 0].set( - # -( - # y_scaled - # - y_sub_half - # + y1 * y_sub_step - # + (y_sub_step / 2.0) - # ) - # ) - # grid_slim = grid_slim.at[sub_index, 1].set( - # x_scaled - # - x_sub_half - # + x1 * x_sub_step - # + (x_sub_step / 2.0) - # ) - # else: grid_slim[sub_index, 0] = -( y_scaled - y_sub_half + y1 * y_sub_step + (y_sub_step / 2.0) ) diff --git a/autoarray/operators/over_sampling/over_sampler.py b/autoarray/operators/over_sampling/over_sampler.py index ccc09f04a..73f7a1aa7 100644 --- a/autoarray/operators/over_sampling/over_sampler.py +++ b/autoarray/operators/over_sampling/over_sampler.py @@ -154,7 +154,7 @@ 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]) + np.isclose(self.sub_size, self.sub_size[0]) ) def tree_flatten(self): @@ -233,12 +233,12 @@ def binned_array_2d_from(self, array: Array2D) -> "Array2D": if self.sub_is_uniform: binned_array_2d = array.reshape( - self.mask.shape_slim, self.sub_size.array[0] ** 2 + 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) + group_sizes = jnp.array(self.sub_size.array ** 2) # Compute the cumulative sum of group sizes to get split points split_indices = jnp.cumsum(group_sizes) @@ -310,7 +310,7 @@ def sub_mask_native_for_sub_mask_slim(self) -> np.ndarray: print(derive_indexes_2d.sub_mask_native_for_sub_mask_slim) """ return over_sample_util.native_sub_index_for_slim_sub_index_2d_from( - mask_2d=self.mask.array, sub_size=np.array(self.sub_size) + mask_2d=self.mask.array, sub_size=self.sub_size.array ).astype("int") @cached_property @@ -363,7 +363,7 @@ def slim_for_sub_slim(self) -> np.ndarray: print(derive_indexes_2d.slim_for_sub_slim) """ return over_sample_util.slim_index_for_sub_slim_index_via_mask_2d_from( - mask_2d=np.array(self.mask), sub_size=np.array(self.sub_size) + mask_2d=np.array(self.mask), sub_size=self.sub_size.array ).astype("int") @property @@ -386,7 +386,7 @@ def uniform_over_sampled(self): grid = over_sample_util.grid_2d_slim_over_sampled_via_mask_from( mask_2d=np.array(self.mask), pixel_scales=self.mask.pixel_scales, - sub_size=np.array(self.sub_size).astype("int"), + sub_size=self.sub_size.array.astype("int"), origin=self.mask.origin, ) diff --git a/autoarray/structures/grids/uniform_2d.py b/autoarray/structures/grids/uniform_2d.py index 10262dd4a..e8a6a530e 100644 --- a/autoarray/structures/grids/uniform_2d.py +++ b/autoarray/structures/grids/uniform_2d.py @@ -185,7 +185,7 @@ def __init__( over_sample_util.grid_2d_slim_over_sampled_via_mask_from( mask_2d=np.array(self.mask), pixel_scales=self.mask.pixel_scales, - sub_size=np.array(self.over_sampler.sub_size._array).astype("int"), + sub_size=self.over_sampler.sub_size.array.astype("int"), origin=self.mask.origin, ) ) From 5d96ba556550654d4317e8ac34f2c1996de63b01 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 6 Apr 2025 18:01:52 +0100 Subject: [PATCH 58/62] check conversion needed --- autoarray/structures/arrays/irregular.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/autoarray/structures/arrays/irregular.py b/autoarray/structures/arrays/irregular.py index 667ddfd43..f4b54ee11 100644 --- a/autoarray/structures/arrays/irregular.py +++ b/autoarray/structures/arrays/irregular.py @@ -38,12 +38,6 @@ def __init__(self, values: Union[List, np.ndarray]): A collection of values. """ - # if len(values) == 0: - # return [] - - # if isinstance(values, ArrayIrregular): - # return values - if type(values) is list: values = np.asarray(values) From ccccf1a162f52becfb681b5831f71b510ef1b554 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Mon, 7 Apr 2025 20:22:31 +0100 Subject: [PATCH 59/62] add small non zero numerical value to make grid radials plot correctly --- autoarray/dataset/imaging/simulator.py | 2 +- autoarray/dataset/preprocess.py | 2 +- autoarray/structures/arrays/kernel_2d.py | 2 +- autoarray/structures/grids/uniform_1d.py | 2 +- autoarray/structures/grids/uniform_2d.py | 5 +++-- 5 files changed, 7 insertions(+), 6 deletions(-) diff --git a/autoarray/dataset/imaging/simulator.py b/autoarray/dataset/imaging/simulator.py index b1bb23198..972e03931 100644 --- a/autoarray/dataset/imaging/simulator.py +++ b/autoarray/dataset/imaging/simulator.py @@ -161,7 +161,7 @@ def via_image_from( image = image - background_sky_map mask = Mask2D.all_false( - shape_native=image.shape_native, pixel_scales=image.pixel_scales + shape_native=image.shape_native, pixel_scales=image.pixel_scales, origin=image.origin ) image = Array2D(values=image, mask=mask) diff --git a/autoarray/dataset/preprocess.py b/autoarray/dataset/preprocess.py index f411b8302..5eb80288f 100644 --- a/autoarray/dataset/preprocess.py +++ b/autoarray/dataset/preprocess.py @@ -521,7 +521,7 @@ def noise_map_with_signal_to_noise_limit_from( ) mask = Mask2D.all_false( - shape_native=data.shape_native, pixel_scales=data.pixel_scales + shape_native=data.shape_native, pixel_scales=data.pixel_scales, origin=data.origin ) if len(noise_map.native) == 1: diff --git a/autoarray/structures/arrays/kernel_2d.py b/autoarray/structures/arrays/kernel_2d.py index b436a4077..9d04c29ed 100644 --- a/autoarray/structures/arrays/kernel_2d.py +++ b/autoarray/structures/arrays/kernel_2d.py @@ -475,7 +475,7 @@ def convolved_array_from(self, array: Array2D) -> Array2D: array_2d = array.native convolved_array_2d = scipy.signal.convolve2d( - array_2d._array, np.array(self.native._array), mode="same" + array_2d.array, np.array(self.native.array), mode="same" ) convolved_array_1d = array_2d_util.array_2d_slim_from( diff --git a/autoarray/structures/grids/uniform_1d.py b/autoarray/structures/grids/uniform_1d.py index 49b748b2e..7fb2e7688 100644 --- a/autoarray/structures/grids/uniform_1d.py +++ b/autoarray/structures/grids/uniform_1d.py @@ -319,4 +319,4 @@ def grid_2d_radial_projected_from(self, angle: float = 0.0) -> Grid2DIrregular: grid_2d=grid, centre=(0.0, 0.0), angle=angle ) - return Grid2DIrregular(values=grid) + return Grid2DIrregular(values=grid + 1e-6) diff --git a/autoarray/structures/grids/uniform_2d.py b/autoarray/structures/grids/uniform_2d.py index e8a6a530e..e6298f0d0 100644 --- a/autoarray/structures/grids/uniform_2d.py +++ b/autoarray/structures/grids/uniform_2d.py @@ -1000,7 +1000,7 @@ def grid_2d_radial_projected_from( if remove_projected_centre: grid_radial_projected_2d = grid_radial_projected_2d[1:, :] - return Grid2DIrregular(values=grid_radial_projected_2d) + return Grid2DIrregular(values=grid_radial_projected_2d + 1e-6) @property def shape_native_scaled_interior(self) -> Tuple[float, float]: @@ -1098,6 +1098,7 @@ def padded_grid_from(self, kernel_shape_native: Tuple[int, int]) -> "Grid2D": padded_mask = Mask2D.all_false( shape_native=padded_shape, pixel_scales=self.mask.pixel_scales, + origin=self.origin ) pad_width = ( @@ -1106,7 +1107,7 @@ def padded_grid_from(self, kernel_shape_native: Tuple[int, int]) -> "Grid2D": ) over_sample_size = np.pad( - self.over_sample_size.native._array, + self.over_sample_size.native.array, pad_width, mode="constant", constant_values=1, From 7948e9089521a14759b98db281e57bfe5c4e9d24 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Mon, 7 Apr 2025 21:33:45 +0100 Subject: [PATCH 60/62] small shift in projected grid --- autoarray/inversion/convolver.py | 4 ++-- autoarray/inversion/inversion/imaging/mapping.py | 4 ++-- autoarray/inversion/inversion/imaging/w_tilde.py | 4 ++-- autoarray/structures/grids/grid_2d_util.py | 2 +- autoarray/structures/grids/uniform_2d.py | 2 +- 5 files changed, 8 insertions(+), 8 deletions(-) diff --git a/autoarray/inversion/convolver.py b/autoarray/inversion/convolver.py index 560009f13..80a131bb0 100644 --- a/autoarray/inversion/convolver.py +++ b/autoarray/inversion/convolver.py @@ -328,11 +328,11 @@ def convolve_image(self, image, blurring_image): ) convolved_image = self.convolve_jit( - image_1d_array=np.array(image.slim), + image_1d_array=np.array(image.slim.array), 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_1d_array=np.array(blurring_image.slim.array), 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, diff --git a/autoarray/inversion/inversion/imaging/mapping.py b/autoarray/inversion/inversion/imaging/mapping.py index 053f1d01a..12881a944 100644 --- a/autoarray/inversion/inversion/imaging/mapping.py +++ b/autoarray/inversion/inversion/imaging/mapping.py @@ -112,8 +112,8 @@ def data_vector(self) -> np.ndarray: return inversion_imaging_util.data_vector_via_blurred_mapping_matrix_from( blurred_mapping_matrix=self.operated_mapping_matrix, - image=np.array(self.data), - noise_map=np.array(self.noise_map), + image=self.data.array, + noise_map=self.noise_map.array, ) @property diff --git a/autoarray/inversion/inversion/imaging/w_tilde.py b/autoarray/inversion/inversion/imaging/w_tilde.py index 21834cdd0..f98a95f51 100644 --- a/autoarray/inversion/inversion/imaging/w_tilde.py +++ b/autoarray/inversion/inversion/imaging/w_tilde.py @@ -207,8 +207,8 @@ def _data_vector_func_list_and_mapper(self) -> np.ndarray: diag = inversion_imaging_util.data_vector_via_blurred_mapping_matrix_from( blurred_mapping_matrix=operated_mapping_matrix, - image=np.array(self.data), - noise_map=np.array(self.noise_map), + image=self.data.array, + noise_map=self.noise_map.array, ) param_range = linear_func_param_range[linear_func_index] diff --git a/autoarray/structures/grids/grid_2d_util.py b/autoarray/structures/grids/grid_2d_util.py index 33ad88eb6..cec240ad4 100644 --- a/autoarray/structures/grids/grid_2d_util.py +++ b/autoarray/structures/grids/grid_2d_util.py @@ -564,7 +564,7 @@ def grid_scaled_2d_slim_radial_projected_from( grid_scaled_2d_slim_radii[slim_index, 1] = radii radii += pixel_scale - return grid_scaled_2d_slim_radii + return grid_scaled_2d_slim_radii + 1e-6 @numba_util.jit() diff --git a/autoarray/structures/grids/uniform_2d.py b/autoarray/structures/grids/uniform_2d.py index e6298f0d0..890301884 100644 --- a/autoarray/structures/grids/uniform_2d.py +++ b/autoarray/structures/grids/uniform_2d.py @@ -1000,7 +1000,7 @@ def grid_2d_radial_projected_from( if remove_projected_centre: grid_radial_projected_2d = grid_radial_projected_2d[1:, :] - return Grid2DIrregular(values=grid_radial_projected_2d + 1e-6) + return Grid2DIrregular(values=grid_radial_projected_2d) @property def shape_native_scaled_interior(self) -> Tuple[float, float]: From 5fc2b63e465bc5c3e165db256f28900c486d4cab Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Tue, 8 Apr 2025 20:11:59 +0100 Subject: [PATCH 61/62] fix some unit tests due to numerical addition on projeted grid --- .../pixelization/border_relocator.py | 10 ++--- autoarray/plot/mat_plot/two_d.py | 5 +-- autoarray/plot/wrap/two_d/array_overlay.py | 2 - autoarray/structures/grids/irregular_2d.py | 4 +- test_autoarray/fit/test_fit_util.py | 2 - test_autoarray/mask/test_mask_2d_util.py | 2 - .../structures/grids/test_grid_2d_util.py | 39 ++++++------------- .../structures/grids/test_uniform_1d.py | 2 +- .../structures/grids/test_uniform_2d.py | 2 +- 9 files changed, 22 insertions(+), 46 deletions(-) diff --git a/autoarray/inversion/pixelization/border_relocator.py b/autoarray/inversion/pixelization/border_relocator.py index 737444b53..88ec78091 100644 --- a/autoarray/inversion/pixelization/border_relocator.py +++ b/autoarray/inversion/pixelization/border_relocator.py @@ -269,13 +269,13 @@ def relocated_grid_from(self, grid: Grid2D) -> Grid2D: return grid values = grid_2d_util.relocated_grid_via_jit_from( - grid=np.array(grid), - border_grid=np.array(grid[self.border_slim]), + grid=np.array(grid.array), + border_grid=np.array(grid.array[self.border_slim]), ) over_sampled = grid_2d_util.relocated_grid_via_jit_from( - grid=np.array(grid.over_sampled), - border_grid=np.array(grid.over_sampled[self.sub_border_slim]), + grid=np.array(grid.over_sampled.array), + border_grid=np.array(grid.over_sampled.array[self.sub_border_slim]), ) return Grid2D( @@ -303,7 +303,7 @@ def relocated_mesh_grid_from( return Grid2DIrregular( values=grid_2d_util.relocated_grid_via_jit_from( - grid=np.array(mesh_grid), + grid=np.array(mesh_grid.array), border_grid=np.array(grid[self.sub_border_slim]), ), ) diff --git a/autoarray/plot/mat_plot/two_d.py b/autoarray/plot/mat_plot/two_d.py index 64811a830..9b162d18a 100644 --- a/autoarray/plot/mat_plot/two_d.py +++ b/autoarray/plot/mat_plot/two_d.py @@ -417,7 +417,6 @@ def plot_grid( if not self.is_for_subplot: fig, ax = self.figure.open() - print(ax) else: ax = self.setup_subplot() @@ -428,10 +427,10 @@ def plot_grid( if color_array is None: if y_errors is None and x_errors is None: - self.grid_scatter.scatter_grid(grid=grid_plot) + self.grid_scatter.scatter_grid(grid=grid_plot.array) else: self.grid_errorbar.errorbar_grid( - grid=grid_plot, y_errors=y_errors, x_errors=x_errors + grid=grid_plot.array, y_errors=y_errors, x_errors=x_errors ) elif color_array is not None: diff --git a/autoarray/plot/wrap/two_d/array_overlay.py b/autoarray/plot/wrap/two_d/array_overlay.py index 5de20b879..e57512579 100644 --- a/autoarray/plot/wrap/two_d/array_overlay.py +++ b/autoarray/plot/wrap/two_d/array_overlay.py @@ -19,6 +19,4 @@ def overlay_array(self, array, figure): aspect = figure.aspect_from(shape_native=array.shape_native) extent = array.extent_of_zoomed_array(buffer=0) - print(type(array)) - plt.imshow(X=array.native._array, aspect=aspect, extent=extent, **self.config_dict) diff --git a/autoarray/structures/grids/irregular_2d.py b/autoarray/structures/grids/irregular_2d.py index 57d4c4a6f..50fdf1a76 100644 --- a/autoarray/structures/grids/irregular_2d.py +++ b/autoarray/structures/grids/irregular_2d.py @@ -233,8 +233,8 @@ def furthest_distances_to_other_coordinates(self) -> ArrayIrregular: radial_distances_max = np.zeros((self.shape[0])) for i in range(self.shape[0]): - x_distances = np.square(np.subtract(self[i, 0], self[:, 0])) - y_distances = np.square(np.subtract(self[i, 1], self[:, 1])) + x_distances = np.square(np.subtract(self.array[i, 0], self.array[:, 0])) + y_distances = np.square(np.subtract(self.array[i, 1], self.array[:, 1])) radial_distances_max[i] = np.sqrt(np.max(np.add(x_distances, y_distances))) diff --git a/test_autoarray/fit/test_fit_util.py b/test_autoarray/fit/test_fit_util.py index 6bbb5f871..b8dcace48 100644 --- a/test_autoarray/fit/test_fit_util.py +++ b/test_autoarray/fit/test_fit_util.py @@ -70,8 +70,6 @@ def test__normalized_residual_map_with_mask_from(): residual_map=residual_map, mask=mask, noise_map=noise_map ) - 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) diff --git a/test_autoarray/mask/test_mask_2d_util.py b/test_autoarray/mask/test_mask_2d_util.py index f3db2938e..ee79a9543 100644 --- a/test_autoarray/mask/test_mask_2d_util.py +++ b/test_autoarray/mask/test_mask_2d_util.py @@ -738,8 +738,6 @@ def test__edge_1d_indexes_from(): edge_pixels = util.mask_2d.edge_1d_indexes_from(mask_2d=mask) - print(edge_pixels) - assert (edge_pixels == np.array([0])).all() mask = np.array( diff --git a/test_autoarray/structures/grids/test_grid_2d_util.py b/test_autoarray/structures/grids/test_grid_2d_util.py index edf1d2497..2d45acb2c 100644 --- a/test_autoarray/structures/grids/test_grid_2d_util.py +++ b/test_autoarray/structures/grids/test_grid_2d_util.py @@ -236,7 +236,7 @@ def test__grid_scaled_2d_slim_radial_projected_from(): pixel_scales=(1.0, 1.0), ) - assert (grid_radii == np.array([[0.0, 0.0], [0.0, 1.0]])).all() + assert grid_radii == pytest.approx(np.array([[0.0, 0.0], [0.0, 1.0]]), abs=1.0e-4) grid_radii = aa.util.grid_2d.grid_scaled_2d_slim_radial_projected_from( extent=np.array([-1.0, 3.0, -1.0, 1.0]), @@ -244,9 +244,7 @@ def test__grid_scaled_2d_slim_radial_projected_from(): pixel_scales=(1.0, 1.0), ) - assert ( - grid_radii == np.array([[0.0, 0.0], [0.0, 1.0], [0.0, 2.0], [0.0, 3.0]]) - ).all() + assert grid_radii == pytest.approx(np.array([[0.0, 0.0], [0.0, 1.0], [0.0, 2.0], [0.0, 3.0]]), abs=1.0e-4) grid_radii = aa.util.grid_2d.grid_scaled_2d_slim_radial_projected_from( extent=np.array([-1.0, 3.0, -1.0, 1.0]), @@ -254,7 +252,7 @@ def test__grid_scaled_2d_slim_radial_projected_from(): pixel_scales=(1.0, 1.0), ) - assert (grid_radii == np.array([[0.0, 1.0], [0.0, 2.0], [0.0, 3.0]])).all() + assert grid_radii == pytest.approx(np.array([[0.0, 1.0], [0.0, 2.0], [0.0, 3.0]]), abs=1.0e-4) grid_radii = aa.util.grid_2d.grid_scaled_2d_slim_radial_projected_from( extent=np.array([-2.0, 1.0, -1.0, 1.0]), @@ -262,9 +260,7 @@ def test__grid_scaled_2d_slim_radial_projected_from(): pixel_scales=(1.0, 1.0), ) - assert ( - grid_radii == np.array([[0.0, 1.0], [0.0, 2.0], [0.0, 3.0], [0.0, 4.0]]) - ).all() + assert grid_radii == pytest.approx(np.array([[0.0, 1.0], [0.0, 2.0], [0.0, 3.0], [0.0, 4.0]]), abs=1.0e-4) grid_radii = aa.util.grid_2d.grid_scaled_2d_slim_radial_projected_from( extent=np.array([-1.0, 1.0, -1.0, 1.0]), @@ -272,10 +268,7 @@ def test__grid_scaled_2d_slim_radial_projected_from(): pixel_scales=(0.1, 0.5), ) - assert ( - grid_radii - == np.array([[0.0, 1.0], [0.0, 1.5], [0.0, 2.0], [0.0, 2.5], [0.0, 3.0]]) - ).all() + assert grid_radii == pytest.approx(np.array([[0.0, 1.0], [0.0, 1.5], [0.0, 2.0], [0.0, 2.5], [0.0, 3.0]]), abs=1.0e-4) grid_radii = aa.util.grid_2d.grid_scaled_2d_slim_radial_projected_from( extent=np.array([5.0, 8.0, 99.9, 100.1]), @@ -283,9 +276,7 @@ def test__grid_scaled_2d_slim_radial_projected_from(): pixel_scales=(10.0, 0.25), ) - assert ( - grid_radii - == np.array( + assert grid_radii == pytest.approx(np.array( [ [100.0, 7.0], [100.0, 7.25], @@ -297,8 +288,7 @@ def test__grid_scaled_2d_slim_radial_projected_from(): [100.0, 8.75], [100.0, 9.0], ] - ) - ).all() + ), abs=1.0e-4) grid_radii = aa.util.grid_2d.grid_scaled_2d_slim_radial_projected_from( extent=np.array([-1.0, 1.0, -1.0, 3.0]), @@ -306,9 +296,7 @@ def test__grid_scaled_2d_slim_radial_projected_from(): pixel_scales=(1.0, 1.0), ) - assert ( - grid_radii == np.array([[0.0, 0.0], [0.0, 1.0], [0.0, 2.0], [0.0, 3.0]]) - ).all() + assert grid_radii == pytest.approx(np.array([[0.0, 0.0], [0.0, 1.0], [0.0, 2.0], [0.0, 3.0]]), abs=1.0e-4) grid_radii = aa.util.grid_2d.grid_scaled_2d_slim_radial_projected_from( extent=np.array([-1.0, 1.0, -2.0, 1.0]), @@ -316,9 +304,7 @@ def test__grid_scaled_2d_slim_radial_projected_from(): pixel_scales=(1.0, 1.0), ) - assert ( - grid_radii == np.array([[1.0, 0.0], [1.0, 1.0], [1.0, 2.0], [1.0, 3.0]]) - ).all() + assert grid_radii == pytest.approx(np.array([[1.0, 0.0], [1.0, 1.0], [1.0, 2.0], [1.0, 3.0]]), abs=1.0e-4) grid_radii = aa.util.grid_2d.grid_scaled_2d_slim_radial_projected_from( extent=np.array([-1.0, 1.0, -1.0, 1.0]), @@ -326,10 +312,7 @@ def test__grid_scaled_2d_slim_radial_projected_from(): pixel_scales=(0.5, 0.1), ) - assert ( - grid_radii - == np.array([[1.0, 0.0], [1.0, 0.5], [1.0, 1.0], [1.0, 1.5], [1.0, 2.0]]) - ).all() + assert grid_radii == pytest.approx(np.array([[1.0, 0.0], [1.0, 0.5], [1.0, 1.0], [1.0, 1.5], [1.0, 2.0]]), abs=1.0e-4) grid_radii = aa.util.grid_2d.grid_scaled_2d_slim_radial_projected_from( extent=np.array([99.9, 100.1, -1.0, 3.0]), @@ -337,7 +320,7 @@ def test__grid_scaled_2d_slim_radial_projected_from(): pixel_scales=(1.5, 10.0), ) - assert (grid_radii == np.array([[-1.0, 100.0], [-1.0, 101.5], [-1.0, 103.0]])).all() + assert grid_radii == pytest.approx(np.array([[-1.0, 100.0], [-1.0, 101.5], [-1.0, 103.0]]), abs=1.0e-4) def test__grid_2d_slim_from(): diff --git a/test_autoarray/structures/grids/test_uniform_1d.py b/test_autoarray/structures/grids/test_uniform_1d.py index b99da3608..e641c4015 100644 --- a/test_autoarray/structures/grids/test_uniform_1d.py +++ b/test_autoarray/structures/grids/test_uniform_1d.py @@ -123,7 +123,7 @@ def test__grid_2d_radial_projected_from(): assert type(grid_2d) == aa.Grid2DIrregular assert grid_2d.slim == pytest.approx( - np.array([[0.0, 1.0], [0.0, 2.0], [0.0, 3.0], [0.0, 4.0]]), 1.0e-4 + np.array([[0.0, 1.0], [0.0, 2.0], [0.0, 3.0], [0.0, 4.0]]), abs=1.0e-4 ) grid_2d = grid_1d.grid_2d_radial_projected_from(angle=90.0) diff --git a/test_autoarray/structures/grids/test_uniform_2d.py b/test_autoarray/structures/grids/test_uniform_2d.py index ae5d53b10..9352dfbb1 100644 --- a/test_autoarray/structures/grids/test_uniform_2d.py +++ b/test_autoarray/structures/grids/test_uniform_2d.py @@ -559,7 +559,7 @@ def test__grid_2d_radial_projected_shape_slim_from(): pixel_scales=grid_2d.pixel_scales, ) - assert (grid_radii == grid_radii_util).all() + assert grid_radii == pytest.approx(grid_radii_util, 1.0e-4) assert grid_radial_shape_slim == grid_radii_util.shape[0] grid_2d = aa.Grid2D.uniform(shape_native=(3, 4), pixel_scales=(3.0, 2.0)) From 135d5b848737f4f1c8881569f0e1bebff9df5593 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Tue, 8 Apr 2025 20:31:31 +0100 Subject: [PATCH 62/62] black --- autoarray/dataset/imaging/dataset.py | 4 +- autoarray/dataset/imaging/simulator.py | 4 +- autoarray/dataset/interferometer/dataset.py | 4 +- autoarray/dataset/preprocess.py | 11 +++-- autoarray/fit/fit_util.py | 21 ++++++---- autoarray/geometry/geometry_util.py | 3 +- autoarray/inversion/convolver.py | 30 +++++++------- .../inversion/inversion/imaging/w_tilde.py | 4 +- .../inversion/interferometer/w_tilde.py | 8 +++- .../pixelization/image_mesh/overlay.py | 12 +++--- autoarray/mask/mask_1d.py | 4 +- autoarray/operators/contour.py | 4 +- .../operators/over_sampling/over_sampler.py | 13 +++--- autoarray/plot/visuals/two_d.py | 4 +- autoarray/plot/wrap/two_d/array_overlay.py | 4 +- autoarray/structures/arrays/array_1d_util.py | 1 + autoarray/structures/arrays/array_2d_util.py | 1 - autoarray/structures/arrays/kernel_2d.py | 4 +- autoarray/structures/grids/grid_2d_util.py | 6 ++- autoarray/structures/grids/uniform_2d.py | 20 ++++----- autoarray/structures/mesh/rectangular_2d.py | 4 +- autoarray/structures/mesh/triangulation_2d.py | 4 +- autoarray/structures/vectors/uniform.py | 5 ++- .../dataset/abstract/test_dataset.py | 8 +++- .../dataset/imaging/test_dataset.py | 15 ++++--- test_autoarray/dataset/test_preprocess.py | 7 +++- test_autoarray/fit/test_fit_util.py | 12 ++++-- .../inversion/imaging/test_imaging.py | 13 +++--- .../imaging/test_inversion_imaging_util.py | 19 ++++++--- test_autoarray/layout/test_region.py | 2 +- .../structures/grids/test_grid_2d_util.py | 41 ++++++++++++++----- 31 files changed, 190 insertions(+), 102 deletions(-) diff --git a/autoarray/dataset/imaging/dataset.py b/autoarray/dataset/imaging/dataset.py index 9fa00da1a..3779a30d6 100644 --- a/autoarray/dataset/imaging/dataset.py +++ b/autoarray/dataset/imaging/dataset.py @@ -227,7 +227,9 @@ def w_tilde(self): ) = inversion_imaging_util.w_tilde_curvature_preload_imaging_from( 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"), + native_index_for_slim_index=np.array( + self.mask.derive_indexes.native_for_slim + ).astype("int"), ) return WTildeImaging( diff --git a/autoarray/dataset/imaging/simulator.py b/autoarray/dataset/imaging/simulator.py index 972e03931..576dc6017 100644 --- a/autoarray/dataset/imaging/simulator.py +++ b/autoarray/dataset/imaging/simulator.py @@ -161,7 +161,9 @@ def via_image_from( image = image - background_sky_map mask = Mask2D.all_false( - shape_native=image.shape_native, pixel_scales=image.pixel_scales, origin=image.origin + shape_native=image.shape_native, + pixel_scales=image.pixel_scales, + origin=image.origin, ) image = Array2D(values=image, mask=mask) diff --git a/autoarray/dataset/interferometer/dataset.py b/autoarray/dataset/interferometer/dataset.py index 4cf7220c2..01b5d84bd 100644 --- a/autoarray/dataset/interferometer/dataset.py +++ b/autoarray/dataset/interferometer/dataset.py @@ -193,7 +193,9 @@ def w_tilde(self): w_matrix = inversion_interferometer_util.w_tilde_via_preload_from( w_tilde_preload=curvature_preload, - native_index_for_slim_index=np.array(self.real_space_mask.derive_indexes.native_for_slim).astype("int"), + 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/dataset/preprocess.py b/autoarray/dataset/preprocess.py index 5eb80288f..f4a7a1231 100644 --- a/autoarray/dataset/preprocess.py +++ b/autoarray/dataset/preprocess.py @@ -149,7 +149,8 @@ def noise_map_via_data_eps_and_exposure_time_map_from(data_eps, exposure_time_ma The exposure time at every data-point of the data. """ return data_eps.with_new_array( - np.abs(data_eps.array * exposure_time_map.array) ** 0.5 / exposure_time_map.array + np.abs(data_eps.array * exposure_time_map.array) ** 0.5 + / exposure_time_map.array ) @@ -263,7 +264,9 @@ def edges_from(image, no_edges): edges = [] for edge_no in range(no_edges): - top_edge = image.native.array[edge_no, edge_no : image.shape_native[1] - edge_no] + 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, @@ -521,7 +524,9 @@ def noise_map_with_signal_to_noise_limit_from( ) mask = Mask2D.all_false( - shape_native=data.shape_native, pixel_scales=data.pixel_scales, origin=data.origin + shape_native=data.shape_native, + pixel_scales=data.pixel_scales, + origin=data.origin, ) if len(noise_map.native) == 1: diff --git a/autoarray/fit/fit_util.py b/autoarray/fit/fit_util.py index 2ad7d0806..c61262c1a 100644 --- a/autoarray/fit/fit_util.py +++ b/autoarray/fit/fit_util.py @@ -174,8 +174,12 @@ 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 * np.array(noise_map).real**2.0)) - noise_normalization_imag = jnp.sum(jnp.log(2 * jnp.pi * np.array(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 @@ -224,6 +228,7 @@ def normalized_residual_map_with_mask_from( """ return jnp.where(jnp.asarray(mask) == 0, jnp.divide(residual_map, noise_map), 0) + @to_new_array def chi_squared_map_with_mask_from( *, residual_map: ty.DataLike, noise_map: ty.DataLike, mask: Mask @@ -244,11 +249,7 @@ def chi_squared_map_with_mask_from( mask The mask applied to the residual-map, where `False` entries are included in the calculation. """ - return jnp.where( - jnp.asarray(mask) == 0, - jnp.square(residual_map / noise_map), - 0 - ) + return jnp.where(jnp.asarray(mask) == 0, jnp.square(residual_map / noise_map), 0) def chi_squared_with_mask_from(*, chi_squared_map: ty.DataLike, mask: Mask) -> float: @@ -321,7 +322,9 @@ 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(jnp.sum(jnp.log(2 * jnp.pi * noise_map[jnp.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( @@ -457,4 +460,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 jnp.where(mask == 0, residual_map / data, 0) \ No newline at end of file + return jnp.where(mask == 0, residual_map / data, 0) diff --git a/autoarray/geometry/geometry_util.py b/autoarray/geometry/geometry_util.py index 7179fb879..4cf32a082 100644 --- a/autoarray/geometry/geometry_util.py +++ b/autoarray/geometry/geometry_util.py @@ -483,7 +483,8 @@ def transform_grid_2d_from_reference_frame( y = jnp.add( jnp.add( - jnp.multiply(grid_2d[:, 1], sin_angle), jnp.multiply(grid_2d[:, 0], cos_angle) + jnp.multiply(grid_2d[:, 1], sin_angle), + jnp.multiply(grid_2d[:, 0], cos_angle), ), centre[0], ) diff --git a/autoarray/inversion/convolver.py b/autoarray/inversion/convolver.py index 80a131bb0..70eb6efdf 100644 --- a/autoarray/inversion/convolver.py +++ b/autoarray/inversion/convolver.py @@ -215,12 +215,12 @@ def __init__(self, mask, kernel): 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_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] @@ -257,15 +257,15 @@ def __init__(self, mask, kernel): 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] + 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 diff --git a/autoarray/inversion/inversion/imaging/w_tilde.py b/autoarray/inversion/inversion/imaging/w_tilde.py index f98a95f51..e5912452f 100644 --- a/autoarray/inversion/inversion/imaging/w_tilde.py +++ b/autoarray/inversion/inversion/imaging/w_tilde.py @@ -74,7 +74,9 @@ def w_tilde_data(self): 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"), + native_index_for_slim_index=np.array( + self.data.mask.derive_indexes.native_for_slim + ).astype("int"), ) @property diff --git a/autoarray/inversion/inversion/interferometer/w_tilde.py b/autoarray/inversion/inversion/interferometer/w_tilde.py index 2e785dfd8..83febb864 100644 --- a/autoarray/inversion/inversion/interferometer/w_tilde.py +++ b/autoarray/inversion/inversion/interferometer/w_tilde.py @@ -131,7 +131,9 @@ 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=np.array(self.transformer.real_space_mask.derive_indexes.native_for_slim).astype("int"), + 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 +145,9 @@ 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=np.array(self.transformer.real_space_mask.derive_indexes.native_for_slim).astype("int"), + 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"), diff --git a/autoarray/inversion/pixelization/image_mesh/overlay.py b/autoarray/inversion/pixelization/image_mesh/overlay.py index de130ee6e..654cf8ec7 100644 --- a/autoarray/inversion/pixelization/image_mesh/overlay.py +++ b/autoarray/inversion/pixelization/image_mesh/overlay.py @@ -220,11 +220,13 @@ def image_plane_mesh_grid_from( origin=origin, ) - 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") + 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") total_pixels = total_pixels_2d_from( mask_2d=mask.array, diff --git a/autoarray/mask/mask_1d.py b/autoarray/mask/mask_1d.py index 6ceacae7f..77f79c01e 100644 --- a/autoarray/mask/mask_1d.py +++ b/autoarray/mask/mask_1d.py @@ -153,7 +153,9 @@ def from_fits( """ return cls( - mask=np.array(array_1d_util.numpy_array_1d_via_fits_from(file_path=file_path, hdu=hdu)), + mask=np.array( + array_1d_util.numpy_array_1d_via_fits_from(file_path=file_path, hdu=hdu) + ), pixel_scales=pixel_scales, origin=origin, ) diff --git a/autoarray/operators/contour.py b/autoarray/operators/contour.py index 804c86a49..abc9c7a25 100644 --- a/autoarray/operators/contour.py +++ b/autoarray/operators/contour.py @@ -61,9 +61,7 @@ def contour_list(self): else: contour_array = np.array(self.contour_array.array) - contour_indices_list = measure.find_contours( - contour_array, 0 - ) + contour_indices_list = measure.find_contours(contour_array, 0) print(contour_indices_list) diff --git a/autoarray/operators/over_sampling/over_sampler.py b/autoarray/operators/over_sampling/over_sampler.py index 73f7a1aa7..b78bb0829 100644 --- a/autoarray/operators/over_sampling/over_sampler.py +++ b/autoarray/operators/over_sampling/over_sampler.py @@ -147,15 +147,12 @@ 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, self.sub_size[0]) - ) + return np.all(np.isclose(self.sub_size, self.sub_size[0])) def tree_flatten(self): return (self.mask, self.sub_size), () @@ -238,7 +235,7 @@ def binned_array_2d_from(self, array: Array2D) -> "Array2D": else: # Define group sizes - group_sizes = jnp.array(self.sub_size.array ** 2) + group_sizes = jnp.array(self.sub_size.array**2) # Compute the cumulative sum of group sizes to get split points split_indices = jnp.cumsum(group_sizes) @@ -248,7 +245,11 @@ def binned_array_2d_from(self, array: Array2D) -> "Array2D": # Compute the group means binned_array_2d = jnp.array( - [array[start:end].mean() for start, end in zip(start_indices, split_indices)]) + [ + array[start:end].mean() + for start, end in zip(start_indices, split_indices) + ] + ) return Array2D( values=binned_array_2d, diff --git a/autoarray/plot/visuals/two_d.py b/autoarray/plot/visuals/two_d.py index eacca5edc..e892a06a7 100644 --- a/autoarray/plot/visuals/two_d.py +++ b/autoarray/plot/visuals/two_d.py @@ -55,7 +55,9 @@ def plot_via_plotter(self, plotter, grid_indexes=None, mapper=None, geometry=Non ) if self.mask is not None: - plotter.mask_scatter.scatter_grid(grid=np.array(self.mask.derive_grid.edge.array)) + plotter.mask_scatter.scatter_grid( + grid=np.array(self.mask.derive_grid.edge.array) + ) if self.border is not None: plotter.border_scatter.scatter_grid(grid=self.border.array) diff --git a/autoarray/plot/wrap/two_d/array_overlay.py b/autoarray/plot/wrap/two_d/array_overlay.py index e57512579..688866ade 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._array, aspect=aspect, extent=extent, **self.config_dict) + plt.imshow( + X=array.native._array, aspect=aspect, extent=extent, **self.config_dict + ) diff --git a/autoarray/structures/arrays/array_1d_util.py b/autoarray/structures/arrays/array_1d_util.py index f7c943724..915a2a4ad 100644 --- a/autoarray/structures/arrays/array_1d_util.py +++ b/autoarray/structures/arrays/array_1d_util.py @@ -127,6 +127,7 @@ def array_1d_native_from( native_index_for_slim_index_1d=native_index_for_slim_index_1d, ) + def array_1d_via_indexes_1d_from( array_1d_slim: np.ndarray, shape: int, diff --git a/autoarray/structures/arrays/array_2d_util.py b/autoarray/structures/arrays/array_2d_util.py index 17c27a29f..f8767e2f9 100644 --- a/autoarray/structures/arrays/array_2d_util.py +++ b/autoarray/structures/arrays/array_2d_util.py @@ -33,7 +33,6 @@ def convert_array(array: Union[np.ndarray, List]) -> np.ndarray: return array - def check_array_2d(array_2d: np.ndarray): if len(array_2d.shape) != 1: raise exc.ArrayException( diff --git a/autoarray/structures/arrays/kernel_2d.py b/autoarray/structures/arrays/kernel_2d.py index 9d04c29ed..c2962dbbe 100644 --- a/autoarray/structures/arrays/kernel_2d.py +++ b/autoarray/structures/arrays/kernel_2d.py @@ -552,7 +552,9 @@ def convolve_image_no_blurring(self, image, mask, jax_method="direct"): if isinstance(image, np.ndarray) or isinstance(image, jnp.ndarray): expanded_array_native = expanded_array_native.at[slim_to_native].set(image) else: - expanded_array_native = expanded_array_native.at[slim_to_native].set(image.array) + expanded_array_native = expanded_array_native.at[slim_to_native].set( + image.array + ) kernel = np.array(self.native.array) diff --git a/autoarray/structures/grids/grid_2d_util.py b/autoarray/structures/grids/grid_2d_util.py index cec240ad4..b547fcb94 100644 --- a/autoarray/structures/grids/grid_2d_util.py +++ b/autoarray/structures/grids/grid_2d_util.py @@ -47,6 +47,7 @@ def check_grid_slim(grid, shape_native): """ ) + def check_grid_2d(grid_2d: np.ndarray): if grid_2d.shape[-1] != 2: raise exc.GridException( @@ -140,6 +141,7 @@ def convert_grid_2d( ) return np.array(grid_2d) if is_numpy else jnp.array(grid_2d) + def convert_grid_2d_to_slim( grid_2d: Union[np.ndarray, List], mask_2d: Mask2D ) -> np.ndarray: @@ -564,7 +566,7 @@ def grid_scaled_2d_slim_radial_projected_from( grid_scaled_2d_slim_radii[slim_index, 1] = radii radii += pixel_scale - return grid_scaled_2d_slim_radii + 1e-6 + return grid_scaled_2d_slim_radii + 1e-6 @numba_util.jit() @@ -849,4 +851,4 @@ def grid_pixels_in_mask_pixels_from( # 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 \ No newline at end of file + return mesh_pixels_per_image_pixel diff --git a/autoarray/structures/grids/uniform_2d.py b/autoarray/structures/grids/uniform_2d.py index 890301884..addfb2ec1 100644 --- a/autoarray/structures/grids/uniform_2d.py +++ b/autoarray/structures/grids/uniform_2d.py @@ -181,13 +181,11 @@ def __init__( self.over_sampler = OverSampler(sub_size=over_sample_size, mask=mask) if over_sampled is None: - over_sampled = ( - over_sample_util.grid_2d_slim_over_sampled_via_mask_from( - mask_2d=np.array(self.mask), - pixel_scales=self.mask.pixel_scales, - sub_size=self.over_sampler.sub_size.array.astype("int"), - origin=self.mask.origin, - ) + over_sampled = over_sample_util.grid_2d_slim_over_sampled_via_mask_from( + mask_2d=np.array(self.mask), + pixel_scales=self.mask.pixel_scales, + sub_size=self.over_sampler.sub_size.array.astype("int"), + origin=self.mask.origin, ) self.over_sampled = Grid2DIrregular(values=over_sampled) @@ -842,9 +840,9 @@ def squared_distances_to_coordinate_from( coordinate The (y,x) coordinate from which the squared distance of every grid (y,x) coordinate is computed. """ - squared_distances = jnp.square( - self.array[:, 0] - coordinate[0] - ) + jnp.square(self.array[:, 1] - coordinate[1]) + squared_distances = jnp.square(self.array[:, 0] - coordinate[0]) + jnp.square( + self.array[:, 1] - coordinate[1] + ) return Array2D(values=squared_distances, mask=self.mask) @@ -1098,7 +1096,7 @@ def padded_grid_from(self, kernel_shape_native: Tuple[int, int]) -> "Grid2D": padded_mask = Mask2D.all_false( shape_native=padded_shape, pixel_scales=self.mask.pixel_scales, - origin=self.origin + origin=self.origin, ) pad_width = ( diff --git a/autoarray/structures/mesh/rectangular_2d.py b/autoarray/structures/mesh/rectangular_2d.py index 7541d592a..c46fc258e 100644 --- a/autoarray/structures/mesh/rectangular_2d.py +++ b/autoarray/structures/mesh/rectangular_2d.py @@ -178,7 +178,9 @@ def interpolated_array_from( shape_native=shape_native, extent=extent ) - interpolated_array = griddata(points=self.array, values=values, xi=interpolation_grid) + interpolated_array = griddata( + points=self.array, values=values, xi=interpolation_grid + ) interpolated_array = interpolated_array.reshape(shape_native) diff --git a/autoarray/structures/mesh/triangulation_2d.py b/autoarray/structures/mesh/triangulation_2d.py index 7d96415b4..635e02f62 100644 --- a/autoarray/structures/mesh/triangulation_2d.py +++ b/autoarray/structures/mesh/triangulation_2d.py @@ -96,7 +96,9 @@ def delaunay(self) -> scipy.spatial.Delaunay: `MeshException`, which helps exception handling in the `inversion` package. """ try: - return scipy.spatial.Delaunay(np.asarray([self.array[:, 0], self.array[:, 1]]).T) + return scipy.spatial.Delaunay( + np.asarray([self.array[:, 0], self.array[:, 1]]).T + ) except (ValueError, OverflowError, scipy.spatial.qhull.QhullError) as e: raise exc.MeshException() from e diff --git a/autoarray/structures/vectors/uniform.py b/autoarray/structures/vectors/uniform.py index 4da4836b2..88dfe6a9c 100644 --- a/autoarray/structures/vectors/uniform.py +++ b/autoarray/structures/vectors/uniform.py @@ -393,7 +393,10 @@ def magnitudes(self) -> Array2D: """ Returns the magnitude of every vector which are computed as sqrt(y**2 + x**2). """ - return Array2D(values=jnp.sqrt(self.array[:, 0] ** 2.0 + self.array[:, 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: diff --git a/test_autoarray/dataset/abstract/test_dataset.py b/test_autoarray/dataset/abstract/test_dataset.py index 770d9442b..794af761c 100644 --- a/test_autoarray/dataset/abstract/test_dataset.py +++ b/test_autoarray/dataset/abstract/test_dataset.py @@ -32,7 +32,9 @@ def test__signal_to_noise_map(): dataset = ds.AbstractDataset(data=array, noise_map=noise_map) - assert dataset.signal_to_noise_map.native == pytest.approx(np.array([[0.1, 0.2], [0.1, 1.0]]), 1.0e-4) + assert dataset.signal_to_noise_map.native == pytest.approx( + np.array([[0.1, 0.2], [0.1, 1.0]]), 1.0e-4 + ) assert dataset.signal_to_noise_max == 1.0 array = aa.Array2D.no_mask([[-1.0, 2.0], [3.0, -4.0]], pixel_scales=1.0) @@ -41,7 +43,9 @@ def test__signal_to_noise_map(): dataset = ds.AbstractDataset(data=array, noise_map=noise_map) - assert dataset.signal_to_noise_map.native == pytest.approx(np.array([[0.0, 0.2], [0.1, 0.0]]), 1.0e-4) + assert dataset.signal_to_noise_map.native == pytest.approx( + np.array([[0.0, 0.2], [0.1, 0.0]]), 1.0e-4 + ) assert dataset.signal_to_noise_max == 0.2 diff --git a/test_autoarray/dataset/imaging/test_dataset.py b/test_autoarray/dataset/imaging/test_dataset.py index 1665865f2..ead9e51f5 100644 --- a/test_autoarray/dataset/imaging/test_dataset.py +++ b/test_autoarray/dataset/imaging/test_dataset.py @@ -140,7 +140,9 @@ 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 == pytest.approx((1.0 / 3.0) * psf_3x3.slim.array, 1.0e-4) + assert masked_imaging_7x7.psf.slim == pytest.approx( + (1.0 / 3.0) * psf_3x3.slim.array, 1.0e-4 + ) assert type(masked_imaging_7x7.psf) == aa.Kernel2D assert masked_imaging_7x7.w_tilde.curvature_preload.shape == (35,) @@ -157,10 +159,12 @@ def test__apply_noise_scaling(imaging_7x7, mask_2d_7x7): assert masked_imaging_7x7.noise_map.native[4, 4] == 1e5 -def test__apply_noise_scaling__use_signal_to_noise_value(image_7x7, psf_3x3, noise_map_7x7, mask_2d_7x7): +def test__apply_noise_scaling__use_signal_to_noise_value( + image_7x7, psf_3x3, noise_map_7x7, mask_2d_7x7 +): image_7x7 = np.array(image_7x7.native.array) - image_7x7[3,3] = 2.0 + image_7x7[3, 3] = 2.0 image_7x7 = aa.Array2D(values=image_7x7, mask=mask_2d_7x7) @@ -258,5 +262,6 @@ def test__psf_not_odd_x_odd_kernel__raises_error(): 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) - + dataset = aa.Imaging( + data=image, noise_map=noise_map, psf=psf, pad_for_psf=False + ) diff --git a/test_autoarray/dataset/test_preprocess.py b/test_autoarray/dataset/test_preprocess.py index 67107be4a..8e99a74e8 100644 --- a/test_autoarray/dataset/test_preprocess.py +++ b/test_autoarray/dataset/test_preprocess.py @@ -133,12 +133,15 @@ def test__noise_map_from_image_exposure_time_map(): data_eps=image, exposure_time_map=exposure_time_map ) - assert poisson_noise_map.native == pytest.approx(np.array( + assert poisson_noise_map.native == pytest.approx( + np.array( [ [np.sqrt(5.0), np.sqrt(6.0) / 2.0], [np.sqrt(30.0) / 3.0, np.sqrt(80.0) / 4.0], ] - ), 1.0e-4) + ), + 1.0e-4, + ) def test__noise_map_from_image_exposure_time_map_and_background_noise_map(): diff --git a/test_autoarray/fit/test_fit_util.py b/test_autoarray/fit/test_fit_util.py index b8dcace48..0cbc811de 100644 --- a/test_autoarray/fit/test_fit_util.py +++ b/test_autoarray/fit/test_fit_util.py @@ -43,7 +43,9 @@ def test__normalized_residual_map_from(): residual_map=residual_map, noise_map=noise_map ) - assert normalized_residual_map == pytest.approx(jnp.array([0.0, 0.0, 0.0, 0.0]), 1.0e-4) + assert normalized_residual_map == pytest.approx( + jnp.array([0.0, 0.0, 0.0, 0.0]), 1.0e-4 + ) model_data = jnp.array([11.0, 10.0, 9.0, 8.0]) @@ -53,7 +55,9 @@ def test__normalized_residual_map_from(): residual_map=residual_map, noise_map=noise_map ) - 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) + 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(): @@ -70,7 +74,9 @@ def test__normalized_residual_map_with_mask_from(): residual_map=residual_map, mask=mask, noise_map=noise_map ) - assert normalized_residual_map == pytest.approx(jnp.array([0.0, 0.0, (1.0 / 2.0), 0.0]), abs=1.0e-4) + 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(): diff --git a/test_autoarray/inversion/inversion/imaging/test_imaging.py b/test_autoarray/inversion/inversion/imaging/test_imaging.py index 8a52e4c35..0fad37e15 100644 --- a/test_autoarray/inversion/inversion/imaging/test_imaging.py +++ b/test_autoarray/inversion/inversion/imaging/test_imaging.py @@ -15,7 +15,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], - convolver=aa.Convolver(kernel=psf_3x3, mask=rectangular_mapper_7x7_3x3.mapper_grids.mask) + 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) @@ -26,7 +28,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))) + 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]]) @@ -57,8 +59,9 @@ 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], - convolver=aa.m.MockConvolver(operated_mapping_matrix=np.ones((2, 2))) + 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]]) @@ -92,7 +95,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))) + 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 fab5b6600..570f79673 100644 --- a/test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py +++ b/test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py @@ -219,7 +219,9 @@ def test__data_vector_via_w_tilde_data_two_methods_agree(): 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"), + native_index_for_slim_index=np.array( + mask.derive_indexes.native_for_slim + ).astype("int"), ) ( @@ -275,14 +277,18 @@ def test__curvature_matrix_via_w_tilde_two_methods_agree(): w_tilde = aa.util.inversion_imaging.w_tilde_curvature_imaging_from( 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"), + 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, mask=mask) + 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, @@ -328,7 +334,9 @@ def test__curvature_matrix_via_w_tilde_preload_two_methods_agree(): ) = aa.util.inversion_imaging.w_tilde_curvature_preload_imaging_from( 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"), + native_index_for_slim_index=np.array( + mask.derive_indexes.native_for_slim + ).astype("int"), ) ( @@ -355,7 +363,8 @@ def test__curvature_matrix_via_w_tilde_preload_two_methods_agree(): ) blurred_mapping_matrix = psf.convolve_mapping_matrix( - mapping_matrix=mapping_matrix, mask=mask, + mapping_matrix=mapping_matrix, + mask=mask, ) curvature_matrix = aa.util.inversion.curvature_matrix_via_mapping_matrix_from( diff --git a/test_autoarray/layout/test_region.py b/test_autoarray/layout/test_region.py index 78bfba73f..690cfa9bb 100644 --- a/test_autoarray/layout/test_region.py +++ b/test_autoarray/layout/test_region.py @@ -362,4 +362,4 @@ def test__exception__negative_coordinates_2d(): aa.Region2D(region=(0, 0, -1, 2)) with pytest.raises(exc.RegionException): - aa.Region2D(region=(0, 1, 2, -1)) \ No newline at end of file + aa.Region2D(region=(0, 1, 2, -1)) diff --git a/test_autoarray/structures/grids/test_grid_2d_util.py b/test_autoarray/structures/grids/test_grid_2d_util.py index 2d45acb2c..463aa4e7b 100644 --- a/test_autoarray/structures/grids/test_grid_2d_util.py +++ b/test_autoarray/structures/grids/test_grid_2d_util.py @@ -244,7 +244,9 @@ def test__grid_scaled_2d_slim_radial_projected_from(): pixel_scales=(1.0, 1.0), ) - assert grid_radii == pytest.approx(np.array([[0.0, 0.0], [0.0, 1.0], [0.0, 2.0], [0.0, 3.0]]), abs=1.0e-4) + assert grid_radii == pytest.approx( + np.array([[0.0, 0.0], [0.0, 1.0], [0.0, 2.0], [0.0, 3.0]]), abs=1.0e-4 + ) grid_radii = aa.util.grid_2d.grid_scaled_2d_slim_radial_projected_from( extent=np.array([-1.0, 3.0, -1.0, 1.0]), @@ -252,7 +254,9 @@ def test__grid_scaled_2d_slim_radial_projected_from(): pixel_scales=(1.0, 1.0), ) - assert grid_radii == pytest.approx(np.array([[0.0, 1.0], [0.0, 2.0], [0.0, 3.0]]), abs=1.0e-4) + assert grid_radii == pytest.approx( + np.array([[0.0, 1.0], [0.0, 2.0], [0.0, 3.0]]), abs=1.0e-4 + ) grid_radii = aa.util.grid_2d.grid_scaled_2d_slim_radial_projected_from( extent=np.array([-2.0, 1.0, -1.0, 1.0]), @@ -260,7 +264,9 @@ def test__grid_scaled_2d_slim_radial_projected_from(): pixel_scales=(1.0, 1.0), ) - assert grid_radii == pytest.approx(np.array([[0.0, 1.0], [0.0, 2.0], [0.0, 3.0], [0.0, 4.0]]), abs=1.0e-4) + assert grid_radii == pytest.approx( + np.array([[0.0, 1.0], [0.0, 2.0], [0.0, 3.0], [0.0, 4.0]]), abs=1.0e-4 + ) grid_radii = aa.util.grid_2d.grid_scaled_2d_slim_radial_projected_from( extent=np.array([-1.0, 1.0, -1.0, 1.0]), @@ -268,7 +274,10 @@ def test__grid_scaled_2d_slim_radial_projected_from(): pixel_scales=(0.1, 0.5), ) - assert grid_radii == pytest.approx(np.array([[0.0, 1.0], [0.0, 1.5], [0.0, 2.0], [0.0, 2.5], [0.0, 3.0]]), abs=1.0e-4) + assert grid_radii == pytest.approx( + np.array([[0.0, 1.0], [0.0, 1.5], [0.0, 2.0], [0.0, 2.5], [0.0, 3.0]]), + abs=1.0e-4, + ) grid_radii = aa.util.grid_2d.grid_scaled_2d_slim_radial_projected_from( extent=np.array([5.0, 8.0, 99.9, 100.1]), @@ -276,7 +285,8 @@ def test__grid_scaled_2d_slim_radial_projected_from(): pixel_scales=(10.0, 0.25), ) - assert grid_radii == pytest.approx(np.array( + assert grid_radii == pytest.approx( + np.array( [ [100.0, 7.0], [100.0, 7.25], @@ -288,7 +298,9 @@ def test__grid_scaled_2d_slim_radial_projected_from(): [100.0, 8.75], [100.0, 9.0], ] - ), abs=1.0e-4) + ), + abs=1.0e-4, + ) grid_radii = aa.util.grid_2d.grid_scaled_2d_slim_radial_projected_from( extent=np.array([-1.0, 1.0, -1.0, 3.0]), @@ -296,7 +308,9 @@ def test__grid_scaled_2d_slim_radial_projected_from(): pixel_scales=(1.0, 1.0), ) - assert grid_radii == pytest.approx(np.array([[0.0, 0.0], [0.0, 1.0], [0.0, 2.0], [0.0, 3.0]]), abs=1.0e-4) + assert grid_radii == pytest.approx( + np.array([[0.0, 0.0], [0.0, 1.0], [0.0, 2.0], [0.0, 3.0]]), abs=1.0e-4 + ) grid_radii = aa.util.grid_2d.grid_scaled_2d_slim_radial_projected_from( extent=np.array([-1.0, 1.0, -2.0, 1.0]), @@ -304,7 +318,9 @@ def test__grid_scaled_2d_slim_radial_projected_from(): pixel_scales=(1.0, 1.0), ) - assert grid_radii == pytest.approx(np.array([[1.0, 0.0], [1.0, 1.0], [1.0, 2.0], [1.0, 3.0]]), abs=1.0e-4) + assert grid_radii == pytest.approx( + np.array([[1.0, 0.0], [1.0, 1.0], [1.0, 2.0], [1.0, 3.0]]), abs=1.0e-4 + ) grid_radii = aa.util.grid_2d.grid_scaled_2d_slim_radial_projected_from( extent=np.array([-1.0, 1.0, -1.0, 1.0]), @@ -312,7 +328,10 @@ def test__grid_scaled_2d_slim_radial_projected_from(): pixel_scales=(0.5, 0.1), ) - assert grid_radii == pytest.approx(np.array([[1.0, 0.0], [1.0, 0.5], [1.0, 1.0], [1.0, 1.5], [1.0, 2.0]]), abs=1.0e-4) + assert grid_radii == pytest.approx( + np.array([[1.0, 0.0], [1.0, 0.5], [1.0, 1.0], [1.0, 1.5], [1.0, 2.0]]), + abs=1.0e-4, + ) grid_radii = aa.util.grid_2d.grid_scaled_2d_slim_radial_projected_from( extent=np.array([99.9, 100.1, -1.0, 3.0]), @@ -320,7 +339,9 @@ def test__grid_scaled_2d_slim_radial_projected_from(): pixel_scales=(1.5, 10.0), ) - assert grid_radii == pytest.approx(np.array([[-1.0, 100.0], [-1.0, 101.5], [-1.0, 103.0]]), abs=1.0e-4) + assert grid_radii == pytest.approx( + np.array([[-1.0, 100.0], [-1.0, 101.5], [-1.0, 103.0]]), abs=1.0e-4 + ) def test__grid_2d_slim_from():