CoCalc -- test_indicators.py

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

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

11
import vectorbt as vbt
12

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

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

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

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

40
seed = 42
41

42

43
# ############# Global ############# #
44

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

53

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

59

60
# ############# factory.py ############# #
61

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

74

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

2264
    def test_get_talib_indicators(self):
2265
        if talib_available:
2266
            assert len(vbt.IndicatorFactory.get_talib_indicators()) > 0
2267

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

2347
    def test_get_pandas_ta_indicators(self):
2348
        if pandas_ta_available:
2349
            assert len(vbt.IndicatorFactory.get_pandas_ta_indicators()) > 0
2350

2351
    def test_from_pandas_ta(self):
2352
        if pandas_ta_available:
2353
            pd.testing.assert_frame_equal(
2354
                vbt.pandas_ta('SMA').run(ts, 2).sma,
2355
                pd.DataFrame(
2356
                    ts.rolling(2).mean().values,
2357
                    index=ts.index,
2358
                    columns=pd.MultiIndex.from_tuples([(2, 'a'), (2, 'b'), (2, 'c')], names=['sma_length', None])
2359
                )
2360
            )
2361
            pd.testing.assert_frame_equal(
2362
                vbt.pandas_ta('SMA').run(ts['a'], [2, 3, 4]).sma,
2363
                pd.DataFrame(
2364
                    np.column_stack((
2365
                        ts['a'].rolling(2).mean().values,
2366
                        ts['a'].rolling(3).mean().values,
2367
                        ts['a'].rolling(4).mean().values
2368
                    )),
2369
                    index=ts.index,
2370
                    columns=pd.Index([2, 3, 4], dtype='int64', name='sma_length')
2371
                )
2372
            )
2373

2374
    def test_get_ta_indicators(self):
2375
        if ta_available:
2376
            assert len(vbt.IndicatorFactory.get_ta_indicators()) > 0
2377

2378
    def test_from_ta(self):
2379
        if ta_available:
2380
            pd.testing.assert_frame_equal(
2381
                vbt.ta('SMAIndicator').run(ts, 2).sma_indicator,
2382
                pd.DataFrame(
2383
                    ts.rolling(2).mean().values,
2384
                    index=ts.index,
2385
                    columns=pd.MultiIndex.from_tuples([(2, 'a'), (2, 'b'), (2, 'c')],
2386
                                                      names=['smaindicator_window', None])
2387
                )
2388
            )
2389
            pd.testing.assert_frame_equal(
2390
                vbt.ta('SMAIndicator').run(ts['a'], [2, 3, 4]).sma_indicator,
2391
                pd.DataFrame(
2392
                    np.column_stack((
2393
                        ts['a'].rolling(2).mean().values,
2394
                        ts['a'].rolling(3).mean().values,
2395
                        ts['a'].rolling(4).mean().values
2396
                    )),
2397
                    index=ts.index,
2398
                    columns=pd.Index([2, 3, 4], dtype='int64', name='smaindicator_window')
2399
                )
2400
            )
2401
            target = pd.DataFrame(
2402
                np.array([
2403
                    [np.nan, np.nan, np.nan],
2404
                    [2.5, 5.5, 2.5],
2405
                    [3.5, 4.5, 3.5],
2406
                    [4.5, 3.5, 3.5],
2407
                    [5.5, 2.5, 2.5]
2408
                ]),
2409
                index=ts.index,
2410
                columns=pd.MultiIndex.from_tuples([
2411
                    (2, 2, 'a'),
2412
                    (2, 2, 'b'),
2413
                    (2, 2, 'c')
2414
                ], names=['bollingerbands_window', 'bollingerbands_window_dev', None])
2415
            )
2416
            BollingerBands = vbt.ta('BollingerBands')
2417
            pd.testing.assert_frame_equal(
2418
                BollingerBands.run(ts, window=2, window_dev=2).bollinger_hband,
2419
                target
2420
            )
2421
            pd.testing.assert_frame_equal(
2422
                BollingerBands.run(ts, window=2, window_dev=2).bollinger_mavg,
2423
                target - 1
2424
            )
