mask array 2d choose type based on if input is jax or numpy

Author: Jammy2211

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,

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:

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(

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"),

test_autoarray/layout/test_region.py

@@ -31,15 +31,15 @@ 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()
 
     arr_1d = aa.Array1D.no_mask(values=np.array([1.0, 2.0, 3.0, 4.0]), pixel_scales=1.0)
 
     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()
 

test_autoarray/structures/arrays/files/array/output_test/array.fits

@@ -150,31 +150,31 @@ 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)
     kernel_2d = kernel_2d.rescaled_with_odd_dimensions_from(
         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)
     kernel_2d = kernel_2d.rescaled_with_odd_dimensions_from(
         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)
     kernel_2d = kernel_2d.rescaled_with_odd_dimensions_from(
         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,23 +201,23 @@ 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(
         rescale_factor=2.0, normalize=True
     )
     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(
         rescale_factor=0.5, normalize=True
     )
 
     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,15 +227,15 @@ 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(
         rescale_factor=0.5, normalize=True
     )
 
     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():

test_autoarray/structures/arrays/test_kernel_2d.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])"