Increase speed of iris.analysis.trajectory.interpolate with linear interpolation

benchmarks/benchmarks/trajectory.py

@@ -0,0 +1,48 @@
+# Copyright Iris contributors
+#
+# This file is part of Iris and is released under the LGPL license.
+# See COPYING and COPYING.LESSER in the root of the repository for full
+# licensing details.
+"""
+Trajectory benchmark test
+
+"""
+
+# import iris tests first so that some things can be initialised before
+# importing anything else
+from iris import tests  # isort:skip
+
+import numpy as np
+
+import iris
+from iris.analysis.trajectory import interpolate
+
+
+class TrajectoryInterpolation:
+    def setup(self) -> None:
+        # Prepare a cube and a template
+
+        cube_file_path = tests.get_data_path(
+            ["NetCDF", "regrid", "regrid_xyt.nc"]
+        )
+        self.cube = iris.load_cube(cube_file_path)
+
+        trajectory = np.array(
+            [np.array((-50 + i, -50 + i)) for i in range(100)]
+        )
+        self.sample_points = [
+            ("longitude", trajectory[:, 0]),
+            ("latitude", trajectory[:, 1]),
+        ]
+
+    def time_trajectory_linear(self) -> None:
+        # Regrid the cube onto the template.
+        out_cube = interpolate(self.cube, self.sample_points, method="linear")
+        # Realise the data
+        out_cube.data
+
+    def time_trajectory_nearest(self) -> None:
+        # Regrid the cube onto the template.
+        out_cube = interpolate(self.cube, self.sample_points, method="nearest")
+        # Realise the data
+        out_cube.data

docs/src/whatsnew/latest.rst

@@ -136,6 +136,10 @@ This document explains the changes made to Iris for this release
    aggregation method to be used on masked data when the missing data tolerance
    is set to 0. (:issue:`4735`, :pull:`4755`)
 
+#. `@wjbenfold`_ improved the speed of linear interpolation using 
+   :meth:`iris.analysis.trajectory.interpolate` (:pull:`4366`)
+
+
 
 🔥 Deprecations
 ===============

lib/iris/analysis/__init__.py

@@ -2696,6 +2696,10 @@ def interpolator(self, cube, coords):
             dimensions in the result cube caused by scalar values in
             `sample_points`.
 
+            The N arrays of values within `sample_points` will be used to
+            create an N-d grid of points that will then be sampled (rather than
+            just N points)
+
             The values for coordinates that correspond to date/times
             may optionally be supplied as datetime.datetime or
             cftime.datetime instances.

lib/iris/analysis/trajectory.py

@@ -320,20 +320,59 @@ def interpolate(cube, sample_points, method=None):
             break
 
     if method in ["linear", None]:
-        for i in range(trajectory_size):
-            point = [(coord, values[i]) for coord, values in sample_points]
-            column = cube.interpolate(point, Linear())
-            new_cube.data[..., i] = column.data
-            # Fill in the empty squashed (non derived) coords.
-            for column_coord in column.dim_coords + column.aux_coords:
-                src_dims = cube.coord_dims(column_coord)
-                if not squish_my_dims.isdisjoint(src_dims):
-                    if len(column_coord.points) != 1:
-                        msg = "Expected to find exactly one point. Found {}."
-                        raise Exception(msg.format(column_coord.points))
-                    new_cube.coord(column_coord.name()).points[
-                        i
-                    ] = column_coord.points[0]
+        # Using cube.interpolate will generate extra values that we don't need
+        # as it makes a grid from the provided coordinates (like a meshgrid)
+        # and then does interpolation for all of them. This is memory
+        # inefficient, but significantly more time efficient than calling
+        # cube.interpolate (or the underlying method on the interpolator)
+        # repeatedly, so using this approach for now. In future, it would be
+        # ideal if we only interpolated at the points we care about
+        columns = cube.interpolate(sample_points, Linear())
+        # np.einsum(a, [0, 0], [0]) is like np.diag(a)
+        # We're using einsum here to do an n-dimensional diagonal, leaving the
+        # other dimensions unaffected and putting the diagonal's direction on
+        # the final axis
+        initial_inds = list(range(1, columns.ndim + 1))
+        for ind in squish_my_dims:
+            initial_inds[ind] = 0
+        final_inds = list(filter(lambda x: x != 0, initial_inds)) + [0]
+        new_cube.data = np.einsum(columns.data, initial_inds, final_inds)
+
+        # Fill in the empty squashed (non derived) coords.
+        # We're using the same einstein summation plan as for the cube, but
+        # redoing those indices to match the indices in the coordinates
+        for columns_coord in columns.dim_coords + columns.aux_coords:
+            src_dims = cube.coord_dims(columns_coord)
+            if not squish_my_dims.isdisjoint(src_dims):
+                # Mapping the cube indicies onto the coord
+                initial_coord_inds = [initial_inds[ind] for ind in src_dims]
+                # Making the final ones the same way as for the cube
+                # 0 will always appear in the initial ones because we know this
+                # coord overlaps the squish dims
+                final_coord_inds = list(
+                    filter(lambda x: x != 0, initial_coord_inds)
+                ) + [0]
+                new_coord_points = np.einsum(
+                    columns_coord.points, initial_coord_inds, final_coord_inds
+                )
+                # Check we're not overwriting coord.points with the wrong shape
+                if (
+                    not new_cube.coord(columns_coord.name()).points.shape
+                    == new_coord_points.shape
+                ):
+                    msg = (
+                        "Coord {} was expected to have new points of shape {}. "
+                        "Found shape of {}."
+                    )
+                    raise ValueError(
+                        msg.format(
+                            columns_coord.name(),
+                            new_cube.coord(columns_coord.name()).points.shape,
+                            new_coord_points.shape,
+                        )
+                    )
+                # Replace the points
+                new_cube.coord(columns_coord.name()).points = new_coord_points
 
     elif method == "nearest":
         # Use a cache with _nearest_neighbour_indices_ndcoords()

