diff --git a/.cirrus.yml b/.cirrus.yml index 61f50740..6eae1cf1 100644 --- a/.cirrus.yml +++ b/.cirrus.yml @@ -22,9 +22,9 @@ env: # Maximum cache period (in weeks) before forcing a new cache upload. CACHE_PERIOD: "2" # Increment the build number to force new conda cache upload. - CONDA_CACHE_BUILD: "0" + CONDA_CACHE_BUILD: "1" # Increment the build number to force new nox cache upload. - NOX_CACHE_BUILD: "0" + NOX_CACHE_BUILD: "2" # Increment the build number to force new pip cache upload. PIP_CACHE_BUILD: "0" # Pip package to be installed. diff --git a/esmf_regrid/experimental/io.py b/esmf_regrid/experimental/io.py new file mode 100644 index 00000000..bb13dadf --- /dev/null +++ b/esmf_regrid/experimental/io.py @@ -0,0 +1,149 @@ +"""Provides load/save functions for regridders.""" + +import iris +from iris.coords import AuxCoord +from iris.cube import Cube, CubeList +from iris.experimental.ugrid import PARSE_UGRID_ON_LOAD +import numpy as np +import scipy.sparse + +from esmf_regrid.experimental.unstructured_scheme import ( + GridToMeshESMFRegridder, + MeshToGridESMFRegridder, +) + + +SUPPORTED_REGRIDDERS = [ + GridToMeshESMFRegridder, + MeshToGridESMFRegridder, +] +REGRIDDER_NAME_MAP = {rg_class.__name__: rg_class for rg_class in SUPPORTED_REGRIDDERS} + + +def save_regridder(rg, filename): + """ + Save a regridder scheme instance. + + Saves either a `GridToMeshESMFRegridder` or a `MeshToGridESMFRegridder`. + + Parameters + ---------- + rg : GridToMeshESMFRegridder, MeshToGridESMFRegridder + The regridder instance to save. + filename : str + The file name to save to. + """ + src_name = "regridder source field" + tgt_name = "regridder target field" + regridder_type = rg.__class__.__name__ + if regridder_type == "GridToMeshESMFRegridder": + src_grid = (rg.grid_y, rg.grid_x) + src_shape = [len(coord.points) for coord in src_grid] + src_data = np.zeros(src_shape) + src_cube = Cube(src_data, long_name=src_name) + src_cube.add_dim_coord(src_grid[0], 0) + src_cube.add_dim_coord(src_grid[1], 1) + + tgt_mesh = rg.mesh + tgt_data = np.zeros(tgt_mesh.face_node_connectivity.indices.shape[0]) + tgt_cube = Cube(tgt_data, long_name=tgt_name) + for coord in tgt_mesh.to_MeshCoords("face"): + tgt_cube.add_aux_coord(coord, 0) + elif regridder_type == "MeshToGridESMFRegridder": + src_mesh = rg.mesh + src_data = np.zeros(src_mesh.face_node_connectivity.indices.shape[0]) + src_cube = Cube(src_data, long_name=src_name) + for coord in src_mesh.to_MeshCoords("face"): + src_cube.add_aux_coord(coord, 0) + + tgt_grid = (rg.grid_y, rg.grid_x) + tgt_shape = [len(coord.points) for coord in tgt_grid] + tgt_data = np.zeros(tgt_shape) + tgt_cube = Cube(tgt_data, long_name=tgt_name) + tgt_cube.add_dim_coord(tgt_grid[0], 0) + tgt_cube.add_dim_coord(tgt_grid[1], 1) + else: + msg = ( + f"Expected a regridder of type `GridToMeshESMFRegridder` or " + f"`MeshToGridESMFRegridder`, got type {regridder_type}." + ) + raise TypeError(msg) + + metadata_name = "regridder weights and metadata" + + weight_matrix = rg.regridder.weight_matrix + reformatted_weight_matrix = weight_matrix.tocoo() + weight_data = reformatted_weight_matrix.data + weight_rows = reformatted_weight_matrix.row + weight_cols = reformatted_weight_matrix.col + weight_shape = reformatted_weight_matrix.shape + + mdtol = rg.mdtol + attributes = { + "regridder type": regridder_type, + "mdtol": mdtol, + "weights shape": weight_shape, + } + + metadata_cube = Cube(weight_data, long_name=metadata_name, attributes=attributes) + row_name = "weight matrix rows" + row_coord = AuxCoord(weight_rows, long_name=row_name) + col_name = "weight matrix columns" + col_coord = AuxCoord(weight_cols, long_name=col_name) + metadata_cube.add_aux_coord(row_coord, 0) + metadata_cube.add_aux_coord(col_coord, 0) + + # Avoid saving bug by placing the mesh cube second. + # TODO: simplify this when this bug is fixed in iris. + if regridder_type == "GridToMeshESMFRegridder": + cube_list = CubeList([src_cube, tgt_cube, metadata_cube]) + elif regridder_type == "MeshToGridESMFRegridder": + cube_list = CubeList([tgt_cube, src_cube, metadata_cube]) + iris.fileformats.netcdf.save(cube_list, filename) + + +def load_regridder(filename): + """ + Load a regridder scheme instance. + + Loads either a `GridToMeshESMFRegridder` or a `MeshToGridESMFRegridder`. + + Parameters + ---------- + filename : str + The file name to load from. + """ + with PARSE_UGRID_ON_LOAD.context(): + cubes = iris.load(filename) + + src_name = "regridder source field" + tgt_name = "regridder target field" + metadata_name = "regridder weights and metadata" + + # Extract the source, target and metadata information. + src_cube = cubes.extract_cube(src_name) + tgt_cube = cubes.extract_cube(tgt_name) + metadata_cube = cubes.extract_cube(metadata_name) + + # Determine the regridder type. + regridder_type = metadata_cube.attributes["regridder type"] + assert regridder_type in REGRIDDER_NAME_MAP.keys() + scheme = REGRIDDER_NAME_MAP[regridder_type] + + # Reconstruct the weight matrix. + weight_data = metadata_cube.data + row_name = "weight matrix rows" + weight_rows = metadata_cube.coord(row_name).points + col_name = "weight matrix columns" + weight_cols = metadata_cube.coord(col_name).points + weight_shape = metadata_cube.attributes["weights shape"] + weight_matrix = scipy.sparse.csr_matrix( + (weight_data, (weight_rows, weight_cols)), shape=weight_shape + ) + + mdtol = metadata_cube.attributes["mdtol"] + + regridder = scheme( + src_cube, tgt_cube, mdtol=mdtol, precomputed_weights=weight_matrix + ) + return regridder diff --git a/esmf_regrid/experimental/unstructured_scheme.py b/esmf_regrid/experimental/unstructured_scheme.py index 7bce8d33..e197f613 100644 --- a/esmf_regrid/experimental/unstructured_scheme.py +++ b/esmf_regrid/experimental/unstructured_scheme.py @@ -228,7 +228,9 @@ def copy_coords(src_coords, add_method): return new_cube -def _regrid_unstructured_to_rectilinear__prepare(src_mesh_cube, target_grid_cube): +def _regrid_unstructured_to_rectilinear__prepare( + src_mesh_cube, target_grid_cube, precomputed_weights=None +): """ First (setup) part of 'regrid_unstructured_to_rectilinear'. @@ -257,7 +259,7 @@ def _regrid_unstructured_to_rectilinear__prepare(src_mesh_cube, target_grid_cube meshinfo = _mesh_to_MeshInfo(mesh) gridinfo = _cube_to_GridInfo(target_grid_cube) - regridder = Regridder(meshinfo, gridinfo) + regridder = Regridder(meshinfo, gridinfo, precomputed_weights) regrid_info = (mesh_dim, grid_x, grid_y, regridder) @@ -350,7 +352,9 @@ def regrid_unstructured_to_rectilinear(src_cube, grid_cube, mdtol=0): class MeshToGridESMFRegridder: """Regridder class for unstructured to rectilinear cubes.""" - def __init__(self, src_mesh_cube, target_grid_cube, mdtol=1): + def __init__( + self, src_mesh_cube, target_grid_cube, mdtol=1, precomputed_weights=None + ): """ Create regridder for conversions between source mesh and target grid. @@ -382,9 +386,12 @@ def __init__(self, src_mesh_cube, target_grid_cube, mdtol=1): self.mdtol = mdtol partial_regrid_info = _regrid_unstructured_to_rectilinear__prepare( - src_mesh_cube, target_grid_cube + src_mesh_cube, target_grid_cube, precomputed_weights=precomputed_weights ) + # Record source mesh. + self.mesh = src_mesh_cube.mesh + # Store regrid info. _, self.grid_x, self.grid_y, self.regridder = partial_regrid_info @@ -491,7 +498,9 @@ def copy_coords(src_coords, add_method): return new_cube -def _regrid_rectilinear_to_unstructured__prepare(src_grid_cube, target_mesh_cube): +def _regrid_rectilinear_to_unstructured__prepare( + src_grid_cube, target_mesh_cube, precomputed_weights=None +): """ First (setup) part of 'regrid_rectilinear_to_unstructured'. @@ -510,7 +519,7 @@ def _regrid_rectilinear_to_unstructured__prepare(src_grid_cube, target_mesh_cube meshinfo = _mesh_to_MeshInfo(mesh) gridinfo = _cube_to_GridInfo(src_grid_cube) - regridder = Regridder(gridinfo, meshinfo) + regridder = Regridder(gridinfo, meshinfo, precomputed_weights) regrid_info = (grid_x_dim, grid_y_dim, grid_x, grid_y, mesh, regridder) @@ -610,7 +619,9 @@ def regrid_rectilinear_to_unstructured(src_cube, mesh_cube, mdtol=0): class GridToMeshESMFRegridder: """Regridder class for rectilinear to unstructured cubes.""" - def __init__(self, src_mesh_cube, target_grid_cube, mdtol=1): + def __init__( + self, src_mesh_cube, target_grid_cube, mdtol=1, precomputed_weights=None + ): """ Create regridder for conversions between source grid and target mesh. @@ -637,7 +648,7 @@ def __init__(self, src_mesh_cube, target_grid_cube, mdtol=1): self.mdtol = mdtol partial_regrid_info = _regrid_rectilinear_to_unstructured__prepare( - src_mesh_cube, target_grid_cube + src_mesh_cube, target_grid_cube, precomputed_weights=precomputed_weights ) # Store regrid info. diff --git a/esmf_regrid/tests/unit/experimental/io/__init__.py b/esmf_regrid/tests/unit/experimental/io/__init__.py new file mode 100644 index 00000000..d7e457c9 --- /dev/null +++ b/esmf_regrid/tests/unit/experimental/io/__init__.py @@ -0,0 +1 @@ +"""Unit tests for :mod:`esmf_regrid.experimental.io`.""" diff --git a/esmf_regrid/tests/unit/experimental/io/test_round_tripping.py b/esmf_regrid/tests/unit/experimental/io/test_round_tripping.py new file mode 100644 index 00000000..f4a1b2e4 --- /dev/null +++ b/esmf_regrid/tests/unit/experimental/io/test_round_tripping.py @@ -0,0 +1,118 @@ +"""Unit tests for round tripping (saving then loading) with :mod:`esmf_regrid.experimental.io`.""" + +from iris.cube import Cube +import numpy as np +from numpy import ma + +from esmf_regrid.experimental.io import load_regridder, save_regridder +from esmf_regrid.experimental.unstructured_scheme import ( + GridToMeshESMFRegridder, + MeshToGridESMFRegridder, +) +from esmf_regrid.tests.unit.experimental.unstructured_scheme.test__cube_to_GridInfo import ( + _grid_cube, +) +from esmf_regrid.tests.unit.experimental.unstructured_scheme.test__mesh_to_MeshInfo import ( + _gridlike_mesh, +) + + +def _make_grid_to_mesh_regridder(): + src_lons = 3 + src_lats = 4 + tgt_lons = 5 + tgt_lats = 6 + lon_bounds = (-180, 180) + lat_bounds = (-90, 90) + # TODO check that circularity is preserved. + src = _grid_cube(src_lons, src_lats, lon_bounds, lat_bounds, circular=True) + src.coord("longitude").var_name = "longitude" + src.coord("latitude").var_name = "latitude" + mesh = _gridlike_mesh(tgt_lons, tgt_lats) + mesh_coord_x, mesh_coord_y = mesh.to_MeshCoords("face") + tgt_data = np.zeros(tgt_lons * tgt_lats) + tgt = Cube(tgt_data) + tgt.add_aux_coord(mesh_coord_x, 0) + tgt.add_aux_coord(mesh_coord_y, 0) + + rg = GridToMeshESMFRegridder(src, tgt, mdtol=0.5) + return rg, src + + +def _make_mesh_to_grid_regridder(): + src_lons = 3 + src_lats = 4 + tgt_lons = 5 + tgt_lats = 6 + lon_bounds = (-180, 180) + lat_bounds = (-90, 90) + # TODO check that circularity is preserved. + tgt = _grid_cube(tgt_lons, tgt_lats, lon_bounds, lat_bounds, circular=True) + tgt.coord("longitude").var_name = "longitude" + tgt.coord("latitude").var_name = "latitude" + mesh = _gridlike_mesh(src_lons, src_lats) + mesh_coord_x, mesh_coord_y = mesh.to_MeshCoords("face") + src_data = np.zeros(src_lons * src_lats) + src = Cube(src_data) + src.add_aux_coord(mesh_coord_x, 0) + src.add_aux_coord(mesh_coord_y, 0) + + rg = MeshToGridESMFRegridder(src, tgt, mdtol=0.5) + return rg, src + + +def test_GridToMeshESMFRegridder_round_trip(tmp_path): + """Test save/load round tripping for `GridToMeshESMFRegridder`.""" + original_rg, src = _make_grid_to_mesh_regridder() + filename = tmp_path / "regridder.nc" + save_regridder(original_rg, filename) + loaded_rg = load_regridder(str(filename)) + + assert original_rg.mdtol == loaded_rg.mdtol + assert original_rg.grid_x == loaded_rg.grid_x + assert original_rg.grid_y == loaded_rg.grid_y + # TODO: uncomment when iris mesh comparison becomes available. + # assert original_rg.mesh == loaded_rg.mesh + + # Compare the weight matrices. + original_matrix = original_rg.regridder.weight_matrix + loaded_matrix = loaded_rg.regridder.weight_matrix + # Ensure the original and loaded weight matrix have identical type. + assert type(original_matrix) is type(loaded_matrix) # noqa E721 + assert np.array_equal(original_matrix.todense(), loaded_matrix.todense()) + + # Demonstrate regridding still gives the same results. + src_data = np.arange(np.product(src.data.shape)).reshape(src.data.shape) + src_mask = np.zeros(src.data.shape) + src_mask[0, 0] = 1 + src.data = ma.array(src_data, mask=src_mask) + # TODO: make this a cube comparison when mesh comparison becomes available. + assert np.array_equal(original_rg(src).data, loaded_rg(src).data) + + +def test_MeshToGridESMFRegridder_round_trip(tmp_path): + """Test save/load round tripping for `MeshToGridESMFRegridder`.""" + original_rg, src = _make_mesh_to_grid_regridder() + filename = tmp_path / "regridder.nc" + save_regridder(original_rg, filename) + loaded_rg = load_regridder(str(filename)) + + assert original_rg.mdtol == loaded_rg.mdtol + assert original_rg.grid_x == loaded_rg.grid_x + assert original_rg.grid_y == loaded_rg.grid_y + # TODO: uncomment when iris mesh comparison becomes available. + # assert original_rg.mesh == loaded_rg.mesh + + # Compare the weight matrices. + original_matrix = original_rg.regridder.weight_matrix + loaded_matrix = loaded_rg.regridder.weight_matrix + # Ensure the original and loaded weight matrix have identical type. + assert type(original_matrix) is type(loaded_matrix) # noqa E721 + assert np.array_equal(original_matrix.todense(), loaded_matrix.todense()) + + # Demonstrate regridding still gives the same results. + src_data = np.arange(np.product(src.data.shape)).reshape(src.data.shape) + src_mask = np.zeros(src.data.shape) + src_mask[0] = 1 + src.data = ma.array(src_data, mask=src_mask) + assert original_rg(src) == loaded_rg(src) diff --git a/esmf_regrid/tests/unit/experimental/io/test_save_regridder.py b/esmf_regrid/tests/unit/experimental/io/test_save_regridder.py new file mode 100644 index 00000000..234fb604 --- /dev/null +++ b/esmf_regrid/tests/unit/experimental/io/test_save_regridder.py @@ -0,0 +1,13 @@ +"""Unit tests for :mod:`esmf_regrid.experimental.io.save_regridder`.""" + +import pytest + +from esmf_regrid.experimental.io import save_regridder + + +def test_invalid_type(tmp_path): + """Test that `save_regridder` raises a TypeError where appropriate.""" + invalid_obj = None + filename = tmp_path / "regridder.nc" + with pytest.raises(TypeError): + save_regridder(invalid_obj, filename)