added height interpolation for the case of just single level data. e.…

…g. u_30m from u_10m and u_100m, with u pressure level array

sup3r/preprocessing/derivers/base.py

@@ -183,6 +183,25 @@ def map_new_name(self, feature, pattern):
         )
         return new_feature
 
+    def has_interp_variables(self, feature):
+        """Check if the given feature can be interpolated from values at nearby
+        heights or from pressure level data. e.g. If ``u_10m`` and ``u_50m``
+        exist then ``u_30m`` can be interpolated from these. If a pressure
+        level array ``u`` is available this can also be used, in conjunction
+        with height data."""
+        fstruct = parse_feature(feature)
+        count = 0
+        for feat in self.data.features:
+            fstruct_check = parse_feature(feat)
+            height = fstruct_check.height
+
+            if (
+                fstruct_check.basename == fstruct.basename
+                and height is not None
+            ):
+                count += 1
+        return count > 1 or fstruct.basename in self.data
+
     def derive(self, feature) -> Union[np.ndarray, da.core.Array]:
         """Routine to derive requested features. Employs a little recursion to
         locate differently named features with a name map in the feature
@@ -195,8 +214,6 @@ def derive(self, feature) -> Union[np.ndarray, da.core.Array]:
         Features are all saved as lower case names and __contains__ checks will
         use feature.lower()
         """
-
-        fstruct = parse_feature(feature)
         if feature not in self.data:
             compute_check = self.check_registry(feature)
             if compute_check is not None and isinstance(compute_check, str):
@@ -206,7 +223,7 @@ def derive(self, feature) -> Union[np.ndarray, da.core.Array]:
             if compute_check is not None:
                 return compute_check
 
-            if fstruct.basename in self.data:
+            if self.has_interp_variables(feature):
                 logger.debug(
                     'Attempting level interpolation for "%s"', feature
                 )
@@ -223,7 +240,7 @@ def derive(self, feature) -> Union[np.ndarray, da.core.Array]:
 
         return self.data[feature]
 
