CoCalc -- test_indicators.py

GitHub Repository: polakowo/vectorbt
Path: blob/master/tests/test_indicators.py
¹⁰⁷¹ views
1
from collections import namedtuple
2
from datetime import datetime
3
from itertools import product
4

5
import numpy as np
6
import pandas as pd
7
import pytest
8
from numba import njit
9

10
import vectorbt as vbt
11

12
ray_available = True
13
try:
14
    import ray
15
except:
16
    ray_available = False
17

18
ta_available = True
19
try:
20
    import ta
21
except:
22
    ta_available = False
23

24
pandas_ta_available = True
25
try:
26
    import pandas_ta
27
except:
28
    try:
29
        import pandas_ta_classic as pandas_ta
30
    except:
31
        pandas_ta_available = False
32

33
talib_available = True
34
try:
35
    import talib
36
except:
37
    talib_available = False
38

39
seed = 42
40

41

42
# ############# Global ############# #
43

44
def setup_module():
45
    vbt.settings.numba['check_func_suffix'] = True
46
    vbt.settings.caching.enabled = False
47
    vbt.settings.caching.whitelist = []
48
    vbt.settings.caching.blacklist = []
49
    if ray_available:
50
        ray.init(local_mode=True, num_cpus=1)
51

52

53
def teardown_module():
54
    if ray_available:
55
        ray.shutdown()
56
    vbt.settings.reset()
57

58

59
# ############# factory.py ############# #
60

61
ts = pd.DataFrame({
62
    'a': [1., 2., 3., 4., 5.],
63
    'b': [5., 4., 3., 2., 1.],
64
    'c': [1., 2., 3., 2., 1.]
65
}, index=pd.DatetimeIndex([
66
    datetime(2018, 1, 1),
67
    datetime(2018, 1, 2),
68
    datetime(2018, 1, 3),
69
    datetime(2018, 1, 4),
70
    datetime(2018, 1, 5)
71
]))
72

73

74
class TestFactory:
75
    def test_config(self, tmp_path):
76
        F = vbt.IndicatorFactory(input_names=['ts'], param_names=['p'], output_names=['out'])
77

78
        def apply_func(ts, p, a, b=10):
79
            return ts * p + a + b
80

81
        I = F.from_apply_func(apply_func, var_args=True)
82
        indicator = I.run(ts, [0, 1], 10, b=100)
83
        assert I.loads(indicator.dumps()) == indicator
84
        indicator.save(tmp_path / 'indicator')
85
        assert I.load(tmp_path / 'indicator') == indicator
86

87
    def test_from_custom_func(self):
88
        F = vbt.IndicatorFactory(input_names=['ts'], param_names=['p'], output_names=['out'])
89

90
        def apply_func(i, ts, p, a, b=10):
91
            return ts * p[i] + a + b
92

93
        @njit
94
        def apply_func_nb(i, ts, p, a, b):
95
            return ts * p[i] + a + b  # numba doesn't support **kwargs
96

97
        def custom_func(ts, p, *args, **kwargs):
98
            return vbt.base.combine_fns.apply_and_concat_one(len(p), apply_func, ts, p, *args, **kwargs)
99

100
        @njit
101
        def custom_func_nb(ts, p, *args):
102
            return vbt.base.combine_fns.apply_and_concat_one_nb(len(p), apply_func_nb, ts, p, *args)
103

104
        target = pd.DataFrame(
105
            np.array([
106
                [110., 110., 110., 111., 115., 111.],
107
                [110., 110., 110., 112., 114., 112.],
108
                [110., 110., 110., 113., 113., 113.],
109
                [110., 110., 110., 114., 112., 112.],
110
                [110., 110., 110., 115., 111., 111.]
111
            ]),
112
            index=ts.index,
113
            columns=pd.MultiIndex.from_tuples([
114
                (0, 'a'),
115
                (0, 'b'),
116
                (0, 'c'),
117
                (1, 'a'),
118
                (1, 'b'),
119
                (1, 'c')
120
            ], names=['custom_p', None])
121
        )
122
        pd.testing.assert_frame_equal(
123
            F.from_custom_func(custom_func, var_args=True).run(ts, [0, 1], 10, b=100).out,
124
            target
125
        )
126
        pd.testing.assert_frame_equal(
127
            F.from_custom_func(custom_func_nb, var_args=True).run(ts, [0, 1], 10, 100).out,
128
            target
129
        )
130
        target = pd.DataFrame(
131
            np.array([
132
                [110., 115., 112.],
133
                [110., 114., 114.],
134
                [110., 113., 116.],
135
                [110., 112., 114.],
136
                [110., 111., 112.]
137
            ]),
138
            index=ts.index,
139
            columns=pd.MultiIndex.from_tuples([
140
                (0, 'a'),
141
                (1, 'b'),
142
                (2, 'c')
143
            ], names=['custom_p', None])
144
        )
145
        pd.testing.assert_frame_equal(
146
            F.from_custom_func(custom_func, var_args=True).run(ts, [0, 1, 2], 10, b=100, per_column=True).out,
147
            target
148
        )
149
        pd.testing.assert_frame_equal(
150
            F.from_custom_func(custom_func_nb, var_args=True).run(ts, [0, 1, 2], 10, 100, per_column=True).out,
151
            target
152
        )
153
        target = pd.DataFrame(
154
            np.array([
155
                [110., 111.],
156
                [110., 112.],
157
                [110., 113.],
158
                [110., 114.],
159
                [110., 115.]
160
            ]),
161
            index=ts.index,
162
            columns=pd.Index([0, 1], dtype='int64', name='custom_p')
163
        )
164
        pd.testing.assert_frame_equal(
165
            F.from_custom_func(custom_func, var_args=True).run(ts['a'], [0, 1], 10, b=100).out,
166
            target
167
        )
168
        pd.testing.assert_frame_equal(
169
            F.from_custom_func(custom_func_nb, var_args=True).run(ts['a'], [0, 1], 10, 100).out,
170
            target
171
        )
172
        target = pd.DataFrame(
173
            np.array([
174
                [110.],
175
                [110.],
176
                [110.],
177
                [110.],
178
                [110.]
179
            ]),
180
            index=ts.index,
181
            columns=pd.MultiIndex.from_tuples([(0, 'a')], names=['custom_p', None])
182
        )
183
        pd.testing.assert_frame_equal(
184
            F.from_custom_func(custom_func, var_args=True).run(ts[['a']], 0, 10, b=100, per_column=True).out,
185
            target
186
        )
187
        pd.testing.assert_frame_equal(
188
            F.from_custom_func(custom_func_nb, var_args=True).run(ts[['a']], 0, 10, 100, per_column=True).out,
189
            target
190
        )
191
        target = pd.Series(
192
            np.array([110., 110., 110., 110., 110.]),
193
            index=ts.index,
194
            name=(0, 'a')
195
        )
196
        pd.testing.assert_series_equal(
197
            F.from_custom_func(custom_func, var_args=True).run(ts['a'], 0, 10, b=100).out,
198
            target
199
        )
200
        pd.testing.assert_series_equal(
201
            F.from_custom_func(custom_func_nb, var_args=True).run(ts['a'], 0, 10, 100).out,
202
            target
203
        )
204
        pd.testing.assert_series_equal(
205
            F.from_custom_func(custom_func, var_args=True).run(ts['a'], 0, 10, b=100, per_column=True).out,
206
            target
207
        )
208
        pd.testing.assert_series_equal(
209
            F.from_custom_func(custom_func_nb, var_args=True).run(ts['a'], 0, 10, 100, per_column=True).out,
210
            target
211
        )
212

213
    def test_from_apply_func(self):
214
        F = vbt.IndicatorFactory(input_names=['ts'], param_names=['p'], output_names=['out'])
215

216
        def apply_func(ts, p, a, b=10):
217
            return ts * p + a + b
218

219
        @njit
220
        def apply_func_nb(ts, p, a, b):
221
            return ts * p + a + b  # numba doesn't support **kwargs
222

223
        target = pd.DataFrame(
224
            np.array([
225
                [110., 110., 110., 111., 115., 111.],
226
                [110., 110., 110., 112., 114., 112.],
227
                [110., 110., 110., 113., 113., 113.],
228
                [110., 110., 110., 114., 112., 112.],
229
                [110., 110., 110., 115., 111., 111.]
230
            ]),
231
            index=ts.index,
232
            columns=pd.MultiIndex.from_tuples([
233
                (0, 'a'),
234
                (0, 'b'),
235
                (0, 'c'),
236
                (1, 'a'),
237
                (1, 'b'),
238
                (1, 'c')
239
            ], names=['custom_p', None])
240
        )
241
        pd.testing.assert_frame_equal(
242
            F.from_apply_func(apply_func, var_args=True).run(ts, [0, 1], 10, b=100).out,
243
            target
244
        )
245
        pd.testing.assert_frame_equal(
246
            F.from_apply_func(apply_func_nb, var_args=True).run(ts, [0, 1], 10, 100).out,
247
            target
248
        )
249
        target = pd.DataFrame(
250
            np.array([
251
                [110., 115., 112.],
252
                [110., 114., 114.],
253
                [110., 113., 116.],
254
                [110., 112., 114.],
255
                [110., 111., 112.]
256
            ]),
257
            index=ts.index,
258
            columns=pd.MultiIndex.from_tuples([
259
                (0, 'a'),
260
                (1, 'b'),
261
                (2, 'c')
262
            ], names=['custom_p', None])
263
        )
264
        pd.testing.assert_frame_equal(
265
            F.from_apply_func(apply_func, var_args=True).run(ts, [0, 1, 2], 10, b=100, per_column=True).out,
266
            target
267
        )
268
        pd.testing.assert_frame_equal(
269
            F.from_apply_func(apply_func, var_args=True).run(ts, [0, 1, 2], 10, 100, per_column=True).out,
270
            target
271
        )
272
        target = pd.DataFrame(
273
            np.array([
274
                [110., 111.],
275
                [110., 112.],
276
                [110., 113.],
277
                [110., 114.],
278
                [110., 115.]
279
            ]),
280
            index=ts.index,
281
            columns=pd.Index([0, 1], dtype='int64', name='custom_p')
282
        )
283
        pd.testing.assert_frame_equal(
284
            F.from_apply_func(apply_func, var_args=True).run(ts['a'], [0, 1], 10, b=100).out,
285
            target
286
        )
287
        pd.testing.assert_frame_equal(
288
            F.from_apply_func(apply_func_nb, numba_loop=True, var_args=True).run(ts['a'], [0, 1], 10, 100).out,
289
            target
290
        )
291
        target = pd.DataFrame(
292
            np.array([
293
                [110.],
294
                [110.],
295
                [110.],
296
                [110.],
297
                [110.]
298
            ]),
299
            index=ts.index,
300
            columns=pd.MultiIndex.from_tuples([(0, 'a')], names=['custom_p', None])
301
        )
302
        pd.testing.assert_frame_equal(
303
            F.from_apply_func(apply_func, var_args=True).run(ts[['a']], 0, 10, b=100, per_column=True).out,
304
            target
305
        )
306
        pd.testing.assert_frame_equal(
307
            F.from_apply_func(apply_func, var_args=True).run(ts[['a']], 0, 10, 100, per_column=True).out,
308
            target
309
        )
310
        target = pd.Series(
311
            np.array([110., 110., 110., 110., 110.]),
312
            index=ts.index,
313
            name=(0, 'a')
314
        )
315
        pd.testing.assert_series_equal(
316
            F.from_apply_func(apply_func, var_args=True).run(ts['a'], 0, 10, b=100).out,
317
            target
318
        )
319
        pd.testing.assert_series_equal(
320
            F.from_apply_func(apply_func_nb, numba_loop=True, var_args=True)
321
                .run(ts['a'], 0, 10, 100).out,
322
            target
323
        )
324
        pd.testing.assert_series_equal(
325
            F.from_apply_func(apply_func, var_args=True).run(ts['a'], 0, 10, b=100, per_column=True).out,
326
            target
327
        )
328
        pd.testing.assert_series_equal(
329
            F.from_apply_func(apply_func_nb, numba_loop=True, var_args=True)
330
                .run(ts['a'], 0, 10, 100, per_column=True).out,
331
            target
332
        )
333

334
    def test_use_ray(self):
335
        if ray_available:
336
            F = vbt.IndicatorFactory(input_names=['ts'], param_names=['p'], output_names=['out'])
337

338
            def apply_func(ts, p, a, b=10):
339
                return ts * p + a + b
340

341
            pd.testing.assert_frame_equal(
342
                F.from_apply_func(apply_func, var_args=True)
343
                    .run(ts, np.arange(10), 10, b=100).out,
344
                F.from_apply_func(apply_func, var_args=True)
345
                    .run(ts, np.arange(10), 10, b=100, use_ray=True).out,
346
            )
347

348
    def test_no_inputs(self):
349
        F = vbt.IndicatorFactory(param_names=['p'], output_names=['out'])
350

351
        def apply_func(p):
352
            return np.full((3, 3), p)
353

354
        @njit
355
        def apply_func_nb(p):
356
            return np.full((3, 3), p)
357

358
        target = pd.DataFrame(
359
            np.array([
360
                [0, 0, 0, 1, 1, 1],
361
                [0, 0, 0, 1, 1, 1],
362
                [0, 0, 0, 1, 1, 1]
363
            ]),
364
            index=pd.RangeIndex(start=0, stop=3, step=1),
365
            columns=pd.MultiIndex.from_tuples([
366
                (0, 0),
367
                (0, 1),
368
                (0, 2),
369
                (1, 0),
370
                (1, 1),
371
                (1, 2)
372
            ], names=['custom_p', None])
373
        )
374
        pd.testing.assert_frame_equal(
375
            F.from_apply_func(apply_func).run([0, 1]).out,
376
            target
377
        )
378
        pd.testing.assert_frame_equal(
379
            F.from_apply_func(apply_func_nb, numba_loop=True).run([0, 1]).out,
380
            target
381
        )
382
        with pytest.raises(Exception):
383
            F.from_apply_func(apply_func).run([0, 1], per_column=True)
384

385
    def test_input_shape(self):
386
        F = vbt.IndicatorFactory(param_names=['p'], output_names=['out'])
387

388
        def apply_func(input_shape, p):
389
            return np.full(input_shape, p)
390

391
        @njit
392
        def apply_func_nb(input_shape, p):
393
            return np.full(input_shape, p)
394

395
        target = pd.Series(
396
            np.array([0, 0, 0, 0, 0]),
397
            index=pd.RangeIndex(start=0, stop=5, step=1)
398
        )
399
        pd.testing.assert_series_equal(
400
            F.from_apply_func(apply_func, require_input_shape=True).run(5, 0).out,
401
            target
402
        )
403
        pd.testing.assert_series_equal(
404
            F.from_apply_func(apply_func_nb, numba_loop=True, require_input_shape=True).run(5, 0).out,
405
            target
406
        )
407
        target = pd.DataFrame(
408
            np.array([
409
                [0, 1],
410
                [0, 1],
411
                [0, 1],
412
                [0, 1],
413
                [0, 1]
414
            ]),
415
            index=pd.RangeIndex(start=0, stop=5, step=1),
416
            columns=pd.Index([0, 1], dtype='int64', name='custom_p')
417
        )
418
        pd.testing.assert_frame_equal(
419
            F.from_apply_func(apply_func, require_input_shape=True).run(5, [0, 1]).out,
420
            target
421
        )
422
        pd.testing.assert_frame_equal(
423
            F.from_apply_func(apply_func_nb, numba_loop=True, require_input_shape=True).run(5, [0, 1]).out,
424
            target
425
        )
426
        target = pd.DataFrame(
427
            np.array([
428
                [0, 0, 0, 1, 1, 1],
429
                [0, 0, 0, 1, 1, 1],
430
                [0, 0, 0, 1, 1, 1],
431
                [0, 0, 0, 1, 1, 1],
432
                [0, 0, 0, 1, 1, 1]
433
            ]),
434
            index=ts.index,
435
            columns=pd.MultiIndex.from_tuples([
436
                (0, 'a'),
437
                (0, 'b'),
438
                (0, 'c'),
439
                (1, 'a'),
440
                (1, 'b'),
441
                (1, 'c')
442
            ], names=['custom_p', None])
443
        )
444
        pd.testing.assert_frame_equal(
445
            F.from_apply_func(apply_func, require_input_shape=True).run(
446
                (5, 3), [0, 1], input_index=ts.index, input_columns=ts.columns).out,
447
            target
448
        )
