Support parallel processing within DistanceTo plugin

improver/generate_ancillaries/generate_distance_to_feature.py

@@ -9,6 +9,7 @@
 import pyproj
 from geopandas import GeoDataFrame, GeoSeries, clip
 from iris.cube import Cube
+from joblib import Parallel, delayed
 from numpy import array, min, round
 from shapely.geometry import Point
 
@@ -43,6 +44,8 @@ def __init__(
         new_name: Optional[str] = None,
         buffer: float = 30000,
         clip_geometry_flag: bool = False,
+        parallel: bool = False,
+        n_parallel_jobs: Optional[int] = 1,
     ) -> None:
         """
         Initialise the DistanceTo plugin.
@@ -65,11 +68,19 @@ def __init__(
                 the site locations with a buffer distance added to the bounds. If set to
                 False, the full geometry will be used to calculate the distance to the
                 nearest feature.
+            parallel:
+                A flag to indicate whether to use parallel processing when calculating
+                distances.
+            n_parallel_jobs:
+                The number of parallel jobs to use when calculating distances.
+                By default, 1 job is used.
         """
         self.epsg_projection = epsg_projection
         self.new_name = new_name
         self.buffer = buffer
         self.clip_geometry_flag = clip_geometry_flag
+        self.parallel = parallel
+        self.n_parallel_jobs = n_parallel_jobs
 
     @staticmethod
     def get_clip_values(points: List[float], buffer: float) -> List[float]:
@@ -200,10 +211,31 @@ def distance_to(self, site_points: GeoSeries, geometry: GeoDataFrame) -> List[in
         Returns:
             A list of distances from each site point to the nearest feature in the
             geometry rounded to the nearest metre."""
-        distance_results = []
-        for point in site_points:
-            distance_to_nearest = min(point.distance(geometry.geometry))
-            distance_results.append(round(distance_to_nearest))
+
+        def _distance_to_nearest(point: Point, geometry: GeoDataFrame) -> float:
+            """Calculate the distance from a point to the nearest feature in the
+            geometry.
+            Args:
+                point:
+                    A shapely Point object representing the site location.
+                geometry:
+                    A GeoDataFrame containing the geometry in the target projection.
+            Returns:
+                The distance from the point to the nearest feature in the geometry
+                rounded to the nearest metre.
+            """
+            return round(min(point.distance(geometry.geometry)))
+
+        if self.parallel:
+            parallel = Parallel(n_jobs=self.n_parallel_jobs, prefer="threads")
+            output_generator = parallel(
+                delayed(_distance_to_nearest)(point, geometry) for point in site_points
+            )
+            distance_results = list(output_generator)
+        else:
+            distance_results = []
+            for point in site_points:
+                distance_results.append(_distance_to_nearest(point, geometry))
 
         return distance_results
 

improver_tests/generate_ancillaries/test_DistanceTo.py

@@ -4,6 +4,8 @@
 # See LICENSE in the root of the repository for full licensing details.
 """Unit tests for the DistanceTo plugin."""
 
+import os
+
 import numpy as np
 import pytest
 from geopandas import GeoDataFrame
@@ -397,27 +399,62 @@ def test_distance_to_with_polygon_geometry(
 
 
 @pytest.mark.parametrize(
-    "geometry_type,expected_distance",
+    "geometry_type,distance_type,if_parallel,expected_distance",
     [
-        ("point", [0, 500, 707, 707]),
-        ("line", [0, 0, 500, 500]),
-        ("polygon", [0, 0, 0, 500]),
+        ("point", "nearest", False, [0, 500, 707, 707]),
+        ("line", "nearest", False, [0, 0, 500, 500]),
+        ("polygon", "nearest", False, [0, 0, 0, 500]),
+        ("point", "nearest", True, [0, 500, 707, 707]),
+        ("line", "nearest", True, [0, 0, 500, 500]),
+        ("polygon", "nearest", True, [0, 0, 0, 500]),
+        ("point", "furthest", True, [1414, 1118, 707, 1581]),
+        ("line", "furthest", True, [0, 0, 500, 500]),
+        ("polygon", "furthest", True, [0, 0, 0, 500]),
     ],
 )
-@pytest.mark.parametrize("geometry_crs", ("laea", "latlon"))
+@pytest.mark.parametrize("geometry_crs", ("laea",))
 def test_distance_to_with_multiple_sites(
     multiple_site_cube,
     geometry_type,
     geometry_crs,
     expected_distance,
+    distance_type,
+    if_parallel,
     request,
 ):
     """Test the DistanceTo plugin works when provided a site cube with multiple sites and
     different types of geometry"""
 
     geometry = request.getfixturevalue(f"geometry_{geometry_type}_{geometry_crs}")
 
-    output_cube = DistanceTo(3035)(multiple_site_cube, geometry)
+    n_jobs = 1
+    if os.cpu_count() > 1:
+        n_jobs = 2
+
+    plugin = DistanceTo(3035, parallel=if_parallel, n_parallel_jobs=n_jobs)
+    # Modify the plugin to use a custom distance calculation method when
+    # parallel processing is requested with the 'furthest' processing mode.
+    # This demonstrates that the plugin must be using the parallel processing pathway.
+    if if_parallel and distance_type == "furthest":
+        from joblib import Parallel, delayed
+
+        def _distance_to(site_points, geometry):
+            def _distance_to_furthest(point, geometry):
+                print(point.distance(geometry.geometry), point, geometry.geometry)
+                return round(max(point.distance(geometry.geometry)))
+
+            parallel = Parallel(n_jobs=n_jobs, prefer="threads")
+            output_generator = parallel(
+                delayed(_distance_to_furthest)(point, geometry) for point in site_points
+            )
+            distance_results = list(output_generator)
+            return distance_results
+
+        plugin.distance_to = _distance_to
+
+    assert plugin.parallel == if_parallel
+    assert plugin.n_parallel_jobs == n_jobs
+    output_cube = plugin(multiple_site_cube, geometry)
     assert output_cube.name() == "rain_rate"
     assert output_cube.units == "m"