Add option to enable/disable regridder weights caching

Author: schlunma

doc/recipe/preprocessor.rst

@@ -881,6 +881,8 @@ The arguments are defined below:
   Otherwise, it cuts off at the previous value.
 * ``step_longitude``: Longitude distance between the centers of two neighbouring cells.
 
+.. _regridding_schemes:
+
 Regridding (interpolation, extrapolation) schemes
 -------------------------------------------------
 
@@ -889,10 +891,26 @@ ESMValCore has a number of built-in regridding schemes, which are presented in
 third party regridding schemes designed for use with :doc:`Iris
 <iris:index>`. This is explained in :ref:`generic regridding schemes`.
 
-Whenever possible, regridding weights are cached to reduce run times (see `here
+If desired, regridding weights can be cached to reduce run times (see `here
 <https://scitools-iris.readthedocs.io/en/latest/userguide/interpolation_and_regridding.html#caching-a-regridder>`__
 for technical details on this).
-An overview of regridding schemes that support weights caching is given `here
+This can speed up the regridding of different datasets with similar source and
+target grids massively, but may take up a lot of memory for extremely
+high-resolution data.
+By default, this feature is disabled; to enable it, use the option
+``cache_weights: true`` in the preprocessor definition:
+
+.. code-block:: yaml
+
+    preprocessors:
+      regrid_preprocessor:
+        regrid:
+          target_grid: 0.1x0.1
+          scheme: linear
+          cache_weights: true
+
+Not all regridding schemes support weights caching. An overview of those that
+do is given `here
 <https://scitools-iris.readthedocs.io/en/latest/further_topics/which_regridder_to_use.html#which-regridder-to-use>`__
 and in the docstrings :ref:`here <regridding_schemes>`.
 

esmvalcore/preprocessor/_regrid.py

