polars-ta 0.4.7__py3-none-any.whl → 0.5.1__py3-none-any.whl

polars_ta/utils/numba_.py

@@ -6,6 +6,7 @@ from typing import List
 
 import numpy as np
 from numba import jit
+from numpy import full
 from numpy.lib.stride_tricks import sliding_window_view
 from polars import Series, Expr, struct, DataFrame
 
@@ -15,6 +16,90 @@ batches_i1_o1这一类的函数输入不支持Series，只支持numpy。设计
 """
 
 
+@jit(nopython=True, nogil=True, fastmath=True, cache=True)
+def isnan(x):
+    # https://github.com/numba/numba/issues/2919#issuecomment-747377615
+    if int(x) == -9223372036854775808:
+        return True
+    else:
+        return False
+
+
+@jit(nopython=True, nogil=True, cache=True)
+def full_with_window_size(arr, fill_value, dtype=None, window_size: int = 1):
+    out = full(arr.shape[0] + window_size - 1, fill_value, dtype=dtype)
+    out[window_size - 1:] = arr
+    return out
+
+
+@jit(nopython=True, nogil=True, cache=True)
+def sliding_window_with_min_periods(arr, window_size: int, min_periods: int):
+    windows = sliding_window_view(arr, window_size)
+    valid_counts = np.sum(~np.isnan(windows), axis=1)
+    # 修改这一行，使用布尔索引而不是np.where
+    result = windows.copy()
+    result[valid_counts < min_periods] = np.nan
+    return result
+
+
+@jit(nopython=True, nogil=True, cache=True)
+def _roll_1(x1: np.ndarray, window: int, min_periods: int, func):
+    """TODO 只是模板演示，无法编译通过"""
+    out1 = full_with_window_size(x1, np.nan, dtype=np.float64, window_size=window)
+    a1 = sliding_window_with_min_periods(out1, window, min_periods)
+    out1[:] = np.nan
+    for i, v1 in enumerate(a1):
+        if np.isnan(v1).all():
+            continue
+        out1[i] = func(v1)
+    return out1[:x1.shape[0]]
+
+
+@jit(nopython=True, nogil=True, cache=True)
+def _roll_2(x1, x2, window, min_periods, func):
+    """TODO 只是模板演示，无法编译通过"""
+    out1 = full_with_window_size(x1, np.nan, dtype=np.float64, window_size=window)
+    out2 = full_with_window_size(x2, np.nan, dtype=np.float64, window_size=window)
+    a1 = sliding_window_with_min_periods(out1, window, min_periods)
+    a2 = sliding_window_with_min_periods(out2, window, min_periods)
+    out1[:] = np.nan
+    for i, (v1, v2) in enumerate(zip(a1, a2)):
+        if np.isnan(v1).all():
+            continue
+        if np.isnan(v2).all():
+            continue
+        out1[i] = func(v1, v2)
+    return out1[:x1.shape[0]]
+
+
+@jit(nopython=True, nogil=True, cache=True)
+def _roll_3(x1, x2, x3, window, min_periods, func):
+    """TODO 只是模板演示，无法编译通过"""
+    out1 = full_with_window_size(x1, np.nan, dtype=np.float64, window_size=window)
+    out2 = full_with_window_size(x2, np.nan, dtype=np.float64, window_size=window)
+    out3 = full_with_window_size(x3, np.nan, dtype=np.float64, window_size=window)
+    a1 = sliding_window_with_min_periods(out1, window, min_periods)
+    a2 = sliding_window_with_min_periods(out2, window, min_periods)
+    a3 = sliding_window_with_min_periods(out3, window, min_periods)
+    out1[:] = np.nan
+    for i, (v1, v2, v3) in enumerate(zip(a1, a2, a3)):
+        if np.isnan(v1).all():
+            continue
+        if np.isnan(v2).all():
+            continue
+        if np.isnan(v3).all():
+            continue
+        out1[i] = func(v1, v2, v3)
+    return out1[:x1.shape[0]]
+
+
+def struct_to_numpy(xx, n: int, dtype=None):
+    if dtype is None:
+        return [xx.struct[i].to_numpy() for i in range(n)]
+    else:
+        return [xx.struct[i].to_numpy().astype(dtype) for i in range(n)]
+
+
 def batches_i1_o1(x1: np.ndarray, func, *args, dtype=None) -> Series:
     return Series(func(x1, *args), nan_to_null=True, dtype=dtype)
 
@@ -24,15 +109,11 @@ def batches_i2_o1(xx: List[np.ndarray], func, *args, dtype=None) -> Series:
 
 
 def batches_i1_o2(x1: np.ndarray, func, *args, dtype=None) -> Series:
-    out = func(x1, *args)
-    ss = [Series(x, nan_to_null=True, dtype=dtype) for x in out]
-    return DataFrame(ss).to_struct()
+    return DataFrame(func(x1, *args), nan_to_null=True).to_struct()
 
 
 def batches_i2_o2(xx: List[np.ndarray], func, *args, dtype=None) -> Series:
-    out = func(*xx, *args)
-    ss = [Series(x, nan_to_null=True, dtype=dtype) for x in out]
-    return DataFrame(ss).to_struct()
+    return DataFrame(func(*xx, *args), nan_to_null=True).to_struct()
 
 
 def batches_i2_o2_v2(xx: List[np.ndarray], func, *args, dtype=None) -> Series:
@@ -76,4 +157,4 @@ def roll_sum(x: Expr, n: int) -> Expr:
 
 
 def roll_cov(a: Expr, b: Expr, n: int) -> Expr:
-    return struct([a, b]).map_batches(lambda xx: batches_i2_o1([xx.struct[i].to_numpy() for i in range(2)], nb_roll_cov, n))
+    return struct([a, b]).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 2), nb_roll_cov, n))

polars_ta/utils/pandas_.py

@@ -35,7 +35,7 @@ def get_window_bounds(
     return start, end
 
 
-def roll_rank(x: Series, d: int, pct: bool = True, method: str = 'average', ascending: bool = True):
+def roll_rank(x: Series, d: int, minp: int, pct: bool = True, method: str = 'average', ascending: bool = True):
     start, end = get_window_bounds(len(x), d)
     """
     https://github.com/pandas-dev/pandas/blob/main/pandas/_libs/window/aggregations.pyx#L1281
