diff --git a/tobac/tests/test_xarray_utils.py b/tobac/tests/test_xarray_utils.py index caabef8c..3f2028fb 100644 --- a/tobac/tests/test_xarray_utils.py +++ b/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( diff --git a/tobac/utils/internal/xarray_utils.py b/tobac/utils/internal/xarray_utils.py index f6802071..5844522a 100644 --- a/tobac/utils/internal/xarray_utils.py +++ b/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