Fix handling of data in "nearest" trajectory interpolate

docs/src/whatsnew/latest.rst

@@ -114,6 +114,11 @@ This document explains the changes made to Iris for this release
 #. `@bjlittle`_ and `@trexfeathers`_ (reviewer) fixed an issue which prevented
    uncompressed PP fields with additional trailing padded words in the field
    data to be loaded and saved. (:pull:`5058`)
+
+#. `@lbdreyer`_ and `@trexfeathers`_ (reviewer) fixed the handling of data when
+   regridding with :class:`~iris.analysis.UnstructuredNearest` or calling
+   :func:`~iris.analysis.trajectory.interpolate` such that the data type and mask is
+   preserved. (:issue:`4463`, :pull:`5062`)
 
 
 💣 Incompatible Changes

lib/iris/analysis/trajectory.py

@@ -443,21 +443,7 @@ def interpolate(cube, sample_points, method=None):
         ]
 
         # Apply the fancy indexing to get all the result data points.
-        source_data = source_data[tuple(fancy_source_indices)]
-
-        # "Fix" problems with missing datapoints producing odd values
-        # when copied from a masked into an unmasked array.
-        # TODO: proper masked data handling.
-        if np.ma.isMaskedArray(source_data):
-            # This is **not** proper mask handling, because we cannot produce a
-            # masked result, but it ensures we use a "filled" version of the
-            # input in this case.
-            source_data = source_data.filled()
-        new_cube.data[:] = source_data
-        # NOTE: we assign to "new_cube.data[:]" and *not* just "new_cube.data",
-        # because the existing code produces a default dtype from 'np.empty'
-        # instead of preserving the input dtype.
-        # TODO: maybe this should be fixed -- i.e. to preserve input dtype ??
+        new_cube.data = source_data[tuple(fancy_source_indices)]
 
         # Fill in the empty squashed (non derived) coords.
         column_coords = [

lib/iris/tests/results/trajectory/hybrid_height.cml

@@ -60,7 +60,7 @@
     <cellMethods/>
     <data checksum="0x2acdccca" dtype="int64" shape="(3, 4, 5, 6)"/>
   </cube>
-  <cube dtype="float64" standard_name="air_temperature" units="K">
+  <cube dtype="int64" standard_name="air_temperature" units="K">
     <coords>
       <coord datadims="[1, 2]">
         <auxCoord id="9041e969" points="[[5090, 5740, 6090, 6440],
@@ -95,6 +95,6 @@
       </coord>
     </coords>
     <cellMethods/>
-    <data checksum="0x11cdc1c8" dtype="float64" shape="(3, 4, 4)"/>
+    <data checksum="0x07fcebe7" dtype="int64" shape="(3, 4, 4)"/>
   </cube>
 </cubes>

lib/iris/tests/results/trajectory/tri_polar_latitude_slice.cml

@@ -1,6 +1,6 @@
 <?xml version="1.0" ?>
 <cubes xmlns="urn:x-iris:cubeml-0.2">
-  <cube dtype="float64" long_name="Temperature" standard_name="sea_water_potential_temperature" units="degC" var_name="votemper">
+  <cube dtype="float32" long_name="Temperature" standard_name="sea_water_potential_temperature" units="degC" var_name="votemper">
     <attributes>
       <attribute name="Conventions" value="CF-1.1"/>
       <attribute name="DOMAIN_DIM_N001" value="x"/>
@@ -144,6 +144,6 @@
         <coord name="time"/>
       </cellMethod>
     </cellMethods>
-    <data checksum="0x3a8195d5" dtype="float64" shape="(1, 31, 85)"/>
+    <data checksum="0x49d623f6" dtype="float32" mask_checksum="0x117b1ebc" shape="(1, 31, 85)"/>
   </cube>
 </cubes>

lib/iris/tests/unit/analysis/trajectory/test_interpolate.py

@@ -12,7 +12,10 @@
 # importing anything else.
 import iris.tests as tests  # isort:skip
 
+from collections import namedtuple
+
 import numpy as np
+import pytest
 
 from iris.analysis.trajectory import interpolate
 from iris.coords import AuxCoord, DimCoord
@@ -38,62 +41,59 @@ def test_unknown_method(self):
             interpolate(cube, sample_point, method="linekar")
 
 
-class TestNearest(tests.IrisTest):
+class TestNearest:
     # Test interpolation with 'nearest' method.
     # This is basically a wrapper to the routine:
     #   'analysis._interpolate_private._nearest_neighbour_indices_ndcoords'.
     # That has its own test, so we don't test the basic calculation
     # exhaustively here.  Instead we check the way it handles the source and
     # result cubes (especially coordinates).
-
-    def setUp(self):
+    @pytest.fixture
+    def src_cube(self):
         cube = iris.tests.stock.simple_3d()
         # Actually, this cube *isn't* terribly realistic, as the lat+lon coords
         # have integer type, which in this case produces some peculiar results.
         # Let's fix that (and not bother to test the peculiar behaviour).
         for coord_name in ("longitude", "latitude"):
             coord = cube.coord(coord_name)
             coord.points = coord.points.astype(float)
-        self.test_cube = cube
+        return cube
+
+    @pytest.fixture
+    def single_point(self, src_cube):
         # Define coordinates for a single-point testcase.
         y_val, x_val = 0, -90
-        # Work out cube indices of the testpoint.
-        self.single_point_iy = np.where(
-            cube.coord("latitude").points == y_val
-        )[0][0]
-        self.single_point_ix = np.where(
-            cube.coord("longitude").points == x_val
-        )[0][0]
+
         # Use slightly-different values to test nearest-neighbour operation.
-        self.single_sample_point = [
+        sample_point = [
             ("latitude", [y_val + 19.23]),
             ("longitude", [x_val - 17.54]),
         ]
 
-    def test_single_point_same_cube(self):
-        # Check exact result matching for a single point.
-        cube = self.test_cube
-        result = interpolate(cube, self.single_sample_point, method="nearest")
-        # Check that the result is a single trajectory point, exactly equal to
-        # the expected part of the original data.
-        self.assertEqual(result.shape[-1], 1)
-        result = result[..., 0]
-        expected = cube[:, self.single_point_iy, self.single_point_ix]
-        self.assertEqual(result, expected)
+        # Work out cube indices of the testpoint.
+        single_point_iy = np.where(src_cube.coord("latitude").points == y_val)[
+            0
+        ][0]
+        single_point_ix = np.where(
+            src_cube.coord("longitude").points == x_val
+        )[0][0]
 
-    def test_multi_point_same_cube(self):
-        # Check an exact result for multiple points.
-        cube = self.test_cube
+        point = namedtuple("point", "ix iy sample_point")
+        return point(single_point_ix, single_point_iy, sample_point)
+
+    @pytest.fixture
+    def multi_sample_points(self):
         # Use latitude selection to recreate a whole row of the original cube.
-        sample_points = [
+        return [
             ("longitude", [-180, -90, 0, 90]),
             ("latitude", [0, 0, 0, 0]),
         ]
-        result = interpolate(cube, sample_points, method="nearest")
 
+    @pytest.fixture
+    def expected_multipoint_cube(self, src_cube):
         # The result should be identical to a single latitude section of the
         # original, but with modified coords (latitude has 4 repeated zeros).
-        expected = cube[:, 1, :]
+        expected = src_cube[:, 1, :]
         # Result 'longitude' is now an aux coord.
         co_x = expected.coord("longitude")
         expected.remove_coord(co_x)
@@ -104,36 +104,86 @@ def test_multi_point_same_cube(self):
             [0, 0, 0, 0], standard_name="latitude", units="degrees"
         )
         expected.add_aux_coord(co_y, 1)
-        self.assertEqual(result, expected)
 
-    def test_aux_coord_noninterpolation_dim(self):
+        return expected
+
+    def test_single_point_same_cube(self, src_cube, single_point):
+        # Check exact result matching for a single point.
+        result = interpolate(
+            src_cube, single_point.sample_point, method="nearest"
+        )
+        # Check that the result is a single trajectory point, exactly equal to
+        # the expected part of the original data.
+        assert result.shape[-1] == 1
+        result = result[..., 0]
+        expected = src_cube[:, single_point.iy, single_point.ix]
+        assert result == expected
+
+    def test_multi_point_same_cube(
+        self, src_cube, multi_sample_points, expected_multipoint_cube
+    ):
+        # Check an exact result for multiple points.
+        result = interpolate(src_cube, multi_sample_points, method="nearest")
+        assert result == expected_multipoint_cube
+
+    def test_mask_preserved(
+        self, src_cube, multi_sample_points, expected_multipoint_cube
+    ):
+        mask = np.zeros_like(src_cube.data)
+        mask[:, :, 1] = 1
+        src_cube.data = np.ma.array(src_cube.data, mask=mask)
+
+        expected_multipoint_cube.data = np.ma.array(
+            expected_multipoint_cube.data, mask=mask[:, 0]
+        )
+
+        result = interpolate(src_cube, multi_sample_points, method="nearest")
+        assert result == expected_multipoint_cube
+        assert np.allclose(
+            result.data.mask, expected_multipoint_cube.data.mask
+        )
+
+    def test_dtype_preserved(
+        self, src_cube, multi_sample_points, expected_multipoint_cube
+    ):
+        src_cube.data = src_cube.data.astype(np.int16)
+
+        result = interpolate(src_cube, multi_sample_points, method="nearest")
+        assert result == expected_multipoint_cube
+        assert np.allclose(result.data, expected_multipoint_cube.data)
+        assert result.data.dtype == np.int16
+
+    def test_aux_coord_noninterpolation_dim(self, src_cube, single_point):
         # Check exact result with an aux-coord mapped to an uninterpolated dim.
-        cube = self.test_cube
-        cube.add_aux_coord(DimCoord([17, 19], long_name="aux0"), 0)
+        src_cube.add_aux_coord(DimCoord([17, 19], long_name="aux0"), 0)
 
         # The result cube should exactly equal a single source point.
-        result = interpolate(cube, self.single_sample_point, method="nearest")
-        self.assertEqual(result.shape[-1], 1)
+        result = interpolate(
+            src_cube, single_point.sample_point, method="nearest"
+        )
+        assert result.shape[-1] == 1
         result = result[..., 0]
-        expected = cube[:, self.single_point_iy, self.single_point_ix]
-        self.assertEqual(result, expected)
+        expected = src_cube[:, single_point.iy, single_point.ix]
+        assert result == expected
 
-    def test_aux_coord_one_interp_dim(self):
+    def test_aux_coord_one_interp_dim(self, src_cube, single_point):
         # Check exact result with an aux-coord over one interpolation dims.
-        cube = self.test_cube
-        cube.add_aux_coord(AuxCoord([11, 12, 13, 14], long_name="aux_x"), 2)
+        src_cube.add_aux_coord(
+            AuxCoord([11, 12, 13, 14], long_name="aux_x"), 2
+        )
 
         # The result cube should exactly equal a single source point.
-        result = interpolate(cube, self.single_sample_point, method="nearest")
-        self.assertEqual(result.shape[-1], 1)
+        result = interpolate(
+            src_cube, single_point.sample_point, method="nearest"
+        )
+        assert result.shape[-1] == 1
         result = result[..., 0]
-        expected = cube[:, self.single_point_iy, self.single_point_ix]
-        self.assertEqual(result, expected)
+        expected = src_cube[:, single_point.iy, single_point.ix]
+        assert result == expected
 
-    def test_aux_coord_both_interp_dims(self):
+    def test_aux_coord_both_interp_dims(self, src_cube, single_point):
         # Check exact result with an aux-coord over both interpolation dims.
-        cube = self.test_cube
-        cube.add_aux_coord(
+        src_cube.add_aux_coord(
             AuxCoord(
                 [[11, 12, 13, 14], [21, 22, 23, 24], [31, 32, 33, 34]],
                 long_name="aux_xy",
@@ -142,17 +192,18 @@ def test_aux_coord_both_interp_dims(self):
         )
 
         # The result cube should exactly equal a single source point.
-        result = interpolate(cube, self.single_sample_point, method="nearest")
-        self.assertEqual(result.shape[-1], 1)
+        result = interpolate(
+            src_cube, single_point.sample_point, method="nearest"
+        )
+        assert result.shape[-1] == 1
         result = result[..., 0]
-        expected = cube[:, self.single_point_iy, self.single_point_ix]
-        self.assertEqual(result, expected)
+        expected = src_cube[:, single_point.iy, single_point.ix]
+        assert result == expected
 
-    def test_aux_coord_fail_mixed_dims(self):
+    def test_aux_coord_fail_mixed_dims(self, src_cube, single_point):
         # Check behaviour with an aux-coord mapped over both interpolation and
         # non-interpolation dims : not supported.
-        cube = self.test_cube
-        cube.add_aux_coord(
+        src_cube.add_aux_coord(
             AuxCoord(
                 [[111, 112, 113, 114], [211, 212, 213, 214]],
                 long_name="aux_0x",
@@ -163,22 +214,23 @@ def test_aux_coord_fail_mixed_dims(self):
             "Coord aux_0x at one x-y position has the shape.*"
             "instead of being a single point"
         )
-        with self.assertRaisesRegex(ValueError, msg):
-            interpolate(cube, self.single_sample_point, method="nearest")
+        with pytest.raises(ValueError, match=msg):
+            interpolate(src_cube, single_point.sample_point, method="nearest")
 
-    def test_metadata(self):
+    def test_metadata(self, src_cube, single_point):
         # Check exact result matching for a single point, with additional
         # attributes and cell-methods.
-        cube = self.test_cube
-        cube.attributes["ODD_ATTR"] = "string-value-example"
-        cube.add_cell_method(iris.coords.CellMethod("mean", "area"))
-        result = interpolate(cube, self.single_sample_point, method="nearest")
+        src_cube.attributes["ODD_ATTR"] = "string-value-example"
+        src_cube.add_cell_method(iris.coords.CellMethod("mean", "area"))
+        result = interpolate(
+            src_cube, single_point.sample_point, method="nearest"
+        )
         # Check that the result is a single trajectory point, exactly equal to
         # the expected part of the original data.
-        self.assertEqual(result.shape[-1], 1)
+        assert result.shape[-1] == 1
         result = result[..., 0]
-        expected = cube[:, self.single_point_iy, self.single_point_ix]
-        self.assertEqual(result, expected)
+        expected = src_cube[:, single_point.iy, single_point.ix]
+        assert result == expected
 
 
 class TestLinear(tests.IrisTest):