Fix handling of data in "nearest" trajectory interpolate

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

@@ -13,6 +13,7 @@
 import iris.tests as tests  # isort:skip
 
 import numpy as np
+import pytest
 
 from iris.analysis.trajectory import interpolate
 from iris.coords import AuxCoord, DimCoord
@@ -38,15 +39,15 @@ 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(autouse=True)
+    def setup(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.
@@ -70,30 +71,17 @@ def setUp(self):
             ("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)
-
-    def test_multi_point_same_cube(self):
-        # Check an exact result for multiple points.
-        cube = self.test_cube
+    @pytest.fixture
+    def multi_point_extra_setup(self, setup):
         # Use latitude selection to recreate a whole row of the original cube.
-        sample_points = [
+        self.sample_points = [
             ("longitude", [-180, -90, 0, 90]),
             ("latitude", [0, 0, 0, 0]),
         ]
-        result = interpolate(cube, sample_points, method="nearest")
 
         # 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 = self.test_cube[:, 1, :]
         # Result 'longitude' is now an aux coord.
         co_x = expected.coord("longitude")
         expected.remove_coord(co_x)
@@ -104,7 +92,51 @@ 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)
+        self.expected_multipoint_cube = expected
+
+    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.
+        assert result.shape[-1] == 1
+        result = result[..., 0]
+        expected = cube[:, self.single_point_iy, self.single_point_ix]
+        assert result == expected
+
+    @pytest.mark.usefixtures("multi_point_extra_setup")
+    def test_multi_point_same_cube(self):
+        # Check an exact result for multiple points.
+        result = interpolate(
+            self.test_cube, self.sample_points, method="nearest"
+        )
+        assert result == self.expected_multipoint_cube
+
+    @pytest.mark.usefixtures("multi_point_extra_setup")
+    def test_mask_preserved(self):
+        cube = self.test_cube
+        mask = np.zeros_like(cube.data)
+        mask[:, :, 1] = 1
+        cube.data = np.ma.array(cube.data, mask=mask)
+
+        expected = self.expected_multipoint_cube
+        expected.data = np.ma.array(expected.data, mask=mask[:, 0])
+
+        result = interpolate(cube, self.sample_points, method="nearest")
+        assert result == expected
+        assert np.allclose(result.data.mask, expected.data.mask)
+
+    @pytest.mark.usefixtures("multi_point_extra_setup")
+    def test_dtype_preserved(self):
+        cube = self.test_cube
+        cube.data = cube.data.astype(np.int16)
+        expected = self.expected_multipoint_cube
+
+        result = interpolate(cube, self.sample_points, method="nearest")
+        assert result == expected
+        assert np.allclose(result.data, expected.data)
+        assert result.data.dtype == np.int16
 
     def test_aux_coord_noninterpolation_dim(self):
         # Check exact result with an aux-coord mapped to an uninterpolated dim.
@@ -113,10 +145,10 @@ def test_aux_coord_noninterpolation_dim(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)
+        assert result.shape[-1] == 1
         result = result[..., 0]
         expected = cube[:, self.single_point_iy, self.single_point_ix]
-        self.assertEqual(result, expected)
+        assert result == expected
 
     def test_aux_coord_one_interp_dim(self):
         # Check exact result with an aux-coord over one interpolation dims.
@@ -125,10 +157,10 @@ def test_aux_coord_one_interp_dim(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)
+        assert result.shape[-1] == 1
         result = result[..., 0]
         expected = cube[:, self.single_point_iy, self.single_point_ix]
-        self.assertEqual(result, expected)
+        assert result == expected
 
     def test_aux_coord_both_interp_dims(self):
         # Check exact result with an aux-coord over both interpolation dims.
@@ -143,10 +175,10 @@ 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)
+        assert result.shape[-1] == 1
         result = result[..., 0]
         expected = cube[:, self.single_point_iy, self.single_point_ix]
-        self.assertEqual(result, expected)
+        assert result == expected
 
     def test_aux_coord_fail_mixed_dims(self):
         # Check behaviour with an aux-coord mapped over both interpolation and
@@ -163,7 +195,7 @@ 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):
+        with pytest.raises(ValueError, match=msg):
             interpolate(cube, self.single_sample_point, method="nearest")
 
     def test_metadata(self):
@@ -175,10 +207,10 @@ def test_metadata(self):
         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)
+        assert result.shape[-1] == 1
         result = result[..., 0]
         expected = cube[:, self.single_point_iy, self.single_point_ix]
-        self.assertEqual(result, expected)
+        assert result == expected
 
 
 class TestLinear(tests.IrisTest):