Merge pull request #517 from tobac-project/main

Merge changes from `main` into `RC_v1.6.x`

Author: w-k-jones

tobac/tests/test_xarray_utils.py

@@ -4,6 +4,7 @@
 
 from typing import Union
 
+import pandas as pd
 import pytest
 import numpy as np
 import xarray as xr
@@ -156,6 +157,31 @@ def test_find_axis_from_dim_coord(
             (1, 1),
             {"test_coord1": (1, 1, 1), "test_coord_time": (5, 6, 7)},
         ),
+        (
+            ["time", "x", "y"],
+            {
+                "test_coord_datetime": (
+                    "time",
+                    pd.date_range(
+                        datetime.datetime(2000, 1, 1),
+                        datetime.datetime(2000, 1, 1, 6),
+                        freq="1h",
+                        inclusive="left",
+                    ),
+                ),
+                "test_coord_time": ("time", [5, 6, 7, 8, 9, 10]),
+            },
+            (1, 1),
+            {
+                "test_coord_datetime": pd.date_range(
+                    datetime.datetime(2000, 1, 1),
+                    datetime.datetime(2000, 1, 1, 3),
+                    freq="1h",
+                    inclusive="left",
+                ),
+                "test_coord_time": (5, 6, 7),
+            },
+        ),
     ],
 )
 def test_add_coordinates_to_features_interpolate_along_other_dims(

tobac/utils/internal/xarray_utils.py

@@ -420,9 +420,20 @@ def add_coordinates_to_features(
             except KeyError:
                 pass
 
-        return_feat_df[interp_coord_name] = renamed_dim_da[interp_coord].interp(
-            coords={
-                dim: dim_interp_coords[dim] for dim in renamed_dim_da[interp_coord].dims
-            }
-        )
+        if renamed_dim_da[interp_coord].dtype.kind in "uifc":
+            # Interpolate over the coordinate
+            return_feat_df[interp_coord_name] = renamed_dim_da[interp_coord].interp(
+                coords={
+                    dim: dim_interp_coords[dim]
+                    for dim in renamed_dim_da[interp_coord].dims
+                }
+            )
+        else:
+            # If non-numeric, we should instead just index the nearest values:
+            return_feat_df[interp_coord_name] = renamed_dim_da[interp_coord].isel(
+                **{
+                    dim: np.round(dim_interp_coords[dim]).astype(int)
+                    for dim in renamed_dim_da[interp_coord].dims
+                }
+            )
     return return_feat_df