Resample all LMS cube planes in one step

scopesim/effects/metis_lms_trace_list.py

@@ -7,7 +7,6 @@
 
 from tqdm.auto import tqdm
 import numpy as np
-from scipy.interpolate import RectBivariateSpline
 
 from astropy.io import fits
 from astropy.io import ascii as ioascii
@@ -29,6 +28,39 @@
 logger = get_logger(__name__)
 
 
+def interpolate_cube_planes(fovcube, yfov, xfov):
+    """
+    Bilinear resampling of every plane of a cube at the same (y, x) points.
+
+    Equivalent to building a ``RectBivariateSpline(kx=1, ky=1)`` over each
+    plane ``fovcube[k]`` and evaluating it at ``(yfov, xfov)`` -- but since
+    the sample coordinates are identical for every plane, the bilinear
+    indices and weights are computed once and applied to the whole cube.
+    Coordinates outside a plane are clamped to the edges, as
+    ``RectBivariateSpline`` does.
+
+    Parameters
+    ----------
+    fovcube : ndarray of shape (n_z, n_y, n_x)
+    yfov, xfov : ndarrays of identical shape, pixel coordinates
+
+    Returns
+    -------
+    ndarray of shape (n_z, *yfov.shape)
+    """
+    n_y, n_x = fovcube.shape[1:]
+    ycl = np.clip(yfov, 0, n_y - 1)
+    xcl = np.clip(xfov, 0, n_x - 1)
+    y_0 = np.clip(ycl.astype(int), 0, n_y - 2)
+    x_0 = np.clip(xcl.astype(int), 0, n_x - 2)
+    w_y = ycl - y_0
+    w_x = xcl - x_0
+    return ((1 - w_y) * (1 - w_x) * fovcube[:, y_0, x_0]
+            + (1 - w_y) * w_x * fovcube[:, y_0, x_0 + 1]
+            + w_y * (1 - w_x) * fovcube[:, y_0 + 1, x_0]
+            + w_y * w_x * fovcube[:, y_0 + 1, x_0 + 1])
+
+
 @lru_cache(maxsize=8)
 def _read_detector_layout(filename):
     """Read (and cache) a detector layout file.
@@ -124,12 +156,9 @@ def apply_to(self, obj, **kwargs):
                 # FOV pixel coordinates for the slice
                 xfov, yfov = fovwcs_spat.all_world2pix(xworld, yworld, 0)
 
-                slicecube = np.zeros((n_z, ny_slice, n_x))
-                for islice in range(n_z):
-                    ifov = RectBivariateSpline(np.arange(n_y),
-                                               np.arange(n_x),
-                                               fovcube[islice], kx=1, ky=1)
-                    slicecube[islice] = ifov(yfov, xfov, grid=False)
+                # Resample all wavelength planes at once; the (yfov, xfov)
+                # sample coordinates are the same for every plane
+                slicecube = interpolate_cube_planes(fovcube, yfov, xfov)
 
                 slicefov = FieldOfView3D(obj.header,
                                          [obj.meta["wave_min"],

scopesim/tests/tests_effects/test_MetisLMSTraceList.py

@@ -16,6 +16,32 @@ def patch_mock_path_metis(mock_dir):
     with patch("scopesim.rc.__search_path__", [metis_dir]):
         yield
 
+class TestInterpolateCubePlanes:
+    def test_matches_spline_per_plane_reference(self):
+        """The one-shot bilinear gather must reproduce the previous
+        RectBivariateSpline-per-plane evaluation exactly."""
+        import numpy as np
+        from scipy.interpolate import RectBivariateSpline
+        from scopesim.effects.metis_lms_trace_list import (
+            interpolate_cube_planes)
+
+        rng = np.random.default_rng(5)
+        n_z, n_y, n_x = 20, 15, 17
+        cube = rng.random((n_z, n_y, n_x))
+        # sample coordinates deliberately extend beyond the plane edges
+        yfov = rng.uniform(-1, n_y, (4, 25))
+        xfov = rng.uniform(-1, n_x, (4, 25))
+
+        reference = np.zeros((n_z, 4, 25))
+        for k in range(n_z):
+            spline = RectBivariateSpline(np.arange(n_y), np.arange(n_x),
+                                         cube[k], kx=1, ky=1)
+            reference[k] = spline(yfov, xfov, grid=False)
+
+        result = interpolate_cube_planes(cube, yfov, xfov)
+        assert_allclose(result, reference, rtol=1e-12)
+
+
 class TestDetectorLayoutCache:
     def test_layout_file_is_read_only_once(self, mock_dir, monkeypatch):
         from scopesim.effects import metis_lms_trace_list as mlt