-    def add_single_level_data(self, feature, lev_array, var_array):
+    def get_single_level_data(self, feature):
         """When doing level interpolation we should include the single level
         data available. e.g. If we have u_100m already and want to interpolate
         u_40m from multi-level data U we should add u_100m at height 100m
@@ -233,6 +250,8 @@ def add_single_level_data(self, feature, lev_array, var_array):
         pattern = fstruct.basename + '_(.*)'
         var_list = []
         lev_list = []
+        lev_array = None
+        var_array = None
         for f in list(self.data.data_vars):
             if re.match(pattern.lower(), f):
                 var_list.append(self.data[f])
@@ -245,22 +264,23 @@ def add_single_level_data(self, feature, lev_array, var_array):
                 lev_list.append(np.float32(lev))
 
         if len(var_list) > 0:
-            var_array = np.concatenate(
-                [var_array, da.stack(var_list, axis=-1)], axis=-1
-            )
+            var_array = da.stack(var_list, axis=-1)
             sl_shape = (*var_array.shape[:-1], len(lev_list))
-            single_levs = da.broadcast_to(da.from_array(lev_list), sl_shape)
-            lev_array = np.concatenate([lev_array, single_levs], axis=-1)
-        return lev_array, var_array
+            lev_array = da.broadcast_to(da.from_array(lev_list), sl_shape)
 
-    def do_level_interpolation(
-        self, feature, interp_method='linear'
-    ) -> xr.DataArray:
-        """Interpolate over height or pressure to derive the given feature."""
+        return var_array, lev_array
+
+    def get_multi_level_data(self, feature):
+        """Get data stored in multi-level arrays, like u stored on pressure
+        levels."""
         fstruct = parse_feature(feature)
-        var_array = self.data[fstruct.basename, ...]
-        if fstruct.height is not None:
-            level = [fstruct.height]
+        var_array = None
+        lev_array = None
+
+        if fstruct.basename in self.data:
+            var_array = self.data[fstruct.basename, ...]
+
+        if fstruct.height is not None and var_array is not None:
             msg = (
                 f'To interpolate {fstruct.basename} to {feature} the loaded '
                 'data needs to include "zg" and "topography" or have a '
@@ -281,8 +301,7 @@ def do_level_interpolation(
                     self.data[Dimension.HEIGHT, ...].astype(np.float32),
                     var_array.shape,
                 )
-        else:
-            level = [fstruct.pressure]
+        elif var_array is not None:
             msg = (
                 f'To interpolate {fstruct.basename} to {feature} the loaded '
                 'data needs to include "level" (a.k.a pressure at multiple '
@@ -293,10 +312,41 @@ def do_level_interpolation(
                 self.data[Dimension.PRESSURE_LEVEL, ...].astype(np.float32),
                 var_array.shape,
             )
+        return var_array, lev_array
+
+    def do_level_interpolation(
+        self, feature, interp_method='linear'
+    ) -> xr.DataArray:
+        """Interpolate over height or pressure to derive the given feature."""
+        ml_var, ml_levs = self.get_multi_level_data(feature)
+        sl_var, sl_levs = self.get_single_level_data(feature)
 
-        lev_array, var_array = self.add_single_level_data(
-            feature, lev_array, var_array
+        fstruct = parse_feature(feature)
+        attrs = {}
+        for feat in self.data.features:
+            if parse_feature(feat).basename == fstruct.basename:
+                attrs = self.data[feat].attrs
+
+        level = (
+            [fstruct.height]
+            if fstruct.height is not None
+            else [fstruct.pressure]
         )
+
+        if ml_var is not None:
+            var_array = ml_var
+            lev_array = ml_levs
+        elif sl_var is not None:
+            var_array = sl_var
+            lev_array = sl_levs
+        elif ml_var is not None and sl_var is not None:
+            var_array = np.concatenate([ml_var, sl_var], axis=-1)
+            lev_array = np.concatenate([ml_levs, sl_levs], axis=-1)
+        else:
+            msg = 'Neither single level nor multi level data was found for %s'
+            logger.error(msg, feature)
+            raise RuntimeError(msg % feature)
+
         out = Interpolator.interp_to_level(
             lev_array=lev_array,
             var_array=var_array,
@@ -306,7 +356,7 @@ def do_level_interpolation(
         return xr.DataArray(
             data=_rechunk_if_dask(out),
             dims=Dimension.dims_3d(),
-            attrs=self.data[fstruct.basename].attrs,
+            attrs=attrs,
         )
 
 

tests/derivers/test_height_interp.py

@@ -19,9 +19,9 @@
         ((10, 10), (37.25, -107), 1000),
     ],
 )
-def test_height_interp_nc(shape, target, height):
-    """Test that variables can be interpolated and extrapolated with height
-    correctly"""
+def test_plevel_height_interp_nc(shape, target, height):
+    """Test that variables on pressure levels can be interpolated and
+    extrapolated with height correctly"""
 
     with TemporaryDirectory() as td:
         wind_file = os.path.join(td, 'wind.nc')
@@ -53,7 +53,42 @@ def test_height_interp_nc(shape, target, height):
     assert np.array_equal(out, transform.data[f'u_{height}m'].data)
 
 
-def test_height_interp_with_single_lev_data_nc(
+def test_single_levels_height_interp_nc(shape=(10, 10), target=(37.25, -107)):
+    """Test that features can be interpolated from only single level
+    variables"""
+
+    with TemporaryDirectory() as td:
+        level_file = os.path.join(td, 'wind_levs.nc')
+        make_fake_nc_file(
+            level_file, shape=(10, 10, 20), features=['u_10m', 'u_100m']
+        )
+
+        derive_features = ['u_30m']
+        no_transform = Rasterizer([level_file], target=target, shape=shape)
+
+        transform = Deriver(
+            no_transform.data, derive_features, interp_method='linear'
+        )
+
+    h10 = np.zeros(transform.shape[:3], dtype=np.float32)[..., None]
+    h10[:] = 10
+    h100 = np.zeros(transform.shape[:3], dtype=np.float32)[..., None]
+    h100[:] = 100
+    hgt_array = np.concatenate([h10, h100], axis=-1)
+    u = np.concatenate(
+        [
+            no_transform['u_10m'].data[..., None],
+            no_transform['u_100m'].data[..., None],
+        ],
+        axis=-1,
+    )
+    out = Interpolator.interp_to_level(hgt_array, u, [np.float32(30)])
+
+    assert transform.data['u_30m'].data.dtype == np.float32
+    assert np.array_equal(out, transform.data['u_30m'].data)
+
+
+def test_plevel_height_interp_with_single_lev_data_nc(
     shape=(10, 10), target=(37.25, -107)
 ):
     """Test that variables can be interpolated with height correctly"""

tests/pipeline/test_pipeline.py

@@ -363,7 +363,8 @@ def test_fwp_pipeline_with_mask(input_files):
         with ResourceX(fp_out) as f:
             assert len(f.time_index) == t_enhance * n_tsteps
 
-            # unmasked gives 4 chunks so without chunk index 2 we have just 3
+            # unmasked gives 4 spatial chunks so without chunk index 2 we
+            # have just 3
             assert len(f.meta) == s_enhance * s_enhance * 3 * np.prod(
                 fp_chunk_shape[:2]
             )