449
        pd.testing.assert_frame_equal(
450
            F.from_apply_func(apply_func_nb, numba_loop=True, require_input_shape=True).run(
451
                (5, 3), [0, 1], input_index=ts.index, input_columns=ts.columns).out,
452
            target
453
        )
454
        target = pd.DataFrame(
455
            np.array([
456
                [0, 1, 2],
457
                [0, 1, 2],
458
                [0, 1, 2],
459
                [0, 1, 2],
460
                [0, 1, 2]
461
            ]),
462
            index=ts.index,
463
            columns=pd.MultiIndex.from_tuples([
464
                (0, 'a'),
465
                (1, 'b'),
466
                (2, 'c')
467
            ], names=['custom_p', None])
468
        )
469
        pd.testing.assert_frame_equal(
470
            F.from_apply_func(apply_func, require_input_shape=True).run(
471
                (5, 3), [0, 1, 2], input_index=ts.index, input_columns=ts.columns, per_column=True).out,
472
            target
473
        )
474
        pd.testing.assert_frame_equal(
475
            F.from_apply_func(apply_func_nb, numba_loop=True, require_input_shape=True).run(
476
                (5, 3), [0, 1, 2], input_index=ts.index, input_columns=ts.columns, per_column=True).out,
477
            target
478
        )
479

480
    def test_multiple_inputs(self):
481
        F = vbt.IndicatorFactory(input_names=['ts1', 'ts2'], param_names=['p'], output_names=['out'])
482

483
        def apply_func(ts1, ts2, p):
484
            return ts1 * ts2 * p
485

486
        @njit
487
        def apply_func_nb(ts1, ts2, p):
488
            return ts1 * ts2 * p
489

490
        target = pd.DataFrame(
491
            np.array([
492
                [0., 0., 0., 1., 25., 1.],
493
                [0., 0., 0., 4., 16., 4.],
494
                [0., 0., 0., 9., 9., 9.],
495
                [0., 0., 0., 16., 4., 4.],
496
                [0., 0., 0., 25., 1., 1.]
497
            ]),
498
            index=ts.index,
499
            columns=pd.MultiIndex.from_tuples([
500
                (0, 'a'),
501
                (0, 'b'),
502
                (0, 'c'),
503
                (1, 'a'),
504
                (1, 'b'),
505
                (1, 'c')
506
            ], names=['custom_p', None])
507
        )
508
        pd.testing.assert_frame_equal(
509
            F.from_apply_func(apply_func).run(ts, ts, [0, 1]).out,
510
            target
511
        )
512
        pd.testing.assert_frame_equal(
513
            F.from_apply_func(apply_func_nb, numba_loop=True).run(ts, ts, [0, 1]).out,
514
            target
515
        )
516
        target = pd.DataFrame(
517
            np.array([
518
                [0., 25., 2.],
519
                [0., 16., 8.],
520
                [0., 9., 18.],
521
                [0., 4., 8.],
522
                [0., 1., 2.]
523
            ]),
524
            index=ts.index,
525
            columns=pd.MultiIndex.from_tuples([
526
                (0, 'a'),
527
                (1, 'b'),
528
                (2, 'c')
529
            ], names=['custom_p', None])
530
        )
531
        pd.testing.assert_frame_equal(
532
            F.from_apply_func(apply_func).run(ts, ts, [0, 1, 2], per_column=True).out,
533
            target
534
        )
535
        pd.testing.assert_frame_equal(
536
            F.from_apply_func(apply_func_nb, numba_loop=True).run(ts, ts, [0, 1, 2], per_column=True).out,
537
            target
538
        )
539

540
    def test_no_params(self):
541
        F = vbt.IndicatorFactory(input_names=['ts'], output_names=['out'])
542

543
        def apply_func(ts):
544
            return ts * 2
545

546
        @njit
547
        def apply_func_nb(ts):
548
            return ts * 2
549

550
        pd.testing.assert_frame_equal(
551
            F.from_apply_func(apply_func).run(ts).out,
552
            ts * 2
553
        )
554
        pd.testing.assert_frame_equal(
555
            F.from_apply_func(apply_func_nb, numba_loop=True).run(ts).out,
556
            ts * 2
557
        )
558
        pd.testing.assert_frame_equal(
559
            F.from_apply_func(apply_func).run(ts, per_column=True).out,
560
            ts * 2
561
        )
562
        pd.testing.assert_frame_equal(
563
            F.from_apply_func(apply_func_nb, numba_loop=True).run(ts, per_column=True).out,
564
            ts * 2
565
        )
566

567
    def test_no_inputs_and_params(self):
568
        F = vbt.IndicatorFactory(output_names=['out'])
569

570
        def apply_func():
571
            return np.full((3, 3), 1)
572

573
        @njit
574
        def apply_func_nb():
575
            return np.full((3, 3), 1)
576

577
        pd.testing.assert_frame_equal(
578
            F.from_apply_func(apply_func).run().out,
579
            pd.DataFrame(np.full((3, 3), 1))
580
        )
581
        pd.testing.assert_frame_equal(
582
            F.from_apply_func(apply_func_nb, numba_loop=True).run().out,
583
            pd.DataFrame(np.full((3, 3), 1))
584
        )
585
        with pytest.raises(Exception):
586
            F.from_apply_func(apply_func).run(per_column=True)
587

588
    def test_multiple_params(self):
589
        F = vbt.IndicatorFactory(input_names=['ts'], param_names=['p1', 'p2'], output_names=['out'])
590

591
        def apply_func(ts, p1, p2):
592
            return ts * (p1 + p2)
593

594
        @njit
595
        def apply_func_nb(ts, p1, p2):
596
            return ts * (p1 + p2)
597

598
        target = pd.DataFrame(
599
            np.array([
600
                [2., 10., 2., 3., 15., 3.],
601
                [4., 8., 4., 6., 12., 6.],
602
                [6., 6., 6., 9., 9., 9.],
603
                [8., 4., 4., 12., 6., 6.],
604
                [10., 2., 2., 15., 3., 3.]
605
            ]),
606
            index=ts.index,
607
            columns=pd.MultiIndex.from_tuples([
608
                (0, 2, 'a'),
609
                (0, 2, 'b'),
610
                (0, 2, 'c'),
611
                (1, 2, 'a'),
612
                (1, 2, 'b'),
613
                (1, 2, 'c')
614
            ], names=['custom_p1', 'custom_p2', None])
615
        )
616
        pd.testing.assert_frame_equal(
617
            F.from_apply_func(apply_func).run(ts, np.asarray([0, 1]), 2).out,
618
            target
619
        )
620
        pd.testing.assert_frame_equal(
621
            F.from_apply_func(apply_func_nb, numba_loop=True).run(ts, np.asarray([0, 1]), 2).out,
622
            target
623
        )
624
        target = pd.DataFrame(
625
            np.array([
626
                [2., 15., 4.],
627
                [4., 12., 8.],
628
                [6., 9., 12.],
629
                [8., 6., 8.],
630
                [10., 3., 4.]
631
            ]),
632
            index=ts.index,
633
            columns=pd.MultiIndex.from_tuples([
634
                (0, 2, 'a'),
635
                (1, 2, 'b'),
636
                (2, 2, 'c')
637
            ], names=['custom_p1', 'custom_p2', None])
638
        )
639
        pd.testing.assert_frame_equal(
640
            F.from_apply_func(apply_func).run(ts, np.asarray([0, 1, 2]), 2, per_column=True).out,
641
            target
642
        )
643
        pd.testing.assert_frame_equal(
644
            F.from_apply_func(apply_func_nb, numba_loop=True).run(ts, np.asarray([0, 1, 2]), 2, per_column=True).out,
645
            target
646
        )
647
        pd.testing.assert_frame_equal(
648
            F.from_apply_func(apply_func).run(ts, np.asarray([0, 1, 2]), [2], per_column=True).out,
649
            target
650
        )
651
        pd.testing.assert_frame_equal(
652
            F.from_apply_func(apply_func).run(ts, np.asarray([0, 1, 2]), np.array([2]), per_column=True).out,
653
            target
654
        )
655
        with pytest.raises(Exception):
656
            F.from_apply_func(apply_func).run(ts, np.asarray([0, 1]), 2, per_column=True)
657
        with pytest.raises(Exception):
658
            F.from_apply_func(apply_func).run(ts, np.asarray([0, 1, 2, 3]), 2, per_column=True)
659

660
    def test_param_settings(self):
661
        F = vbt.IndicatorFactory(input_names=['ts'], param_names=['p'], output_names=['out'])
662

663
        def apply_func(ts, p):
664
            return ts * p
665

666
        @njit
667
        def apply_func_nb(ts, p):
668
            return ts * p
669

670
        target = pd.DataFrame(
671
            np.array([
672
                [0., 5., 2.],
673
                [0., 4., 4.],
674
                [0., 3., 6.],
675
                [0., 2., 4.],
676
                [0., 1., 2.]
677
            ]),
678
            index=ts.index,
679
            columns=pd.MultiIndex.from_tuples([
680
                ('array_0', 'a'),
681
                ('array_0', 'b'),
682
                ('array_0', 'c'),
683
            ], names=['custom_p', None])
684
        )
685
        pd.testing.assert_frame_equal(
686
            F.from_apply_func(apply_func).run(ts, np.asarray([0, 1, 2]), param_settings={'p': {
687
                'is_array_like': True
688
            }}).out,
689
            target
690
        )
691
        pd.testing.assert_frame_equal(
692
            F.from_apply_func(apply_func_nb, numba_loop=True).run(ts, np.asarray([0, 1, 2]), param_settings={'p': {
693
                'is_array_like': True
694
            }}).out,
695
            target
696
        )
697
        target = pd.DataFrame(
698
            np.array([
699
                [0., 5., 2.],
700
                [0., 4., 4.],
701
                [0., 3., 6.],
702
                [0., 2., 4.],
703
                [0., 1., 2.]
704
            ]),
705
            index=ts.index,
706
            columns=pd.MultiIndex.from_tuples([
707
                (0, 'a'),
708
                (1, 'b'),
709
                (2, 'c'),
710
            ], names=['custom_p', None])
711
        )
712
        pd.testing.assert_frame_equal(
713
            F.from_apply_func(apply_func).run(ts, np.asarray([0, 1, 2]), param_settings={'p': {
714
                'is_array_like': True,
715
                'bc_to_input': 1,
716
                'per_column': True
717
            }}).out,
718
            target
719
        )
720
        pd.testing.assert_frame_equal(
721
            F.from_apply_func(apply_func_nb, numba_loop=True).run(ts, np.asarray([0, 1, 2]), param_settings={'p': {
722
                'is_array_like': True,
723
                'bc_to_input': 1,
724
                'per_column': True
725
            }}).out,
726
            target
727
        )
728

729
        def apply_func2(ts, p):
730
            return ts * np.expand_dims(p, 1)
731

732
        @njit
733
        def apply_func2_nb(ts, p):
734
            return ts * np.expand_dims(p, 1)
735

736
        target = pd.DataFrame(
737
            np.array([
738
                [0., 0., 0.],
739
                [2., 4., 2.],
740
                [6., 6., 6.],
741
                [12., 6., 6.],
742
                [20., 4., 4.]
743
            ]),
744
            index=ts.index,
745
            columns=pd.MultiIndex.from_tuples([
746
                ('array_0', 'a'),
747
                ('array_0', 'b'),
748
                ('array_0', 'c'),
749
            ], names=['custom_p', None])
750
        )
751
        pd.testing.assert_frame_equal(
752
            F.from_apply_func(apply_func2).run(ts, np.asarray([0, 1, 2, 3, 4]), param_settings={'p': {
753
                'is_array_like': True,
754
                'bc_to_input': 0
755
            }}).out,
756
            target
757
        )
758
        pd.testing.assert_frame_equal(
759
            F.from_apply_func(apply_func2_nb).run(ts, np.asarray([0, 1, 2, 3, 4]), param_settings={'p': {
760
                'is_array_like': True,
761
                'bc_to_input': 0
762
            }}).out,
763
            target
764
        )
765

766
        def apply_func3(ts, p):
767
            return ts * (p[0] + p[1])
768

769
        @njit
770
        def apply_func3_nb(ts, p):
771
            return ts * (p[0] + p[1])
772

773
        target = pd.DataFrame(
774
            np.array([
775
                [1., 5., 1.],
776
                [2., 4., 2.],
777
                [3., 3., 3.],
778
                [4., 2., 2.],
779
                [5., 1., 1.]
780
            ]),
781
            index=ts.index,
782
            columns=pd.MultiIndex.from_tuples([
783
                ('tuple_0', 'a'),
784
                ('tuple_0', 'b'),
785
                ('tuple_0', 'c'),
786
            ], names=['custom_p', None])
787
        )
788
        pd.testing.assert_frame_equal(
789
            F.from_apply_func(apply_func3).run(ts, (0, 1), param_settings={'p': {
790
                'is_tuple': True
791
            }}).out,
792
            target
793
        )
794
        pd.testing.assert_frame_equal(
795
            F.from_apply_func(apply_func3_nb).run(ts, (0, 1), param_settings={'p': {
796
                'is_tuple': True
797
            }}).out,
798
            target
799
        )
800

801
    def test_param_product(self):
802
        F = vbt.IndicatorFactory(input_names=['ts'], param_names=['p1', 'p2'], output_names=['out'])
803

804
        def apply_func(ts, p1, p2):
805
            return ts * (p1 + p2)
806

807
        @njit
808
        def apply_func_nb(ts, p1, p2):
809
            return ts * (p1 + p2)
810

811
        target = pd.DataFrame(
812
            np.array([
813
                [2., 10., 2., 3., 15., 3., 3., 15., 3., 4., 20., 4.],
814
                [4., 8., 4., 6., 12., 6., 6., 12., 6., 8., 16., 8.],
815
                [6., 6., 6., 9., 9., 9., 9., 9., 9., 12., 12., 12.],
816
                [8., 4., 4., 12., 6., 6., 12., 6., 6., 16., 8., 8.],
817
                [10., 2., 2., 15., 3., 3., 15., 3., 3., 20., 4., 4.]
818
            ]),
819
            index=ts.index,
820
            columns=pd.MultiIndex.from_tuples([
821
                (0, 2, 'a'),
822
                (0, 2, 'b'),
823
                (0, 2, 'c'),
824
                (0, 3, 'a'),
825
                (0, 3, 'b'),
826
                (0, 3, 'c'),
827
                (1, 2, 'a'),
828
                (1, 2, 'b'),
829
                (1, 2, 'c'),
830
                (1, 3, 'a'),
831
                (1, 3, 'b'),
832
                (1, 3, 'c')
833
            ], names=['custom_p1', 'custom_p2', None])
834
        )
835
        pd.testing.assert_frame_equal(
836
            F.from_apply_func(apply_func).run(ts, [0, 1], [2, 3], param_product=True).out,
837
            target
838
        )
839
        pd.testing.assert_frame_equal(
840
            F.from_apply_func(apply_func_nb, numba_loop=True).run(ts, [0, 1], [2, 3], param_product=True).out,
841
            target
842
        )
843

844
    def test_default(self):
845
        F = vbt.IndicatorFactory(
846
            input_names=['ts1', 'ts2'],
847
            param_names=['p1', 'p2'],
848
            in_output_names=['in_out1', 'in_out2'],
849
            output_names=['out']
850
        )
851

852
        def apply_func(ts1, ts2, in_out1, in_out2, p1, p2):
853
            in_out1[::2] = ts1[::2] * ts2[::2] * (p1 + p2)
854
            in_out2[::2] = ts1[::2] * ts2[::2] * (p1 + p2)
855
            return ts1 * ts2 * (p1 + p2)
856

857
        # default inputs
858
        pd.testing.assert_frame_equal(
859
            F.from_apply_func(apply_func, ts2=0).run(ts, [1, 2], 3).out,
860
            F.from_apply_func(apply_func).run(ts, 0, [1, 2], 3).out
861
        )
862
        pd.testing.assert_frame_equal(
863
            F.from_apply_func(apply_func, ts2='ts1').run(ts, [1, 2], 3).out,
864
            F.from_apply_func(apply_func).run(ts, ts, [1, 2], 3).out
865
        )
866
        # default params
867
        pd.testing.assert_frame_equal(
868
            F.from_apply_func(apply_func, p2=0, hide_default=False)
869
                .run(ts, ts, [1, 2]).out,
870
            F.from_apply_func(apply_func, hide_default=False)
871
                .run(ts, ts, [1, 2], 0).out
872
        )
873
        pd.testing.assert_frame_equal(
874
            F.from_apply_func(apply_func, p2='p1', hide_default=False)
875
                .run(ts, ts, [1, 2]).out,
876
            F.from_apply_func(apply_func, hide_default=False)
877
                .run(ts, ts, [1, 2], [1, 2]).out
878
        )
