polars-ta 0.5.6__tar.gz → 0.5.9__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (80) hide show
  1. {polars_ta-0.5.6 → polars_ta-0.5.9}/PKG-INFO +2 -1
  2. polars_ta-0.5.9/polars_ta/_version.py +1 -0
  3. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/prefix/tdx.py +2 -0
  4. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/tdx/__init__.py +1 -0
  5. polars_ta-0.5.9/polars_ta/tdx/times.py +74 -0
  6. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/utils/numba_.py +10 -0
  7. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/wq/__init__.py +1 -0
  8. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/wq/_nb.py +14 -2
  9. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/wq/cross_sectional.py +20 -3
  10. polars_ta-0.5.9/polars_ta/wq/half_life.py +26 -0
  11. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/wq/preprocess.py +9 -0
  12. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/wq/time_series.py +184 -13
  13. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/wq/transformational.py +5 -0
  14. {polars_ta-0.5.6 → polars_ta-0.5.9}/pyproject.toml +1 -0
  15. polars_ta-0.5.6/polars_ta/_version.py +0 -1
  16. {polars_ta-0.5.6 → polars_ta-0.5.9}/.gitignore +0 -0
  17. {polars_ta-0.5.6 → polars_ta-0.5.9}/LICENSE +0 -0
  18. {polars_ta-0.5.6 → polars_ta-0.5.9}/README.md +0 -0
  19. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/__init__.py +0 -0
  20. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/candles/__init__.py +0 -0
  21. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/candles/cdl1.py +0 -0
  22. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/candles/cdl1_limit.py +0 -0
  23. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/candles/cdl2.py +0 -0
  24. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/labels/__init__.py +0 -0
  25. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/labels/_nb.py +0 -0
  26. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/labels/future.py +0 -0
  27. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/noise.py +0 -0
  28. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/performance/__init__.py +0 -0
  29. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/performance/drawdown.py +0 -0
  30. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/performance/returns.py +0 -0
  31. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/prefix/__init__.py +0 -0
  32. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/prefix/cdl.py +0 -0
  33. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/prefix/labels.py +0 -0
  34. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/prefix/reports.py +0 -0
  35. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/prefix/ta.py +0 -0
  36. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/prefix/talib.py +0 -0
  37. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/prefix/vec.py +0 -0
  38. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/prefix/wq.py +0 -0
  39. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/reports/__init__.py +0 -0
  40. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/reports/cicc.py +0 -0
  41. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/ta/README.md +0 -0
  42. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/ta/__init__.py +0 -0
  43. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/ta/momentum.py +0 -0
  44. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/ta/operators.py +0 -0
  45. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/ta/overlap.py +0 -0
  46. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/ta/price.py +0 -0
  47. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/ta/statistic.py +0 -0
  48. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/ta/transform.py +0 -0
  49. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/ta/volatility.py +0 -0
  50. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/ta/volume.py +0 -0
  51. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/talib/README.md +0 -0
  52. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/talib/__init__.py +0 -0
  53. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/tdx/README.md +0 -0
  54. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/tdx/_chip.py +0 -0
  55. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/tdx/_nb.py +0 -0
  56. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/tdx/_slow.py +0 -0
  57. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/tdx/arithmetic.py +0 -0
  58. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/tdx/choice.py +0 -0
  59. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/tdx/energy.py +0 -0
  60. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/tdx/logical.py +0 -0
  61. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/tdx/moving_average.py +0 -0
  62. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/tdx/over_bought_over_sold.py +0 -0
  63. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/tdx/pattern.py +0 -0
  64. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/tdx/pattern_feature.py +0 -0
  65. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/tdx/pressure_support.py +0 -0
  66. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/tdx/reference.py +0 -0
  67. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/tdx/statistic.py +0 -0
  68. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/tdx/trend.py +0 -0
  69. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/tdx/trend_feature.py +0 -0
  70. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/tdx/volume.py +0 -0
  71. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/utils/__init__.py +0 -0
  72. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/utils/helper.py +0 -0
  73. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/utils/pandas_.py +0 -0
  74. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/utils/pit.py +0 -0
  75. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/utils/withs.py +0 -0
  76. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/utils/wrapper.py +0 -0
  77. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/wq/_slow.py +0 -0
  78. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/wq/arithmetic.py +0 -0
  79. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/wq/logical.py +0 -0
  80. {polars_ta-0.5.6 → polars_ta-0.5.9}/polars_ta/wq/vector.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: polars_ta
