添加筹码分布函数

docs/tdx/pattern.md

@@ -0,0 +1 @@
+::: polars_ta.tdx.pattern

polars_ta/_version.py

@@ -1 +1 @@
-__version__ = "0.4.4"
+__version__ = "0.4.5"

polars_ta/prefix/tdx.py

@@ -56,28 +56,41 @@
 from polars_ta.tdx.over_bought_over_sold import MTM as ts_MTM  # noqa
 from polars_ta.tdx.over_bought_over_sold import RSI as ts_RSI  # noqa
 from polars_ta.tdx.over_bought_over_sold import RSV as ts_RSV  # noqa
+from polars_ta.tdx.pattern import ts_WINNER_COST  # noqa
 from polars_ta.tdx.pressure_support import BOLL as ts_BOLL  # noqa
 from polars_ta.tdx.pressure_support import BOLL_M as ts_BOLL_M  # noqa
 from polars_ta.tdx.reference import BARSLAST as ts_BARSLAST  # noqa
 from polars_ta.tdx.reference import BARSLASTCOUNT as ts_BARSLASTCOUNT  # noqa
 from polars_ta.tdx.reference import BARSSINCE as ts_BARSSINCE  # noqa
 from polars_ta.tdx.reference import BARSSINCEN as ts_BARSSINCEN  # noqa
+from polars_ta.tdx.reference import COUNT as ts_COUNT  # noqa
 from polars_ta.tdx.reference import CUMSUM as ts_CUMSUM  # noqa
+from polars_ta.tdx.reference import DIFF as ts_DIFF  # noqa
 from polars_ta.tdx.reference import DMA as ts_DMA  # noqa
 from polars_ta.tdx.reference import EMA as ts_EMA  # noqa
 from polars_ta.tdx.reference import EXPMA as ts_EXPMA  # noqa
 from polars_ta.tdx.reference import EXPMEMA as ts_EXPMEMA  # noqa
 from polars_ta.tdx.reference import FILTER as ts_FILTER  # noqa
+from polars_ta.tdx.reference import HHV as ts_HHV  # noqa
+from polars_ta.tdx.reference import HHVBARS as ts_HHVBARS  # noqa
 from polars_ta.tdx.reference import HOD as ts_HOD  # noqa
+from polars_ta.tdx.reference import LLV as ts_LLV  # noqa
+from polars_ta.tdx.reference import LLVBARS as ts_LLVBARS  # noqa
 from polars_ta.tdx.reference import LOD as ts_LOD  # noqa
+from polars_ta.tdx.reference import MA as ts_MA  # noqa
 from polars_ta.tdx.reference import MAX  # noqa
 from polars_ta.tdx.reference import MEMA as ts_MEMA  # noqa
 from polars_ta.tdx.reference import MIN  # noqa
+from polars_ta.tdx.reference import MULAR as ts_MULAR  # noqa
 from polars_ta.tdx.reference import RANGE  # noqa
+from polars_ta.tdx.reference import REF as ts_REF  # noqa
+from polars_ta.tdx.reference import REFX as ts_REFX  # noqa
 from polars_ta.tdx.reference import SMA_CN as ts_SMA_CN  # noqa
+from polars_ta.tdx.reference import SUM as ts_SUM  # noqa
 from polars_ta.tdx.reference import SUMIF as ts_SUMIF  # noqa
 from polars_ta.tdx.reference import TMA as ts_TMA  # noqa
 from polars_ta.tdx.reference import TR as ts_TR  # noqa
+from polars_ta.tdx.reference import WMA as ts_WMA  # noqa
 from polars_ta.tdx.statistic import AVEDEV as ts_AVEDEV  # noqa
 from polars_ta.tdx.statistic import COVAR as ts_COVAR  # noqa
 from polars_ta.tdx.statistic import DEVSQ as ts_DEVSQ  # noqa

polars_ta/tdx/__init__.py

@@ -4,6 +4,7 @@
 from polars_ta.tdx.logical import *  # noqa
 from polars_ta.tdx.moving_average import *  # noqa
 from polars_ta.tdx.over_bought_over_sold import *  # noqa
+from polars_ta.tdx.pattern import *  # noqa
 from polars_ta.tdx.pressure_support import *  # noqa
 from polars_ta.tdx.reference import *  # noqa
 from polars_ta.tdx.statistic import *  # noqa

polars_ta/tdx/_chip.py

