check p2w in wcs

Author: nabobalis

ndcube/tests/helpers.py

@@ -13,6 +13,7 @@
 from numpy.testing import assert_equal
 
 import astropy
+from astropy.wcs.wcsapi import SlicedLowLevelWCS
 from astropy.wcs.wcsapi.fitswcs import SlicedFITSWCS
 from astropy.wcs.wcsapi.low_level_api import BaseLowLevelWCS
 from astropy.wcs.wcsapi.wrappers.sliced_wcs import sanitize_slices
@@ -95,9 +96,11 @@ def assert_metas_equal(test_input, expected_output):
             assert test_input[key] == expected_output[key]
 
 
-def assert_cubes_equal(test_input, expected_cube):
+def assert_cubes_equal(test_input, expected_cube, check_data=False):
     assert isinstance(test_input, type(expected_cube))
     assert np.all(test_input.mask == expected_cube.mask)
+    if check_data:
+        np.testing.assert_array_equal(test_input.data, expected_cube.data)
     assert_wcs_are_equal(test_input.wcs, expected_cube.wcs)
     if test_input.uncertainty:
         assert test_input.uncertainty.array.shape == expected_cube.uncertainty.array.shape
@@ -110,12 +113,12 @@ def assert_cubes_equal(test_input, expected_cube):
         assert_extra_coords_equal(test_input.extra_coords, expected_cube.extra_coords)
 
 
-def assert_cubesequences_equal(test_input, expected_sequence):
+def assert_cubesequences_equal(test_input, expected_sequence, check_data=False):
     assert isinstance(test_input, type(expected_sequence))
     assert_metas_equal(test_input.meta, expected_sequence.meta)
     assert test_input._common_axis == expected_sequence._common_axis
     for i, cube in enumerate(test_input.data):
-        assert_cubes_equal(cube, expected_sequence.data[i])
+        assert_cubes_equal(cube, expected_sequence.data[i], check_data=check_data)
 
 
 def assert_wcs_are_equal(wcs1, wcs2):
@@ -140,7 +143,13 @@ def assert_wcs_are_equal(wcs1, wcs2):
     assert wcs1.world_axis_units == wcs2.world_axis_units
     assert_equal(wcs1.axis_correlation_matrix, wcs2.axis_correlation_matrix)
     assert wcs1.pixel_bounds == wcs2.pixel_bounds