3
- Version: 0.5.6
3
+ Version: 0.5.9
4
4
  Summary: polars expressions
5
5
  Author-email: wukan <wu-kan@163.com>
6
6
  License: MIT License
@@ -29,6 +29,7 @@ Keywords: expression,polars,talib
29
29
  Classifier: Development Status :: 4 - Beta
30
30
  Classifier: Programming Language :: Python
31
31
  Requires-Python: >=3.8
32
+ Requires-Dist: more-itertools
32
33
  Requires-Dist: numba
33
34
  Requires-Dist: numpy
34
35
  Requires-Dist: pandas
@@ -0,0 +1 @@
1
+ __version__ = "0.5.9"
@@ -128,6 +128,8 @@ from polars_ta.tdx.statistic import STDP as ts_STDP # noqa
128
128
  from polars_ta.tdx.statistic import VAR as ts_VAR # noqa
129
129
  from polars_ta.tdx.statistic import VARP as ts_VARP # noqa
130
130
  from polars_ta.tdx.statistic import ts_up_stat # noqa
131
+ from polars_ta.tdx.times import FROMOPEN # noqa
132
+ from polars_ta.tdx.times import FROMOPEN_1 # noqa
131
133
  from polars_ta.tdx.trend import ADX as ts_ADX # noqa
132
134
  from polars_ta.tdx.trend import ADXR as ts_ADXR # noqa
133
135
  from polars_ta.tdx.trend import DPO as ts_DPO # noqa
@@ -9,6 +9,7 @@ from polars_ta.tdx.pattern_feature import * # noqa
9
9
  from polars_ta.tdx.pressure_support import * # noqa
10
10
  from polars_ta.tdx.reference import * # noqa
11
11
  from polars_ta.tdx.statistic import * # noqa
12
+ from polars_ta.tdx.times import * # noqa
12
13
  from polars_ta.tdx.trend import * # noqa
13
14
  from polars_ta.tdx.trend_feature import * # noqa
14
15
  from polars_ta.tdx.volume import * # noqa
