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

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.
@@ -1,14 +1,17 @@
1
+ from typing import Optional
2
+
1
3
  import polars_ols as pls
2
4
  from polars import Expr, UInt16, struct, when, Struct, Field, Float64, Boolean, UInt32
3
5
  from polars import rolling_corr, rolling_cov
4
6
  from polars_ols import RollingKwargs
5
7
 
6
- from polars_ta.utils.numba_ import batches_i1_o1, batches_i2_o1, batches_i2_o2
7
- from polars_ta.utils.pandas_ import roll_kurt, roll_rank
8
- 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, roll_decay_linear, roll_decay_exp_window
8
+ import polars_ta
9
+ from polars_ta.utils.numba_ import batches_i1_o1, batches_i2_o1, batches_i2_o2, struct_to_numpy
10
+ 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
9
12
 
10
13
 
11
- def ts_arg_max(x: Expr, d: int = 5, reverse: bool = True) -> Expr:
14
+ def ts_arg_max(x: Expr, d: int = 5, reverse: bool = True, min_samples: Optional[int] = None) -> Expr:
12
15
  """Returns the relative index of the max value in the time series for the past d days.
13
16
  If the current day has the max value for the past d days, it returns 0.
14
17
  If previous day has the max value for the past d days, it returns 1.
@@ -19,6 +22,7 @@ def ts_arg_max(x: Expr, d: int = 5, reverse: bool = True) -> Expr:
19
22
  d
20
23
  reverse
21
24
  反向
25
+ min_samples
22
26
 
23
27
  See Also
24
28
  --------
@@ -52,10 +56,11 @@ def ts_arg_max(x: Expr, d: int = 5, reverse: bool = True) -> Expr:
52
56
  https://platform.worldquantbrain.com/learn/operators/detailed-operator-descriptions#ts_arg_maxx-d
53
57
 
54
58
  """
55
- return x.map_batches(lambda x1: batches_i1_o1(x1.to_numpy(), roll_argmax, d, reverse, dtype=UInt16))
59
+ minp = min_samples or polars_ta.MIN_SAMPLES or d
60
+ return x.map_batches(lambda x1: batches_i1_o1(x1.to_numpy(), roll_argmax, d, minp, reverse, dtype=UInt16))
56
61
 
57
62
 
58
- def ts_arg_min(x: Expr, d: int = 5, reverse: bool = True) -> Expr:
63
+ def ts_arg_min(x: Expr, d: int = 5, reverse: bool = True, min_samples: Optional[int] = None) -> Expr:
59
64
  """
60
65
 
61
66
  Parameters
@@ -64,6 +69,7 @@ def ts_arg_min(x: Expr, d: int = 5, reverse: bool = True) -> Expr:
64
69
  d
65
70
  reverse
66
71
  反向
72
+ min_samples
67
73
 
68
74
  See Also
69
75
  --------
@@ -74,18 +80,21 @@ def ts_arg_min(x: Expr, d: int = 5, reverse: bool = True) -> Expr:
74
80
  https://platform.worldquantbrain.com/learn/operators/detailed-operator-descriptions#ts_arg_minx-d
75
81
 
76
82
  """
77
- return x.map_batches(lambda x1: batches_i1_o1(x1.to_numpy(), roll_argmin, d, reverse, dtype=UInt16))
83
+ minp = min_samples or polars_ta.MIN_SAMPLES or d
84
+ return x.map_batches(lambda x1: batches_i1_o1(x1.to_numpy(), roll_argmin, d, minp, reverse, dtype=UInt16))
78
85
 
79
86
 
80
- def ts_co_kurtosis(x: Expr, y: Expr, d: int = 5, ddof: int = 0) -> Expr:
81
- return struct([x, y]).map_batches(lambda xx: batches_i2_o1([xx.struct[i].to_numpy() for i in range(2)], roll_co_kurtosis, d))
87
+ def ts_co_kurtosis(x: Expr, y: Expr, d: int = 5, ddof: int = 0, min_samples: Optional[int] = None) -> Expr:
88
+ minp = min_samples or polars_ta.MIN_SAMPLES or d
89
+ return struct([x, y]).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 2), roll_co_kurtosis, d, minp))
82
90
 
83
91
 
84
- def ts_co_skewness(x: Expr, y: Expr, d: int = 5, ddof: int = 0) -> Expr:
85
- return struct([x, y]).map_batches(lambda xx: batches_i2_o1([xx.struct[i].to_numpy() for i in range(2)], roll_co_skewness, d))
92
+ def ts_co_skewness(x: Expr, y: Expr, d: int = 5, ddof: int = 0, min_samples: Optional[int] = None) -> Expr:
93
+ minp = min_samples or polars_ta.MIN_SAMPLES or d
94
+ return struct([x, y]).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 2), roll_co_skewness, d, minp))
86
95
 
