Merge pull request #194 from Jammy2211/feature/simplify_slim_to_native_tuple

Feature/simplify slim to native tuple

Author: Jammy2211

autoarray/abstract_ndarray.py

@@ -338,6 +338,7 @@ def __getitem__(self, item):
 
         try:
             import jax.numpy as jnp
+
             if isinstance(result, jnp.ndarray):
                 result = self.with_new_array(result)
         except ImportError:
@@ -351,6 +352,7 @@ def __setitem__(self, key, value):
             self._array[key] = value
         else:
             import jax.numpy as jnp
+
             self._array = jnp.where(key, value, self._array)
 
     def __repr__(self):

autoarray/dataset/imaging/dataset.py

@@ -3,6 +3,8 @@
 from pathlib import Path
 from typing import Optional, Union
 
+from autoconf import cached_property
+
 from autoarray.dataset.abstract.dataset import AbstractDataset
 from autoarray.dataset.grids import GridsDataset
 from autoarray.dataset.imaging.w_tilde import WTildeImaging
@@ -191,7 +193,7 @@ def __init__(
             psf=self.psf,
         )
 
-    @property
+    @cached_property
     def w_tilde(self):
         """
         The w_tilde formalism of the linear algebra equations precomputes the convolution of every pair of masked

autoarray/dataset/interferometer/dataset.py

@@ -2,6 +2,8 @@
 import numpy as np
 from pathlib import Path
 
+from autoconf import cached_property
+
 from autoconf.fitsable import ndarray_via_fits_from, output_to_fits
 
 from autoarray.dataset.abstract.dataset import AbstractDataset
@@ -14,6 +16,8 @@
 
 from autoarray.inversion.inversion.interferometer import inversion_interferometer_util
 
+from autoarray import exc
+
 logger = logging.getLogger(__name__)
 
 
@@ -165,7 +169,7 @@ def w_tilde_preprocessing(self):
 
                 fits.writeto(filename, data=curvature_preload)
 
-    @property
+    @cached_property
     def w_tilde(self):
         """
         The w_tilde formalism of the linear algebra equations precomputes the Fourier Transform of all the visibilities

autoarray/fit/fit_util.py

@@ -198,7 +198,7 @@ def residual_map_with_mask_from(
     model_data
         The model data used to fit the data.
     """
-    return xp.where(xp.asarray(mask) == 0, xp.subtract(data, model_data), 0)
+    return xp.where(mask == 0, xp.subtract(data, model_data), 0)
 
 
 @to_new_array
@@ -221,7 +221,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 xp.where(xp.asarray(mask) == 0, xp.divide(residual_map, noise_map), 0)
+    return xp.where(mask == 0, xp.divide(residual_map, noise_map), 0)
 
 
 @to_new_array
@@ -244,7 +244,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 xp.where(xp.asarray(mask) == 0, xp.square(residual_map / noise_map), 0)
+    return xp.where(mask == 0, xp.square(residual_map / noise_map), 0)
 
 
 def chi_squared_with_mask_from(
@@ -263,7 +263,7 @@ def chi_squared_with_mask_from(
     mask
         The mask applied to the chi-squared-map, where `False` entries are included in the calculation.
     """
-    return float(xp.sum(chi_squared_map[xp.asarray(mask) == 0]))
+    return float(xp.sum(chi_squared_map[mask == 0]))
 
 
 def chi_squared_with_mask_fast_from(
@@ -301,8 +301,8 @@ def chi_squared_with_mask_fast_from(
                     xp.subtract(
                         data,
                         model_data,
-                    )[xp.asarray(mask) == 0],
-                    noise_map[xp.asarray(mask) == 0],
+                    )[mask == 0],
+                    noise_map[mask == 0],
                 )
             )
         )
