lots of JAX stuff cleaned up

Author: Jammy2211

autoarray/inversion/inversion/abstract.py

@@ -75,6 +75,8 @@ def __init__(
 
         self.xp = xp
 
+
+
     @property
     def data(self):
         return self.dataset.data

autoarray/structures/grids/uniform_2d.py

@@ -1,5 +1,4 @@
 from __future__ import annotations
-import jax.numpy as jnp
 import numpy as np
 from pathlib import Path
 from typing import List, Optional, Tuple, Union
@@ -28,6 +27,7 @@ def __init__(
         store_native: bool = False,
         over_sample_size: Union[int, Array2D] = 4,
         over_sampled: Optional[Grid2D] = None,
+        xp=np,
         *args,
         **kwargs,
     ):
@@ -163,9 +163,10 @@ def __init__(
             grid_2d=values,
             mask_2d=mask,
             store_native=store_native,
+            xp=xp
         )
 
-        super().__init__(values)
+        super().__init__(values, xp=xp)
 
         self.mask = mask
 
@@ -690,7 +691,7 @@ def blurring_grid_from(
             over_sample_size=over_sample_size,
         )
 
-    def subtracted_from(self, offset: Tuple[(float, float), np.ndarray]) -> "Grid2D":
+    def subtracted_from(self, offset: Tuple[(float, float), np.ndarray], xp=np) -> "Grid2D":
 
         mask = Mask2D(
             mask=self.mask,
@@ -699,10 +700,10 @@ def subtracted_from(self, offset: Tuple[(float, float), np.ndarray]) -> "Grid2D"
         )
 
         return Grid2D(
-            values=self - jnp.array(offset),
+            values=self - xp.array(offset),
             mask=mask,
             over_sample_size=self.over_sample_size,
-            over_sampled=self.over_sampled - jnp.array(offset),
+            over_sampled=self.over_sampled - xp.array(offset),
         )
 
     @property
@@ -847,7 +848,7 @@ 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(
+        squared_distances = self.xp.square(self.array[:, 0] - coordinate[0]) + self.xp.square(
             self.array[:, 1] - coordinate[1]
         )
 
@@ -867,7 +868,7 @@ def distances_to_coordinate_from(
         squared_distance = self.squared_distances_to_coordinate_from(
             coordinate=coordinate
         )
-        distances = jnp.sqrt(squared_distance.array)
+        distances = self.xp.sqrt(squared_distance.array)
         return Array2D(values=distances, mask=self.mask)
 
     def grid_2d_radial_projected_shape_slim_from(
@@ -1017,50 +1018,32 @@ def shape_native_scaled_interior(self) -> Tuple[float, float]:
         of the grid's (y,x) values, whereas the `shape_native_scaled` uses the uniform geometry of the grid and its
         ``pixel_scales``, which means it has a buffer at each edge of half a ``pixel_scale``.
         """
-        if isinstance(self, jnp.ndarray):
-            return (
-                np.amax(self.array[:, 0]) - np.amin(self.array[:, 0]),
-                np.amax(self.array[:, 1]) - np.amin(self.array[:, 1]),
-            )
-        else:
-            return (
-                np.amax(self[:, 0]) - np.amin(self[:, 0]),
-                np.amax(self[:, 1]) - np.amin(self[:, 1]),
-            )
+        return (
+            np.amax(self[:, 0]) - np.amin(self[:, 0]),
+            np.amax(self[:, 1]) - np.amin(self[:, 1]),
+        )
 
     @property
     def scaled_minima(self) -> Tuple:
         """
         The (y,x) minimum values of the grid in scaled units, buffed such that their extent is further than the grid's
         extent.
         """
-        if isinstance(self, jnp.ndarray):
-            return (
-                jnp.amin(self.array[:, 0]).astype("float"),
-                jnp.amin(self.array[:, 1]).astype("float"),
-            )
-        else:
-            return (
-                np.amin(self[:, 0]).astype("float"),
-                np.amin(self[:, 1]).astype("float"),
-            )
+        return (
+            np.amin(self[:, 0]).astype("float"),
+            np.amin(self[:, 1]).astype("float"),
+        )
 
     @property
     def scaled_maxima(self) -> Tuple:
         """
         The (y,x) maximum values of the grid in scaled units, buffed such that their extent is further than the grid's
         extent.
         """
-        if isinstance(self, jnp.ndarray):
-            return (
-                jnp.amax(self.array[:, 0]).astype("float"),
-                jnp.amax(self.array[:, 1]).astype("float"),
-            )
-        else:
-            return (
-                np.amax(self[:, 0]).astype("float"),
-                np.amax(self[:, 1]).astype("float"),
-            )
+        return (
+            np.amax(self[:, 0]).astype("float"),
+            np.amax(self[:, 1]).astype("float"),
+        )
 
     def extent_with_buffer_from(self, buffer: float = 1.0e-8) -> List[float]:
         """

test_autoarray/inversion/inversion/test_factory.py

@@ -114,6 +114,8 @@ def test__inversion_imaging__via_mapper(
     assert inversion.log_det_curvature_reg_matrix_term == pytest.approx(10.6674, 1.0e-4)
     assert inversion.mapped_reconstructed_image == pytest.approx(np.ones(9), 1.0e-4)
 
+    fff
+
 
 def test__inversion_imaging__via_regularizations(
     masked_imaging_7x7_no_blur,