879
        with pytest.raises(Exception):
880
            pd.testing.assert_frame_equal(
881
                F.from_apply_func(apply_func, in_out1=1, in_out2=2)
882
                    .run(ts, ts, [1, 2], 3).in_out1,
883
                F.from_apply_func(apply_func, in_out1=1, in_out2=2)
884
                    .run(ts, ts, [1, 2], 3).in_out2
885
            )
886
        pd.testing.assert_frame_equal(
887
            F.from_apply_func(apply_func, in_out1=1, in_out2='in_out1')
888
                .run(ts, ts, [1, 2], 3).in_out2,
889
            F.from_apply_func(apply_func, in_out1=1, in_out2=1)
890
                .run(ts, ts, [1, 2], 3).in_out2
891
        )
892
        pd.testing.assert_frame_equal(
893
            F.from_apply_func(apply_func, in_out1=1, in_out2='ts1')
894
                .run(ts, ts, [1, 2], 3).in_out2,
895
            F.from_apply_func(apply_func, in_out1=1, in_out2=ts)
896
                .run(ts, ts, [1, 2], 3).in_out2
897
        )
898

899
    def test_hide_params(self):
900
        F = vbt.IndicatorFactory(input_names=['ts'], param_names=['p1', 'p2'], output_names=['out'])
901

902
        assert F.from_apply_func(lambda ts, p1, p2: ts * p1 * p2, hide_params=[]) \
903
                   .run(ts, [0, 1], 2) \
904
                   .out.columns.names == ['custom_p1', 'custom_p2', None]
905
        assert F.from_apply_func(lambda ts, p1, p2: ts * p1 * p2, hide_params=['p2']) \
906
                   .run(ts, [0, 1], 2) \
907
                   .out.columns.names == ['custom_p1', None]
908

909
    def test_hide_default(self):
910
        F = vbt.IndicatorFactory(input_names=['ts'], param_names=['p1', 'p2'], output_names=['out'])
911

912
        assert F.from_apply_func(lambda ts, p1, p2: ts * p1 * p2, p2=2, hide_default=False) \
913
                   .run(ts, [0, 1]) \
914
                   .out.columns.names == ['custom_p1', 'custom_p2', None]
915
        assert F.from_apply_func(lambda ts, p1, p2: ts * p1 * p2, p2=2, hide_default=True) \
916
                   .run(ts, [0, 1]) \
917
                   .out.columns.names == ['custom_p1', None]
918

919
    def test_multiple_outputs(self):
920
        F = vbt.IndicatorFactory(input_names=['ts'], param_names=['p'], output_names=['o1', 'o2'])
921

922
        def apply_func(ts, p):
923
            return (ts * p, (ts * p) ** 2)
924

925
        @njit
926
        def apply_func_nb(ts, p):
927
            return (ts * p, (ts * p) ** 2)
928

929
        target = pd.DataFrame(
930
            np.array([
931
                [0., 0., 0., 1., 5., 1.],
932
                [0., 0., 0., 2., 4., 2.],
933
                [0., 0., 0., 3., 3., 3.],
934
                [0., 0., 0., 4., 2., 2.],
935
                [0., 0., 0., 5., 1., 1.]
936
            ]),
937
            index=ts.index,
938
            columns=pd.MultiIndex.from_tuples([
939
                (0, 'a'),
940
                (0, 'b'),
941
                (0, 'c'),
942
                (1, 'a'),
943
                (1, 'b'),
944
                (1, 'c')
945
            ], names=['custom_p', None])
946
        )
947
        pd.testing.assert_frame_equal(
948
            F.from_apply_func(apply_func).run(ts, [0, 1]).o1,
949
            target
950
        )
951
        pd.testing.assert_frame_equal(
952
            F.from_apply_func(apply_func_nb, numba_loop=True).run(ts, [0, 1]).o1,
953
            target
954
        )
955
        pd.testing.assert_frame_equal(
956
            F.from_apply_func(apply_func).run(ts.vbt.tile(2), [0, 0, 0, 1, 1, 1], per_column=True).o1,
957
            target
958
        )
959
        pd.testing.assert_frame_equal(
960
            F.from_apply_func(apply_func_nb, numba_loop=True)
961
                .run(ts.vbt.tile(2), [0, 0, 0, 1, 1, 1], per_column=True).o1,
962
            target
963
        )
964
        target = pd.DataFrame(
965
            np.array([
966
                [0., 0., 0., 1., 25., 1.],
967
                [0., 0., 0., 4., 16., 4.],
968
                [0., 0., 0., 9., 9., 9.],
969
                [0., 0., 0., 16., 4., 4.],
970
                [0., 0., 0., 25., 1., 1.]
971
            ]),
972
            index=target.index,
973
            columns=target.columns
974
        )
975
        pd.testing.assert_frame_equal(
976
            F.from_apply_func(apply_func).run(ts, [0, 1]).o2,
977
            target
978
        )
979
        pd.testing.assert_frame_equal(
980
            F.from_apply_func(apply_func_nb, numba_loop=True).run(ts, [0, 1]).o2,
981
            target
982
        )
983
        pd.testing.assert_frame_equal(
984
            F.from_apply_func(apply_func).run(ts.vbt.tile(2), [0, 0, 0, 1, 1, 1], per_column=True).o2,
985
            target
986
        )
987
        pd.testing.assert_frame_equal(
988
            F.from_apply_func(apply_func_nb, numba_loop=True)
989
                .run(ts.vbt.tile(2), [0, 0, 0, 1, 1, 1], per_column=True).o2,
990
            target
991
        )
992

993
    def test_in_outputs(self):
994
        F = vbt.IndicatorFactory(
995
            input_names=['ts'], param_names=['p'],
996
            output_names=['out'], in_output_names=['in_out']
997
        )
998

999
        def apply_func(ts, in_out, p):
1000
            in_out[:, 0] = p
1001
            return ts * p
1002

1003
        @njit
1004
        def apply_func_nb(ts, in_out, p):
1005
            in_out[:, 0] = p
1006
            return ts * p
1007

1008
        target = pd.DataFrame(
1009
            np.array([
1010
                [0, -1, -1, 1, -1, -1],
1011
                [0, -1, -1, 1, -1, -1],
1012
                [0, -1, -1, 1, -1, -1],
1013
                [0, -1, -1, 1, -1, -1],
1014
                [0, -1, -1, 1, -1, -1]
1015
            ]),
1016
            index=ts.index,
1017
            columns=pd.MultiIndex.from_tuples([
1018
                (0, 'a'),
1019
                (0, 'b'),
1020
                (0, 'c'),
1021
                (1, 'a'),
1022
                (1, 'b'),
1023
                (1, 'c')
1024
            ], names=['custom_p', None])
1025
        )
1026
        assert F.from_apply_func(apply_func).run(ts, [0, 1])._in_out.dtype == np.float64
1027
        assert F.from_apply_func(apply_func, in_output_settings={'in_out': {'dtype': np.int64}}) \
1028
                   .run(ts, [0, 1])._in_out.dtype == np.int64
1029
        pd.testing.assert_frame_equal(
1030
            F.from_apply_func(apply_func, in_out=-1).run(ts, [0, 1]).in_out,
1031
            target
1032
        )
1033
        pd.testing.assert_frame_equal(
1034
            F.from_apply_func(apply_func_nb, numba_loop=True, in_out=-1).run(ts, [0, 1]).in_out,
1035
            target
1036
        )
1037
        pd.testing.assert_frame_equal(
1038
            F.from_apply_func(apply_func).run(ts, [0, 1], in_out=-1).in_out,
1039
            target
1040
        )
1041
        pd.testing.assert_frame_equal(
1042
            F.from_apply_func(apply_func_nb, numba_loop=True).run(ts, [0, 1], in_out=-1).in_out,
1043
            target
1044
        )
1045
        pd.testing.assert_frame_equal(
1046
            F.from_apply_func(apply_func).run(ts, [0, 1], in_out=np.full(ts.shape, -1, dtype=int)).in_out,
1047
            target
1048
        )
1049
        pd.testing.assert_frame_equal(
1050
            F.from_apply_func(apply_func_nb, numba_loop=True)
1051
                .run(ts, [0, 1], in_out=np.full(ts.shape, -1, dtype=int)).in_out,
1052
            target
1053
        )
1054
        target = pd.DataFrame(
1055
            np.array([
1056
                [0, 1, 2],
1057
                [0, 1, 2],
1058
                [0, 1, 2],
1059
                [0, 1, 2],
1060
                [0, 1, 2]
1061
            ]),
1062
            index=ts.index,
1063
            columns=pd.MultiIndex.from_tuples([
1064
                (0, 'a'),
1065
                (1, 'b'),
1066
                (2, 'c')
1067
            ], names=['custom_p', None])
1068
        )
1069
        pd.testing.assert_frame_equal(
1070
            F.from_apply_func(apply_func).run(ts, [0, 1, 2], in_out=-1, per_column=True).in_out,
1071
            target
1072
        )
1073
        pd.testing.assert_frame_equal(
1074
            F.from_apply_func(apply_func_nb, numba_loop=True).run(ts, [0, 1, 2], in_out=-1, per_column=True).in_out,
1075
            target
1076
        )
1077
        pd.testing.assert_frame_equal(
1078
            F.from_apply_func(apply_func)
1079
                .run(ts, [0, 1, 2], in_out=np.full(ts.shape, -1, dtype=int), per_column=True).in_out,
1080
            target
1081
        )
1082
        pd.testing.assert_frame_equal(
1083
            F.from_apply_func(apply_func_nb, numba_loop=True)
1084
                .run(ts, [0, 1, 2], in_out=np.full(ts.shape, -1, dtype=int), per_column=True).in_out,
1085
            target
1086
        )
1087

1088
    def test_no_outputs(self):
1089
        F = vbt.IndicatorFactory(
1090
            param_names=['p'], in_output_names=['in_out']
1091
        )
1092

1093
        def apply_func(in_out, p):
1094
            in_out[:] = p
1095

1096
        @njit
1097
        def apply_func_nb(in_out, p):
1098
            in_out[:] = p
1099

1100
        target = pd.DataFrame(
1101
            np.array([
1102
                [0, 0, 0, 1, 1, 1],
1103
                [0, 0, 0, 1, 1, 1],
1104
                [0, 0, 0, 1, 1, 1],
1105
                [0, 0, 0, 1, 1, 1],
1106
                [0, 0, 0, 1, 1, 1]
1107
            ]),
1108
            index=ts.index,
1109
            columns=pd.MultiIndex.from_tuples([
1110
                (0, 'a'),
1111
                (0, 'b'),
1112
                (0, 'c'),
1113
                (1, 'a'),
1114
                (1, 'b'),
1115
                (1, 'c')
1116
            ], names=['custom_p', None])
1117
        )
1118
        pd.testing.assert_frame_equal(
1119
            F.from_apply_func(apply_func, in_output_settings=dict(in_out=dict(dtype=np.int64)))
1120
                .run([0, 1], input_shape=ts.shape, input_index=ts.index, input_columns=ts.columns).in_out,
1121
            target
1122
        )
1123
        pd.testing.assert_frame_equal(
1124
            F.from_apply_func(apply_func_nb, numba_loop=True, in_output_settings=dict(in_out=dict(dtype=np.int64)))
1125
                .run([0, 1], input_shape=ts.shape, input_index=ts.index, input_columns=ts.columns).in_out,
1126
            target
1127
        )
1128

1129
    def test_kwargs_to_args(self):
1130
        F = vbt.IndicatorFactory(input_names=['ts'], output_names=['out'])
1131

1132
        def apply_func(ts, kw):
1133
            return ts * kw
1134

1135
        @njit
1136
        def apply_func_nb(ts, kw):
1137
            return ts * kw
1138

1139
        pd.testing.assert_frame_equal(
1140
            F.from_apply_func(apply_func, kwargs_to_args=['kw']).run(ts, kw=2).out,
1141
            ts * 2
1142
        )
1143
        pd.testing.assert_frame_equal(
1144
            F.from_apply_func(apply_func_nb, numba_loop=True, kwargs_to_args=['kw']).run(ts, kw=2).out,
1145
            ts * 2
1146
        )
1147

1148
    def test_cache(self):
1149
        F = vbt.IndicatorFactory(input_names=['ts'], param_names=['p'], output_names=['out'])
1150

1151
        def cache_func(ts, ps):
1152
            np.random.seed(seed)
1153
            return np.random.uniform(0, 1)
1154

1155
        @njit
1156
        def cache_func_nb(ts, ps):
1157
            np.random.seed(seed)
1158
            return np.random.uniform(0, 1)
1159

1160
        def apply_func(ts, p, c):
1161
            return ts * p + c
1162

1163
        @njit
1164
        def apply_func_nb(ts, p, c):
1165
            return ts * p + c
1166

1167
        target = pd.DataFrame(
1168
            np.array([
1169
                [0.37454012, 0.37454012, 0.37454012, 1.37454012, 5.37454012, 1.37454012],
1170
                [0.37454012, 0.37454012, 0.37454012, 2.37454012, 4.37454012, 2.37454012],
1171
                [0.37454012, 0.37454012, 0.37454012, 3.37454012, 3.37454012, 3.37454012],
1172
                [0.37454012, 0.37454012, 0.37454012, 4.37454012, 2.37454012, 2.37454012],
1173
                [0.37454012, 0.37454012, 0.37454012, 5.37454012, 1.37454012, 1.37454012]
1174
            ]),
1175
            index=ts.index,
1176
            columns=pd.MultiIndex.from_tuples([
1177
                (0, 'a'),
1178
                (0, 'b'),
1179
                (0, 'c'),
1180
                (1, 'a'),
1181
                (1, 'b'),
1182
                (1, 'c')
1183
            ], names=['custom_p', None])
1184
        )
1185
        pd.testing.assert_frame_equal(
1186
            F.from_apply_func(
1187
                apply_func,
1188
                cache_func=cache_func
1189
            ).run(ts, [0, 1]).out,
1190
            target
1191
        )
1192
        pd.testing.assert_frame_equal(
1193
            F.from_apply_func(
1194
                apply_func_nb,
1195
                cache_func=cache_func_nb
1196
            ).run(ts, [0, 1]).out,
1197
            target
1198
        )
1199
        # return_cache
1200
        cache = F.from_apply_func(
1201
            apply_func,
1202
            cache_func=cache_func
1203
        ).run(ts, [0, 1], return_cache=True)
1204
        assert cache == 0.3745401188473625
1205
        cache = F.from_apply_func(
1206
            apply_func_nb,
1207
            cache_func=cache_func_nb
1208
        ).run(ts, [0, 1], return_cache=True)
1209
        assert cache == 0.3745401188473625
1210
        # use_cache
1211
        pd.testing.assert_frame_equal(
1212
            F.from_apply_func(
1213
                apply_func
1214
            ).run(ts, [0, 1], use_cache=cache).out,
1215
            target
1216
        )
1217
        pd.testing.assert_frame_equal(
1218
            F.from_apply_func(
1219
                apply_func_nb
1220
            ).run(ts, [0, 1], use_cache=cache).out,
1221
            target
1222
        )
1223

1224
        # per_column
1225
        def cache_func(col, ts, ps):
1226
            np.random.seed(seed + col)
1227
            return np.random.uniform(0, 1)
1228

1229
        @njit
1230
        def cache_func_nb(col, ts, ps):
1231
            np.random.seed(seed + col)
1232
            return np.random.uniform(0, 1)
1233

1234
        cache = F.from_apply_func(
1235
            apply_func,
1236
            cache_func=cache_func,
1237
            pass_col=True
1238
        ).run(ts, [0, 1, 2], return_cache=True, per_column=True)
1239
        assert cache == [0.3745401188473625, 0.11505456638977896, 0.8348421486656494]
1240
        cache = F.from_apply_func(
1241
            apply_func_nb,
1242
            cache_func=cache_func_nb,
1243
            pass_col=True
1244
        ).run(ts, [0, 1, 2], return_cache=True, per_column=True)
1245
        assert cache == [0.3745401188473625, 0.11505456638977896, 0.8348421486656494]
1246
        target = pd.DataFrame(
1247
            np.array([
1248
                [0.37454012, 5.115054566389779, 2.8348421486656497],
1249
                [0.37454012, 4.115054566389779, 4.8348421486656497],
1250
                [0.37454012, 3.115054566389779, 6.8348421486656497],
1251
                [0.37454012, 2.115054566389779, 4.8348421486656497],
1252
                [0.37454012, 1.115054566389779, 2.8348421486656497]
1253
            ]),
1254
            index=ts.index,
1255
            columns=pd.MultiIndex.from_tuples([
1256
                (0, 'a'),
1257
                (1, 'b'),
1258
                (2, 'c')
1259
            ], names=['custom_p', None])
1260
        )
