Fix get_grid when region_mask_regions data is present

.gitignore

@@ -102,3 +102,6 @@ venv.bak/
 
 # mypy
 .mypy_cache/
+
+# VSCode
+.vscode

pop_tools/grid.py

@@ -315,8 +315,18 @@ def get_grid(grid_name, scrip=False):
         )
 
     grid_attrs.update({'title': f'{grid_name} grid'})
+
     dso.attrs = grid_attrs
 
+    # Remove region_mask_regions
+    if 'region_mask_regions' in dso.attrs:
+        regions = dso.attrs.pop('region_mask_regions')
+        region_names, region_vals = list(zip(*regions.items()))
+        region_coord = list(range(len(regions)))
+        dso['region_name'] = xr.DataArray(list(region_names), coords=[region_coord], dims=['nreg'])
+        dso['region_val'] = xr.DataArray(list(region_vals), coords=[region_coord], dims=['nreg'])
+        dso['region_val'].attrs['coordinate'] = 'region_name'
+
     return dso
 
 

pop_tools/region_masks.py

@@ -57,7 +57,7 @@ def region_mask_3d(grid_name, mask_name=None, region_defs=None):
         mask_name = 'user-defined'
 
     if region_defs is None:
-        region_defs = _get_region_defititions(ds, grid_name, mask_name)
+        region_defs = _get_region_definitions(ds, grid_name, mask_name)
 
     _verify_region_def_schema(region_defs)
 
@@ -131,7 +131,7 @@ def _verify_region_def_schema(region_defs):
                 assert crit_type in ['bounds', 'match']
 
 
-def _get_region_defititions(ds, grid_name, mask_name):
+def _get_region_definitions(ds, grid_name, mask_name):
     """Return the region mask definition from `region_mask_definitions.yaml`."""
     if mask_name == 'default':
         grid_attrs = grid_defs[grid_name]

tests/test_grid.py

@@ -26,3 +26,20 @@ def test_get_grid_scrip():
     ds_test = pop_tools.get_grid('POP_gx3v7', scrip=True)
     ds_ref = xr.open_dataset(DATASETS.fetch('POP_gx3v7.nc'))
     assert ds_compare(ds_test, ds_ref, assertion='allclose', rtol=1e-14, atol=1e-14)
+
+
+def test_get_grid_twice():
+    ds1 = pop_tools.get_grid('POP_gx1v7')
+    ds2 = pop_tools.get_grid('POP_gx1v7')
+    xr.testing.assert_identical(ds1, ds2)
+
+
+def test_get_grid_to_netcdf():
+    for grid in pop_tools.grid_defs.keys():
+        print('-' * 80)
+        print(grid)
+        ds = pop_tools.get_grid(grid)
+        for format in ['NETCDF4', 'NETCDF3_64BIT']:
+            gridfile = f'{grid}_{format}.nc'
+            ds.to_netcdf(gridfile, format=format)
+            os.system(f'rm -f {gridfile}')