@@ -0,0 +1,121 @@
+import numba
+import numpy as np
+
+
+@numba.jit(nopython=True, nogil=True, fastmath=True, cache=True)
+def nb_chip(high, low, avg, turnover,
+            start=None, stop=None, step=0.2):
+    """筹码分布，可用于WINNER或COST指标
+
+    不可能完全还原真实的筹码分布，只能接近。所以做了一下特别处理
+
+    1. 三角分布，比平均分布更接近
+    2. 步长。没有必要每个价格都统计，特别是复权后价格也无法正好是0.01间隔
+        高价股建议步长设大些，低价股步长需设小些
+
+    Parameters
+    ----------
+    high
+    low
+    avg
+        一维序列
+    turnover:
+        换手率，需要在外转成0~1范围内
+    start
+        开始价格
+    stop
+        结束价格
+    step
+        步长。一字涨停时，三角分布的底为1，高为2。但无法当成梯形计算面积，所以从中用半步长切开计算
+
+    Returns
+    -------
+    out
+        筹码分布
+    columns
+        价格表头
+
+    """
+    # 网格范围
+    if start is None:
+        start = np.min(low)
+    if stop is None:
+        stop = np.max(high)
+
+    left = round(start / step) * 2 - 1
+    right = round(stop / step) * 2 + 1
+
+    # 最小最大值左右要留半格，range是左闭右开，长度必须为2n+1
+    columns = np.arange(left, right + 1)
+    grid_shape = (len(turnover), len(columns))
+
+    # numba中round写法特殊
+    _high = np.empty_like(high)
+    _low = np.empty_like(low)
+    _avg = np.empty_like(avg)
+
+    # high和low必须落到边缘上
+    _high = np.round(high / step, 0, _high) * 2 + 1
+    _low = np.round(low / step, 0, _low) * 2 - 1
+    # avg必须落在实体上
+    _avg = np.round(avg / step, 0, _avg) * 2
+    tri_height = 2 / ((_high - _low) // 2)  # 三角形高度
+
+    # 得到三组值在网格中的位置
+    high_arg = np.argwhere(columns == _high.reshape(-1, 1))[:, 1]
+    avg_arg = np.argwhere(columns == _avg.reshape(-1, 1))[:, 1]
+    low_arg = np.argwhere(columns == _low.reshape(-1, 1))[:, 1]
+
+    # 高度表
+    height = np.zeros(grid_shape)
+    for i in range(len(height)):
+        la = low_arg[i]
+        aa = avg_arg[i]
+        ha = high_arg[i]
+        th = tri_height[i]
+        height[i, la:aa + 1] = np.linspace(0, th, aa - la + 1)
+        height[i, aa:ha + 1] = np.linspace(th, 0, ha - aa + 1)
+
+    # 计算半块面积, 三角形的高变成了梯形的上下底，梯形高固定为0.5，*0.5/2=/4
+    # 宽度-1，例如，原长度为5，-1后为4
+    area = (height[:, :-1] + height[:, 1:]) / 4
+    # 合成一块。宽度/2，例如原长度为4，/2后为2
+    weight = area[:, ::2] + area[:, 1::2]
+
+    # 输出
+    out = np.zeros_like(weight)
+    # 剩余换手率
+    turnover2 = 1 - turnover
+    # 第一天其实应当用上市发行价，过于麻烦，还是将第一天等权
+    # 取巧方法，利用-1的特性，可减少if判断，
+    out[-1] = weight[0]
+    # 这里现在用的numpy, 还要快可考虑numba
+    for i in range(len(turnover)):
+        out[i] = out[i - 1] * turnover2[i] + weight[i] * turnover[i]
+
+    # print(out.sum(axis=1))
+    return out, (step / 2) * columns[1::2]
+
+
+@numba.jit(nopython=True, nogil=True, fastmath=True, cache=True)
+def _WINNER_COST(high, low, avg, turnover, close, cost, step):
+    out, columns = nb_chip(high, low, avg, turnover, step=step)
+
+    # WINNER
+    cheap = np.where(columns <= close.reshape(-1, 1), out, 0)
+    sum_cheap = np.sum(cheap, axis=1)
+
+    # COST
+    # cum = np.cumsum(out, axis=1)
+    cum = np.copy(out)
+    for i in range(0, out.shape[0]):
+        cum[i, :] = np.cumsum(out[i, :])
+
+    prices = np.where(cum <= cost.reshape(-1, 1), columns, 0)
+
+    # np.max(prices, axis=1)
+    max_price = prices[:, 0]
+    for i in range(0, out.shape[0]):
+        max_price[i] = np.max(prices[i, :])
+
+    return sum_cheap, max_price

polars_ta/tdx/logical.py

@@ -15,22 +15,25 @@ def CROSS(a: Expr, b: Expr) -> Expr:
         'a': [None, -1, 1, 1, 2],
         'b': [None, -1, 0, 1, 2],
         'c': [None, 0, 0, 0, 0],
+        'd': [None, False, False, True, True],
     }).with_columns(
         out1=CROSS(pl.col('a'), pl.col('c')),
         out2=CROSS(pl.col('b'), pl.col('c')),
+        out3=CROSS(0, pl.col('b')),
+        out4=CROSS(pl.col('d'), 0.5),
     )
