Partition MVP

docs/src/userguide/examples.rst

@@ -54,6 +54,97 @@ certain regridders. We can do this as follows::
     # Use loaded regridder.
     result = loaded_regridder(source_mesh_cube)
 
+Partitioning a Regridder
+------------------------
+
+If a regridder would be too large to handle in memory, it can be broken down
+into smaller regridders which can collectively do the job of the larger regridder.
+This is done using a `Partition` object.
+
+.. note:: Currently, it is only possible to partition regridding when the source is
+          a large grid and the target is small enough to fit in memory.
+
+A `Partition` is made by specifying a source, a target, a list of files, and a way
+to divide the source grid into blocks whose regridders are saved to those files::
+
+    from iris.util import make_gridcube
+
+    from esmf_regrid import ESMFAreaWeighted
+    from esmf_regrid.experimental.partition import Partition
+
+    # Create a large source cube.
+    source_cube = make_gridcube(nx=800, ny=800)
+
+    # Create a small target cube.
+    target_cube = make_gridcube(nx=100, ny=100)
+
+    # Set the regridding scheme.
+    scheme = AreaWeighted()
+
+    # List a collection of file names/paths to save partial regridders to.
+    files = ["file_1", "file_2", "file_3", "file_4"]
+
+    # Set the size of each block of the partition. For the keyword `src_chunks`
+    # this follows the dask chunking API.
+    src_chunks = (400, 400)
+
+    # Initialise the partition.
+    partition = Partition(
+        source_cube,
+        target_cube,
+        scheme,
+        files,
+        src_chunks=src_chunks
+    )
+
+.. note:: there are several different ways of specifying the division of the
+          source into blocks :
+          see :class:`~esmf_regrid.experimental.partition.Partition`.
+
+Initialising the `Partition` will not generate the files automatically unless
+the `auto_generate` keyword is set to `True`. In order for this `Partition` to
+function, the regridder files must be generated by calling the `generate_files`
+method::
+
+    # Generate partial regridders and save them to the list of files.
+    partition.generate_files()
+
+    # Once the files have been generated, they can be used for regridding.
+    result = partition.apply_regridders(source_cube)
+
+.. note:: Not all files need to be generated at once, if you have a grid which
+          needs to be split into very many files, it is possible to generate only
+          a portion of those files within a given session by passing the number
+          of files to generate as an argument to the regridder. It is then possible
+          to split the file generation in batches across multiple python sessions.
+
+Once the files for a regridder have been generated, they can be used to reconstruct
+the partition object in a later session. This is done by passing in the list of
+files which have already been generated::
+
+    # Use the same arguments which constructed the original partition.
+    source_cube = make_gridcube(nx=800, ny=800)
+    target_cube = make_gridcube(nx=100, ny=100)
+    scheme = AreaWeighted()
+    files = ["file_1", "file_2", "file_3", "file_4"]
+    src_chunks = (400, 400)
+
+    # List the files which have already been generated.
+    saved_files = ["file_1", "file_2", "file_3", "file_4"]
+
+    # Reconstruct Partition from pre-generated files.
+    partition = Partition(
+        source_cube,
+        target_cube,
+        scheme,
+        files,
+        src_chunks=src_chunks
+        saved_files=saved_files  # Pass in the list of saved files.
+    )
+
+    # The new Partition can now be used without the need for generating files.
+    result = partition.apply_regridders(source_cube)
+
 .. todo:
     Add more examples.
 

pyproject.toml

@@ -298,5 +298,5 @@ convention = "numpy"
 [tool.ruff.lint.pylint]
 # TODO: refactor to reduce complexity, if possible
 max-args = 10
-max-branches = 21
+max-branches = 25
 max-statements = 110

src/esmf_regrid/esmf_regridder.py

@@ -175,6 +175,46 @@ def _out_dtype(self, in_dtype):
             ).dtype
         return out_dtype
 
