polars-ta 0.5.2__py3-none-any.whl → 0.5.4__py3-none-any.whl

polars_ta/_version.py

@@ -1 +1 @@
-__version__ = "0.5.2"
+__version__ = "0.5.4"

polars_ta/labels/__init__.py

@@ -0,0 +1 @@
+from polars_ta.labels.future import *  # noqa

polars_ta/labels/_nb.py

@@ -0,0 +1,40 @@
+import numpy as np
+from numba import jit
+from numpy import full
+
+
+@jit(nopython=True, nogil=True, cache=True)
+def _triple_barrier(close: np.ndarray, high: np.ndarray, low: np.ndarray, window: int, take_profit: float, stop_loss: float) -> np.ndarray:
+    """三重障碍打标法"""
+    out = full(close.shape[0], np.nan, dtype=np.float64)
+    for i in range(close.shape[0] - window + 1):
+        entry_price = close[i]
+        if np.isnan(entry_price):
+            # out[i] = 0
+            continue
+        upper_barrier = entry_price * (1 + take_profit)
+        lower_barrier = entry_price * (1 - stop_loss)
+        for j in range(i + 1, i + window):
+            hit_upper = high[j] >= upper_barrier
+            hit_lower = low[j] <= lower_barrier
+            if hit_upper and hit_lower:
+                # TODO 同一天无法知道是先触发止损还是先触发止盈
+                # 1. 假定离收盘价远的先触发
+                if high[j] - close[j] > close[j] - low[j]:
+                    out[i] = 1  # 最高价更远，触发止盈
+                else:
+                    out[i] = -1  # 最低价更远，触发止损
+
+                # out[i] = -1 # 2. 简化处理认为先触发止损
+                break
+            if hit_upper:
+                out[i] = 1  # 止盈
+                break
+            if hit_lower:
+                out[i] = -1  # 止损
+                break
+        else:
+            # out[i] = 0  # 1. 时间到了触发平仓
+            out[i] = np.sign(close[j] / entry_price - 1)  # 2. 时间到了触发平仓
+
+    return out

polars_ta/labels/future.py

