From 11f5df745a09ab0bea455af77dd43a9f24bd7d61 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Wed, 17 Dec 2025 15:12:15 +0000 Subject: [PATCH 01/41] Partition MVP --- src/esmf_regrid/esmf_regridder.py | 29 +++ src/esmf_regrid/experimental/_partial.py | 41 ++++ src/esmf_regrid/experimental/io.py | 175 ++++++++++----- src/esmf_regrid/experimental/partition.py | 206 ++++++++++++++++++ src/esmf_regrid/schemes.py | 81 ++++--- .../experimental/io/partition/__init__.py | 1 + .../io/partition/test_Partition.py | 112 ++++++++++ 7 files changed, 552 insertions(+), 93 deletions(-) create mode 100644 src/esmf_regrid/experimental/_partial.py create mode 100644 src/esmf_regrid/experimental/partition.py create mode 100644 src/esmf_regrid/tests/unit/experimental/io/partition/__init__.py create mode 100644 src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py diff --git a/src/esmf_regrid/esmf_regridder.py b/src/esmf_regrid/esmf_regridder.py index 800fdbd3..73bc13a0 100644 --- a/src/esmf_regrid/esmf_regridder.py +++ b/src/esmf_regrid/esmf_regridder.py @@ -175,6 +175,35 @@ 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] + + flat_src = self.src._array_to_matrix(ma.filled(src_array, 0.0)) + flat_tgt = self.weight_matrix @ flat_src + + src_inverted_mask = self.src._array_to_matrix(~ma.getmaskarray(src_array)) + weight_sums = self.weight_matrix @ src_inverted_mask + + tgt_data = self.tgt._matrix_to_array(flat_tgt, extra_shape) + tgt_weights = self.tgt._matrix_to_array(weight_sums, extra_shape) + return tgt_weights, tgt_data + + def _regrid_from_weights_and_data(self, tgt_weights, tgt_data, 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 + masked_weight_sums = tgt_weights * tgt_mask + normalisations = np.ones_like(tgt_data) + 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 + return tgt_array + def regrid(self, src_array, norm_type=Constants.NormType.FRACAREA, mdtol=1): """Perform regridding on an array of data. diff --git a/src/esmf_regrid/experimental/_partial.py b/src/esmf_regrid/experimental/_partial.py new file mode 100644 index 00000000..c8067158 --- /dev/null +++ b/src/esmf_regrid/experimental/_partial.py @@ -0,0 +1,41 @@ +"""Provides a regridder class compatible with Partition.""" + +from esmf_regrid.schemes import ( + _ESMFRegridder, + _create_cube, +) + +class PartialRegridder(_ESMFRegridder): + def __init__(self, src, tgt, src_slice, tgt_slice, weights, scheme, **kwargs): + self.src_slice = src_slice # this will be tuple-like + self.tgt_slice = tgt_slice + self.scheme = scheme + # TODO: consider disallowing ESMFNearest (unless out of bounds can be made masked) + + # Pop duplicate kwargs. + for arg in set(kwargs.keys()).intersection(vars(self.scheme)): + kwargs.pop(arg) + + self._regridder = scheme.regridder( + src, + tgt, + precomputed_weights=weights, + **kwargs, + ) + self.__dict__.update(self._regridder.__dict__) + + def partial_regrid(self, src): + return self.regridder._gen_weights_and_data(src.data) + + def finish_regridding(self, src_cube, weights, data): + dims = self._get_cube_dims(src_cube) + + result_data = self.regridder._regrid_from_weights_and_data(weights, data) + result_cube = _create_cube( + result_data, + src_cube, + dims, + self._tgt, + len(self._tgt) + ) + return result_cube diff --git a/src/esmf_regrid/experimental/io.py b/src/esmf_regrid/experimental/io.py index e71d7365..e7161aa7 100644 --- a/src/esmf_regrid/experimental/io.py +++ b/src/esmf_regrid/experimental/io.py @@ -14,10 +14,14 @@ GridToMeshESMFRegridder, MeshToGridESMFRegridder, ) +from esmf_regrid.experimental._partial import PartialRegridder from esmf_regrid.schemes import ( ESMFAreaWeightedRegridder, ESMFBilinearRegridder, ESMFNearestRegridder, + ESMFAreaWeighted, + ESMFBilinear, + ESMFNearest, GridRecord, MeshRecord, ) @@ -28,6 +32,7 @@ ESMFNearestRegridder, GridToMeshESMFRegridder, MeshToGridESMFRegridder, + PartialRegridder, ] _REGRIDDER_NAME_MAP = {rg_class.__name__: rg_class for rg_class in SUPPORTED_REGRIDDERS} _SOURCE_NAME = "regridder_source_field" @@ -47,6 +52,8 @@ _SOURCE_RESOLUTION = "src_resolution" _TARGET_RESOLUTION = "tgt_resolution" _ESMF_ARGS = "esmf_args" +_SRC_SLICE_NAME = "src_slice" +_TGT_SLICE_NAME = "tgt_slice" _VALID_ESMF_KWARGS = [ "pole_method", "regrid_pole_npoints", @@ -118,54 +125,39 @@ def _clean_var_names(cube): con.var_name = None -def save_regridder(rg, filename): - """Save a regridder scheme instance. - - Saves any of the regridder classes, i.e. - :class:`~esmf_regrid.experimental.unstructured_scheme.GridToMeshESMFRegridder`, - :class:`~esmf_regrid.experimental.unstructured_scheme.MeshToGridESMFRegridder`, - :class:`~esmf_regrid.schemes.ESMFAreaWeightedRegridder`, - :class:`~esmf_regrid.schemes.ESMFBilinearRegridder` or - :class:`~esmf_regrid.schemes.ESMFNearestRegridder`. - . - - Parameters - ---------- - rg : :class:`~esmf_regrid.schemes._ESMFRegridder` - The regridder instance to save. - filename : str - The file name to save to. - """ - regridder_type = rg.__class__.__name__ - - def _standard_grid_cube(grid, name): - if grid[0].ndim == 1: - shape = [coord.points.size for coord in grid] - else: - shape = grid[0].shape - data = np.zeros(shape) - cube = Cube(data, var_name=name, long_name=name) - if grid[0].ndim == 1: - cube.add_dim_coord(grid[0], 0) - cube.add_dim_coord(grid[1], 1) - else: - cube.add_aux_coord(grid[0], [0, 1]) - cube.add_aux_coord(grid[1], [0, 1]) - return cube - - def _standard_mesh_cube(mesh, location, name): - mesh_coords = mesh.to_MeshCoords(location) - data = np.zeros(mesh_coords[0].points.shape[0]) - cube = Cube(data, var_name=name, long_name=name) - for coord in mesh_coords: - cube.add_aux_coord(coord, 0) - return cube - +def _standard_grid_cube(grid, name): + if grid[0].ndim == 1: + shape = [coord.points.size for coord in grid] + else: + shape = grid[0].shape + data = np.zeros(shape) + cube = Cube(data, var_name=name, long_name=name) + if grid[0].ndim == 1: + cube.add_dim_coord(grid[0], 0) + cube.add_dim_coord(grid[1], 1) + else: + cube.add_aux_coord(grid[0], [0, 1]) + cube.add_aux_coord(grid[1], [0, 1]) + return cube + +def _standard_mesh_cube(mesh, location, name): + mesh_coords = mesh.to_MeshCoords(location) + data = np.zeros(mesh_coords[0].points.shape[0]) + cube = Cube(data, var_name=name, long_name=name) + for coord in mesh_coords: + cube.add_aux_coord(coord, 0) + return cube + +def _generate_src_tgt(regridder_type, rg, allow_partial): if regridder_type in [ "ESMFAreaWeightedRegridder", "ESMFBilinearRegridder", "ESMFNearestRegridder", + "PartialRegridder", ]: + if regridder_type == "PartialRegridder" and not allow_partial: + e_msg = "To save a PartialRegridder, `allow_partial=True` must be set." + raise ValueError(e_msg) src_grid = rg._src if isinstance(src_grid, GridRecord): src_cube = _standard_grid_cube( @@ -210,12 +202,38 @@ def _standard_mesh_cube(mesh, location, name): tgt_grid = (rg.grid_y, rg.grid_x) tgt_cube = _standard_grid_cube(tgt_grid, _TARGET_NAME) _add_mask_to_cube(rg.tgt_mask, tgt_cube, _TARGET_MASK_NAME) + else: e_msg = ( - f"Expected a regridder of type `GridToMeshESMFRegridder` or " - f"`MeshToGridESMFRegridder`, got type {regridder_type}." + f"Unexpected regridder type {regridder_type}." ) raise TypeError(e_msg) + return src_cube, tgt_cube + + + + +def save_regridder(rg, filename, allow_partial=False): + """Save a regridder scheme instance. + + Saves any of the regridder classes, i.e. + :class:`~esmf_regrid.experimental.unstructured_scheme.GridToMeshESMFRegridder`, + :class:`~esmf_regrid.experimental.unstructured_scheme.MeshToGridESMFRegridder`, + :class:`~esmf_regrid.schemes.ESMFAreaWeightedRegridder`, + :class:`~esmf_regrid.schemes.ESMFBilinearRegridder` or + :class:`~esmf_regrid.schemes.ESMFNearestRegridder`. + . + + Parameters + ---------- + rg : :class:`~esmf_regrid.schemes._ESMFRegridder` + The regridder instance to save. + filename : str + The file name to save to. + """ + regridder_type = rg.__class__.__name__ + + src_cube, tgt_cube = _generate_src_tgt(regridder_type, rg, allow_partial) method = str(check_method(rg.method).name) @@ -223,7 +241,7 @@ def _standard_mesh_cube(mesh, location, name): resolution = rg.resolution src_resolution = None tgt_resolution = None - elif regridder_type == "ESMFAreaWeightedRegridder": + elif method == "CONSERVATIVE": resolution = None src_resolution = rg.src_resolution tgt_resolution = rg.tgt_resolution @@ -264,6 +282,19 @@ def _standard_mesh_cube(mesh, location, name): if tgt_resolution is not None: attributes[_TARGET_RESOLUTION] = tgt_resolution + extra_cubes = [] + if regridder_type == "PartialRegridder": + src_slice = rg.src_slice # this slice is described by a tuple + if src_slice is None: + src_slice = [] + src_slice_cube = Cube(src_slice, long_name=_SRC_SLICE_NAME, var_name=_SRC_SLICE_NAME) + tgt_slice = rg.tgt_slice # this slice is described by a tuple + if tgt_slice is None: + tgt_slice = [] + tgt_slice_cube = Cube(src_slice, long_name=_TGT_SLICE_NAME, var_name=_TGT_SLICE_NAME) + extra_cubes = [src_slice_cube, tgt_slice_cube] + + weights_cube = Cube(weight_data, var_name=_WEIGHTS_NAME, long_name=_WEIGHTS_NAME) row_coord = AuxCoord( weight_rows, var_name=_WEIGHTS_ROW_NAME, long_name=_WEIGHTS_ROW_NAME @@ -298,7 +329,7 @@ def _standard_mesh_cube(mesh, location, name): # Save cubes while ensuring var_names do not conflict for the sake of consistency. with _managed_var_name(src_cube, tgt_cube): - cube_list = CubeList([src_cube, tgt_cube, weights_cube, weight_shape_cube]) + cube_list = CubeList([src_cube, tgt_cube, weights_cube, weight_shape_cube, *extra_cubes]) for cube in cube_list: cube.attributes = attributes @@ -306,7 +337,7 @@ def _standard_mesh_cube(mesh, location, name): iris.fileformats.netcdf.save(cube_list, filename) -def load_regridder(filename): +def load_regridder(filename, allow_partial=False): """Load a regridder scheme instance. Loads any of the regridder classes, i.e. @@ -343,6 +374,10 @@ def load_regridder(filename): raise TypeError(e_msg) scheme = _REGRIDDER_NAME_MAP[regridder_type] + if regridder_type == "PartialRegridder" and not allow_partial: + e_msg = "PartialRegridder cannot be loaded without setting `allow_partial=True`." + raise ValueError(e_msg) + # Determine the regridding method, allowing for files created when # conservative regridding was the only method. method_string = weights_cube.attributes.get(_METHOD, "CONSERVATIVE") @@ -396,26 +431,48 @@ def load_regridder(filename): elif scheme is MeshToGridESMFRegridder: resolution_keyword = _TARGET_RESOLUTION kwargs = {resolution_keyword: resolution, "method": method, "mdtol": mdtol} - elif scheme is ESMFAreaWeightedRegridder: + elif method is Constants.Method.CONSERVATIVE: kwargs = { _SOURCE_RESOLUTION: src_resolution, _TARGET_RESOLUTION: tgt_resolution, "mdtol": mdtol, } - elif scheme is ESMFBilinearRegridder: + elif method is Constants.Method.BILINEAR: kwargs = {"mdtol": mdtol} else: kwargs = {} - regridder = scheme( - src_cube, - tgt_cube, - precomputed_weights=weight_matrix, - use_src_mask=use_src_mask, - use_tgt_mask=use_tgt_mask, - esmf_args=esmf_args, - **kwargs, - ) + if scheme is PartialRegridder: + src_slice = cubes.extract_cube(_SRC_SLICE_NAME).data.tolist() + if src_slice == []: + src_slice = None + tgt_slice = cubes.extract_cube(_TGT_SLICE_NAME).data.tolist() + if tgt_slice == []: + tgt_slice = None + sub_scheme = { + Constants.Method.CONSERVATIVE: ESMFAreaWeighted, + Constants.Method.BILINEAR: ESMFBilinear, + Constants.Method.NEAREST: ESMFNearest, + }[method] + regridder = scheme( + src_cube, + tgt_cube, + src_slice, + tgt_slice, + weight_matrix, + sub_scheme(), + **kwargs, + ) + else: + regridder = scheme( + src_cube, + tgt_cube, + precomputed_weights=weight_matrix, + use_src_mask=use_src_mask, + use_tgt_mask=use_tgt_mask, + esmf_args=esmf_args, + **kwargs, + ) esmf_version = weights_cube.attributes[_VERSION_ESMF] regridder.regridder.esmf_version = esmf_version diff --git a/src/esmf_regrid/experimental/partition.py b/src/esmf_regrid/experimental/partition.py new file mode 100644 index 00000000..22ed9381 --- /dev/null +++ b/src/esmf_regrid/experimental/partition.py @@ -0,0 +1,206 @@ +"""Provides an interface for splitting up a large regridding task.""" + +import numpy as np + +from esmf_regrid.constants import Constants +from esmf_regrid.experimental.io import load_regridder, save_regridder +from esmf_regrid.experimental._partial import PartialRegridder +from esmf_regrid.schemes import _get_grid_dims + + +def _get_chunk(cube, sl): + if cube.mesh is None: + grid_dims = _get_grid_dims(cube) + else: + grid_dims = (cube.mesh_dim(),) + slice = [np.s_[:]] * len(cube.shape) + for s, d in zip(sl, grid_dims): + slice[d] = np.s_[s[0]:s[1]] + return cube[*slice] + +def _determine_blocks(shape, chunks, num_chunks, explicit_chunks): + which_inputs = (chunks is not None, num_chunks is not None, explicit_chunks is not None) + if sum(which_inputs) == 0: + raise ValueError("Partition blocks must must be specified by either chunks, num_chunks, or explicit_chunks.") + if sum(which_inputs) > 1: + raise ValueError("Potentially conflicting partition block definitions.") + if num_chunks is not None: + chunks = [s//n for s, n in zip(shape, num_chunks)] + for chunk in chunks: + if chunk == 0: + raise ValueError("`num_chunks` cannot divide a dimension into more blocks than the size of that dimension.") + if chunks is not None: + if all(isinstance(x, int)for x in chunks): + proper_chunks = [] + for s, c in zip(shape, chunks): + proper_chunk = [c] * (s//c) + if s%c != 0: + proper_chunk += [s%c] + proper_chunks.append(proper_chunk) + chunks = proper_chunks + for s, chunk in zip(shape, chunks): + if sum(chunk) != s: + raise ValueError("Chunks must sum to the size of their respective dimension.") + bounds = [np.cumsum([0] + list(chunk)) for chunk in chunks] + if len(bounds) == 1: + explicit_chunks = [[[int(lower), int(upper)]] for lower, upper in zip(bounds[0][:-1], bounds[0][1:])] + elif len(bounds) == 2: + explicit_chunks = [ + [[int(ly), int(uy)], [int(lx), int(ux)]] for ly, uy in zip(bounds[0][:-1], bounds[0][1:]) for lx, ux in zip(bounds[1][:-1], bounds[1][1:]) + ] + else: + raise ValueError(f"Chunks must not exceed two dimensions.") + return explicit_chunks + +class Partition: + def __init__( + self, + src, + tgt, + scheme, + file_names, + use_dask_src_chunks=False, + src_chunks=None, + num_src_chunks=None, + explicit_src_chunks=None, + # tgt_chunks=None, + # num_tgt_chunks=None, + auto_generate=False, + saved_files=None, + ): + """Class for breaking down regridding into manageable chunks. + + Note + ---- + Currently, it is only possible to divide the source grid into chunks. + Meshes are not yet supported as a source. + + Parameters + ---------- + src : cube + Source cube. + tgt : cube + Target cube. + scheme : + scheme + file_names : iterable of str + A list of file names to save/load parts of the regridder to/from. + use_dask_src_chunks : bool, default=False + If true, partition using the same chunks from the source cube. + src_chunks : numpy array, tuple of int or str, default=None + Specify the size of blocks to use to divide up the cube. Demensions are specified + in y,x axis order. If `src_chunks` is a tuple of int, each integer describes + the maximum size of a block in that dimension. If `src_chunks` is a tuple of int, + each tuple describes the size of each successive block in that dimension. These + block sizes should add up to the total size of that dimension or else an error + is raised. + num_src_chunks : tuple of int + Specify the number of blocks to use to divide up the cube. Demensions are specified + in y,x axis order. Each integer describes the number of blocks that dimension will + be divided into. + explicit_src_chunks : arraylike NxMx2 + Explicitly specify the bounds of each block in the partition. + # tgt_chunks : ???, default=None + # ??? + # num_tgt_chunks : tuple of int + # ??? + auto_generate : bool, default=False + When true, start generating files on initialisation. + saved_files : iterable of str + A list of paths to previously saved files. + """ + if scheme._method == Constants.Method.NEAREST: + raise NotImplementedError("The `Nearest` method is not implemented.") + if src.mesh is not None: + raise NotImplementedError("Partition does not yet support source meshes.") + # TODO Extract a slice of the cube. + self.src = src + if src.mesh is None: + grid_dims = _get_grid_dims(src) + else: + grid_dims = (src.mesh_dim(),) + shape = tuple(src.shape[i] for i in grid_dims) + self.tgt = tgt + self.scheme = scheme + # TODO: consider abstracting away the idea of files + self.file_names = file_names + if use_dask_src_chunks: + assert num_src_chunks is None and src_chunks is None + assert src.has_lazy_data() + src_chunks = src.slices(grid_dims).next().lazy_data().chunks + self.src_chunks = _determine_blocks(shape, src_chunks, num_src_chunks, explicit_src_chunks) + assert len(self.src_chunks) == len(file_names) + # This will be controllable in future + tgt_chunks = None + self.tgt_chunks = tgt_chunks + assert tgt_chunks is None # We don't handle big targets currently + + # Note: this may need to become more sophisticated when both src and tgt are large + self.file_chunk_dict = {file: chunk for file, chunk in zip(self.file_names, self.src_chunks)} + + if saved_files is None: + self.saved_files = [] + else: + self.saved_files = saved_files + if auto_generate: + self.generate_files(self.file_names) + + def __repr__(self): + """Return a representation of the class.""" + result = ( + f"Partition(" + f"src={self.src}, " + f"tgt={self.tgt}, " + f"scheme={self.scheme}, " + f"num file_names={len(self.file_names)}," + f"num saved_files={len(self.saved_files)})" + ) + return result + + @property + def unsaved_files(self): + files = set(self.file_names) - set(self.saved_files) + return [file for file in self.file_names if file in files] + + def generate_files(self, files_to_generate=None): + if files_to_generate is None: + files = self.unsaved_files + else: + assert isinstance(files_to_generate, int) + files = self.unsaved_files[:files_to_generate] + + for file in files: + src_chunk = self.file_chunk_dict[file] + src = _get_chunk(self.src, src_chunk) + tgt = self.tgt + regridder = self.scheme.regridder(src, tgt) + weights = regridder.regridder.weight_matrix + regridder = PartialRegridder(src, self.tgt, src_chunk, None, weights, self.scheme) + save_regridder(regridder, file, allow_partial=True) + self.saved_files.append(file) + + def apply_regridders(self, cube, allow_incomplete=False): + # for each target chunk, iterate through each associated regridder + # for now, assume one target chunk + if not allow_incomplete: + assert len(self.unsaved_files) == 0 + current_result = None + current_weights = None + files = self.saved_files + + for file, chunk in zip(self.file_names, self.src_chunks): + if file in files: + next_regridder = load_regridder(file, allow_partial=True) + # cube_chunk = cube[*_interpret_slice(chunk)] + cube_chunk = _get_chunk(cube, chunk) + next_weights, next_result = next_regridder.partial_regrid(cube_chunk) + if current_weights is None: + current_weights = next_weights + else: + current_weights += next_weights + if current_result is None: + current_result = next_result + else: + current_result += next_result + + return next_regridder.finish_regridding(cube_chunk, current_weights, current_result) diff --git a/src/esmf_regrid/schemes.py b/src/esmf_regrid/schemes.py index adfe87b9..739fa41c 100644 --- a/src/esmf_regrid/schemes.py +++ b/src/esmf_regrid/schemes.py @@ -606,6 +606,18 @@ def _make_meshinfo(cube_or_mesh, method, mask, src_or_tgt, location=None): return _mesh_to_MeshInfo(mesh, location, mask=mask) +def _get_grid_dims(cube): + src_x = _get_coord(cube, "x") + src_y = _get_coord(cube, "y") + + if len(src_x.shape) == 1: + grid_x_dim = cube.coord_dims(src_x)[0] + grid_y_dim = cube.coord_dims(src_y)[0] + else: + grid_y_dim, grid_x_dim = cube.coord_dims(src_x) + return grid_y_dim, grid_x_dim + + def _regrid_rectilinear_to_rectilinear__prepare( src_grid_cube, tgt_grid_cube, @@ -625,14 +637,8 @@ def _regrid_rectilinear_to_rectilinear__prepare( """ tgt_x = _get_coord(tgt_grid_cube, "x") tgt_y = _get_coord(tgt_grid_cube, "y") - src_x = _get_coord(src_grid_cube, "x") - src_y = _get_coord(src_grid_cube, "y") - if len(src_x.shape) == 1: - grid_x_dim = src_grid_cube.coord_dims(src_x)[0] - grid_y_dim = src_grid_cube.coord_dims(src_y)[0] - else: - grid_y_dim, grid_x_dim = src_grid_cube.coord_dims(src_x) + grid_y_dim, grid_x_dim = _get_grid_dims(src_grid_cube) srcinfo = _make_gridinfo(src_grid_cube, method, src_resolution, src_mask) tgtinfo = _make_gridinfo(tgt_grid_cube, method, tgt_resolution, tgt_mask) @@ -805,8 +811,6 @@ def _regrid_rectilinear_to_unstructured__prepare( The 'regrid info' returned can be reused over many 2d slices. """ - grid_x = _get_coord(src_grid_cube, "x") - grid_y = _get_coord(src_grid_cube, "y") if isinstance(tgt_cube_or_mesh, MeshXY): mesh = tgt_cube_or_mesh location = tgt_location @@ -814,11 +818,7 @@ def _regrid_rectilinear_to_unstructured__prepare( mesh = tgt_cube_or_mesh.mesh location = tgt_cube_or_mesh.location - if grid_x.ndim == 1: - (grid_x_dim,) = src_grid_cube.coord_dims(grid_x) - (grid_y_dim,) = src_grid_cube.coord_dims(grid_y) - else: - grid_y_dim, grid_x_dim = src_grid_cube.coord_dims(grid_x) + grid_y_dim, grid_x_dim = _get_grid_dims(src_grid_cube) meshinfo = _make_meshinfo( tgt_cube_or_mesh, method, tgt_mask, "target", location=tgt_location @@ -1087,6 +1087,7 @@ def __init__( the regridder is saved . """ + self._method = Constants.Method.CONSERVATIVE if not (0 <= mdtol <= 1): msg = "Value for mdtol must be in range 0 - 1, got {}." raise ValueError(msg.format(mdtol)) @@ -1123,6 +1124,7 @@ def regridder( use_tgt_mask=None, tgt_location="face", esmf_args=None, + precomputed_weights=None, ): """Create regridder to perform regridding from ``src_grid`` to ``tgt_grid``. @@ -1191,6 +1193,7 @@ def regridder( use_tgt_mask=use_tgt_mask, tgt_location="face", esmf_args=esmf_args, + precomputed_weights=precomputed_weights, ) @@ -1240,6 +1243,7 @@ def __init__( the regridder is saved . """ + self._method = Constants.Method.BILINEAR if not (0 <= mdtol <= 1): msg = "Value for mdtol must be in range 0 - 1, got {}." raise ValueError(msg.format(mdtol)) @@ -1274,6 +1278,7 @@ def regridder( tgt_location=None, extrapolate_gaps=False, esmf_args=None, + precomputed_weights=None, ): """Create regridder to perform regridding from ``src_grid`` to ``tgt_grid``. @@ -1336,6 +1341,7 @@ def regridder( use_tgt_mask=use_tgt_mask, tgt_location=tgt_location, esmf_args=esmf_args, + precomputed_weights=precomputed_weights, ) @@ -1389,6 +1395,7 @@ def __init__( arguments are recorded as a property of this regridder and are stored when the regridder is saved . """ + self._method = Constants.Method.NEAREST self.use_src_mask = use_src_mask self.use_tgt_mask = use_tgt_mask self.tgt_location = tgt_location @@ -1415,6 +1422,7 @@ def regridder( use_tgt_mask=None, tgt_location=None, esmf_args=None, + precomputed_weights=None, ): """Create regridder to perform regridding from ``src_grid`` to ``tgt_grid``. @@ -1468,6 +1476,7 @@ def regridder( use_tgt_mask=use_tgt_mask, tgt_location=tgt_location, esmf_args=esmf_args, + precomputed_weights=precomputed_weights, ) @@ -1566,26 +1575,7 @@ def __init__( else: self._src = GridRecord(_get_coord(src, "x"), _get_coord(src, "y")) - def __call__(self, cube): - """Regrid this :class:`~iris.cube.Cube` onto the target grid of this regridder instance. - - The given :class:`~iris.cube.Cube` must be defined with the same grid as the source - :class:`~iris.cube.Cube` used to create this :class:`_ESMFRegridder` instance. - - Parameters - ---------- - cube : :class:`iris.cube.Cube` - A :class:`~iris.cube.Cube` instance to be regridded. - - Returns - ------- - :class:`iris.cube.Cube` - A :class:`~iris.cube.Cube` defined with the horizontal dimensions of the target - and the other dimensions from this :class:`~iris.cube.Cube`. The data values of - this :class:`~iris.cube.Cube` will be converted to values on the new grid using - regridding via :mod:`esmpy` generated weights. - - """ + def _get_cube_dims(self, cube): if cube.mesh is not None: # TODO: replace temporary hack when iris issues are sorted. # Ignore differences in var_name that might be caused by saving. @@ -1629,6 +1619,29 @@ def __call__(self, cube): else: # Due to structural reasons, the order here must be reversed. dims = cube.coord_dims(new_src_x)[::-1] + return dims + + def __call__(self, cube): + """Regrid this :class:`~iris.cube.Cube` onto the target grid of this regridder instance. + + The given :class:`~iris.cube.Cube` must be defined with the same grid as the source + :class:`~iris.cube.Cube` used to create this :class:`_ESMFRegridder` instance. + + Parameters + ---------- + cube : :class:`iris.cube.Cube` + A :class:`~iris.cube.Cube` instance to be regridded. + + Returns + ------- + :class:`iris.cube.Cube` + A :class:`~iris.cube.Cube` defined with the horizontal dimensions of the target + and the other dimensions from this :class:`~iris.cube.Cube`. The data values of + this :class:`~iris.cube.Cube` will be converted to values on the new grid using + regridding via :mod:`esmpy` generated weights. + + """ + dims = self._get_cube_dims(cube) regrid_info = RegridInfo( dims=dims, diff --git a/src/esmf_regrid/tests/unit/experimental/io/partition/__init__.py b/src/esmf_regrid/tests/unit/experimental/io/partition/__init__.py new file mode 100644 index 00000000..773355ff --- /dev/null +++ b/src/esmf_regrid/tests/unit/experimental/io/partition/__init__.py @@ -0,0 +1 @@ +"""Unit tests for :mod:`esmf_regrid.experimental.partition`.""" \ No newline at end of file diff --git a/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py b/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py new file mode 100644 index 00000000..99083f68 --- /dev/null +++ b/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py @@ -0,0 +1,112 @@ +"""Unit tests for :mod:`esmf_regrid.experimental.partition`.""" + +import dask.array as da +import numpy as np +import pytest + +import esmf_regrid +print(esmf_regrid.__file__) +from esmf_regrid import ESMFAreaWeighted +from esmf_regrid.experimental.partition import Partition + +from esmf_regrid.tests.unit.schemes.test__cube_to_GridInfo import ( + _curvilinear_cube, + _grid_cube, +) +from esmf_regrid.tests.unit.schemes.test__mesh_to_MeshInfo import ( + _gridlike_mesh_cube, +) + +def test_Partition(tmp_path): + src = _grid_cube(150, 500, (-180, 180), (-90, 90), circular=True) + src.data = np.arange(150*500).reshape([500, 150]) + tgt = _grid_cube(16, 36, (-180, 180), (-90, 90), circular=True) + + files = [tmp_path / f"partial_{x}.nc" for x in range(5)] + scheme = ESMFAreaWeighted(mdtol=1) + chunks = [[[0, 100], [0, 150]], [[100, 200], [0, 150]], [[200, 300], [0, 150]], [[300, 400], [0, 150]], [[400, 500], [0, 150]]] + + partition = Partition(src, tgt, scheme, files, explicit_src_chunks=chunks) + + partition.generate_files() + + result = partition.apply_regridders(src) + expected = src.regrid(tgt, scheme) + assert np.allclose(result.data, expected.data) + assert result == expected + +def test_Partition_block_api(tmp_path): + src = _grid_cube(150, 500, (-180, 180), (-90, 90), circular=True) + src.data = np.arange(150*500).reshape([500, 150]) + tgt = _grid_cube(16, 36, (-180, 180), (-90, 90), circular=True) + + files = [tmp_path / f"partial_{x}.nc" for x in range(5)] + scheme = ESMFAreaWeighted(mdtol=1) + num_src_chunks = (5, 1) + partition = Partition(src, tgt, scheme, files, num_src_chunks=num_src_chunks) + + expected_chunks = [[[0, 100], [0, 150]], [[100, 200], [0, 150]], [[200, 300], [0, 150]], [[300, 400], [0, 150]], [[400, 500], [0, 150]]] + assert partition.src_chunks == expected_chunks + + src_chunks = (100, 150) + partition = Partition(src, tgt, scheme, files, src_chunks=src_chunks) + + expected_chunks = [[[0, 100], [0, 150]], [[100, 200], [0, 150]], [[200, 300], [0, 150]], [[300, 400], [0, 150]], [[400, 500], [0, 150]]] + assert partition.src_chunks == expected_chunks + + + src_chunks = ((100, 100, 100, 100, 100), (150,)) + partition = Partition(src, tgt, scheme, files, src_chunks=src_chunks) + + expected_chunks = [[[0, 100], [0, 150]], [[100, 200], [0, 150]], [[200, 300], [0, 150]], [[300, 400], [0, 150]], [[400, 500], [0, 150]]] + assert partition.src_chunks == expected_chunks + + src.data = da.from_array(src.data, chunks=src_chunks) + partition = Partition(src, tgt, scheme, files, use_dask_src_chunks=True) + + expected_chunks = [[[0, 100], [0, 150]], [[100, 200], [0, 150]], [[200, 300], [0, 150]], [[300, 400], [0, 150]], [[400, 500], [0, 150]]] + assert partition.src_chunks == expected_chunks + +def test_Partition_mesh_src(tmp_path): + src = _gridlike_mesh_cube(150, 500) + src.data = np.arange(150*500) + tgt = _grid_cube(16, 36, (-180, 180), (-90, 90), circular=True) + + files = [tmp_path / f"partial_{x}.nc" for x in range(5)] + scheme = ESMFAreaWeighted(mdtol=1) + + src_chunks = (15000,) + with pytest.raises(NotImplementedError): + partition = Partition(src, tgt, scheme, files, src_chunks=src_chunks) + + # TODO when mesh partitioning becomes possible, uncomment. + # expected_src_chunks = [[[0, 15000]], [[15000, 30000]], [[30000, 45000]], [[45000, 60000]], [[60000, 75000]]] + # assert partition.src_chunks == expected_src_chunks + # + # partition.generate_files() + # + # result = partition.apply_regridders(src) + # expected = src.regrid(tgt, scheme) + # assert np.allclose(result.data, expected.data) + # assert result == expected + +def test_Partition_curv_src(tmp_path): + src = _curvilinear_cube(150, 500, (-180, 180), (-90, 90)) + src.data = np.arange(150*500).reshape([500, 150]) + tgt = _grid_cube(16, 36, (-180, 180), (-90, 90), circular=True) + + files = [tmp_path / f"partial_{x}.nc" for x in range(5)] + scheme = ESMFAreaWeighted(mdtol=1) + + src_chunks = (100, 150) + partition = Partition(src, tgt, scheme, files, src_chunks=src_chunks) + + expected_src_chunks = [[[0, 100], [0, 150]], [[100, 200], [0, 150]], [[200, 300], [0, 150]], [[300, 400], [0, 150]], [[400, 500], [0, 150]]] + assert partition.src_chunks == expected_src_chunks + + partition.generate_files() + + result = partition.apply_regridders(src) + expected = src.regrid(tgt, scheme) + assert np.allclose(result.data, expected.data) + assert result == expected \ No newline at end of file From 8ad69837f78549ab74e88d14a7148598385c1b81 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Wed, 17 Dec 2025 15:31:23 +0000 Subject: [PATCH 02/41] lint fixes --- pyproject.toml | 3 +- src/esmf_regrid/experimental/partition.py | 28 ++++++++++++------- .../experimental/io/partition/__init__.py | 2 +- .../io/partition/test_Partition.py | 2 +- 4 files changed, 22 insertions(+), 13 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 6d1a53c6..5d42c581 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -180,6 +180,7 @@ ignore = [ "ANN202", "ANN204", + "B905", # Zip strictness should be explicit "D104", # Misssing docstring "E501", # Line too long "ERA001", # Commented out code @@ -297,5 +298,5 @@ convention = "numpy" [tool.ruff.lint.pylint] # TODO: refactor to reduce complexity, if possible max-args = 10 -max-branches = 21 +max-branches = 22 max-statements = 110 diff --git a/src/esmf_regrid/experimental/partition.py b/src/esmf_regrid/experimental/partition.py index 22ed9381..189da1ae 100644 --- a/src/esmf_regrid/experimental/partition.py +++ b/src/esmf_regrid/experimental/partition.py @@ -13,22 +13,25 @@ def _get_chunk(cube, sl): grid_dims = _get_grid_dims(cube) else: grid_dims = (cube.mesh_dim(),) - slice = [np.s_[:]] * len(cube.shape) + full_slice = [np.s_[:]] * len(cube.shape) for s, d in zip(sl, grid_dims): - slice[d] = np.s_[s[0]:s[1]] - return cube[*slice] + full_slice[d] = np.s_[s[0]:s[1]] + return cube[*full_slice] def _determine_blocks(shape, chunks, num_chunks, explicit_chunks): which_inputs = (chunks is not None, num_chunks is not None, explicit_chunks is not None) if sum(which_inputs) == 0: - raise ValueError("Partition blocks must must be specified by either chunks, num_chunks, or explicit_chunks.") + msg = "Partition blocks must must be specified by either chunks, num_chunks, or explicit_chunks." + raise ValueError(msg) if sum(which_inputs) > 1: - raise ValueError("Potentially conflicting partition block definitions.") + msg = "Potentially conflicting partition block definitions." + raise ValueError(msg) if num_chunks is not None: chunks = [s//n for s, n in zip(shape, num_chunks)] for chunk in chunks: if chunk == 0: - raise ValueError("`num_chunks` cannot divide a dimension into more blocks than the size of that dimension.") + msg = "`num_chunks` cannot divide a dimension into more blocks than the size of that dimension." + raise ValueError(msg) if chunks is not None: if all(isinstance(x, int)for x in chunks): proper_chunks = [] @@ -40,7 +43,8 @@ def _determine_blocks(shape, chunks, num_chunks, explicit_chunks): chunks = proper_chunks for s, chunk in zip(shape, chunks): if sum(chunk) != s: - raise ValueError("Chunks must sum to the size of their respective dimension.") + msg = "Chunks must sum to the size of their respective dimension." + raise ValueError(msg) bounds = [np.cumsum([0] + list(chunk)) for chunk in chunks] if len(bounds) == 1: explicit_chunks = [[[int(lower), int(upper)]] for lower, upper in zip(bounds[0][:-1], bounds[0][1:])] @@ -49,10 +53,12 @@ def _determine_blocks(shape, chunks, num_chunks, explicit_chunks): [[int(ly), int(uy)], [int(lx), int(ux)]] for ly, uy in zip(bounds[0][:-1], bounds[0][1:]) for lx, ux in zip(bounds[1][:-1], bounds[1][1:]) ] else: - raise ValueError(f"Chunks must not exceed two dimensions.") + msg = "Chunks must not exceed two dimensions." + raise ValueError(msg) return explicit_chunks class Partition: + """Class for breaking down regridding into manageable chunks.""" def __init__( self, src, @@ -110,9 +116,11 @@ def __init__( A list of paths to previously saved files. """ if scheme._method == Constants.Method.NEAREST: - raise NotImplementedError("The `Nearest` method is not implemented.") + msg = "The `Nearest` method is not implemented." + raise NotImplementedError(msg) if src.mesh is not None: - raise NotImplementedError("Partition does not yet support source meshes.") + msg = "Partition does not yet support source meshes." + raise NotImplementedError(msg) # TODO Extract a slice of the cube. self.src = src if src.mesh is None: diff --git a/src/esmf_regrid/tests/unit/experimental/io/partition/__init__.py b/src/esmf_regrid/tests/unit/experimental/io/partition/__init__.py index 773355ff..656fc3a9 100644 --- a/src/esmf_regrid/tests/unit/experimental/io/partition/__init__.py +++ b/src/esmf_regrid/tests/unit/experimental/io/partition/__init__.py @@ -1 +1 @@ -"""Unit tests for :mod:`esmf_regrid.experimental.partition`.""" \ No newline at end of file +"""Unit tests for :mod:`esmf_regrid.experimental.partition`.""" diff --git a/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py b/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py index 99083f68..d4159580 100644 --- a/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py +++ b/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py @@ -109,4 +109,4 @@ def test_Partition_curv_src(tmp_path): result = partition.apply_regridders(src) expected = src.regrid(tgt, scheme) assert np.allclose(result.data, expected.data) - assert result == expected \ No newline at end of file + assert result == expected From bc2961415dbeadca50de2869c527f985f6ef8aaa Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Wed, 17 Dec 2025 15:41:53 +0000 Subject: [PATCH 03/41] lint fixes --- src/esmf_regrid/experimental/partition.py | 6 +++--- .../tests/unit/experimental/io/partition/test_Partition.py | 3 --- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/src/esmf_regrid/experimental/partition.py b/src/esmf_regrid/experimental/partition.py index 189da1ae..2bb4b1af 100644 --- a/src/esmf_regrid/experimental/partition.py +++ b/src/esmf_regrid/experimental/partition.py @@ -45,7 +45,7 @@ def _determine_blocks(shape, chunks, num_chunks, explicit_chunks): if sum(chunk) != s: msg = "Chunks must sum to the size of their respective dimension." raise ValueError(msg) - bounds = [np.cumsum([0] + list(chunk)) for chunk in chunks] + bounds = [np.cumsum([0, *chunk]) for chunk in chunks] if len(bounds) == 1: explicit_chunks = [[[int(lower), int(upper)]] for lower, upper in zip(bounds[0][:-1], bounds[0][1:])] elif len(bounds) == 2: @@ -121,7 +121,7 @@ def __init__( if src.mesh is not None: msg = "Partition does not yet support source meshes." raise NotImplementedError(msg) - # TODO Extract a slice of the cube. + # TODO: Extract a slice of the cube. self.src = src if src.mesh is None: grid_dims = _get_grid_dims(src) @@ -144,7 +144,7 @@ def __init__( assert tgt_chunks is None # We don't handle big targets currently # Note: this may need to become more sophisticated when both src and tgt are large - self.file_chunk_dict = {file: chunk for file, chunk in zip(self.file_names, self.src_chunks)} + self.file_chunk_dict = dict(zip(self.file_names, self.src_chunks)) if saved_files is None: self.saved_files = [] diff --git a/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py b/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py index d4159580..fe9d708e 100644 --- a/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py +++ b/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py @@ -4,11 +4,8 @@ import numpy as np import pytest -import esmf_regrid -print(esmf_regrid.__file__) from esmf_regrid import ESMFAreaWeighted from esmf_regrid.experimental.partition import Partition - from esmf_regrid.tests.unit.schemes.test__cube_to_GridInfo import ( _curvilinear_cube, _grid_cube, From e9fa31aba85359e2b9dbe5ec692da99b568e96c7 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Thu, 18 Dec 2025 12:10:35 +0000 Subject: [PATCH 04/41] lint fixes and name changes --- src/esmf_regrid/experimental/partition.py | 46 +++++++++++-------- .../io/partition/test_Partition.py | 18 ++++---- 2 files changed, 37 insertions(+), 27 deletions(-) diff --git a/src/esmf_regrid/experimental/partition.py b/src/esmf_regrid/experimental/partition.py index 2bb4b1af..3bf3ae1b 100644 --- a/src/esmf_regrid/experimental/partition.py +++ b/src/esmf_regrid/experimental/partition.py @@ -68,7 +68,7 @@ def __init__( use_dask_src_chunks=False, src_chunks=None, num_src_chunks=None, - explicit_src_chunks=None, + explicit_src_blocks=None, # tgt_chunks=None, # num_tgt_chunks=None, auto_generate=False, @@ -94,17 +94,17 @@ def __init__( use_dask_src_chunks : bool, default=False If true, partition using the same chunks from the source cube. src_chunks : numpy array, tuple of int or str, default=None - Specify the size of blocks to use to divide up the cube. Demensions are specified + Specify the size of blocks to use to divide up the cube. Dimensions are specified in y,x axis order. If `src_chunks` is a tuple of int, each integer describes the maximum size of a block in that dimension. If `src_chunks` is a tuple of int, each tuple describes the size of each successive block in that dimension. These block sizes should add up to the total size of that dimension or else an error is raised. num_src_chunks : tuple of int - Specify the number of blocks to use to divide up the cube. Demensions are specified + Specify the number of blocks to use to divide up the cube. Dimensions are specified in y,x axis order. Each integer describes the number of blocks that dimension will be divided into. - explicit_src_chunks : arraylike NxMx2 + explicit_src_blocks : arraylike NxMx2 Explicitly specify the bounds of each block in the partition. # tgt_chunks : ???, default=None # ??? @@ -133,18 +133,26 @@ def __init__( # TODO: consider abstracting away the idea of files self.file_names = file_names if use_dask_src_chunks: - assert num_src_chunks is None and src_chunks is None - assert src.has_lazy_data() + if src_chunks is not None: + msg = "Potentially conflicting partition block definitions." + raise ValueError(msg) + if not src.has_lazy_data(): + msg = "If `use_dask_src_chunks=True`, the source cube must be lazy." + raise TypeError(msg) src_chunks = src.slices(grid_dims).next().lazy_data().chunks - self.src_chunks = _determine_blocks(shape, src_chunks, num_src_chunks, explicit_src_chunks) - assert len(self.src_chunks) == len(file_names) + self.src_blocks = _determine_blocks(shape, src_chunks, num_src_chunks, explicit_src_blocks) + if len(self.src_blocks) != len(file_names): + msg = "Number of source blocks does not match number of file names." + raise ValueError(msg) # This will be controllable in future tgt_chunks = None self.tgt_chunks = tgt_chunks - assert tgt_chunks is None # We don't handle big targets currently + if tgt_chunks is not None: + msg = "Target chunking not yet implemented." + raise NotImplementedError(msg) # Note: this may need to become more sophisticated when both src and tgt are large - self.file_chunk_dict = dict(zip(self.file_names, self.src_chunks)) + self.file_block_dict = dict(zip(self.file_names, self.src_blocks)) if saved_files is None: self.saved_files = [] @@ -174,32 +182,34 @@ def generate_files(self, files_to_generate=None): if files_to_generate is None: files = self.unsaved_files else: - assert isinstance(files_to_generate, int) + if not isinstance(files_to_generate, int): + msg = "`files_to_generate` must be an integer." + raise ValueError(msg) files = self.unsaved_files[:files_to_generate] for file in files: - src_chunk = self.file_chunk_dict[file] - src = _get_chunk(self.src, src_chunk) + src_block = self.file_block_dict[file] + src = _get_chunk(self.src, src_block) tgt = self.tgt regridder = self.scheme.regridder(src, tgt) weights = regridder.regridder.weight_matrix - regridder = PartialRegridder(src, self.tgt, src_chunk, None, weights, self.scheme) + regridder = PartialRegridder(src, self.tgt, src_block, None, weights, self.scheme) save_regridder(regridder, file, allow_partial=True) self.saved_files.append(file) def apply_regridders(self, cube, allow_incomplete=False): # for each target chunk, iterate through each associated regridder # for now, assume one target chunk - if not allow_incomplete: - assert len(self.unsaved_files) == 0 + if not allow_incomplete and len(self.unsaved_files) != 0: + msg = "Not all files have been constructed." + raise OSError(msg) current_result = None current_weights = None files = self.saved_files - for file, chunk in zip(self.file_names, self.src_chunks): + for file, chunk in zip(self.file_names, self.src_blocks): if file in files: next_regridder = load_regridder(file, allow_partial=True) - # cube_chunk = cube[*_interpret_slice(chunk)] cube_chunk = _get_chunk(cube, chunk) next_weights, next_result = next_regridder.partial_regrid(cube_chunk) if current_weights is None: diff --git a/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py b/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py index fe9d708e..4010613a 100644 --- a/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py +++ b/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py @@ -23,7 +23,7 @@ def test_Partition(tmp_path): scheme = ESMFAreaWeighted(mdtol=1) chunks = [[[0, 100], [0, 150]], [[100, 200], [0, 150]], [[200, 300], [0, 150]], [[300, 400], [0, 150]], [[400, 500], [0, 150]]] - partition = Partition(src, tgt, scheme, files, explicit_src_chunks=chunks) + partition = Partition(src, tgt, scheme, files, explicit_src_blocks=chunks) partition.generate_files() @@ -43,26 +43,26 @@ def test_Partition_block_api(tmp_path): partition = Partition(src, tgt, scheme, files, num_src_chunks=num_src_chunks) expected_chunks = [[[0, 100], [0, 150]], [[100, 200], [0, 150]], [[200, 300], [0, 150]], [[300, 400], [0, 150]], [[400, 500], [0, 150]]] - assert partition.src_chunks == expected_chunks + assert partition.src_blocks == expected_chunks src_chunks = (100, 150) partition = Partition(src, tgt, scheme, files, src_chunks=src_chunks) expected_chunks = [[[0, 100], [0, 150]], [[100, 200], [0, 150]], [[200, 300], [0, 150]], [[300, 400], [0, 150]], [[400, 500], [0, 150]]] - assert partition.src_chunks == expected_chunks + assert partition.src_blocks == expected_chunks src_chunks = ((100, 100, 100, 100, 100), (150,)) partition = Partition(src, tgt, scheme, files, src_chunks=src_chunks) expected_chunks = [[[0, 100], [0, 150]], [[100, 200], [0, 150]], [[200, 300], [0, 150]], [[300, 400], [0, 150]], [[400, 500], [0, 150]]] - assert partition.src_chunks == expected_chunks + assert partition.src_blocks == expected_chunks src.data = da.from_array(src.data, chunks=src_chunks) partition = Partition(src, tgt, scheme, files, use_dask_src_chunks=True) expected_chunks = [[[0, 100], [0, 150]], [[100, 200], [0, 150]], [[200, 300], [0, 150]], [[300, 400], [0, 150]], [[400, 500], [0, 150]]] - assert partition.src_chunks == expected_chunks + assert partition.src_blocks == expected_chunks def test_Partition_mesh_src(tmp_path): src = _gridlike_mesh_cube(150, 500) @@ -74,11 +74,11 @@ def test_Partition_mesh_src(tmp_path): src_chunks = (15000,) with pytest.raises(NotImplementedError): - partition = Partition(src, tgt, scheme, files, src_chunks=src_chunks) + _ = Partition(src, tgt, scheme, files, src_chunks=src_chunks) - # TODO when mesh partitioning becomes possible, uncomment. + # TODO: when mesh partitioning becomes possible, uncomment. # expected_src_chunks = [[[0, 15000]], [[15000, 30000]], [[30000, 45000]], [[45000, 60000]], [[60000, 75000]]] - # assert partition.src_chunks == expected_src_chunks + # assert partition.src_blocks == expected_src_chunks # # partition.generate_files() # @@ -99,7 +99,7 @@ def test_Partition_curv_src(tmp_path): partition = Partition(src, tgt, scheme, files, src_chunks=src_chunks) expected_src_chunks = [[[0, 100], [0, 150]], [[100, 200], [0, 150]], [[200, 300], [0, 150]], [[300, 400], [0, 150]], [[400, 500], [0, 150]]] - assert partition.src_chunks == expected_src_chunks + assert partition.src_blocks == expected_src_chunks partition.generate_files() From e62585635a0fdc5768719d164a2006bb75c2ca80 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Thu, 18 Dec 2025 13:44:56 +0000 Subject: [PATCH 05/41] lint fixes --- src/esmf_regrid/esmf_regridder.py | 4 +- src/esmf_regrid/experimental/_partial.py | 7 +- src/esmf_regrid/experimental/io.py | 25 ++++--- src/esmf_regrid/experimental/partition.py | 75 ++++++++++++------- .../io/partition/test_Partition.py | 61 ++++++++++++--- 5 files changed, 118 insertions(+), 54 deletions(-) diff --git a/src/esmf_regrid/esmf_regridder.py b/src/esmf_regrid/esmf_regridder.py index 73bc13a0..0225cd3c 100644 --- a/src/esmf_regrid/esmf_regridder.py +++ b/src/esmf_regrid/esmf_regridder.py @@ -188,7 +188,9 @@ def _gen_weights_and_data(self, src_array): tgt_weights = self.tgt._matrix_to_array(weight_sums, extra_shape) return tgt_weights, tgt_data - def _regrid_from_weights_and_data(self, tgt_weights, tgt_data, norm_type=Constants.NormType.FRACAREA, mdtol=1): + def _regrid_from_weights_and_data( + self, tgt_weights, tgt_data, 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) diff --git a/src/esmf_regrid/experimental/_partial.py b/src/esmf_regrid/experimental/_partial.py index c8067158..f761c88f 100644 --- a/src/esmf_regrid/experimental/_partial.py +++ b/src/esmf_regrid/experimental/_partial.py @@ -5,6 +5,7 @@ _create_cube, ) + class PartialRegridder(_ESMFRegridder): def __init__(self, src, tgt, src_slice, tgt_slice, weights, scheme, **kwargs): self.src_slice = src_slice # this will be tuple-like @@ -32,10 +33,6 @@ def finish_regridding(self, src_cube, weights, data): result_data = self.regridder._regrid_from_weights_and_data(weights, data) result_cube = _create_cube( - result_data, - src_cube, - dims, - self._tgt, - len(self._tgt) + result_data, src_cube, dims, self._tgt, len(self._tgt) ) return result_cube diff --git a/src/esmf_regrid/experimental/io.py b/src/esmf_regrid/experimental/io.py index e7161aa7..885e9364 100644 --- a/src/esmf_regrid/experimental/io.py +++ b/src/esmf_regrid/experimental/io.py @@ -140,6 +140,7 @@ def _standard_grid_cube(grid, name): cube.add_aux_coord(grid[1], [0, 1]) return cube + def _standard_mesh_cube(mesh, location, name): mesh_coords = mesh.to_MeshCoords(location) data = np.zeros(mesh_coords[0].points.shape[0]) @@ -148,6 +149,7 @@ def _standard_mesh_cube(mesh, location, name): cube.add_aux_coord(coord, 0) return cube + def _generate_src_tgt(regridder_type, rg, allow_partial): if regridder_type in [ "ESMFAreaWeightedRegridder", @@ -204,15 +206,11 @@ def _generate_src_tgt(regridder_type, rg, allow_partial): _add_mask_to_cube(rg.tgt_mask, tgt_cube, _TARGET_MASK_NAME) else: - e_msg = ( - f"Unexpected regridder type {regridder_type}." - ) + e_msg = f"Unexpected regridder type {regridder_type}." raise TypeError(e_msg) return src_cube, tgt_cube - - def save_regridder(rg, filename, allow_partial=False): """Save a regridder scheme instance. @@ -287,14 +285,17 @@ def save_regridder(rg, filename, allow_partial=False): src_slice = rg.src_slice # this slice is described by a tuple if src_slice is None: src_slice = [] - src_slice_cube = Cube(src_slice, long_name=_SRC_SLICE_NAME, var_name=_SRC_SLICE_NAME) + src_slice_cube = Cube( + src_slice, long_name=_SRC_SLICE_NAME, var_name=_SRC_SLICE_NAME + ) tgt_slice = rg.tgt_slice # this slice is described by a tuple if tgt_slice is None: tgt_slice = [] - tgt_slice_cube = Cube(src_slice, long_name=_TGT_SLICE_NAME, var_name=_TGT_SLICE_NAME) + tgt_slice_cube = Cube( + src_slice, long_name=_TGT_SLICE_NAME, var_name=_TGT_SLICE_NAME + ) extra_cubes = [src_slice_cube, tgt_slice_cube] - weights_cube = Cube(weight_data, var_name=_WEIGHTS_NAME, long_name=_WEIGHTS_NAME) row_coord = AuxCoord( weight_rows, var_name=_WEIGHTS_ROW_NAME, long_name=_WEIGHTS_ROW_NAME @@ -329,7 +330,9 @@ def save_regridder(rg, filename, allow_partial=False): # Save cubes while ensuring var_names do not conflict for the sake of consistency. with _managed_var_name(src_cube, tgt_cube): - cube_list = CubeList([src_cube, tgt_cube, weights_cube, weight_shape_cube, *extra_cubes]) + cube_list = CubeList( + [src_cube, tgt_cube, weights_cube, weight_shape_cube, *extra_cubes] + ) for cube in cube_list: cube.attributes = attributes @@ -375,7 +378,9 @@ def load_regridder(filename, allow_partial=False): scheme = _REGRIDDER_NAME_MAP[regridder_type] if regridder_type == "PartialRegridder" and not allow_partial: - e_msg = "PartialRegridder cannot be loaded without setting `allow_partial=True`." + e_msg = ( + "PartialRegridder cannot be loaded without setting `allow_partial=True`." + ) raise ValueError(e_msg) # Determine the regridding method, allowing for files created when diff --git a/src/esmf_regrid/experimental/partition.py b/src/esmf_regrid/experimental/partition.py index 3bf3ae1b..a34e92dd 100644 --- a/src/esmf_regrid/experimental/partition.py +++ b/src/esmf_regrid/experimental/partition.py @@ -15,11 +15,16 @@ def _get_chunk(cube, sl): grid_dims = (cube.mesh_dim(),) full_slice = [np.s_[:]] * len(cube.shape) for s, d in zip(sl, grid_dims): - full_slice[d] = np.s_[s[0]:s[1]] + full_slice[d] = np.s_[s[0] : s[1]] return cube[*full_slice] + def _determine_blocks(shape, chunks, num_chunks, explicit_chunks): - which_inputs = (chunks is not None, num_chunks is not None, explicit_chunks is not None) + which_inputs = ( + chunks is not None, + num_chunks is not None, + explicit_chunks is not None, + ) if sum(which_inputs) == 0: msg = "Partition blocks must must be specified by either chunks, num_chunks, or explicit_chunks." raise ValueError(msg) @@ -27,18 +32,18 @@ def _determine_blocks(shape, chunks, num_chunks, explicit_chunks): msg = "Potentially conflicting partition block definitions." raise ValueError(msg) if num_chunks is not None: - chunks = [s//n for s, n in zip(shape, num_chunks)] + chunks = [s // n for s, n in zip(shape, num_chunks)] for chunk in chunks: if chunk == 0: msg = "`num_chunks` cannot divide a dimension into more blocks than the size of that dimension." raise ValueError(msg) if chunks is not None: - if all(isinstance(x, int)for x in chunks): + if all(isinstance(x, int) for x in chunks): proper_chunks = [] for s, c in zip(shape, chunks): - proper_chunk = [c] * (s//c) - if s%c != 0: - proper_chunk += [s%c] + proper_chunk = [c] * (s // c) + if s % c != 0: + proper_chunk += [s % c] proper_chunks.append(proper_chunk) chunks = proper_chunks for s, chunk in zip(shape, chunks): @@ -47,32 +52,39 @@ def _determine_blocks(shape, chunks, num_chunks, explicit_chunks): raise ValueError(msg) bounds = [np.cumsum([0, *chunk]) for chunk in chunks] if len(bounds) == 1: - explicit_chunks = [[[int(lower), int(upper)]] for lower, upper in zip(bounds[0][:-1], bounds[0][1:])] + explicit_chunks = [ + [[int(lower), int(upper)]] + for lower, upper in zip(bounds[0][:-1], bounds[0][1:]) + ] elif len(bounds) == 2: explicit_chunks = [ - [[int(ly), int(uy)], [int(lx), int(ux)]] for ly, uy in zip(bounds[0][:-1], bounds[0][1:]) for lx, ux in zip(bounds[1][:-1], bounds[1][1:]) + [[int(ly), int(uy)], [int(lx), int(ux)]] + for ly, uy in zip(bounds[0][:-1], bounds[0][1:]) + for lx, ux in zip(bounds[1][:-1], bounds[1][1:]) ] else: msg = "Chunks must not exceed two dimensions." raise ValueError(msg) return explicit_chunks + class Partition: """Class for breaking down regridding into manageable chunks.""" + def __init__( - self, - src, - tgt, - scheme, - file_names, - use_dask_src_chunks=False, - src_chunks=None, - num_src_chunks=None, - explicit_src_blocks=None, - # tgt_chunks=None, - # num_tgt_chunks=None, - auto_generate=False, - saved_files=None, + self, + src, + tgt, + scheme, + file_names, + use_dask_src_chunks=False, + src_chunks=None, + num_src_chunks=None, + explicit_src_blocks=None, + # tgt_chunks=None, + # num_tgt_chunks=None, + auto_generate=False, + saved_files=None, ): """Class for breaking down regridding into manageable chunks. @@ -87,8 +99,8 @@ def __init__( Source cube. tgt : cube Target cube. - scheme : - scheme + scheme : regridding scheme + Regridding scheme to generate regridders, either ESMFAreaWeighted or ESMFBilinear. file_names : iterable of str A list of file names to save/load parts of the regridder to/from. use_dask_src_chunks : bool, default=False @@ -140,7 +152,9 @@ def __init__( msg = "If `use_dask_src_chunks=True`, the source cube must be lazy." raise TypeError(msg) src_chunks = src.slices(grid_dims).next().lazy_data().chunks - self.src_blocks = _determine_blocks(shape, src_chunks, num_src_chunks, explicit_src_blocks) + self.src_blocks = _determine_blocks( + shape, src_chunks, num_src_chunks, explicit_src_blocks + ) if len(self.src_blocks) != len(file_names): msg = "Number of source blocks does not match number of file names." raise ValueError(msg) @@ -175,10 +189,12 @@ def __repr__(self): @property def unsaved_files(self): + """List of files not yet generated.""" files = set(self.file_names) - set(self.saved_files) return [file for file in self.file_names if file in files] def generate_files(self, files_to_generate=None): + """Generate files with regridding information.""" if files_to_generate is None: files = self.unsaved_files else: @@ -193,11 +209,14 @@ def generate_files(self, files_to_generate=None): tgt = self.tgt regridder = self.scheme.regridder(src, tgt) weights = regridder.regridder.weight_matrix - regridder = PartialRegridder(src, self.tgt, src_block, None, weights, self.scheme) + regridder = PartialRegridder( + src, self.tgt, src_block, None, weights, self.scheme + ) save_regridder(regridder, file, allow_partial=True) self.saved_files.append(file) def apply_regridders(self, cube, allow_incomplete=False): + """Apply the saved regridders to a cube.""" # for each target chunk, iterate through each associated regridder # for now, assume one target chunk if not allow_incomplete and len(self.unsaved_files) != 0: @@ -221,4 +240,6 @@ def apply_regridders(self, cube, allow_incomplete=False): else: current_result += next_result - return next_regridder.finish_regridding(cube_chunk, current_weights, current_result) + return next_regridder.finish_regridding( + cube_chunk, current_weights, current_result + ) diff --git a/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py b/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py index 4010613a..ba04bb83 100644 --- a/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py +++ b/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py @@ -14,14 +14,21 @@ _gridlike_mesh_cube, ) + def test_Partition(tmp_path): src = _grid_cube(150, 500, (-180, 180), (-90, 90), circular=True) - src.data = np.arange(150*500).reshape([500, 150]) + src.data = np.arange(150 * 500).reshape([500, 150]) tgt = _grid_cube(16, 36, (-180, 180), (-90, 90), circular=True) files = [tmp_path / f"partial_{x}.nc" for x in range(5)] scheme = ESMFAreaWeighted(mdtol=1) - chunks = [[[0, 100], [0, 150]], [[100, 200], [0, 150]], [[200, 300], [0, 150]], [[300, 400], [0, 150]], [[400, 500], [0, 150]]] + chunks = [ + [[0, 100], [0, 150]], + [[100, 200], [0, 150]], + [[200, 300], [0, 150]], + [[300, 400], [0, 150]], + [[400, 500], [0, 150]], + ] partition = Partition(src, tgt, scheme, files, explicit_src_blocks=chunks) @@ -32,9 +39,10 @@ def test_Partition(tmp_path): assert np.allclose(result.data, expected.data) assert result == expected + def test_Partition_block_api(tmp_path): src = _grid_cube(150, 500, (-180, 180), (-90, 90), circular=True) - src.data = np.arange(150*500).reshape([500, 150]) + src.data = np.arange(150 * 500).reshape([500, 150]) tgt = _grid_cube(16, 36, (-180, 180), (-90, 90), circular=True) files = [tmp_path / f"partial_{x}.nc" for x in range(5)] @@ -42,31 +50,55 @@ def test_Partition_block_api(tmp_path): num_src_chunks = (5, 1) partition = Partition(src, tgt, scheme, files, num_src_chunks=num_src_chunks) - expected_chunks = [[[0, 100], [0, 150]], [[100, 200], [0, 150]], [[200, 300], [0, 150]], [[300, 400], [0, 150]], [[400, 500], [0, 150]]] + expected_chunks = [ + [[0, 100], [0, 150]], + [[100, 200], [0, 150]], + [[200, 300], [0, 150]], + [[300, 400], [0, 150]], + [[400, 500], [0, 150]], + ] assert partition.src_blocks == expected_chunks src_chunks = (100, 150) partition = Partition(src, tgt, scheme, files, src_chunks=src_chunks) - expected_chunks = [[[0, 100], [0, 150]], [[100, 200], [0, 150]], [[200, 300], [0, 150]], [[300, 400], [0, 150]], [[400, 500], [0, 150]]] + expected_chunks = [ + [[0, 100], [0, 150]], + [[100, 200], [0, 150]], + [[200, 300], [0, 150]], + [[300, 400], [0, 150]], + [[400, 500], [0, 150]], + ] assert partition.src_blocks == expected_chunks - src_chunks = ((100, 100, 100, 100, 100), (150,)) partition = Partition(src, tgt, scheme, files, src_chunks=src_chunks) - expected_chunks = [[[0, 100], [0, 150]], [[100, 200], [0, 150]], [[200, 300], [0, 150]], [[300, 400], [0, 150]], [[400, 500], [0, 150]]] + expected_chunks = [ + [[0, 100], [0, 150]], + [[100, 200], [0, 150]], + [[200, 300], [0, 150]], + [[300, 400], [0, 150]], + [[400, 500], [0, 150]], + ] assert partition.src_blocks == expected_chunks src.data = da.from_array(src.data, chunks=src_chunks) partition = Partition(src, tgt, scheme, files, use_dask_src_chunks=True) - expected_chunks = [[[0, 100], [0, 150]], [[100, 200], [0, 150]], [[200, 300], [0, 150]], [[300, 400], [0, 150]], [[400, 500], [0, 150]]] + expected_chunks = [ + [[0, 100], [0, 150]], + [[100, 200], [0, 150]], + [[200, 300], [0, 150]], + [[300, 400], [0, 150]], + [[400, 500], [0, 150]], + ] assert partition.src_blocks == expected_chunks + def test_Partition_mesh_src(tmp_path): src = _gridlike_mesh_cube(150, 500) - src.data = np.arange(150*500) + src.data = np.arange(150 * 500) tgt = _grid_cube(16, 36, (-180, 180), (-90, 90), circular=True) files = [tmp_path / f"partial_{x}.nc" for x in range(5)] @@ -87,9 +119,10 @@ def test_Partition_mesh_src(tmp_path): # assert np.allclose(result.data, expected.data) # assert result == expected + def test_Partition_curv_src(tmp_path): src = _curvilinear_cube(150, 500, (-180, 180), (-90, 90)) - src.data = np.arange(150*500).reshape([500, 150]) + src.data = np.arange(150 * 500).reshape([500, 150]) tgt = _grid_cube(16, 36, (-180, 180), (-90, 90), circular=True) files = [tmp_path / f"partial_{x}.nc" for x in range(5)] @@ -98,7 +131,13 @@ def test_Partition_curv_src(tmp_path): src_chunks = (100, 150) partition = Partition(src, tgt, scheme, files, src_chunks=src_chunks) - expected_src_chunks = [[[0, 100], [0, 150]], [[100, 200], [0, 150]], [[200, 300], [0, 150]], [[300, 400], [0, 150]], [[400, 500], [0, 150]]] + expected_src_chunks = [ + [[0, 100], [0, 150]], + [[100, 200], [0, 150]], + [[200, 300], [0, 150]], + [[300, 400], [0, 150]], + [[400, 500], [0, 150]], + ] assert partition.src_blocks == expected_src_chunks partition.generate_files() From 60fc6758a45686b47be548eaad4a4c82c1d0592b Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Thu, 18 Dec 2025 14:01:29 +0000 Subject: [PATCH 06/41] fix import order --- src/esmf_regrid/experimental/_partial.py | 2 +- src/esmf_regrid/experimental/io.py | 2 +- src/esmf_regrid/experimental/partition.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/esmf_regrid/experimental/_partial.py b/src/esmf_regrid/experimental/_partial.py index f761c88f..3ba26673 100644 --- a/src/esmf_regrid/experimental/_partial.py +++ b/src/esmf_regrid/experimental/_partial.py @@ -1,8 +1,8 @@ """Provides a regridder class compatible with Partition.""" from esmf_regrid.schemes import ( - _ESMFRegridder, _create_cube, + _ESMFRegridder, ) diff --git a/src/esmf_regrid/experimental/io.py b/src/esmf_regrid/experimental/io.py index 885e9364..49dd52ad 100644 --- a/src/esmf_regrid/experimental/io.py +++ b/src/esmf_regrid/experimental/io.py @@ -10,11 +10,11 @@ import esmf_regrid from esmf_regrid import Constants, _load_context, check_method, esmpy +from esmf_regrid.experimental._partial import PartialRegridder from esmf_regrid.experimental.unstructured_scheme import ( GridToMeshESMFRegridder, MeshToGridESMFRegridder, ) -from esmf_regrid.experimental._partial import PartialRegridder from esmf_regrid.schemes import ( ESMFAreaWeightedRegridder, ESMFBilinearRegridder, diff --git a/src/esmf_regrid/experimental/partition.py b/src/esmf_regrid/experimental/partition.py index a34e92dd..faba245a 100644 --- a/src/esmf_regrid/experimental/partition.py +++ b/src/esmf_regrid/experimental/partition.py @@ -3,8 +3,8 @@ import numpy as np from esmf_regrid.constants import Constants -from esmf_regrid.experimental.io import load_regridder, save_regridder from esmf_regrid.experimental._partial import PartialRegridder +from esmf_regrid.experimental.io import load_regridder, save_regridder from esmf_regrid.schemes import _get_grid_dims From decc42ff50ca3d635e35864ff85da970fb6c1457 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Thu, 18 Dec 2025 14:04:24 +0000 Subject: [PATCH 07/41] fix import order --- src/esmf_regrid/experimental/io.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/esmf_regrid/experimental/io.py b/src/esmf_regrid/experimental/io.py index 49dd52ad..3a70c861 100644 --- a/src/esmf_regrid/experimental/io.py +++ b/src/esmf_regrid/experimental/io.py @@ -16,12 +16,12 @@ MeshToGridESMFRegridder, ) from esmf_regrid.schemes import ( - ESMFAreaWeightedRegridder, - ESMFBilinearRegridder, - ESMFNearestRegridder, ESMFAreaWeighted, + ESMFAreaWeightedRegridder, ESMFBilinear, + ESMFBilinearRegridder, ESMFNearest, + ESMFNearestRegridder, GridRecord, MeshRecord, ) From 237cef8a4436c4f8b46c1e72a595bd039588fc01 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Thu, 18 Dec 2025 15:49:25 +0000 Subject: [PATCH 08/41] add tests, docstrings --- .../io/partition/test_Partition.py | 43 +++++++++++++++++-- 1 file changed, 40 insertions(+), 3 deletions(-) diff --git a/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py b/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py index ba04bb83..a52d83b2 100644 --- a/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py +++ b/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py @@ -16,13 +16,14 @@ def test_Partition(tmp_path): + """Test basic implementation of Partition class.""" src = _grid_cube(150, 500, (-180, 180), (-90, 90), circular=True) src.data = np.arange(150 * 500).reshape([500, 150]) tgt = _grid_cube(16, 36, (-180, 180), (-90, 90), circular=True) files = [tmp_path / f"partial_{x}.nc" for x in range(5)] scheme = ESMFAreaWeighted(mdtol=1) - chunks = [ + blocks = [ [[0, 100], [0, 150]], [[100, 200], [0, 150]], [[200, 300], [0, 150]], @@ -30,7 +31,7 @@ def test_Partition(tmp_path): [[400, 500], [0, 150]], ] - partition = Partition(src, tgt, scheme, files, explicit_src_blocks=chunks) + partition = Partition(src, tgt, scheme, files, explicit_src_blocks=blocks) partition.generate_files() @@ -41,8 +42,8 @@ def test_Partition(tmp_path): def test_Partition_block_api(tmp_path): + """Test API for controlling block shape for Partition class.""" src = _grid_cube(150, 500, (-180, 180), (-90, 90), circular=True) - src.data = np.arange(150 * 500).reshape([500, 150]) tgt = _grid_cube(16, 36, (-180, 180), (-90, 90), circular=True) files = [tmp_path / f"partial_{x}.nc" for x in range(5)] @@ -97,6 +98,7 @@ def test_Partition_block_api(tmp_path): def test_Partition_mesh_src(tmp_path): + """Test Partition class when the source has a mesh.""" src = _gridlike_mesh_cube(150, 500) src.data = np.arange(150 * 500) tgt = _grid_cube(16, 36, (-180, 180), (-90, 90), circular=True) @@ -121,6 +123,7 @@ def test_Partition_mesh_src(tmp_path): def test_Partition_curv_src(tmp_path): + """Test Partition class when the source has a curvilinear grid.""" src = _curvilinear_cube(150, 500, (-180, 180), (-90, 90)) src.data = np.arange(150 * 500).reshape([500, 150]) tgt = _grid_cube(16, 36, (-180, 180), (-90, 90), circular=True) @@ -146,3 +149,37 @@ def test_Partition_curv_src(tmp_path): expected = src.regrid(tgt, scheme) assert np.allclose(result.data, expected.data) assert result == expected + + +def test_conflicting_chunks(tmp_path): + """Test error handling of Partition class.""" + src = _grid_cube(150, 500, (-180, 180), (-90, 90), circular=True) + tgt = _grid_cube(16, 36, (-180, 180), (-90, 90), circular=True) + + files = [tmp_path / f"partial_{x}.nc" for x in range(5)] + scheme = ESMFAreaWeighted(mdtol=1) + num_src_chunks = (5, 1) + src_chunks = (100, 150) + blocks = [ + [[0, 100], [0, 150]], + [[100, 200], [0, 150]], + [[200, 300], [0, 150]], + [[300, 400], [0, 150]], + [[400, 500], [0, 150]], + ] + + with pytest.raises(ValueError): + _ = Partition(src, tgt, scheme, files, src_chunks=src_chunks, num_src_chunks=num_src_chunks) + with pytest.raises(ValueError): + _ = Partition(src, tgt, scheme, files, src_chunks=src_chunks, explicit_src_blocks=blocks) + with pytest.raises(ValueError): + _ = Partition(src, tgt, scheme, files) + with pytest.raises(TypeError): + _ = Partition(src, tgt, scheme, files, use_dask_src_chunks=True) + with pytest.raises(ValueError): + _ = Partition(src, tgt, scheme, files[:-1], src_chunks=src_chunks) + +def test_multidimensional_cube(tmp_path): + """Test Partition class when the source has a multidimensional cube.""" + #TODO: add test. + pass From 7a8affae41028fec0590c2f52807cad9b8bce490 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Fri, 19 Dec 2025 09:27:26 +0000 Subject: [PATCH 09/41] add multidimensional test --- .../io/partition/test_Partition.py | 16 +- .../test_regrid_rectilinear_to_rectilinear.py | 303 ++++-------------- 2 files changed, 75 insertions(+), 244 deletions(-) diff --git a/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py b/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py index a52d83b2..1477485c 100644 --- a/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py +++ b/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py @@ -13,6 +13,9 @@ from esmf_regrid.tests.unit.schemes.test__mesh_to_MeshInfo import ( _gridlike_mesh_cube, ) +from esmf_regrid.tests.unit.schemes.test_regrid_rectilinear_to_rectilinear import ( + _make_full_cubes, +) def test_Partition(tmp_path): @@ -181,5 +184,14 @@ def test_conflicting_chunks(tmp_path): def test_multidimensional_cube(tmp_path): """Test Partition class when the source has a multidimensional cube.""" - #TODO: add test. - pass + src_cube, tgt_grid, expected_cube = _make_full_cubes() + files = [tmp_path / f"partial_{x}.nc" for x in range(4)] + scheme = ESMFAreaWeighted(mdtol=1) + chunks = (2,3) + + partition = Partition(src_cube, tgt_grid, scheme, files, src_chunks=chunks) + + partition.generate_files() + + result = partition.apply_regridders(src_cube) + assert result == expected_cube diff --git a/src/esmf_regrid/tests/unit/schemes/test_regrid_rectilinear_to_rectilinear.py b/src/esmf_regrid/tests/unit/schemes/test_regrid_rectilinear_to_rectilinear.py index 548cc1f4..05e40350 100644 --- a/src/esmf_regrid/tests/unit/schemes/test_regrid_rectilinear_to_rectilinear.py +++ b/src/esmf_regrid/tests/unit/schemes/test_regrid_rectilinear_to_rectilinear.py @@ -68,12 +68,18 @@ def test_rotated_regridding(): assert np.allclose(expected_data, full_mdtol_result.data) -def test_extra_dims(): - """Test for :func:`esmf_regrid.schemes.regrid_rectilinear_to_rectilinear`. - - Tests the handling of extra dimensions and metadata. Ensures that proper - coordinates, attributes, names and units are copied over. - """ +def _add_metadata(cube): + result = cube.copy() + result.units = "K" + result.attributes = {"a": 1} + result.standard_name = "air_temperature" + scalar_height = AuxCoord([5], units="m", standard_name="height") + scalar_time = DimCoord([10], units="s", standard_name="time") + result.add_aux_coord(scalar_height) + result.add_aux_coord(scalar_time) + return result + +def _make_full_cubes(src_rectilinear=True, tgt_rectilinear=True): h = 2 t = 4 e = 6 @@ -86,13 +92,22 @@ def test_extra_dims(): lon_bounds = (-180, 180) lat_bounds = (-90, 90) - src_grid = _grid_cube( + if src_rectilinear: + src_func = _grid_cube + else: + src_func = _curvilinear_cube + if tgt_rectilinear: + tgt_func = _grid_cube + else: + tgt_func = _curvilinear_cube + src_grid = src_func( src_lons, src_lats, lon_bounds, lat_bounds, ) - tgt_grid = _grid_cube( + + tgt_grid = tgt_func( tgt_lons, tgt_lats, lon_bounds, @@ -110,47 +125,55 @@ def test_extra_dims(): ] src_cube = Cube(src_data) + if src_rectilinear: + src_cube.add_dim_coord(src_grid.coord("latitude"), 1) + src_cube.add_dim_coord(src_grid.coord("longitude"), 3) + else: + src_cube.add_aux_coord(src_grid.coord("latitude"), (1, 3)) + src_cube.add_aux_coord(src_grid.coord("longitude"), (1, 3)) src_cube.add_dim_coord(height, 0) - src_cube.add_dim_coord(src_grid.coord("latitude"), 1) src_cube.add_dim_coord(time, 2) - src_cube.add_dim_coord(src_grid.coord("longitude"), 3) src_cube.add_aux_coord(extra, 4) src_cube.add_aux_coord(spanning, [0, 2, 4]) - def _add_metadata(cube): - result = cube.copy() - result.units = "K" - result.attributes = {"a": 1} - result.standard_name = "air_temperature" - scalar_height = AuxCoord([5], units="m", standard_name="height") - scalar_time = DimCoord([10], units="s", standard_name="time") - result.add_aux_coord(scalar_height) - result.add_aux_coord(scalar_time) - return result - src_cube = _add_metadata(src_cube) - result = regrid_rectilinear_to_rectilinear(src_cube, tgt_grid) - expected_data = np.empty([h, tgt_lats, t, tgt_lons, e]) expected_data[:] = np.arange(t * h * e).reshape([h, t, e])[ - :, np.newaxis, :, np.newaxis, : - ] + :, np.newaxis, :, np.newaxis, : + ] + expected_cube = Cube(expected_data) + if tgt_rectilinear: + expected_cube.add_dim_coord(tgt_grid.coord("latitude"), 1) + expected_cube.add_dim_coord(tgt_grid.coord("longitude"), 3) + else: + expected_cube.add_aux_coord(tgt_grid.coord("latitude"), (1, 3)) + expected_cube.add_aux_coord(tgt_grid.coord("longitude"), (1, 3)) expected_cube.add_dim_coord(height, 0) - expected_cube.add_dim_coord(tgt_grid.coord("latitude"), 1) expected_cube.add_dim_coord(time, 2) - expected_cube.add_dim_coord(tgt_grid.coord("longitude"), 3) expected_cube.add_aux_coord(extra, 4) expected_cube.add_aux_coord(spanning, [0, 2, 4]) expected_cube = _add_metadata(expected_cube) + return src_cube, tgt_grid, expected_cube + + +def test_extra_dims(): + """Test for :func:`esmf_regrid.schemes.regrid_rectilinear_to_rectilinear`. + + Tests the handling of extra dimensions and metadata. Ensures that proper + coordinates, attributes, names and units are copied over. + """ + src_cube, tgt_grid, expected_cube = _make_full_cubes(src_rectilinear=True, tgt_rectilinear=True) + + result = regrid_rectilinear_to_rectilinear(src_cube, tgt_grid) # Lenient check for data. - assert np.allclose(expected_data, result.data) + assert np.allclose(expected_cube.data, result.data) # Check metadata and coords. - result.data = expected_data + result.data = expected_cube.data assert expected_cube == result @@ -266,83 +289,15 @@ def test_extra_dims_curvilinear(): Tests the handling of extra dimensions and metadata. Ensures that proper coordinates, attributes, names and units are copied over. """ - h = 2 - t = 4 - e = 6 - src_lats = 3 - src_lons = 5 - - tgt_lats = 5 - tgt_lons = 3 - - lon_bounds = (-180, 180) - lat_bounds = (-90, 90) - - src_grid = _curvilinear_cube( - src_lons, - src_lats, - lon_bounds, - lat_bounds, - ) - tgt_grid = _curvilinear_cube( - tgt_lons, - tgt_lats, - lon_bounds, - lat_bounds, - ) - - height = DimCoord(np.arange(h), standard_name="height") - time = DimCoord(np.arange(t), standard_name="time") - extra = AuxCoord(np.arange(e), long_name="extra dim") - spanning = AuxCoord(np.ones([h, t, e]), long_name="spanning dim") - - src_data = np.empty([h, src_lats, t, src_lons, e]) - src_data[:] = np.arange(t * h * e).reshape([h, t, e])[ - :, np.newaxis, :, np.newaxis, : - ] - - src_cube = Cube(src_data) - src_cube.add_dim_coord(height, 0) - src_cube.add_aux_coord(src_grid.coord("latitude"), (1, 3)) - src_cube.add_dim_coord(time, 2) - src_cube.add_aux_coord(src_grid.coord("longitude"), (1, 3)) - src_cube.add_aux_coord(extra, 4) - src_cube.add_aux_coord(spanning, [0, 2, 4]) - - def _add_metadata(cube): - result = cube.copy() - result.units = "K" - result.attributes = {"a": 1} - result.standard_name = "air_temperature" - scalar_height = AuxCoord([5], units="m", standard_name="height") - scalar_time = DimCoord([10], units="s", standard_name="time") - result.add_aux_coord(scalar_height) - result.add_aux_coord(scalar_time) - return result - - src_cube = _add_metadata(src_cube) + src_cube, tgt_grid, expected_cube = _make_full_cubes(src_rectilinear=False, tgt_rectilinear=False) result = regrid_rectilinear_to_rectilinear(src_cube, tgt_grid) - expected_data = np.empty([h, tgt_lats, t, tgt_lons, e]) - expected_data[:] = np.arange(t * h * e).reshape([h, t, e])[ - :, np.newaxis, :, np.newaxis, : - ] - - expected_cube = Cube(expected_data) - expected_cube.add_dim_coord(height, 0) - expected_cube.add_aux_coord(tgt_grid.coord("latitude"), (1, 3)) - expected_cube.add_dim_coord(time, 2) - expected_cube.add_aux_coord(tgt_grid.coord("longitude"), (1, 3)) - expected_cube.add_aux_coord(extra, 4) - expected_cube.add_aux_coord(spanning, [0, 2, 4]) - expected_cube = _add_metadata(expected_cube) - # Lenient check for data. - assert np.allclose(expected_data, result.data) + assert np.allclose(expected_cube.data, result.data) # Check metadata and coords. - result.data = expected_data + result.data = expected_cube.data assert expected_cube == result @@ -352,83 +307,15 @@ def test_extra_dims_curvilinear_to_rectilinear(): Tests the handling of extra dimensions and metadata. Ensures that proper coordinates, attributes, names and units are copied over. """ - h = 2 - t = 4 - e = 6 - src_lats = 3 - src_lons = 5 - - tgt_lats = 5 - tgt_lons = 3 - - lon_bounds = (-180, 180) - lat_bounds = (-90, 90) - - src_grid = _curvilinear_cube( - src_lons, - src_lats, - lon_bounds, - lat_bounds, - ) - tgt_grid = _grid_cube( - tgt_lons, - tgt_lats, - lon_bounds, - lat_bounds, - ) - - height = DimCoord(np.arange(h), standard_name="height") - time = DimCoord(np.arange(t), standard_name="time") - extra = AuxCoord(np.arange(e), long_name="extra dim") - spanning = AuxCoord(np.ones([h, t, e]), long_name="spanning dim") - - src_data = np.empty([h, src_lats, t, src_lons, e]) - src_data[:] = np.arange(t * h * e).reshape([h, t, e])[ - :, np.newaxis, :, np.newaxis, : - ] - - src_cube = Cube(src_data) - src_cube.add_dim_coord(height, 0) - src_cube.add_aux_coord(src_grid.coord("latitude"), (1, 3)) - src_cube.add_dim_coord(time, 2) - src_cube.add_aux_coord(src_grid.coord("longitude"), (1, 3)) - src_cube.add_aux_coord(extra, 4) - src_cube.add_aux_coord(spanning, [0, 2, 4]) - - def _add_metadata(cube): - result = cube.copy() - result.units = "K" - result.attributes = {"a": 1} - result.standard_name = "air_temperature" - scalar_height = AuxCoord([5], units="m", standard_name="height") - scalar_time = DimCoord([10], units="s", standard_name="time") - result.add_aux_coord(scalar_height) - result.add_aux_coord(scalar_time) - return result - - src_cube = _add_metadata(src_cube) + src_cube, tgt_grid, expected_cube = _make_full_cubes(src_rectilinear=False, tgt_rectilinear=True) result = regrid_rectilinear_to_rectilinear(src_cube, tgt_grid) - expected_data = np.empty([h, tgt_lats, t, tgt_lons, e]) - expected_data[:] = np.arange(t * h * e).reshape([h, t, e])[ - :, np.newaxis, :, np.newaxis, : - ] - - expected_cube = Cube(expected_data) - expected_cube.add_dim_coord(height, 0) - expected_cube.add_dim_coord(tgt_grid.coord("latitude"), 1) - expected_cube.add_dim_coord(time, 2) - expected_cube.add_dim_coord(tgt_grid.coord("longitude"), 3) - expected_cube.add_aux_coord(extra, 4) - expected_cube.add_aux_coord(spanning, [0, 2, 4]) - expected_cube = _add_metadata(expected_cube) - # Lenient check for data. - assert np.allclose(expected_data, result.data) + assert np.allclose(expected_cube.data, result.data) # Check metadata and coords. - result.data = expected_data + result.data = expected_cube.data assert expected_cube == result @@ -438,81 +325,13 @@ def test_extra_dims_rectilinear_to_curvilinear(): Tests the handling of extra dimensions and metadata. Ensures that proper coordinates, attributes, names and units are copied over. """ - h = 2 - t = 4 - e = 6 - src_lats = 3 - src_lons = 5 - - tgt_lats = 5 - tgt_lons = 3 - - lon_bounds = (-180, 180) - lat_bounds = (-90, 90) - - src_grid = _grid_cube( - src_lons, - src_lats, - lon_bounds, - lat_bounds, - ) - tgt_grid = _curvilinear_cube( - tgt_lons, - tgt_lats, - lon_bounds, - lat_bounds, - ) - - height = DimCoord(np.arange(h), standard_name="height") - time = DimCoord(np.arange(t), standard_name="time") - extra = AuxCoord(np.arange(e), long_name="extra dim") - spanning = AuxCoord(np.ones([h, t, e]), long_name="spanning dim") - - src_data = np.empty([h, src_lats, t, src_lons, e]) - src_data[:] = np.arange(t * h * e).reshape([h, t, e])[ - :, np.newaxis, :, np.newaxis, : - ] - - src_cube = Cube(src_data) - src_cube.add_dim_coord(height, 0) - src_cube.add_dim_coord(src_grid.coord("latitude"), 1) - src_cube.add_dim_coord(time, 2) - src_cube.add_dim_coord(src_grid.coord("longitude"), 3) - src_cube.add_aux_coord(extra, 4) - src_cube.add_aux_coord(spanning, [0, 2, 4]) - - def _add_metadata(cube): - result = cube.copy() - result.units = "K" - result.attributes = {"a": 1} - result.standard_name = "air_temperature" - scalar_height = AuxCoord([5], units="m", standard_name="height") - scalar_time = DimCoord([10], units="s", standard_name="time") - result.add_aux_coord(scalar_height) - result.add_aux_coord(scalar_time) - return result - - src_cube = _add_metadata(src_cube) + src_cube, tgt_grid, expected_cube = _make_full_cubes(src_rectilinear=True, tgt_rectilinear=False) result = regrid_rectilinear_to_rectilinear(src_cube, tgt_grid) - expected_data = np.empty([h, tgt_lats, t, tgt_lons, e]) - expected_data[:] = np.arange(t * h * e).reshape([h, t, e])[ - :, np.newaxis, :, np.newaxis, : - ] - - expected_cube = Cube(expected_data) - expected_cube.add_dim_coord(height, 0) - expected_cube.add_aux_coord(tgt_grid.coord("latitude"), (1, 3)) - expected_cube.add_dim_coord(time, 2) - expected_cube.add_aux_coord(tgt_grid.coord("longitude"), (1, 3)) - expected_cube.add_aux_coord(extra, 4) - expected_cube.add_aux_coord(spanning, [0, 2, 4]) - expected_cube = _add_metadata(expected_cube) - # Lenient check for data. - assert np.allclose(expected_data, result.data) + assert np.allclose(expected_cube.data, result.data) # Check metadata and coords. - result.data = expected_data + result.data = expected_cube.data assert expected_cube == result From 8abfed2e597fd7c6b72923bf86993d2c874db0c3 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Fri, 19 Dec 2025 18:12:38 +0000 Subject: [PATCH 10/41] add multidimensional cube support --- src/esmf_regrid/esmf_regridder.py | 28 ++----------------- src/esmf_regrid/experimental/_partial.py | 17 ++++++++++- .../io/partition/test_Partition.py | 4 +++ 3 files changed, 22 insertions(+), 27 deletions(-) diff --git a/src/esmf_regrid/esmf_regridder.py b/src/esmf_regrid/esmf_regridder.py index 0225cd3c..a53ae86a 100644 --- a/src/esmf_regrid/esmf_regridder.py +++ b/src/esmf_regrid/esmf_regridder.py @@ -243,30 +243,6 @@ 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 = self._gen_weights_and_data(src_array) + tgt_array = self._regrid_from_weights_and_data(tgt_weights, tgt_data, norm_type=norm_type, mdtol=mdtol) return tgt_array diff --git a/src/esmf_regrid/experimental/_partial.py b/src/esmf_regrid/experimental/_partial.py index 3ba26673..7ef54b23 100644 --- a/src/esmf_regrid/experimental/_partial.py +++ b/src/esmf_regrid/experimental/_partial.py @@ -1,5 +1,7 @@ """Provides a regridder class compatible with Partition.""" +import numpy as np + from esmf_regrid.schemes import ( _create_cube, _ESMFRegridder, @@ -26,7 +28,20 @@ def __init__(self, src, tgt, src_slice, tgt_slice, weights, scheme, **kwargs): self.__dict__.update(self._regridder.__dict__) def partial_regrid(self, src): - return self.regridder._gen_weights_and_data(src.data) + dims = self._get_cube_dims(src) + num_out_dims = self.regridder.tgt.dims + 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) + + standard_out_dims = [-1, -2][:num_out_dims] + if num_dims == 2 and num_out_dims == 1: + dims = [min(dims)] + if num_dims == 1 and num_out_dims == 2: + dims = [dims[0] + 1, dims[0]] + result = tuple(np.moveaxis(r, standard_out_dims, dims) for r in result) + return result def finish_regridding(self, src_cube, weights, data): dims = self._get_cube_dims(src_cube) diff --git a/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py b/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py index 1477485c..20a50cbb 100644 --- a/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py +++ b/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py @@ -194,4 +194,8 @@ def test_multidimensional_cube(tmp_path): partition.generate_files() result = partition.apply_regridders(src_cube) + assert np.allclose(result.data, expected_cube.data) + result.data = expected_cube.data + + # TODO: The resulting longitude coordinate has become circular. Investigate. assert result == expected_cube From 9b3fc781a1a3652db5c51bd58decbe6d3f2e4132 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Fri, 19 Dec 2025 18:41:52 +0000 Subject: [PATCH 11/41] fix test --- .../tests/unit/experimental/io/partition/test_Partition.py | 6 ++++-- .../unit/schemes/test_regrid_rectilinear_to_rectilinear.py | 6 ++++-- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py b/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py index 20a50cbb..defae2a0 100644 --- a/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py +++ b/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py @@ -194,8 +194,10 @@ def test_multidimensional_cube(tmp_path): partition.generate_files() result = partition.apply_regridders(src_cube) + + # Lenient check for data. assert np.allclose(result.data, expected_cube.data) - result.data = expected_cube.data - # TODO: The resulting longitude coordinate has become circular. Investigate. + # Check metadata and coords. + result.data = expected_cube.data assert result == expected_cube diff --git a/src/esmf_regrid/tests/unit/schemes/test_regrid_rectilinear_to_rectilinear.py b/src/esmf_regrid/tests/unit/schemes/test_regrid_rectilinear_to_rectilinear.py index 05e40350..8a037248 100644 --- a/src/esmf_regrid/tests/unit/schemes/test_regrid_rectilinear_to_rectilinear.py +++ b/src/esmf_regrid/tests/unit/schemes/test_regrid_rectilinear_to_rectilinear.py @@ -1,5 +1,7 @@ """Unit tests for :func:`esmf_regrid.schemes.regrid_rectilinear_to_rectilinear`.""" +from functools import partial + import dask.array as da from iris.coord_systems import RotatedGeogCS from iris.coords import AuxCoord, DimCoord @@ -93,11 +95,11 @@ def _make_full_cubes(src_rectilinear=True, tgt_rectilinear=True): lat_bounds = (-90, 90) if src_rectilinear: - src_func = _grid_cube + src_func = partial(_grid_cube, circular=True) else: src_func = _curvilinear_cube if tgt_rectilinear: - tgt_func = _grid_cube + src_func = partial(_grid_cube, circular=True) else: tgt_func = _curvilinear_cube src_grid = src_func( From 78e447975d6841ce715e8b70d6fa6f7f4d32e8fe Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Fri, 19 Dec 2025 19:14:27 +0000 Subject: [PATCH 12/41] fix test --- .../unit/schemes/test_regrid_rectilinear_to_rectilinear.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/esmf_regrid/tests/unit/schemes/test_regrid_rectilinear_to_rectilinear.py b/src/esmf_regrid/tests/unit/schemes/test_regrid_rectilinear_to_rectilinear.py index 8a037248..f47c4650 100644 --- a/src/esmf_regrid/tests/unit/schemes/test_regrid_rectilinear_to_rectilinear.py +++ b/src/esmf_regrid/tests/unit/schemes/test_regrid_rectilinear_to_rectilinear.py @@ -97,7 +97,7 @@ def _make_full_cubes(src_rectilinear=True, tgt_rectilinear=True): if src_rectilinear: src_func = partial(_grid_cube, circular=True) else: - src_func = _curvilinear_cube + tgt_func = _curvilinear_cube if tgt_rectilinear: src_func = partial(_grid_cube, circular=True) else: From 147abbf5da80c17e03bbd15ad9f8bd693de74286 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Fri, 19 Dec 2025 19:15:07 +0000 Subject: [PATCH 13/41] fix test --- .../unit/schemes/test_regrid_rectilinear_to_rectilinear.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/esmf_regrid/tests/unit/schemes/test_regrid_rectilinear_to_rectilinear.py b/src/esmf_regrid/tests/unit/schemes/test_regrid_rectilinear_to_rectilinear.py index f47c4650..296a3ad4 100644 --- a/src/esmf_regrid/tests/unit/schemes/test_regrid_rectilinear_to_rectilinear.py +++ b/src/esmf_regrid/tests/unit/schemes/test_regrid_rectilinear_to_rectilinear.py @@ -97,9 +97,9 @@ def _make_full_cubes(src_rectilinear=True, tgt_rectilinear=True): if src_rectilinear: src_func = partial(_grid_cube, circular=True) else: - tgt_func = _curvilinear_cube + src_func = _curvilinear_cube if tgt_rectilinear: - src_func = partial(_grid_cube, circular=True) + tgt_func = partial(_grid_cube, circular=True) else: tgt_func = _curvilinear_cube src_grid = src_func( From 36a9b0554881dac0007c96913ed54559840e8be6 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Fri, 19 Dec 2025 19:48:47 +0000 Subject: [PATCH 14/41] fix ESMFNearest behaviour --- src/esmf_regrid/esmf_regridder.py | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/src/esmf_regrid/esmf_regridder.py b/src/esmf_regrid/esmf_regridder.py index a53ae86a..d2cc65c3 100644 --- a/src/esmf_regrid/esmf_regridder.py +++ b/src/esmf_regrid/esmf_regridder.py @@ -178,11 +178,16 @@ def _out_dtype(self, in_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 = self.weight_matrix @ flat_src + flat_tgt = weight_matrix @ flat_src src_inverted_mask = self.src._array_to_matrix(~ma.getmaskarray(src_array)) - weight_sums = self.weight_matrix @ src_inverted_mask + weight_sums = weight_matrix @ src_inverted_mask tgt_data = self.tgt._matrix_to_array(flat_tgt, extra_shape) tgt_weights = self.tgt._matrix_to_array(weight_sums, extra_shape) @@ -195,12 +200,13 @@ def _regrid_from_weights_and_data( # errors. mdtol = max(mdtol, 1e-8) tgt_mask = tgt_weights > 1 - mdtol - masked_weight_sums = tgt_weights * tgt_mask normalisations = np.ones_like(tgt_data) - if norm_type == Constants.NormType.FRACAREA: - normalisations[tgt_mask] /= masked_weight_sums[tgt_mask] - elif norm_type == Constants.NormType.DSTAREA: - pass + 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 From 1dd92e66af17dbee62fddae073b004d05894b80d Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Wed, 14 Jan 2026 14:44:56 +0000 Subject: [PATCH 15/41] add test, improve docstrings --- src/esmf_regrid/experimental/partition.py | 32 +++++++++----- .../io/partition/test_Partition.py | 42 ++++++++++++++++++- 2 files changed, 63 insertions(+), 11 deletions(-) diff --git a/src/esmf_regrid/experimental/partition.py b/src/esmf_regrid/experimental/partition.py index faba245a..4a8a3d5d 100644 --- a/src/esmf_regrid/experimental/partition.py +++ b/src/esmf_regrid/experimental/partition.py @@ -105,11 +105,11 @@ def __init__( A list of file names to save/load parts of the regridder to/from. use_dask_src_chunks : bool, default=False If true, partition using the same chunks from the source cube. - src_chunks : numpy array, tuple of int or str, default=None + src_chunks : numpy array, tuple of int or tuple of tuple of int, default=None Specify the size of blocks to use to divide up the cube. Dimensions are specified in y,x axis order. If `src_chunks` is a tuple of int, each integer describes - the maximum size of a block in that dimension. If `src_chunks` is a tuple of int, - each tuple describes the size of each successive block in that dimension. These + the maximum size of a block in that dimension. If `src_chunks` is a tuple of tuples, + each sub-tuple describes the size of each successive block in that dimension. These block sizes should add up to the total size of that dimension or else an error is raised. num_src_chunks : tuple of int @@ -118,10 +118,6 @@ def __init__( be divided into. explicit_src_blocks : arraylike NxMx2 Explicitly specify the bounds of each block in the partition. - # tgt_chunks : ???, default=None - # ??? - # num_tgt_chunks : tuple of int - # ??? auto_generate : bool, default=False When true, start generating files on initialisation. saved_files : iterable of str @@ -194,7 +190,13 @@ def unsaved_files(self): return [file for file in self.file_names if file in files] def generate_files(self, files_to_generate=None): - """Generate files with regridding information.""" + """Generate files with regridding information. + + Parameters + ---------- + files_to_generate : int, default=None + Specify the number of files to generate, default behaviour is to generate all files. + """ if files_to_generate is None: files = self.unsaved_files else: @@ -216,11 +218,21 @@ def generate_files(self, files_to_generate=None): self.saved_files.append(file) def apply_regridders(self, cube, allow_incomplete=False): - """Apply the saved regridders to a cube.""" + """Apply the saved regridders to a cube. + + Parameters + ---------- + allow_incomplete : bool, default=False + If False, raise an error if not all files have been generated. If True, perform + regridding using the files which have been generated. + """ # for each target chunk, iterate through each associated regridder # for now, assume one target chunk + if len(self.unsaved_files) == 0: + msg = "No files have been generated." + raise OSError(msg) if not allow_incomplete and len(self.unsaved_files) != 0: - msg = "Not all files have been constructed." + msg = "Not all files have been generated." raise OSError(msg) current_result = None current_weights = None diff --git a/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py b/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py index defae2a0..5708362a 100644 --- a/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py +++ b/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py @@ -26,6 +26,7 @@ def test_Partition(tmp_path): files = [tmp_path / f"partial_{x}.nc" for x in range(5)] scheme = ESMFAreaWeighted(mdtol=1) + blocks = [ [[0, 100], [0, 150]], [[100, 200], [0, 150]], @@ -33,7 +34,6 @@ def test_Partition(tmp_path): [[300, 400], [0, 150]], [[400, 500], [0, 150]], ] - partition = Partition(src, tgt, scheme, files, explicit_src_blocks=blocks) partition.generate_files() @@ -51,6 +51,7 @@ def test_Partition_block_api(tmp_path): files = [tmp_path / f"partial_{x}.nc" for x in range(5)] scheme = ESMFAreaWeighted(mdtol=1) + num_src_chunks = (5, 1) partition = Partition(src, tgt, scheme, files, num_src_chunks=num_src_chunks) @@ -201,3 +202,42 @@ def test_multidimensional_cube(tmp_path): # Check metadata and coords. result.data = expected_cube.data assert result == expected_cube + +def test_save_incomplete(tmp_path): + """Test Partition class when a limited number of files are saved.""" + src = _grid_cube(150, 500, (-180, 180), (-90, 90), circular=True) + tgt = _grid_cube(16, 36, (-180, 180), (-90, 90), circular=True) + + files = [tmp_path / f"partial_{x}.nc" for x in range(5)] + src_chunks = (100, 150) + scheme = ESMFAreaWeighted(mdtol=1) + num_initial_chunks = 3 + expected_files = files[:num_initial_chunks] + + partition = Partition(src, tgt, scheme, files, src_chunks=src_chunks) + with pytest.raises(OSError): + _ = partition.apply_regridders(src, allow_incomplete=True) + + partition.generate_files(files_to_generate=num_initial_chunks) + assert partition.saved_files == expected_files + + expected_array_partial = np.ma.zeros([36, 16]) + expected_array_partial[22:] = np.ma.masked + + with pytest.raises(OSError): + _ = partition.apply_regridders(src) + partial_result = partition.apply_regridders(src, allow_incomplete=True) + assert np.ma.allclose(partial_result.data, expected_array_partial) + + loaded_partition = Partition(src, tgt, scheme, files, src_chunks=src_chunks, saved_files=expected_files) + + with pytest.raises(OSError): + _ = loaded_partition.apply_regridders(src) + partial_result_2 = partition.apply_regridders(src, allow_incomplete=True) + assert np.ma.allclose(partial_result_2.data, expected_array_partial) + + loaded_partition.generate_files() + + result = loaded_partition.apply_regridders(src) + expected_array = np.ma.zeros([36, 16]) + assert np.ma.allclose(result.data, expected_array) From 45d449cc6ef78cccb413fd813ee1d2661296801e Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Thu, 15 Jan 2026 15:49:43 +0000 Subject: [PATCH 16/41] fix test --- src/esmf_regrid/experimental/partition.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/esmf_regrid/experimental/partition.py b/src/esmf_regrid/experimental/partition.py index 4a8a3d5d..c6f846dc 100644 --- a/src/esmf_regrid/experimental/partition.py +++ b/src/esmf_regrid/experimental/partition.py @@ -155,9 +155,9 @@ def __init__( msg = "Number of source blocks does not match number of file names." raise ValueError(msg) # This will be controllable in future - tgt_chunks = None - self.tgt_chunks = tgt_chunks - if tgt_chunks is not None: + tgt_blocks = None + self.tgt_blocks = tgt_blocks + if tgt_blocks is not None: msg = "Target chunking not yet implemented." raise NotImplementedError(msg) @@ -228,7 +228,7 @@ def apply_regridders(self, cube, allow_incomplete=False): """ # for each target chunk, iterate through each associated regridder # for now, assume one target chunk - if len(self.unsaved_files) == 0: + if len(self.saved_files) == 0: msg = "No files have been generated." raise OSError(msg) if not allow_incomplete and len(self.unsaved_files) != 0: From 9bc0dc96541546eff683ff81ed0e93d47955bb10 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Fri, 16 Jan 2026 13:05:45 +0000 Subject: [PATCH 17/41] add test --- .../unit/experimental/io/partition/test_Partition.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py b/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py index 5708362a..8e259f01 100644 --- a/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py +++ b/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py @@ -17,6 +17,8 @@ _make_full_cubes, ) +from src.esmf_regrid import ESMFNearest + def test_Partition(tmp_path): """Test basic implementation of Partition class.""" @@ -241,3 +243,13 @@ def test_save_incomplete(tmp_path): result = loaded_partition.apply_regridders(src) expected_array = np.ma.zeros([36, 16]) assert np.ma.allclose(result.data, expected_array) + +def test_nearest_invalid(tmp_path): + """Test Partition class when initialised with an invalid scheme.""" + src_cube, tgt_grid, expected_cube = _make_full_cubes() + files = [tmp_path / f"partial_{x}.nc" for x in range(4)] + scheme = ESMFNearest() + chunks = (2,3) + + with pytest.raises(NotImplementedError): + _ = Partition(src_cube, tgt_grid, scheme, files, src_chunks=chunks) From 004e895774350bf7f9901591712919e8fc60db18 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Fri, 16 Jan 2026 13:13:40 +0000 Subject: [PATCH 18/41] fix test --- .../tests/unit/experimental/io/partition/test_Partition.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py b/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py index 8e259f01..b74b2df0 100644 --- a/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py +++ b/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py @@ -4,7 +4,7 @@ import numpy as np import pytest -from esmf_regrid import ESMFAreaWeighted +from esmf_regrid import ESMFAreaWeighted, ESMFNearest from esmf_regrid.experimental.partition import Partition from esmf_regrid.tests.unit.schemes.test__cube_to_GridInfo import ( _curvilinear_cube, @@ -17,8 +17,6 @@ _make_full_cubes, ) -from src.esmf_regrid import ESMFNearest - def test_Partition(tmp_path): """Test basic implementation of Partition class.""" From f7ce8ca68f70b5dbad38c7e12dda3dd4d7b73043 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Fri, 16 Jan 2026 13:22:55 +0000 Subject: [PATCH 19/41] ruff format --- src/esmf_regrid/esmf_regridder.py | 4 ++- .../io/partition/test_Partition.py | 26 ++++++++++++++----- .../test_regrid_rectilinear_to_rectilinear.py | 22 +++++++++++----- 3 files changed, 38 insertions(+), 14 deletions(-) diff --git a/src/esmf_regrid/esmf_regridder.py b/src/esmf_regrid/esmf_regridder.py index d2cc65c3..06114c32 100644 --- a/src/esmf_regrid/esmf_regridder.py +++ b/src/esmf_regrid/esmf_regridder.py @@ -250,5 +250,7 @@ def regrid(self, src_array, norm_type=Constants.NormType.FRACAREA, mdtol=1): ) raise ValueError(e_msg) tgt_weights, tgt_data = self._gen_weights_and_data(src_array) - tgt_array = self._regrid_from_weights_and_data(tgt_weights, tgt_data, norm_type=norm_type, mdtol=mdtol) + tgt_array = self._regrid_from_weights_and_data( + tgt_weights, tgt_data, norm_type=norm_type, mdtol=mdtol + ) return tgt_array diff --git a/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py b/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py index b74b2df0..399dcc49 100644 --- a/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py +++ b/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py @@ -173,9 +173,18 @@ def test_conflicting_chunks(tmp_path): ] with pytest.raises(ValueError): - _ = Partition(src, tgt, scheme, files, src_chunks=src_chunks, num_src_chunks=num_src_chunks) + _ = Partition( + src, + tgt, + scheme, + files, + src_chunks=src_chunks, + num_src_chunks=num_src_chunks, + ) with pytest.raises(ValueError): - _ = Partition(src, tgt, scheme, files, src_chunks=src_chunks, explicit_src_blocks=blocks) + _ = Partition( + src, tgt, scheme, files, src_chunks=src_chunks, explicit_src_blocks=blocks + ) with pytest.raises(ValueError): _ = Partition(src, tgt, scheme, files) with pytest.raises(TypeError): @@ -183,12 +192,13 @@ def test_conflicting_chunks(tmp_path): with pytest.raises(ValueError): _ = Partition(src, tgt, scheme, files[:-1], src_chunks=src_chunks) + def test_multidimensional_cube(tmp_path): """Test Partition class when the source has a multidimensional cube.""" src_cube, tgt_grid, expected_cube = _make_full_cubes() files = [tmp_path / f"partial_{x}.nc" for x in range(4)] scheme = ESMFAreaWeighted(mdtol=1) - chunks = (2,3) + chunks = (2, 3) partition = Partition(src_cube, tgt_grid, scheme, files, src_chunks=chunks) @@ -203,6 +213,7 @@ def test_multidimensional_cube(tmp_path): result.data = expected_cube.data assert result == expected_cube + def test_save_incomplete(tmp_path): """Test Partition class when a limited number of files are saved.""" src = _grid_cube(150, 500, (-180, 180), (-90, 90), circular=True) @@ -229,7 +240,9 @@ def test_save_incomplete(tmp_path): partial_result = partition.apply_regridders(src, allow_incomplete=True) assert np.ma.allclose(partial_result.data, expected_array_partial) - loaded_partition = Partition(src, tgt, scheme, files, src_chunks=src_chunks, saved_files=expected_files) + loaded_partition = Partition( + src, tgt, scheme, files, src_chunks=src_chunks, saved_files=expected_files + ) with pytest.raises(OSError): _ = loaded_partition.apply_regridders(src) @@ -242,12 +255,13 @@ def test_save_incomplete(tmp_path): expected_array = np.ma.zeros([36, 16]) assert np.ma.allclose(result.data, expected_array) + def test_nearest_invalid(tmp_path): """Test Partition class when initialised with an invalid scheme.""" - src_cube, tgt_grid, expected_cube = _make_full_cubes() + src_cube, tgt_grid, _ = _make_full_cubes() files = [tmp_path / f"partial_{x}.nc" for x in range(4)] scheme = ESMFNearest() - chunks = (2,3) + chunks = (2, 3) with pytest.raises(NotImplementedError): _ = Partition(src_cube, tgt_grid, scheme, files, src_chunks=chunks) diff --git a/src/esmf_regrid/tests/unit/schemes/test_regrid_rectilinear_to_rectilinear.py b/src/esmf_regrid/tests/unit/schemes/test_regrid_rectilinear_to_rectilinear.py index 296a3ad4..1233b696 100644 --- a/src/esmf_regrid/tests/unit/schemes/test_regrid_rectilinear_to_rectilinear.py +++ b/src/esmf_regrid/tests/unit/schemes/test_regrid_rectilinear_to_rectilinear.py @@ -81,6 +81,7 @@ def _add_metadata(cube): result.add_aux_coord(scalar_time) return result + def _make_full_cubes(src_rectilinear=True, tgt_rectilinear=True): h = 2 t = 4 @@ -142,9 +143,8 @@ def _make_full_cubes(src_rectilinear=True, tgt_rectilinear=True): expected_data = np.empty([h, tgt_lats, t, tgt_lons, e]) expected_data[:] = np.arange(t * h * e).reshape([h, t, e])[ - :, np.newaxis, :, np.newaxis, : - ] - + :, np.newaxis, :, np.newaxis, : + ] expected_cube = Cube(expected_data) if tgt_rectilinear: @@ -167,7 +167,9 @@ def test_extra_dims(): Tests the handling of extra dimensions and metadata. Ensures that proper coordinates, attributes, names and units are copied over. """ - src_cube, tgt_grid, expected_cube = _make_full_cubes(src_rectilinear=True, tgt_rectilinear=True) + src_cube, tgt_grid, expected_cube = _make_full_cubes( + src_rectilinear=True, tgt_rectilinear=True + ) result = regrid_rectilinear_to_rectilinear(src_cube, tgt_grid) @@ -291,7 +293,9 @@ def test_extra_dims_curvilinear(): Tests the handling of extra dimensions and metadata. Ensures that proper coordinates, attributes, names and units are copied over. """ - src_cube, tgt_grid, expected_cube = _make_full_cubes(src_rectilinear=False, tgt_rectilinear=False) + src_cube, tgt_grid, expected_cube = _make_full_cubes( + src_rectilinear=False, tgt_rectilinear=False + ) result = regrid_rectilinear_to_rectilinear(src_cube, tgt_grid) @@ -309,7 +313,9 @@ def test_extra_dims_curvilinear_to_rectilinear(): Tests the handling of extra dimensions and metadata. Ensures that proper coordinates, attributes, names and units are copied over. """ - src_cube, tgt_grid, expected_cube = _make_full_cubes(src_rectilinear=False, tgt_rectilinear=True) + src_cube, tgt_grid, expected_cube = _make_full_cubes( + src_rectilinear=False, tgt_rectilinear=True + ) result = regrid_rectilinear_to_rectilinear(src_cube, tgt_grid) @@ -327,7 +333,9 @@ def test_extra_dims_rectilinear_to_curvilinear(): Tests the handling of extra dimensions and metadata. Ensures that proper coordinates, attributes, names and units are copied over. """ - src_cube, tgt_grid, expected_cube = _make_full_cubes(src_rectilinear=True, tgt_rectilinear=False) + src_cube, tgt_grid, expected_cube = _make_full_cubes( + src_rectilinear=True, tgt_rectilinear=False + ) result = regrid_rectilinear_to_rectilinear(src_cube, tgt_grid) From 6f332517c4dca6abd2cd16846a6298c19e705a81 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Tue, 20 Jan 2026 16:42:59 +0000 Subject: [PATCH 20/41] add documentation --- docs/src/userguide/examples.rst | 87 +++++++++++++++++++++++ src/esmf_regrid/experimental/partition.py | 2 - 2 files changed, 87 insertions(+), 2 deletions(-) diff --git a/docs/src/userguide/examples.rst b/docs/src/userguide/examples.rst index e66a5a14..a3e4917b 100644 --- a/docs/src/userguide/examples.rst +++ b/docs/src/userguide/examples.rst @@ -54,6 +54,93 @@ 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 to save +the parts of the regridder and a way to divide the source grid into blocks:: + + 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 + ) + +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. + +.. 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 possible to + generate the remaining files in a different python session.:: + + # 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) + +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. diff --git a/src/esmf_regrid/experimental/partition.py b/src/esmf_regrid/experimental/partition.py index c6f846dc..31229042 100644 --- a/src/esmf_regrid/experimental/partition.py +++ b/src/esmf_regrid/experimental/partition.py @@ -81,8 +81,6 @@ def __init__( src_chunks=None, num_src_chunks=None, explicit_src_blocks=None, - # tgt_chunks=None, - # num_tgt_chunks=None, auto_generate=False, saved_files=None, ): From c3673253796e05c7edd56c609f3f14af971e08f8 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Wed, 21 Jan 2026 00:03:07 +0000 Subject: [PATCH 21/41] add to docstrings --- src/esmf_regrid/experimental/io.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/esmf_regrid/experimental/io.py b/src/esmf_regrid/experimental/io.py index 3a70c861..56c5958b 100644 --- a/src/esmf_regrid/experimental/io.py +++ b/src/esmf_regrid/experimental/io.py @@ -228,6 +228,8 @@ def save_regridder(rg, filename, allow_partial=False): The regridder instance to save. filename : str The file name to save to. + allow_partial : bool, default=False + If True, allow the saving of :class:`~esmf_regrid.experimental._partial.PartialRegridder` instances. """ regridder_type = rg.__class__.__name__ @@ -354,6 +356,8 @@ def load_regridder(filename, allow_partial=False): ---------- filename : str The file name to load from. + allow_partial : bool, default=False + If True, allow the loading of :class:`~esmf_regrid.experimental._partial.PartialRegridder` instances. Returns ------- From 9cbb76ec59a10eb69552500adcbbde1e60cefba8 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Wed, 21 Jan 2026 11:22:23 +0000 Subject: [PATCH 22/41] add docstrings and repr to PartialRegridder --- src/esmf_regrid/experimental/_partial.py | 34 ++++++++++++++++++++++- src/esmf_regrid/experimental/partition.py | 2 +- src/esmf_regrid/schemes.py | 16 +++++------ 3 files changed, 42 insertions(+), 10 deletions(-) diff --git a/src/esmf_regrid/experimental/_partial.py b/src/esmf_regrid/experimental/_partial.py index 7ef54b23..c0377aa1 100644 --- a/src/esmf_regrid/experimental/_partial.py +++ b/src/esmf_regrid/experimental/_partial.py @@ -9,11 +9,31 @@ class PartialRegridder(_ESMFRegridder): + """Regridder class designed for use in :class:`~esmf_regrid.experimental._partial.Partial`.""" + def __init__(self, src, tgt, src_slice, tgt_slice, weights, scheme, **kwargs): + """Create a regridder instance for a block of :class:`~esmf_regrid.experimental._partial.Partial`. + + 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. + tgt_slice : tuple + The upper and lower bounds of the block taken from the original target from which the + ``tgt`` was derived. + 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. + """ self.src_slice = src_slice # this will be tuple-like self.tgt_slice = tgt_slice self.scheme = scheme - # TODO: consider disallowing ESMFNearest (unless out of bounds can be made masked) # Pop duplicate kwargs. for arg in set(kwargs.keys()).intersection(vars(self.scheme)): @@ -27,7 +47,18 @@ def __init__(self, src, tgt, src_slice, tgt_slice, weights, scheme, **kwargs): ) self.__dict__.update(self._regridder.__dict__) + def __repr__(self): + """Return a representation of the class.""" + result = ( + f"PartialRegridder(" + 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_out_dims = self.regridder.tgt.dims num_dims = len(dims) @@ -44,6 +75,7 @@ def partial_regrid(self, src): return result def finish_regridding(self, src_cube, weights, data): + """Perform the second half of regridding, combining weights and data.""" dims = self._get_cube_dims(src_cube) result_data = self.regridder._regrid_from_weights_and_data(weights, data) diff --git a/src/esmf_regrid/experimental/partition.py b/src/esmf_regrid/experimental/partition.py index 31229042..657fd8a6 100644 --- a/src/esmf_regrid/experimental/partition.py +++ b/src/esmf_regrid/experimental/partition.py @@ -210,7 +210,7 @@ def generate_files(self, files_to_generate=None): regridder = self.scheme.regridder(src, tgt) weights = regridder.regridder.weight_matrix regridder = PartialRegridder( - src, self.tgt, src_block, None, weights, self.scheme + src, tgt, src_block, None, weights, self.scheme ) save_regridder(regridder, file, allow_partial=True) self.saved_files.append(file) diff --git a/src/esmf_regrid/schemes.py b/src/esmf_regrid/schemes.py index 739fa41c..25f8c9e1 100644 --- a/src/esmf_regrid/schemes.py +++ b/src/esmf_regrid/schemes.py @@ -1500,9 +1500,9 @@ def __init__( Parameters ---------- src : :class:`iris.cube.Cube` - The rectilinear :class:`~iris.cube.Cube` providing the source grid. + The :class:`~iris.cube.Cube` providing the source grid. tgt : :class:`iris.cube.Cube` or :class:`iris.mesh.MeshXY` - The rectilinear :class:`~iris.cube.Cube` providing the target grid. + The :class:`~iris.cube.Cube` providing the target grid. method : :class:`Constants.Method` The method to be used to calculate weights. mdtol : float, default=None @@ -1686,11 +1686,11 @@ def __init__( Parameters ---------- src : :class:`iris.cube.Cube` - The rectilinear :class:`~iris.cube.Cube` providing the source. + The :class:`~iris.cube.Cube` providing the source. If this cube has a grid defined by latitude/longitude coordinates, those coordinates must have bounds. tgt : :class:`iris.cube.Cube` or :class:`iris.mesh.MeshXY` - The unstructured :class:`~iris.cube.Cube`or + The :class:`~iris.cube.Cube`or :class:`~iris.mesh.MeshXY` defining the target. If this cube has a grid defined by latitude/longitude coordinates, those coordinates must have bounds. @@ -1773,9 +1773,9 @@ def __init__( Parameters ---------- src : :class:`iris.cube.Cube` - The rectilinear :class:`~iris.cube.Cube` providing the source. + The :class:`~iris.cube.Cube` providing the source. tgt : :class:`iris.cube.Cube` or :class:`iris.mesh.MeshXY` - The unstructured :class:`~iris.cube.Cube`or + The :class:`~iris.cube.Cube`or :class:`~iris.mesh.MeshXY` defining the target. mdtol : float, default=0 Tolerance of missing data. The value returned in each element of @@ -1847,9 +1847,9 @@ def __init__( Parameters ---------- src : :class:`iris.cube.Cube` - The rectilinear :class:`~iris.cube.Cube` providing the source. + The :class:`~iris.cube.Cube` providing the source. tgt : :class:`iris.cube.Cube` or :class:`iris.mesh.MeshXY` - The unstructured :class:`~iris.cube.Cube`or + The :class:`~iris.cube.Cube`or :class:`~iris.mesh.MeshXY` defining the target. precomputed_weights : :class:`scipy.sparse.spmatrix`, optional If ``None``, :mod:`esmpy` will be used to From 79b56a1d5caa9fa299fa9379a84997c5a1d31dd9 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Wed, 21 Jan 2026 12:00:19 +0000 Subject: [PATCH 23/41] repr testing --- src/esmf_regrid/experimental/partition.py | 4 ++-- .../{io => }/partition/__init__.py | 0 .../partition/test_PartialRegridder.py | 22 +++++++++++++++++++ .../{io => }/partition/test_Partition.py | 16 ++++++++++++++ 4 files changed, 40 insertions(+), 2 deletions(-) rename src/esmf_regrid/tests/unit/experimental/{io => }/partition/__init__.py (100%) create mode 100644 src/esmf_regrid/tests/unit/experimental/partition/test_PartialRegridder.py rename src/esmf_regrid/tests/unit/experimental/{io => }/partition/test_Partition.py (91%) diff --git a/src/esmf_regrid/experimental/partition.py b/src/esmf_regrid/experimental/partition.py index 657fd8a6..3bb99e1c 100644 --- a/src/esmf_regrid/experimental/partition.py +++ b/src/esmf_regrid/experimental/partition.py @@ -173,8 +173,8 @@ def __repr__(self): """Return a representation of the class.""" result = ( f"Partition(" - f"src={self.src}, " - f"tgt={self.tgt}, " + f"src={repr(self.src)}, " + f"tgt={repr(self.tgt)}, " f"scheme={self.scheme}, " f"num file_names={len(self.file_names)}," f"num saved_files={len(self.saved_files)})" diff --git a/src/esmf_regrid/tests/unit/experimental/io/partition/__init__.py b/src/esmf_regrid/tests/unit/experimental/partition/__init__.py similarity index 100% rename from src/esmf_regrid/tests/unit/experimental/io/partition/__init__.py rename to src/esmf_regrid/tests/unit/experimental/partition/__init__.py diff --git a/src/esmf_regrid/tests/unit/experimental/partition/test_PartialRegridder.py b/src/esmf_regrid/tests/unit/experimental/partition/test_PartialRegridder.py new file mode 100644 index 00000000..9b0e2199 --- /dev/null +++ b/src/esmf_regrid/tests/unit/experimental/partition/test_PartialRegridder.py @@ -0,0 +1,22 @@ +"""Unit tests for :mod:`esmf_regrid.experimental.partition`.""" + +from esmf_regrid import ESMFAreaWeighted +from esmf_regrid.experimental._partial import PartialRegridder +from esmf_regrid.tests.unit.schemes.test__cube_to_GridInfo import ( + _grid_cube, +) + +def test_PartialRegridder_repr(): + """Test repr of PartialRegridder instance.""" + src = _grid_cube(10, 15, (-180, 180), (-90, 90), circular=True) + tgt = _grid_cube(5, 10, (-180, 180), (-90, 90), circular=True) + src_slice = ((10,20), (15, 30)) + tgt_slice = ((0, 5), (0, 10)) + weights = None + scheme = ESMFAreaWeighted(mdtol=0.5) + + pr = PartialRegridder(src, tgt, src_slice, tgt_slice, weights, scheme) + + expected_repr = ("PartialRegridder(src_slice=((10, 20), (15, 30)), tgt_slice=((0, 5), (0, 10)), " + "scheme=ESMFAreaWeighted(mdtol=0.5, use_src_mask=False, use_tgt_mask=False, esmf_args={}))") + assert repr(pr) == expected_repr diff --git a/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py b/src/esmf_regrid/tests/unit/experimental/partition/test_Partition.py similarity index 91% rename from src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py rename to src/esmf_regrid/tests/unit/experimental/partition/test_Partition.py index 399dcc49..0c616800 100644 --- a/src/esmf_regrid/tests/unit/experimental/io/partition/test_Partition.py +++ b/src/esmf_regrid/tests/unit/experimental/partition/test_Partition.py @@ -265,3 +265,19 @@ def test_nearest_invalid(tmp_path): with pytest.raises(NotImplementedError): _ = Partition(src_cube, tgt_grid, scheme, files, src_chunks=chunks) + +def test_Partition_repr(tmp_path): + """Test repr of Partition instance.""" + src_cube, tgt_grid, _ = _make_full_cubes() + files = [tmp_path / f"partial_{x}.nc" for x in range(4)] + scheme = ESMFAreaWeighted() + chunks = (2, 3) + + partition = Partition(src_cube, tgt_grid, scheme, files, src_chunks=chunks) + + expected_repr = ("Partition(src=, " + "tgt=, " + "scheme=ESMFAreaWeighted(mdtol=0, use_src_mask=False, use_tgt_mask=False, esmf_args={}), " + "num file_names=4,num saved_files=0)") + assert repr(partition) == expected_repr From 90968b6d3831770af527144d9d11fa90096f7f7e Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Wed, 21 Jan 2026 13:23:46 +0000 Subject: [PATCH 24/41] ruff fix --- src/esmf_regrid/experimental/partition.py | 4 ++-- .../experimental/partition/test_PartialRegridder.py | 9 ++++++--- .../unit/experimental/partition/test_Partition.py | 13 ++++++++----- 3 files changed, 16 insertions(+), 10 deletions(-) diff --git a/src/esmf_regrid/experimental/partition.py b/src/esmf_regrid/experimental/partition.py index 3bb99e1c..b46f73f3 100644 --- a/src/esmf_regrid/experimental/partition.py +++ b/src/esmf_regrid/experimental/partition.py @@ -173,8 +173,8 @@ def __repr__(self): """Return a representation of the class.""" result = ( f"Partition(" - f"src={repr(self.src)}, " - f"tgt={repr(self.tgt)}, " + f"src={self.src!r}, " + f"tgt={self.tgt!r}, " f"scheme={self.scheme}, " f"num file_names={len(self.file_names)}," f"num saved_files={len(self.saved_files)})" diff --git a/src/esmf_regrid/tests/unit/experimental/partition/test_PartialRegridder.py b/src/esmf_regrid/tests/unit/experimental/partition/test_PartialRegridder.py index 9b0e2199..37ad2ecc 100644 --- a/src/esmf_regrid/tests/unit/experimental/partition/test_PartialRegridder.py +++ b/src/esmf_regrid/tests/unit/experimental/partition/test_PartialRegridder.py @@ -6,17 +6,20 @@ _grid_cube, ) + def test_PartialRegridder_repr(): """Test repr of PartialRegridder instance.""" src = _grid_cube(10, 15, (-180, 180), (-90, 90), circular=True) tgt = _grid_cube(5, 10, (-180, 180), (-90, 90), circular=True) - src_slice = ((10,20), (15, 30)) + src_slice = ((10, 20), (15, 30)) tgt_slice = ((0, 5), (0, 10)) weights = None scheme = ESMFAreaWeighted(mdtol=0.5) pr = PartialRegridder(src, tgt, src_slice, tgt_slice, weights, scheme) - expected_repr = ("PartialRegridder(src_slice=((10, 20), (15, 30)), tgt_slice=((0, 5), (0, 10)), " - "scheme=ESMFAreaWeighted(mdtol=0.5, use_src_mask=False, use_tgt_mask=False, esmf_args={}))") + expected_repr = ( + "PartialRegridder(src_slice=((10, 20), (15, 30)), tgt_slice=((0, 5), (0, 10)), " + "scheme=ESMFAreaWeighted(mdtol=0.5, use_src_mask=False, use_tgt_mask=False, esmf_args={}))" + ) assert repr(pr) == expected_repr diff --git a/src/esmf_regrid/tests/unit/experimental/partition/test_Partition.py b/src/esmf_regrid/tests/unit/experimental/partition/test_Partition.py index 0c616800..0d72b457 100644 --- a/src/esmf_regrid/tests/unit/experimental/partition/test_Partition.py +++ b/src/esmf_regrid/tests/unit/experimental/partition/test_Partition.py @@ -266,6 +266,7 @@ def test_nearest_invalid(tmp_path): with pytest.raises(NotImplementedError): _ = Partition(src_cube, tgt_grid, scheme, files, src_chunks=chunks) + def test_Partition_repr(tmp_path): """Test repr of Partition instance.""" src_cube, tgt_grid, _ = _make_full_cubes() @@ -275,9 +276,11 @@ def test_Partition_repr(tmp_path): partition = Partition(src_cube, tgt_grid, scheme, files, src_chunks=chunks) - expected_repr = ("Partition(src=, " - "tgt=, " - "scheme=ESMFAreaWeighted(mdtol=0, use_src_mask=False, use_tgt_mask=False, esmf_args={}), " - "num file_names=4,num saved_files=0)") + expected_repr = ( + "Partition(src=, " + "tgt=, " + "scheme=ESMFAreaWeighted(mdtol=0, use_src_mask=False, use_tgt_mask=False, esmf_args={}), " + "num file_names=4,num saved_files=0)" + ) assert repr(partition) == expected_repr From 15c1f5a308270d4e055a0039dedc25ffaeaa9ecb Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Wed, 21 Jan 2026 13:27:04 +0000 Subject: [PATCH 25/41] fix docs --- docs/src/userguide/examples.rst | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/src/userguide/examples.rst b/docs/src/userguide/examples.rst index a3e4917b..c774cf12 100644 --- a/docs/src/userguide/examples.rst +++ b/docs/src/userguide/examples.rst @@ -106,7 +106,8 @@ method. 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 possible to - generate the remaining files in a different python session.:: + generate the remaining files in a different python session. +:: # Generate partial regridders and save them to the list of files. partition.generate_files() From 84c9482853489728c43d528c1449dc4fcbc641cd Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Wed, 21 Jan 2026 16:53:35 +0000 Subject: [PATCH 26/41] docs grammar --- docs/src/userguide/examples.rst | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/src/userguide/examples.rst b/docs/src/userguide/examples.rst index c774cf12..9eb15980 100644 --- a/docs/src/userguide/examples.rst +++ b/docs/src/userguide/examples.rst @@ -64,8 +64,8 @@ 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 to save -the parts of the regridder and a way to divide the source grid into blocks:: +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 @@ -117,7 +117,7 @@ method. 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.:: +files which have already been generated:: # Use the same arguments which constructed the original partition. source_cube = make_gridcube(nx=800, ny=800) From b6aca52a11df48e52ae9ce7e78f1293cb1d073e0 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Thu, 22 Jan 2026 16:26:06 +0000 Subject: [PATCH 27/41] attempt benchmarks slowdown fix --- src/esmf_regrid/esmf_regridder.py | 13 +++++++------ src/esmf_regrid/experimental/_partial.py | 23 +++++++++++++++++------ src/esmf_regrid/experimental/partition.py | 4 ++-- 3 files changed, 26 insertions(+), 14 deletions(-) diff --git a/src/esmf_regrid/esmf_regridder.py b/src/esmf_regrid/esmf_regridder.py index 06114c32..9344d88b 100644 --- a/src/esmf_regrid/esmf_regridder.py +++ b/src/esmf_regrid/esmf_regridder.py @@ -189,12 +189,12 @@ def _gen_weights_and_data(self, src_array): src_inverted_mask = self.src._array_to_matrix(~ma.getmaskarray(src_array)) weight_sums = weight_matrix @ src_inverted_mask - tgt_data = self.tgt._matrix_to_array(flat_tgt, extra_shape) - tgt_weights = self.tgt._matrix_to_array(weight_sums, extra_shape) - return tgt_weights, tgt_data + # tgt_data = self.tgt._matrix_to_array(flat_tgt, extra_shape) + # tgt_weights = self.tgt._matrix_to_array(weight_sums, extra_shape) + return weight_sums, flat_tgt, extra_shape def _regrid_from_weights_and_data( - self, tgt_weights, tgt_data, norm_type=Constants.NormType.FRACAREA, mdtol=1 + 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. @@ -210,6 +210,7 @@ def _regrid_from_weights_and_data( 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): @@ -249,8 +250,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) - tgt_weights, tgt_data = self._gen_weights_and_data(src_array) + tgt_weights, tgt_data, extra = self._gen_weights_and_data(src_array) tgt_array = self._regrid_from_weights_and_data( - tgt_weights, tgt_data, norm_type=norm_type, mdtol=mdtol + tgt_weights, tgt_data, extra, norm_type=norm_type, mdtol=mdtol ) return tgt_array diff --git a/src/esmf_regrid/experimental/_partial.py b/src/esmf_regrid/experimental/_partial.py index c0377aa1..087d2b5c 100644 --- a/src/esmf_regrid/experimental/_partial.py +++ b/src/esmf_regrid/experimental/_partial.py @@ -66,19 +66,30 @@ def partial_regrid(self, src): data = np.moveaxis(src.data, dims, standard_in_dims) result = self.regridder._gen_weights_and_data(data) + # standard_out_dims = [-1, -2][:num_out_dims] + # if num_dims == 2 and num_out_dims == 1: + # dims = [min(dims)] + # if num_dims == 1 and num_out_dims == 2: + # dims = [dims[0] + 1, dims[0]] + # result = tuple(np.moveaxis(r, standard_out_dims, dims) for r in result) + return result + + def finish_regridding(self, src_cube, weights, data, extra): + """Perform the second half of regridding, combining weights and data.""" + 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(dims) standard_out_dims = [-1, -2][:num_out_dims] if num_dims == 2 and num_out_dims == 1: dims = [min(dims)] if num_dims == 1 and num_out_dims == 2: dims = [dims[0] + 1, dims[0]] - result = tuple(np.moveaxis(r, standard_out_dims, dims) for r in result) - return result - def finish_regridding(self, src_cube, weights, data): - """Perform the second half of regridding, combining weights and data.""" - dims = self._get_cube_dims(src_cube) + result_data = np.moveaxis(result_data, standard_out_dims, dims) - result_data = self.regridder._regrid_from_weights_and_data(weights, data) result_cube = _create_cube( result_data, src_cube, dims, self._tgt, len(self._tgt) ) diff --git a/src/esmf_regrid/experimental/partition.py b/src/esmf_regrid/experimental/partition.py index b46f73f3..59fb7330 100644 --- a/src/esmf_regrid/experimental/partition.py +++ b/src/esmf_regrid/experimental/partition.py @@ -240,7 +240,7 @@ def apply_regridders(self, cube, allow_incomplete=False): if file in files: next_regridder = load_regridder(file, allow_partial=True) cube_chunk = _get_chunk(cube, chunk) - next_weights, next_result = next_regridder.partial_regrid(cube_chunk) + next_weights, next_result, extra = next_regridder.partial_regrid(cube_chunk) if current_weights is None: current_weights = next_weights else: @@ -251,5 +251,5 @@ def apply_regridders(self, cube, allow_incomplete=False): current_result += next_result return next_regridder.finish_regridding( - cube_chunk, current_weights, current_result + cube_chunk, current_weights, current_result, extra, ) From 2f9ba4eaf9c8acf61ff8f235c69e506b7708a03f Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Thu, 22 Jan 2026 22:54:32 +0000 Subject: [PATCH 28/41] tidy unused code --- src/esmf_regrid/esmf_regridder.py | 3 --- src/esmf_regrid/experimental/_partial.py | 8 -------- 2 files changed, 11 deletions(-) diff --git a/src/esmf_regrid/esmf_regridder.py b/src/esmf_regrid/esmf_regridder.py index 9344d88b..9ded1511 100644 --- a/src/esmf_regrid/esmf_regridder.py +++ b/src/esmf_regrid/esmf_regridder.py @@ -188,9 +188,6 @@ def _gen_weights_and_data(self, src_array): src_inverted_mask = self.src._array_to_matrix(~ma.getmaskarray(src_array)) weight_sums = weight_matrix @ src_inverted_mask - - # tgt_data = self.tgt._matrix_to_array(flat_tgt, extra_shape) - # tgt_weights = self.tgt._matrix_to_array(weight_sums, extra_shape) return weight_sums, flat_tgt, extra_shape def _regrid_from_weights_and_data( diff --git a/src/esmf_regrid/experimental/_partial.py b/src/esmf_regrid/experimental/_partial.py index 087d2b5c..11e5b90c 100644 --- a/src/esmf_regrid/experimental/_partial.py +++ b/src/esmf_regrid/experimental/_partial.py @@ -60,18 +60,10 @@ def __repr__(self): def partial_regrid(self, src): """Perform the first half of regridding, generating weights and data.""" dims = self._get_cube_dims(src) - num_out_dims = self.regridder.tgt.dims 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) - - # standard_out_dims = [-1, -2][:num_out_dims] - # if num_dims == 2 and num_out_dims == 1: - # dims = [min(dims)] - # if num_dims == 1 and num_out_dims == 2: - # dims = [dims[0] + 1, dims[0]] - # result = tuple(np.moveaxis(r, standard_out_dims, dims) for r in result) return result def finish_regridding(self, src_cube, weights, data, extra): From ae5b6388ef07adea6cf0acc3217f1d0f4bc22962 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Mon, 26 Jan 2026 14:03:07 +0000 Subject: [PATCH 29/41] ruff fix --- src/esmf_regrid/esmf_regridder.py | 7 ++++++- src/esmf_regrid/experimental/partition.py | 9 +++++++-- 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/src/esmf_regrid/esmf_regridder.py b/src/esmf_regrid/esmf_regridder.py index 9ded1511..c335320c 100644 --- a/src/esmf_regrid/esmf_regridder.py +++ b/src/esmf_regrid/esmf_regridder.py @@ -191,7 +191,12 @@ def _gen_weights_and_data(self, src_array): 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 + 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. diff --git a/src/esmf_regrid/experimental/partition.py b/src/esmf_regrid/experimental/partition.py index 59fb7330..1c73a28e 100644 --- a/src/esmf_regrid/experimental/partition.py +++ b/src/esmf_regrid/experimental/partition.py @@ -240,7 +240,9 @@ def apply_regridders(self, cube, allow_incomplete=False): if file in files: next_regridder = load_regridder(file, allow_partial=True) cube_chunk = _get_chunk(cube, chunk) - next_weights, next_result, extra = next_regridder.partial_regrid(cube_chunk) + next_weights, next_result, extra = next_regridder.partial_regrid( + cube_chunk + ) if current_weights is None: current_weights = next_weights else: @@ -251,5 +253,8 @@ def apply_regridders(self, cube, allow_incomplete=False): current_result += next_result return next_regridder.finish_regridding( - cube_chunk, current_weights, current_result, extra, + cube_chunk, + current_weights, + current_result, + extra, ) From 905bbfe26cf407de720eca68e5a8d5d3e99c02a7 Mon Sep 17 00:00:00 2001 From: stephenworsley <49274989+stephenworsley@users.noreply.github.com> Date: Tue, 3 Feb 2026 14:08:06 +0000 Subject: [PATCH 30/41] Update src/esmf_regrid/experimental/partition.py Co-authored-by: Patrick Peglar --- src/esmf_regrid/experimental/partition.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/esmf_regrid/experimental/partition.py b/src/esmf_regrid/experimental/partition.py index 1c73a28e..1416f12b 100644 --- a/src/esmf_regrid/experimental/partition.py +++ b/src/esmf_regrid/experimental/partition.py @@ -184,8 +184,7 @@ def __repr__(self): @property def unsaved_files(self): """List of files not yet generated.""" - files = set(self.file_names) - set(self.saved_files) - return [file for file in self.file_names if file in files] + return [file for file in self.file_names if file not in self.saved_files] def generate_files(self, files_to_generate=None): """Generate files with regridding information. From a9f4397cd70b83f6f617f9aaaf6438aec1b3d741 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Wed, 4 Feb 2026 16:04:11 +0000 Subject: [PATCH 31/41] review actions --- docs/src/userguide/examples.rst | 19 ++++--- pyproject.toml | 1 - src/esmf_regrid/experimental/_partial.py | 49 ++++++++++++---- src/esmf_regrid/experimental/partition.py | 56 ++++++++++++------- .../experimental/partition/test_Partition.py | 20 +++++++ 5 files changed, 105 insertions(+), 40 deletions(-) diff --git a/docs/src/userguide/examples.rst b/docs/src/userguide/examples.rst index 9eb15980..a24edd3c 100644 --- a/docs/src/userguide/examples.rst +++ b/docs/src/userguide/examples.rst @@ -97,17 +97,14 @@ to divide the source grid into blocks whose regridders are saved to those 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. - -.. 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 possible to - generate the remaining files in a different python session. -:: +method:: # Generate partial regridders and save them to the list of files. partition.generate_files() @@ -115,6 +112,12 @@ method. # 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:: diff --git a/pyproject.toml b/pyproject.toml index e7b22a62..92a2cd3e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -181,7 +181,6 @@ ignore = [ "ANN202", "ANN204", - "B905", # Zip strictness should be explicit "D104", # Misssing docstring "E501", # Line too long "ERA001", # Commented out code diff --git a/src/esmf_regrid/experimental/_partial.py b/src/esmf_regrid/experimental/_partial.py index 11e5b90c..381d245a 100644 --- a/src/esmf_regrid/experimental/_partial.py +++ b/src/esmf_regrid/experimental/_partial.py @@ -5,14 +5,16 @@ from esmf_regrid.schemes import ( _create_cube, _ESMFRegridder, + GridRecord, + MeshRecord, ) class PartialRegridder(_ESMFRegridder): - """Regridder class designed for use in :class:`~esmf_regrid.experimental._partial.Partial`.""" + """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._partial.Partial`. + """Create a regridder instance for a block of :class:`~esmf_regrid.experimental.Partition`. Parameters ---------- @@ -67,22 +69,49 @@ def partial_regrid(self, src): return result def finish_regridding(self, src_cube, weights, data, extra): - """Perform the second half of regridding, combining weights and data.""" - dims = self._get_cube_dims(src_cube) + """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. + """ + old_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(dims) + num_dims = len(old_dims) standard_out_dims = [-1, -2][:num_out_dims] if num_dims == 2 and num_out_dims == 1: - dims = [min(dims)] - if num_dims == 1 and num_out_dims == 2: - dims = [dims[0] + 1, dims[0]] + new_dims = [min(old_dims)] + elif num_dims == 1 and num_out_dims == 2: + new_dims = [old_dims[0] + 1, old_dims[0]] + else: + new_dims = old_dims + + result_data = np.moveaxis(result_data, standard_out_dims, new_dims) - result_data = np.moveaxis(result_data, standard_out_dims, 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 ValueError(msg) result_cube = _create_cube( - result_data, src_cube, dims, self._tgt, len(self._tgt) + result_data, src_cube, old_dims, tgt_coords, out_dims ) return result_cube diff --git a/src/esmf_regrid/experimental/partition.py b/src/esmf_regrid/experimental/partition.py index 1416f12b..2c949f88 100644 --- a/src/esmf_regrid/experimental/partition.py +++ b/src/esmf_regrid/experimental/partition.py @@ -19,11 +19,11 @@ def _get_chunk(cube, sl): return cube[*full_slice] -def _determine_blocks(shape, chunks, num_chunks, explicit_chunks): +def _determine_blocks(shape, chunks, num_chunks, explicit_blocks): which_inputs = ( chunks is not None, num_chunks is not None, - explicit_chunks is not None, + explicit_blocks is not None, ) if sum(which_inputs) == 0: msg = "Partition blocks must must be specified by either chunks, num_chunks, or explicit_chunks." @@ -32,7 +32,7 @@ def _determine_blocks(shape, chunks, num_chunks, explicit_chunks): msg = "Potentially conflicting partition block definitions." raise ValueError(msg) if num_chunks is not None: - chunks = [s // n for s, n in zip(shape, num_chunks)] + chunks = [s // n for s, n in zip(shape, num_chunks, strict=True)] for chunk in chunks: if chunk == 0: msg = "`num_chunks` cannot divide a dimension into more blocks than the size of that dimension." @@ -40,24 +40,29 @@ def _determine_blocks(shape, chunks, num_chunks, explicit_chunks): if chunks is not None: if all(isinstance(x, int) for x in chunks): proper_chunks = [] - for s, c in zip(shape, chunks): + for s, c in zip(shape, chunks, strict=True): proper_chunk = [c] * (s // c) if s % c != 0: proper_chunk += [s % c] proper_chunks.append(proper_chunk) chunks = proper_chunks - for s, chunk in zip(shape, chunks): + for s, chunk in zip(shape, chunks, strict=True): if sum(chunk) != s: msg = "Chunks must sum to the size of their respective dimension." raise ValueError(msg) bounds = [np.cumsum([0, *chunk]) for chunk in chunks] if len(bounds) == 1: - explicit_chunks = [ - [[int(lower), int(upper)]] - for lower, upper in zip(bounds[0][:-1], bounds[0][1:]) - ] + msg = "Chunks must have exactly two dimensions." + raise ValueError(msg) + # TODO: This is currently blocked by the fact that slicing an Iris cube on its mesh dimension + # does not currently yield another cube with a mesh. When this is fixed, the following + # code can be uncommented. + # explicit_blocks = [ + # [[int(lower), int(upper)]] + # for lower, upper in zip(bounds[0][:-1], bounds[0][1:]) + # ] elif len(bounds) == 2: - explicit_chunks = [ + explicit_blocks = [ [[int(ly), int(uy)], [int(lx), int(ux)]] for ly, uy in zip(bounds[0][:-1], bounds[0][1:]) for lx, ux in zip(bounds[1][:-1], bounds[1][1:]) @@ -65,7 +70,10 @@ def _determine_blocks(shape, chunks, num_chunks, explicit_chunks): else: msg = "Chunks must not exceed two dimensions." raise ValueError(msg) - return explicit_chunks + if len(explicit_blocks[0]) != len(shape): + msg = "Dimensionality of blocks does not match the number of dimensions." + raise ValueError(msg) + return explicit_blocks class Partition: @@ -107,9 +115,9 @@ def __init__( Specify the size of blocks to use to divide up the cube. Dimensions are specified in y,x axis order. If `src_chunks` is a tuple of int, each integer describes the maximum size of a block in that dimension. If `src_chunks` is a tuple of tuples, - each sub-tuple describes the size of each successive block in that dimension. These - block sizes should add up to the total size of that dimension or else an error - is raised. + each sub-tuple describes the size of each successive block in that dimension. The sum + of these block sizes in each of the sub-tuples should add up to the total size of that + dimension or else an error is raised. num_src_chunks : tuple of int Specify the number of blocks to use to divide up the cube. Dimensions are specified in y,x axis order. Each integer describes the number of blocks that dimension will @@ -124,15 +132,17 @@ def __init__( if scheme._method == Constants.Method.NEAREST: msg = "The `Nearest` method is not implemented." raise NotImplementedError(msg) - if src.mesh is not None: - msg = "Partition does not yet support source meshes." - raise NotImplementedError(msg) # TODO: Extract a slice of the cube. self.src = src if src.mesh is None: grid_dims = _get_grid_dims(src) else: - grid_dims = (src.mesh_dim(),) + msg = "Partition does not yet support source meshes." + raise NotImplementedError(msg) + # TODO: This is currently blocked by the fact that slicing an Iris cube on its mesh dimension + # does not currently yield another cube with a mesh. When this is fixed, the following + # code can be uncommented. + # grid_dims = (src.mesh_dim(),) shape = tuple(src.shape[i] for i in grid_dims) self.tgt = tgt self.scheme = scheme @@ -235,7 +245,7 @@ def apply_regridders(self, cube, allow_incomplete=False): current_weights = None files = self.saved_files - for file, chunk in zip(self.file_names, self.src_blocks): + for file, chunk in zip(self.file_names, self.src_blocks, strict=True): if file in files: next_regridder = load_regridder(file, allow_partial=True) cube_chunk = _get_chunk(cube, chunk) @@ -251,9 +261,13 @@ def apply_regridders(self, cube, allow_incomplete=False): else: current_result += next_result + # NOTE: the final "finish_regridding" operation could be performed using any one + # of the partial regridders,but the correct "corresponding" slice of the source + # must be passed. + # See :meth:`~esmf_regrid.experimental._partial.PartialRegridder.finish_regridding`. return next_regridder.finish_regridding( - cube_chunk, + cube_chunk, # matches *this* partial regridder current_weights, current_result, - extra, + extra, # should be *the same* for all the partial results ) diff --git a/src/esmf_regrid/tests/unit/experimental/partition/test_Partition.py b/src/esmf_regrid/tests/unit/experimental/partition/test_Partition.py index 0d72b457..d109091c 100644 --- a/src/esmf_regrid/tests/unit/experimental/partition/test_Partition.py +++ b/src/esmf_regrid/tests/unit/experimental/partition/test_Partition.py @@ -155,6 +155,26 @@ def test_Partition_curv_src(tmp_path): assert result == expected +def test_Partition_mesh_tgt(tmp_path): + """Test Partition class when the target has a mesh.""" + src = _grid_cube(150, 500, (-180, 180), (-90, 90), circular=True) + src.data = np.arange(150 * 500).reshape([500, 150]) + tgt = _gridlike_mesh_cube(16, 36) + + files = [tmp_path / f"partial_{x}.nc" for x in range(5)] + scheme = ESMFAreaWeighted(mdtol=1) + + src_chunks = (100, 150) + partition = Partition(src, tgt, scheme, files, src_chunks=src_chunks) + + partition.generate_files() + + result = partition.apply_regridders(src) + expected = src.regrid(tgt, scheme) + assert np.allclose(result.data, expected.data) + assert result == expected + + def test_conflicting_chunks(tmp_path): """Test error handling of Partition class.""" src = _grid_cube(150, 500, (-180, 180), (-90, 90), circular=True) From 9e5c6aaab3186a8e1c6b161a1d4c6f89c4ebe023 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Wed, 4 Feb 2026 16:22:50 +0000 Subject: [PATCH 32/41] ruff fixes --- src/esmf_regrid/experimental/_partial.py | 6 +++--- src/esmf_regrid/experimental/partition.py | 14 +++++++------- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/src/esmf_regrid/experimental/_partial.py b/src/esmf_regrid/experimental/_partial.py index 381d245a..3c19219d 100644 --- a/src/esmf_regrid/experimental/_partial.py +++ b/src/esmf_regrid/experimental/_partial.py @@ -3,10 +3,10 @@ import numpy as np from esmf_regrid.schemes import ( - _create_cube, - _ESMFRegridder, GridRecord, MeshRecord, + _create_cube, + _ESMFRegridder, ) @@ -109,7 +109,7 @@ def finish_regridding(self, src_cube, weights, data, extra): out_dims = 1 else: msg = "Unrecognised target information." - raise ValueError(msg) + raise TypeError(msg) result_cube = _create_cube( result_data, src_cube, old_dims, tgt_coords, out_dims diff --git a/src/esmf_regrid/experimental/partition.py b/src/esmf_regrid/experimental/partition.py index 2c949f88..e290d010 100644 --- a/src/esmf_regrid/experimental/partition.py +++ b/src/esmf_regrid/experimental/partition.py @@ -14,7 +14,7 @@ def _get_chunk(cube, sl): else: grid_dims = (cube.mesh_dim(),) full_slice = [np.s_[:]] * len(cube.shape) - for s, d in zip(sl, grid_dims): + for s, d in zip(sl, grid_dims, strict=True): full_slice[d] = np.s_[s[0] : s[1]] return cube[*full_slice] @@ -56,16 +56,16 @@ def _determine_blocks(shape, chunks, num_chunks, explicit_blocks): raise ValueError(msg) # TODO: This is currently blocked by the fact that slicing an Iris cube on its mesh dimension # does not currently yield another cube with a mesh. When this is fixed, the following - # code can be uncommented. + # code can be uncommented and the noqa on the following line can be removed. # explicit_blocks = [ # [[int(lower), int(upper)]] - # for lower, upper in zip(bounds[0][:-1], bounds[0][1:]) + # for lower, upper in zip(bounds[0][:-1], bounds[0][1:], strict=True) # ] - elif len(bounds) == 2: + elif len(bounds) == 2: # noqa: RET506 explicit_blocks = [ [[int(ly), int(uy)], [int(lx), int(ux)]] - for ly, uy in zip(bounds[0][:-1], bounds[0][1:]) - for lx, ux in zip(bounds[1][:-1], bounds[1][1:]) + for ly, uy in zip(bounds[0][:-1], bounds[0][1:], strict=True) + for lx, ux in zip(bounds[1][:-1], bounds[1][1:], strict=True) ] else: msg = "Chunks must not exceed two dimensions." @@ -170,7 +170,7 @@ def __init__( raise NotImplementedError(msg) # Note: this may need to become more sophisticated when both src and tgt are large - self.file_block_dict = dict(zip(self.file_names, self.src_blocks)) + self.file_block_dict = dict(zip(self.file_names, self.src_blocks, strict=True)) if saved_files is None: self.saved_files = [] From 32e5bafd1267c28cfa38a0631f4b56f59c424627 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Thu, 5 Feb 2026 11:26:37 +0000 Subject: [PATCH 33/41] add bilinear regridding support --- src/esmf_regrid/experimental/partition.py | 9 +++++- .../experimental/partition/test_Partition.py | 29 ++++++++++++++++++- 2 files changed, 36 insertions(+), 2 deletions(-) diff --git a/src/esmf_regrid/experimental/partition.py b/src/esmf_regrid/experimental/partition.py index e290d010..9548b2a1 100644 --- a/src/esmf_regrid/experimental/partition.py +++ b/src/esmf_regrid/experimental/partition.py @@ -1,11 +1,12 @@ """Provides an interface for splitting up a large regridding task.""" - +import esmpy import numpy as np from esmf_regrid.constants import Constants from esmf_regrid.experimental._partial import PartialRegridder from esmf_regrid.experimental.io import load_regridder, save_regridder from esmf_regrid.schemes import _get_grid_dims +from pyarrow import dictionary def _get_chunk(cube, sl): @@ -132,6 +133,12 @@ def __init__( if scheme._method == Constants.Method.NEAREST: msg = "The `Nearest` method is not implemented." raise NotImplementedError(msg) + if scheme._method == Constants.Method.BILINEAR: + pole_method = scheme.esmf_args.get("pole_method") + if pole_method != esmpy.PoleMethod.NONE: + msg = ("Bilinear regridding must have a `pole_method` of `esmpy.PoleMethod.NONE` in " + "the `esmf_args` in order for Partition to work.`") + raise ValueError(msg) # TODO: Extract a slice of the cube. self.src = src if src.mesh is None: diff --git a/src/esmf_regrid/tests/unit/experimental/partition/test_Partition.py b/src/esmf_regrid/tests/unit/experimental/partition/test_Partition.py index d109091c..ed42231e 100644 --- a/src/esmf_regrid/tests/unit/experimental/partition/test_Partition.py +++ b/src/esmf_regrid/tests/unit/experimental/partition/test_Partition.py @@ -1,10 +1,11 @@ """Unit tests for :mod:`esmf_regrid.experimental.partition`.""" import dask.array as da +import esmpy import numpy as np import pytest -from esmf_regrid import ESMFAreaWeighted, ESMFNearest +from esmf_regrid import ESMFAreaWeighted, ESMFBilinear, ESMFNearest from esmf_regrid.experimental.partition import Partition from esmf_regrid.tests.unit.schemes.test__cube_to_GridInfo import ( _curvilinear_cube, @@ -155,6 +156,32 @@ def test_Partition_curv_src(tmp_path): assert result == expected +def test_Partition_bilinear(tmp_path): + """Test Partition class for bilinear regridding.""" + src = _grid_cube(150, 500, (-180, 180), (-90, 90), circular=True) + src.data = np.arange(150 * 500).reshape([500, 150]) + tgt = _grid_cube(16, 36, (-180, 180), (-90, 90), circular=True) + + files = [tmp_path / f"partial_{x}.nc" for x in range(5)] + src_chunks = (100, 150) + + bad_scheme = ESMFBilinear() + with pytest.raises(ValueError): + _ = Partition(src, tgt, bad_scheme, files, src_chunks=src_chunks) + + # The pole_method must be NONE for bilinear regridding partitions to work. + scheme = ESMFBilinear(esmf_args={"pole_method": esmpy.PoleMethod.NONE}) + + partition = Partition(src, tgt, scheme, files, src_chunks=src_chunks) + + partition.generate_files() + + result = partition.apply_regridders(src) + expected = src.regrid(tgt, scheme) + assert np.allclose(result.data, expected.data) + assert result == expected + + def test_Partition_mesh_tgt(tmp_path): """Test Partition class when the target has a mesh.""" src = _grid_cube(150, 500, (-180, 180), (-90, 90), circular=True) From 21a4d54a83cf4066b21804c4cc149aa924f02c0a Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Thu, 5 Feb 2026 11:58:46 +0000 Subject: [PATCH 34/41] address review comments --- src/esmf_regrid/experimental/io.py | 1 - src/esmf_regrid/experimental/partition.py | 9 ++++++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/esmf_regrid/experimental/io.py b/src/esmf_regrid/experimental/io.py index 56c5958b..5d2f04ee 100644 --- a/src/esmf_regrid/experimental/io.py +++ b/src/esmf_regrid/experimental/io.py @@ -461,7 +461,6 @@ def load_regridder(filename, allow_partial=False): sub_scheme = { Constants.Method.CONSERVATIVE: ESMFAreaWeighted, Constants.Method.BILINEAR: ESMFBilinear, - Constants.Method.NEAREST: ESMFNearest, }[method] regridder = scheme( src_cube, diff --git a/src/esmf_regrid/experimental/partition.py b/src/esmf_regrid/experimental/partition.py index 9548b2a1..9583eb42 100644 --- a/src/esmf_regrid/experimental/partition.py +++ b/src/esmf_regrid/experimental/partition.py @@ -6,7 +6,6 @@ from esmf_regrid.experimental._partial import PartialRegridder from esmf_regrid.experimental.io import load_regridder, save_regridder from esmf_regrid.schemes import _get_grid_dims -from pyarrow import dictionary def _get_chunk(cube, sl): @@ -124,7 +123,10 @@ def __init__( in y,x axis order. Each integer describes the number of blocks that dimension will be divided into. explicit_src_blocks : arraylike NxMx2 - Explicitly specify the bounds of each block in the partition. + Explicitly specify the bounds of each block in the partition. Describes N blocks + along M dimensions with a pair of upper and lower bounds. The upper and lower bounds + describe a slice of an array, e.g. the bounds (3, 6) describe the indices 3, 4, 5 in + a particular dimension. auto_generate : bool, default=False When true, start generating files on initialisation. saved_files : iterable of str @@ -157,7 +159,8 @@ def __init__( self.file_names = file_names if use_dask_src_chunks: if src_chunks is not None: - msg = "Potentially conflicting partition block definitions." + msg = ("`src_chunks` and `use_dask_src_chunks` may provide conflicting" + "partition block definitions.") raise ValueError(msg) if not src.has_lazy_data(): msg = "If `use_dask_src_chunks=True`, the source cube must be lazy." From ef2883a0f5e862d2188a8777193a364ab6d3cba4 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Thu, 5 Feb 2026 12:32:31 +0000 Subject: [PATCH 35/41] ruff format --- src/esmf_regrid/experimental/partition.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/src/esmf_regrid/experimental/partition.py b/src/esmf_regrid/experimental/partition.py index 9583eb42..cf71aefc 100644 --- a/src/esmf_regrid/experimental/partition.py +++ b/src/esmf_regrid/experimental/partition.py @@ -1,4 +1,5 @@ """Provides an interface for splitting up a large regridding task.""" + import esmpy import numpy as np @@ -138,8 +139,10 @@ def __init__( if scheme._method == Constants.Method.BILINEAR: pole_method = scheme.esmf_args.get("pole_method") if pole_method != esmpy.PoleMethod.NONE: - msg = ("Bilinear regridding must have a `pole_method` of `esmpy.PoleMethod.NONE` in " - "the `esmf_args` in order for Partition to work.`") + msg = ( + "Bilinear regridding must have a `pole_method` of `esmpy.PoleMethod.NONE` in " + "the `esmf_args` in order for Partition to work.`" + ) raise ValueError(msg) # TODO: Extract a slice of the cube. self.src = src @@ -159,8 +162,10 @@ def __init__( self.file_names = file_names if use_dask_src_chunks: if src_chunks is not None: - msg = ("`src_chunks` and `use_dask_src_chunks` may provide conflicting" - "partition block definitions.") + msg = ( + "`src_chunks` and `use_dask_src_chunks` may provide conflicting" + "partition block definitions." + ) raise ValueError(msg) if not src.has_lazy_data(): msg = "If `use_dask_src_chunks=True`, the source cube must be lazy." From 598d2d33f68617537de95e2b91e9e21e6cec70c1 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Thu, 5 Feb 2026 13:42:09 +0000 Subject: [PATCH 36/41] more sensible PartialRegridder kwarg management --- pyproject.toml | 2 +- src/esmf_regrid/experimental/_partial.py | 6 ++---- src/esmf_regrid/experimental/io.py | 6 +++++- 3 files changed, 8 insertions(+), 6 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 92a2cd3e..b21faca0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -298,5 +298,5 @@ convention = "numpy" [tool.ruff.lint.pylint] # TODO: refactor to reduce complexity, if possible max-args = 10 -max-branches = 22 +max-branches = 23 max-statements = 110 diff --git a/src/esmf_regrid/experimental/_partial.py b/src/esmf_regrid/experimental/_partial.py index 3c19219d..ddf451e2 100644 --- a/src/esmf_regrid/experimental/_partial.py +++ b/src/esmf_regrid/experimental/_partial.py @@ -32,15 +32,13 @@ def __init__(self, src, tgt, src_slice, tgt_slice, weights, scheme, **kwargs): 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 - # Pop duplicate kwargs. - for arg in set(kwargs.keys()).intersection(vars(self.scheme)): - kwargs.pop(arg) - self._regridder = scheme.regridder( src, tgt, diff --git a/src/esmf_regrid/experimental/io.py b/src/esmf_regrid/experimental/io.py index 5d2f04ee..8e532304 100644 --- a/src/esmf_regrid/experimental/io.py +++ b/src/esmf_regrid/experimental/io.py @@ -462,13 +462,17 @@ def load_regridder(filename, allow_partial=False): Constants.Method.CONSERVATIVE: ESMFAreaWeighted, Constants.Method.BILINEAR: ESMFBilinear, }[method] + mdtol = kwargs.pop(_MDTOL, None) + mdtol_dict = {} + if mdtol is not None: + mdtol_dict[_MDTOL] = mdtol regridder = scheme( src_cube, tgt_cube, src_slice, tgt_slice, weight_matrix, - sub_scheme(), + sub_scheme(**mdtol_dict), **kwargs, ) else: From 2c53d31472b407a01080c17f688898572d1ead75 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Thu, 5 Feb 2026 15:46:17 +0000 Subject: [PATCH 37/41] make round tripping more thorough --- pyproject.toml | 2 +- src/esmf_regrid/experimental/io.py | 22 +++++++++++----- .../partition/test_PartialRegridder.py | 26 +++++++++++++++++++ 3 files changed, 42 insertions(+), 8 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b21faca0..10ada6d2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -298,5 +298,5 @@ convention = "numpy" [tool.ruff.lint.pylint] # TODO: refactor to reduce complexity, if possible max-args = 10 -max-branches = 23 +max-branches = 25 max-statements = 110 diff --git a/src/esmf_regrid/experimental/io.py b/src/esmf_regrid/experimental/io.py index 8e532304..62646e01 100644 --- a/src/esmf_regrid/experimental/io.py +++ b/src/esmf_regrid/experimental/io.py @@ -20,7 +20,6 @@ ESMFAreaWeightedRegridder, ESMFBilinear, ESMFBilinearRegridder, - ESMFNearest, ESMFNearestRegridder, GridRecord, MeshRecord, @@ -76,6 +75,7 @@ "extrap_method": _EXTRAP_METHOD_DICT, "unmapped_action": _UNMAPPED_ACTION_DICT, } +_ESMF_BOOL_ARGS = ["ignore_degenerate", "large_file"] def _add_mask_to_cube(mask, cube, name): @@ -294,7 +294,7 @@ def save_regridder(rg, filename, allow_partial=False): if tgt_slice is None: tgt_slice = [] tgt_slice_cube = Cube( - src_slice, long_name=_TGT_SLICE_NAME, var_name=_TGT_SLICE_NAME + tgt_slice, long_name=_TGT_SLICE_NAME, var_name=_TGT_SLICE_NAME ) extra_cubes = [src_slice_cube, tgt_slice_cube] @@ -416,11 +416,11 @@ def load_regridder(filename, allow_partial=False): mdtol = weights_cube.attributes[_MDTOL] if src_cube.coords(_SOURCE_MASK_NAME): - use_src_mask = src_cube.coord(_SOURCE_MASK_NAME).points + use_src_mask = src_cube.coord(_SOURCE_MASK_NAME).points.astype(bool) else: use_src_mask = False if tgt_cube.coords(_TARGET_MASK_NAME): - use_tgt_mask = tgt_cube.coord(_TARGET_MASK_NAME).points + use_tgt_mask = tgt_cube.coord(_TARGET_MASK_NAME).points.astype(bool) else: use_tgt_mask = False @@ -433,6 +433,9 @@ def load_regridder(filename, allow_partial=False): for arg, arg_dict in _ESMF_ENUM_ARGS.items(): if arg in esmf_args: esmf_args[arg] = arg_dict[esmf_args[arg]] + for arg in _ESMF_BOOL_ARGS: + if arg in esmf_args: + esmf_args[arg] = bool(esmf_args[arg]) if scheme is GridToMeshESMFRegridder: resolution_keyword = _SOURCE_RESOLUTION @@ -463,16 +466,21 @@ def load_regridder(filename, allow_partial=False): Constants.Method.BILINEAR: ESMFBilinear, }[method] mdtol = kwargs.pop(_MDTOL, None) - mdtol_dict = {} + sub_kwargs = {} if mdtol is not None: - mdtol_dict[_MDTOL] = mdtol + sub_kwargs[_MDTOL] = mdtol regridder = scheme( src_cube, tgt_cube, src_slice, tgt_slice, weight_matrix, - sub_scheme(**mdtol_dict), + sub_scheme( + use_src_mask=use_src_mask, + use_tgt_mask=use_tgt_mask, + esmf_args=esmf_args, + **sub_kwargs, + ), **kwargs, ) else: diff --git a/src/esmf_regrid/tests/unit/experimental/partition/test_PartialRegridder.py b/src/esmf_regrid/tests/unit/experimental/partition/test_PartialRegridder.py index 37ad2ecc..09f95c65 100644 --- a/src/esmf_regrid/tests/unit/experimental/partition/test_PartialRegridder.py +++ b/src/esmf_regrid/tests/unit/experimental/partition/test_PartialRegridder.py @@ -1,7 +1,10 @@ """Unit tests for :mod:`esmf_regrid.experimental.partition`.""" +import numpy as np + from esmf_regrid import ESMFAreaWeighted from esmf_regrid.experimental._partial import PartialRegridder +from esmf_regrid.experimental.io import load_regridder, save_regridder from esmf_regrid.tests.unit.schemes.test__cube_to_GridInfo import ( _grid_cube, ) @@ -23,3 +26,26 @@ def test_PartialRegridder_repr(): "scheme=ESMFAreaWeighted(mdtol=0.5, use_src_mask=False, use_tgt_mask=False, esmf_args={}))" ) assert repr(pr) == expected_repr + + +def test_PartialRegridder_roundtrip(tmp_path): + """Test load/save for PartialRegridder instance.""" + src = _grid_cube(10, 15, (-180, 180), (-90, 90), circular=True) + mask = np.zeros_like(src.data) + mask[0, 0] = 1 + src.data = np.ma.array(src.data, mask=mask) + tgt = _grid_cube(5, 10, (-180, 180), (-90, 90), circular=True) + src_slice = [[10, 20], [15, 30]] + tgt_slice = [[0, 5], [0, 10]] + weights = None + scheme = ESMFAreaWeighted( + mdtol=0.5, use_src_mask=src.data.mask, esmf_args={"ignore_degenerate": True} + ) + + pr = PartialRegridder(src, tgt, src_slice, tgt_slice, weights, scheme) + file = tmp_path / "partial.nc" + + save_regridder(pr, file, allow_partial=True) + loaded_pr = load_regridder(file, allow_partial=True) + + assert repr(loaded_pr) == repr(pr) From 1c1f513dc4fee0338018be985811f0cb9967b31f Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Sun, 8 Feb 2026 23:45:59 +0000 Subject: [PATCH 38/41] address review comments --- src/esmf_regrid/experimental/_partial.py | 12 ++++++++++-- .../experimental/partition/test_PartialRegridder.py | 9 +++++++-- 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/src/esmf_regrid/experimental/_partial.py b/src/esmf_regrid/experimental/_partial.py index ddf451e2..7868ada6 100644 --- a/src/esmf_regrid/experimental/_partial.py +++ b/src/esmf_regrid/experimental/_partial.py @@ -24,10 +24,16 @@ def __init__(self, src, tgt, src_slice, tgt_slice, weights, scheme, **kwargs): 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. + ``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. + ``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` @@ -51,6 +57,8 @@ 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})" diff --git a/src/esmf_regrid/tests/unit/experimental/partition/test_PartialRegridder.py b/src/esmf_regrid/tests/unit/experimental/partition/test_PartialRegridder.py index 09f95c65..d418f66f 100644 --- a/src/esmf_regrid/tests/unit/experimental/partition/test_PartialRegridder.py +++ b/src/esmf_regrid/tests/unit/experimental/partition/test_PartialRegridder.py @@ -22,8 +22,13 @@ def test_PartialRegridder_repr(): pr = PartialRegridder(src, tgt, src_slice, tgt_slice, weights, scheme) expected_repr = ( - "PartialRegridder(src_slice=((10, 20), (15, 30)), tgt_slice=((0, 5), (0, 10)), " - "scheme=ESMFAreaWeighted(mdtol=0.5, use_src_mask=False, use_tgt_mask=False, esmf_args={}))" + "PartialRegridder(src=GridRecord(" + "grid_x=, " + "grid_y=), " + "tgt_slice=GridRecord(grid_x=, " + "grid_y=), " + "src_slice=((10, 20), (15, 30)), tgt_slice=((0, 5), (0, 10)), scheme=ESMFAreaWeighted(mdtol=0.5, " + "use_src_mask=False, use_tgt_mask=False, esmf_args={}))" ) assert repr(pr) == expected_repr From 10f12395c6d719e68e8318311a27d0ce4e432ecb Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Mon, 9 Feb 2026 12:23:48 +0000 Subject: [PATCH 39/41] address review comments --- src/esmf_regrid/experimental/_partial.py | 1 + src/esmf_regrid/experimental/partition.py | 4 ++++ 2 files changed, 5 insertions(+) diff --git a/src/esmf_regrid/experimental/_partial.py b/src/esmf_regrid/experimental/_partial.py index 7868ada6..5d379ea1 100644 --- a/src/esmf_regrid/experimental/_partial.py +++ b/src/esmf_regrid/experimental/_partial.py @@ -101,6 +101,7 @@ def finish_regridding(self, src_cube, weights, data, extra): if num_dims == 2 and num_out_dims == 1: new_dims = [min(old_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 = [old_dims[0] + 1, old_dims[0]] else: new_dims = old_dims diff --git a/src/esmf_regrid/experimental/partition.py b/src/esmf_regrid/experimental/partition.py index cf71aefc..59e78352 100644 --- a/src/esmf_regrid/experimental/partition.py +++ b/src/esmf_regrid/experimental/partition.py @@ -97,6 +97,10 @@ def __init__( Note ---- + The source is partitioned into blocks using one of the four mutually exclusive arguments, + `use_dask_src_chunks`, `src_chunks`, `num_src_chunks`, or `explicit_src_blocks`. These + describe a partition into a numnber of blocks which must equal the number of `file_names`. + Currently, it is only possible to divide the source grid into chunks. Meshes are not yet supported as a source. From bfbea29bdb2ca33aa041a8fa82ce8b4ad6d512ac Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Mon, 9 Feb 2026 12:40:54 +0000 Subject: [PATCH 40/41] fix typo --- src/esmf_regrid/experimental/partition.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/esmf_regrid/experimental/partition.py b/src/esmf_regrid/experimental/partition.py index 59e78352..7a564458 100644 --- a/src/esmf_regrid/experimental/partition.py +++ b/src/esmf_regrid/experimental/partition.py @@ -99,7 +99,7 @@ def __init__( ---- The source is partitioned into blocks using one of the four mutually exclusive arguments, `use_dask_src_chunks`, `src_chunks`, `num_src_chunks`, or `explicit_src_blocks`. These - describe a partition into a numnber of blocks which must equal the number of `file_names`. + describe a partition into a number of blocks which must equal the number of `file_names`. Currently, it is only possible to divide the source grid into chunks. Meshes are not yet supported as a source. From 503a3098535ad33f19535873ce51c850926da2c8 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Mon, 9 Feb 2026 16:28:39 +0000 Subject: [PATCH 41/41] address review comments --- src/esmf_regrid/experimental/_partial.py | 12 ++++++------ src/esmf_regrid/experimental/partition.py | 5 +---- 2 files changed, 7 insertions(+), 10 deletions(-) diff --git a/src/esmf_regrid/experimental/_partial.py b/src/esmf_regrid/experimental/_partial.py index 5d379ea1..ffaf8708 100644 --- a/src/esmf_regrid/experimental/_partial.py +++ b/src/esmf_regrid/experimental/_partial.py @@ -91,20 +91,20 @@ def finish_regridding(self, src_cube, weights, data, extra): of the 'src' reference cube provided in regridder creation (`self._src`). So, it must be the correct "corresponding slice" of the source cube. """ - old_dims = self._get_cube_dims(src_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(old_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(old_dims)] + 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 = [old_dims[0] + 1, old_dims[0]] + new_dims = [src_dims[0] + 1, src_dims[0]] else: - new_dims = old_dims + new_dims = src_dims result_data = np.moveaxis(result_data, standard_out_dims, new_dims) @@ -119,6 +119,6 @@ def finish_regridding(self, src_cube, weights, data, extra): raise TypeError(msg) result_cube = _create_cube( - result_data, src_cube, old_dims, tgt_coords, out_dims + result_data, src_cube, src_dims, tgt_coords, out_dims ) return result_cube diff --git a/src/esmf_regrid/experimental/partition.py b/src/esmf_regrid/experimental/partition.py index 7a564458..6513fa3a 100644 --- a/src/esmf_regrid/experimental/partition.py +++ b/src/esmf_regrid/experimental/partition.py @@ -166,10 +166,7 @@ def __init__( self.file_names = file_names if use_dask_src_chunks: if src_chunks is not None: - msg = ( - "`src_chunks` and `use_dask_src_chunks` may provide conflicting" - "partition block definitions." - ) + msg = "`src_chunks` and `use_dask_src_chunks` cannot be used at the same time." raise ValueError(msg) if not src.has_lazy_data(): msg = "If `use_dask_src_chunks=True`, the source cube must be lazy."