87
96
 
88
- def ts_corr(x: Expr, y: Expr, d: int = 5, ddof: int = 1) -> Expr:
97
+ def ts_corr(x: Expr, y: Expr, d: int = 5, ddof: int = 1, min_samples: Optional[int] = None) -> Expr:
89
98
  """rolling correlation between two columns
90
99
 
91
100
  时序滚动相关系数
@@ -97,22 +106,25 @@ def ts_corr(x: Expr, y: Expr, d: int = 5, ddof: int = 1) -> Expr:
97
106
  d
98
107
  ddof
99
108
  自由度
109
+ min_samples
100
110
 
101
111
  Notes
102
112
  -----
103
113
  x、y不区分先后
104
114
 
105
115
  """
106
- return rolling_corr(x, y, window_size=d, ddof=ddof)
116
+ minp = min_samples or polars_ta.MIN_SAMPLES
117
+ return rolling_corr(x, y, window_size=d, ddof=ddof, min_samples=minp)
107
118
 
108
119
 
109
- def ts_count(x: Expr, d: int = 30) -> Expr:
120
+ def ts_count(x: Expr, d: int = 30, min_samples: Optional[int] = None) -> Expr:
110
121
  """时序滚动计数
111
122
 
112
123
  Parameters
113
124
  ----------
114
125
  x
115
126
  d
127
+ min_samples
116
128
 
117
129
  Examples
118
130
  --------
@@ -141,16 +153,18 @@ def ts_count(x: Expr, d: int = 30) -> Expr:
141
153
  ```
142
154
 
143
155
  """
144
- return x.cast(Boolean).cast(UInt32).rolling_sum(d)
156
+ minp = min_samples or polars_ta.MIN_SAMPLES
157
+ return x.cast(Boolean).cast(UInt32).rolling_sum(d, min_samples=minp)
145
158
 
146
159
 
147
- def ts_count_nans(x: Expr, d: int = 5) -> Expr:
160
+ def ts_count_nans(x: Expr, d: int = 5, min_samples: Optional[int] = None) -> Expr:
148
161
  """时序滚动统计nan出现次数
149
162
 
150
163
  Parameters
151
164
  ----------
152
165
  x
153
166
  d
167
+ min_samples
154
168
 
155
169
  Examples
156
170
  --------
@@ -178,16 +192,18 @@ def ts_count_nans(x: Expr, d: int = 5) -> Expr:
178
192
  ```
179
193
 
180
194
  """
181
- return x.is_nan().cast(UInt32).rolling_sum(d)
195
+ minp = min_samples or polars_ta.MIN_SAMPLES
196
+ return x.is_nan().cast(UInt32).rolling_sum(d, min_samples=minp)
182
197
 
183
198
 
184
- def ts_count_nulls(x: Expr, d: int = 5) -> Expr:
199
+ def ts_count_nulls(x: Expr, d: int = 5, min_samples: Optional[int] = None) -> Expr:
185
200
  """时序滚动统计null出现次数
186
201
 
187
202
  Parameters
188
203
  ----------
189
204
  x
190
205
  d
206
+ min_samples
191
207
 
192
208
  Examples
193
209
  --------
@@ -215,10 +231,11 @@ def ts_count_nulls(x: Expr, d: int = 5) -> Expr:
215
231
  ```
216
232
 
217
233
  """
218
- return x.is_null().cast(UInt32).rolling_sum(d)
234
+ minp = min_samples or polars_ta.MIN_SAMPLES
235
+ return x.is_null().cast(UInt32).rolling_sum(d, min_samples=minp)
219
236
 
220
237
 
221
- def ts_covariance(x: Expr, y: Expr, d: int = 5, ddof: int = 1) -> Expr:
238
+ def ts_covariance(x: Expr, y: Expr, d: int = 5, ddof: int = 1, min_samples: Optional[int] = None) -> Expr:
222
239
  """rolling covariance between two columns
223
240
 
224
241
  时序协方差
@@ -230,13 +247,15 @@ def ts_covariance(x: Expr, y: Expr, d: int = 5, ddof: int = 1) -> Expr:
230
247
  d
231
248
  ddof
232
249
  自由度
250
+ min_samples
233
251
 
234
252
  Notes
235
253
  -----
236
254
  x、y不区分先后
237
255
 
238
256
  """
239
- return rolling_cov(x, y, window_size=d, ddof=ddof)
257
+ minp = min_samples or polars_ta.MIN_SAMPLES
258
+ return rolling_cov(x, y, window_size=d, ddof=ddof, min_samples=minp)
240
259
 