1261
        pd.testing.assert_frame_equal(
1262
            F.from_apply_func(
1263
                apply_func
1264
            ).run(ts, [0, 1, 2], use_cache=cache, per_column=True).out,
1265
            target
1266
        )
1267
        pd.testing.assert_frame_equal(
1268
            F.from_apply_func(
1269
                apply_func_nb
1270
            ).run(ts, [0, 1, 2], use_cache=cache, per_column=True).out,
1271
            target
1272
        )
1273

1274
    def test_raw(self):
1275
        F = vbt.IndicatorFactory(input_names=['ts'], param_names=['p'], output_names=['out'])
1276

1277
        def apply_func(ts, p, a, b=10):
1278
            return ts * p + a + b
1279

1280
        @njit
1281
        def apply_func_nb(ts, p, a, b):
1282
            return ts * p + a + b
1283

1284
        target = np.array([
1285
            [110., 110., 110., 111., 115., 111.],
1286
            [110., 110., 110., 112., 114., 112.],
1287
            [110., 110., 110., 113., 113., 113.],
1288
            [110., 110., 110., 114., 112., 112.],
1289
            [110., 110., 110., 115., 111., 111.]
1290
        ])
1291
        np.testing.assert_array_equal(
1292
            F.from_apply_func(
1293
                apply_func, var_args=True
1294
            ).run(ts, [0, 1], 10, b=100, return_raw=True)[0][0],
1295
            target
1296
        )
1297
        np.testing.assert_array_equal(
1298
            F.from_apply_func(
1299
                apply_func_nb, var_args=True
1300
            ).run(ts, [0, 1], 10, 100, return_raw=True)[0][0],
1301
            target
1302
        )
1303
        np.testing.assert_array_equal(
1304
            F.from_apply_func(
1305
                apply_func, var_args=True
1306
            ).run(ts, [0, 1], 10, b=100, return_raw=True)[1],
1307
            [(0,), (1,)]
1308
        )
1309
        np.testing.assert_array_equal(
1310
            F.from_apply_func(
1311
                apply_func_nb, var_args=True
1312
            ).run(ts, [0, 1], 10, 100, return_raw=True)[1],
1313
            [(0,), (1,)]
1314
        )
1315
        assert F.from_apply_func(
1316
            apply_func, var_args=True
1317
        ).run(ts, [0, 1], 10, b=100, return_raw=True)[2] == 3
1318
        assert F.from_apply_func(
1319
            apply_func_nb, var_args=True
1320
        ).run(ts, [0, 1], 10, 100, return_raw=True)[2] == 3
1321
        assert F.from_apply_func(
1322
            apply_func, var_args=True
1323
        ).run(ts, [0, 1], 10, b=100, return_raw=True)[3] == []
1324
        assert F.from_apply_func(
1325
            apply_func_nb, var_args=True
1326
        ).run(ts, [0, 1], 10, 100, return_raw=True)[3] == []
1327
        raw_results = F.from_apply_func(
1328
            apply_func, var_args=True
1329
        ).run(ts, [0, 1, 2], 10, b=100, return_raw=True)
1330
        pd.testing.assert_frame_equal(
1331
            F.from_apply_func(
1332
                apply_func, var_args=True
1333
            ).run(ts, [0, 1], 10, b=100, use_raw=raw_results).out,
1334
            F.from_apply_func(
1335
                apply_func_nb, var_args=True
1336
            ).run(ts, [0, 1], 10, 100).out
1337
        )
1338

1339
        # per_column
1340
        target = np.array([
1341
            [110., 115., 112.],
1342
            [110., 114., 114.],
1343
            [110., 113., 116.],
1344
            [110., 112., 114.],
1345
            [110., 111., 112.]
1346
        ])
1347
        np.testing.assert_array_equal(
1348
            F.from_apply_func(
1349
                apply_func, var_args=True
1350
            ).run(ts, [0, 1, 2], 10, b=100, return_raw=True, per_column=True)[0][0],
1351
            target
1352
        )
1353
        np.testing.assert_array_equal(
1354
            F.from_apply_func(
1355
                apply_func_nb, var_args=True
1356
            ).run(ts, [0, 1, 2], 10, 100, return_raw=True, per_column=True)[0][0],
1357
            target
1358
        )
1359
        np.testing.assert_array_equal(
1360
            F.from_apply_func(
1361
                apply_func, var_args=True
1362
            ).run(ts, [0, 1, 2], 10, b=100, return_raw=True, per_column=True)[1],
1363
            [(0,), (1,), (2,)]
1364
        )
1365
        np.testing.assert_array_equal(
1366
            F.from_apply_func(
1367
                apply_func_nb, var_args=True
1368
            ).run(ts, [0, 1, 2], 10, 100, return_raw=True, per_column=True)[1],
1369
            [(0,), (1,), (2,)]
1370
        )
1371
        assert F.from_apply_func(
1372
            apply_func, var_args=True
1373
        ).run(ts, [0, 1, 2], 10, b=100, return_raw=True)[2] == 3
1374
        assert F.from_apply_func(
1375
            apply_func_nb, var_args=True
1376
        ).run(ts, [0, 1, 2], 10, 100, return_raw=True)[2] == 3
1377
        assert F.from_apply_func(
1378
            apply_func, var_args=True
1379
        ).run(ts, [0, 1, 2], 10, b=100, return_raw=True)[3] == []
1380
        assert F.from_apply_func(
1381
            apply_func_nb, var_args=True
1382
        ).run(ts, [0, 1, 2], 10, 100, return_raw=True)[3] == []
1383
        raw_results = F.from_apply_func(
1384
            apply_func, var_args=True
1385
        ).run(ts, [0, 1, 2], 10, b=100, return_raw=True)
1386
        pd.testing.assert_frame_equal(
1387
            F.from_apply_func(
1388
                apply_func, var_args=True
1389
            ).run(ts, [0, 0, 0], 10, b=100, use_raw=raw_results).out,
1390
            F.from_apply_func(
1391
                apply_func_nb, var_args=True
1392
            ).run(ts, [0, 0, 0], 10, 100).out
1393
        )
1394

1395
    @pytest.mark.parametrize(
1396
        "test_to_2d,test_keep_pd",
1397
        [
1398
            (False, False),
1399
            (False, True),
1400
            (True, False),
1401
            (True, True)
1402
        ]
1403
    )
1404
    def test_to_2d_and_keep_pd(self, test_to_2d, test_keep_pd):
1405
        F = vbt.IndicatorFactory(input_names=['ts'], in_output_names=['in_out'], output_names=['out'])
1406

1407
        def custom_func(_ts, _in_out):
1408
            if test_to_2d:
1409
                assert _ts.ndim == 2
1410
                for __in_out in _in_out:
1411
                    assert __in_out.ndim == 2
1412
                if test_keep_pd:
1413
                    pd.testing.assert_frame_equal(_ts, ts[['a']].vbt.wrapper.wrap(_ts.values))
1414
                    for __in_out in _in_out:
1415
                        pd.testing.assert_frame_equal(__in_out, ts[['a']].vbt.wrapper.wrap(__in_out.values))
1416
            else:
1417
                assert _ts.ndim == 1
1418
                for __in_out in _in_out:
1419
                    assert __in_out.ndim == 1
1420
                if test_keep_pd:
1421
                    pd.testing.assert_series_equal(_ts, ts['a'].vbt.wrapper.wrap(_ts.values))
1422
                    for __in_out in _in_out:
1423
                        pd.testing.assert_series_equal(__in_out, ts['a'].vbt.wrapper.wrap(__in_out.values))
1424
            return _ts
1425

1426
        def apply_func(_ts, _in_out):
1427
            if test_to_2d:
1428
                assert _ts.ndim == 2
1429
                assert _in_out.ndim == 2
1430
                if test_keep_pd:
1431
                    pd.testing.assert_frame_equal(_ts, ts[['a']].vbt.wrapper.wrap(_ts.values))
1432
                    pd.testing.assert_frame_equal(_in_out, ts[['a']].vbt.wrapper.wrap(_in_out.values))
1433
            else:
1434
                assert _ts.ndim == 1
1435
                assert _in_out.ndim == 1
1436
                if test_keep_pd:
1437
                    pd.testing.assert_series_equal(_ts, ts['a'].vbt.wrapper.wrap(_ts.values))
1438
                    pd.testing.assert_series_equal(_in_out, ts['a'].vbt.wrapper.wrap(_in_out.values))
1439
            return _ts
1440

1441
        _ = F.from_custom_func(custom_func, to_2d=test_to_2d, keep_pd=test_keep_pd, var_args=True) \
1442
            .run(ts['a'])
1443
        _ = F.from_apply_func(apply_func, to_2d=test_to_2d, keep_pd=test_keep_pd, var_args=True) \
1444
            .run(ts['a'])
1445

1446
        def custom_func(_ts, _in_out, col=None):
1447
            if test_to_2d:
1448
                assert _ts.ndim == 2
1449
                for __in_out in _in_out:
1450
                    assert __in_out.ndim == 2
1451
                if test_keep_pd:
1452
                    pd.testing.assert_frame_equal(_ts, ts.iloc[:, [col]].vbt.wrapper.wrap(_ts.values))
1453
                    for __in_out in _in_out:
1454
                        pd.testing.assert_frame_equal(__in_out, ts.iloc[:, [col]].vbt.wrapper.wrap(__in_out.values))
1455
            else:
1456
                assert _ts.ndim == 1
1457
                for __in_out in _in_out:
1458
                    assert __in_out.ndim == 1
1459
                if test_keep_pd:
1460
                    pd.testing.assert_series_equal(_ts, ts.iloc[:, col].vbt.wrapper.wrap(_ts.values))
1461
                    for __in_out in _in_out:
1462
                        pd.testing.assert_series_equal(__in_out, ts.iloc[:, col].vbt.wrapper.wrap(__in_out.values))
1463
            return _ts
1464

1465
        def apply_func(col, _ts, _in_out):
1466
            if test_to_2d:
1467
                assert _ts.ndim == 2
1468
                assert _in_out.ndim == 2
1469
                if test_keep_pd:
1470
                    pd.testing.assert_frame_equal(_ts, ts.iloc[:, [col]].vbt.wrapper.wrap(_ts.values))
1471
                    pd.testing.assert_frame_equal(_in_out, ts.iloc[:, [col]].vbt.wrapper.wrap(_in_out.values))
1472
            else:
1473
                assert _ts.ndim == 1
1474
                assert _in_out.ndim == 1
1475
                if test_keep_pd:
1476
                    pd.testing.assert_series_equal(_ts, ts.iloc[:, col].vbt.wrapper.wrap(_ts.values))
1477
                    pd.testing.assert_series_equal(_in_out, ts.iloc[:, col].vbt.wrapper.wrap(_in_out.values))
1478
            return _ts
1479

1480
        _ = F.from_custom_func(custom_func, to_2d=test_to_2d, keep_pd=test_keep_pd, var_args=True) \
1481
            .run(ts['a'], per_column=True, pass_col=True)
1482
        _ = F.from_apply_func(apply_func, to_2d=test_to_2d, keep_pd=test_keep_pd, var_args=True) \
1483
            .run(ts['a'], per_column=True, pass_col=True)
1484

1485
    def test_as_lists(self):
1486
        F = vbt.IndicatorFactory(input_names=['ts'], param_names=['p'], output_names=['out'])
1487

1488
        def custom_func(input_list, in_output_list, param_list):
1489
            return input_list[0] * param_list[0][0]
1490

1491
        @njit
1492
        def custom_func_nb(input_list, in_output_list, param_list):
1493
            return input_list[0] * param_list[0][0]
1494

1495
        target = pd.DataFrame(
1496
            ts.values * 2,
1497
            index=ts.index,
1498
            columns=pd.MultiIndex.from_tuples([
1499
                (2, 'a'),
1500
                (2, 'b'),
1501
                (2, 'c')
1502
            ], names=['custom_p', None])
1503
        )
1504
        pd.testing.assert_frame_equal(
1505
            F.from_custom_func(custom_func, as_lists=True).run(ts, 2).out,
1506
            target
1507
        )
1508
        pd.testing.assert_frame_equal(
1509
            F.from_custom_func(custom_func_nb, as_lists=True).run(ts, 2).out,
1510
            target
1511
        )
1512
        pd.testing.assert_frame_equal(
1513
            F.from_custom_func(custom_func, as_lists=True).run(ts, 2, per_column=True).out,
1514
            target
1515
        )
1516
        pd.testing.assert_frame_equal(
1517
            F.from_custom_func(custom_func_nb, as_lists=True).run(ts, 2, per_column=True).out,
1518
            target
1519
        )
1520

1521
    def test_other(self):
1522
        F = vbt.IndicatorFactory(input_names=['ts'], output_names=['o1', 'o2'])
1523

1524
        def custom_func(ts):
1525
            return ts, ts + 1, ts + 2
1526

1527
        @njit
1528
        def custom_func_nb(ts):
1529
            return ts, ts + 1, ts + 2
1530

1531
        obj, other = F.from_custom_func(custom_func).run(ts)
1532
        np.testing.assert_array_equal(other, ts + 2)
1533
        obj, other = F.from_custom_func(custom_func_nb).run(ts)
1534
        np.testing.assert_array_equal(other, ts + 2)
1535
        obj, *others = F.from_custom_func(custom_func).run(ts, per_column=True)
1536
        for i, other in enumerate(others):
1537
            np.testing.assert_array_equal(other[0], ts.iloc[:, [i]] + 2)
1538
        obj, *others = F.from_custom_func(custom_func_nb).run(ts, per_column=True)
1539
        for i, other in enumerate(others):
1540
            np.testing.assert_array_equal(other[0], ts.iloc[:, [i]] + 2)
1541

1542
    def test_run_unique(self):
1543
        F = vbt.IndicatorFactory(input_names=['ts'], param_names=['p1', 'p2'], output_names=['out'])
1544

1545
        def apply_func(ts, p1, p2):
1546
            return ts * (p1 + p2)
1547

1548
        pd.testing.assert_series_equal(
1549
            F.from_apply_func(apply_func).run(ts['a'], 2, 3, run_unique=True).out,
1550
            F.from_apply_func(apply_func).run(ts['a'], 2, 3, run_unique=False).out
1551
        )
1552
        raw = F.from_apply_func(apply_func).run(ts['a'], [2, 2, 2], [3, 3, 3], run_unique=True, return_raw=True)
1553
        np.testing.assert_array_equal(
1554
            raw[0][0],
1555
            np.array([[5.], [10.], [15.], [20.], [25.]])
1556
        )
1557
        assert raw[1] == [(2, 3)]
1558
        assert raw[2] == 1
1559
        assert raw[3] == []
1560
        pd.testing.assert_frame_equal(
1561
            F.from_apply_func(apply_func).run(ts['a'], [2, 2, 2], [3, 3, 3], run_unique=True).out,
1562
            F.from_apply_func(apply_func).run(ts['a'], [2, 2, 2], [3, 3, 3], run_unique=False).out
1563
        )
1564
        pd.testing.assert_frame_equal(
1565
            F.from_apply_func(apply_func).run(ts, 2, 3, run_unique=True).out,
1566
            F.from_apply_func(apply_func).run(ts, 2, 3, run_unique=False).out
1567
        )
1568
        pd.testing.assert_frame_equal(
1569
            F.from_apply_func(apply_func).run(ts, [2, 2, 2], [3, 3, 3], run_unique=True).out,
1570
            F.from_apply_func(apply_func).run(ts, [2, 2, 2], [3, 3, 3], run_unique=False).out
1571
        )
1572
        pd.testing.assert_frame_equal(
1573
            F.from_apply_func(apply_func).run(ts, [2, 3, 4], [4, 3, 2], run_unique=True).out,
1574
            F.from_apply_func(apply_func).run(ts, [2, 3, 4], [4, 3, 2], run_unique=False).out
1575
        )
1576

1577
    def test_run_combs(self):
1578
        # itertools.combinations
1579
        F = vbt.IndicatorFactory(input_names=['ts'], param_names=['p1', 'p2'], output_names=['out'])
1580

1581
        ind1 = F.from_apply_func(lambda ts, p1, p2: ts * p1 * p2) \
1582
            .run(ts, [2, 2, 3], [10, 10, 11], short_name='custom_1')