-
+    if wcs1.pixel_shape is not None:
+        random_idx = np.random.randint(wcs1.pixel_shape,size=[10,wcs1.pixel_n_dim])
+        # SlicedLowLevelWCS vs HighLevelWCS don't have the same pixel_to_world method
+        if isinstance(wcs1, SlicedLowLevelWCS):
+            np.testing.assert_array_equal(wcs1.pixel_to_world_values(*random_idx.T), wcs2.pixel_to_world_values(*random_idx.T))
+        else:
+            np.testing.assert_array_equal(wcs1.pixel_to_world(*random_idx.T), wcs2.pixel_to_world(*random_idx.T))
 
 def create_sliced_wcs(wcs, item, dim):
     """
@@ -152,15 +161,15 @@ def create_sliced_wcs(wcs, item, dim):
     return SlicedFITSWCS(wcs, item)
 
 
-def assert_collections_equal(collection1, collection2):
+def assert_collections_equal(collection1, collection2, check_data=False):
     assert collection1.keys() == collection2.keys()
     assert collection1.aligned_axes == collection2.aligned_axes
     for cube1, cube2 in zip(collection1.values(), collection2.values()):
         # Check cubes are same type.
         assert type(cube1) is type(cube2)
         if isinstance(cube1, NDCube):
-            assert_cubes_equal(cube1, cube2)
+            assert_cubes_equal(cube1, cube2, check_data=check_data)
         elif isinstance(cube1, NDCubeSequence):
-            assert_cubesequences_equal(cube1, cube2)
+            assert_cubesequences_equal(cube1, cube2, check_data=check_data)
         else:
             raise TypeError(f"Unsupported Type in NDCollection: {type(cube1)}")

ndcube/tests/test_ndcube.py

@@ -410,7 +410,7 @@ def test_crop(ndcube_4d_ln_lt_l_t):
     upper_corner = [coord[-1] for coord in intervals]
     expected = cube[1:3, 0:2, 0:2, 0:3]
     output = cube.crop(lower_corner, upper_corner)
-    helpers.assert_cubes_equal(output, expected)
+    helpers.assert_cubes_equal(output, expected, check_data=True)
 
 
 def test_crop_tuple_non_tuple_input(ndcube_2d_ln_lt):
@@ -420,7 +420,7 @@ def test_crop_tuple_non_tuple_input(ndcube_2d_ln_lt):
     upper_corner = SkyCoord(Tx=0.0044444444, Ty=0.0011111111, unit="deg", frame=frame)
     cropped_by_tuples = cube.crop((lower_corner,), (upper_corner,))
     cropped_by_coords = cube.crop(lower_corner, upper_corner)
-    helpers.assert_cubes_equal(cropped_by_tuples, cropped_by_coords)
+    helpers.assert_cubes_equal(cropped_by_tuples, cropped_by_coords, check_data=True)
 
 
 def test_crop_with_nones(ndcube_4d_ln_lt_l_t):
@@ -432,31 +432,31 @@ def test_crop_with_nones(ndcube_4d_ln_lt_l_t):
     upper_corner[0] = interval0[-1]
     expected = cube[:, :, :, 0:3]
     output = cube.crop(lower_corner, upper_corner)
-    helpers.assert_cubes_equal(output, expected)
+    helpers.assert_cubes_equal(output, expected, check_data=True)
 
 
 def test_crop_1d_independent(ndcube_4d_ln_lt_l_t):
     cube_1d = ndcube_4d_ln_lt_l_t[0, 0, :, 0]
     wl_range = SpectralCoord([3e-11, 4.5e-11], unit=u.m)
     expected = cube_1d[0:2]
     output = cube_1d.crop([wl_range[0]], [wl_range[-1]])
-    helpers.assert_cubes_equal(output, expected)
+    helpers.assert_cubes_equal(output, expected, check_data=True)
 
 
 def test_crop_1d_dependent(ndcube_4d_ln_lt_l_t):
     cube_1d = ndcube_4d_ln_lt_l_t[0, :, 0, 0]
     sky_range = cube_1d.wcs.array_index_to_world([0, 1])
     expected = cube_1d[0:2]
     output = cube_1d.crop([sky_range[0]], [sky_range[-1]])
-    helpers.assert_cubes_equal(output, expected)
+    helpers.assert_cubes_equal(output, expected, check_data=True)
 
 
 def test_crop_reduces_dimensionality(ndcube_4d_ln_lt_l_t):
     cube = ndcube_4d_ln_lt_l_t
     point = (None, SpectralCoord([3e-11], unit=u.m), None)
     expected = cube[:, :, 0, :]
     output = cube.crop(point)
-    helpers.assert_cubes_equal(output, expected)
+    helpers.assert_cubes_equal(output, expected, check_data=True)
 
 
 def test_crop_scalar_valuerror(ndcube_2d_ln_lt):
@@ -502,7 +502,7 @@ def test_crop_by_values(ndcube_4d_ln_lt_l_t):
     upper_corner[-1] = upper_corner[-1].to(units[-1])
     expected = cube[1:3, 0:2, 0:2, 0:3]
     output = cube.crop_by_values(lower_corner, upper_corner)
-    helpers.assert_cubes_equal(output, expected)
+    helpers.assert_cubes_equal(output, expected, check_data=True)
 
 
 def test_crop_by_values_with_units(ndcube_4d_ln_lt_l_t):
@@ -518,7 +518,7 @@ def test_crop_by_values_with_units(ndcube_4d_ln_lt_l_t):
     units[0] = None
     expected = ndcube_4d_ln_lt_l_t[1:3, 0:2, 0:2, 0:3]
     output = ndcube_4d_ln_lt_l_t.crop_by_values(lower_corner, upper_corner, units=units)
-    helpers.assert_cubes_equal(output, expected)
+    helpers.assert_cubes_equal(output, expected, check_data=True)
 
 
 def test_crop_by_values_with_equivalent_units(ndcube_2d_ln_lt):
@@ -528,7 +528,7 @@ def test_crop_by_values_with_equivalent_units(ndcube_2d_ln_lt):
     upper_corner = [(coord[-1]*u.deg).to(u.arcsec) for coord in intervals]
     expected = ndcube_2d_ln_lt[0:4, 1:7]
     output = ndcube_2d_ln_lt.crop_by_values(lower_corner, upper_corner)
-    helpers.assert_cubes_equal(output, expected)
+    helpers.assert_cubes_equal(output, expected, check_data=True)
 
 
 def test_crop_by_values_with_nones(ndcube_4d_ln_lt_l_t):
@@ -539,15 +539,15 @@ def test_crop_by_values_with_nones(ndcube_4d_ln_lt_l_t):
     upper_corner[0] = 1.1 * u.min
     expected = cube[:, :, :, 0:3]
     output = cube.crop_by_values(lower_corner, upper_corner)
-    helpers.assert_cubes_equal(output, expected)
+    helpers.assert_cubes_equal(output, expected, check_data=True)
 
 
 def test_crop_by_values_all_nones(ndcube_4d_ln_lt_l_t):
     cube = ndcube_4d_ln_lt_l_t
     lower_corner = [None] * 4
     upper_corner = [None] * 4
     output = cube.crop_by_values(lower_corner, upper_corner)
-    helpers.assert_cubes_equal(output, cube)
+    helpers.assert_cubes_equal(output, cube, check_data=True)
 
 
 def test_crop_by_values_valueerror1(ndcube_4d_ln_lt_l_t):
@@ -594,7 +594,7 @@ def test_crop_by_values_1d_dependent(ndcube_4d_ln_lt_l_t):
     upper_corner = [lat_range[-1] * u.deg, lon_range[-1] * u.deg]
     expected = cube_1d[0:2]
     output = cube_1d.crop_by_values(lower_corner, upper_corner)
-    helpers.assert_cubes_equal(output, expected)
+    helpers.assert_cubes_equal(output, expected, check_data=True)
 
 
 def test_crop_by_extra_coords(ndcube_3d_ln_lt_l_ec_time):
@@ -603,7 +603,7 @@ def test_crop_by_extra_coords(ndcube_3d_ln_lt_l_ec_time):
     upper_corner = (Time("2000-01-01T20:00:00", scale="utc", format="fits"), None)
     output = cube.crop(lower_corner, upper_corner, wcs=cube.extra_coords)
     expected = cube[0]
-    helpers.assert_cubes_equal(output, expected)
+    helpers.assert_cubes_equal(output, expected, check_data=True)
 
 
 def test_crop_by_extra_coords_values(ndcube_3d_ln_lt_l_ec_time):
@@ -612,7 +612,7 @@ def test_crop_by_extra_coords_values(ndcube_3d_ln_lt_l_ec_time):
     upper_corner = (8 * 60 * 60 * u.s, 2 * u.pix)
     output = cube.crop_by_values(lower_corner, upper_corner, wcs=cube.extra_coords)
     expected = cube[0]
-    helpers.assert_cubes_equal(output, expected)
+    helpers.assert_cubes_equal(output, expected, check_data=True)
 
 
 def test_crop_by_extra_coords_all_axes_with_coord(ndcube_3d_ln_lt_l_ec_all_axes):
@@ -624,7 +624,7 @@ def test_crop_by_extra_coords_all_axes_with_coord(ndcube_3d_ln_lt_l_ec_all_axes)
     upper_corner = (interval0[1], interval1[1], interval2[1])
     output = cube.crop(lower_corner, upper_corner, wcs=cube.extra_coords)
     expected = cube[0, 0:2, 1:4]
-    helpers.assert_cubes_equal(output, expected)
+    helpers.assert_cubes_equal(output, expected, check_data=True)
 
 
 def test_crop_by_extra_coords_values_all_axes_with_coord(ndcube_3d_ln_lt_l_ec_all_axes):
@@ -636,7 +636,7 @@ def test_crop_by_extra_coords_values_all_axes_with_coord(ndcube_3d_ln_lt_l_ec_al
     upper_corner = (interval0[1], interval1[1], interval2[1])
     output = cube.crop_by_values(lower_corner, upper_corner, wcs=cube.extra_coords)
     expected = cube[0, 0:2, 1:4]
-    helpers.assert_cubes_equal(output, expected)
+    helpers.assert_cubes_equal(output, expected, check_data=True)
 
 
 def test_crop_by_extra_coords_shared_axis(ndcube_3d_ln_lt_l_ec_sharing_axis):
@@ -645,7 +645,7 @@ def test_crop_by_extra_coords_shared_axis(ndcube_3d_ln_lt_l_ec_sharing_axis):
     upper_corner = (2 * u.m, 2 * u.keV)
     output = cube.crop(lower_corner, upper_corner, wcs=cube.extra_coords)
     expected = cube[:, 1:3]
-    helpers.assert_cubes_equal(output, expected)
+    helpers.assert_cubes_equal(output, expected, check_data=True)
 
 
 def test_crop_by_extra_coords_values_shared_axis(ndcube_3d_ln_lt_l_ec_sharing_axis):
@@ -654,7 +654,7 @@ def test_crop_by_extra_coords_values_shared_axis(ndcube_3d_ln_lt_l_ec_sharing_ax
     upper_corner = (2 * u.m, 2 * u.keV)
     output = cube.crop_by_values(lower_corner, upper_corner, wcs=cube.extra_coords)
     expected = cube[:, 1:3]
-    helpers.assert_cubes_equal(output, expected)
+    helpers.assert_cubes_equal(output, expected, check_data=True)
 
 
 def test_crop_rotated_celestial(ndcube_4d_ln_lt_l_t):