lib/iris/cube.py

@@ -4359,6 +4359,8 @@ def interpolate(self, sample_points, scheme, collapse_scalar=True):
             interpolate. The values for coordinates that correspond to
             dates or times may optionally be supplied as datetime.datetime or
             cftime.datetime instances.
+            The N pairs supplied will be used to create an N-d grid of points
+            that will then be sampled (rather than just N points).
         * scheme:
             An instance of the type of interpolation to use to interpolate from this
             :class:`~iris.cube.Cube` to the given sample points. The

lib/iris/tests/results/netcdf/save_load_traj.cml

@@ -1,6 +1,6 @@
 <?xml version="1.0" ?>
 <cubes xmlns="urn:x-iris:cubeml-0.2">
-  <cube dtype="float64" standard_name="air_temperature" units="K" var_name="air_temperature">
+  <cube dtype="float32" standard_name="air_temperature" units="K" var_name="air_temperature">
     <attributes>
       <attribute name="Conventions" value="CF-1.7"/>
       <attribute name="STASH" value="m01s16i203"/>
@@ -36,6 +36,6 @@
         <coord name="time"/>
       </cellMethod>
     </cellMethods>
-    <data dtype="float64" shape="(10,)" state="deferred"/>
+    <data dtype="float32" shape="(10,)" state="deferred"/>
   </cube>
 </cubes>

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

@@ -1,6 +1,6 @@
 <?xml version="1.0" ?>
 <cubes xmlns="urn:x-iris:cubeml-0.2">
-  <cube dtype="float64" standard_name="air_potential_temperature" units="K">
+  <cube dtype="float32" standard_name="air_potential_temperature" units="K">
     <attributes>
       <attribute name="STASH" value="m01s00i004"/>
       <attribute name="source" value="Data from Met Office Unified Model"/>
@@ -16,13 +16,13 @@
       </coord>
       <coord datadims="[2]">
         <auxCoord id="f1ab3066" points="[-0.1188, -0.1188, -0.1188, ..., -0.1188,
-		-0.1188, -0.1188]" shape="(100,)" standard_name="grid_latitude" units="Unit('degrees')" value_type="float32">
+		-0.1188, -0.1188]" shape="(100,)" standard_name="grid_latitude" units="Unit('degrees')" value_type="float64">
           <rotatedGeogCS ellipsoid="GeogCS(6371229.0)" grid_north_pole_latitude="37.5" grid_north_pole_longitude="177.5"/>
         </auxCoord>
       </coord>
       <coord datadims="[2]">
-        <auxCoord id="d3676747" points="[359.58, 359.58, 359.581, ..., 359.667, 359.668,
-		359.669]" shape="(100,)" standard_name="grid_longitude" units="Unit('degrees')" value_type="float32">
+        <auxCoord id="d3676747" points="[359.57958984, 359.580489954, 359.581390069, ...,
+		359.666900941, 359.667801056, 359.66870117]" shape="(100,)" standard_name="grid_longitude" units="Unit('degrees')" value_type="float64">
           <rotatedGeogCS ellipsoid="GeogCS(6371229.0)" grid_north_pole_latitude="37.5" grid_north_pole_longitude="177.5"/>
         </auxCoord>
       </coord>
@@ -94,6 +94,6 @@
       </coord>
     </coords>
     <cellMethods/>
-    <data checksum="0x66df698f" dtype="float64" shape="(6, 70, 100)"/>
+    <data checksum="0x04132092" dtype="float32" shape="(6, 70, 100)"/>
   </cube>
 </cubes>

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

@@ -1,6 +1,6 @@
 <?xml version="1.0" ?>
 <cubes xmlns="urn:x-iris:cubeml-0.2">
-  <cube dtype="float64" standard_name="air_potential_temperature" units="K">
+  <cube dtype="float32" standard_name="air_potential_temperature" units="K">
     <attributes>
       <attribute name="STASH" value="m01s00i004"/>
       <attribute name="source" value="Data from Met Office Unified Model"/>
@@ -15,12 +15,12 @@
 		347921.666667, 347921.833333, 347922.0]" shape="(6,)" standard_name="forecast_reference_time" units="Unit('hours since 1970-01-01 00:00:00', calendar='gregorian')" value_type="float64"/>
       </coord>
       <coord datadims="[2]">