1583
        ind2 = F.from_apply_func(lambda ts, p1, p2: ts * p1 * p2) \
1584
            .run(ts, [3, 4, 4], [11, 12, 12], short_name='custom_2')
1585
        ind1_1, ind2_1 = F.from_apply_func(lambda ts, p1, p2: ts * p1 * p2) \
1586
            .run_combs(ts, [2, 3, 4], [10, 11, 12], r=2, run_unique=False)
1587
        ind1_2, ind2_2 = F.from_apply_func(lambda ts, p1, p2: ts * p1 * p2) \
1588
            .run_combs(ts, [2, 3, 4], [10, 11, 12], r=2, run_unique=True)
1589
        pd.testing.assert_frame_equal(
1590
            ind1.out,
1591
            ind1_1.out
1592
        )
1593
        pd.testing.assert_frame_equal(
1594
            ind2.out,
1595
            ind2_1.out
1596
        )
1597
        pd.testing.assert_frame_equal(
1598
            ind1.out,
1599
            ind1_2.out
1600
        )
1601
        pd.testing.assert_frame_equal(
1602
            ind2.out,
1603
            ind2_2.out
1604
        )
1605
        # itertools.product
1606
        ind3 = F.from_apply_func(lambda ts, p1, p2: ts * p1 * p2) \
1607
            .run(ts, [2, 2, 2, 3, 3, 3, 4, 4, 4], [10, 10, 10, 11, 11, 11, 12, 12, 12], short_name='custom_1')
1608
        ind4 = F.from_apply_func(lambda ts, p1, p2: ts * p1 * p2) \
1609
            .run(ts, [2, 3, 4, 2, 3, 4, 2, 3, 4], [10, 11, 12, 10, 11, 12, 10, 11, 12], short_name='custom_2')
1610
        ind3_1, ind4_1 = F.from_apply_func(lambda ts, p1, p2: ts * p1 * p2) \
1611
            .run_combs(ts, [2, 3, 4], [10, 11, 12], r=2, comb_func=product, run_unique=False)
1612
        ind3_2, ind4_2 = F.from_apply_func(lambda ts, p1, p2: ts * p1 * p2) \
1613
            .run_combs(ts, [2, 3, 4], [10, 11, 12], r=2, comb_func=product, run_unique=True)
1614
        pd.testing.assert_frame_equal(
1615
            ind3.out,
1616
            ind3_1.out
1617
        )
1618
        pd.testing.assert_frame_equal(
1619
            ind4.out,
1620
            ind4_1.out
1621
        )
1622
        pd.testing.assert_frame_equal(
1623
            ind3.out,
1624
            ind3_2.out
1625
        )
1626
        pd.testing.assert_frame_equal(
1627
            ind4.out,
1628
            ind4_2.out
1629
        )
1630

1631
    def test_wrapper(self):
1632
        F = vbt.IndicatorFactory(input_names=['ts'], param_names=['p1', 'p2'], output_names=['out'])
1633

1634
        obj = F.from_apply_func(lambda ts, p1, p2: ts * p1 * p2).run(ts['a'], 0, 1)
1635
        assert obj.wrapper.ndim == 1
1636
        pd.testing.assert_index_equal(obj.wrapper.index, ts.index)
1637
        pd.testing.assert_index_equal(
1638
            obj.wrapper.columns,
1639
            pd.MultiIndex.from_tuples([
1640
                (0, 1, 'a')
1641
            ], names=['custom_p1', 'custom_p2', None])
1642
        )
1643
        obj = F.from_apply_func(lambda ts, p1, p2: ts * p1 * p2).run(ts['a'], [0, 1], 2)
1644
        assert obj.wrapper.ndim == 2
1645
        pd.testing.assert_index_equal(obj.wrapper.index, ts.index)
1646
        pd.testing.assert_index_equal(
1647
            obj.wrapper.columns,
1648
            pd.MultiIndex.from_tuples([
1649
                (0, 2),
1650
                (1, 2),
1651
            ], names=['custom_p1', 'custom_p2'])
1652
        )
1653
        obj = F.from_apply_func(lambda ts, p1, p2: ts * p1 * p2).run(ts, 0, 1)
1654
        assert obj.wrapper.ndim == 2
1655
        pd.testing.assert_index_equal(obj.wrapper.index, ts.index)
1656
        pd.testing.assert_index_equal(
1657
            obj.wrapper.columns,
1658
            pd.MultiIndex.from_tuples([
1659
                (0, 1, 'a'),
1660
                (0, 1, 'b'),
1661
                (0, 1, 'c')
1662
            ], names=['custom_p1', 'custom_p2', None])
1663
        )
1664
        obj = F.from_apply_func(lambda ts, p1, p2: ts * p1 * p2).run(ts, [1, 2], 3)
1665
        assert obj.wrapper.ndim == 2
1666
        pd.testing.assert_index_equal(obj.wrapper.index, ts.index)
1667
        pd.testing.assert_index_equal(
1668
            obj.wrapper.columns,
1669
            pd.MultiIndex.from_tuples([
1670
                (1, 3, 'a'),
1671
                (1, 3, 'b'),
1672
                (1, 3, 'c'),
1673
                (2, 3, 'a'),
1674
                (2, 3, 'b'),
1675
                (2, 3, 'c')
1676
            ], names=['custom_p1', 'custom_p2', None])
1677
        )
1678
        obj = F.from_apply_func(lambda ts, p1, p2: ts * p1 * p2).run(ts['a'], 0, 1, per_column=True)
1679
        assert obj.wrapper.ndim == 1
1680
        pd.testing.assert_index_equal(obj.wrapper.index, ts.index)
1681
        pd.testing.assert_index_equal(
1682
            obj.wrapper.columns,
1683
            pd.MultiIndex.from_tuples([
1684
                (0, 1, 'a')
1685
            ], names=['custom_p1', 'custom_p2', None])
1686
        )
1687
        obj = F.from_apply_func(lambda ts, p1, p2: ts * p1 * p2).run(ts[['a']], 0, 1, per_column=True)
1688
        assert obj.wrapper.ndim == 2
1689
        pd.testing.assert_index_equal(obj.wrapper.index, ts.index)
1690
        pd.testing.assert_index_equal(
1691
            obj.wrapper.columns,
1692
            pd.MultiIndex.from_tuples([
1693
                (0, 1, 'a')
1694
            ], names=['custom_p1', 'custom_p2', None])
1695
        )
1696
        obj = F.from_apply_func(lambda ts, p1, p2: ts * p1 * p2).run(ts, 0, 1, per_column=True)
1697
        assert obj.wrapper.ndim == 2
1698
        pd.testing.assert_index_equal(obj.wrapper.index, ts.index)
1699
        pd.testing.assert_index_equal(
1700
            obj.wrapper.columns,
1701
            pd.MultiIndex.from_tuples([
1702
                (0, 1, 'a'),
1703
                (0, 1, 'b'),
1704
                (0, 1, 'c')
1705
            ], names=['custom_p1', 'custom_p2', None])
1706
        )
1707

1708
    @pytest.mark.parametrize(
1709
        "test_config",
1710
        [
1711
            lambda F, shape, *args, **kwargs: F.from_apply_func(
1712
                lambda input_shape, p1, p2: np.empty(input_shape) * p1 * p2, require_input_shape=True)
1713
                .run(shape, *args, **kwargs),
1714
            lambda F, shape, *args, **kwargs: F.from_apply_func(
1715
                lambda p1, p2: np.full(shape, p1 + p2))
1716
                .run(*args, **kwargs)
1717
        ]
1718
    )
1719
    def test_no_inputs_wrapper(self, test_config):
1720
        F = vbt.IndicatorFactory(param_names=['p1', 'p2'], output_names=['out'])
1721

1722
        obj = test_config(F, (5,), 0, 1)
1723
        assert obj.wrapper.ndim == 1
1724
        pd.testing.assert_index_equal(obj.wrapper.index, pd.RangeIndex(start=0, stop=5, step=1))
1725
        pd.testing.assert_index_equal(
1726
            obj.wrapper.columns,
1727
            pd.MultiIndex.from_tuples([
1728
                (0, 1),
1729
            ], names=['custom_p1', 'custom_p2'])
1730
        )
1731
        obj = test_config(F, (5,), [0, 1], 2)
1732
        assert obj.wrapper.ndim == 2
1733
        pd.testing.assert_index_equal(obj.wrapper.index, pd.RangeIndex(start=0, stop=5, step=1))
1734
        pd.testing.assert_index_equal(
1735
            obj.wrapper.columns,
1736
            pd.MultiIndex.from_tuples([
1737
                (0, 2),
1738
                (1, 2)
1739
            ], names=['custom_p1', 'custom_p2'])
1740
        )
1741
        obj = test_config(F, (5, 3), [0, 1], 2)
1742
        assert obj.wrapper.ndim == 2
1743
        pd.testing.assert_index_equal(obj.wrapper.index, pd.RangeIndex(start=0, stop=5, step=1))
1744
        pd.testing.assert_index_equal(
1745
            obj.wrapper.columns,
1746
            pd.MultiIndex.from_tuples([
1747
                (0, 2, 0),
1748
                (0, 2, 1),
1749
                (0, 2, 2),
1750
                (1, 2, 0),
1751
                (1, 2, 1),
1752
                (1, 2, 2)
1753
            ], names=['custom_p1', 'custom_p2', None])
1754
        )
1755
        obj = test_config(F, ts.shape, [0, 1], 2, input_index=ts.index, input_columns=ts.columns)
1756
        assert obj.wrapper.ndim == ts.ndim
1757
        pd.testing.assert_index_equal(obj.wrapper.index, ts.index)
1758
        pd.testing.assert_index_equal(
1759
            obj.wrapper.columns,
1760
            pd.MultiIndex.from_tuples([
1761
                (0, 2, 'a'),
1762
                (0, 2, 'b'),
1763
                (0, 2, 'c'),
1764
                (1, 2, 'a'),
1765
                (1, 2, 'b'),
1766
                (1, 2, 'c')
1767
            ], names=['custom_p1', 'custom_p2', None])
1768
        )
1769

1770
    def test_mappers(self):
1771
        F = vbt.IndicatorFactory(input_names=['ts'], param_names=['p1', 'p2'], output_names=['out'])
1772

1773
        obj = F.from_apply_func(lambda ts, p1, p2: ts * (p1 + p2)) \
1774
            .run(ts, 0, 2)
1775
        np.testing.assert_array_equal(
1776
            obj._input_mapper,
1777
            np.array([0, 1, 2])
1778
        )
1779
        np.testing.assert_array_equal(
1780
            obj._p1_mapper,
1781
            np.array([0, 0, 0])
1782
        )
1783
        np.testing.assert_array_equal(
1784
            obj._p2_mapper,
1785
            np.array([2, 2, 2])
1786
        )
1787
        assert obj._tuple_mapper == [(0, 2), (0, 2), (0, 2)]
1788
        obj = F.from_apply_func(lambda ts, p1, p2: ts * (p1 + p2)) \
1789
            .run(ts, [0, 1], [1, 2])
1790
        np.testing.assert_array_equal(
1791
            obj._input_mapper,
1792
            np.array([0, 1, 2, 0, 1, 2])
1793
        )
1794
        np.testing.assert_array_equal(
1795
            obj._p1_mapper,
1796
            np.array([0, 0, 0, 1, 1, 1])
1797
        )
1798
        np.testing.assert_array_equal(
1799
            obj._p2_mapper,
1800
            np.array([1, 1, 1, 2, 2, 2])
1801
        )
1802
        assert obj._tuple_mapper == [(0, 1), (0, 1), (0, 1), (1, 2), (1, 2), (1, 2)]
1803
        obj = F.from_apply_func(lambda ts, p1, p2: ts * (p1 + p2)) \
1804
            .run(ts, [0, 1, 2], 3, per_column=True)
1805
        np.testing.assert_array_equal(
1806
            obj._input_mapper,
1807
            np.array([0, 1, 2])
1808
        )
1809
        np.testing.assert_array_equal(
1810
            obj._p1_mapper,
1811
            np.array([0, 1, 2])
1812
        )
1813
        np.testing.assert_array_equal(
1814
            obj._p2_mapper,
1815
            np.array([3, 3, 3])
1816
        )
1817
        assert obj._tuple_mapper == [(0, 3), (1, 3), (2, 3)]
1818

1819
    def test_properties(self):
1820
        I = vbt.IndicatorFactory(
1821
            input_names=['ts1', 'ts2'],
1822
            param_names=['p1', 'p2'],
1823
            output_names=['o1', 'o2'],
1824
            in_output_names=['in_o1', 'in_o2'],
1825
            output_flags={'o1': 'Hello'}
1826
        ).from_apply_func(lambda ts1, ts2, p1, p2, in_o1, in_o2: (ts1, ts2))
1827
        obj = I.run(ts, ts, [0, 1], 2)
1828

1829
        # Class properties
1830
        assert I.input_names == ('ts1', 'ts2')
1831
        assert I.param_names == ('p1', 'p2')
1832
        assert I.output_names == ('o1', 'o2')
1833
        assert I.in_output_names == ('in_o1', 'in_o2')
1834
        assert I.output_flags == {'o1': 'Hello'}
1835

1836
        # Instance properties
1837
        assert obj.input_names == ('ts1', 'ts2')
1838
        assert obj.param_names == ('p1', 'p2')
1839
        assert obj.output_names == ('o1', 'o2')
1840
        assert obj.in_output_names == ('in_o1', 'in_o2')
1841
        assert obj.output_flags == {'o1': 'Hello'}
1842
        assert obj.short_name == 'custom'
1843
        assert obj.level_names == ('custom_p1', 'custom_p2')
1844
        assert obj.p1_list == [0, 1]
1845
        assert obj.p2_list == [2, 2]
1846

1847
    @pytest.mark.parametrize(
1848
        "test_attr",
1849
        ['ts1', 'ts2', 'o1', 'o2', 'in_o1', 'in_o2', 'co1', 'co2']
1850
    )
1851
    def test_indexing(self, test_attr):
1852
        obj = vbt.IndicatorFactory(
1853
            input_names=['ts1', 'ts2'],
1854
            param_names=['p1', 'p2'],
1855
            output_names=['o1', 'o2'],
1856
            in_output_names=['in_o1', 'in_o2'],
1857
            custom_output_props={
1858
                'co1': lambda self: self.ts1 + self.ts2,
1859
                'co2': lambda self: self.o1 + self.o2
1860
            }
1861
        ).from_apply_func(lambda ts1, ts2, p1, p2, in_o1, in_o2: (ts1, ts2)).run(ts, ts + 1, [1, 2], 3)
1862

1863
        pd.testing.assert_frame_equal(
1864
            getattr(obj.iloc[np.arange(3), np.arange(3)], test_attr),
1865
            getattr(obj, test_attr).iloc[np.arange(3), np.arange(3)]
1866
        )
1867
        pd.testing.assert_series_equal(
1868
            getattr(obj.loc[:, (1, 3, 'a')], test_attr),
1869
            getattr(obj, test_attr).loc[:, (1, 3, 'a')]
1870
        )
1871
        pd.testing.assert_frame_equal(
1872
            getattr(obj.loc[:, (1, 3)], test_attr),
1873
            getattr(obj, test_attr).loc[:, (1, 3)]
1874
        )
1875
        pd.testing.assert_frame_equal(
1876
            getattr(obj[(1, 3)], test_attr),
1877
            getattr(obj, test_attr)[(1, 3)]
1878
        )
1879
        pd.testing.assert_frame_equal(
1880
            getattr(obj.xs(1, axis=1, level=0), test_attr),
1881
            getattr(obj, test_attr).xs(1, axis=1, level=0)
1882
        )
1883
        pd.testing.assert_frame_equal(
1884
            getattr(obj.p1_loc[2], test_attr),
1885
            getattr(obj, test_attr).xs(2, level='custom_p1', axis=1)
1886
        )
1887
        pd.testing.assert_frame_equal(
1888
            getattr(obj.p1_loc[1:2], test_attr),
1889
            pd.concat((
1890
                getattr(obj, test_attr).xs(1, level='custom_p1', drop_level=False, axis=1),
1891
                getattr(obj, test_attr).xs(2, level='custom_p1', drop_level=False, axis=1)
1892
            ), axis=1)
1893
        )
1894
        pd.testing.assert_frame_equal(
1895
            getattr(obj.p1_loc[[1, 1, 1]], test_attr),
1896
            pd.concat((
1897
                getattr(obj, test_attr).xs(1, level='custom_p1', drop_level=False, axis=1),
1898
                getattr(obj, test_attr).xs(1, level='custom_p1', drop_level=False, axis=1),
1899
                getattr(obj, test_attr).xs(1, level='custom_p1', drop_level=False, axis=1)
1900
            ), axis=1)
1901
        )