@@ -0,0 +1,74 @@
1
+ from datetime import time, timedelta
2
+
3
+ from polars import Expr, when
4
+
5
+
6
+ def FROMOPEN(t: Expr) -> Expr:
7
+ """返回当前时刻距开盘有多少分钟
8
+
9
+ 范围0~240,开盘前为0,10点为31
10
+
11
+ Examples
12
+ --------
13
+ from datetime import datetime
14
+
15
+ import polars as pl
16
+
17
+ from polars_ta.tdx.times import FROMOPEN, FROMOPEN1
18
+
19
+ df = pl.DataFrame({'datetime': [
20
+ datetime(2025, 1, 1, 0, 0),
21
+ datetime(2025, 1, 1, 9, 25),
22
+ datetime(2025, 1, 1, 9, 30, 57),
23
+ datetime(2025, 1, 1, 9, 31),
24
+ datetime(2025, 1, 1, 10, 0),
25
+ datetime(2025, 1, 1, 13, 0),
26
+ ]})
27
+
28
+ df = df.with_columns(
29
+ FROMOPEN=FROMOPEN(pl.col('datetime')),
30
+ FROMOPEN1=FROMOPEN_1(pl.col('datetime'), 0),
31
+ FROMOPEN2=FROMOPEN_1(pl.col('datetime'), 60),
32
+ )
33
+
34
+ shape: (6, 4)
35
+ ┌─────────────────────┬──────────┬───────────┬───────────┐
36
+ │ datetime ┆ FROMOPEN ┆ FROMOPEN1 ┆ FROMOPEN2 │
37
+ │ --- ┆ --- ┆ --- ┆ --- │
38
+ │ datetime[μs] ┆ i64 ┆ i64 ┆ i64 │
39
+ ╞═════════════════════╪══════════╪═══════════╪═══════════╡
40
+ │ 2025-01-01 00:00:00 ┆ 0 ┆ 240 ┆ 240 │
41
+ │ 2025-01-01 09:25:00 ┆ 0 ┆ 1 ┆ 1 │
42
+ │ 2025-01-01 09:30:57 ┆ 1 ┆ 1 ┆ 2 │
43
+ │ 2025-01-01 09:31:00 ┆ 2 ┆ 2 ┆ 3 │
44
+ │ 2025-01-01 10:00:00 ┆ 31 ┆ 31 ┆ 32 │
45
+ │ 2025-01-01 13:00:00 ┆ 121 ┆ 121 ┆ 122 │
46
+ └─────────────────────┴──────────┴───────────┴───────────┘
47
+
48
+ """
49
+ am = (t.dt.time() - time(9, 29)).dt.total_minutes().clip(0, 120)
50
+ pm = (t.dt.time() - time(12, 59)).dt.total_minutes().clip(0, 120)
51
+ return am + pm
52
+
53
+
54
+ def FROMOPEN_1(t: Expr, offset: int) -> Expr:
55
+ """返回当前时刻距开盘有多少分钟。范围1~240
56
+
57
+ 用于计算量比
58
+ 1. 竞价量比,分母应当为1
59
+ 2. 日线数据0~8点时,返回240
60
+ 3. 日线数据9点时,返回1
61
+
62
+ Parameters
63
+ ----------
64
+ t : Expr
65
+ 时间列
66
+ offset : int
67
+ 偏移量,单位秒
68
+
69
+ Notes
70
+ -----
71
+ 每根K线结束时,标签是当前时间的50多秒,而结束时时间已经到下以分钟了,所以建议加60秒
72
+
73
+ """
74
+ return when(t.dt.time() >= time(8, 45)).then(FROMOPEN(t + timedelta(seconds=offset)).clip(1, 240)).otherwise(240)
@@ -2,6 +2,7 @@
2
2
  Demo for using numba to implement rolling functions.
3
3
  本文件是使用numba实现rolling的函数,演示用
