11import warnings
22from functools import partial
3- from typing import Any , Callable , Dict , Sequence
3+ from numbers import Number
4+ from typing import Any , Callable , Dict , Hashable , Sequence , Union
45
56import numpy as np
67import pandas as pd
78
89from . import utils
9- from .common import _contains_datetime_like_objects
10+ from .common import _contains_datetime_like_objects , ones_like
1011from .computation import apply_ufunc
1112from .duck_array_ops import dask_array_type
1213from .utils import OrderedSet , is_scalar
1314from .variable import Variable , broadcast_variables
1415
1516
17+ def _get_nan_block_lengths (obj , dim : Hashable , index : Variable ):
18+ """
19+ Return an object where each NaN element in 'obj' is replaced by the
20+ length of the gap the element is in.
21+ """
22+
23+ # make variable so that we get broadcasting for free
24+ index = Variable ([dim ], index )
25+
26+ # algorithm from https://github.com/pydata/xarray/pull/3302#discussion_r324707072
27+ arange = ones_like (obj ) * index
28+ valid = obj .notnull ()
29+ valid_arange = arange .where (valid )
30+ cumulative_nans = valid_arange .ffill (dim = dim ).fillna (index [0 ])
31+
32+ nan_block_lengths = (
33+ cumulative_nans .diff (dim = dim , label = "upper" )
34+ .reindex ({dim : obj [dim ]})
35+ .where (valid )
36+ .bfill (dim = dim )
37+ .where (~ valid , 0 )
38+ .fillna (index [- 1 ] - valid_arange .max ())
39+ )
40+
41+ return nan_block_lengths
42+
43+
1644class BaseInterpolator :
1745 """Generic interpolator class for normalizing interpolation methods
1846 """
@@ -178,7 +206,7 @@ def _apply_over_vars_with_dim(func, self, dim=None, **kwargs):
178206 return ds
179207
180208
181- def get_clean_interp_index (arr , dim , use_coordinate = True ):
209+ def get_clean_interp_index (arr , dim : Hashable , use_coordinate : Union [ str , bool ] = True ):
182210 """get index to use for x values in interpolation.
183211
184212 If use_coordinate is True, the coordinate that shares the name of the
@@ -195,23 +223,33 @@ def get_clean_interp_index(arr, dim, use_coordinate=True):
195223 index = arr .coords [use_coordinate ]
196224 if index .ndim != 1 :
197225 raise ValueError (
198- "Coordinates used for interpolation must be 1D, "
199- "%s is %dD." % ( use_coordinate , index .ndim )
226+ f "Coordinates used for interpolation must be 1D, "
227+ f" { use_coordinate } { index .ndim } D."
200228 )
229+ index = index .to_index ()
230+
231+ # TODO: index.name is None for multiindexes
232+ # set name for nice error messages below
233+ if isinstance (index , pd .MultiIndex ):
234+ index .name = dim
235+
236+ if not index .is_monotonic :
237+ raise ValueError (f"Index { index .name !r} )
238+
239+ if not index .is_unique :
240+ raise ValueError (f"Index { index .name !r} )
201241
202242 # raise if index cannot be cast to a float (e.g. MultiIndex)
203243 try :
204244 index = index .values .astype (np .float64 )
205245 except (TypeError , ValueError ):
206246 # pandas raises a TypeError
207- # xarray/nuppy raise a ValueError
247+ # xarray/numpy raise a ValueError
208248 raise TypeError (
209- "Index must be castable to float64 to support"
210- "interpolation, got: %s" % type (index )
249+ f "Index { index . name !r } must be castable to float64 to support "
250+ f "interpolation, got { type (index ). __name__ } ."
211251 )
212- # check index sorting now so we can skip it later
213- if not (np .diff (index ) > 0 ).all ():
214- raise ValueError ("Index must be monotonicly increasing" )
252+
215253 else :
216254 axis = arr .get_axis_num (dim )
217255 index = np .arange (arr .shape [axis ], dtype = np .float64 )
@@ -220,7 +258,13 @@ def get_clean_interp_index(arr, dim, use_coordinate=True):
220258
221259
222260def interp_na (
223- self , dim = None , use_coordinate = True , method = "linear" , limit = None , ** kwargs
261+ self ,
262+ dim : Hashable = None ,
263+ use_coordinate : Union [bool , str ] = True ,
264+ method : str = "linear" ,
265+ limit : int = None ,
266+ max_gap : Union [int , float , str , pd .Timedelta , np .timedelta64 ] = None ,
267+ ** kwargs ,
224268):
225269 """Interpolate values according to different methods.
226270 """
@@ -230,6 +274,40 @@ def interp_na(
230274 if limit is not None :
231275 valids = _get_valid_fill_mask (self , dim , limit )
232276
277+ if max_gap is not None :
278+ max_type = type (max_gap ).__name__
279+ if not is_scalar (max_gap ):
280+ raise ValueError ("max_gap must be a scalar." )
281+
282+ if (
283+ dim in self .indexes
284+ and isinstance (self .indexes [dim ], pd .DatetimeIndex )
285+ and use_coordinate
286+ ):
287+ if not isinstance (max_gap , (np .timedelta64 , pd .Timedelta , str )):
288+ raise TypeError (
289+ f"Underlying index is DatetimeIndex. Expected max_gap of type str, pandas.Timedelta or numpy.timedelta64 but received { max_type }
290+ )
291+
292+ if isinstance (max_gap , str ):
293+ try :
294+ max_gap = pd .to_timedelta (max_gap )
295+ except ValueError :
296+ raise ValueError (
297+ f"Could not convert { max_gap !r}
298+ )
299+
300+ if isinstance (max_gap , pd .Timedelta ):
301+ max_gap = np .timedelta64 (max_gap .value , "ns" )
302+
303+ max_gap = np .timedelta64 (max_gap , "ns" ).astype (np .float64 )
304+
305+ if not use_coordinate :
306+ if not isinstance (max_gap , (Number , np .number )):
307+ raise TypeError (
308+ f"Expected integer or floating point max_gap since use_coordinate=False. Received { max_type }
309+ )
310+
233311 # method
234312 index = get_clean_interp_index (self , dim , use_coordinate = use_coordinate )
235313 interp_class , kwargs = _get_interpolator (method , ** kwargs )
@@ -253,6 +331,14 @@ def interp_na(
253331 if limit is not None :
254332 arr = arr .where (valids )
255333
334+ if max_gap is not None :
335+ if dim not in self .coords :
336+ raise NotImplementedError (
337+ "max_gap not implemented for unlabeled coordinates yet."
338+ )
339+ nan_block_lengths = _get_nan_block_lengths (self , dim , index )
340+ arr = arr .where (nan_block_lengths <= max_gap )
341+
256342 return arr
257343
258344
0 commit comments