polars-ta 0.5.3__py3-none-any.whl → 0.5.5__py3-none-any.whl

polars_ta/_version.py

@@ -1 +1 @@
-__version__ = "0.5.3"
+__version__ = "0.5.5"

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/prefix/vec.py

@@ -0,0 +1,17 @@
+# this code is auto generated by tools/prefix_vec.py
+
+from polars_ta.wq.vector import vec_avg as cs_vec_avg  # noqa
+from polars_ta.wq.vector import vec_choose as cs_vec_choose  # noqa
+from polars_ta.wq.vector import vec_count as cs_vec_count  # noqa
+from polars_ta.wq.vector import vec_ir as cs_vec_ir  # noqa
+from polars_ta.wq.vector import vec_kurtosis as cs_vec_kurtosis  # noqa
+from polars_ta.wq.vector import vec_l2_norm as cs_vec_l2_norm  # noqa
+from polars_ta.wq.vector import vec_max as cs_vec_max  # noqa
+from polars_ta.wq.vector import vec_min as cs_vec_min  # noqa
+from polars_ta.wq.vector import vec_norm as cs_vec_norm  # noqa
+from polars_ta.wq.vector import vec_percentage as cs_vec_percentage  # noqa
+from polars_ta.wq.vector import vec_powersum as cs_vec_powersum  # noqa
+from polars_ta.wq.vector import vec_range as cs_vec_range  # noqa
+from polars_ta.wq.vector import vec_skewness as cs_vec_skewness  # noqa
+from polars_ta.wq.vector import vec_stddev as cs_vec_stddev  # noqa
+from polars_ta.wq.vector import vec_sum as cs_vec_sum  # 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/ta/overlap.py

@@ -39,7 +39,7 @@ def EMA(close: Expr, timeperiod: int = 30) -> Expr:
 
     """
     # 相当于alpha=2/(1+timeperiod)
-    return close.ewm_mean(span=timeperiod, adjust=False, min_periods=timeperiod)
+    return close.ewm_mean(span=timeperiod, adjust=False, min_samples=timeperiod)
 
 
 def KAMA(close: Expr, timeperiod: int = 30) -> Expr:
@@ -78,7 +78,7 @@ def RMA(close: Expr, timeperiod: int = 30) -> Expr:
     https://github.com/twopirllc/pandas-ta/blob/main/pandas_ta/overlap/rma.py
 
     """
-    return close.ewm_mean(alpha=1 / timeperiod, adjust=False, min_periods=timeperiod)
+    return close.ewm_mean(alpha=1 / timeperiod, adjust=False, min_samples=timeperiod)
 
 
 def TEMA(close: Expr, timeperiod: int = 30) -> Expr:

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/tdx/reference.py

@@ -85,7 +85,7 @@ def DMA(close: Expr, alpha: float = 0.5) -> Expr:
 
     求X的动态移动平均.
     算法:Y=A*X+(1-A)*Y',其中Y'表示上一周期Y值,A必须大于0且小于1.A支持变量"""
-    return close.ewm_mean(alpha=alpha, adjust=False, min_periods=1)
+    return close.ewm_mean(alpha=alpha, adjust=False, min_samples=1)
 
 
 def EMA(close: Expr, N: int = 30) -> Expr:
@@ -113,7 +113,7 @@ def EXPMEMA(close: Expr, N: int = 30) -> Expr:
     """
     sma = MA(close, N)
     x = when(close.cum_count() < N).then(sma).otherwise(close)
-    return x.ewm_mean(span=N, adjust=False, min_periods=1)
+    return x.ewm_mean(span=N, adjust=False, min_samples=1)
 
 
 def HOD(close: Expr, N: int = 30) -> Expr:
@@ -156,7 +156,7 @@ def SMA_CN(X: Expr, N: int, M: int) -> Expr:
 
     !!!为防止与talib版SMA误用，这里去了默认值1
     """
-    return X.ewm_mean(alpha=M / N, adjust=False, min_periods=1)
+    return X.ewm_mean(alpha=M / N, adjust=False, min_samples=1)
 
 
 def SUMIF(condition: Expr, close: Expr, N: int = 30) -> Expr:

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/utils/withs.py

@@ -0,0 +1,44 @@
+import re
+
+import polars as pl
+
+
+def with_industry(df: pl.DataFrame, industry_name: str, drop_first: bool, keep_col: bool) -> pl.DataFrame:
+    """添加行业哑元变量
+
+    Parameters
+    ----------
+    df
+    industry_name:str
+        行业列名
+    drop_first
+        丢弃第一列
+    keep_col
+        是否保留源列
+
+    Returns
+    -------
+    pl.DataFrame
+
+    """
+    df = df.with_columns([
+        # 行业处理，由浮点改成整数
+        pl.col(industry_name).cast(pl.UInt32),
+    ])
+
+    # TODO 没有行业的也过滤，这会不会有问题？已经退市了，但引入了未来数据
+    df = df.filter(
+        pl.col(industry_name).is_not_null(),
+    )
+
+    if keep_col:
+        df = df.with_columns(df.to_dummies(industry_name, drop_first=False))
+    else:
+        df = df.to_dummies(industry_name, drop_first=False)
+
+    if drop_first:
+        # drop_first丢弃哪个字段是随机的，非常不友好，只能在行业中性化时动态修改代码
+        industry_columns = sorted(filter(lambda x: re.search(rf"^{industry_name}_\d+$", x), df.columns))
+        return df.drop(industry_columns[0])
+    else:
+        return df

polars_ta/wq/arithmetic.py

@@ -5,7 +5,7 @@ from polars import max_horizontal, sum_horizontal, min_horizontal, mean_horizont
 
 
 def abs_(x: Expr) -> Expr:
-    """绝对值
+    """求绝对值
 
     Examples
     --------
@@ -39,7 +39,7 @@ def abs_(x: Expr) -> Expr:
 
 
 def add(a: Expr, b: Expr, *args) -> Expr:
-    """水平多列加
+    """水平多列相加
 
     Examples
     --------
@@ -133,12 +133,12 @@ def degrees(x: Expr) -> Expr:
     return x.degrees()
 
 
-def densify(x: Expr) -> Expr:
+def _densify(x: Expr) -> Expr:
     raise
 
 
 def div(x: Expr, y: Expr) -> Expr:
-    """x/y的整数部分
+    """x除以y的整数部分
 
     Examples
     --------
@@ -170,7 +170,9 @@ def div(x: Expr, y: Expr) -> Expr:
 
 
 def divide(x: Expr, y: Expr) -> Expr:
-    """x/y
+    """除法
+
+    x/y
 
     Examples
     --------
@@ -231,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
     --------
@@ -266,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
@@ -306,7 +311,9 @@ def fraction(x: Expr) -> Expr:
 
 
 def inverse(x: Expr) -> Expr:
-    """1/x
+    """倒数
+
+    1/x
 
     Examples
     --------
@@ -335,7 +342,7 @@ def inverse(x: Expr) -> Expr:
 
 
 def log(x: Expr) -> Expr:
-    """e为底的对数
+    """以e为底的对数
 
     Examples
     --------
@@ -367,7 +374,7 @@ def log(x: Expr) -> Expr:
 
 
 def log10(x: Expr) -> Expr:
-    """10为底的对数
+    """以10为底的对数
 
     Examples
     --------
@@ -396,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)
 
@@ -427,7 +436,7 @@ def log1p(x: Expr) -> Expr:
 
 
 def log2(x: Expr) -> Expr:
-    """2为底的对数
+    """以2为底的对数
 
     Examples
     --------
@@ -456,12 +465,14 @@ def log2(x: Expr) -> Expr:
 
 
 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
     --------
@@ -492,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
     --------
@@ -528,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
     --------
@@ -569,7 +586,9 @@ def multiply(a: Expr, b: Expr, *args) -> Expr:
 
 
 def power(x: Expr, y: Expr) -> Expr:
-    """x ** y
+    """乘幂
+
+    x ** y
 
     Examples
     --------
@@ -608,12 +627,14 @@ def radians(x: Expr) -> Expr:
 
 
 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
     ----------
@@ -653,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
     ----------
@@ -693,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
@@ -729,7 +754,7 @@ def s_log_1p(x: Expr) -> Expr:
 
 
 def sign(x: Expr) -> Expr:
-    """符号"""
+    """符号函数"""
     if isinstance(x, (Expr, Series)):
         return x.sign()
     else:
@@ -737,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
     --------
@@ -793,7 +820,7 @@ def sinh(x: Expr) -> Expr:
 
 
 def softsign(x: Expr) -> Expr:
-    """softsign是 tanh激活函数的另一个替代选择
+    """softsign激活函数
 
     Examples
     --------
@@ -833,7 +860,9 @@ def square(x: Expr) -> Expr:
 
 
 def subtract(x: Expr, y: Expr) -> Expr:
-    """x-y"""
+    """减法
+
+    x-y"""
     return x - y
 
 
@@ -898,7 +927,7 @@ def tanh(x: Expr) -> Expr:
 
 
 def var(a: Expr, b: Expr, *args) -> Expr:
-    """水平多列方差
+    """水平多列求方差
 
     Examples
     --------
@@ -932,7 +961,7 @@ def var(a: Expr, b: Expr, *args) -> Expr:
 
 
 def std(a: Expr, b: Expr, *args) -> Expr:
-    """水平多列标准差
+    """水平多列求标准差
 
     Examples
     --------