jnp revmoed from vectors

Author: Jammy2211

autoarray/structures/grids/grid_1d_util.py

@@ -1,6 +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:
@@ -40,30 +40,27 @@ def convert_grid_1d(
 
     grid_1d = grid_2d_util.convert_grid(grid=grid_1d)
 
-    is_numpy = True if isinstance(grid_1d, np.ndarray) else False
-
     is_native = grid_1d.shape[0] == mask_1d.shape_native[0]
 
     if is_native == store_native:
-        grid_1d = grid_1d
+        return grid_1d
     elif not store_native:
-        grid_1d = grid_1d_slim_from(
+        return grid_1d_slim_from(
             grid_1d_native=grid_1d,
             mask_1d=mask_1d,
         )
-    else:
-        grid_1d = grid_1d_native_from(
-            grid_1d_slim=grid_1d,
-            mask_1d=mask_1d,
-            xp=xp
-        )
-    return np.array(grid_1d) if is_numpy else jnp.array(grid_1d)
+    return grid_1d_native_from(
+        grid_1d_slim=grid_1d,
+        mask_1d=mask_1d,
+        xp=xp
+    )
 
 
 def grid_1d_slim_via_shape_slim_from(
     shape_slim: Tuple[int],
     pixel_scales: ty.PixelScales,
     origin: Tuple[float] = (0.0,),
+    xp=np
 ) -> np.ndarray:
     """
     This routine computes the (x) scaled coordinates at the centre of every pixel defined by a 1D shape of the
@@ -96,13 +93,15 @@ def grid_1d_slim_via_shape_slim_from(
         mask_1d=np.full(fill_value=False, shape=shape_slim),
         pixel_scales=pixel_scales,
         origin=origin,
+        xp=xp
     )
 
 
 def grid_1d_slim_via_mask_from(
     mask_1d: np.ndarray,
     pixel_scales: ty.PixelScales,
     origin: Tuple[float] = (0.0,),
+    xp=np
 ) -> np.ndarray:
     """
     For a grid, every unmasked pixel of its 1D mask with shape (total_pixels,) is divided into a finer uniform
@@ -137,8 +136,8 @@ def grid_1d_slim_via_mask_from(
     centres_scaled = geometry_util.central_scaled_coordinate_1d_from(
         shape_slim=mask_1d.shape, pixel_scales=pixel_scales, origin=origin
     )
-    indices = jnp.arange(mask_1d.shape[0])
-    unmasked = jnp.logical_not(mask_1d)
+    indices = xp.arange(mask_1d.shape[0])
+    unmasked = xp.logical_not(mask_1d)
     coords = (indices - centres_scaled[0]) * pixel_scales[0]
     return coords[unmasked]
 

autoarray/structures/grids/uniform_2d.py

@@ -5,7 +5,6 @@
 from typing import List, Optional, Tuple, Union
 
 from autoconf import conf
-from autoconf import cached_property
 from autoconf.fitsable import ndarray_via_fits_from
 
 from autoarray.mask.mask_2d import Mask2D

autoarray/structures/vectors/irregular.py

@@ -1,6 +1,5 @@
 import logging
 import numpy as np
-import jax.numpy as jnp
 from typing import List, Tuple, Union
 
 from autoarray.structures.vectors.abstract import AbstractVectorYX2D
@@ -120,13 +119,13 @@ def vectors_within_radius(
         squared_distances = self.grid.distances_to_coordinate_from(coordinate=centre)
         mask = squared_distances < radius
 
-        if jnp.all(mask == False):
+        if np.all(mask == False):
             raise exc.VectorYXException(
                 "The input radius removed all vectors / points on the grid."
             )
 
         return VectorYX2DIrregular(
-            values=jnp.array(self.array)[mask], grid=Grid2DIrregular(self.grid[mask])
+            values=np.array(self.array)[mask], grid=Grid2DIrregular(self.grid[mask])
         )
 
     def vectors_within_annulus(

autoarray/structures/vectors/uniform.py

@@ -1,7 +1,6 @@
 import logging
 
 import numpy as np
-import jax.numpy as jnp
 from typing import List, Optional, Tuple, Union
 
 from autoarray.structures.arrays.uniform_2d import Array2D
@@ -394,7 +393,7 @@ 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),
+            values=np.sqrt(self.array[:, 0] ** 2.0 + self.array[:, 1] ** 2.0),
             mask=self.mask,
         )