Resampled wcs

changelog/449.feature.rst

@@ -0,0 +1 @@
+Introduce optional offset between old and new pixel grids in `~ndcube.wcs.wrappers.resampled_wcs.ResampledLowLevelWCS`.

ndcube/wcs/wrappers/resampled_wcs.py

@@ -13,34 +13,73 @@ class ResampledLowLevelWCS(BaseWCSWrapper):
     ----------
     wcs : `~astropy.wcs.wcsapi.BaseLowLevelWCS`
         The original WCS for which to reorder axes
-    factor : int or float or iterable
+
+    factor : `int` or `float` or iterable of the same
         The factor by which to increase the pixel size for each pixel
         axis. If a scalar, the same factor is used for all axes.
+
+    offset: `int` or `float` or iterable of the same
+        The location on the underlying pixel grid which corresponds
+        to zero on the top level pixel grid. If a scalar, the grid will be
+        shifted by the same amount in all dimensions.
     """
-    def __init__(self, wcs, factor):
+    def __init__(self, wcs, factor, offset=0):
         self._wcs = wcs
         if np.isscalar(factor):
             factor = [factor] * self.pixel_n_dim
-        self._factor = factor
+        self._factor = np.array(factor)
+        if len(self._factor) !=  self.pixel_n_dim:
+            raise ValueError(f"Length of factor must equal number of dimensions {self.pixel_n_dim}.")
+        if np.isscalar(offset):
+            offset = [offset] * self.pixel_n_dim
+        self._offset = np.array(offset)
+        if len(self._offset) != self.pixel_n_dim:
+            raise ValueError(f"Length of offset must equal number of dimensions {self.pixel_n_dim}.")
+
+    def _top_to_underlying_pixels(self, top_pixels):
+        # Convert user-facing pixel indices to the pixel grid of underlying WCS.
+        factor = self._pad_dims(self._factor, top_pixels.ndim)
+        offset = self._pad_dims(self._offset, top_pixels.ndim)
+        return top_pixels * factor + offset
+
+    def _underlying_to_top_pixels(self, underlying_pixels):
+        # Convert pixel indices of underlying pixel grid to user-facing grid.
+        factor = self._pad_dims(self._factor, underlying_pixels.ndim)
+        offset = self._pad_dims(self._offset, underlying_pixels.ndim)
+        return (underlying_pixels - offset) / factor
+
+    def _pad_dims(self, arr, ndim):
+        # Pad array with trailing degenerate dimensions.
+        # This make scaling with pixel arrays easier.
+        shape = np.ones(ndim, dtype=int)
+        shape[0] = len(arr)
+        return arr.reshape(tuple(shape))
 
     def pixel_to_world_values(self, *pixel_arrays):
-        pixel_arrays = [np.asarray(pixel_arrays[i]) * self._factor[i]
-                        for i in range(self.pixel_n_dim)]
-        return self._wcs.pixel_to_world_values(*pixel_arrays)
+        underlying_pixel_arrays = self._top_to_underlying_pixels(np.asarray(pixel_arrays))
+        return self._wcs.pixel_to_world_values(*underlying_pixel_arrays)
 
     def world_to_pixel_values(self, *world_arrays):
-        pixel_arrays = self._wcs.world_to_pixel_values(*world_arrays)
-        pixel_arrays = [np.asarray(pixel_arrays[i]) / self._factor[i]
-                        for i in range(self.pixel_n_dim)]
-        return pixel_arrays
+        underlying_pixel_arrays = self._wcs.world_to_pixel_values(*world_arrays)
+        top_pixel_arrays = self._underlying_to_top_pixels(np.asarray(underlying_pixel_arrays))
+        return tuple(array for array in top_pixel_arrays)
 
     @property
     def pixel_shape(self):
-        return tuple(self._wcs.pixel_shape[i] / self._factor[i]
-                     for i in range(self.pixel_n_dim))
+        # Return pixel shape of resampled grid.
+        # Where shape is an integer, return an int type as its required for some uses.
+        if self._wcs.pixel_shape is None:
+            return self._wcs.pixel_shape
+        underlying_shape = np.asarray(self._wcs.pixel_shape)
+        int_elements = np.isclose(np.mod(underlying_shape, self._factor), 0,
+                                  atol=np.finfo(float).resolution)
+        pixel_shape = underlying_shape / self._factor
+        return tuple(int(np.rint(i)) if is_int else i
+                     for i, is_int in zip(pixel_shape, int_elements))
 
     @property
     def pixel_bounds(self):
-        return tuple((self._wcs.pixel_bounds[i][0] / self._factor[i],
-                      self._wcs.pixel_bounds[i][1] / self._factor[i])
-                     for i in range(self.pixel_n_dim))
+        if self._wcs.pixel_bounds is None:
+            return self._wcs.pixel_bounds
+        top_level_bounds = self._underlying_to_top_pixels(np.asarray(self._wcs.pixel_bounds))
+        return [tuple(bounds) for bounds in top_level_bounds]

ndcube/wcs/wrappers/tests/test_resampled_wcs.py

@@ -1,3 +1,5 @@
+import numbers
+
 import numpy as np
 import pytest
 from astropy import units as u
@@ -108,3 +110,50 @@ def test_scalar_factor(celestial_wcs):
     assert_allclose(wcs.array_index_to_world_values(*pixel_scalar[::-1]), world_scalar)
     assert_allclose(wcs.world_to_pixel_values(*world_scalar), pixel_scalar)
     assert_allclose(wcs.world_to_array_index_values(*world_scalar), [4, 2])
+
+
+@pytest.mark.parametrize('celestial_wcs',
+                         ['celestial_2d_ape14_wcs', 'celestial_2d_fitswcs'],
+                         indirect=True)
+def test_offset(celestial_wcs):
+    offset = 1
+    factor = 2
+    wcs = ResampledLowLevelWCS(celestial_wcs, factor, offset=offset)
+
+    pixel_scalar = (2.3, 4.3)
+    world_scalar = celestial_wcs.pixel_to_world_values(*[p * factor + offset
+                                                         for p in pixel_scalar])
+
+    assert_allclose(wcs.pixel_to_world_values(*pixel_scalar), world_scalar)
+    assert_allclose(wcs.array_index_to_world_values(*pixel_scalar[::-1]), world_scalar)
+    assert_allclose(wcs.world_to_pixel_values(*world_scalar), pixel_scalar)
+    assert_allclose(wcs.world_to_array_index_values(*world_scalar), [4, 2])
+
+
+@pytest.mark.parametrize('celestial_wcs',
+                         ['celestial_2d_ape14_wcs'],
+                         indirect=True)
+def test_factor_wrong_length_error(celestial_wcs):
+    with pytest.raises(ValueError):
+        ResampledLowLevelWCS(celestial_wcs, [2] * 3)
+
+
+@pytest.mark.parametrize('celestial_wcs',
+                         ['celestial_2d_ape14_wcs'],
+                         indirect=True)
+def test_scalar_wrong_length_error(celestial_wcs):
+    with pytest.raises(ValueError):
+        wcs = ResampledLowLevelWCS(celestial_wcs, 2, offset=[1] * 3)
+
+
+@pytest.mark.parametrize('celestial_wcs',
+                         ['celestial_2d_ape14_wcs', 'celestial_2d_fitswcs'],
+                         indirect=True)
+def test_int_fraction_pixel_shape(celestial_wcs):
+    # Some fractional factors are not representable by exact floats, e.g. 1/3.
+    # However, it is still desirable for the pixel shape to return ints in these cases.
+    # This test checks that this is the case.
+    wcs = ResampledLowLevelWCS(celestial_wcs, 1/3)
+    assert wcs.pixel_shape == (18, 21)
+    for dim in wcs.pixel_shape:
+        assert isinstance(dim, numbers.Integral)