From fcee4ee3db08be6dfec05f0e00479fbbfb048f89 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Mon, 29 Mar 2021 14:27:24 +0100 Subject: [PATCH 1/4] handle extra dims in _create_cube --- .../experimental/unstructured_scheme.py | 31 +++++++++++++++---- 1 file changed, 25 insertions(+), 6 deletions(-) diff --git a/esmf_regrid/experimental/unstructured_scheme.py b/esmf_regrid/experimental/unstructured_scheme.py index 6ef03cd1..50777833 100644 --- a/esmf_regrid/experimental/unstructured_scheme.py +++ b/esmf_regrid/experimental/unstructured_scheme.py @@ -79,10 +79,6 @@ def _create_cube(data, src_cube, mesh_dim, grid_x, grid_y): new_cube = iris.cube.Cube(data) - # TODO: The following code assumes a 1D source cube and mesh_dim = 0. - # This is therefore simple code which should be updated when we start - # supporting the regridding of extra dimensions. - # TODO: The following code is rigid with respect to which dimensions # the x coord and y coord are assigned to. We should decide if it is # appropriate to copy the dimension ordering from the target cube @@ -92,8 +88,31 @@ def _create_cube(data, src_cube, mesh_dim, grid_x, grid_y): new_cube.metadata = copy.deepcopy(src_cube.metadata) - for coord in src_cube.coords(dimensions=()): - new_cube.add_aux_coord(coord.copy()) + coord_mapping = {} + + def copy_coords(src_coords, add_method): + for coord in src_coords: + dims = src_cube.coord_dims(coord) + if hasattr(coord, "mesh") or mesh_dim in dims: + continue + dims = [dim if dim < mesh_dim else dim + 1 for dim in dims] + result_coord = coord.copy() + add_method(result_coord, dims) + coord_mapping[id(coord)] = result_coord + + copy_coords(src_cube.dim_coords, new_cube.add_dim_coord) + copy_coords(src_cube.aux_coords, new_cube.add_aux_coord) + + for factory in src_cube.aux_factories: + # TODO: Regrid dependant coordinates which span mesh_dim. + try: + result.add_aux_factory(factory.updated(coord_mapping)) + except KeyError: + msg = ( + "Cannot update aux_factory {!r} because of dropped" + " coordinates.".format(factory.name()) + ) + warnings.warn(msg) return new_cube From ea8007aa276e92f137b6f5a66ed605e82435745d Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Mon, 29 Mar 2021 17:02:34 +0100 Subject: [PATCH 2/4] add test --- .../experimental/unstructured_scheme.py | 27 +++++----- .../unstructured_scheme/test__create_cube.py | 53 +++++++++++++++++-- 2 files changed, 63 insertions(+), 17 deletions(-) diff --git a/esmf_regrid/experimental/unstructured_scheme.py b/esmf_regrid/experimental/unstructured_scheme.py index 50777833..0a97b6a5 100644 --- a/esmf_regrid/experimental/unstructured_scheme.py +++ b/esmf_regrid/experimental/unstructured_scheme.py @@ -88,7 +88,10 @@ def _create_cube(data, src_cube, mesh_dim, grid_x, grid_y): new_cube.metadata = copy.deepcopy(src_cube.metadata) - coord_mapping = {} + # TODO: Handle derived coordinates. The following code is taken from + # iris, the parts dealing with derived coordinates have been + # commented out for the time being. + # coord_mapping = {} def copy_coords(src_coords, add_method): for coord in src_coords: @@ -98,21 +101,21 @@ def copy_coords(src_coords, add_method): dims = [dim if dim < mesh_dim else dim + 1 for dim in dims] result_coord = coord.copy() add_method(result_coord, dims) - coord_mapping[id(coord)] = result_coord + # coord_mapping[id(coord)] = result_coord copy_coords(src_cube.dim_coords, new_cube.add_dim_coord) copy_coords(src_cube.aux_coords, new_cube.add_aux_coord) - for factory in src_cube.aux_factories: - # TODO: Regrid dependant coordinates which span mesh_dim. - try: - result.add_aux_factory(factory.updated(coord_mapping)) - except KeyError: - msg = ( - "Cannot update aux_factory {!r} because of dropped" - " coordinates.".format(factory.name()) - ) - warnings.warn(msg) + # for factory in src_cube.aux_factories: + # # TODO: Regrid dependant coordinates which span mesh_dim. + # try: + # result.add_aux_factory(factory.updated(coord_mapping)) + # except KeyError: + # msg = ( + # "Cannot update aux_factory {!r} because of dropped" + # " coordinates.".format(factory.name()) + # ) + # warnings.warn(msg) return new_cube diff --git a/esmf_regrid/tests/unit/experimental/unstructured_scheme/test__create_cube.py b/esmf_regrid/tests/unit/experimental/unstructured_scheme/test__create_cube.py index 1d9cc33d..d9994839 100644 --- a/esmf_regrid/tests/unit/experimental/unstructured_scheme/test__create_cube.py +++ b/esmf_regrid/tests/unit/experimental/unstructured_scheme/test__create_cube.py @@ -1,6 +1,8 @@ """Unit tests for miscellaneous helper functions in `esmf_regrid.experimental.unstructured_scheme`.""" import iris +from iris.coords import AuxCoord, DimCoord +from iris.cube import Cube import numpy as np from esmf_regrid.experimental.unstructured_scheme import _create_cube @@ -11,17 +13,55 @@ def test_create_cube_2D(): data = np.ones([2, 3]) # Create a source cube with metadata and scalar coords - src_cube = iris.cube.Cube(np.zeros(5)) + src_cube = Cube(np.zeros(5)) src_cube.units = "K" src_cube.attributes = {"a": 1} src_cube.standard_name = "air_temperature" - scalar_height = iris.coords.AuxCoord([5], units="m", standard_name="height") - scalar_time = iris.coords.DimCoord([10], units="s", standard_name="time") + scalar_height = AuxCoord([5], units="m", standard_name="height") + scalar_time = DimCoord([10], units="s", standard_name="time") src_cube.add_aux_coord(scalar_height) src_cube.add_aux_coord(scalar_time) mesh_dim = 0 + grid_x = DimCoord(np.arange(3), standard_name="longitude") + grid_y = DimCoord(np.arange(2), standard_name="latitude") + + cube = _create_cube(data, src_cube, mesh_dim, grid_x, grid_y) + src_metadata = src_cube.metadata + + expected_cube = Cube(data) + expected_cube.metadata = src_metadata + expected_cube.add_dim_coord(grid_x, 1) + expected_cube.add_dim_coord(grid_y, 0) + expected_cube.add_aux_coord(scalar_height) + expected_cube.add_aux_coord(scalar_time) + assert expected_cube == cube + +def test_create_cube_4D(): + """Test creation of 2D output grid.""" + data = np.ones([4, 2, 3, 5]) + + # Create a source cube with metadata and scalar coords + src_cube = Cube(np.zeros([4, 5, 5])) + src_cube.units = "K" + src_cube.attributes = {"a": 1} + src_cube.standard_name = "air_temperature" + scalar_height = AuxCoord([5], units="m", standard_name="height") + scalar_time = DimCoord([10], units="s", standard_name="time") + src_cube.add_aux_coord(scalar_height) + src_cube.add_aux_coord(scalar_time) + first_coord = DimCoord(np.arange(4), standard_name="air_pressure") + src_cube.add_dim_coord(first_coord, 0) + last_coord = AuxCoord(np.arange(5), long_name="last_coord") + src_cube.add_aux_coord(last_coord, 2) + multidim_coord = AuxCoord(np.ones([4, 5]), long_name="2d_coord") + src_cube.add_aux_coord(multidim_coord, (0, 2)) + ignored_coord = AuxCoord(np.arange(5), long_name="ignore") + src_cube.add_aux_coord(ignored_coord, 1) + + mesh_dim = 1 + grid_x = iris.coords.DimCoord(np.arange(3), standard_name="longitude") grid_y = iris.coords.DimCoord(np.arange(2), standard_name="latitude") @@ -30,8 +70,11 @@ def test_create_cube_2D(): expected_cube = iris.cube.Cube(data) expected_cube.metadata = src_metadata - expected_cube.add_dim_coord(grid_x, 1) - expected_cube.add_dim_coord(grid_y, 0) + expected_cube.add_dim_coord(grid_x, 2) + expected_cube.add_dim_coord(grid_y, 1) + expected_cube.add_dim_coord(first_coord, 0) + expected_cube.add_aux_coord(last_coord, 3) + expected_cube.add_aux_coord(multidim_coord, (0, 3)) expected_cube.add_aux_coord(scalar_height) expected_cube.add_aux_coord(scalar_time) assert expected_cube == cube From 57838082c85e7e26e6051f580f4aedb3d3c3dbb9 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Tue, 30 Mar 2021 09:26:20 +0100 Subject: [PATCH 3/4] lint fix --- .../unit/experimental/unstructured_scheme/test__create_cube.py | 1 + 1 file changed, 1 insertion(+) diff --git a/esmf_regrid/tests/unit/experimental/unstructured_scheme/test__create_cube.py b/esmf_regrid/tests/unit/experimental/unstructured_scheme/test__create_cube.py index d9994839..50536059 100644 --- a/esmf_regrid/tests/unit/experimental/unstructured_scheme/test__create_cube.py +++ b/esmf_regrid/tests/unit/experimental/unstructured_scheme/test__create_cube.py @@ -38,6 +38,7 @@ def test_create_cube_2D(): expected_cube.add_aux_coord(scalar_time) assert expected_cube == cube + def test_create_cube_4D(): """Test creation of 2D output grid.""" data = np.ones([4, 2, 3, 5]) From 1bb659fc1b4319e7371e4173a21877489c57db76 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Thu, 8 Apr 2021 11:53:12 +0100 Subject: [PATCH 4/4] address review comment --- esmf_regrid/experimental/unstructured_scheme.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/esmf_regrid/experimental/unstructured_scheme.py b/esmf_regrid/experimental/unstructured_scheme.py index 0a97b6a5..677163f3 100644 --- a/esmf_regrid/experimental/unstructured_scheme.py +++ b/esmf_regrid/experimental/unstructured_scheme.py @@ -73,7 +73,6 @@ def _create_cube(data, src_cube, mesh_dim, grid_x, grid_y): # data: a masked array containing the result of the regridding operation # src_cube: the source cube which data is regrid from # mesh_dim: the dimension on src_cube which the mesh belongs to - # mesh: the Mesh (or MeshCoord) object belonging to src_cube # grid_x: the coordinate on the target cube representing the x axis # grid_y: the coordinate on the target cube representing the y axis @@ -98,8 +97,11 @@ def copy_coords(src_coords, add_method): dims = src_cube.coord_dims(coord) if hasattr(coord, "mesh") or mesh_dim in dims: continue + # Since the mesh will be replaced by a 2D grid, dims which are + # beyond the mesh_dim are increased by one. dims = [dim if dim < mesh_dim else dim + 1 for dim in dims] result_coord = coord.copy() + # Add result_coord to the owner of add_method. add_method(result_coord, dims) # coord_mapping[id(coord)] = result_coord