4
4
  """
5
+ from functools import lru_cache
5
6
  from typing import List
6
7
 
7
8
  import numpy as np
@@ -160,3 +161,12 @@ def roll_sum(x: Expr, n: int) -> Expr:
160
161
 
161
162
  def roll_cov(a: Expr, b: Expr, n: int) -> Expr:
162
163
  return struct([a, b]).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 2), nb_roll_cov, n))
164
+
165
+
166
+ @lru_cache
167
+ @jit(nopython=True, nogil=True, fastmath=True, cache=True)
168
+ def get_exponent_weights(
169
+ window: int = 10,
170
+ half_life: int = 5,
171
+ ) -> np.ndarray:
172
+ return np.repeat(0.5 ** (1 / half_life), window) ** np.arange(window - 1, -1, -1)
@@ -1,5 +1,6 @@
1
1
  from polars_ta.wq.arithmetic import * # noqa
2
2
  from polars_ta.wq.cross_sectional import * # noqa
3
+ from polars_ta.wq.half_life import * # noqa
3
4
  from polars_ta.wq.logical import * # noqa
4
5
  from polars_ta.wq.preprocess import * # noqa
5
6
  from polars_ta.wq.time_series import * # noqa
@@ -312,7 +312,13 @@ def _signals_to_size(is_long_entry: np.ndarray, is_long_exit: np.ndarray,
312
312
 
313
313
 
314
314
  @jit(nopython=True, nogil=True, cache=True)
315
- def roll_decay_linear(x1, window, min_periods):
315
+ def _roll_decay_linear(x1, window, min_periods):
316
+ """
317
+ def ts_decay_linear(x: Expr, d: int = 30, min_samples: Optional[int] = None) -> Expr:
318
+ minp = min_samples or polars_ta.MIN_SAMPLES or d
319
+ return x.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _roll_decay_linear, d, minp))
320
+
321
+ """
316
322
  weights = np.arange(1., window + 1)
317
323
 
318
324
  out1 = full_with_window_size(x1, np.nan, dtype=np.float64, window_size=window)
@@ -327,7 +333,13 @@ def roll_decay_linear(x1, window, min_periods):
327
333
 
328
334
 
329
335
  @jit(nopython=True, nogil=True, cache=True)
330
- def roll_decay_exp_window(x1, window, min_periods, factor):
336
+ def _roll_decay_exp_window(x1, window, min_periods, factor):
337
+ """
338
+ def ts_decay_exp_window(x: Expr, d: int = 30, factor: float = 1.0, min_samples: Optional[int] = None) -> Expr:
339
+ minp = min_samples or polars_ta.MIN_SAMPLES or d
340
+ return x.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _roll_decay_exp_window, d, minp, factor))
341
+
342
+ """
331
343
  weights = factor ** np.arange(window - 1, -1, -1)
332
344
 
333
345
  out1 = full_with_window_size(x1, np.nan, dtype=np.float64, window_size=window)
@@ -290,7 +290,7 @@ def cs_rank(x: Expr, pct: bool = True) -> Expr:
290
290
 
291
291
 
292
292
  def cs_rank_if(condition: Expr, x: Expr, pct: bool = True) -> Expr:
293
- """动态票池过滤排名
293
+ """横截面筛选排名。可实现动态票池
294
294
 
295
295
  Parameters
296
296
  ----------
@@ -299,17 +299,34 @@ def cs_rank_if(condition: Expr, x: Expr, pct: bool = True) -> Expr:
299
299
  x:Expr
300
300
  因子
301
301
  pct:bool
302
+ 排名百分比。范围:[0,1]
302
303
 
303
304
  Examples
304
305
  --------
305
306
  ```python
306
307
  df = pl.DataFrame({
307
308
  'a': [None, 1, 1, 1, 2, 2, 3, 10],
308
- 'b': [1, 2, 3, 4, 5, 6, 7, 8],
309
+ 'b': [1, 2, 3, 4, 5, 6, None, 8],
309
310
  }).with_columns(
310
- out1=cs_rank_if(True, pl.col('a'), True), # 与cs_rank等价
311
+ out1=cs_rank_if(True, pl.col('a'), True), # 与cs_rank等价
311
312
  out2=cs_rank_if(pl.col('b') > 3, -pl.col('a'), False),
312
313
  )
314
+
315
+ shape: (8, 4)
316
+ ┌──────┬──────┬──────────┬──────┐
317
+ │ a ┆ b ┆ out1 ┆ out2 │
318
+ │ --- ┆ --- ┆ --- ┆ --- │
319
+ │ i64 ┆ i64 ┆ f64 ┆ u32 │
320
+ ╞══════╪══════╪══════════╪══════╡
321
+ │ null ┆ 1 ┆ null ┆ null │
322
+ │ 1 ┆ 2 ┆ 0.0 ┆ null │
323
+ │ 1 ┆ 3 ┆ 0.0 ┆ null │
324
+ │ 1 ┆ 4 ┆ 0.0 ┆ 3 │
325
+ │ 2 ┆ 5 ┆ 0.333333 ┆ 2 │
326
+ │ 2 ┆ 6 ┆ 0.333333 ┆ 2 │
327
+ │ 3 ┆ null ┆ 0.666667 ┆ null │
328
+ │ 10 ┆ 8 ┆ 1.0 ┆ 1 │
329
+ └──────┴──────┴──────────┴──────┘
313
330
  ```
314
331
 
315
332
  Notes
@@ -0,0 +1,26 @@
1
+ import numpy as np
2
+ from polars import Expr
3
+
4
+ from polars_ta.utils.numba_ import get_exponent_weights
5
+
6
+
7
+ def ts_mean_hl(x: Expr, d: int, half_life: int):
8
+ """滚动均值。带半衰期"""
9
+ return x.fill_null(np.nan).rolling_mean(d, weights=get_exponent_weights(d, half_life)).fill_nan(None)
10
+
11
+
12
+ def ts_sum_hl(x: Expr, d: int, half_life: int):
13
+ """滚动求和。带半衰期"""
14
+ return x.fill_null(np.nan).rolling_sum(d, weights=get_exponent_weights(d, half_life)).fill_nan(None)
15
+
16
+
17
+ def ts_std_hl(x: Expr, d: int, half_life: int):
18
+ """滚动标准差。带半衰期"""
19
+ return x.fill_null(np.nan).rolling_std(d, weights=get_exponent_weights(d, half_life)).fill_nan(None)
20
+
21
+
22
+ def ts_var_hl(x: Expr, d: int, half_life: int):
23
+ """滚动方差。带半衰期"""
24
+ return x.fill_null(np.nan).rolling_var(d, weights=get_exponent_weights(d, half_life)).fill_nan(None)
25
+
26
+ # TODO 混动时序回归,带半衰期
@@ -119,3 +119,12 @@ def cs_mad_zscore_resid_zscore(y: Expr, *more_x: Expr) -> Expr:
119
119
  def cs_quantile_zscore(y: Expr, low_limit: float = 0.025, up_limit: float = 0.975) -> Expr:
120
120
  """横截面分位数去极值、标准化"""
121
121
  return cs_zscore(cs_quantile(y, low_limit, up_limit))
122
+
123
+
124
+ # ==========================
125
+ def cs_resid_w(w: Expr, y: Expr, *more_x: Expr) -> Expr:
126
+ """横截面加权多元回归取残差
127
+
128
+ Barra中权重采用流通市值的平方根
129
+ """
130
+ return pls.compute_least_squares(y, *more_x, sample_weights=w, mode='residuals', ols_kwargs=_ols_kwargs)
@@ -1,14 +1,18 @@
1
+ import itertools
1
2
  from typing import Optional
2
3
 
4
+ import more_itertools
5
+ import numpy as np
3
6
  import polars_ols as pls
4
- from polars import Expr, UInt16, struct, when, Struct, Field, Float64, Boolean, UInt32
7
+ from polars import Expr, UInt16, struct, when, Struct, Field, Float64, Boolean, UInt32, all_horizontal, any_horizontal
5
8
  from polars import rolling_corr, rolling_cov
6
9
  from polars_ols import RollingKwargs
7
10
 
8
11
  import polars_ta
9
12
  from polars_ta.utils.numba_ import batches_i1_o1, batches_i2_o1, batches_i2_o2, struct_to_numpy
10
13
  from polars_ta.utils.pandas_ import roll_rank
11
- from polars_ta.wq._nb import roll_argmax, roll_argmin, 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, roll_decay_linear, roll_decay_exp_window, roll_prod
14
+ from polars_ta.wq._nb import roll_argmax, roll_argmin, roll_co_kurtosis, roll_co_skewness, roll_moment, roll_partial_corr, roll_triple_corr, _cum_prod_by, _cum_sum_by, _signals_to_size, \
15
+ _cum_sum_reset, _sum_split_by, roll_prod
12
16
 
13
17
 
14
18
  def ts_arg_max(x: Expr, d: int = 5, reverse: bool = True, min_samples: Optional[int] = None) -> Expr:
@@ -165,6 +169,40 @@ def ts_count(x: Expr, d: int = 30, min_samples: Optional[int] = None) -> Expr:
165
169
  return x.cast(Boolean).cast(UInt32).rolling_sum(d, min_samples=minp)
166
170
 
167
171
 
172
+ def ts_count_eq(x: Expr, d: int = 30, n: int = 10, min_samples: Optional[int] = None) -> Expr:
173
+ """D天内最近连续出现N次
174
+
175
+ Parameters
176
+ ----------
177
+ x
178
+ d: int
179
+ 窗口大小
180
+ n: int
181
+ 连续出现次数
182
+
183
+ """
184
+ minp = min_samples or polars_ta.MIN_SAMPLES
185
+ xx = x.cast(Boolean).cast(UInt32)
186
+ return (xx.rolling_sum(n) == n) & (xx.rolling_sum(d, min_samples=minp) == n)
187
+
188
+
189
+ def ts_count_ge(x: Expr, d: int = 30, n: int = 10, min_samples: Optional[int] = None) -> Expr:
190
+ """D天内最近连续出现至少N次
191
+
192
+ Parameters
193
+ ----------
194
+ x
195
+ d: int
196
+ 窗口大小
197
+ n: int
198
+ 至少连续出现次数
199
+
200
+ """
201
+ minp = min_samples or polars_ta.MIN_SAMPLES
202
+ xx = x.cast(Boolean).cast(UInt32)
203
+ return (xx.rolling_sum(n) == n) & (xx.rolling_sum(d, min_samples=minp) >= n)
204
+
205
+
168
206
  def ts_count_nans(x: Expr, d: int = 5, min_samples: Optional[int] = None) -> Expr:
169
207
  """时序滚动统计nan出现次数