241
260
 
242
261
  def ts_cum_count(x: Expr) -> Expr:
@@ -379,12 +398,14 @@ def ts_cum_sum_reset(x: Expr) -> Expr:
379
398
  return x.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _cum_sum_reset))
380
399
 
381
400
 
382
- def ts_decay_exp_window(x: Expr, d: int = 30, factor: float = 1.0) -> Expr:
401
+ def ts_decay_exp_window(x: Expr, d: int = 30, factor: float = 1.0, min_samples: Optional[int] = None) -> Expr:
383
402
  """指数衰减移动平均
384
403
 
385
404
  Examples
386
405
  --------
387
406
  ```python
407
+ from polars_ta.wq.time_series import ts_decay_linear, ts_decay_exp_window
408
+
388
409
  df = pl.DataFrame({
389
410
  'a': [None, 6, 5, 4, 5, 30],
390
411
  }).with_columns(
@@ -412,6 +433,7 @@ def ts_decay_exp_window(x: Expr, d: int = 30, factor: float = 1.0) -> Expr:
412
433
  d
413
434
  factor
414
435
  衰减系数
436
+ min_samples
415
437
 
416
438
  References
417
439
  ----------
@@ -422,21 +444,26 @@ def ts_decay_exp_window(x: Expr, d: int = 30, factor: float = 1.0) -> Expr:
422
444
  # weights = repeat(factor, d, eager=True).pow(y)
423
445
  # print(weights)
424
446
  # return x.rolling_mean(d, weights=weights)
425
- return x.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), roll_decay_exp_window, d, factor))
447
+ minp = min_samples or polars_ta.MIN_SAMPLES or d
448
+ return x.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), roll_decay_exp_window, d, minp, factor))
426
449
 
427
450
 
428
- def ts_decay_linear(x: Expr, d: int = 30) -> Expr:
451
+ def ts_decay_linear(x: Expr, d: int = 30, min_samples: Optional[int] = None) -> Expr:
429
452
  """线性衰减移动平均
430
453
 
431
454
  Examples
432
455
  --------