@@ -0,0 +1,165 @@
+"""
+
+由于标签的定义比较灵活，所以以下代码主要用于参考
+
+Notes
+-----
+标签都是未来数据,在机器学习中，只能用于`y`,不能用于`X`。
+
+References
+----------
+https://mp.weixin.qq.com/s/XtgYezFsslOfW-QyIMr0VA
+https://github.com/Rachnog/Advanced-Deep-Trading/blob/master/bars-labels-diff/Labeling.ipynb
+
+"""
+from polars import Expr, struct
+
+from polars_ta.labels._nb import _triple_barrier
+from polars_ta.utils.numba_ import batches_i2_o1, struct_to_numpy
+from polars_ta.wq import cut, ts_delay, ts_log_diff, log
+
+
+def ts_log_return(close: Expr, n: int = 5) -> Expr:
+    """将未来数据当成卖出价后移到买入价位置，计算对数收益率
+
+    Examples
+    --------
+    ```python
+    df = pl.DataFrame({
+        'a': [None, 10, 11, 12, 9, 12, 13],
+    }).with_columns(
+        out1=ts_log_return(pl.col('a'), 3),
+        out2=_ts_log_return(pl.col('a'), 3),
+    )
+
+    shape: (7, 3)
+    ┌──────┬───────────┬───────────┐
+    │ a    ┆ out1      ┆ out2      │
+    │ ---  ┆ ---       ┆ ---       │
+    │ i64  ┆ f64       ┆ f64       │
+    ╞══════╪═══════════╪═══════════╡
+    │ null ┆ null      ┆ null      │
+    │ 10   ┆ -0.105361 ┆ -0.105361 │
+    │ 11   ┆ 0.087011  ┆ 0.087011  │
+    │ 12   ┆ 0.080043  ┆ 0.080043  │
+    │ 9    ┆ null      ┆ null      │
+    │ 12   ┆ null      ┆ null      │
+    │ 13   ┆ null      ┆ null      │
+    └──────┴───────────┴───────────┘
+    ```
+
+    """
+    # return (close.shift(-n) / close).log()
+    return log(ts_delay(close, -n) / close)
+
+
+def _ts_log_return(close: Expr, n: int = 5) -> Expr:
+    """计算对数收益率，但将结果后移
+
+    如果打标签方式复杂，这种最终结果后移的方法更方便
+    """
+    # return (close / close.shift(n)).log().shift(-n)
+    return ts_delay(ts_log_diff(close, n), -n)
+
+
+def ts_simple_return(close: Expr, n: int = 5, threshold: float = 0.0, *more_threshold) -> Expr:
+    """简单收益率标签。支持二分类、三分类等。对收益率使用`cut`进行分类
+    
+    Parameters
+    ----------
+    close
+    n:int
+        未来n天
+    threshold:float
+        收益率阈值，小于该值为0，大于等于该值为1
+    more_threshold:float
+        更多的阈值，用于三分类等。小于该值为1，大于等于该值为2，以此类推
+
+    Returns
+    -------
+    Expr
+        标签列, 类型为UInt32, 取值为0, 1, 2, ...
+
+    Examples
+    --------
+    ```python
+    df = pl.DataFrame({
+        'a': [None, 10., 9.99, 9., 10., 11., 11.],
+    }).with_columns(
+        out1=label_simple_return(pl.col('a'), 1, 0),
+        out2=label_simple_return(pl.col('a'), 1, -0.001, 0.001),
+    )
+
+    shape: (7, 3)
+    ┌──────┬──────┬──────┐
+    │ a    ┆ out1 ┆ out2 │
+    │ ---  ┆ ---  ┆ ---  │
+    │ f64  ┆ u32  ┆ u32  │
+    ╞══════╪══════╪══════╡
+    │ null ┆ null ┆ null │
+    │ 10.0 ┆ 0    ┆ 0    │
+    │ 9.99 ┆ 0    ┆ 0    │
+    │ 9.0  ┆ 1    ┆ 2    │
+    │ 10.0 ┆ 1    ┆ 2    │
+    │ 11.0 ┆ 0    ┆ 1    │
+    │ 11.0 ┆ null ┆ null │
+    └──────┴──────┴──────┘
+
+    """
+    return cut(close.pct_change(n).shift(-n), threshold, *more_threshold)
+
+
+def ts_triple_barrier(close: Expr, high: Expr, low: Expr, d: int = 5, take_profit: float = 0.1, stop_loss: float = 0.05) -> Expr:
+    """三重障碍打标法
+
+    Parameters
+    ----------
+    close:Expr
+        收盘价
+    high:Expr
+        最高价
+    low:Expr
+        最低价
+    d:int
+        时间窗口
+    take_profit:float
+        止盈比例
+    stop_loss:float
+        止损比例
+
+    Returns
+    -------
+    Expr
+        标签列。取值为-1止损, 1止盈，0时间到期
+
+    Notes
+    -----
+    1. `high`, `low`在粗略情况下可用`close`代替
+    2. 时间到期时，根据盈亏返回不同的标签
+
+    Examples
+    --------
+    ```python
+    df = pl.DataFrame({
+        "close": [np.nan, 1, 1, 1.0],
+        "high": [np.nan, 1, 1.1, 1],
+        "low": [np.nan, 1, 1, 0.95],
+    }).with_columns(
+        out=ts_triple_barrier(pl.col("close"), pl.col("high"), pl.col("low"), 2, 0.1, 0.05)
+    )
+
+    shape: (4, 4)
+    ┌───────┬──────┬──────┬──────┐
+    │ close ┆ high ┆ low  ┆ out  │
+    │ ---   ┆ ---  ┆ ---  ┆ ---  │
+    │ f64   ┆ f64  ┆ f64  ┆ f64  │
+    ╞═══════╪══════╪══════╪══════╡
+    │ NaN   ┆ NaN  ┆ NaN  ┆ null │
+    │ 1.0   ┆ 1.0  ┆ 1.0  ┆ 1.0  │
+    │ 1.0   ┆ 1.1  ┆ 1.0  ┆ -1.0 │
+    │ 1.0   ┆ 1.0  ┆ 0.95 ┆ null │
+    └───────┴──────┴──────┴──────┘
+    ```
+
+    """
+    return struct([close, high, low]).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 3), _triple_barrier, d, take_profit, stop_loss))

polars_ta/performance/returns.py

@@ -1,12 +1,6 @@
 from polars import Expr
 
-from polars_ta.wq.arithmetic import log1p, expm1  # noqa
-# log return
-# 对数收益
-from polars_ta.wq.time_series import ts_log_diff as ts_log_return  # noqa
-# simple percentage return
-# 简单收益
-from polars_ta.wq.time_series import ts_returns as ts_percent_return  # noqa
+from polars_ta.wq.arithmetic import log1p, expm1
 
 
 def ts_cum_return(close: Expr) -> Expr:

polars_ta/prefix/labels.py

@@ -0,0 +1 @@
+from polars_ta.labels import *  # noqa

polars_ta/ta/README.md

@@ -0,0 +1,12 @@
+# polars_ta.ta
+
+1. Files in this folder mimic `talib`, and implement `polars` versions for the same functions
+2. Since we reduce the functino calls between `Python` and `C` code, it should be faster than `talib`.
+3. We first try to import from `ta`, then from `wq`, and only implement the function if it is not available.
+4. When there is a circular dependency, we use `polars` instead.
+
+
+1. 本文件夹中模仿`talib`，实现同名函数的`polars`版
+2. 由于减少了python与c来回调用，理论上比直接调用`talib`快
+3. 优先从`ta`中导入，然后从`wq`中导入，没有的才实现
+4. 出现循环依赖时，使用`polars`

polars_ta/talib/README.md

@@ -0,0 +1,12 @@
+# polars_ta.talib
+
+Inside this package, files are generated by `tools.codegen_talib2`.
+It is a wrapper of `talib` functions, with the following features:
+
+1. Input and output are `Expr` instead of `Series`
+2. ~~Add skipna feature (not efficient, will update when `polars` support backward fill)~~
+
+本包由`tools.codegen_talib2`自动生成，是对`talib`代码的封装。实现了以下功能
+1. 输入输出由`Series`改`Expr`
+2. ~~添加跳过空值功能（效率不高，等`polars`支持反向填充，此部分将更新）~~
+

polars_ta/tdx/README.md

@@ -0,0 +1,10 @@
+# polars_ta.tdx
+
+1. Follows the `tdx` naming convention
+2. Except for some element-wise functions, all functions are time-series functions. Pay special attention to `MAX` and similar functions.
+3. First import from `tdx`, then from `wq`, and finally from `ta`. Only implement the function if it is not available.
+
+
+1. 函数名称按照通达信来
+2. 除了部分按元素计算的函数，其它都为时序函数，特别注意`MAX`等一类不要混淆
+3. 优先从`tdx`中导入，然后从`wq`中导入，最后从`ta`，没有的才实现

polars_ta/utils/numba_.py

@@ -27,6 +27,7 @@ def isnan(x):
 
 @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
@@ -34,6 +35,7 @@ def full_with_window_size(arr, fill_value, dtype=None, window_size: int = 1):
 
 @jit(nopython=True, nogil=True, cache=True)
 def sliding_window_with_min_periods(arr, window_size: int, min_periods: int):
+    """为rolling准备的数据，当数据长度不足时，用nan填充"""
     windows = sliding_window_view(arr, window_size)
     valid_counts = np.sum(~np.isnan(windows), axis=1)
     # 修改这一行，使用布尔索引而不是np.where

polars_ta/wq/arithmetic.py

@@ -1,10 +1,11 @@
 import numpy as np
-from polars import Expr, Series, fold, any_horizontal, Float64
+from polars import Expr, Series, fold, any_horizontal, Float64, Int64
+from polars import arctan2 as _arctan2
 from polars import max_horizontal, sum_horizontal, min_horizontal, mean_horizontal
 
 
 def abs_(x: Expr) -> Expr:
-    """绝对值
+    """求绝对值
 
     Examples
     --------
@@ -38,7 +39,7 @@ def abs_(x: Expr) -> Expr:
 
 
 def add(a: Expr, b: Expr, *args) -> Expr:
-    """水平多列加
+    """水平多列相加
 
     Examples
     --------
@@ -92,6 +93,16 @@ def arc_tan(x: Expr) -> Expr:
     return x.arctan()
 
 
+def arc_tan2(y: Expr, x: Expr) -> Expr:
+    """反正切二值函数"""
+    return _arctan2(y, x)
+
+
+def cbrt(x: Expr) -> Expr:
+    """立方根"""
+    return x.cbrt()
+
+
 def ceiling(x: Expr) -> Expr:
     """向上取整"""
     return x.ceil()
@@ -107,12 +118,61 @@ def cosh(x: Expr) -> Expr:
     return x.cosh()
 
 
-def densify(x: Expr) -> Expr:
+def cot(x: Expr) -> Expr:
+    """余切"""
+    return x.cot()
+
+
+def cube(x: Expr) -> Expr:
+    """立方"""
+    return x.pow(3)
+
+
+def degrees(x: Expr) -> Expr:
+    """弧度转角度"""
+    return x.degrees()
+
+
+def _densify(x: Expr) -> Expr:
     raise
 
 