170
208
 
@@ -448,12 +486,11 @@ def ts_decay_exp_window(x: Expr, d: int = 30, factor: float = 1.0, min_samples:
448
486
  https://platform.worldquantbrain.com/learn/operators/detailed-operator-descriptions#ts_decay_exp_windowx-d-factor-10-nan-true
449
487
 
450
488
  """
451
- # y = arange(d - 1, -1, step=-1, eager=False)
452
- # weights = repeat(factor, d, eager=True).pow(y)
489
+ minp = min_samples or polars_ta.MIN_SAMPLES
490
+ weights = np.repeat(factor, d) ** np.arange(d - 1, -1, -1)
453
491
  # print(weights)
454
- # return x.rolling_mean(d, weights=weights)
455
- minp = min_samples or polars_ta.MIN_SAMPLES or d
456
- return x.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), roll_decay_exp_window, d, minp, factor))
492
+ # pyo3_runtime.PanicException: weights not yet supported on array with null values
493
+ return x.fill_null(np.nan).rolling_mean(d, weights=weights, min_samples=minp).fill_nan(None)
457
494
 
458
495
 
459
496
  def ts_decay_linear(x: Expr, d: int = 30, min_samples: Optional[int] = None) -> Expr:
@@ -492,12 +529,12 @@ def ts_decay_linear(x: Expr, d: int = 30, min_samples: Optional[int] = None) ->
492
529
  https://platform.worldquantbrain.com/learn/operators/detailed-operator-descriptions#ts_decay_linearx-d-dense-false
493
530
 
494
531
  """
495
- # # weights not yet supported on array with null values
496
- # weights = arange(1, d + 1, eager=True)
497
- # # print(weights)
498
- # return x.rolling_mean(d, weights=weights)
499
- minp = min_samples or polars_ta.MIN_SAMPLES or d
500
- return x.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), roll_decay_linear, d, minp))
532
+ minp = min_samples or polars_ta.MIN_SAMPLES
533
+ weights = np.arange(1, d + 1)
534
+ # print(weights)
535
+ # pyo3_runtime.PanicException: weights not yet supported on array with null values
536
+ # null换成NaN就不报错了,再换回来
537
+ return x.fill_null(np.nan).rolling_mean(d, weights=weights, min_samples=minp).fill_nan(None)
501
538
 
