Merge pull request #23 from jhardenberg/fix/level_idx

Prefix to identify helper index and fix single level

CHANGELOG.md

@@ -3,6 +3,11 @@
 All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 
+## [v0.0.5]
+
+- Allow also single 3D level
+- Prefix used to identify helper index coordinate
+
 ## [v0.0.4]
 
 - Additional helper coordinate to select correct levels from full 3D weights

smmregrid/__init__.py

@@ -4,4 +4,4 @@
 from .cdo_weights import cdo_generate_weights
 # from .util import find_vert_coords
 
-__version__ = '0.0.4'
+__version__ = '0.0.5'

smmregrid/regrid.py

@@ -247,8 +247,8 @@ class Regridder(object):
         weights (:class:`xarray.Dataset`): Pre-computed interpolation weights
         vert_coord (str): Name of the vertical coordinate.
                           If provided, 3D weights are generated (default: None)
-        level_idx (str): Name of helper vertical coordinate with original level indices.
-                         If provided, 3D weights are selected from those levels (default: "idx_3d")
+        level_idx (str): Prefix of helper vertical coordinate with original level indices.
+                         If provided, 3D weights are selected from those levels (default: "idx_")
         method (str): Method to use for interpolation (default: 'con')
         space_dims (list): list of dimensions to interpolate (default: None)
         transpose (bool): transpose the output so that the vertical coordinate is
@@ -257,7 +257,7 @@ class Regridder(object):
 
     def __init__(self, source_grid=None, target_grid=None, weights=None,
                  method='con', space_dims=None, vert_coord=None, transpose=True,
-                 cdo='cdo', level_idx="idx_3d"):
+                 cdo='cdo', level_idx="idx_"):
 
         if (source_grid is None or target_grid is None) and (weights is None):
             raise ValueError(
@@ -368,10 +368,12 @@ def regrid3d(self, source_data):
             return source_data
 
         # If a special additional coordinate is present pick correct levels from weights
-        if self.level_idx in source_data.coords:
-            levlist = source_data.coords[self.level_idx].values
+        coord = next((coord for coord in source_data.coords if coord.startswith(self.level_idx)), None)
+        if coord:  # if a coordinate starting with level_idx is found
+            levlist = source_data.coords[coord].values.tolist()
+            levlist = [levlist] if numpy.isscalar(levlist) else levlist
         else:
-            levlist = range(0, source_data.coords[self.vert_coord].values.size)
+            levlist = list(range(0, source_data.coords[self.vert_coord].values.size))
 
         data3d_list = []
         for lev in range(0, len(levlist)):

tests/levels_test.py

@@ -19,4 +19,9 @@ def test_plev_gaussian_levels(method):
 def test_nemo_3d_levels(method):
     fff = check_cdo_regrid_levels(os.path.join(INDIR, 'so3d-nemo.nc'), tfile,
                            "lev", [14,15,17], remap_method=method)
-    assert fff is True
+    assert fff is True,  "Multiple levels test failed"
+
+    fff = check_cdo_regrid_levels(os.path.join(INDIR, 'so3d-nemo.nc'), tfile,
+                           "lev", [15], remap_method=method) # single level
+    assert fff is True, "Single 3D level test failed"
+