1902
        pd.testing.assert_frame_equal(
1903
            getattr(obj.tuple_loc[(1, 3)], test_attr),
1904
            getattr(obj, test_attr).xs((1, 3), level=('custom_p1', 'custom_p2'), axis=1)
1905
        )
1906
        pd.testing.assert_frame_equal(
1907
            getattr(obj.tuple_loc[(1, 3):(2, 3)], test_attr),
1908
            pd.concat((
1909
                getattr(obj, test_attr).xs((1, 3), level=('custom_p1', 'custom_p2'), drop_level=False, axis=1),
1910
                getattr(obj, test_attr).xs((2, 3), level=('custom_p1', 'custom_p2'), drop_level=False, axis=1)
1911
            ), axis=1)
1912
        )
1913
        pd.testing.assert_frame_equal(
1914
            getattr(obj.tuple_loc[[(1, 3), (1, 3), (1, 3)]], test_attr),
1915
            pd.concat((
1916
                getattr(obj, test_attr).xs((1, 3), level=('custom_p1', 'custom_p2'), drop_level=False, axis=1),
1917
                getattr(obj, test_attr).xs((1, 3), level=('custom_p1', 'custom_p2'), drop_level=False, axis=1),
1918
                getattr(obj, test_attr).xs((1, 3), level=('custom_p1', 'custom_p2'), drop_level=False, axis=1)
1919
            ), axis=1)
1920
        )
1921

1922
    def test_numeric_attr(self):
1923
        obj = vbt.IndicatorFactory(input_names=['ts'], param_names=['p'], output_names=['out']) \
1924
            .from_apply_func(lambda ts, p: ts * p).run(ts, 1)
1925

1926
        pd.testing.assert_frame_equal(obj.out_above(2), obj.out > 2)
1927
        target = pd.DataFrame(
1928
            np.array([
1929
                [False, True, False, False, True, False],
1930
                [False, True, False, False, True, False],
1931
                [True, True, True, False, False, False],
1932
                [True, False, False, True, False, False],
1933
                [True, False, False, True, False, False]
1934
            ]),
1935
            index=ts.index,
1936
            columns=pd.MultiIndex.from_tuples([
1937
                (2, 1, 'a'),
1938
                (2, 1, 'b'),
1939
                (2, 1, 'c'),
1940
                (3, 1, 'a'),
1941
                (3, 1, 'b'),
1942
                (3, 1, 'c'),
1943
            ], names=['custom_out_above', 'custom_p', None])
1944
        )
1945
        pd.testing.assert_frame_equal(
1946
            obj.out_above([2, 3]),
1947
            target
1948
        )
1949
        columns = target.columns.set_names('my_above', level=0)
1950
        pd.testing.assert_frame_equal(
1951
            obj.out_above([2, 3], level_name='my_above'),
1952
            pd.DataFrame(
1953
                target.values,
1954
                index=target.index,
1955
                columns=columns
1956
            )
1957
        )
1958
        pd.testing.assert_frame_equal(
1959
            obj.out_crossed_above(2),
1960
            pd.DataFrame(
1961
                np.array([
1962
                    [False, False, False],
1963
                    [False, False, False],
1964
                    [True, False, True],
1965
                    [False, False, False],
1966
                    [False, False, False]
1967
                ]),
1968
                index=ts.index,
1969
                columns=pd.MultiIndex.from_tuples([
1970
                    (1, 'a'),
1971
                    (1, 'b'),
1972
                    (1, 'c'),
1973
                ], names=['custom_p', None])
1974
            )
1975
        )
1976
        pd.testing.assert_series_equal(
1977
            obj.out_stats(),
1978
            pd.Series([
1979
                pd.Timestamp('2018-01-01 00:00:00'),
1980
                pd.Timestamp('2018-01-05 00:00:00'),
1981
                pd.Timedelta('5 days 00:00:00'),
1982
                5.0, 2.6, 1.3329792289008184, 1.0, 2.6666666666666665, 4.333333333333333
1983
            ],
1984
                index=pd.Index([
1985
                    'Start', 'End', 'Period', 'Count', 'Mean', 'Std', 'Min', 'Median', 'Max'
1986
                ], dtype='object'),
1987
                name='agg_func_mean'
1988
            )
1989
        )
1990

1991
    def test_boolean_attr(self):
1992
        obj = vbt.IndicatorFactory(
1993
            input_names=['ts'], param_names=['p'], output_names=['out'],
1994
            attr_settings=dict(out=dict(dtype=np.bool_))) \
1995
            .from_apply_func(lambda ts, p: ts > p).run(ts, 2)
1996

1997
        pd.testing.assert_frame_equal(obj.out_and(True), obj.out)
1998
        target = pd.DataFrame(
1999
            np.array([
2000
                [False, False, False, False, True, False],
2001
                [False, False, False, False, True, False],
2002
                [False, False, False, True, True, True],
2003
                [False, False, False, True, False, False],
2004
                [False, False, False, True, False, False]
2005
            ]),
2006
            index=ts.index,
2007
            columns=pd.MultiIndex.from_tuples([
2008
                (False, 2, 'a'),
2009
                (False, 2, 'b'),
2010
                (False, 2, 'c'),
2011
                (True, 2, 'a'),
2012
                (True, 2, 'b'),
2013
                (True, 2, 'c'),
2014
            ], names=['custom_out_and', 'custom_p', None])
2015
        )
2016
        pd.testing.assert_frame_equal(
2017
            obj.out_and([False, True]),
2018
            target
2019
        )
2020
        columns = target.columns.set_names('my_and', level=0)
2021
        pd.testing.assert_frame_equal(
2022
            obj.out_and([False, True], level_name='my_and'),
2023
            pd.DataFrame(
2024
                target.values,
2025
                index=target.index,
2026
                columns=columns
2027
            )
2028
        )
2029
        pd.testing.assert_series_equal(
2030
            obj.out_stats(),
2031
            pd.Series([
2032
                pd.Timestamp('2018-01-01 00:00:00'), pd.Timestamp('2018-01-05 00:00:00'),
2033
                pd.Timedelta('5 days 00:00:00'), 2.3333333333333335, 46.666666666666664,
2034
                pd.Timestamp('2018-01-02 08:00:00'), pd.Timestamp('2018-01-03 16:00:00'), 0.0,
2035
                pd.Timedelta('1 days 00:00:00'), pd.Timedelta('1 days 00:00:00'),
2036
                pd.Timedelta('1 days 00:00:00'), pd.Timedelta('0 days 00:00:00'),
2037
                1.0, 55.55555555555555, pd.Timedelta('2 days 08:00:00'),
2038
                pd.Timedelta('2 days 08:00:00'), pd.Timedelta('2 days 08:00:00'),
2039
                pd.NaT, pd.NaT, pd.NaT, pd.NaT, pd.NaT
2040
            ],
2041
                index=pd.Index([
2042
                    'Start', 'End', 'Period', 'Total', 'Rate [%]', 'First Index',
2043
                    'Last Index', 'Norm Avg Index [-1, 1]', 'Distance: Min',
2044
                    'Distance: Max', 'Distance: Mean', 'Distance: Std', 'Total Partitions',
2045
                    'Partition Rate [%]', 'Partition Length: Min', 'Partition Length: Max',
2046
                    'Partition Length: Mean', 'Partition Length: Std',
2047
                    'Partition Distance: Min', 'Partition Distance: Max',
2048
                    'Partition Distance: Mean', 'Partition Distance: Std'
2049
                ], dtype='object'),
2050
                name='agg_func_mean'
2051
            )
2052
        )
2053

2054
    def test_mapping_attr(self):
2055
        TestEnum = namedtuple('TestEnum', ['Hello', 'World'])(0, 1)
2056
        obj = vbt.IndicatorFactory(
2057
            output_names=['out'],
2058
            attr_settings=dict(out=dict(dtype=TestEnum))) \
2059
            .from_apply_func(lambda: np.array([[0, 1], [1, -1]])).run()
2060

2061
        pd.testing.assert_frame_equal(
2062
            obj.out_readable,
2063
            pd.DataFrame([
2064
                ['Hello', 'World'],
2065
                ['World', None]
2066
            ])
2067
        )
2068
        pd.testing.assert_series_equal(
2069
            obj.out_stats(),
2070
            pd.Series([
2071
                0.0, 1.0, 2.0, 0.5, 0.5, 1.0
2072
            ],
2073
                index=pd.Index([
2074
                    'Start', 'End', 'Period', 'Value Counts: None', 'Value Counts: Hello', 'Value Counts: World'
2075
                ], dtype='object'),
2076
                name='agg_func_mean'
2077
            )
2078
        )
2079

2080
    def test_stats(self):
2081
        @njit
2082
        def apply_func_nb(ts):
2083
            return ts ** 2, ts ** 3
2084

2085
        MyInd = vbt.IndicatorFactory(
2086
            input_names=['ts'],
2087
            output_names=['out1', 'out2'],
2088
            metrics=dict(
2089
                sum_diff=dict(
2090
                    calc_func=lambda self, const: self.out2.sum() * self.out1.sum() + const
2091
                )
2092
            ),
2093
            stats_defaults=dict(settings=dict(const=1000))
2094
        ).from_apply_func(
2095
            apply_func_nb
2096
        )
2097

2098
        myind = MyInd.run(ts)
2099
        pd.testing.assert_series_equal(
2100
            myind.stats(),
2101
            pd.Series([9535.0], index=['sum_diff'], name='agg_func_mean')
2102
        )
2103

2104
    def test_dir(self):
2105
        TestEnum = namedtuple('TestEnum', ['Hello', 'World'])(0, 1)
2106
        F = vbt.IndicatorFactory(
2107
            input_names=['ts'], output_names=['o1', 'o2'], in_output_names=['in_out'], param_names=['p1', 'p2'],
2108
            attr_settings={
2109
                'ts': {'dtype': None},
2110
                'o1': {'dtype': np.float64},
2111
                'o2': {'dtype': np.bool_},
2112
                'in_out': {'dtype': TestEnum}
2113
            }
2114
        )
2115
        ind = F.from_apply_func(lambda ts, in_out, p1, p2: (ts + in_out, ts + in_out)).run(ts, 100, 200)
2116
        test_attr_list = dir(ind)
2117
        assert test_attr_list == [
2118
            '__annotations__',
2119
            '__class__',
2120
            '__delattr__',
2121
            '__dict__',
2122
            '__dir__',
2123
            '__doc__',
2124
            '__eq__',
2125
            '__format__',
2126
            '__ge__',
2127
            '__getattribute__',
2128
            '__getitem__',
2129
            '__gt__',
2130
            '__hash__',
2131
            '__init__',
2132
            '__init_subclass__',
2133
            '__le__',
2134
            '__lt__',
2135
            '__module__',
2136
            '__ne__',
2137
            '__new__',
2138
            '__reduce__',
2139
            '__reduce_ex__',
2140
            '__repr__',
2141
            '__setattr__',
2142
            '__sizeof__',
2143
            '__str__',
2144
            '__subclasshook__',
2145
            '__weakref__',
2146
            '_config',
2147
            '_iloc',
2148
            '_in_out',
2149
            '_in_output_names',
2150
            '_indexing_kwargs',
2151
            '_input_mapper',
2152
            '_input_names',
2153
            '_level_names',
2154
            '_loc',
2155
            '_metrics',
2156
            '_o1',
2157
            '_o2',
2158
            '_output_flags',
2159
            '_output_names',
2160
            '_p1_list',
2161
            '_p1_loc',
2162
            '_p1_mapper',
2163
            '_p2_list',
2164
            '_p2_loc',
2165
            '_p2_mapper',
2166
            '_param_names',
2167
            '_run',
2168
            '_run_combs',
2169
            '_short_name',
2170
            '_subplots',
2171
            '_ts',
2172
            '_tuple_loc',
2173
            '_tuple_mapper',
2174
            '_wrapper',
2175
            'apply_func',
2176
            'build_metrics_doc',
2177
            'build_subplots_doc',
2178
            'config',
2179
            'copy',
2180
            'custom_func',
2181
            'deep_getattr',
2182
            'dumps',
2183
            'iloc',
2184
            'in_out',
2185
            'in_out_readable',
2186
            'in_out_stats',
2187
            'in_output_names',
2188
            'indexing_func',
2189
            'indexing_kwargs',
2190
            'input_names',
2191
            'level_names',
2192
            'load',
2193
            'loads',
2194
            'loc',
2195
            'metrics',
2196
            'o1',
2197
            'o1_above',
2198
            'o1_below',
2199
            'o1_crossed_above',
2200
            'o1_crossed_below',
2201
            'o1_equal',
2202
            'o1_stats',
2203
            'o2',
2204
            'o2_and',
2205
            'o2_or',
2206
            'o2_stats',
2207
            'o2_xor',
2208
            'output_flags',
2209
            'output_names',
2210
            'override_metrics_doc',
2211
            'override_subplots_doc',
2212
            'p1_list',
2213
            'p1_loc',
2214
            'p2_list',
2215
            'p2_loc',
2216
            'param_names',
2217
            'plots',
2218
            'plots_defaults',
2219
            'post_resolve_attr',
2220
            'pre_resolve_attr',
2221
            'regroup',
2222
            'replace',
2223
            'resolve_attr',
2224
            'resolve_self',
2225
            'run',
2226
            'run_combs',
2227
            'save',
2228
            'select_one',
2229
            'select_one_from_obj',
2230
            'self_aliases',
2231
            'short_name',
2232
            'stats',
2233
            'stats_defaults',
2234
            'subplots',
2235
            'to_doc',
2236
            'ts',
2237
            'ts_above',
2238
            'ts_below',
2239
            'ts_crossed_above',
2240
            'ts_crossed_below',
2241
            'ts_equal',
2242
            'ts_stats',
2243
            'tuple_loc',
2244
            'update_config',
2245
            'wrapper',
2246
            'writeable_attrs',
2247
            'xs'
2248
        ]
2249

2250
    def test_get_talib_indicators(self):
2251
        if talib_available:
2252
            assert len(vbt.IndicatorFactory.get_talib_indicators()) > 0
2253

2254
    def test_from_talib(self):
2255
        if talib_available:
2256
            # with params
2257
            target = pd.DataFrame(
2258
                np.array([
2259
                    [np.nan, np.nan, np.nan],
2260
                    [2.5, 5.5, 2.5],
2261
                    [3.5, 4.5, 3.5],
2262
                    [4.5, 3.5, 3.5],
2263
                    [5.5, 2.5, 2.5]
2264
                ]),
2265
                index=ts.index,
2266
                columns=pd.MultiIndex.from_tuples([
2267
                    (2, 2, 2, 'a'),
2268
                    (2, 2, 2, 'b'),
2269
                    (2, 2, 2, 'c')
2270
                ], names=['bbands_timeperiod', 'bbands_nbdevup', 'bbands_nbdevdn', None])
2271
            )
2272
            BBANDS = vbt.talib('BBANDS')
2273
            pd.testing.assert_frame_equal(
2274
                BBANDS.run(ts, timeperiod=2, nbdevup=2, nbdevdn=2).upperband,
2275
                target
2276
            )
2277
            pd.testing.assert_frame_equal(
2278
                BBANDS.run(ts, timeperiod=2, nbdevup=2, nbdevdn=2).middleband,
2279
                target - 1
2280
            )
2281
            pd.testing.assert_frame_equal(
2282
                BBANDS.run(ts, timeperiod=2, nbdevup=2, nbdevdn=2).lowerband,
2283
                target - 2
2284
            )
2285
            target = pd.DataFrame(
2286
                np.array([
2287
                    [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan],
2288
                    [2.5, 5.5, 2.5, 2.5, 5.5, 2.5],
2289
                    [3.5, 4.5, 3.5, 3.5, 4.5, 3.5],
2290
                    [4.5, 3.5, 3.5, 4.5, 3.5, 3.5],
2291
                    [5.5, 2.5, 2.5, 5.5, 2.5, 2.5]
2292
                ]),
2293
                index=ts.index,
2294
                columns=pd.MultiIndex.from_tuples([
2295
                    (2, 2, 2, 'a'),
2296
                    (2, 2, 2, 'b'),
2297
                    (2, 2, 2, 'c'),
2298
                    (2, 2, 2, 'a'),
2299
                    (2, 2, 2, 'b'),
2300
                    (2, 2, 2, 'c')
2301
                ], names=['bbands_timeperiod', 'bbands_nbdevup', 'bbands_nbdevdn', None])
2302
            )
