Unstructured Scheme - Cube Creation 3D (#47)

* handle extra dims in _create_cube

* add test

* lint fix

* address review comment

esmf_regrid/experimental/unstructured_scheme.py

@@ -73,16 +73,11 @@ 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
 
     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 +87,37 @@ 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())
+    # 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:
+            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
+
+    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
 

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,56 @@ 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 +71,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