2425
            pd.testing.assert_frame_equal(
2426
                BollingerBands.run(ts, window=2, window_dev=2).bollinger_lband,
2427
                target - 2
2428
            )
2429

2430

2431
# ############# basic.py ############# #
2432

2433
close_ts = pd.Series([1, 2, 3, 4, 3, 2, 1], index=pd.DatetimeIndex([
2434
    datetime(2018, 1, 1),
2435
    datetime(2018, 1, 2),
2436
    datetime(2018, 1, 3),
2437
    datetime(2018, 1, 4),
2438
    datetime(2018, 1, 5),
2439
    datetime(2018, 1, 6),
2440
    datetime(2018, 1, 7)
2441
]))
2442
high_ts = close_ts * 1.1
2443
low_ts = close_ts * 0.9
2444
volume_ts = pd.Series([4, 3, 2, 1, 2, 3, 4], index=close_ts.index)
2445

2446

2447
class TestBasic:
2448
    def test_MA(self):
2449
        pd.testing.assert_frame_equal(
2450
            vbt.MA.run(close_ts, window=(2, 3), ewm=(False, True), param_product=True).ma,
2451
            pd.DataFrame(
2452
                np.array([
2453
                    [np.nan, np.nan, np.nan, np.nan],
2454
                    [1.5, 1.66666667, np.nan, np.nan],
2455
                    [2.5, 2.55555556, 2., 2.25],
2456
                    [3.5, 3.51851852, 3., 3.125],
2457
                    [3.5, 3.17283951, 3.33333333, 3.0625],
2458
                    [2.5, 2.3909465, 3., 2.53125],
2459
                    [1.5, 1.46364883, 2., 1.765625]
2460
                ]),
2461
                index=close_ts.index,
2462
                columns=pd.MultiIndex.from_tuples([
2463
                    (2, False),
2464
                    (2, True),
2465
                    (3, False),
2466
                    (3, True)
2467
                ], names=['ma_window', 'ma_ewm'])
2468
            )
2469
        )
2470

2471
    def test_MSTD(self):
2472
        pd.testing.assert_frame_equal(
2473
            vbt.MSTD.run(close_ts, window=(2, 3), ewm=(False, True), param_product=True).mstd,
2474
            pd.DataFrame(
2475
                np.array([
2476
                    [np.nan, np.nan, np.nan, np.nan],
2477
                    [0.5, 0.70710678, np.nan, np.nan],
2478
                    [0.5, 0.97467943, 0.81649658, 1.04880885],
2479
                    [0.5, 1.11434207, 0.81649658, 1.30018314],
2480
                    [0.5, 0.73001838, 0.47140452, 0.91715673],
2481
                    [0.5, 0.88824841, 0.81649658, 0.9182094],
2482
                    [0.5, 1.05965735, 0.81649658, 1.14049665]
2483
                ]),
2484
                index=close_ts.index,
2485
                columns=pd.MultiIndex.from_tuples([
2486
                    (2, False),
2487
                    (2, True),
2488
                    (3, False),
2489
                    (3, True)
2490
                ], names=['mstd_window', 'mstd_ewm'])
2491
            )
2492
        )
2493

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

2651
    def test_RSI(self):
2652
        pd.testing.assert_frame_equal(
2653
            vbt.RSI.run(close_ts, window=(2, 3), ewm=(False, True), param_product=True).rsi,
2654
            pd.DataFrame(
2655
                np.array([
2656
                    [np.nan, np.nan, np.nan, np.nan],
2657
                    [np.nan, np.nan, np.nan, np.nan],
2658
                    [100., 100., np.nan, np.nan],
2659
                    [100., 100., 100., 100.],
2660
                    [50., 33.33333333, 66.66666667, 50.],
2661
                    [0., 11.11111111, 33.33333333, 25.],
2662
                    [0., 3.7037037, 0., 12.5]
2663
                ]),
2664
                index=close_ts.index,
2665
                columns=pd.MultiIndex.from_tuples([
2666
                    (2, False),
2667
                    (2, True),
2668
                    (3, False),
2669
                    (3, True)
2670
                ], names=['rsi_window', 'rsi_ewm'])
2671
            )
2672
        )