2303
            BBANDS = vbt.talib('BBANDS')
2304
            pd.testing.assert_frame_equal(
2305
                BBANDS.run(ts, timeperiod=[2, 2], nbdevup=2, nbdevdn=2).upperband,
2306
                target
2307
            )
2308
            pd.testing.assert_frame_equal(
2309
                BBANDS.run(ts, timeperiod=[2, 2], nbdevup=2, nbdevdn=2).middleband,
2310
                target - 1
2311
            )
2312
            pd.testing.assert_frame_equal(
2313
                BBANDS.run(ts, timeperiod=[2, 2], nbdevup=2, nbdevdn=2).lowerband,
2314
                target - 2
2315
            )
2316
            # without params
2317
            OBV = vbt.talib('OBV')
2318
            pd.testing.assert_frame_equal(
2319
                OBV.run(ts, ts * 2).real,
2320
                pd.DataFrame(
2321
                    np.array([
2322
                        [2., 10., 2.],
2323
                        [6., 2., 6.],
2324
                        [12., -4., 12.],
2325
                        [20., -8., 8.],
2326
                        [30., -10., 6.]
2327
                    ]),
2328
                    index=ts.index,
2329
                    columns=ts.columns
2330
                )
2331
            )
2332

2333
    def test_get_pandas_ta_indicators(self):
2334
        if pandas_ta_available:
2335
            assert len(vbt.IndicatorFactory.get_pandas_ta_indicators()) > 0
2336

2337
    def test_from_pandas_ta(self):
2338
        if pandas_ta_available:
2339
            pd.testing.assert_frame_equal(
2340
                vbt.pandas_ta('SMA').run(ts, 2).sma,
2341
                pd.DataFrame(
2342
                    ts.rolling(2).mean().values,
2343
                    index=ts.index,
2344
                    columns=pd.MultiIndex.from_tuples([(2, 'a'), (2, 'b'), (2, 'c')], names=['sma_length', None])
2345
                )
2346
            )
2347
            pd.testing.assert_frame_equal(
2348
                vbt.pandas_ta('SMA').run(ts['a'], [2, 3, 4]).sma,
2349
                pd.DataFrame(
2350
                    np.column_stack((
2351
                        ts['a'].rolling(2).mean().values,
2352
                        ts['a'].rolling(3).mean().values,
2353
                        ts['a'].rolling(4).mean().values
2354
                    )),
2355
                    index=ts.index,
2356
                    columns=pd.Index([2, 3, 4], dtype='int64', name='sma_length')
2357
                )
2358
            )
2359

2360
    def test_get_ta_indicators(self):
2361
        if ta_available:
2362
            assert len(vbt.IndicatorFactory.get_ta_indicators()) > 0
2363

2364
    def test_from_ta(self):
2365
        if ta_available:
2366
            pd.testing.assert_frame_equal(
2367
                vbt.ta('SMAIndicator').run(ts, 2).sma_indicator,
2368
                pd.DataFrame(
2369
                    ts.rolling(2).mean().values,
2370
                    index=ts.index,
2371
                    columns=pd.MultiIndex.from_tuples([(2, 'a'), (2, 'b'), (2, 'c')],
2372
                                                      names=['smaindicator_window', None])
2373
                )
2374
            )
2375
            pd.testing.assert_frame_equal(
2376
                vbt.ta('SMAIndicator').run(ts['a'], [2, 3, 4]).sma_indicator,
2377
                pd.DataFrame(
2378
                    np.column_stack((
2379
                        ts['a'].rolling(2).mean().values,
2380
                        ts['a'].rolling(3).mean().values,
2381
                        ts['a'].rolling(4).mean().values
2382
                    )),
2383
                    index=ts.index,
2384
                    columns=pd.Index([2, 3, 4], dtype='int64', name='smaindicator_window')
2385
                )
2386
            )
2387
            target = pd.DataFrame(
2388
                np.array([
2389
                    [np.nan, np.nan, np.nan],
2390
                    [2.5, 5.5, 2.5],
2391
                    [3.5, 4.5, 3.5],
2392
                    [4.5, 3.5, 3.5],
2393
                    [5.5, 2.5, 2.5]
2394
                ]),
2395
                index=ts.index,
2396
                columns=pd.MultiIndex.from_tuples([
2397
                    (2, 2, 'a'),
2398
                    (2, 2, 'b'),
2399
                    (2, 2, 'c')
2400
                ], names=['bollingerbands_window', 'bollingerbands_window_dev', None])
2401
            )
2402
            BollingerBands = vbt.ta('BollingerBands')
2403
            pd.testing.assert_frame_equal(
2404
                BollingerBands.run(ts, window=2, window_dev=2).bollinger_hband,
2405
                target
2406
            )
2407
            pd.testing.assert_frame_equal(
2408
                BollingerBands.run(ts, window=2, window_dev=2).bollinger_mavg,
2409
                target - 1
2410
            )
2411
            pd.testing.assert_frame_equal(
2412
                BollingerBands.run(ts, window=2, window_dev=2).bollinger_lband,
2413
                target - 2
2414
            )
2415

2416

2417
# ############# basic.py ############# #
2418

2419
close_ts = pd.Series([1, 2, 3, 4, 3, 2, 1], index=pd.DatetimeIndex([
2420
    datetime(2018, 1, 1),
2421
    datetime(2018, 1, 2),
2422
    datetime(2018, 1, 3),
2423
    datetime(2018, 1, 4),
2424
    datetime(2018, 1, 5),
2425
    datetime(2018, 1, 6),
2426
    datetime(2018, 1, 7)
2427
]))
2428
high_ts = close_ts * 1.1
2429
low_ts = close_ts * 0.9
2430
volume_ts = pd.Series([4, 3, 2, 1, 2, 3, 4], index=close_ts.index)
2431

2432

2433
class TestBasic:
2434
    def test_MA(self):
2435
        pd.testing.assert_frame_equal(
2436
            vbt.MA.run(close_ts, window=(2, 3), ewm=(False, True), param_product=True).ma,
2437
            pd.DataFrame(
2438
                np.array([
2439
                    [np.nan, np.nan, np.nan, np.nan],
2440
                    [1.5, 1.66666667, np.nan, np.nan],
2441
                    [2.5, 2.55555556, 2., 2.25],
2442
                    [3.5, 3.51851852, 3., 3.125],
2443
                    [3.5, 3.17283951, 3.33333333, 3.0625],
2444
                    [2.5, 2.3909465, 3., 2.53125],
2445
                    [1.5, 1.46364883, 2., 1.765625]
2446
                ]),
2447
                index=close_ts.index,
2448
                columns=pd.MultiIndex.from_tuples([
2449
                    (2, False),
2450
                    (2, True),
2451
                    (3, False),
2452
                    (3, True)
2453
                ], names=['ma_window', 'ma_ewm'])
2454
            )
2455
        )
2456

2457
    def test_MSTD(self):
2458
        pd.testing.assert_frame_equal(
2459
            vbt.MSTD.run(close_ts, window=(2, 3), ewm=(False, True), param_product=True).mstd,
2460
            pd.DataFrame(
2461
                np.array([
2462
                    [np.nan, np.nan, np.nan, np.nan],
2463
                    [0.5, 0.70710678, np.nan, np.nan],
2464
                    [0.5, 0.97467943, 0.81649658, 1.04880885],
2465
                    [0.5, 1.11434207, 0.81649658, 1.30018314],
2466
                    [0.5, 0.73001838, 0.47140452, 0.91715673],
2467
                    [0.5, 0.88824841, 0.81649658, 0.9182094],
2468
                    [0.5, 1.05965735, 0.81649658, 1.14049665]
2469
                ]),
2470
                index=close_ts.index,
2471
                columns=pd.MultiIndex.from_tuples([
2472
                    (2, False),
2473
                    (2, True),
2474
                    (3, False),
2475
                    (3, True)
2476
                ], names=['mstd_window', 'mstd_ewm'])
2477
            )
2478
        )
2479

2480
    def test_BBANDS(self):
2481
        columns = pd.MultiIndex.from_tuples([
2482
            (2, False, 2.0),
2483
            (2, False, 3.0),
2484
            (2, True, 2.0),
2485
            (2, True, 3.0),
2486
            (3, False, 2.0),
2487
            (3, False, 3.0),
2488
            (3, True, 2.0),
2489
            (3, True, 3.0)
2490
        ], names=['bb_window', 'bb_ewm', 'bb_alpha'])
2491
        pd.testing.assert_frame_equal(
2492
            vbt.BBANDS.run(
2493
                close_ts,
2494
                window=(2, 3),
2495
                alpha=(2., 3.),
2496
                ewm=(False, True),
2497
                param_product=True
2498
            ).middle,
2499
            pd.DataFrame(
2500
                np.array([
2501
                    [np.nan, np.nan, np.nan, np.nan, np.nan,
2502
                     np.nan, np.nan, np.nan],
2503
                    [1.5, 1.5, 1.66666667, 1.66666667, np.nan,
2504
                     np.nan, np.nan, np.nan],
2505
                    [2.5, 2.5, 2.55555556, 2.55555556, 2.,
2506
                     2., 2.25, 2.25],
2507
                    [3.5, 3.5, 3.51851852, 3.51851852, 3.,
2508
                     3., 3.125, 3.125],
2509
                    [3.5, 3.5, 3.17283951, 3.17283951, 3.33333333,
2510
                     3.33333333, 3.0625, 3.0625],
2511
                    [2.5, 2.5, 2.3909465, 2.3909465, 3.,
2512
                     3., 2.53125, 2.53125],
2513
                    [1.5, 1.5, 1.46364883, 1.46364883, 2.,
2514
                     2., 1.765625, 1.765625]
2515
                ]),
2516
                index=close_ts.index,
2517
                columns=columns
2518
            )
2519
        )
2520
        pd.testing.assert_frame_equal(
2521
            vbt.BBANDS.run(
2522
                close_ts,
2523
                window=(2, 3),
2524
                alpha=(2., 3.),
2525
                ewm=(False, True),
2526
                param_product=True
2527
            ).upper,
2528
            pd.DataFrame(
2529
                np.array([
2530
                    [np.nan, np.nan, np.nan, np.nan, np.nan,
2531
                     np.nan, np.nan, np.nan],
2532
                    [2.5, 3., 3.08088023, 3.78798701, np.nan,
2533
                     np.nan, np.nan, np.nan],
2534
                    [3.5, 4., 4.50491442, 5.47959386, 3.63299316,
2535
                     4.44948974, 4.3476177, 5.39642654],
2536
                    [4.5, 5., 5.74720265, 6.86154472, 4.63299316,
2537
                     5.44948974, 5.72536627, 7.02554941],
2538
                    [4.5, 5., 4.63287626, 5.36289463, 4.27614237,
2539
                     4.7475469, 4.89681346, 5.8139702],
2540
                    [3.5, 4., 4.16744332, 5.05569172, 4.63299316,
2541
                     5.44948974, 4.3676688, 5.2858782],
2542
                    [2.5, 3., 3.58296354, 4.64262089, 3.63299316,
2543
                     4.44948974, 4.04661829, 5.18711494]
2544
                ]),
2545
                index=close_ts.index,
2546
                columns=columns
2547
            )
2548
        )
2549
        pd.testing.assert_frame_equal(
2550
            vbt.BBANDS.run(
2551
                close_ts,
2552
                window=(2, 3),
2553
                alpha=(2., 3.),
2554
                ewm=(False, True),
2555
                param_product=True
2556
            ).lower,
2557
            pd.DataFrame(
2558
                np.array([
2559
                    [np.nan, np.nan, np.nan, np.nan, np.nan,
2560
                     np.nan, np.nan, np.nan],
2561
                    [0.5, 0., 0.2524531, -0.45465368, np.nan,
2562
                     np.nan, np.nan, np.nan],
2563
                    [1.5, 1., 0.60619669, -0.36848275, 0.36700684,
2564
                     -0.44948974, 0.1523823, -0.89642654],
2565
                    [2.5, 2., 1.28983438, 0.17549232, 1.36700684,
2566
                     0.55051026, 0.52463373, -0.77554941],
2567
                    [2.5, 2., 1.71280275, 0.98278438, 2.39052429,
2568
                     1.91911977, 1.22818654, 0.3110298],
2569
                    [1.5, 1., 0.61444969, -0.27379872, 1.36700684,
2570
                     0.55051026, 0.6948312, -0.2233782],
2571
                    [0.5, 0., -0.65566587, -1.71532322, 0.36700684,
2572
                     -0.44948974, -0.51536829, -1.65586494]
2573
                ]),
2574
                index=close_ts.index,
2575
                columns=columns
2576
            )
2577
        )
2578
        pd.testing.assert_frame_equal(
2579
            vbt.BBANDS.run(
2580
                close_ts,
2581
                window=(2, 3),
2582
                alpha=(2., 3.),
2583
                ewm=(False, True),
2584
                param_product=True
2585
            ).percent_b,
2586
            pd.DataFrame(
2587
                np.array([
2588
                    [np.nan, np.nan, np.nan, np.nan, np.nan,
2589
                     np.nan, np.nan, np.nan],
2590
                    [0.75, 0.66666667, 0.61785113, 0.57856742, np.nan,
2591
                     np.nan, np.nan, np.nan],
2592
                    [0.75, 0.66666667, 0.61399759, 0.5759984, 0.80618622,
2593
                     0.70412415, 0.67877424, 0.61918282],
2594
                    [0.75, 0.66666667, 0.60801923, 0.57201282, 0.80618622,
2595
                     0.70412415, 0.66824553, 0.61216369],
2596
                    [0.25, 0.33333333, 0.44080988, 0.46053992, 0.3232233,
2597
                     0.38214887, 0.48296365, 0.48864244],
2598
                    [0.25, 0.33333333, 0.38996701, 0.42664468, 0.19381378,
2599
                     0.29587585, 0.35535707, 0.40357138],
2600
                    [0.25, 0.33333333, 0.3906135, 0.42707567, 0.19381378,
2601
                     0.29587585, 0.3321729, 0.38811526]
2602
                ]),
2603
                index=close_ts.index,
2604
                columns=columns
2605
            )
2606
        )
2607
        pd.testing.assert_frame_equal(
2608
            vbt.BBANDS.run(
2609
                close_ts,
2610
                window=(2, 3),
2611
                alpha=(2., 3.),
2612
                ewm=(False, True),
2613
                param_product=True
2614
            ).bandwidth,
2615
            pd.DataFrame(
2616
                np.array([
2617
                    [np.nan, np.nan, np.nan, np.nan, np.nan,
2618
                     np.nan, np.nan, np.nan],
2619
                    [1.33333333, 2., 1.69705627, 2.54558441, np.nan,
2620
                     np.nan, np.nan, np.nan],
2621
                    [0.8, 1.2, 1.5255852, 2.2883778, 1.63299316,
2622
                     2.44948974, 1.86454906, 2.7968236],
2623
                    [0.57142857, 0.85714286, 1.26683098, 1.90024647, 1.08866211,
2624
                     1.63299316, 1.66423442, 2.49635162],
2625
                    [0.57142857, 0.85714286, 0.92033445, 1.38050168, 0.56568542,
2626
                     0.84852814, 1.197919, 1.79687849],
2627
                    [0.8, 1.2, 1.48601971, 2.22902956, 1.08866211,
2628
                     1.63299316, 1.45099757, 2.17649636],
2629
                    [1.33333333, 2., 2.8959333, 4.34389996, 1.63299316,
2630
                     2.44948974, 2.58378001, 3.87567002]
2631
                ]),
2632
                index=close_ts.index,
2633
                columns=columns
2634
            )
2635
        )
2636

2637
    def test_RSI(self):
2638
        pd.testing.assert_frame_equal(
2639
            vbt.RSI.run(close_ts, window=(2, 3), ewm=(False, True), param_product=True).rsi,
2640
            pd.DataFrame(
2641
                np.array([
2642
                    [np.nan, np.nan, np.nan, np.nan],
2643
                    [np.nan, np.nan, np.nan, np.nan],
2644
                    [100., 100., np.nan, np.nan],
2645
                    [100., 100., 100., 100.],
2646
                    [50., 33.33333333, 66.66666667, 50.],
2647
                    [0., 11.11111111, 33.33333333, 25.],
2648
                    [0., 3.7037037, 0., 12.5]
2649
                ]),
2650
                index=close_ts.index,
2651
                columns=pd.MultiIndex.from_tuples([
2652
                    (2, False),
2653
                    (2, True),
2654
                    (3, False),
2655
                    (3, True)
2656
                ], names=['rsi_window', 'rsi_ewm'])
2657
            )
2658
        )
