Added coupled WCS from Will

Author: nabobalis

ndcube/conftest.py

@@ -292,6 +292,47 @@ def wcs_3d_ln_lt_t_rotated():
     return WCS(header=h_rotated)
 
 
+@pytest.fixture
+def wcs_3d_ln_lt_l_coupled():
+    # WCS for a 3D data cube with two celestial axes and one wavelength axis.
+    # The latitudinal dimension is coupled to the third pixel dimension through
+    # a single off diagonal element in the PCij matrix
+    header = {
+        'CTYPE1': 'HPLN-TAN',
+        'CRPIX1': 5,
+        'CDELT1': 5,
+        'CUNIT1': 'arcsec',
+        'CRVAL1': 0.0,
+
+        'CTYPE2': 'HPLT-TAN',
+        'CRPIX2': 5,
+        'CDELT2': 5,
+        'CUNIT2': 'arcsec',
+        'CRVAL2': 0.0,
+
+        'CTYPE3': 'WAVE',
+        'CRPIX3': 1.0,
+        'CDELT3': 1,
+        'CUNIT3': 'Angstrom',
+        'CRVAL3': 1.0,
+
+        'PC1_1': 1,
+        'PC1_2': 0,
+        'PC1_3': 0,
+        'PC2_1': 0,
+        'PC2_2': 1,
+        'PC2_3': -1.0,
+        'PC3_1': 0.0,
+        'PC3_2': 0.0,
+        'PC3_3': 1.0,
+
+        'WCSAXES': 3,
+
+        'DATEREF': "2020-01-01T00:00:00"
+    }
+    return WCS(header=header)
+
+
 ################################################################################
 # Extra and Global Coords Fixtures
 ################################################################################
@@ -519,6 +560,25 @@ def ndcube_3d_rotated(wcs_3d_ln_lt_t_rotated, simple_extra_coords_3d):
     return cube
 
 
+
+@pytest.fixture
+def ndcube_3d_coupled(wcs_3d_ln_lt_l_coupled):
+    shape = (10, 10, 5)
+    wcs_3d_ln_lt_l_coupled.array_shape = shape
+    data = data_nd(shape)
+    mask = data > 0
+    cube = NDCube(
+        data,
+        wcs_3d_ln_lt_l_coupled,
+        mask=mask,
+        uncertainty=data,
+    )
+    base_time = Time('2000-01-01', format='fits', scale='utc')
+    timestamps = Time([base_time + TimeDelta(60 * i, format='sec') for i in range(data.shape[0])])
+    cube.extra_coords.add('time', 0, timestamps)
+    return cube
+
+
 @pytest.fixture
 def ndcube_3d_l_ln_lt_ectime(wcs_3d_lt_ln_l):
     return gen_ndcube_3d_l_ln_lt_ectime(wcs_3d_lt_ln_l,

ndcube/tests/test_ndcube.py

@@ -230,6 +230,40 @@ def test_axis_world_coords_wave_ec(ndcube_3d_l_ln_lt_ectime):
     assert coords[0].shape == (5,)
 
 
+def test_axis_world_coords_wave_coupled_dims(ndcube_3d_coupled):
+    cube = ndcube_3d_coupled
+
+    coords = cube.axis_world_coords('em.wl')
+    assert u.allclose(coords, [1.e-10, 2.e-10, 3.e-10, 4.e-10, 5.e-10,
+                               6.e-10, 7.e-10, 8.e-10, 9.e-10, 1.e-09] * u.m)
+
+    coords = cube.axis_world_coords()
+    assert len(coords) == 2
+    assert isinstance(coords[0], SkyCoord)
+    assert coords[0].shape == (10, 10, 5)
+    assert isinstance(coords[1], SpectralCoord)
+    assert coords[1].shape == (10,)
+
+    coords = cube.axis_world_coords(wcs=cube.combined_wcs)
+    assert len(coords) == 3
+    assert isinstance(coords[0], SkyCoord)
+    assert coords[0].shape == (10, 10, 5)
+    assert isinstance(coords[1], SpectralCoord)
+    assert coords[1].shape == (10,)
+    assert isinstance(coords[2], Time)
+    assert coords[2].shape == (10,)
+
+    coords = cube.axis_world_coords(wcs=cube.extra_coords)
+    assert len(coords) == 1
+    assert isinstance(coords[0], Time)
+    assert coords[0].shape == (10,)
+
+    coords = cube.axis_world_coords_values(wcs=cube.extra_coords)
+    assert len(coords) == 1
+    assert isinstance(coords[0], u.Quantity)
+    assert coords[0].shape == (10,)
+
+
 def test_axis_world_coords_empty_ec(ndcube_3d_l_ln_lt_ectime):
     cube = ndcube_3d_l_ln_lt_ectime
     sub_cube = cube[:, 0]