Unify unstructured and rectilinear schemes (curvilinear support) (#175)

* have _create_cube and _cube_to_GridInfo derive from common source * fix _cube_to_GridInfo * fix _create_cube * fix _create_cube * fix _create_cube * fix _create_cube * fix test__create_cube * add curvilinear support to unstructured schemes * lint fix * allow saving of curvilinear regridders * edit docstrings * lint fix * tidy code * address review comments * address review comments * fix test
SciTools-incubator · Jun 23, 2022 · 6624e35 · 6624e35
1 parent 24997c7
commit 6624e35
Show file tree

Hide file tree

Showing 7 changed files with 339 additions and 338 deletions.
diff --git a/esmf_regrid/experimental/io.py b/esmf_regrid/experimental/io.py
@@ -50,13 +50,25 @@ def save_regridder(rg, filename):
         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
+
     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, var_name=SOURCE_NAME, long_name=SOURCE_NAME)
-        src_cube.add_dim_coord(src_grid[0], 0)
-        src_cube.add_dim_coord(src_grid[1], 1)
+        src_cube = _standard_grid_cube(src_grid, SOURCE_NAME)
 
         tgt_mesh = rg.mesh
         tgt_location = rg.location
@@ -75,11 +87,7 @@ def save_regridder(rg, filename):
             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, var_name=TARGET_NAME, long_name=TARGET_NAME)
-        tgt_cube.add_dim_coord(tgt_grid[0], 0)
-        tgt_cube.add_dim_coord(tgt_grid[1], 1)
+        tgt_cube = _standard_grid_cube(tgt_grid, TARGET_NAME)
     else:
         msg = (
             f"Expected a regridder of type `GridToMeshESMFRegridder` or "