Feature/xp namespace

autoarray/abstract_ndarray.py

@@ -6,7 +6,8 @@
 from abc import abstractmethod
 import jax.numpy as jnp
 from jax._src.tree_util import register_pytree_node
-from jax import Array
+
+import numpy as np
 
 from autoconf.fitsable import output_to_fits
 
@@ -64,7 +65,11 @@ def wrapper(self, other):
 
 
 class AbstractNDArray(ABC):
-    def __init__(self, array):
+
+    __no_flatten__ = ()
+
+    def __init__(self, array, xp=np):
+
         self._is_transformed = False
 
         while isinstance(array, AbstractNDArray):
@@ -79,7 +84,7 @@ def __init__(self, array):
         except ValueError:
             pass
 
-    __no_flatten__ = ()
+        self._xp = xp
 
     def invert(self):
         new = self.copy()
@@ -102,12 +107,6 @@ def instance_flatten(cls, instance):
         )
         return values, keys
 
-    @staticmethod
-    def flip_hdu_for_ds9(values):
-        if conf.instance["general"]["fits"]["flip_for_ds9"]:
-            return jnp.flipud(values)
-        return values
-
     @classmethod
     def instance_unflatten(cls, aux_data, children):
         """
@@ -138,6 +137,12 @@ def with_new_array(self, array: jnp.ndarray) -> "AbstractNDArray":
         new_array._array = array
         return new_array
 
+    @staticmethod
+    def flip_hdu_for_ds9(values):
+        if conf.instance["general"]["fits"]["flip_for_ds9"]:
+            return jnp.flipud(values)
+        return values
+
     def copy(self):
         new = copy(self)
         return new
@@ -336,6 +341,7 @@ def __getitem__(self, item):
         return result
 
     def __setitem__(self, key, value):
+        from jax import Array
         if isinstance(key, (jnp.ndarray, AbstractNDArray, Array)):
             self._array = jnp.where(key, value, self._array)
         else:

autoarray/config/general.yaml

@@ -13,7 +13,7 @@ inversion:
   reconstruction_vmax_factor: 0.5     # Plots of an Inversion's reconstruction use the reconstructed data's bright value multiplied by this factor.
 numba:
   use_numba: true
-  cache: false
+  cache: true
   nopython: true
   parallel: false
 pixelization:

autoarray/dataset/abstract/dataset.py

@@ -4,10 +4,6 @@
 import warnings
 from typing import Optional, Union
 
-from autoconf import cached_property
-
-from autoarray.dataset.grids import GridsDataset
-
 from autoarray import exc
 from autoarray.mask.mask_1d import Mask1D
 from autoarray.mask.mask_2d import Mask2D
@@ -140,14 +136,6 @@ def __init__(
     def grid(self):
         return self.grids.lp
 
-    @cached_property
-    def grids(self):
-        return GridsDataset(
-            mask=self.data.mask,
-            over_sample_size_lp=self.over_sample_size_lp,
-            over_sample_size_pixelization=self.over_sample_size_pixelization,
-        )
-
     @property
     def shape_native(self):
         return self.mask.shape_native
@@ -188,7 +176,7 @@ def signal_to_noise_max(self) -> float:
         """
         return np.max(self.signal_to_noise_map)
 
-    @cached_property
+    @property
     def noise_covariance_matrix_inv(self) -> np.ndarray:
         """
         Returns the inverse of the noise covariance matrix, which is used when computing a chi-squared which accounts

autoarray/dataset/imaging/dataset.py

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

autoarray/dataset/imaging/w_tilde.py

@@ -1,8 +1,6 @@
 import logging
 import numpy as np
 
-from autoconf import cached_property
-
 from autoarray.dataset.abstract.w_tilde import AbstractWTilde
 
 from autoarray.inversion.inversion.imaging import inversion_imaging_util
@@ -55,7 +53,7 @@ def __init__(
         self.psf = psf
         self.mask = mask
 
-    @cached_property
+    @property
     def w_matrix(self):
         """
         The matrix `w_tilde_curvature` is a matrix of dimensions [image_pixels, image_pixels] that encodes the PSF
@@ -93,7 +91,7 @@ def w_matrix(self):
             ).astype("int"),
         )
 
-    @cached_property
+    @property
     def psf_operator_matrix_dense(self):
 
         return inversion_imaging_util.psf_operator_matrix_dense_from(

autoarray/dataset/interferometer/dataset.py

@@ -2,7 +2,6 @@
 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
@@ -166,7 +165,7 @@ def w_tilde_preprocessing(self):
 
                 fits.writeto(filename, data=curvature_preload)
 
-    @cached_property
+    @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_dataset.py

@@ -5,8 +5,6 @@
 
 import numpy as np
 
-from autoconf import cached_property
-
 from autoarray.dataset.grids import GridsInterface
 from autoarray.dataset.dataset_model import DatasetModel
 from autoarray.fit import fit_util
@@ -85,7 +83,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.array)
+        return fit_util.chi_squared_from(chi_squared_map=self.chi_squared_map.array, xp=self._xp)
 
     @property
     def noise_normalization(self) -> float:
@@ -94,7 +92,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.array)
+        return fit_util.noise_normalization_from(noise_map=self.noise_map.array, xp=self._xp)
 
     @property
     def log_likelihood(self) -> float:
@@ -115,6 +113,7 @@ def __init__(
         dataset,
         use_mask_in_fit: bool = False,
         dataset_model: DatasetModel = None,
+        xp=np
     ):
         """Class to fit a masked dataset where the dataset's data structures are any dimension.
 
@@ -147,12 +146,13 @@ def __init__(
         self.dataset = dataset
         self.use_mask_in_fit = use_mask_in_fit
         self.dataset_model = dataset_model or DatasetModel()
+        self._xp = xp
 
     @property
     def mask(self) -> Mask2D:
         return self.dataset.mask
 
-    @cached_property
+    @property
     def grids(self) -> GridsInterface:
 
         def subtracted_from(grid, offset):
@@ -196,7 +196,7 @@ def residual_map(self) -> ty.DataLike:
 
         if self.use_mask_in_fit:
             return fit_util.residual_map_with_mask_from(
-                data=self.data, model_data=self.model_data, mask=self.mask
+                data=self.data, model_data=self.model_data, mask=self.mask, xp=self._xp
             )
         return super().residual_map
 
@@ -209,7 +209,7 @@ def normalized_residual_map(self) -> ty.DataLike:
         """
         if self.use_mask_in_fit:
             return fit_util.normalized_residual_map_with_mask_from(
-                residual_map=self.residual_map, noise_map=self.noise_map, mask=self.mask
+                residual_map=self.residual_map, noise_map=self.noise_map, mask=self.mask, xp=self._xp
             )
         return super().normalized_residual_map
 
@@ -222,7 +222,7 @@ def chi_squared_map(self) -> ty.DataLike:
         """
         if self.use_mask_in_fit:
             return fit_util.chi_squared_map_with_mask_from(
-                residual_map=self.residual_map, noise_map=self.noise_map, mask=self.mask
+                residual_map=self.residual_map, noise_map=self.noise_map, mask=self.mask, xp=self._xp
             )
         return super().chi_squared_map
 
@@ -243,7 +243,7 @@ def chi_squared(self) -> float:
 
         if self.use_mask_in_fit:
             return fit_util.chi_squared_with_mask_from(
-                chi_squared_map=self.chi_squared_map, mask=self.mask
+                chi_squared_map=self.chi_squared_map, mask=self.mask, xp=self._xp
             )
         return super().chi_squared
 
@@ -256,7 +256,7 @@ def noise_normalization(self) -> float:
         """
         if self.use_mask_in_fit:
             return fit_util.noise_normalization_with_mask_from(
-                noise_map=self.noise_map, mask=self.mask
+                noise_map=self.noise_map, mask=self.mask, xp=self._xp
             )
         return super().noise_normalization
 

autoarray/fit/fit_imaging.py

@@ -1,4 +1,4 @@
-from typing import Dict, Optional
+import numpy as np
 
 from autoarray.dataset.imaging.dataset import Imaging
 from autoarray.dataset.dataset_model import DatasetModel
@@ -14,6 +14,7 @@ def __init__(
         dataset: Imaging,
         use_mask_in_fit: bool = False,
         dataset_model: DatasetModel = None,
+        xp=np
     ):
         """
         Class to fit a masked imaging dataset.
@@ -49,6 +50,7 @@ def __init__(
             dataset=dataset,
             use_mask_in_fit=use_mask_in_fit,
             dataset_model=dataset_model,
+            xp=xp
         )
 
     @property

autoarray/fit/fit_interferometer.py

@@ -1,5 +1,4 @@
 import numpy as np
-from typing import Dict, Optional
 
 from autoarray.dataset.interferometer.dataset import Interferometer
 
@@ -18,6 +17,7 @@ def __init__(
         dataset: Interferometer,
         dataset_model: DatasetModel = None,
         use_mask_in_fit: bool = False,
+        xp=np
     ):
         """
         Class to fit a masked interferometer dataset.
@@ -58,6 +58,7 @@ def __init__(
             dataset=dataset,
             dataset_model=dataset_model,
             use_mask_in_fit=use_mask_in_fit,
+            xp=xp
         )
 
     @property