+def div(x: Expr, y: Expr) -> Expr:
+    """x除以y的整数部分
+
+    Examples
+    --------
+    ```python
+    df = pl.DataFrame({
+        'a': [None, -1.5, 0., 1.5, 2.5],
+        'b': [None, -1, 0, 1, 2],
+    }).with_columns(
+        out1=div(pl.col('a'), 0),
+        out2=div(pl.col('a'), 1),
+        out3=div(pl.col('a'), pl.col('b')),
+    )
+    shape: (5, 5)
+    ┌──────┬──────┬──────┬──────┬──────┐
+    │ a    ┆ b    ┆ out1 ┆ out2 ┆ out3 │
+    │ ---  ┆ ---  ┆ ---  ┆ ---  ┆ ---  │
+    │ f64  ┆ i64  ┆ i64  ┆ i64  ┆ i64  │
+    ╞══════╪══════╪══════╪══════╪══════╡
+    │ null ┆ null ┆ null ┆ null ┆ null │
+    │ -1.5 ┆ -1   ┆ null ┆ -2   ┆ 1    │
+    │ 0.0  ┆ 0    ┆ null ┆ 0    ┆ null │
+    │ 1.5  ┆ 1    ┆ null ┆ 1    ┆ 1    │
+    │ 2.5  ┆ 2    ┆ null ┆ 2    ┆ 1    │
+    └──────┴──────┴──────┴──────┴──────┘
+    ```
+
+    """
+    return x.floordiv(y).cast(Int64, strict=False)
+
+
 def divide(x: Expr, y: Expr) -> Expr:
-    """x/y
+    """除法
+
+    x/y
 
     Examples
     --------
@@ -173,7 +233,9 @@ def exp(x: Expr) -> Expr:
 
 
 def expm1(x: Expr) -> Expr:
-    """对数收益率 转 简单收益率 convert log return to simple return
+    """对数收益率 转 简单收益率
+
+    convert log return to simple return
 
     Examples
     --------