433
456
  ```python
457
+ from polars_ta.talib import WMA as ts_WMA
458
+ from polars_ta.wq.time_series import ts_decay_linear
459
+
434
460
  df = pl.DataFrame({
435
461
  'a': [None, 6, 5, 4, 5, 30],
436
462
  }).with_columns(
437
463
  out1=ts_decay_linear(pl.col('a'), 5),
438
464
  out2=ts_WMA(pl.col('a'), 5),
439
465
  )
466
+
440
467
  shape: (6, 3)
441
468
  ┌──────┬──────┬──────┐
442
469
  │ a ┆ out1 ┆ out2 │
@@ -461,7 +488,8 @@ def ts_decay_linear(x: Expr, d: int = 30) -> Expr:
461
488
  # weights = arange(1, d + 1, eager=True)
462
489
  # # print(weights)
463
490
  # return x.rolling_mean(d, weights=weights)
464
- return x.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), roll_decay_linear, d))
491
+ minp = min_samples or polars_ta.MIN_SAMPLES or d
492
+ return x.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), roll_decay_linear, d, minp))
465
493
 
466
494
 
467
495
  def ts_delay(x: Expr, d: int = 1, fill_value=None) -> Expr:
@@ -522,29 +550,64 @@ def ts_fill_null(x: Expr, limit: int = None) -> Expr:
522
550
  return x.forward_fill(limit)
523
551
 
524
552
 
525
- def ts_ir(x: Expr, d: int = 1) -> Expr:
553
+ def ts_ir(x: Expr, d: int = 1, min_samples: Optional[int] = None) -> Expr:
526
554
  """时序滚动信息系数rolling information ratio"""
527
- return ts_mean(x, d) / ts_std_dev(x, d, 0)
555
+ return ts_mean(x, d, min_samples) / ts_std_dev(x, d, 0, min_samples)
528
556
 
529
557
 
530
- def ts_kurtosis(x: Expr, d: int = 5) -> Expr:
558
+ def ts_kurtosis(x: Expr, d: int = 5, bias: bool = False, min_samples: Optional[int] = None) -> Expr:
531
559
  """kurtosis of x for the last d days
532
560
 
533
561
  时序滚动峰度
534
562
 
535
- Warnings
563
+ Parameters
564
+ ----------
565
+ x
566
+ d
567
+ bias
568
+ 有偏
569
+ min_samples
570
+
571
+ Notes
572
+ -----
573
+ `bias=False`时与`pandas`结果一样
574
+
575
+ Examples
536
576
  --------
537
- 等待polars官方出rolling_kurt
577
+ ```python
578
+ df = pl.DataFrame({
579
+ 'a': [None, 1, 2, 3, 4, 999],
580
+ }).with_columns(
581
+ out1=pl.col('a').map_batches(lambda x: pl.Series(pd.Series(x).rolling(4).kurt())),
582
+ out2=ts_kurtosis(pl.col('a'), 4),
583
+ )
584
+
585
+ shape: (6, 3)
586
+ ┌──────┬──────────┬──────────┐
587
+ │ a ┆ out1 ┆ out2 │
588
+ │ --- ┆ --- ┆ --- │
589
+ │ i64 ┆ f64 ┆ f64 │
590
+ ╞══════╪══════════╪══════════╡
591
+ │ null ┆ null ┆ null │
592
+ │ 1 ┆ null ┆ null │
593
+ │ 2 ┆ null ┆ null │
594
+ │ 3 ┆ null ┆ null │
595
+ │ 4 ┆ -1.2 ┆ -1.2 │
596
+ │ 999 ┆ 3.999946 ┆ 3.999946 │
597
+ └──────┴──────────┴──────────┘
598
+ ```
538
599
 
539
600
  """
540
- return x.map_batches(lambda a: roll_kurt(a, d))
601
+ minp = min_samples or polars_ta.MIN_SAMPLES
602
+ return x.rolling_kurtosis(d, min_samples=minp, bias=bias)
541
603
 
542
604
 
543
- def ts_l2_norm(x: Expr, d: int = 5) -> Expr:
605
+ def ts_l2_norm(x: Expr, d: int = 5, min_samples: Optional[int] = None) -> Expr:
544
606
  """Euclidean norm
545
607
 
546
608
  欧几里得范数"""
547
- return x.pow(2).rolling_sum(d).sqrt()
609
+ minp = min_samples or polars_ta.MIN_SAMPLES
610
+ return x.pow(2).rolling_sum(d, min_samples=minp).sqrt()
548
611
 
549
612
 
550
613
  def ts_log_diff(x: Expr, d: int = 1) -> Expr:
@@ -554,47 +617,51 @@ def ts_log_diff(x: Expr, d: int = 1) -> Expr:
554
617
  return x.log().diff(d)
555
618
 
556
619
 
557
- def ts_max(x: Expr, d: int = 30) -> Expr:
620
+ def ts_max(x: Expr, d: int = 30, min_samples: Optional[int] = None) -> Expr:
558
621
  """时序滚动最大值"""
559
- return x.rolling_max(d)
622
+ minp = min_samples or polars_ta.MIN_SAMPLES
623
+ return x.rolling_max(d, min_samples=minp)
560
624
 
561
625
 
562
- def ts_max_diff(x: Expr, d: int = 30) -> Expr:
626
+ def ts_max_diff(x: Expr, d: int = 30, min_samples: Optional[int] = None) -> Expr:
563
627
  """Returns x - ts_max(x, d)"""
564
- return x - ts_max(x, d)
628
+ return x - ts_max(x, d, min_samples)
565
629
 
566
630
 
567
- def ts_mean(x: Expr, d: int = 5) -> Expr:
631
+ def ts_mean(x: Expr, d: int = 5, min_samples: Optional[int] = None) -> Expr:
568
632
  """简单移动平均"""
569
- return x.rolling_mean(d)
633
+ minp = min_samples or polars_ta.MIN_SAMPLES
634
+ return x.rolling_mean(d, min_samples=minp)
570
635
 
571
636
 
572
- def ts_median(x: Expr, d: int = 5) -> Expr:
637
+ def ts_median(x: Expr, d: int = 5, min_samples: Optional[int] = None) -> Expr:
573
638
  """时序滚动中位数"""
574
- return x.rolling_median(d)
639
+ minp = min_samples or polars_ta.MIN_SAMPLES
640
+ return x.rolling_median(d, min_samples=minp)
575
641
 
576
642
 
577
- def ts_min(x: Expr, d: int = 30) -> Expr:
643
+ def ts_min(x: Expr, d: int = 30, min_samples: Optional[int] = None) -> Expr:
578
644
  """时序滚动最小值"""
579
- return x.rolling_min(d)
645
+ minp = min_samples or polars_ta.MIN_SAMPLES
646
+ return x.rolling_min(d, min_samples=minp)
580
647
 
581
648
 
582
- def ts_min_diff(x: Expr, d: int = 30) -> Expr:
649
+ def ts_min_diff(x: Expr, d: int = 30, min_samples: Optional[int] = None) -> Expr:
583
650
  """Returns x - ts_min(x, d)"""
584
- return x - ts_min(x, d)
651
+ return x - ts_min(x, d, min_samples)
585
652
 
586
653
 
587
- def ts_min_max_cps(x: Expr, d: int, f: float = 2.0) -> Expr:
654
+ def ts_min_max_cps(x: Expr, d: int, f: float = 2.0, min_samples: Optional[int] = None) -> Expr:
588
655
  """Returns (ts_min(x, d) + ts_max(x, d)) - f * x"""
589
- return (ts_min(x, d) + ts_max(x, d)) - f * x
656
+ return (ts_min(x, d, min_samples) + ts_max(x, d, min_samples)) - f * x
590
657
 
591
658
 
592
- def ts_min_max_diff(x: Expr, d: int, f: float = 0.5) -> Expr:
659
+ def ts_min_max_diff(x: Expr, d: int, f: float = 0.5, min_samples: Optional[int] = None) -> Expr:
593
660
  """Returns x - f * (ts_min(x, d) + ts_max(x, d))"""
594
- return x - f * (ts_min(x, d) + ts_max(x, d))
661
+ return x - f * (ts_min(x, d, min_samples) + ts_max(x, d, min_samples))
595
662
 
596
663
 
597
- def ts_moment(x: Expr, d: int, k: int = 0) -> Expr:
664
+ def ts_moment(x: Expr, d: int, k: int = 0, min_samples: Optional[int] = None) -> Expr:
598
665
  """Returns K-th central moment of x for the past d days.
599
666
 
600
667
  滚动k阶中心距
@@ -604,20 +671,23 @@ def ts_moment(x: Expr, d: int, k: int = 0) -> Expr:
604
671
  x
605
672
  d
606
673
  k
674
+ min_samples
607
675
 
608
676
  """
609
- return x.map_batches(lambda x1: batches_i1_o1(x1.to_numpy(), roll_moment, d, k))
677
+ minp = min_samples or polars_ta.MIN_SAMPLES or d
678
+ return x.map_batches(lambda x1: batches_i1_o1(x1.to_numpy(), roll_moment, d, minp, k))
610
679
 
611
680
 
612
- def ts_partial_corr(x: Expr, y: Expr, z: Expr, d: int) -> Expr:
681
+ def ts_partial_corr(x: Expr, y: Expr, z: Expr, d: int, min_samples: Optional[int] = None) -> Expr:
613
682
  """Returns partial correlation of x, y, z for the past d days.
614
683
 
615
684
  滚动偏相关
616
685
  """
617
- return struct([x, y, z]).map_batches(lambda xx: batches_i2_o1([xx.struct[i].to_numpy() for i in range(3)], roll_partial_corr, d))
686
+ minp = min_samples or polars_ta.MIN_SAMPLES or d
687
+ return struct([x, y, z]).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 3), roll_partial_corr, d, minp))
618
688
 
619
689
 
620
- def ts_percentage(x: Expr, d: int, percentage: float = 0.5) -> Expr:
690
+ def ts_percentage(x: Expr, d: int, percentage: float = 0.5, min_samples: Optional[int] = None) -> Expr:
621
691
  """Returns percentile value of x for the past d days.
622
692
 
623
693
  滚动百分位数
@@ -627,17 +697,20 @@ def ts_percentage(x: Expr, d: int, percentage: float = 0.5) -> Expr:
627
697
  x
628
698
  d
629
699
  percentage
700
+ min_samples
630
701
 
631
702
  """
632
- return x.rolling_quantile(percentage, window_size=d)
703
+ minp = min_samples or polars_ta.MIN_SAMPLES
704
+ return x.rolling_quantile(percentage, window_size=d, min_samples=minp)
633
705
 
634
706
 
635
- def ts_product(x: Expr, d: int = 5) -> Expr:
707
+ def ts_product(x: Expr, d: int = 5, min_samples: Optional[int] = None) -> Expr:
636
708
  """时序滚动乘"""
637
- return x.map_batches(lambda x1: batches_i1_o1(x1.to_numpy(), roll_prod, d))
709
+ minp = min_samples or polars_ta.MIN_SAMPLES or d
710
+ return x.map_batches(lambda x1: batches_i1_o1(x1.to_numpy(), roll_prod, d, minp))
638
711
 
639
712
 
640
- def ts_rank(x: Expr, d: int = 5) -> Expr:
713
+ def ts_rank(x: Expr, d: int = 5, min_samples: Optional[int] = None) -> Expr:
641
714
  """时序滚动排名
642
715
 
643
716
  Warnings
@@ -645,7 +718,8 @@ def ts_rank(x: Expr, d: int = 5) -> Expr:
645
718
  等待polars官方出rolling_rank
646
719
 
647
720
  """
648
- return x.map_batches(lambda a: roll_rank(a, d, True))
721
+ minp = min_samples or polars_ta.MIN_SAMPLES or d
722
+ return x.map_batches(lambda a: roll_rank(a, d, minp, True))
649
723
 
650
724
 
651
725
  def ts_returns(x: Expr, d: int = 1) -> Expr:
@@ -653,18 +727,18 @@ def ts_returns(x: Expr, d: int = 1) -> Expr:
653
727
  return x.pct_change(d)
654
728
 
655
729
 
656
- def ts_scale(x: Expr, d: int = 5) -> Expr:
730
+ def ts_scale(x: Expr, d: int = 5, min_samples: Optional[int] = None) -> Expr:
657
731
  """Returns (x – ts_min(x, d)) / (ts_max(x, d) – ts_min(x, d)) + constant
658
732
 
659
733
  时序滚动缩放
660
734
  """
661
- a = ts_min(x, d)
662
- b = ts_max(x, d)
735
+ a = ts_min(x, d, min_samples)
736
+ b = ts_max(x, d, min_samples)
663
737
  # return (x - a) / (b - a + TA_EPSILON)
664
738
  return when(a != b).then((x - a) / (b - a)).otherwise(0)
665
739
 
666
740
 
667
- def ts_skewness(x: Expr, d: int = 5, bias: bool = False) -> Expr:
741
+ def ts_skewness(x: Expr, d: int = 5, bias: bool = False, min_samples: Optional[int] = None) -> Expr:
668
742
  """Return skewness of x for the past d days
669
743
 
670
744
  时序滚动偏度
@@ -675,16 +749,18 @@ def ts_skewness(x: Expr, d: int = 5, bias: bool = False) -> Expr:
675
749
  d
676
750
  bias
677
751
  有偏
752
+ min_samples
678
753
 
679
754
  Notes
680
755
  -----
681
756
  `bias=False`时与`pandas`结果一样
682
757
 
683
758
  """
684
- return x.rolling_skew(d, bias=bias)
759
+ minp = min_samples or polars_ta.MIN_SAMPLES
760
+ return x.rolling_skew(d, min_samples=minp, bias=bias)
685
761
 
686
762
 
687
- def ts_std_dev(x: Expr, d: int = 5, ddof: int = 0) -> Expr:
763
+ def ts_std_dev(x: Expr, d: int = 5, ddof: int = 0, min_samples: Optional[int] = None) -> Expr:
688
764
  """时序滚动标准差
689
765
 
690
766
  Parameters
@@ -693,14 +769,17 @@ def ts_std_dev(x: Expr, d: int = 5, ddof: int = 0) -> Expr:
693
769
  d
694
770
  ddof
695
771
  自由度
772
+ min_samples
696
773
 
697
774
  """
698
- return x.rolling_std(d, ddof=ddof)
775
+ minp = min_samples or polars_ta.MIN_SAMPLES
776
+ return x.rolling_std(d, ddof=ddof, min_samples=minp)
699
777
 
700
778
 
701
- def ts_sum(x: Expr, d: int = 30) -> Expr:
779
+ def ts_sum(x: Expr, d: int = 30, min_samples: Optional[int] = None) -> Expr:
702
780
  """时序滚动求和"""
703
- return x.rolling_sum(d)
781
+ minp = min_samples or polars_ta.MIN_SAMPLES
782
+ return x.rolling_sum(d, min_samples=minp)
704
783
 
705
784
 
706
785
  def ts_sum_split_by(x: Expr, by: Expr, d: int = 30, k: int = 10) -> Expr:
@@ -753,18 +832,19 @@ def ts_sum_split_by(x: Expr, by: Expr, d: int = 30, k: int = 10) -> Expr:
753
832
 
754
833
  """
755
834
  dtype = Struct([Field(f"column_{i}", Float64) for i in range(2)])
756
- return struct([x, by]).map_batches(lambda xx: batches_i2_o2([xx.struct[i].to_numpy() for i in range(2)], _sum_split_by, d, k), return_dtype=dtype)
835
+ return struct([x, by]).map_batches(lambda xx: batches_i2_o2(struct_to_numpy(xx, 2), _sum_split_by, d, k), return_dtype=dtype)
757
836
 
758
837
 
759
- def ts_triple_corr(x: Expr, y: Expr, z: Expr, d: int) -> Expr:
838
+ def ts_triple_corr(x: Expr, y: Expr, z: Expr, d: int, min_samples: Optional[int] = None) -> Expr:
760
839
  """时序滚动三重相关系数 Returns triple correlation of x, y, z for the past d days.
761
840
 
762
841
 
763
842
  """
764
- return struct([x, y, z]).map_batches(lambda xx: batches_i2_o1([xx.struct[i].to_numpy() for i in range(3)], roll_triple_corr, d))
843
+ minp = min_samples or polars_ta.MIN_SAMPLES or d
844
+ return struct([x, y, z]).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 3), roll_triple_corr, d, minp))
765
845
 
766
846
 
767
- def ts_weighted_decay(x: Expr, k: float = 0.5) -> Expr:
847
+ def ts_weighted_decay(x: Expr, k: float = 0.5, min_samples: Optional[int] = None) -> Expr:
768
848
  """加权衰减 Instead of replacing today’s value with yesterday’s as in ts_delay(x, 1),
769
849
  it assigns weighted average of today’s and yesterday’s values with weight on today’s value being k and yesterday’s being (1-k).
770
850
 
@@ -773,14 +853,16 @@ def ts_weighted_decay(x: Expr, k: float = 0.5) -> Expr:
773
853
  x
774
854
  k
775
855
  衰减系数
856
+ min_samples
776
857
 
777
858
  """
778
- return x.rolling_sum(2, weights=[1 - k, k])
859
+ minp = min_samples or polars_ta.MIN_SAMPLES
860
+ return x.rolling_sum(2, weights=[1 - k, k], min_samples=minp)
779
861
 
780
862
 
781
- def ts_zscore(x: Expr, d: int = 5) -> Expr:
863
+ def ts_zscore(x: Expr, d: int = 5, min_samples: Optional[int] = None) -> Expr:
782
864
  """时序滚动zscore"""
783
- return (x - ts_mean(x, d)) / ts_std_dev(x, d, 0)
865
+ return (x - ts_mean(x, d, min_samples)) / ts_std_dev(x, d, 0, min_samples)
784
866
 
785
867
 
786
868
  def ts_cum_prod_by(r: Expr, v: Expr) -> Expr:
@@ -834,7 +916,7 @@ def ts_cum_prod_by(r: Expr, v: Expr) -> Expr:
834
916
 
835
917
 
836
918
  """
837
- return struct([r, v]).map_batches(lambda xx: batches_i2_o1([xx.struct[i].to_numpy().astype(float) for i in range(2)], _cum_prod_by))
919
+ return struct([r, v]).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 2), _cum_prod_by))
838
920
 
839
921
 
840
922
  def ts_cum_sum_by(r: Expr, v: Expr) -> Expr:
@@ -883,32 +965,36 @@ def ts_cum_sum_by(r: Expr, v: Expr) -> Expr:
883
965
  ```
