interpolation adjustment: option to include single levels in interpol…

sup3r/preprocessing/derivers/base.py

@@ -55,17 +55,21 @@ def __init__(
             that is not found in the :class:`Rasterizer` data it will look for
             a method to derive the feature in the registry.
         interp_kwargs : dict | None
-            Dictionary of kwargs for level interpolation. Can include "method"
-            and "run_level_check" keys. Method specifies how to perform height
-            interpolation. e.g. Deriving u_20m from u_10m and u_100m. Options
-            are "linear" and "log". See
+            Dictionary of kwargs for level interpolation. Can include "method",
+            "run_level_check", and "include_single_levels" keys. Method
+            specifies how to perform height interpolation. e.g. Deriving
+            u_20m from u_10m and u_100m. Options are "linear" and "log".
+            ``include_single_levels = True`` will include single level
+            variables in addition to pressure level variables in the
+            interpolation. e.g. the 3D array of ``temperature_2m`` along
+            with the 4D array of ``temperature``. ``See
             :py:meth:`sup3r.preprocessing.derivers.Deriver.do_level_interpolation`
         """  # pylint: disable=line-too-long
         if FeatureRegistry is not None:
             self.FEATURE_REGISTRY = FeatureRegistry
 
         super().__init__(data=data)
-        self.interp_kwargs = interp_kwargs
+        self.interp_kwargs = interp_kwargs or {}
         features = parse_to_list(data=data, features=features)
         new_features = [f for f in features if f not in self.data]
         for f in new_features:
@@ -269,7 +273,6 @@ def get_single_level_data(self, feature):
             var_array = da.stack(var_list, axis=-1)
             sl_shape = (*var_array.shape[:-1], len(lev_list))
             lev_array = da.broadcast_to(da.from_array(lev_list), sl_shape)
-
         return var_array, lev_array
 
     def get_multi_level_data(self, feature):
@@ -296,8 +299,8 @@ def get_multi_level_data(self, feature):
             assert can_calc_height or have_height, msg
 
             if can_calc_height:
-                lev_array = self.data[['zg', 'topography']].as_array()
-                lev_array = lev_array[..., 0] - lev_array[..., 1]
+                lev_array = self.data['zg'] - self.data['topography']
+                lev_array = lev_array.data
             else:
                 lev_array = da.broadcast_to(
                     self.data[Dimension.HEIGHT].astype(np.float32),
@@ -321,7 +324,10 @@ def do_level_interpolation(
     ) -> 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)
+        if interp_kwargs.get('include_single_levels', False):
+            sl_var, sl_levs = self.get_single_level_data(feature)
+        else:
+            sl_var, sl_levs = None, None
 
         fstruct = parse_feature(feature)
         attrs = {}

sup3r/utilities/interpolation.py

@@ -20,9 +20,6 @@ def get_level_masks(cls, lev_array, level):
 
         Parameters
         ----------
-        var_array : Union[np.ndarray, da.core.Array]
-            Array of variable data, for example u-wind in a 4D array of shape
-            (lat, lon, time, level)
         lev_array : Union[np.ndarray, da.core.Array]
             Height or pressure values for the corresponding entries in
             var_array, in the same shape as var_array. If this is height and
@@ -45,14 +42,14 @@ def get_level_masks(cls, lev_array, level):
             to the one requested.
             (lat, lon, time, level)
         """
-        argmin1 = da.argmin(da.abs(lev_array - level), axis=-1, keepdims=True)
+        lev_diff = np.abs(lev_array - level)
+        argmin1 = da.argmin(lev_diff, axis=-1, keepdims=True)
         lev_indices = da.broadcast_to(
             da.arange(lev_array.shape[-1]), lev_array.shape
         )
         mask1 = lev_indices == argmin1
-
-        other_levs = da.ma.masked_array(lev_array, mask1)
-        argmin2 = da.argmin(da.abs(other_levs - level), axis=-1, keepdims=True)
+        lev_diff = da.abs(da.ma.masked_array(lev_array, mask1) - level)
+        argmin2 = da.argmin(lev_diff, axis=-1, keepdims=True)
         mask2 = lev_indices == argmin2
         return mask1, mask2
 
@@ -109,16 +106,16 @@ def interp_to_level(
 
         Parameters
         ----------
-        var_array : xr.DataArray
-            Array of variable data, for example u-wind in a 4D array of shape
-            (lat, lon, time, level)
         lev_array : xr.DataArray
             Height or pressure values for the corresponding entries in
             var_array, in the same shape as var_array. If this is height and
             the requested levels are hub heights above surface, lev_array
             should be the geopotential height corresponding to every var_array
             index relative to the surface elevation (subtract the elevation at
             the surface from the geopotential height)
+        var_array : xr.DataArray
+            Array of variable data, for example u-wind in a 4D array of shape
+            (lat, lon, time, level)
         level : float
             level or levels to interpolate to (e.g. final desired hub height
             above surface elevation)

tests/data_handlers/test_dh_nc_cc.py

@@ -69,6 +69,7 @@ def test_reload_cache():
             target=target,
             shape=(20, 20),
             cache_kwargs=cache_kwargs,
+            interp_kwargs={'include_single_levels': True}
         )
 
         # reload from cache
@@ -80,7 +81,9 @@ def test_reload_cache():
             cache_kwargs=cache_kwargs,
         )
         assert all(f in cached for f in features)
-        assert np.array_equal(handler.as_array(), cached.as_array())
+        harr = handler.as_array().compute()
+        carr = cached.as_array().compute()
+        assert np.array_equal(harr, carr)
 
 
 @pytest.mark.parametrize(

tests/derivers/test_height_interp.py

@@ -156,7 +156,7 @@ def test_single_levels_height_interp_nc(shape=(10, 10), target=(37.25, -107)):
         transform = Deriver(
             no_transform.data,
             derive_features,
-            interp_kwargs={'method': 'linear'},
+            interp_kwargs={'method': 'linear', 'include_single_levels': True},
         )
 
     h10 = np.zeros(transform.shape[:3], dtype=np.float32)[..., None]
@@ -197,7 +197,11 @@ def test_plevel_height_interp_with_single_lev_data_nc(
             [wind_file, level_file], target=target, shape=shape
         )
 
-        transform = Deriver(no_transform.data, derive_features)
+        transform = Deriver(
+            no_transform.data,
+            derive_features,
+            interp_kwargs={'include_single_levels': True},
+        )
 
     hgt_array = (
         no_transform['zg'].data - no_transform['topography'].data[..., None]
@@ -237,7 +241,7 @@ def test_log_interp(shape=(10, 10), target=(37.25, -107)):
         transform = Deriver(
             no_transform.data,
             derive_features,
-            interp_kwargs={'method': 'log'},
+            interp_kwargs={'method': 'log', 'include_single_levels': True},
         )
 
     hgt_array = (

tests/forward_pass/test_forward_pass.py

@@ -491,9 +491,7 @@ def test_fwp_chunking(input_files):
             data_chunked[hr_slice][..., t_hr_slice, :] = out
 
         err = data_chunked - data_nochunk
-        err /= data_nochunk
-
-        assert np.mean(np.abs(err.flatten())) < 0.01
+        assert np.mean(np.abs(err)) < 1e-6
 
 
 def test_fwp_nochunking(input_files):