2673

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

2748
    def test_MACD(self):
2749
        columns = pd.MultiIndex.from_tuples([
2750
            (2, 3, 2, False, False),
2751
            (2, 3, 2, False, True),
2752
            (2, 3, 2, True, False),
2753
            (2, 3, 2, True, True),
2754
            (2, 3, 3, False, False),
2755
            (2, 3, 3, False, True),
2756
            (2, 3, 3, True, False),
2757
            (2, 3, 3, True, True),
2758
            (2, 4, 2, False, False),
2759
            (2, 4, 2, False, True),
2760
            (2, 4, 2, True, False),
2761
            (2, 4, 2, True, True),
2762
            (2, 4, 3, False, False),
2763
            (2, 4, 3, False, True),
2764
            (2, 4, 3, True, False),
2765
            (2, 4, 3, True, True),
2766
            (3, 3, 2, False, False),
2767
            (3, 3, 2, False, True),
2768
            (3, 3, 2, True, False),
2769
            (3, 3, 2, True, True),
2770
            (3, 3, 3, False, False),
2771
            (3, 3, 3, False, True),
2772
            (3, 3, 3, True, False),
2773
            (3, 3, 3, True, True),
2774
            (3, 4, 2, False, False),
2775
            (3, 4, 2, False, True),
2776
            (3, 4, 2, True, False),
2777
            (3, 4, 2, True, True),
2778
            (3, 4, 3, False, False),
2779
            (3, 4, 3, False, True),
2780
            (3, 4, 3, True, False),
2781
            (3, 4, 3, True, True)
2782
        ], names=['macd_fast_window', 'macd_slow_window', 'macd_signal_window', 'macd_macd_ewm', 'macd_signal_ewm'])
2783

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

2983
    def test_ATR(self):
2984
        columns = pd.MultiIndex.from_tuples([
2985
            (2, False),
2986
            (2, True),
2987
            (3, False),
2988
            (3, True)
2989
        ], names=['atr_window', 'atr_ewm'])
2990
        pd.testing.assert_frame_equal(
2991
            vbt.ATR.run(high_ts, low_ts, close_ts, window=(2, 3), ewm=(False, True), param_product=True).tr,
2992
            pd.DataFrame(
2993
                np.array([
2994
                    [0.2, 0.2, 0.2, 0.2],
2995
                    [1.2, 1.2, 1.2, 1.2],
2996
                    [1.3, 1.3, 1.3, 1.3],
2997
                    [1.4, 1.4, 1.4, 1.4],
2998
                    [1.3, 1.3, 1.3, 1.3],
2999
                    [1.2, 1.2, 1.2, 1.2],
3000
                    [1.1, 1.1, 1.1, 1.1]
3001
                ]),
3002
                index=close_ts.index,
3003
                columns=columns
3004
            )
3005
        )
3006
        pd.testing.assert_frame_equal(
3007
            vbt.ATR.run(high_ts, low_ts, close_ts, window=(2, 3), ewm=(False, True), param_product=True).atr,
3008
            pd.DataFrame(
3009
                np.array([
3010
                    [np.nan, np.nan, np.nan, np.nan],
3011
                    [0.7, 0.86666667, np.nan, np.nan],
3012
                    [1.25, 1.15555556, 0.9, 1.],
3013
                    [1.35, 1.31851852, 1.3, 1.2],
3014
                    [1.35, 1.30617284, 1.33333333, 1.25],
3015
                    [1.25, 1.23539095, 1.3, 1.225],
3016
                    [1.15, 1.14513032, 1.2, 1.1625]
3017
                ]),
3018
                index=close_ts.index,
3019
                columns=columns
3020
            )
3021
        )
3022

3023
    def test_OBV(self):
3024
        pd.testing.assert_series_equal(
3025
            vbt.OBV.run(close_ts, volume_ts).obv,
3026
            pd.Series(
3027
                np.array([4, 7, 9, 10, 8, 5, 1]),
3028
                index=close_ts.index,
3029
                name=close_ts.name
3030
            )
3031
        )
3032

3033
Product

Resources

Company