884
966
 
885
967
  """
886
- return struct([r, v]).map_batches(lambda xx: batches_i2_o1([xx.struct[i].to_numpy().astype(float) for i in range(2)], _cum_sum_by))
968
+ return struct([r, v]).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 2), _cum_sum_by))
887
969
 
888
970
 
889
- def ts_regression_resid(y: Expr, x: Expr, d: int) -> Expr:
971
+ def ts_regression_resid(y: Expr, x: Expr, d: int, min_samples: Optional[int] = None) -> Expr:
890
972
  """时序滚动回归取残差"""
891
- return pls.compute_rolling_least_squares(y, x, mode='residuals', add_intercept=True, rolling_kwargs=RollingKwargs(window_size=d, min_periods=d))
973
+ minp = min_samples or polars_ta.MIN_SAMPLES or d
974
+ return pls.compute_rolling_least_squares(y, x, mode='residuals', add_intercept=True, rolling_kwargs=RollingKwargs(window_size=d, min_periods=minp))
892
975
 
893
976
 
894
- def ts_regression_pred(y: Expr, x: Expr, d: int) -> Expr:
977
+ def ts_regression_pred(y: Expr, x: Expr, d: int, min_samples: Optional[int] = None) -> Expr:
895
978
  """时序滚动回归取y的预测值
