diff --git a/docs/talib/index.md b/docs/talib/index.md index e45450a..ad60fef 100644 --- a/docs/talib/index.md +++ b/docs/talib/index.md @@ -1,3 +1,3 @@ -直接调用`TA-Lib` +直接调用`TA-Lib`, 多输出时返回`struct`, 可以使用`.struct[0]`取到对应字段的`Series`. ::: polars_ta.talib \ No newline at end of file diff --git a/docs/tdx/choice.md b/docs/tdx/choice.md new file mode 100644 index 0000000..54bb548 --- /dev/null +++ b/docs/tdx/choice.md @@ -0,0 +1 @@ +::: polars_ta.tdx.choice \ No newline at end of file diff --git a/docs/tdx/energy.md b/docs/tdx/energy.md new file mode 100644 index 0000000..906088e --- /dev/null +++ b/docs/tdx/energy.md @@ -0,0 +1 @@ +::: polars_ta.tdx.energy \ No newline at end of file diff --git a/docs/tdx/logical.md b/docs/tdx/logical.md new file mode 100644 index 0000000..5635d6b --- /dev/null +++ b/docs/tdx/logical.md @@ -0,0 +1 @@ +::: polars_ta.tdx.logical \ No newline at end of file diff --git a/docs/tdx/moving_average.md b/docs/tdx/moving_average.md new file mode 100644 index 0000000..d4a1277 --- /dev/null +++ b/docs/tdx/moving_average.md @@ -0,0 +1 @@ +::: polars_ta.tdx.moving_average \ No newline at end of file diff --git a/docs/tdx/over_bought_over_sold.md b/docs/tdx/over_bought_over_sold.md new file mode 100644 index 0000000..683daac --- /dev/null +++ b/docs/tdx/over_bought_over_sold.md @@ -0,0 +1 @@ +::: polars_ta.tdx.over_bought_over_sold \ No newline at end of file diff --git a/docs/tdx/pressure_support.md b/docs/tdx/pressure_support.md new file mode 100644 index 0000000..3bccc43 --- /dev/null +++ b/docs/tdx/pressure_support.md @@ -0,0 +1 @@ +::: polars_ta.tdx.pressure_support \ No newline at end of file diff --git a/docs/tdx/reference.md b/docs/tdx/reference.md new file mode 100644 index 0000000..13da2c4 --- /dev/null +++ b/docs/tdx/reference.md @@ -0,0 +1 @@ +::: polars_ta.tdx.reference \ No newline at end of file diff --git a/docs/tdx/statistic.md b/docs/tdx/statistic.md new file mode 100644 index 0000000..1627de6 --- /dev/null +++ b/docs/tdx/statistic.md @@ -0,0 +1 @@ +::: polars_ta.tdx.statistic \ No newline at end of file diff --git a/docs/tdx/trend.md b/docs/tdx/trend.md new file mode 100644 index 0000000..1dc1dbc --- /dev/null +++ b/docs/tdx/trend.md @@ -0,0 +1 @@ +::: polars_ta.tdx.trend \ No newline at end of file diff --git a/docs/tdx/volume.md b/docs/tdx/volume.md new file mode 100644 index 0000000..8b19e20 --- /dev/null +++ b/docs/tdx/volume.md @@ -0,0 +1 @@ +::: polars_ta.tdx.volume \ No newline at end of file diff --git a/mkdocs.yml b/mkdocs.yml index 35fcbb4..4ca542d 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -23,6 +23,16 @@ nav: - ta/volume.md - polars_ta.tdx: - tdx/arithmetic.md + - tdx/choice.md + - tdx/energy.md + - tdx/logical.md + - tdx/moving_average.md + - tdx/over_bought_over_sold.md + - tdx/pressure_support.md + - tdx/reference.md + - tdx/statistic.md + - tdx/trend.md + - tdx/volume.md - polars_ta.talib: talib/index.md theme: diff --git a/polars_ta/_version.py b/polars_ta/_version.py index 260c070..6a9beea 100644 --- a/polars_ta/_version.py +++ b/polars_ta/_version.py @@ -1 +1 @@ -__version__ = "0.3.1" +__version__ = "0.4.0" diff --git a/polars_ta/prefix/ta.py b/polars_ta/prefix/ta.py index 475e0af..f33752c 100644 --- a/polars_ta/prefix/ta.py +++ b/polars_ta/prefix/ta.py @@ -17,6 +17,15 @@ from polars_ta.ta.momentum import STOCHF_fastd as ts_STOCHF_fastd # noqa from polars_ta.ta.momentum import TRIX as ts_TRIX # noqa from polars_ta.ta.momentum import WILLR as ts_WILLR # noqa +from polars_ta.ta.operators import ADD # noqa +from polars_ta.ta.operators import DIV # noqa +from polars_ta.ta.operators import MAX as ts_MAX # noqa +from polars_ta.ta.operators import MAXINDEX as ts_MAXINDEX # noqa +from polars_ta.ta.operators import MIN as ts_MIN # noqa +from polars_ta.ta.operators import MININDEX as ts_MININDEX # noqa +from polars_ta.ta.operators import MUL # noqa +from polars_ta.ta.operators import SUB # noqa +from polars_ta.ta.operators import SUM as ts_SUM # noqa from polars_ta.ta.overlap import BBANDS_upperband as ts_BBANDS_upperband # noqa from polars_ta.ta.overlap import DEMA as ts_DEMA # noqa from polars_ta.ta.overlap import EMA as ts_EMA # noqa diff --git a/polars_ta/prefix/tdx.py b/polars_ta/prefix/tdx.py index 1872d32..096e81e 100644 --- a/polars_ta/prefix/tdx.py +++ b/polars_ta/prefix/tdx.py @@ -11,7 +11,6 @@ from polars_ta.tdx.arithmetic import EXP # noqa from polars_ta.tdx.arithmetic import FLOOR # noqa from polars_ta.tdx.arithmetic import FRACPART # noqa -from polars_ta.tdx.arithmetic import INTPART # noqa from polars_ta.tdx.arithmetic import LN # noqa from polars_ta.tdx.arithmetic import LOG # noqa from polars_ta.tdx.arithmetic import MAX # noqa diff --git a/polars_ta/ta/momentum.py b/polars_ta/ta/momentum.py index b2d820b..4fbeee2 100644 --- a/polars_ta/ta/momentum.py +++ b/polars_ta/ta/momentum.py @@ -1,10 +1,4 @@ -""" - -```python -``` - -""" -from polars import Expr +from polars import Expr, when from polars_ta import TA_EPSILON from polars_ta.ta.operators import MAX @@ -149,50 +143,57 @@ def RSI(close: Expr, timeperiod: int = 14) -> Expr: return RMA(max_(dif, 0), timeperiod) / (RMA(dif.abs(), timeperiod) + TA_EPSILON) # * 100 +def STOCHF_fastd(high: Expr, low: Expr, close: Expr, fastk_period: int = 5, fastd_period: int = 3) -> Expr: + return SMA(RSV(high, low, close, fastk_period), fastd_period) + + +def TRIX(close: Expr, timeperiod: int = 30) -> Expr: + EMA1 = EMA(close, timeperiod) + EMA2 = EMA(EMA1, timeperiod) + EMA3 = EMA(EMA2, timeperiod) + return ROCP(EMA3, 1) + + def RSV(high: Expr, low: Expr, close: Expr, timeperiod: int = 5) -> Expr: - """ + """RSV=STOCHF_FASTK + + (Today's Close - LowestLow) + FASTK(Kperiod) = --------------------------- * 100 + (HighestHigh - LowestLow) Notes ----- - Also called `STOCHF_fastk`, `talib.STOCHF` is multiplied by 100, - and it is very similar to the `WILLR` indicator + RSV = STOCHF_FASTK。没乘以100 - 又名STOCHF_fastk, talib.STOCHF版相当于多乘了100,与WILLR指标又很像 + References + ---------- + https://github.com/TA-Lib/ta-lib/blob/main/src/ta_func/ta_STOCHF.c#L279 """ a = MAX(high, timeperiod) b = MIN(low, timeperiod) - return (close - b) / (a - b + TA_EPSILON) - - -def STOCHF_fastd(high: Expr, low: Expr, close: Expr, fastk_period: int = 5, fastd_period: int = 3) -> Expr: - return SMA(RSV(high, low, close, fastk_period), fastd_period) - - -def TRIX(close: Expr, timeperiod: int = 30) -> Expr: - EMA1 = EMA(close, timeperiod) - EMA2 = EMA(EMA1, timeperiod) - EMA3 = EMA(EMA2, timeperiod) - return ROCP(EMA3, 1) + # return (close - b) / (a - b + TA_EPSILON) + return when(a != b).then((close - b) / (a - b)).otherwise(0) def WILLR(high: Expr, low: Expr, close: Expr, timeperiod: int = 14) -> Expr: - """ + """威廉指标 Notes ----- - `talib.WILLR` is multiplied by -100, but I think it is unnecessary - - talib.WILLR版相当于多乘了-100,但个人认为没有必要 + WILLR=1-RSV References ---------- + - https://github.com/TA-Lib/ta-lib/blob/main/src/ta_func/ta_WILLR.c#L294 - https://www.investopedia.com/terms/w/williamsr.asp - https://school.stockcharts.com/doku.php?id=technical_indicators:williams_r + """ a = MAX(high, timeperiod) b = MIN(low, timeperiod) - return (a - close) / (a - b + TA_EPSILON) + # return (a - close) / (a - b + TA_EPSILON) + return when(a != b).then((a - close) / (a - b)).otherwise(0) diff --git a/polars_ta/ta/operators.py b/polars_ta/ta/operators.py index 5c29d79..c542bec 100644 --- a/polars_ta/ta/operators.py +++ b/polars_ta/ta/operators.py @@ -1,61 +1,109 @@ +""" +通过`import`直接导入或更名的函数 + +```python +from polars_ta.wq.arithmetic import add as ADD # noqa +from polars_ta.wq.arithmetic import divide as DIV # noqa +from polars_ta.wq.arithmetic import multiply as MUL # noqa +from polars_ta.wq.arithmetic import subtract as SUB # noqa +from polars_ta.wq.time_series import ts_max as MAX # noqa +from polars_ta.wq.time_series import ts_min as MIN # noqa +from polars_ta.wq.time_series import ts_sum as SUM # noqa +``` + +""" from polars import Expr +from polars_ta.wq.arithmetic import add as ADD # noqa +from polars_ta.wq.arithmetic import divide as DIV # noqa +from polars_ta.wq.arithmetic import multiply as MUL # noqa +from polars_ta.wq.arithmetic import subtract as SUB # noqa from polars_ta.wq.time_series import ts_arg_max from polars_ta.wq.time_series import ts_arg_min -from polars_ta.wq.time_series import ts_max -from polars_ta.wq.time_series import ts_min -from polars_ta.wq.time_series import ts_sum +from polars_ta.wq.time_series import ts_max as MAX # noqa +from polars_ta.wq.time_series import ts_min as MIN # noqa +from polars_ta.wq.time_series import ts_sum as SUM # noqa -def ADD(high: Expr, low: Expr) -> Expr: - return high + low - - -def DIV(high: Expr, low: Expr) -> Expr: - return high / low - - -def MAX(close: Expr, timeperiod: int = 30) -> Expr: +def MAXINDEX(close: Expr, timeperiod: int = 30) -> Expr: """ Notes ----- - It is the maximum value of the time series, not the maximum value of multiple columns (max_horizontal) - - 时序上窗口最大,不要与多列最大搞混 - - """ - return ts_max(close, timeperiod) - - -def MAXINDEX(close: Expr, timeperiod: int = 30) -> Expr: - """ Comparing to `ts_arg_max` this also marks the abs. position of the max value 与ts_arg_max的区别是,标记了每个区间最大值的绝对位置,可用来画图标记 + + + Examples + -------- + ```python + from polars_ta.ta import MAXINDEX as ta_MAXINDEX + from polars_ta.talib import MAXINDEX as talib_MAXINDEX + from polars_ta.wq import ts_arg_max + + df = pl.DataFrame({ + 'a': [6, 2, 8, 5, 9, 4], + }).with_columns( + out1=ts_arg_max(pl.col('a'), 3), + out2=ta_MAXINDEX(pl.col('a'), 3), + out3=talib_MAXINDEX(pl.col('a'), 3), + ) + shape: (6, 4) + ┌─────┬──────┬──────┬──────┐ + │ a ┆ out1 ┆ out2 ┆ out3 │ + │ --- ┆ --- ┆ --- ┆ --- │ + │ i64 ┆ u16 ┆ i64 ┆ i32 │ + ╞═════╪══════╪══════╪══════╡ + │ 6 ┆ null ┆ null ┆ 0 │ + │ 2 ┆ null ┆ null ┆ 0 │ + │ 8 ┆ 0 ┆ 2 ┆ 2 │ + │ 5 ┆ 1 ┆ 2 ┆ 2 │ + │ 9 ┆ 0 ┆ 4 ┆ 4 │ + │ 4 ┆ 1 ┆ 4 ┆ 4 │ + └─────┴──────┴──────┴──────┘ + ``` """ a = close.cum_count() b = ts_arg_max(close, timeperiod) return a - b - 1 -def MIN(close: Expr, timeperiod: int = 30) -> Expr: - return ts_min(close, timeperiod) +def MININDEX(close: Expr, timeperiod: int = 30) -> Expr: + """ + -def MININDEX(close: Expr, timeperiod: int = 30) -> Expr: + Examples + -------- + ```python + from polars_ta.ta import MININDEX as ta_MININDEX + from polars_ta.talib import MININDEX as talib_MININDEX + from polars_ta.wq import ts_arg_min + + df = pl.DataFrame({ + 'a': [6, 2, 8, 5, 9, 4], + + }).with_columns( + out1=ts_arg_min(pl.col('a'), 3), + out2=ta_MININDEX(pl.col('a'), 3), + out3=talib_MININDEX(pl.col('a'), 3), + ) + shape: (6, 4) + ┌─────┬──────┬──────┬──────┐ + │ a ┆ out1 ┆ out2 ┆ out3 │ + │ --- ┆ --- ┆ --- ┆ --- │ + │ i64 ┆ u16 ┆ i64 ┆ i32 │ + ╞═════╪══════╪══════╪══════╡ + │ 6 ┆ null ┆ null ┆ 0 │ + │ 2 ┆ null ┆ null ┆ 0 │ + │ 8 ┆ 1 ┆ 1 ┆ 1 │ + │ 5 ┆ 2 ┆ 1 ┆ 1 │ + │ 9 ┆ 1 ┆ 3 ┆ 3 │ + │ 4 ┆ 0 ┆ 5 ┆ 5 │ + └─────┴──────┴──────┴──────┘ + ``` + """ a = close.cum_count() b = ts_arg_min(close, timeperiod) return a - b - 1 - - -def MULT(high: Expr, low: Expr) -> Expr: - return high * low - - -def SUB(high: Expr, low: Expr) -> Expr: - return high - low - - -def SUM(close: Expr, timeperiod: int = 30) -> Expr: - return ts_sum(close, timeperiod) diff --git a/polars_ta/ta/overlap.py b/polars_ta/ta/overlap.py index 9140b3c..d5c2def 100644 --- a/polars_ta/ta/overlap.py +++ b/polars_ta/ta/overlap.py @@ -47,10 +47,24 @@ def KAMA(close: Expr, timeperiod: int = 30) -> Expr: def MIDPOINT(close: Expr, timeperiod: int = 14) -> Expr: + """MIDPOINT = (Highest Value + Lowest Value)/2 + + References + ---------- + https://github.com/TA-Lib/ta-lib/blob/main/src/ta_func/ta_MIDPOINT.c#L198 + + """ return (MAX(close, timeperiod) + MIN(close, timeperiod)) / 2 def MIDPRICE(high: Expr, low: Expr, timeperiod: int = 14) -> Expr: + """MIDPRICE = (Highest High + Lowest Low)/2 + + References + ---------- + https://github.com/TA-Lib/ta-lib/blob/main/src/ta_func/ta_MIDPRICE.c#L202 + + """ return (MAX(high, timeperiod) + MIN(low, timeperiod)) / 2 diff --git a/polars_ta/ta/price.py b/polars_ta/ta/price.py index 2b85486..c44362f 100644 --- a/polars_ta/ta/price.py +++ b/polars_ta/ta/price.py @@ -2,20 +2,44 @@ def AVGPRICE(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: - """(open + high + low + close) / 4""" + """(open + high + low + close) / 4 + + References + ---------- + https://github.com/TA-Lib/ta-lib/blob/main/src/ta_func/ta_AVGPRICE.c#L187 + + """ return (open + high + low + close) / 4 def MEDPRICE(high: Expr, low: Expr) -> Expr: - """(high + low) / 2""" + """(high + low) / 2 + + References + ---------- + https://github.com/TA-Lib/ta-lib/blob/main/src/ta_func/ta_MEDPRICE.c#L180 + + """ return (high + low) / 2 def TYPPRICE(high: Expr, low: Expr, close: Expr) -> Expr: - """(high + low + close) / 3""" + """(high + low + close) / 3 + + References + ---------- + https://github.com/TA-Lib/ta-lib/blob/main/src/ta_func/ta_TYPPRICE.c#L185 + + """ return (high + low + close) / 3 def WCLPRICE(high: Expr, low: Expr, close: Expr) -> Expr: - """(high + low + close * 2) / 4""" + """(high + low + close * 2) / 4 + + References + ---------- + https://github.com/TA-Lib/ta-lib/blob/main/src/ta_func/ta_WCLPRICE.c#L184 + + """ return (high + low + close * 2) / 4 diff --git a/polars_ta/ta/statistic.py b/polars_ta/ta/statistic.py index 354229a..bbdda0f 100644 --- a/polars_ta/ta/statistic.py +++ b/polars_ta/ta/statistic.py @@ -1,6 +1,15 @@ +""" +通过`import`直接导入或更名的函数 + +```python +from polars_ta.wq.time_series import ts_corr as CORREL # noqa +``` + +""" + from polars import Expr -from polars_ta.wq.time_series import ts_corr +from polars_ta.wq.time_series import ts_corr as CORREL # noqa from polars_ta.wq.time_series import ts_std_dev @@ -8,10 +17,6 @@ def BETA(high: Expr, low: Expr, timeperiod: int = 5) -> Expr: raise -def CORREL(high: Expr, low: Expr, timeperiod: int = 30) -> Expr: - return ts_corr(high, low, timeperiod, 1) - - def LINEARREG(close: Expr, timeperiod: int = 14) -> Expr: raise diff --git a/polars_ta/ta/transform.py b/polars_ta/ta/transform.py index 28810a6..6351452 100644 --- a/polars_ta/ta/transform.py +++ b/polars_ta/ta/transform.py @@ -1,5 +1,24 @@ """ -基本上都来自于`polars_ta.wq.arithmetic` +通过`import`直接导入或更名的函数 + +```python +from polars_ta.wq.arithmetic import arc_cos as ACOS # noqa +from polars_ta.wq.arithmetic import arc_sin as ASIN # noqa +from polars_ta.wq.arithmetic import arc_tan as ATAN # noqa +from polars_ta.wq.arithmetic import ceiling as CEIL # noqa +from polars_ta.wq.arithmetic import cos as COS # noqa +from polars_ta.wq.arithmetic import cosh as COSH # noqa +from polars_ta.wq.arithmetic import exp as EXP # noqa +from polars_ta.wq.arithmetic import floor as FLOOR # noqa +from polars_ta.wq.arithmetic import log as LN # noqa +from polars_ta.wq.arithmetic import log10 as LOG10 # noqa +from polars_ta.wq.arithmetic import sin as SIN # noqa +from polars_ta.wq.arithmetic import sinh as SINH # noqa +from polars_ta.wq.arithmetic import sqrt as SQRT # noqa +from polars_ta.wq.arithmetic import tan as TAN # noqa +from polars_ta.wq.arithmetic import tanh as TANH # noqa +``` + """ from polars_ta.wq.arithmetic import arc_cos as ACOS # noqa from polars_ta.wq.arithmetic import arc_sin as ASIN # noqa diff --git a/polars_ta/ta/volume.py b/polars_ta/ta/volume.py index b02a540..35ffa17 100644 --- a/polars_ta/ta/volume.py +++ b/polars_ta/ta/volume.py @@ -1,12 +1,11 @@ -from polars import Expr +from polars import Expr, when -from polars_ta import TA_EPSILON from polars_ta.ta.overlap import EMA def AD(high: Expr, low: Expr, close: Expr, volume: Expr) -> Expr: - ad = ((close - low) - (high - close)) / (high - low + TA_EPSILON) * volume - return ad.cum_sum() + ad = when(high != low).then(((close - low) - (high - close)) / (high - low)).otherwise(0) + return (ad * volume).cum_sum() def ADOSC(high: Expr, low: Expr, close: Expr, volume: Expr, fastperiod: int = 3, slowperiod: int = 10) -> Expr: diff --git a/polars_ta/tdx/arithmetic.py b/polars_ta/tdx/arithmetic.py index 255f9f9..74a6424 100644 --- a/polars_ta/tdx/arithmetic.py +++ b/polars_ta/tdx/arithmetic.py @@ -1,3 +1,34 @@ +""" +通过`import`直接导入或更名的函数 + +```python +from polars_ta.wq.arithmetic import abs_ as ABS # noqa +from polars_ta.wq.arithmetic import add as ADD # noqa +from polars_ta.wq.arithmetic import arc_cos as ACOS # noqa +from polars_ta.wq.arithmetic import arc_sin as ASIN # noqa +from polars_ta.wq.arithmetic import arc_tan as ATAN # noqa +from polars_ta.wq.arithmetic import ceiling as CEILING # noqa +from polars_ta.wq.arithmetic import cos as COS # noqa +from polars_ta.wq.arithmetic import exp as EXP # noqa +from polars_ta.wq.arithmetic import floor as FLOOR # noqa +from polars_ta.wq.arithmetic import fraction as FRACPART # noqa +from polars_ta.wq.arithmetic import log as LN # noqa # 自然对数 (log base e) +from polars_ta.wq.arithmetic import log10 as LOG # noqa # 10为底的对数 (log base 10) +from polars_ta.wq.arithmetic import max_ as MAX # noqa +from polars_ta.wq.arithmetic import min_ as MIN # noqa +from polars_ta.wq.arithmetic import mod as MOD # noqa +from polars_ta.wq.arithmetic import power as POW # noqa +from polars_ta.wq.arithmetic import reverse as REVERSE # noqa +from polars_ta.wq.arithmetic import round_ as _round # noqa +from polars_ta.wq.arithmetic import sign as SIGN # noqa +from polars_ta.wq.arithmetic import sin as SIN # noqa +from polars_ta.wq.arithmetic import sqrt as SQRT # noqa +from polars_ta.wq.arithmetic import subtract as SUB # noqa +from polars_ta.wq.arithmetic import tan as TAN # noqa +from polars_ta.wq.transformational import int_ as INTPART # noqa +``` + +""" from polars import Expr from polars_ta.wq.arithmetic import abs_ as ABS # noqa @@ -23,7 +54,7 @@ from polars_ta.wq.arithmetic import sqrt as SQRT # noqa from polars_ta.wq.arithmetic import subtract as SUB # noqa from polars_ta.wq.arithmetic import tan as TAN # noqa -from polars_ta.wq.arithmetic import truncate as INTPART # noqa +from polars_ta.wq.transformational import int_ as INTPART # noqa SGN = SIGN diff --git a/polars_ta/tdx/reference.py b/polars_ta/tdx/reference.py index e862f7c..502b5e6 100644 --- a/polars_ta/tdx/reference.py +++ b/polars_ta/tdx/reference.py @@ -1,3 +1,24 @@ +""" +通过`import`直接导入或更名的函数 + +```python +from polars_ta.ta.overlap import SMA as MA +from polars_ta.ta.volatility import TRANGE as TR # noqa +from polars_ta.wq.arithmetic import max_ as MAX # noqa +from polars_ta.wq.arithmetic import min_ as MIN # noqa +from polars_ta.wq.time_series import ts_arg_max as HHVBARS # noqa +from polars_ta.wq.time_series import ts_arg_min as LLVBARS # noqa +from polars_ta.wq.time_series import ts_count as COUNT # noqa +from polars_ta.wq.time_series import ts_decay_linear as WMA # noqa +from polars_ta.wq.time_series import ts_delay as REF # noqa +from polars_ta.wq.time_series import ts_delta as DIFF # noqa +from polars_ta.wq.time_series import ts_max as HHV # noqa +from polars_ta.wq.time_series import ts_min as LLV # noqa +from polars_ta.wq.time_series import ts_product as MULAR # noqa +from polars_ta.wq.time_series import ts_sum as SUM +``` + +""" from polars import Boolean, Int32, UInt16 from polars import Expr from polars import when diff --git a/polars_ta/wq/arithmetic.py b/polars_ta/wq/arithmetic.py index 7b51d1a..056c63c 100644 --- a/polars_ta/wq/arithmetic.py +++ b/polars_ta/wq/arithmetic.py @@ -52,14 +52,17 @@ def add(a: Expr, b: Expr, *args) -> Expr: def arc_cos(x: Expr) -> Expr: + """反余弦""" return x.arccos() def arc_sin(x: Expr) -> Expr: + """反正弦""" return x.arcsin() def arc_tan(x: Expr) -> Expr: + """反正切""" return x.arctan() @@ -69,10 +72,12 @@ def ceiling(x: Expr) -> Expr: def cos(x: Expr) -> Expr: + """余弦""" return x.cos() def cosh(x: Expr) -> Expr: + """双曲余弦""" return x.cosh() @@ -86,6 +91,7 @@ def divide(x: Expr, y: Expr) -> Expr: def exp(x: Expr) -> Expr: + """自然指数函数""" return x.exp() @@ -100,7 +106,8 @@ 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. + """小数部分 + This operator removes the whole number part and returns the remaining fraction part with sign. Examples -------- @@ -145,6 +152,7 @@ def inverse(x: Expr) -> Expr: def log(x: Expr) -> Expr: + """e为底的对数""" if isinstance(x, (Expr, Series)): return x.log() else: @@ -152,25 +160,30 @@ def log(x: Expr) -> Expr: def log10(x: Expr) -> Expr: + """10为底的对数""" return x.log10() def log1p(x: Expr) -> Expr: - """简单收益率 转 对数收益率 convert simple return to log return""" + """简单收益率 转 对数收益率 convert simple return to log return + + log(x+1) + """ return x.log1p() 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: + """水平多列均值""" return mean_horizontal(a, b, *args) 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) @@ -271,7 +284,7 @@ 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 ---------- @@ -347,6 +360,7 @@ def s_log_1p(x: Expr) -> Expr: def sign(x: Expr) -> Expr: + """符号""" if isinstance(x, (Expr, Series)): return x.sign() else: @@ -372,10 +386,12 @@ def signed_power(x: Expr, y: Expr) -> Expr: def sin(x: Expr) -> Expr: + """正弦""" return x.sin() def sinh(x: Expr) -> Expr: + """双曲正弦""" return x.sinh() @@ -385,6 +401,7 @@ def softsign(x: Expr) -> Expr: def sqrt(x: Expr) -> Expr: + """平方根""" return x.sqrt() @@ -394,11 +411,12 @@ def subtract(x: Expr, y: Expr) -> Expr: def tan(x: Expr) -> Expr: + """正切""" return x.tan() def tanh(x: Expr) -> Expr: - """Hyperbolic tangent of x""" + """双曲正切""" return x.tanh() diff --git a/polars_ta/wq/preprocess.py b/polars_ta/wq/preprocess.py index bd3729a..e64603a 100644 --- a/polars_ta/wq/preprocess.py +++ b/polars_ta/wq/preprocess.py @@ -1,8 +1,7 @@ import polars_ols as pls -from polars import Expr +from polars import Expr, when from polars_ols.least_squares import OLSKwargs -from polars_ta import TA_EPSILON from polars_ta.wq.cross_sectional import cs_rank @@ -15,7 +14,9 @@ def cs_zscore(x: Expr, ddof: int = 0) -> Expr: def cs_minmax(x: Expr) -> Expr: a = x.min() b = x.max() - return (x - a) / (b - a + TA_EPSILON) + # 这个版本在b-a为整数时,得到的结果不好看 + # return (x - a) / (b - a + TA_EPSILON) + return when(a != b).then((x - a) / (b - a)).otherwise(0) # ====================== diff --git a/polars_ta/wq/time_series.py b/polars_ta/wq/time_series.py index 898ee04..0bc0ff1 100644 --- a/polars_ta/wq/time_series.py +++ b/polars_ta/wq/time_series.py @@ -1,10 +1,9 @@ import polars_ols as pls -from polars import Expr, Int32, UInt16, struct +from polars import Expr, Int32, UInt16, struct, when from polars import arange, repeat from polars import rolling_corr, rolling_cov from polars_ols import RollingKwargs -from polars_ta import TA_EPSILON from polars_ta.utils.numba_ import batches_i1_o1, batches_i2_o1, batches_i2_o2 from polars_ta.utils.pandas_ import roll_kurt, roll_rank from polars_ta.wq._nb import roll_argmax, roll_argmin, roll_prod, roll_co_kurtosis, roll_co_skewness, roll_moment, roll_partial_corr, roll_triple_corr, _cum_prod_by, _cum_sum_by, _signals_to_size, _cum_sum_reset, _sum_split_by @@ -398,7 +397,8 @@ def ts_scale(x: Expr, d: int = 5) -> Expr: """ a = ts_min(x, d) b = ts_max(x, d) - return (x - a) / (b - a + TA_EPSILON) + # return (x - a) / (b - a + TA_EPSILON) + return when(a != b).then((x - a) / (b - a)).otherwise(0) def ts_skewness(x: Expr, d: int = 5, bias: bool = False) -> Expr: diff --git a/polars_ta/wq/transformational.py b/polars_ta/wq/transformational.py index 9329bb2..1bb07f3 100644 --- a/polars_ta/wq/transformational.py +++ b/polars_ta/wq/transformational.py @@ -1,4 +1,4 @@ -from polars import Expr, when, Int8, Boolean +from polars import Expr, when, Boolean, Int32 def cut(x: Expr, b: float, *more_bins) -> Expr: @@ -259,7 +259,7 @@ def tail(x: Expr, lower: float = 0, upper: float = 0, newval: float = 0) -> Expr def int_(a: Expr) -> Expr: """convert bool to int""" - return a.cast(Int8) + return a.cast(Int32) def bool_(a: Expr) -> Expr: diff --git a/tools/prefix_ta.py b/tools/prefix_ta.py index f82b825..d9a7889 100644 --- a/tools/prefix_ta.py +++ b/tools/prefix_ta.py @@ -3,10 +3,11 @@ lines = ["""# this code is auto generated by tools/prefix_ta.py """] lines += codegen_import_as('polars_ta.ta.momentum', include_parameter=['timeperiod', 'fastperiod', 'fastk_period']) -lines += codegen_import_as('polars_ta.ta.overlap', include_modules=['polars_ta.wq.time_series'], include_func=['SMA', 'WMA'], include_parameter=['timeperiod']) +lines += codegen_import_as('polars_ta.ta.operators', include_modules=['polars_ta.wq.arithmetic', 'polars_ta.wq.time_series'], include_parameter=['timeperiod', 'd'], exclude_func=['ts_arg_max', 'ts_arg_min']) +lines += codegen_import_as('polars_ta.ta.overlap', include_modules=['polars_ta.wq.time_series'], include_func=['SMA', 'WMA'], include_parameter=['timeperiod', 'd'], exclude_func=['MAX', 'MIN']) lines += codegen_import_as('polars_ta.ta.price', include_parameter=['timeperiod']) -lines += codegen_import_as('polars_ta.ta.statistic', include_parameter=['timeperiod']) -lines += codegen_import_as('polars_ta.ta.transform', include_modules=['polars_ta.wq.arithmetic'], include_parameter=['timeperiod']) +lines += codegen_import_as('polars_ta.ta.statistic', include_modules=['polars_ta.wq.time_series'], include_parameter=['timeperiod', 'd'], exclude_func=['ts_std_dev']) +lines += codegen_import_as('polars_ta.ta.transform', include_modules=['polars_ta.wq.arithmetic'], include_parameter=['timeperiod', 'd']) lines += codegen_import_as('polars_ta.ta.volatility', include_func=['TRANGE'], include_parameter=['timeperiod']) lines += codegen_import_as('polars_ta.ta.volume', include_func=['AD', 'OBV'], include_parameter=['fastperiod']) save(lines, module='polars_ta.prefix.ta', write=True)