@@ -46,11 +46,11 @@ def roll_rank(x: Series, d: int, pct: bool = True, method: str = 'average', asce
 
     O(N log(window)) implementation using skip list
     """
-    ret = _roll_rank(x.to_numpy().astype(float), start, end, d, pct, method, ascending)
+    ret = _roll_rank(x.to_numpy().astype(float), start, end, minp, pct, method, ascending)
     return Series(ret, nan_to_null=True)
 
 
-def roll_kurt(x, d):
+def roll_kurt(x, d: int, minp: int):
     start, end = get_window_bounds(len(x), d)
     """
     https://github.com/pandas-dev/pandas/blob/main/pandas/_libs/window/aggregations.pyx#L803
@@ -58,5 +58,5 @@ def roll_kurt(x, d):
     def roll_kurt(ndarray[float64_t] values, ndarray[int64_t] start,
               ndarray[int64_t] end, int64_t minp) -> np.ndarray:
     """
-    ret = _roll_kurt(x.to_numpy().astype(float), start, end, d)
+    ret = _roll_kurt(x.to_numpy().astype(float), start, end, minp)
     return Series(ret, nan_to_null=True)

polars_ta/wq/_nb.py

@@ -1,107 +1,121 @@
 import numpy as np
-from numba import jit, float64, bool_
-from numpy import argmax, argmin, prod, mean, std, full, vstack, corrcoef
+from numba import jit, float64, boolean
+from numpy import argmax, argmin, full, vstack, corrcoef, nanprod, nanmean, nanstd
 from numpy.lib.stride_tricks import sliding_window_view
 
+from polars_ta.utils.numba_ import isnan, full_with_window_size, sliding_window_with_min_periods
 
-@jit(nopython=True, nogil=True, fastmath=True, cache=True)
-def roll_argmax(x1, window, reverse):
-    out = full(x1.shape, np.nan, dtype=np.float64)
-    if len(x1) < window:
-        return out
-    a1 = sliding_window_view(x1, window)
+
+@jit(nopython=True, nogil=True, cache=True)
+def roll_argmax(x1, window, min_periods, reverse):
+    out1 = full_with_window_size(x1, np.nan, dtype=np.float64, window_size=window)
+    a1 = sliding_window_with_min_periods(out1, window, min_periods)
+    out1[:] = np.nan
     if reverse:
         a1 = a1[:, ::-1]
     for i, v1 in enumerate(a1):
-        out[i + window - 1] = argmax(v1)
-    return out
+        if np.isnan(v1).all():
+            continue
+        out1[i] = argmax(v1)
+    return out1[:x1.shape[0]]
 
 
-@jit(nopython=True, nogil=True, fastmath=True, cache=True)
-def roll_argmin(x1, window, reverse):
-    out = full(x1.shape, np.nan, dtype=np.float64)
-    if len(x1) < window:
-        return out
-    a1 = sliding_window_view(x1, window)
+@jit(nopython=True, nogil=True, cache=True)
+def roll_argmin(x1, window, min_periods, reverse):
+    out1 = full_with_window_size(x1, np.nan, dtype=np.float64, window_size=window)
+    a1 = sliding_window_with_min_periods(out1, window, min_periods)
+    out1[:] = np.nan
     if reverse:
         a1 = a1[:, ::-1]
     for i, v1 in enumerate(a1):
-        out[i + window - 1] = argmin(v1)
-    return out
+        if np.isnan(v1).all():
+            continue
+        out1[i] = argmin(v1)
+    return out1[:x1.shape[0]]
 
 
-@jit(nopython=True, nogil=True, fastmath=True, cache=True)
-def roll_prod(x1, window):
-    out = full(x1.shape, np.nan, dtype=np.float64)
-    if len(x1) < window:
-        return out
-    a1 = sliding_window_view(x1, window)
+@jit(nopython=True, nogil=True, cache=True)
+def roll_prod(x1, window, min_periods):
+    out1 = full_with_window_size(x1, np.nan, dtype=np.float64, window_size=window)
+    a1 = sliding_window_with_min_periods(out1, window, min_periods)
+    out1[:] = np.nan
     for i, v1 in enumerate(a1):
-        out[i + window - 1] = prod(v1)
-    return out
+        if np.isnan(v1).all():
+            continue
+        out1[i] = nanprod(v1)
+    return out1[:x1.shape[0]]
 
 
-@jit(nopython=True, nogil=True, fastmath=True, cache=True)
+@jit(nopython=True, nogil=True, cache=True)
 def _co_kurtosis(a1, a2):
-    t1 = a1 - mean(a1)
-    t2 = a2 - mean(a2)
-    t3 = std(a1)
-    t4 = std(a2)
-    return mean(t1 * (t2 ** 3)) / (t3 * (t4 ** 3))
+    t1 = a1 - nanmean(a1)
+    t2 = a2 - nanmean(a2)
+    t3 = nanstd(a1)
+    t4 = nanstd(a2)
+    return nanmean(t1 * (t2 ** 3)) / (t3 * (t4 ** 3))
 
 
-@jit(nopython=True, nogil=True, fastmath=True, cache=True)
-def roll_co_kurtosis(x1, x2, window):
-    out = full(x1.shape, np.nan, dtype=np.float64)
-    if len(x1) < window:
-        return out
-    a1 = sliding_window_view(x1, window)
-    a2 = sliding_window_view(x2, window)
+@jit(nopython=True, nogil=True, cache=True)
+def roll_co_kurtosis(x1, x2, window, min_periods):
+    out1 = full_with_window_size(x1, np.nan, dtype=np.float64, window_size=window)
+    out2 = full_with_window_size(x2, np.nan, dtype=np.float64, window_size=window)
+    a1 = sliding_window_with_min_periods(out1, window, min_periods)
+    a2 = sliding_window_with_min_periods(out2, window, min_periods)
+    out1[:] = np.nan
     for i, (v1, v2) in enumerate(zip(a1, a2)):
-        out[i + window - 1] = _co_kurtosis(v1, v2)
-    return out
+        if np.isnan(v1).all():
+            continue
+        if np.isnan(v2).all():
+            continue
+        out1[i] = _co_kurtosis(v1, v2)
+    return out1[:x1.shape[0]]
 
 
-@jit(nopython=True, nogil=True, fastmath=True, cache=True)
+@jit(nopython=True, nogil=True, cache=True)
 def _co_skewness(a1, a2):
-    t1 = a1 - mean(a1)
-    t2 = a2 - mean(a2)
-    t3 = std(a1)
-    t4 = std(a2)
-    return mean(t1 * (t2 ** 2)) / (t3 * (t4 ** 2))
+    t1 = a1 - nanmean(a1)
+    t2 = a2 - nanmean(a2)
+    t3 = nanstd(a1)
+    t4 = nanstd(a2)
+    return nanmean(t1 * (t2 ** 2)) / (t3 * (t4 ** 2))
 
 
-@jit(nopython=True, nogil=True, fastmath=True, cache=True)
-def roll_co_skewness(x1, x2, window):
-    out = full(x1.shape, np.nan, dtype=np.float64)
-    if len(x1) < window:
-        return out
-    a1 = sliding_window_view(x1, window)
-    a2 = sliding_window_view(x2, window)
+@jit(nopython=True, nogil=True, cache=True)
+def roll_co_skewness(x1, x2, window, min_periods):
+    out1 = full_with_window_size(x1, np.nan, dtype=np.float64, window_size=window)
+    out2 = full_with_window_size(x2, np.nan, dtype=np.float64, window_size=window)
+    a1 = sliding_window_with_min_periods(out1, window, min_periods)
+    a2 = sliding_window_with_min_periods(out2, window, min_periods)
+    out1[:] = np.nan
     for i, (v1, v2) in enumerate(zip(a1, a2)):
-        out[i + window - 1] = _co_skewness(v1, v2)
-    return out
+        if np.isnan(v1).all():
+            continue
+        if np.isnan(v2).all():
+            continue
+        out1[i] = _co_skewness(v1, v2)
+    return out1[:x1.shape[0]]
 
 
-@jit(nopython=True, nogil=True, fastmath=True, cache=True)
+@jit(nopython=True, nogil=True, cache=True)
 def _moment(a1, k):
     """centeral moment
     中心矩阵"""
-    return mean((a1 - mean(a1)) ** k)
+    return nanmean((a1 - nanmean(a1)) ** k)
 
 
-@jit(nopython=True, nogil=True, fastmath=True, cache=True)
-def roll_moment(x1, window, k):
-    out = full(x1.shape, np.nan, dtype=np.float64)
-    if len(x1) < window:
-        return out
-    a1 = sliding_window_view(x1, window)
+@jit(nopython=True, nogil=True, cache=True)
+def roll_moment(x1, window, min_periods, k):
+    out1 = full_with_window_size(x1, np.nan, dtype=np.float64, window_size=window)
+    a1 = sliding_window_with_min_periods(out1, window, min_periods)
+    out1[:] = np.nan
     for i, v1 in enumerate(a1):
-        out[i + window - 1] = _moment(v1, k)
-    return out
+        if np.isnan(v1).all():
+            continue
+        out1[i] = _moment(v1, k)
+    return out1[:x1.shape[0]]
 
 
-@jit(nopython=True, nogil=True, fastmath=True, cache=True)
+@jit(nopython=True, nogil=True, cache=True)
 def _partial_corr(a1, a2, a3):
     """TODO 不知是否正确，需要检查"""
     c = corrcoef(vstack((a1, a2, a3)))
@@ -114,50 +128,55 @@ def _partial_corr(a1, a2, a3):
     return t1 / (t2 * t3)
 
 
-@jit(nopython=True, nogil=True, fastmath=True, cache=True)
-def roll_partial_corr(x1, x2, x3, window):
-    out = full(x1.shape, np.nan, dtype=np.float64)
-    if len(x1) < window:
-        return out
-    a1 = sliding_window_view(x1, window)
-    a2 = sliding_window_view(x2, window)
-    a3 = sliding_window_view(x3, window)
+@jit(nopython=True, nogil=True, cache=True)
+def roll_partial_corr(x1, x2, x3, window, min_periods):
+    out1 = full_with_window_size(x1, np.nan, dtype=np.float64, window_size=window)
+    out2 = full_with_window_size(x2, np.nan, dtype=np.float64, window_size=window)
+    out3 = full_with_window_size(x3, np.nan, dtype=np.float64, window_size=window)
+    a1 = sliding_window_with_min_periods(out1, window, min_periods)
+    a2 = sliding_window_with_min_periods(out2, window, min_periods)
+    a3 = sliding_window_with_min_periods(out3, window, min_periods)
+    out1[:] = np.nan
     for i, (v1, v2, v3) in enumerate(zip(a1, a2, a3)):
-        out[i + window - 1] = _partial_corr(v1, v2, v3)
-    return out
+        if np.isnan(v1).all():
+            continue
+        if np.isnan(v2).all():
+            continue
+        if np.isnan(v3).all():
+            continue
+        out1[i] = _partial_corr(v1, v2, v3)
+    return out1[:x1.shape[0]]
 
 
-@jit(nopython=True, nogil=True, fastmath=True, cache=True)
+@jit(nopython=True, nogil=True, cache=True)
 def _triple_corr(a1, a2, a3):
-    t1 = a1 - mean(a1)
-    t2 = a2 - mean(a2)
-    t3 = a3 - mean(a3)
-    t4 = std(a1)
-    t5 = std(a2)
-    t6 = std(a3)
-    return mean(t1 * t2 * t3) / (t4 * t5 * t6)
+    t1 = a1 - nanmean(a1)
+    t2 = a2 - nanmean(a2)
+    t3 = a3 - nanmean(a3)
+    t4 = nanstd(a1)
+    t5 = nanstd(a2)
+    t6 = nanstd(a3)
+    return nanmean(t1 * t2 * t3) / (t4 * t5 * t6)
 
 
-@jit(nopython=True, nogil=True, fastmath=True, cache=True)
-def roll_triple_corr(x1, x2, x3, window):
-    out = full(x1.shape, np.nan, dtype=np.float64)
-    if len(x1) < window:
-        return out
-    a1 = sliding_window_view(x1, window)
-    a2 = sliding_window_view(x2, window)
-    a3 = sliding_window_view(x3, window)
+@jit(nopython=True, nogil=True, cache=True)
+def roll_triple_corr(x1, x2, x3, window, min_periods):
+    out1 = full_with_window_size(x1, np.nan, dtype=np.float64, window_size=window)
+    out2 = full_with_window_size(x2, np.nan, dtype=np.float64, window_size=window)
+    out3 = full_with_window_size(x3, np.nan, dtype=np.float64, window_size=window)
+    a1 = sliding_window_with_min_periods(out1, window, min_periods)
+    a2 = sliding_window_with_min_periods(out2, window, min_periods)
+    a3 = sliding_window_with_min_periods(out3, window, min_periods)
+    out1[:] = np.nan
     for i, (v1, v2, v3) in enumerate(zip(a1, a2, a3)):
-        out[i + window - 1] = _triple_corr(v1, v2, v3)
-    return out
-
-
-@jit(nopython=True, nogil=True, fastmath=True, cache=True)
-def isnan(x):
-    # https://github.com/numba/numba/issues/2919#issuecomment-747377615
-    if int(x) == -9223372036854775808:
-        return True
-    else:
-        return False
+        if np.isnan(v1).all():
+            continue
+        if np.isnan(v2).all():
+            continue
+        if np.isnan(v3).all():
+            continue
+        out1[i] = _triple_corr(v1, v2, v3)
+    return out1[:x1.shape[0]]
 
 
 @jit(nopython=True, nogil=True, fastmath=True, cache=True)
@@ -217,7 +236,8 @@ def _sum_split_by(x1, x2, window=10, n=2):
     return out1, out2
 
 
-@jit(float64[:](bool_[:], bool_[:], bool_[:], bool_[:], bool_, bool_), nopython=True, fastmath=True, nogil=True, cache=True)
+@jit(float64[:](boolean[:], boolean[:], boolean[:], boolean[:], boolean, boolean),
+     nopython=True, nogil=True, cache=True)
 def _signals_to_size(is_long_entry: np.ndarray, is_long_exit: np.ndarray,
                      is_short_entry: np.ndarray, is_short_exit: np.ndarray,
                      accumulate: bool = False,
@@ -289,27 +309,31 @@ def _signals_to_size(is_long_entry: np.ndarray, is_long_exit: np.ndarray,
     return out
 
 
-@jit(nopython=True, nogil=True, fastmath=True, cache=True)
-def roll_decay_linear(x1, window):
-    out = full(x1.shape, np.nan, dtype=np.float64)
-    if len(x1) < window:
-        return out
+@jit(nopython=True, nogil=True, cache=True)
+def roll_decay_linear(x1, window, min_periods):
     weights = np.arange(1., window + 1)
-    # print(weights)
-    a1 = sliding_window_view(x1, window)
+
+    out1 = full_with_window_size(x1, np.nan, dtype=np.float64, window_size=window)
+    a1 = sliding_window_with_min_periods(out1, window, min_periods)
+    out1[:] = np.nan
     for i, v1 in enumerate(a1):
-        out[i + window - 1] = np.average(v1, weights=weights)
-    return out
+        if np.isnan(v1).all():
+            continue
+        mask = ~np.isnan(v1)
+        out1[i] = np.average(v1[mask], weights=weights[mask])
+    return out1[:x1.shape[0]]
 
 
-@jit(nopython=True, nogil=True, fastmath=True, cache=True)
-def roll_decay_exp_window(x1, window, factor):
-    out = full(x1.shape, np.nan, dtype=np.float64)
-    if len(x1) < window:
-        return out
+@jit(nopython=True, nogil=True, cache=True)
+def roll_decay_exp_window(x1, window, min_periods, factor):
     weights = factor ** np.arange(window - 1, -1, -1)
-    # print(weights)
-    a1 = sliding_window_view(x1, window)
+
+    out1 = full_with_window_size(x1, np.nan, dtype=np.float64, window_size=window)
+    a1 = sliding_window_with_min_periods(out1, window, min_periods)
+    out1[:] = np.nan
     for i, v1 in enumerate(a1):
-        out[i + window - 1] = np.average(v1, weights=weights)
-    return out
+        if np.isnan(v1).all():
+            continue
+        mask = ~np.isnan(v1)
+        out1[i] = np.average(v1[mask], weights=weights[mask])
+    return out1[:x1.shape[0]]