-    shape: (5, 5)
-    ┌──────┬──────┬──────┬───────┬───────┐
-    │ a    ┆ b    ┆ c    ┆ out1  ┆ out2  │
-    │ ---  ┆ ---  ┆ ---  ┆ ---   ┆ ---   │
-    │ i64  ┆ i64  ┆ i64  ┆ bool  ┆ bool  │
-    ╞══════╪══════╪══════╪═══════╪═══════╡
-    │ null ┆ null ┆ null ┆ null  ┆ null  │
-    │ -1   ┆ -1   ┆ 0    ┆ false ┆ false │
-    │ 1    ┆ 0    ┆ 0    ┆ true  ┆ false │
-    │ 1    ┆ 1    ┆ 0    ┆ false ┆ true  │
-    │ 2    ┆ 2    ┆ 0    ┆ false ┆ false │
-    └──────┴──────┴──────┴───────┴───────┘
+    shape: (5, 8)
+    ┌──────┬──────┬──────┬───────┬───────┬───────┬───────┬───────┐
+    │ a    ┆ b    ┆ c    ┆ d     ┆ out1  ┆ out2  ┆ out3  ┆ out4  │
+    │ ---  ┆ ---  ┆ ---  ┆ ---   ┆ ---   ┆ ---   ┆ ---   ┆ ---   │
+    │ i64  ┆ i64  ┆ i64  ┆ bool  ┆ bool  ┆ bool  ┆ bool  ┆ bool  │
+    ╞══════╪══════╪══════╪═══════╪═══════╪═══════╪═══════╪═══════╡
+    │ null ┆ null ┆ null ┆ null  ┆ null  ┆ null  ┆ null  ┆ null  │
+    │ -1   ┆ -1   ┆ 0    ┆ false ┆ false ┆ false ┆ null  ┆ false │
+    │ 1    ┆ 0    ┆ 0    ┆ false ┆ true  ┆ false ┆ false ┆ false │
+    │ 1    ┆ 1    ┆ 0    ┆ true  ┆ false ┆ true  ┆ false ┆ true  │
+    │ 2    ┆ 2    ┆ 0    ┆ true  ┆ false ┆ false ┆ false ┆ false │
+    └──────┴──────┴──────┴───────┴───────┴───────┴───────┴───────┘
     ```
 
     """

polars_ta/tdx/pattern.py

@@ -0,0 +1,49 @@
+from polars import Expr, Struct, Field, Float64, struct
+
+from polars_ta.tdx._chip import _WINNER_COST
+from polars_ta.utils.numba_ import batches_i2_o2
+
+
+def ts_WINNER_COST(high: Expr, low: Expr, avg: Expr, turnover: Expr, close: Expr, cost: Expr = 0.5, step: float = 0.1) -> Expr:
+    """
+    获利盘比例
+        WINNER(CLOSE),表示以当前收市价卖出的获利盘比例,例如返回0.1表示10%获利盘;WINNER(10.5)表示10.5元价格的获利盘比例
+
+    成本分布价
+         COST(10),表示10%获利盘的价格是多少,即有10%的持仓量在该价格以下,其余90%在该价格以上,为套牢盘
+
+    Parameters
+    ----------
+    high
+        最高价
+    low
+        最低价
+    avg
+        平均价。可以用vwap
+    turnover:
+        换手率。需要在外转成0~1范围内
+    close
+        判断获利比例的价格，可以用收盘价，也可以用均价
+    cost
+        成本比例，0~1
+    step
+        步长。一字涨停时，三角分布的底为1，高为2。但无法当成梯形计算面积，所以从中用半步长切开计算
+
+    Returns
+    -------
+    winner
+        获利盘比例
+    cost
+        成本分布价
+
+    Examples
+    --------
+    >>> WINNER_COST = ts_WINNER_COST(HIGH, LOW, VWAP, turnover_ratio / 100, CLOSE, 0.5)
+
+    Notes
+    -----
+    该函数仅对日线分析周期有效
+
+    """
+    dtype = Struct([Field(f"column_{i}", Float64) for i in range(2)])
+    return struct([high, low, avg, turnover, close, cost]).map_batches(lambda xx: batches_i2_o2([xx.struct[i].to_numpy().astype(float) for i in range(6)], _WINNER_COST, step), return_dtype=dtype)

polars_ta/tdx/reference.py

@@ -154,7 +154,7 @@ def SMA_CN(X: Expr, N: int, M: int) -> Expr:
 
 
 def SUMIF(condition: Expr, close: Expr, N: int = 30) -> Expr:
-    return SUM(condition.cast(Int32) * close, N)
+    return SUM(condition.cast(Boolean).cast(Int32) * close, N)
 
 
 def TMA(close: Expr, N: int = 30) -> Expr:
@@ -164,3 +164,8 @@ def TMA(close: Expr, N: int = 30) -> Expr:
 
 def FILTER(close: Expr, N: int = 30) -> Expr:
     raise
+
+
+def REFX(close: Expr, N: int = 30) -> Expr:
+    """属于未来函数,引用若干周期后的数据"""
+    return REF(close, -N)

tests/tdx/chip_test.py

@@ -0,0 +1,17 @@
+import numpy as np
+import polars as pl
+
+from polars_ta.tdx.pattern import ts_WINNER_COST
+
+high = np.array([10.4, 10.2, 10.2, 10.4, 10.5, 10.7, 10.7, 10.7, 10.8, 10.9])
+low = np.array([10.3, 10.1, 10.2, 10.3, 10.4, 10.5, 10.6, 10.7, 10.8, 10.9])
+avg = np.array([10.3, 10.1, 10.2, 10.3, 10.4, 10.5, 10.6, 10.7, 10.8, 10.9])
+close = np.array([10.3, 10.1, 10.2, 10.3, 10.4, 10.5, 10.6, 10.7, 10.8, 10.9])
+turnover = np.array([0.3, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9])
+cost = np.array([0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5])
+
+step = 0.1
+
+df = pl.DataFrame({'high': high, 'low': low, 'avg': avg, 'close': close, 'turnover': turnover, 'cost': cost})
+df = df.with_columns(WINNER=ts_WINNER_COST(pl.col('high'), pl.col('low'), pl.col('avg'), pl.col('turnover'), pl.col('close'), 0.5, step=step))
+print(df)

tools/prefix_tdx.py

@@ -8,8 +8,9 @@
 lines += codegen_import_as('polars_ta.tdx.logical', include_func=['CROSS', 'LONGCROSS', 'EXISTR', 'LAST'], include_parameter=['N'])
 lines += codegen_import_as('polars_ta.tdx.moving_average', include_parameter=['M1'])
 lines += codegen_import_as('polars_ta.tdx.over_bought_over_sold', include_modules=['polars_ta.ta.momentum'], include_parameter=['N', 'timeperiod'])
+lines += codegen_import_as('polars_ta.tdx.pattern', include_parameter=['ts_WINNER_COST'])
 lines += codegen_import_as('polars_ta.tdx.pressure_support', include_parameter=['N'])
-lines += codegen_import_as('polars_ta.tdx.reference', include_modules=['polars_ta.ta.overlap', 'polars_ta.ta.volatility', 'polars_ta.wq.arithmetic', 'polars_ta.wq.time_serie'], include_func=['BARSLAST', 'BARSLASTCOUNT', 'BARSSINCE', 'DMA', 'CUMSUM', 'TR'], include_parameter=['N', 'd', 'timeperiod'])
+lines += codegen_import_as('polars_ta.tdx.reference', include_modules=['polars_ta.ta.overlap', 'polars_ta.ta.volatility', 'polars_ta.wq.arithmetic', 'polars_ta.wq.time_series'], include_func=['BARSLAST', 'BARSLASTCOUNT', 'BARSSINCE', 'DMA', 'CUMSUM', 'TR'], include_parameter=['N', 'd', 'timeperiod'])
 lines += codegen_import_as('polars_ta.tdx.statistic', include_modules=['polars_ta.wq.time_series'], include_parameter=['timeperiod', 'd'])
 lines += codegen_import_as('polars_ta.tdx.trend', include_parameter=['N'])
 lines += codegen_import_as('polars_ta.tdx.volume', include_func=['OBV'], include_parameter=['N'])