@@ -326,7 +326,7 @@ def noise_normalization_with_mask_from(
     mask
         The mask applied to the noise-map, where `False` entries are included in the calculation.
     """
-    return float(xp.sum(xp.log(2 * xp.pi * noise_map[xp.asarray(mask) == 0] ** 2.0)))
+    return float(xp.sum(xp.log(2 * xp.pi * noise_map[mask == 0] ** 2.0)))
 
 
 def chi_squared_with_noise_covariance_from(

autoarray/inversion/pixelization/mesh/rectangular.py

@@ -158,7 +158,12 @@ def requires_image_mesh(self):
 
 class RectangularSource(RectangularMagnification):
 
-    def __init__(self, shape: Tuple[int, int] = (3, 3), weight_power: float = 1.0, weight_floor : float = 0.0):
+    def __init__(
+        self,
+        shape: Tuple[int, int] = (3, 3),
+        weight_power: float = 1.0,
+        weight_floor: float = 0.0,
+    ):
         """
         A uniform mesh of rectangular pixels, which without interpolation are paired with a 2D grid of (y,x)
         coordinates.
@@ -203,9 +208,9 @@ def mesh_weight_map_from(self, adapt_data, xp=np) -> np.ndarray:
         xp
             The array library to use.
         """
-        mesh_weight_map = xp.asarray(adapt_data.array)
+        mesh_weight_map = adapt_data.array
         mesh_weight_map = xp.clip(mesh_weight_map, 1e-12, None)
-        mesh_weight_map = mesh_weight_map ** self.weight_power
+        mesh_weight_map = mesh_weight_map**self.weight_power
 
         # Apply floor using xp.where (safe for NumPy and JAX)
         mesh_weight_map = xp.where(

autoarray/inversion/regularization/gaussian_kernel.py

@@ -33,7 +33,7 @@ def gauss_cov_matrix_from(
         The Gaussian covariance matrix.
     """
     # Ensure array:
-    pts = xp.asarray(pixel_points)  # (N, 2)
+    pts = pixel_points  # (N, 2)
     # Compute squared distances: ||p_i - p_j||^2
     diffs = pts[:, None, :] - pts[None, :, :]  # (N, N, 2)
     d2 = xp.sum(diffs**2, axis=-1)  # (N, N)

autoarray/mask/mask_2d.py

@@ -217,6 +217,9 @@ def __init__(
             xp=xp,
         )
 
+        slim_to_native = self.derive_indexes.native_for_slim.astype("int32")
+        self.slim_to_native_tuple = (slim_to_native[:, 0], slim_to_native[:, 1])
+
     @property
     def native_for_slim(self):
         return self.derive_indexes.native_for_slim

autoarray/operators/over_sampling/over_sampler.py

@@ -258,10 +258,14 @@ def binned_array_2d_from(self, array: Array2D, xp=np) -> "Array2D":
             else:
 
                 # Sum values per segment
-                sums = np.bincount(self.segment_ids, weights=array, minlength=self.mask.pixels_in_mask)
+                sums = np.bincount(
+                    self.segment_ids, weights=array, minlength=self.mask.pixels_in_mask
+                )
 
                 # Count number of items per segment
-                counts = np.bincount(self.segment_ids, minlength=self.mask.pixels_in_mask)
+                counts = np.bincount(
+                    self.segment_ids, minlength=self.mask.pixels_in_mask
+                )
 
                 # Avoid division by zero
                 counts[counts == 0] = 1

autoarray/operators/transformer.py

@@ -142,7 +142,7 @@ def visibilities_from(self, image: Array2D, xp=np) -> Visibilities:
                 image_1d=image.slim.array,
                 grid_radians=self.grid.array,
                 uv_wavelengths=self.uv_wavelengths,
-                xp=xp
+                xp=xp,
             )
 
         return Visibilities(visibilities=xp.array(visibilities))

autoarray/operators/transformer_util.py

@@ -247,7 +247,10 @@ def transformed_mapping_matrix_via_preload_from(
 
 
 def transformed_mapping_matrix_from(
-    mapping_matrix: np.ndarray, grid_radians: np.ndarray, uv_wavelengths: np.ndarray, xp=np
+    mapping_matrix: np.ndarray,
+    grid_radians: np.ndarray,
+    uv_wavelengths: np.ndarray,
+    xp=np,
 ) -> np.ndarray:
     """
     Computes the Fourier-transformed mapping matrix used in radio interferometric imaging.