@@ -500,8 +500,7 @@ def _get_target_grid_cube(
     elif isinstance(target_grid, (str, Path)) and os.path.isfile(target_grid):
         target_grid_cube = iris.load_cube(target_grid)
     elif isinstance(target_grid, str):
-        # Generate a target grid from the provided cell-specification,
-        # and cache the resulting stock cube for later use.
+        # Generate a target grid from the provided cell-specification
         target_grid_cube = _global_stock_cube(
             target_grid, lat_offset, lon_offset
         )
@@ -642,7 +641,12 @@ def _load_generic_scheme(scheme: dict):
 _CACHED_REGRIDDERS: dict[tuple, dict] = {}
 
 
-def _get_regridder(src_cube: Cube, tgt_cube: Cube, scheme: str | dict):
+def _get_regridder(
+    src_cube: Cube,
+    tgt_cube: Cube,
+    scheme: str | dict,
+    cache_weights: bool,
+):
     """Get regridder to actually perform regridding.
 
     Note
@@ -652,25 +656,33 @@ def _get_regridder(src_cube: Cube, tgt_cube: Cube, scheme: str | dict):
     interpolation_and_regridding.html#caching-a-regridder.)
 
     """
-    # Use existing regridder if possible. This first checks the regridding
-    # scheme name and shapes of source and target coordinates, and only if
-    # these match check coordinates themselves (this is much more expensive).
-    coord_key = _get_coord_key(src_cube, tgt_cube)
-    name_shape_key = _get_name_and_shape_key(src_cube, tgt_cube, scheme)
-    if name_shape_key in _CACHED_REGRIDDERS:
-        # We cannot simply do a test for `coord_key in
-        # _CACHED_REGRIDDERS[shape_key]` below since the hash() of a coordinate
-        # is simply its id() (thus, coordinates loaded from two different files
-        # would never be considered equal)
-        for (key, regridder) in _CACHED_REGRIDDERS[name_shape_key].items():
-            if key == coord_key:
-                return regridder
-
-    # If regridder is not cached yet, return a new one and cache it
-    loaded_scheme = _load_scheme(src_cube, scheme)
-    regridder = loaded_scheme.regridder(src_cube, tgt_cube)
-    _CACHED_REGRIDDERS.setdefault(name_shape_key, {})
-    _CACHED_REGRIDDERS[name_shape_key][coord_key] = regridder
+    # (1) Weights caching enabled
+    if cache_weights:
+        # To search for a matching regridder in the cache, first check the
+        # regridding scheme name and shapes of source and target coordinates.
+        # Only if these match, check coordinates themselves (this is much more
+        # expensive).
+        coord_key = _get_coord_key(src_cube, tgt_cube)
+        name_shape_key = _get_name_and_shape_key(src_cube, tgt_cube, scheme)
+        if name_shape_key in _CACHED_REGRIDDERS:
+            # We cannot simply do a test for `coord_key in
+            # _CACHED_REGRIDDERS[shape_key]` below since the hash() of a
+            # coordinate is simply its id() (thus, coordinates loaded from two
+            # different files would never be considered equal)
+            for (key, regridder) in _CACHED_REGRIDDERS[name_shape_key].items():
+                if key == coord_key:
+                    return regridder
+
+        # Regridder is not in cached -> return a new one and cache it
+        loaded_scheme = _load_scheme(src_cube, scheme)
+        regridder = loaded_scheme.regridder(src_cube, tgt_cube)
+        _CACHED_REGRIDDERS.setdefault(name_shape_key, {})
+        _CACHED_REGRIDDERS[name_shape_key][coord_key] = regridder
+
+    # (2) Weights caching disabled
+    else:
+        loaded_scheme = _load_scheme(src_cube, scheme)
+        regridder = loaded_scheme.regridder(src_cube, tgt_cube)
 
     return regridder
 
@@ -701,6 +713,7 @@ def regrid(
     scheme: str | dict,
     lat_offset: bool = True,
     lon_offset: bool = True,
+    cache_weights: bool = False,
 ) -> Cube:
     """Perform horizontal regridding.
 
@@ -747,6 +760,13 @@ def regrid(
         Offset the grid centers of the longitude coordinate w.r.t. Greenwich
         meridian by half a grid step. This argument is ignored if
         `target_grid` is a cube or file.
+    cache_weights:
+        If ``True``, cache regridding weights for later usage. This can speed
+        up the regridding of different datasets with similar source and target
+        grids massively, but may take up a lot of memory for extremely
+        high-resolution data. This option is ignored for schemes that do not
+        support weights caching. A list of supported schemes is available in
+        the section on :ref:`regridding_schemes`.
 
     Returns
     -------
@@ -816,7 +836,7 @@ def regrid(
     # Load scheme and reuse existing regridder if possible
     if isinstance(scheme, str):
         scheme = scheme.lower()
-    regridder = _get_regridder(cube, target_grid_cube, scheme)
+    regridder = _get_regridder(cube, target_grid_cube, scheme, cache_weights)
 
     # Rechunk and actually perform the regridding
     cube = _rechunk(cube, target_grid_cube)

tests/unit/preprocessor/_regrid/test_regrid.py

@@ -79,22 +79,26 @@ def cube_30x30():
 SCHEMES = ['area_weighted', 'linear', 'nearest']
 
 
+@pytest.mark.parametrize('cache_weights', [True, False])
 @pytest.mark.parametrize('scheme', SCHEMES)
-def test_builtin_regridding(scheme, cube_10x10, cube_30x30):
+def test_builtin_regridding(scheme, cache_weights, cube_10x10, cube_30x30):
     """Test `regrid.`"""
     _cached_regridders = esmvalcore.preprocessor._regrid._CACHED_REGRIDDERS
     assert _cached_regridders == {}
 
-    res = regrid(cube_10x10, cube_30x30, scheme)
+    res = regrid(cube_10x10, cube_30x30, scheme, cache_weights=cache_weights)
 
     assert res.coord('latitude') == cube_30x30.coord('latitude')
     assert res.coord('longitude') == cube_30x30.coord('longitude')
     assert res.dtype == np.float32
     assert np.allclose(res.data, 0.0)
 
-    assert len(_cached_regridders) == 1
-    key = (scheme, (18,), (36,), (30,), (30,))
-    assert key in _cached_regridders
+    if cache_weights:
+        assert len(_cached_regridders) == 1
+        key = (scheme, (18,), (36,), (30,), (30,))
+        assert key in _cached_regridders
+    else:
+        assert not _cached_regridders
 
 
 @pytest.mark.parametrize('scheme', SCHEMES)
@@ -127,7 +131,8 @@ def test_regrid_generic_invalid_reference(cube_10x10, cube_30x30):
         regrid(cube_10x10, cube_30x30, {'reference': 'this.does:not.exist'})
 
 
-def test_regrid_generic_regridding(cube_10x10, cube_30x30):
+@pytest.mark.parametrize('cache_weights', [True, False])
+def test_regrid_generic_regridding(cache_weights, cube_10x10, cube_30x30):
     """Test `regrid.`"""
     _cached_regridders = esmvalcore.preprocessor._regrid._CACHED_REGRIDDERS
     assert _cached_regridders == {}
@@ -139,23 +144,30 @@ def test_regrid_generic_regridding(cube_10x10, cube_30x30):
             'reference': 'iris.analysis:Linear',
             'extrapolation_mode': 'mask',
         },
+        cache_weights=cache_weights,
+    )
+    cube_lin = regrid(
+        cube_10x10, cube_30x30, 'linear', cache_weights=cache_weights
     )
-    cube_lin = regrid(cube_10x10, cube_30x30, 'linear')
     assert cube_gen.dtype == np.float32
     assert cube_lin.dtype == np.float32
     assert cube_gen == cube_lin
 
-    assert len(_cached_regridders) == 2
-    key_1 = (
-        "{'reference': 'iris.analysis:Linear', 'extrapolation_mode': 'mask'}",
-        (18,),
-        (36,),
-        (30,),
-        (30,),
-    )
-    key_2 = ('linear', (18,), (36,), (30,), (30,))
-    assert key_1 in _cached_regridders
-    assert key_2 in _cached_regridders
+    if cache_weights:
+        assert len(_cached_regridders) == 2
+        key_1 = (
+            "{'reference': 'iris.analysis:Linear', 'extrapolation_mode': "
+            "'mask'}",
+            (18,),
+            (36,),
+            (30,),
+            (30,),
+        )
+        key_2 = ('linear', (18,), (36,), (30,), (30,))
+        assert key_1 in _cached_regridders
+        assert key_2 in _cached_regridders
+    else:
+        assert not _cached_regridders
 
 
 @pytest.mark.parametrize(
@@ -360,7 +372,7 @@ def test_regridding_weights_use_cache(scheme, cube_10x10, cube_30x30, mocker):
         esmvalcore.preprocessor._regrid, '_load_scheme', autospec=True
     )
 
-    reg = _get_regridder(cube_10x10, cube_30x30, scheme)
+    reg = _get_regridder(cube_10x10, cube_30x30, scheme, cache_weights=True)
 
     assert reg == mocker.sentinel.regridder