+    def _gen_weights_and_data(self, src_array):
+        extra_shape = src_array.shape[: -self.src.dims]
+
+        if self.method == Constants.Method.NEAREST:
+            weight_matrix = self.weight_matrix.astype(src_array.dtype)
+        else:
+            weight_matrix = self.weight_matrix
+
+        flat_src = self.src._array_to_matrix(ma.filled(src_array, 0.0))
+        flat_tgt = weight_matrix @ flat_src
+
+        src_inverted_mask = self.src._array_to_matrix(~ma.getmaskarray(src_array))
+        weight_sums = weight_matrix @ src_inverted_mask
+        return weight_sums, flat_tgt, extra_shape
+
+    def _regrid_from_weights_and_data(
+        self,
+        tgt_weights,
+        tgt_data,
+        extra,
+        norm_type=Constants.NormType.FRACAREA,
+        mdtol=1,
+    ):
+        # Set the minimum mdtol to be slightly higher than 0 to account for rounding
+        # errors.
+        mdtol = max(mdtol, 1e-8)
+        tgt_mask = tgt_weights > 1 - mdtol
+        normalisations = np.ones_like(tgt_data)
+        if self.method != Constants.Method.NEAREST:
+            masked_weight_sums = tgt_weights * tgt_mask
+            if norm_type == Constants.NormType.FRACAREA:
+                normalisations[tgt_mask] /= masked_weight_sums[tgt_mask]
+            elif norm_type == Constants.NormType.DSTAREA:
+                pass
+        normalisations = ma.array(normalisations, mask=np.logical_not(tgt_mask))
+
+        tgt_array = tgt_data * normalisations
+        tgt_array = self.tgt._matrix_to_array(tgt_array, extra)
+        return tgt_array
+
     def regrid(self, src_array, norm_type=Constants.NormType.FRACAREA, mdtol=1):
         """Perform regridding on an array of data.
 
@@ -212,30 +252,8 @@ def regrid(self, src_array, norm_type=Constants.NormType.FRACAREA, mdtol=1):
                 f"got an array with shape ending in {main_shape}."
             )
             raise ValueError(e_msg)
-        extra_shape = array_shape[: -self.src.dims]
-        extra_size = max(1, np.prod(extra_shape))
-        src_inverted_mask = self.src._array_to_matrix(~ma.getmaskarray(src_array))
-        weight_matrix = self.weight_matrix
-        if self.method == Constants.Method.NEAREST:
-            # force out_dtype := in_dtype
-            weight_matrix = weight_matrix.astype(src_array.dtype)
-        weight_sums = weight_matrix @ src_inverted_mask
-        out_dtype = self._out_dtype(src_array.dtype)
-        # Set the minimum mdtol to be slightly higher than 0 to account for rounding
-        # errors.
-        mdtol = max(mdtol, 1e-8)
-        tgt_mask = weight_sums > 1 - mdtol
-        normalisations = np.ones([self.tgt.size, extra_size], dtype=out_dtype)
-        if self.method != Constants.Method.NEAREST:
-            masked_weight_sums = weight_sums * tgt_mask
-            if norm_type == Constants.NormType.FRACAREA:
-                normalisations[tgt_mask] /= masked_weight_sums[tgt_mask]
-            elif norm_type == Constants.NormType.DSTAREA:
-                pass
-        normalisations = ma.array(normalisations, mask=np.logical_not(tgt_mask))
-
-        flat_src = self.src._array_to_matrix(ma.filled(src_array, 0.0))
-        flat_tgt = weight_matrix @ flat_src
-        flat_tgt = flat_tgt * normalisations
-        tgt_array = self.tgt._matrix_to_array(flat_tgt, extra_shape)
+        tgt_weights, tgt_data, extra = self._gen_weights_and_data(src_array)
+        tgt_array = self._regrid_from_weights_and_data(
+            tgt_weights, tgt_data, extra, norm_type=norm_type, mdtol=mdtol
+        )
         return tgt_array

src/esmf_regrid/experimental/_partial.py

@@ -0,0 +1,124 @@
+"""Provides a regridder class compatible with Partition."""
+
+import numpy as np
+
+from esmf_regrid.schemes import (
+    GridRecord,
+    MeshRecord,
+    _create_cube,
+    _ESMFRegridder,
+)
+
+
+class PartialRegridder(_ESMFRegridder):
+    """Regridder class designed for use in :class:`~esmf_regrid.experimental.Partition`."""
+
+    def __init__(self, src, tgt, src_slice, tgt_slice, weights, scheme, **kwargs):
+        """Create a regridder instance for a block of :class:`~esmf_regrid.experimental.Partition`.
+
+        Parameters
+        ----------
+        src : :class:`iris.cube.Cube`
+            The :class:`~iris.cube.Cube` providing the source.
+        tgt : :class:`iris.cube.Cube` or :class:`iris.mesh.MeshXY`
+            The :class:`~iris.cube.Cube` or :class:`~iris.mesh.MeshXY` providing the target.
+        src_slice : tuple
+            The upper and lower bounds of the block taken from the original source from which the
+            ``src`` was derived. In the form ((x_low, x_high), ...) where x_low and x_high are the
+            upper and lower bounds of the slice (in the x dimension) taken from the original source.
+            There are as many tuples of upper and lower bounds as there are horizontal dimensions in
+            the source cube (currently this is always 2 as Meshes are not yet supported for sources).
+        tgt_slice : tuple
+            The upper and lower bounds of the block taken from the original target from which the
+            ``tgt`` was derived. In the form ((x_low, x_high), ...) where x_low and x_high are the
+            upper and lower bounds of the slice (in the x dimension) taken from the original target.
+            There are as many tuples of upper and lower bounds as there are horizontal dimensions in
+            the target cube.
+        weights : :class:`scipy.sparse.spmatrix`
+            The weights to use for regridding.
+        scheme : :class:`~esmf_regrid.schemes.ESMFAreaWeighted` or :class:`~esmf_regrid.schemes.ESMFBilinear`
+            The scheme used to construct the regridder.
+        kwargs : dict
+            Additional keyword arguments to pass to the `scheme`s regridder method.
+        """
+        self.src_slice = src_slice  # this will be tuple-like
+        self.tgt_slice = tgt_slice
+        self.scheme = scheme
+
+        self._regridder = scheme.regridder(
+            src,
+            tgt,
+            precomputed_weights=weights,
+            **kwargs,
+        )
+        self.__dict__.update(self._regridder.__dict__)
+
+    def __repr__(self):
+        """Return a representation of the class."""
+        result = (
+            f"PartialRegridder("
+            f"src={self._src}, "
+            f"tgt_slice={self._tgt}, "
+            f"src_slice={self.src_slice}, "
+            f"tgt_slice={self.tgt_slice}, "
+            f"scheme={self.scheme})"
+        )
+        return result
+
+    def partial_regrid(self, src):
+        """Perform the first half of regridding, generating weights and data."""
+        dims = self._get_cube_dims(src)
+        num_dims = len(dims)
+        standard_in_dims = [-1, -2][:num_dims]
+        data = np.moveaxis(src.data, dims, standard_in_dims)
+        result = self.regridder._gen_weights_and_data(data)
+        return result
+
+    def finish_regridding(self, src_cube, weights, data, extra):
+        """Perform the second half of regridding, combining weights and data.
+
+        This operation is used to process the combined results from all the partial
+        regridders in a Partition.
+        Since all the combined data is passed in, this operation can be done using
+        *any one* of the individual PartialRegridders.
+        However, the passed "src_cube" must be the "correct" slice of the source
+        data cube, corresponding to the 'tgt_slice' slice params it was created with.
+        It is also implicit that the 'extra' arg (additional dimensions) will be the
+        same for all partial results.
+        The `src_cube` provides coordinates for the non-horizontal dimensions of the
+        result cube, matching the dimensions of the `data` array.
+        For technical convenience, its *horizontal* coordinates need to match those
+        of the 'src' reference cube provided in regridder creation (`self._src`).
+        So, it must be the correct "corresponding slice" of the source cube.
+        """
+        src_dims = self._get_cube_dims(src_cube)
+
+        result_data = self.regridder._regrid_from_weights_and_data(weights, data, extra)
+
+        num_out_dims = self.regridder.tgt.dims
+        num_dims = len(src_dims)
+        standard_out_dims = [-1, -2][:num_out_dims]
+        if num_dims == 2 and num_out_dims == 1:
+            new_dims = [min(src_dims)]
+        elif num_dims == 1 and num_out_dims == 2:
+            # Note: this code is currently inaccessible since src_cube can't have a Mesh.
+            new_dims = [src_dims[0] + 1, src_dims[0]]
+        else:
+            new_dims = src_dims
+
+        result_data = np.moveaxis(result_data, standard_out_dims, new_dims)
+
+        if isinstance(self._tgt, GridRecord):
+            tgt_coords = self._tgt
+            out_dims = 2
+        elif isinstance(self._tgt, MeshRecord):
+            tgt_coords = self._tgt.mesh.to_MeshCoords(self._tgt.location)
+            out_dims = 1
+        else:
+            msg = "Unrecognised target information."
+            raise TypeError(msg)
+
+        result_cube = _create_cube(
+            result_data, src_cube, src_dims, tgt_coords, out_dims
+        )
+        return result_cube