896
979
  """
897
- return pls.compute_rolling_least_squares(y, x, mode='predictions', add_intercept=True, rolling_kwargs=RollingKwargs(window_size=d, min_periods=d))
980
+ minp = min_samples or polars_ta.MIN_SAMPLES or d
981
+ return pls.compute_rolling_least_squares(y, x, mode='predictions', add_intercept=True, rolling_kwargs=RollingKwargs(window_size=d, min_periods=minp))
898
982
 
899
983
 
900
- def ts_regression_intercept(y: Expr, x: Expr, d: int) -> Expr:
984
+ def ts_regression_intercept(y: Expr, x: Expr, d: int, min_samples: Optional[int] = None) -> Expr:
901
985
  """时序滚动回归取截距
902
986
  """
903
- return pls.compute_rolling_least_squares(y, x, mode='coefficients', add_intercept=True, rolling_kwargs=RollingKwargs(window_size=d, min_periods=d)).struct[1]
987
+ minp = min_samples or polars_ta.MIN_SAMPLES or d
988
+ return pls.compute_rolling_least_squares(y, x, mode='coefficients', add_intercept=True, rolling_kwargs=RollingKwargs(window_size=d, min_periods=minp)).struct[1]
904
989
 
905
990
 
906
- def ts_regression_slope(y: Expr, x: Expr, d: int) -> Expr:
991
+ def ts_regression_slope(y: Expr, x: Expr, d: int, min_samples: Optional[int] = None) -> Expr:
907
992
  """时序滚动回归取斜率"""
908
- return pls.compute_rolling_least_squares(y, x, mode='coefficients', add_intercept=True, rolling_kwargs=RollingKwargs(window_size=d, min_periods=d)).struct[0]
993
+ minp = min_samples or polars_ta.MIN_SAMPLES or d
994
+ return pls.compute_rolling_least_squares(y, x, mode='coefficients', add_intercept=True, rolling_kwargs=RollingKwargs(window_size=d, min_periods=minp)).struct[0]
909
995
 
910
996
 
911
- def ts_resid(y: Expr, *more_x: Expr, d: int = 30) -> Expr:
997
+ def ts_resid(y: Expr, *more_x: Expr, d: int = 30, min_samples: Optional[int] = None) -> Expr:
912
998
  """多元时序滚动回归取残差