502
539
 
503
540
  def ts_delay(x: Expr, d: int = 1, fill_value: float = None) -> Expr:
@@ -768,6 +805,140 @@ def ts_scale(x: Expr, d: int = 5, min_samples: Optional[int] = None) -> Expr:
768
805
  return when(a != b).then((x - a) / (b - a)).otherwise(0)
769
806
 
770
807
 
808
+ def ts_shifts_v1(*args) -> Expr:
809
+ """时序上按顺序进行平移,然后逻辑与
810
+
811
+ 比如今天反包,昨天阴跌,前天涨停。只需要写`ts_shifts_v1(反包,阴跌,涨停)`。使用此函数能一定程度上简化代码
812
+
813
+ Parameters
814
+ ----------
815
+ args
816
+ pl.col按照顺序排列
817
+
818
+ Examples
819
+ --------
820
+ ```python
821
+ df = pl.DataFrame({
822
+ 'a': [None, False, False, True, True, True],
823
+ 'b': [None, False, True, True, True, True],
824
+ 'c': [None, True, True, True, True, True],
825
+ }).with_columns(
826
+ out1=ts_shifts_v1(pl.col('a'), pl.col('b'), pl.col('c')),
827
+ out2=pl.col('a').shift(0) & pl.col('b').shift(1) & pl.col('c').shift(2),
828
+ )
829
+
830
+ shape: (6, 5)
831
+ ┌──────┬───────┬───────┬───────┬───────┐
832
+ │ c ┆ b ┆ a ┆ out1 ┆ out2 │
833
+ │ --- ┆ --- ┆ --- ┆ --- ┆ --- │
834
+ │ bool ┆ bool ┆ bool ┆ bool ┆ bool │
835
+ ╞══════╪═══════╪═══════╪═══════╪═══════╡
836
+ │ null ┆ null ┆ null ┆ null ┆ null │
837
+ │ true ┆ false ┆ false ┆ false ┆ false │
838
+ │ true ┆ true ┆ false ┆ false ┆ false │
839
+ │ true ┆ true ┆ true ┆ true ┆ true │
840
+ │ true ┆ true ┆ true ┆ true ┆ true │
841
+ │ true ┆ true ┆ true ┆ true ┆ true │
842
+ └──────┴───────┴───────┴───────┴───────┘
843
+ ```
844
+
845
+ """
846
+ return all_horizontal(arg.shift(i) for i, arg in enumerate(args))
847
+
848
+
849
+ def ts_shifts_v2(*args) -> Expr:
850
+ """时序上按顺序进行平移,然后逻辑与
851
+
852
+ 遇到连续条件时可简化参数。如连续3天涨停后连续2天跌停,可用`ts_shifts_v2(跌停,2,涨停,3)`
853
+ 另一种实现方法是:`ts_delay((ts_count(涨停,3)==3),2)&(ts_count(跌停,2)==2)`
854
+
855
+ Parameters
856
+ ----------
857
+ args
858
+ pl.col, repeat。两个参数循环
859
+
860
+ Examples
861
+ --------
862
+ ```python
863
+ df = pl.DataFrame({
864
+ 'a': [None, False, False, True, True, True],
865
+ 'b': [None, False, True, True, True, True],
866
+ 'c': [None, True, True, True, True, True],
867
+ }).with_columns(
868
+ out1=ts_shifts_v1(pl.col('a'), pl.col('b'), pl.col('b')),
869
+ out2=ts_shifts_v2(pl.col('a'), 1, pl.col('b'), 2),
870
+ out3=pl.col('a').shift(0) & pl.col('b').shift(1) & pl.col('b').shift(2),
871
+ )
872
+
873
+ shape: (6, 6)
874
+ ┌───────┬───────┬──────┬───────┬───────┬───────┐
875
+ │ a ┆ b ┆ c ┆ out1 ┆ out2 ┆ out3 │
876
+ │ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │
877
+ │ bool ┆ bool ┆ bool ┆ bool ┆ bool ┆ bool │
878
+ ╞═══════╪═══════╪══════╪═══════╪═══════╪═══════╡
879
+ │ null ┆ null ┆ null ┆ null ┆ null ┆ null │
880
+ │ false ┆ false ┆ true ┆ false ┆ false ┆ false │
881
+ │ false ┆ true ┆ true ┆ false ┆ false ┆ false │
882
+ │ true ┆ true ┆ true ┆ false ┆ false ┆ false │
883
+ │ true ┆ true ┆ true ┆ true ┆ true ┆ true │
884
+ │ true ┆ true ┆ true ┆ true ┆ true ┆ true │
885
+ └───────┴───────┴──────┴───────┴───────┴───────┘
886
+ ```
887
+
888
+ """
889
+ return ts_shifts_v1(*itertools.chain.from_iterable([item] * count for item, count in more_itertools.chunked(args, 2)))
890
+
891
+
892
+ def ts_shifts_v3(*args) -> Expr:
893
+ """时序上按顺序进行平移,然后逻辑或
894
+
895
+ 如:涨停后连续下跌1~3天。这个案例会导致涨停可能出现在动态的一天。`ts_shifts_v3(下跌,1,3,涨停,1,1)`
896
+ 它本质上是`ts_shifts_v2(下跌,1,涨停,1)|ts_shifts_v2(下跌,2,涨停,1)|ts_shifts_v2(下跌,3,涨停,1)|`
897
+
898
+ Parameters
899
+ ----------
900
+ args
901
+ pl.col, [start, end]。三个参数循环。start/end都是闭区间
902
+
903
+ Examples
904
+ --------
905
+ ```python
906
+ df = pl.DataFrame({
907
+ 'a': [None, False, False, True, True, True],
908
+ 'b': [None, False, True, True, True, True],
909
+ 'c': [None, True, True, True, True, True],
910
+ }).with_columns(
911
+ out0=ts_shifts_v3(pl.col('a'), 1, 3, pl.col('b'), 2, 2),
912
+ out1=ts_shifts_v2(pl.col('a'), 1, pl.col('b'), 2),
913
+ out2=ts_shifts_v2(pl.col('a'), 2, pl.col('b'), 2),
914
+ out3=ts_shifts_v2(pl.col('a'), 3, pl.col('b'), 2)
915
+ ).with_columns(
916
+ out4=pl.col('out1') | pl.col('out2') | pl.col('out3')
917
+ )
918
+
919
+ shape: (6, 8)
920
+ ┌───────┬───────┬──────┬───────┬───────┬───────┬───────┬───────┐
921
+ │ a ┆ b ┆ c ┆ out0 ┆ out1 ┆ out2 ┆ out3 ┆ out4 │
922
+ │ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │
923
+ │ bool ┆ bool ┆ bool ┆ bool ┆ bool ┆ bool ┆ bool ┆ bool │
924
+ ╞═══════╪═══════╪══════╪═══════╪═══════╪═══════╪═══════╪═══════╡
925
+ │ null ┆ null ┆ null ┆ null ┆ null ┆ null ┆ null ┆ null │
926
+ │ false ┆ false ┆ true ┆ false ┆ false ┆ false ┆ false ┆ false │
927
+ │ false ┆ true ┆ true ┆ false ┆ false ┆ false ┆ false ┆ false │
928
+ │ true ┆ true ┆ true ┆ false ┆ false ┆ false ┆ false ┆ false │
929
+ │ true ┆ true ┆ true ┆ true ┆ true ┆ false ┆ false ┆ true │
930
+ │ true ┆ true ┆ true ┆ true ┆ true ┆ true ┆ false ┆ true │
931
+ └───────┴───────┴──────┴───────┴───────┴───────┴───────┴───────┘
932
+ ```
933
+
934
+ """
935
+ exprs = [a for a, b, c in more_itertools.chunked(args, 3)]
936
+ ranges = [range(b, c + 1) for a, b, c in more_itertools.chunked(args, 3)]
937
+ # 参数整理成col,repeat模式
938
+ outputs = [itertools.chain.from_iterable(zip(exprs, d)) for d in itertools.product(*ranges)]
939
+ return any_horizontal(ts_shifts_v2(*_) for _ in outputs)
940
+
941
+
771
942
  def ts_skewness(x: Expr, d: int = 5, bias: bool = False, min_samples: Optional[int] = None) -> Expr:
772
943
  """时序滚动偏度
773
944
 
@@ -275,3 +275,8 @@ def bool_(a: Expr) -> Expr:
275
275
  def float_(a: Expr) -> Expr:
276
276
  """int转成float"""
277
277
  return a.cast(Float32)
278
+
279
+
280
+ def nop(x: Expr) -> Expr:
281
+ """空操作"""
282
+ return x
@@ -18,6 +18,7 @@ dependencies = [
18
18
  "numpy",
19
19
  "numba",
20
20
  "pandas",
21
+ "more_itertools",
21
22
  ]
22
23
  dynamic = ["version"]
23
24
 
@@ -1 +0,0 @@
1
- __version__ = "0.5.6"
File without changes
File without changes
File without changes
File without changes