9
9
from xarray .core .options import _get_keep_attrs
10
10
from xarray .core .pdcompat import count_not_none
11
11
from xarray .core .pycompat import is_duck_dask_array
12
- from xarray .core .types import T_DataWithCoords
12
+ from xarray .core .types import T_DataWithCoords , T_DuckArray
13
13
14
14
15
- def _get_alpha (com = None , span = None , halflife = None , alpha = None ):
15
+ def _get_alpha (
16
+ com : float | None = None ,
17
+ span : float | None = None ,
18
+ halflife : float | None = None ,
19
+ alpha : float | None = None ,
20
+ ) -> float :
16
21
# pandas defines in terms of com (converting to alpha in the algo)
17
22
# so use its function to get a com and then convert to alpha
18
23
19
24
com = _get_center_of_mass (com , span , halflife , alpha )
20
25
return 1 / (1 + com )
21
26
22
27
23
- def move_exp_nanmean (array , * , axis , alpha ) :
28
+ def move_exp_nanmean (array : T_DuckArray , * , axis : int , alpha : float ) -> np . ndarray :
24
29
if is_duck_dask_array (array ):
25
30
raise TypeError ("rolling_exp is not currently support for dask-like arrays" )
26
31
import numbagg
@@ -32,15 +37,20 @@ def move_exp_nanmean(array, *, axis, alpha):
32
37
return numbagg .move_exp_nanmean (array , axis = axis , alpha = alpha )
33
38
34
39
35
- def move_exp_nansum (array , * , axis , alpha ) :
40
+ def move_exp_nansum (array : T_DuckArray , * , axis : int , alpha : float ) -> np . ndarray :
36
41
if is_duck_dask_array (array ):
37
42
raise TypeError ("rolling_exp is not currently supported for dask-like arrays" )
38
43
import numbagg
39
44
40
45
return numbagg .move_exp_nansum (array , axis = axis , alpha = alpha )
41
46
42
47
43
- def _get_center_of_mass (comass , span , halflife , alpha ):
48
+ def _get_center_of_mass (
49
+ comass : float | None ,
50
+ span : float | None ,
51
+ halflife : float | None ,
52
+ alpha : float | None ,
53
+ ) -> float :
44
54
"""
45
55
Vendored from pandas.core.window.common._get_center_of_mass
46
56
0 commit comments