913
999
 
914
1000
  Parameters
@@ -917,12 +1003,14 @@ def ts_resid(y: Expr, *more_x: Expr, d: int = 30) -> Expr:
917
1003
  *more_x
918
1004
  多个x
919
1005
  d
1006
+ min_samples
920
1007
 
921
1008
  """
922
- return pls.compute_rolling_least_squares(y, *more_x, mode='residuals', rolling_kwargs=RollingKwargs(window_size=d, min_periods=d))
1009
+ minp = min_samples or polars_ta.MIN_SAMPLES or d
1010
+ return pls.compute_rolling_least_squares(y, *more_x, mode='residuals', rolling_kwargs=RollingKwargs(window_size=d, min_periods=minp))
923
1011
 
924
1012
 
925
- def ts_pred(y: Expr, *more_x: Expr, d: int = 30) -> Expr:
1013
+ def ts_pred(y: Expr, *more_x: Expr, d: int = 30, min_samples: Optional[int] = None) -> Expr:
926
1014
  """多元时序滚动回归预测
927
1015
 
928
1016
  Parameters
@@ -931,19 +1019,23 @@ def ts_pred(y: Expr, *more_x: Expr, d: int = 30) -> Expr:
931
1019
  *more_x
932
1020
  多个x
933
1021
  d