-        <auxCoord id="f1ab3066" points="[-0.1188]" shape="(1,)" standard_name="grid_latitude" units="Unit('degrees')" value_type="float32">
+        <auxCoord id="f1ab3066" points="[-0.1188]" shape="(1,)" standard_name="grid_latitude" units="Unit('degrees')" value_type="float64">
           <rotatedGeogCS ellipsoid="GeogCS(6371229.0)" grid_north_pole_latitude="37.5" grid_north_pole_longitude="177.5"/>
         </auxCoord>
       </coord>
       <coord datadims="[2]">
-        <auxCoord id="d3676747" points="[359.58]" shape="(1,)" standard_name="grid_longitude" units="Unit('degrees')" value_type="float32">
+        <auxCoord id="d3676747" points="[359.57958984]" shape="(1,)" standard_name="grid_longitude" units="Unit('degrees')" value_type="float64">
           <rotatedGeogCS ellipsoid="GeogCS(6371229.0)" grid_north_pole_latitude="37.5" grid_north_pole_longitude="177.5"/>
         </auxCoord>
       </coord>
@@ -92,6 +92,6 @@
       </coord>
     </coords>
     <cellMethods/>
-    <data dtype="float64" shape="(6, 70, 1)" state="loaded"/>
+    <data dtype="float32" shape="(6, 70, 1)" state="loaded"/>
   </cube>
 </cubes>

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

@@ -1,6 +1,6 @@
 <?xml version="1.0" ?>
 <cubes xmlns="urn:x-iris:cubeml-0.2">
-  <cube dtype="float64" standard_name="air_potential_temperature" units="K">
+  <cube dtype="float32" standard_name="air_potential_temperature" units="K">
     <attributes>
       <attribute name="STASH" value="m01s00i004"/>
       <attribute name="source" value="Data from Met Office Unified Model"/>
@@ -14,19 +14,26 @@
         <dimCoord id="9c8bdf81" points="[347921.166667]" shape="(1,)" standard_name="forecast_reference_time" units="Unit('hours since 1970-01-01 00:00:00', calendar='gregorian')" value_type="float64"/>
       </coord>
       <coord datadims="[0]">
-        <auxCoord id="f1ab3066" points="[-0.1188, -0.115707, -0.112614, -0.109521,
-		-0.106428, -0.103335, -0.100242, -0.0971485,
-		-0.0940554, -0.0909623, -0.0878693, -0.0847762,
-		-0.081034, -0.0761435, -0.0712529, -0.0663623,
-		-0.0614717, -0.0565812, -0.0516906, -0.0468]" shape="(20,)" standard_name="grid_latitude" units="Unit('degrees')" value_type="float32">
+        <auxCoord id="f1ab3066" points="[-0.1188, -0.11570692636, -0.112613852721,
+		-0.109520779081, -0.106427705441,
+		-0.103334631802, -0.100241558162,
+		-0.0971484845225, -0.0940554108829,
+		-0.0909623372432, -0.0878692636036,
+		-0.0847761899639, -0.0810340518516,
+		-0.0761434730157, -0.0712528941797,
+		-0.0663623153438, -0.0614717365078,
+		-0.0565811576719, -0.0516905788359, -0.0468]" shape="(20,)" standard_name="grid_latitude" units="Unit('degrees')" value_type="float64">
           <rotatedGeogCS ellipsoid="GeogCS(6371229.0)" grid_north_pole_latitude="37.5" grid_north_pole_longitude="177.5"/>
         </auxCoord>
       </coord>
       <coord datadims="[0]">
-        <auxCoord id="d3676747" points="[359.589, 359.595, 359.601, 359.607, 359.613,
-		359.62, 359.626, 359.632, 359.638, 359.644,
-		359.65, 359.657, 359.659, 359.654, 359.649,
-		359.644, 359.639, 359.634, 359.629, 359.625]" shape="(20,)" standard_name="grid_longitude" units="Unit('degrees')" value_type="float32">
+        <auxCoord id="d3676747" points="[359.5886, 359.594786147, 359.600972295,
+		359.607158442, 359.613344589, 359.619530736,
+		359.625716884, 359.631903031, 359.638089178,
+		359.644275326, 359.650461473, 359.65664762,
+		359.658834052, 359.653943473, 359.649052894,
+		359.644162315, 359.639271737, 359.634381158,
+		359.629490579, 359.6246]" shape="(20,)" standard_name="grid_longitude" units="Unit('degrees')" value_type="float64">
           <rotatedGeogCS ellipsoid="GeogCS(6371229.0)" grid_north_pole_latitude="37.5" grid_north_pole_longitude="177.5"/>
         </auxCoord>
       </coord>
@@ -52,6 +59,6 @@
       </coord>
     </coords>
     <cellMethods/>
-    <data dtype="float64" shape="(20,)" state="loaded"/>
+    <data dtype="float32" shape="(20,)" state="loaded"/>
   </cube>
 </cubes>