@@ -208,6 +270,7 @@ def floor(x: Expr) -> Expr:
 
 def fraction(x: Expr) -> Expr:
     """小数部分
+
     This operator removes the whole number part and returns the remaining fraction part with sign.
 
     Examples
@@ -248,7 +311,9 @@ def fraction(x: Expr) -> Expr:
 
 
 def inverse(x: Expr) -> Expr:
-    """1/x
+    """倒数
+
+    1/x
 
     Examples
     --------
@@ -277,7 +342,7 @@ def inverse(x: Expr) -> Expr:
 
 
 def log(x: Expr) -> Expr:
-    """e为底的对数
+    """以e为底的对数
 
     Examples
     --------
@@ -309,7 +374,7 @@ def log(x: Expr) -> Expr:
 
 
 def log10(x: Expr) -> Expr:
-    """10为底的对数
+    """以10为底的对数
 
     Examples
     --------
@@ -338,7 +403,9 @@ def log10(x: Expr) -> Expr:
 
 
 def log1p(x: Expr) -> Expr:
-    """简单收益率 转 对数收益率 convert simple return to log return
+    """简单收益率 转 对数收益率
+
+    convert simple return to log return
 
     log(x+1)
 
@@ -368,13 +435,44 @@ def log1p(x: Expr) -> Expr:
     return x.log1p()
 
 
+def log2(x: Expr) -> Expr:
+    """以2为底的对数
+
+    Examples
+    --------
+    ```python
+    df = pl.DataFrame({
+        'a': [None, -1, 0, 1, 2],
+    }).with_columns(
+        out1=log2(pl.col('a')),
+    )
+    shape: (5, 2)
+    ┌──────┬──────┐
+    │ a    ┆ out1 │
+    │ ---  ┆ ---  │
+    │ i64  ┆ f64  │
+    ╞══════╪══════╡
+    │ null ┆ null │
+    │ -1   ┆ NaN  │
+    │ 0    ┆ -inf │
+    │ 1    ┆ 0.0  │
+    │ 2    ┆ 1.0  │
+    └──────┴──────┘
+    ```
+
+    """
+    return x.log(2)
+
+
 def max_(a: Expr, b: Expr, *args) -> Expr:
-    """水平多列最大值 Maximum value of all inputs. At least 2 inputs are required."""
+    """水平多列求最大值
+
+    Maximum value of all inputs. At least 2 inputs are required."""
     return max_horizontal(a, b, *args)
 
 
 def mean(a: Expr, b: Expr, *args) -> Expr:
-    """水平多列均值
+    """水平多列求均值
 
     Examples
     --------
@@ -405,12 +503,16 @@ def mean(a: Expr, b: Expr, *args) -> Expr:
 
 
 def min_(a: Expr, b: Expr, *args) -> Expr:
-    """水平多列最小值 Maximum value of all inputs. At least 2 inputs are required."""
+    """水平多列求最小值
+
+    Maximum value of all inputs. At least 2 inputs are required."""
     return min_horizontal(a, b, *args)
 
 
 def mod(x: Expr, y: Expr) -> Expr:
-    """x%y
+    """求余
+
+    x%y
 
     Examples
     --------
@@ -441,7 +543,9 @@ def mod(x: Expr, y: Expr) -> Expr:
 
 
 def multiply(a: Expr, b: Expr, *args) -> Expr:
-    """水平多列乘 Multiply all inputs. At least 2 inputs are required.
+    """水平多列相乘
+
+    Multiply all inputs. At least 2 inputs are required.
 
     Examples
     --------
@@ -482,7 +586,9 @@ def multiply(a: Expr, b: Expr, *args) -> Expr:
 
 
 def power(x: Expr, y: Expr) -> Expr:
-    """x ** y
+    """乘幂
+
+    x ** y
 
     Examples
     --------
@@ -515,13 +621,20 @@ def power(x: Expr, y: Expr) -> Expr:
     return x.pow(y.cast(Float64))
 
 
+def radians(x: Expr) -> Expr:
+    """角度转弧度"""
+    return x.radians()
+
+
 def reverse(x: Expr) -> Expr:
-    """-x"""
+    """求相反数"""
     return -x
 
 
 def round_(x: Expr, decimals: int = 0) -> Expr:
-    """四舍五入 Round input to closest integer.
+    """四舍五入
+
+    Round input to closest integer.
 
     Parameters
     ----------
@@ -561,7 +674,9 @@ def round_(x: Expr, decimals: int = 0) -> Expr:
 
 
 def round_down(x: Expr, f: int = 1) -> Expr:
-    """小于输入的f的最大倍数 Round input to greatest multiple of f less than input
+    """小于输入的f的最大倍数
+
+    Round input to greatest multiple of f less than input
 
     Parameters
     ----------
@@ -601,6 +716,8 @@ def round_down(x: Expr, f: int = 1) -> Expr:
 def s_log_1p(x: Expr) -> Expr:
     """sign(x) * log10(1 + abs(x))
 
+    一种结合符号函数和对数变换的复合函数，常用于‌保留数据符号的同时压缩数值范围‌
+
     Examples
     --------
     ```python
@@ -637,7 +754,7 @@ def s_log_1p(x: Expr) -> Expr:
 
 
 def sign(x: Expr) -> Expr:
-    """符号"""
+    """符号函数"""
     if isinstance(x, (Expr, Series)):
         return x.sign()
     else:
@@ -645,7 +762,9 @@ def sign(x: Expr) -> Expr:
 
 
 def signed_power(x: Expr, y: Expr) -> Expr:
-    """x raised to the power of y such that final result preserves sign of x.
+    """x的y次幂，符号保留
+
+    x raised to the power of y such that final result preserves sign of x.
 
     Examples
     --------
@@ -701,7 +820,7 @@ def sinh(x: Expr) -> Expr:
 
 
 def softsign(x: Expr) -> Expr:
-    """softsign是 tanh激活函数的另一个替代选择
+    """softsign激活函数
 
     Examples
     --------
@@ -735,8 +854,15 @@ def sqrt(x: Expr) -> Expr:
     return x.sqrt()
 
 
+def square(x: Expr) -> Expr:
+    """平方"""
+    return x.pow(2)
+
+
 def subtract(x: Expr, y: Expr) -> Expr:
-    """x-y"""
+    """减法
+
+    x-y"""
     return x - y
 
 
@@ -801,7 +927,7 @@ def tanh(x: Expr) -> Expr:
 
 
 def var(a: Expr, b: Expr, *args) -> Expr:
-    """水平多列方差
+    """水平多列求方差
 
     Examples
     --------
@@ -835,7 +961,7 @@ def var(a: Expr, b: Expr, *args) -> Expr:
 
 
 def std(a: Expr, b: Expr, *args) -> Expr:
-    """水平多列标准差
+    """水平多列求标准差
 
     Examples
     --------