1022
+ min_samples
934
1023
 
935
1024
  """
936
- return pls.compute_rolling_least_squares(y, *more_x, mode='predictions', rolling_kwargs=RollingKwargs(window_size=d, min_periods=d))
1025
+ minp = min_samples or polars_ta.MIN_SAMPLES or d
1026
+ return pls.compute_rolling_least_squares(y, *more_x, mode='predictions', rolling_kwargs=RollingKwargs(window_size=d, min_periods=minp))
937
1027
 
938
1028
 
939
- def ts_weighted_mean(x: Expr, w: Expr, d: int) -> Expr:
1029
+ def ts_weighted_mean(x: Expr, w: Expr, d: int, min_samples: Optional[int] = None) -> Expr:
940
1030
  """时序滚动加权平均"""
941
- return (x * w).rolling_sum(d) / w.rolling_sum(d)
1031
+ minp = min_samples or polars_ta.MIN_SAMPLES
1032
+ return (x * w).rolling_sum(d, min_samples=minp) / w.rolling_sum(d, min_samples=minp)
942
1033
 
943
1034
 
944
- def ts_weighted_sum(x: Expr, w: Expr, d: int) -> Expr:
1035
+ def ts_weighted_sum(x: Expr, w: Expr, d: int, min_samples: Optional[int] = None) -> Expr:
945
1036
  """时序滚动加权求和"""
946
- return (x * w).rolling_sum(d)
1037
+ minp = min_samples or polars_ta.MIN_SAMPLES
1038
+ return (x * w).rolling_sum(d, min_samples=minp)
947
1039
 
948
1040
 
949
1041
  def ts_signals_to_size(long_entry: Expr, long_exit: Expr,
@@ -969,5 +1061,4 @@ def ts_signals_to_size(long_entry: Expr, long_exit: Expr,
969
1061
 
970
1062
  """
971
1063
  return struct([long_entry, long_exit, short_entry, short_exit]).map_batches(
972
- lambda xx: batches_i2_o1([xx.struct[i].to_numpy().astype(float) for i in range(4)],
973
- _signals_to_size, accumulate, action))
1064
+ lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _signals_to_size, accumulate, action))