2659

2660
    def test_STOCH(self):
2661
        columns = pd.MultiIndex.from_tuples([
2662
            (2, 2, False),
2663
            (2, 2, True),
2664
            (2, 3, False),
2665
            (2, 3, True),
2666
            (3, 2, False),
2667
            (3, 2, True),
2668
            (3, 3, False),
2669
            (3, 3, True)
2670
        ], names=['stoch_k_window', 'stoch_d_window', 'stoch_d_ewm'])
2671
        pd.testing.assert_frame_equal(
2672
            vbt.STOCH.run(
2673
                high_ts,
2674
                low_ts,
2675
                close_ts,
2676
                k_window=(2, 3),
2677
                d_window=(2, 3),
2678
                d_ewm=(False, True),
2679
                param_product=True
2680
            ).percent_k,
2681
            pd.DataFrame(
2682
                np.array([
2683
                    [np.nan, np.nan, np.nan, np.nan, np.nan,
2684
                     np.nan, np.nan, np.nan],
2685
                    [84.61538462, 84.61538462, 84.61538462, 84.61538462, np.nan,
2686
                     np.nan, np.nan, np.nan],
2687
                    [80., 80., 80., 80., 87.5,
2688
                     87.5, 87.5, 87.5],
2689
                    [76.47058824, 76.47058824, 76.47058824, 76.47058824, 84.61538462,
2690
                     84.61538462, 84.61538462, 84.61538462],
2691
                    [17.64705882, 17.64705882, 17.64705882, 17.64705882, 17.64705882,
2692
                     17.64705882, 17.64705882, 17.64705882],
2693
                    [13.33333333, 13.33333333, 13.33333333, 13.33333333, 7.69230769,
2694
                     7.69230769, 7.69230769, 7.69230769],
2695
                    [7.69230769, 7.69230769, 7.69230769, 7.69230769, 4.16666667,
2696
                     4.16666667, 4.16666667, 4.16666667]
2697
                ]),
2698
                index=close_ts.index,
2699
                columns=columns
2700
            )
2701
        )
2702
        pd.testing.assert_frame_equal(
2703
            vbt.STOCH.run(
2704
                high_ts,
2705
                low_ts,
2706
                close_ts,
2707
                k_window=(2, 3),
2708
                d_window=(2, 3),
2709
                d_ewm=(False, True),
2710
                param_product=True
2711
            ).percent_d,
2712
            pd.DataFrame(
2713
                np.array([
2714
                    [np.nan, np.nan, np.nan, np.nan, np.nan,
2715
                     np.nan, np.nan, np.nan],
2716
                    [np.nan, np.nan, np.nan, np.nan, np.nan,
2717
                     np.nan, np.nan, np.nan],
2718
                    [82.30769231, 81.53846154, np.nan, np.nan, np.nan,
2719
                     np.nan, np.nan, np.nan],
2720
                    [78.23529412, 78.15987934, 80.36199095, 79.38914027, 86.05769231,
2721
                     85.57692308, np.nan, np.nan],
2722
                    [47.05882353, 37.81799899, 58.03921569, 48.51809955, 51.13122172,
2723
                     40.29034691, 63.25414781, 51.85237557],
2724
                    [15.49019608, 21.49488855, 35.81699346, 30.92571644, 12.66968326,
2725
                     18.55832076, 36.65158371, 29.77234163],
2726
                    [10.51282051, 12.29316798, 12.89089995, 19.30901207, 5.92948718,
2727
                     8.9638847, 9.83534439, 16.96950415]
2728
                ]),
2729
                index=close_ts.index,
2730
                columns=columns
2731
            )
2732
        )
2733

2734
    def test_MACD(self):
2735
        columns = pd.MultiIndex.from_tuples([
2736
            (2, 3, 2, False, False),
2737
            (2, 3, 2, False, True),
2738
            (2, 3, 2, True, False),
2739
            (2, 3, 2, True, True),
2740
            (2, 3, 3, False, False),
2741
            (2, 3, 3, False, True),
2742
            (2, 3, 3, True, False),
2743
            (2, 3, 3, True, True),
2744
            (2, 4, 2, False, False),
2745
            (2, 4, 2, False, True),
2746
            (2, 4, 2, True, False),
2747
            (2, 4, 2, True, True),
2748
            (2, 4, 3, False, False),
2749
            (2, 4, 3, False, True),
2750
            (2, 4, 3, True, False),
2751
            (2, 4, 3, True, True),
2752
            (3, 3, 2, False, False),
2753
            (3, 3, 2, False, True),
2754
            (3, 3, 2, True, False),
2755
            (3, 3, 2, True, True),
2756
            (3, 3, 3, False, False),
2757
            (3, 3, 3, False, True),
2758
            (3, 3, 3, True, False),
2759
            (3, 3, 3, True, True),
2760
            (3, 4, 2, False, False),
2761
            (3, 4, 2, False, True),
2762
            (3, 4, 2, True, False),
2763
            (3, 4, 2, True, True),
2764
            (3, 4, 3, False, False),
2765
            (3, 4, 3, False, True),
2766
            (3, 4, 3, True, False),
2767
            (3, 4, 3, True, True)
2768
        ], names=['macd_fast_window', 'macd_slow_window', 'macd_signal_window', 'macd_macd_ewm', 'macd_signal_ewm'])
2769

2770
        pd.testing.assert_frame_equal(
2771
            vbt.MACD.run(
2772
                close_ts,
2773
                fast_window=(2, 3),
2774
                slow_window=(3, 4),
2775
                signal_window=(2, 3),
2776
                macd_ewm=(False, True),
2777
                signal_ewm=(False, True),
2778
                param_product=True
2779
            ).macd,
2780
            pd.DataFrame(
2781
                np.array([
2782
                    [np.nan, np.nan, np.nan, np.nan, np.nan,
2783
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2784
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2785
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2786
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2787
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2788
                     np.nan, np.nan],
2789
                    [np.nan, np.nan, np.nan, np.nan, np.nan,
2790
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2791
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2792
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2793
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2794
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2795
                     np.nan, np.nan],
2796
                    [0.5, 0.5, 0.30555556, 0.30555556, 0.5,
2797
                     0.5, 0.30555556, 0.30555556, np.nan, np.nan,
2798
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2799
                     np.nan, 0., 0., 0., 0.,
2800
                     0., 0., 0., 0., np.nan,
2801
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2802
                     np.nan, np.nan],
2803
                    [0.5, 0.5, 0.39351852, 0.39351852, 0.5,
2804
                     0.5, 0.39351852, 0.39351852, 1., 1.,
2805
                     0.69451852, 0.69451852, 1., 1., 0.69451852,
2806
                     0.69451852, 0., 0., 0., 0.,
2807
                     0., 0., 0., 0., 0.5,
2808
                     0.5, 0.301, 0.301, 0.5, 0.5,
2809
                     0.301, 0.301],
2810
                    [0.16666667, 0.16666667, 0.11033951, 0.11033951, 0.16666667,
2811
                     0.16666667, 0.11033951, 0.11033951, 0.5, 0.5,
2812
                     0.27843951, 0.27843951, 0.5, 0.5, 0.27843951,
2813
                     0.27843951, 0., 0., 0., 0.,
2814
                     0., 0., 0., 0., 0.33333333,
2815
                     0.33333333, 0.1681, 0.1681, 0.33333333, 0.33333333,
2816
                     0.1681, 0.1681],
2817
                    [-0.5, -0.5, -0.1403035, -0.1403035, -0.5,
2818
                     -0.5, -0.1403035, -0.1403035, -0.5, -0.5,
2819
                     -0.1456935, -0.1456935, -0.5, -0.5, -0.1456935,
2820
                     -0.1456935, 0., 0., 0., 0.,
2821
                     0., 0., 0., 0., 0.,
2822
                     0., -0.00539, -0.00539, 0., 0.,
2823
                     -0.00539, -0.00539],
2824
                    [-0.5, -0.5, -0.30197617, -0.30197617, -0.5,
2825
                     -0.5, -0.30197617, -0.30197617, -1., -1.,
2826
                     -0.45833517, -0.45833517, -1., -1., -0.45833517,
2827
                     -0.45833517, 0., 0., 0., 0.,
2828
                     0., 0., 0., 0., -0.5,
2829
                     -0.5, -0.156359, -0.156359, -0.5, -0.5,
2830
                     -0.156359, -0.156359]
2831
                ]),
2832
                index=close_ts.index,
2833
                columns=columns
2834
            )
2835
        )
2836
        pd.testing.assert_frame_equal(
2837
            vbt.MACD.run(
2838
                close_ts,
2839
                fast_window=(2, 3),
2840
                slow_window=(3, 4),
2841
                signal_window=(2, 3),
2842
                macd_ewm=(False, True),
2843
                signal_ewm=(False, True),
2844
                param_product=True
2845
            ).signal,
2846
            pd.DataFrame(
2847
                np.array([
2848
                    [np.nan, np.nan, np.nan, np.nan, np.nan,
2849
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2850
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2851
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2852
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2853
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2854
                     np.nan, np.nan],
2855
                    [np.nan, np.nan, np.nan, np.nan, np.nan,
2856
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2857
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2858
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2859
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2860
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2861
                     np.nan, np.nan],
2862
                    [np.nan, np.nan, np.nan, np.nan, np.nan,
2863
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2864
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2865
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2866
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2867
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2868
                     np.nan, np.nan],
2869
                    [0.5, 0.5, 0.34953704, 0.36419753, np.nan,
2870
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2871
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2872
                     np.nan, 0., 0., 0., 0.,
2873
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2874
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2875
                     np.nan, np.nan],
2876
                    [0.33333333, 0.27777778, 0.25192901, 0.19495885, 0.38888889,
2877
                     0.33333333, 0.26980453, 0.22993827, 0.75, 0.66666667,
2878
                     0.48647901, 0.41713251, np.nan, np.nan, np.nan,
2879
                     np.nan, 0., 0., 0., 0.,
2880
                     0., 0., 0., 0., 0.41666667,
2881
                     0.38888889, 0.23455, 0.2124, np.nan, np.nan,
2882
                     np.nan, np.nan],
2883
                    [-0.16666667, -0.24074074, -0.014982, -0.02854938, 0.05555556,
2884
                     -0.08333333, 0.12118484, 0.04481739, 0., -0.11111111,
2885
                     0.066373, 0.04191517, 0.33333333, 0.125, 0.27575484,
2886
                     0.17039276, 0., 0., 0., 0.,
2887
                     0., 0., 0., 0., 0.16666667,
2888
                     0.12962963, 0.081355, 0.06720667, 0.27777778, 0.20833333,
2889
                     0.15457, 0.11458],
2890
                    [-0.5, -0.41358025, -0.22113983, -0.2108339, -0.27777778,
2891
                     -0.29166667, -0.11064672, -0.12857939, -0.75, -0.7037037,
2892
                     -0.30201433, -0.29158505, -0.33333333, -0.4375, -0.10852972,
2893
                     -0.1439712, 0., 0., 0., 0.,
2894
                     0., 0., 0., 0., -0.25,
2895
                     -0.29012346, -0.0808745, -0.08183711, -0.05555556, -0.14583333,
2896
                     0.002117, -0.0208895]
2897
                ]),
2898
                index=close_ts.index,
2899
                columns=columns
2900
            )
2901
        )
2902
        pd.testing.assert_frame_equal(
2903
            vbt.MACD.run(
2904
                close_ts,
2905
                fast_window=(2, 3),
2906
                slow_window=(3, 4),
2907
                signal_window=(2, 3),
2908
                macd_ewm=(False, True),
2909
                signal_ewm=(False, True),
2910
                param_product=True
2911
            ).hist,
2912
            pd.DataFrame(
2913
                np.array([
2914
                    [np.nan, np.nan, np.nan, np.nan, np.nan,
2915
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2916
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2917
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2918
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2919
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2920
                     np.nan, np.nan],
2921
                    [np.nan, np.nan, np.nan, np.nan, np.nan,
2922
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2923
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2924
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2925
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2926
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2927
                     np.nan, np.nan],
2928
                    [np.nan, np.nan, np.nan, np.nan, np.nan,
2929
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2930
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2931
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2932
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2933
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2934
                     np.nan, np.nan],
2935
                    [0., 0., 0.04398148, 0.02932099, np.nan,
2936
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2937
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2938
                     np.nan, 0., 0., 0., 0.,
2939
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2940
                     np.nan, np.nan, np.nan, np.nan, np.nan,
2941
                     np.nan, np.nan],
2942
                    [-0.16666667, -0.11111111, -0.14158951, -0.08461934, -0.22222222,
2943
                     -0.16666667, -0.15946502, -0.11959877, -0.25, -0.16666667,
2944
                     -0.20803951, -0.138693, np.nan, np.nan, np.nan,
2945
                     np.nan, 0., 0., 0., 0.,
2946
                     0., 0., 0., 0., -0.08333333,
2947
                     -0.05555556, -0.06645, -0.0443, np.nan, np.nan,
2948
                     np.nan, np.nan],
2949
                    [-0.33333333, -0.25925926, -0.1253215, -0.11175412, -0.55555556,
2950
                     -0.41666667, -0.26148834, -0.18512088, -0.5, -0.38888889,
2951
                     -0.2120665, -0.18760867, -0.83333333, -0.625, -0.42144834,
2952
                     -0.31608626, 0., 0., 0., 0.,
2953
                     0., 0., 0., 0., -0.16666667,
2954
                     -0.12962963, -0.086745, -0.07259667, -0.27777778, -0.20833333,
2955
                     -0.15996, -0.11997],
2956
                    [0., -0.08641975, -0.08083633, -0.09114226, -0.22222222,
2957
                     -0.20833333, -0.19132945, -0.17339678, -0.25, -0.2962963,
2958
                     -0.15632083, -0.16675011, -0.66666667, -0.5625, -0.34980545,
2959
                     -0.31436396, 0., 0., 0., 0.,
2960
                     0., 0., 0., 0., -0.25,
2961
                     -0.20987654, -0.0754845, -0.07452189, -0.44444444, -0.35416667,
2962
                     -0.158476, -0.1354695]
2963
                ]),
2964
                index=close_ts.index,
2965
                columns=columns
2966
            )
2967
        )
2968

2969
    def test_ATR(self):
2970
        columns = pd.MultiIndex.from_tuples([
2971
            (2, False),
2972
            (2, True),
2973
            (3, False),
2974
            (3, True)
2975
        ], names=['atr_window', 'atr_ewm'])
2976
        pd.testing.assert_frame_equal(
2977
            vbt.ATR.run(high_ts, low_ts, close_ts, window=(2, 3), ewm=(False, True), param_product=True).tr,
2978
            pd.DataFrame(
2979
                np.array([
2980
                    [0.2, 0.2, 0.2, 0.2],
2981
                    [1.2, 1.2, 1.2, 1.2],
2982
                    [1.3, 1.3, 1.3, 1.3],
2983
                    [1.4, 1.4, 1.4, 1.4],
2984
                    [1.3, 1.3, 1.3, 1.3],
2985
                    [1.2, 1.2, 1.2, 1.2],
2986
                    [1.1, 1.1, 1.1, 1.1]
2987
                ]),
2988
                index=close_ts.index,
2989
                columns=columns
2990
            )
2991
        )
2992
        pd.testing.assert_frame_equal(
2993
            vbt.ATR.run(high_ts, low_ts, close_ts, window=(2, 3), ewm=(False, True), param_product=True).atr,
2994
            pd.DataFrame(
2995
                np.array([
2996
                    [np.nan, np.nan, np.nan, np.nan],
2997
                    [0.7, 0.86666667, np.nan, np.nan],
2998
                    [1.25, 1.15555556, 0.9, 1.],
2999
                    [1.35, 1.31851852, 1.3, 1.2],
3000
                    [1.35, 1.30617284, 1.33333333, 1.25],
3001
                    [1.25, 1.23539095, 1.3, 1.225],
3002
                    [1.15, 1.14513032, 1.2, 1.1625]
3003
                ]),
3004
                index=close_ts.index,
3005
                columns=columns
3006
            )
3007
        )
3008

3009
    def test_OBV(self):
3010
        pd.testing.assert_series_equal(
3011
            vbt.OBV.run(close_ts, volume_ts).obv,
3012
            pd.Series(
3013
                np.array([4, 7, 9, 10, 8, 5, 1]),
3014
                index=close_ts.index,
3015
                name=close_ts.name
3016
            )
3017
        )
3018

3019
Product

Resources

Company
