fix dataset unit tests

Author: Jammy2211

autoarray/dataset/imaging/dataset.py

@@ -135,8 +135,6 @@ def __init__(
                     f"the image and noise-map yourself."
                 )
 
-        print(type(data.array))
-
         super().__init__(
             data=data,
             noise_map=noise_map,

autoarray/structures/grids/uniform_2d.py

@@ -320,7 +320,7 @@ def from_yx_1d(
             x = np.asarray(x)
 
         return cls.no_mask(
-            values=jnp.stack((y, x), axis=-1),
+            values=np.stack((y, x), axis=-1),
             shape_native=shape_native,
             pixel_scales=pixel_scales,
             origin=origin,
@@ -547,14 +547,14 @@ def from_mask(
             The mask whose masked pixels are used to setup the grid.
         """
 
-        grid_1d = grid_2d_util.grid_2d_slim_via_mask_from(
-            mask_2d=mask._array,
+        grid_2d = grid_2d_util.grid_2d_slim_via_mask_from(
+            mask_2d=mask,
             pixel_scales=mask.pixel_scales,
             origin=mask.origin,
         )
 
         return Grid2D(
-            values=grid_1d,
+            values=np.array(grid_2d),
             mask=mask,
             over_sample_size=over_sample_size,
         )

test_autoarray/dataset/abstract/test_dataset.py

@@ -32,9 +32,7 @@ def test__signal_to_noise_map():
 
     dataset = ds.AbstractDataset(data=array, noise_map=noise_map)
 
-    assert (
-        dataset.signal_to_noise_map.native == np.array([[0.1, 0.2], [0.1, 1.0]])
-    ).all()
+    assert  dataset.signal_to_noise_map.native == pytest.approx(np.array([[0.1, 0.2], [0.1, 1.0]]), 1.0e-4)
     assert dataset.signal_to_noise_max == 1.0
 
     array = aa.Array2D.no_mask([[-1.0, 2.0], [3.0, -4.0]], pixel_scales=1.0)
@@ -43,9 +41,7 @@ def test__signal_to_noise_map():
 
     dataset = ds.AbstractDataset(data=array, noise_map=noise_map)
 
-    assert (
-        dataset.signal_to_noise_map.native == np.array([[0.0, 0.2], [0.1, 0.0]])
-    ).all()
+    assert  dataset.signal_to_noise_map.native == pytest.approx(np.array([[0.0, 0.2], [0.1, 0.0]]), 1.0e-4)
     assert dataset.signal_to_noise_max == 0.2
 
 
@@ -149,8 +145,8 @@ def test__apply_over_sampling(image_7x7, noise_map_7x7):
     grid_sub_2 = dataset_7x7.grids.lp
     grids_pixelization_sub_2 = dataset_7x7.grids.pixelization
 
-    dataset_7x7.grids.lp = dataset_7x7.grids.lp.at[0, 0].set(100.0)
-    dataset_7x7.grids.pixelization = dataset_7x7.grids.pixelization.at[0, 0].set(100.0)
+    dataset_7x7.grids.__dict__["lp"][0][0] = 100.0
+    dataset_7x7.grids.__dict__["pixelization"][0][0] = 100.0
 
     assert dataset_7x7.grids.lp[0][0] == pytest.approx(100.0, 1.0e-4)
     assert dataset_7x7.grids.pixelization[0][0] == pytest.approx(100.0, 1.0e-4)

test_autoarray/dataset/imaging/test_dataset.py

@@ -78,7 +78,7 @@ def test__from_fits():
     )
 
     assert (dataset.data.native == np.ones((3, 3))).all()
-    assert (dataset.psf.native == (1.0 / 9.0) * np.ones((3, 3))).all()
+    assert dataset.psf.native == pytest.approx((1.0 / 9.0) * np.ones((3, 3)), 1.0e-4)
     assert (dataset.noise_map.native == 3.0 * np.ones((3, 3))).all()
 
     assert dataset.pixel_scales == (0.1, 0.1)
@@ -96,7 +96,7 @@ def test__from_fits():
     )
 
     assert (dataset.data.native == np.ones((3, 3))).all()
-    assert (dataset.psf.native == (1.0 / 9.0) * np.ones((3, 3))).all()
+    assert dataset.psf.native == pytest.approx((1.0 / 9.0) * np.ones((3, 3)), 1.0e-4)
     assert (dataset.noise_map.native == 3.0 * np.ones((3, 3))).all()
 
     assert dataset.pixel_scales == (0.1, 0.1)

test_autoarray/dataset/test_preprocess.py

@@ -133,15 +133,12 @@ def test__noise_map_from_image_exposure_time_map():
         data_eps=image, exposure_time_map=exposure_time_map
     )
 
-    assert (
-        poisson_noise_map.native
-        == np.array(
+    assert poisson_noise_map.native == pytest.approx(np.array(
             [
                 [np.sqrt(5.0), np.sqrt(6.0) / 2.0],
                 [np.sqrt(30.0) / 3.0, np.sqrt(80.0) / 4.0],
             ]
-        )
-    ).all()
+        ), 1.0e-4)
 
 
 def test__noise_map_from_image_exposure_time_map_and_background_noise_map():