CoCalc -- test_signals.py

GitHub Repository: polakowo/vectorbt
Path: blob/master/tests/test_signals.py
¹⁰⁷¹ views
1
from datetime import datetime
2

3
import numpy as np
4
import pandas as pd
5
import pytest
6
from numba import njit
7

8
import vectorbt as vbt
9
from tests.utils import record_arrays_close
10
from vectorbt.generic import nb as generic_nb
11
from vectorbt.generic.enums import range_dt
12

13
seed = 42
14

15
day_dt = np.timedelta64(86400000000000)
16

17
mask = pd.DataFrame([
18
    [True, False, False],
19
    [False, True, False],
20
    [False, False, True],
21
    [True, False, False],
22
    [False, True, False]
23
], index=pd.Index([
24
    datetime(2020, 1, 1),
25
    datetime(2020, 1, 2),
26
    datetime(2020, 1, 3),
27
    datetime(2020, 1, 4),
28
    datetime(2020, 1, 5)
29
]), columns=['a', 'b', 'c'])
30

31
ts = pd.Series([1., 2., 3., 2., 1.], index=mask.index)
32

33
price = pd.DataFrame({
34
    'open': [10, 11, 12, 11, 10],
35
    'high': [11, 12, 13, 12, 11],
36
    'low': [9, 10, 11, 10, 9],
37
    'close': [11, 12, 11, 10, 9]
38
})
39

40
group_by = pd.Index(['g1', 'g1', 'g2'])
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

51

52
def teardown_module():
53
    vbt.settings.reset()
54

55

56
# ############# accessors.py ############# #
57

58

59
class TestAccessors:
60
    def test_indexing(self):
61
        assert mask.vbt.signals['a'].total() == mask['a'].vbt.signals.total()
62

63
    def test_freq(self):
64
        assert mask.vbt.signals.wrapper.freq == day_dt
65
        assert mask['a'].vbt.signals.wrapper.freq == day_dt
66
        assert mask.vbt.signals(freq='2D').wrapper.freq == day_dt * 2
67
        assert mask['a'].vbt.signals(freq='2D').wrapper.freq == day_dt * 2
68
        assert pd.Series([False, True]).vbt.signals.wrapper.freq is None
69
        assert pd.Series([False, True]).vbt.signals(freq='3D').wrapper.freq == day_dt * 3
70
        assert pd.Series([False, True]).vbt.signals(freq=np.timedelta64(4, 'D')).wrapper.freq == day_dt * 4
71

72
    @pytest.mark.parametrize(
73
        "test_n",
74
        [1, 2, 3, 4, 5],
75
    )
76
    def test_fshift(self, test_n):
77
        pd.testing.assert_series_equal(mask['a'].vbt.signals.fshift(test_n), mask['a'].shift(test_n, fill_value=False))
78
        np.testing.assert_array_equal(
79
            mask['a'].vbt.signals.fshift(test_n).values,
80
            generic_nb.fshift_1d_nb(mask['a'].values, test_n, fill_value=False)
81
        )
82
        pd.testing.assert_frame_equal(mask.vbt.signals.fshift(test_n), mask.shift(test_n, fill_value=False))
83

84
    @pytest.mark.parametrize(
85
        "test_n",
86
        [1, 2, 3, 4, 5],
87
    )
88
    def test_bshift(self, test_n):
89
        pd.testing.assert_series_equal(
90
            mask['a'].vbt.signals.bshift(test_n),
91
            mask['a'].shift(-test_n, fill_value=False))
92
        np.testing.assert_array_equal(
93
            mask['a'].vbt.signals.bshift(test_n).values,
94
            generic_nb.bshift_1d_nb(mask['a'].values, test_n, fill_value=False)
95
        )
96
        pd.testing.assert_frame_equal(mask.vbt.signals.bshift(test_n), mask.shift(-test_n, fill_value=False))
97

98
    def test_empty(self):
99
        pd.testing.assert_series_equal(
100
            pd.Series.vbt.signals.empty(5, index=np.arange(10, 15), name='a'),
101
            pd.Series(np.full(5, False), index=np.arange(10, 15), name='a')
102
        )
103
        pd.testing.assert_frame_equal(
104
            pd.DataFrame.vbt.signals.empty((5, 3), index=np.arange(10, 15), columns=['a', 'b', 'c']),
105
            pd.DataFrame(np.full((5, 3), False), index=np.arange(10, 15), columns=['a', 'b', 'c'])
106
        )
107
        pd.testing.assert_series_equal(
108
            pd.Series.vbt.signals.empty_like(mask['a']),
109
            pd.Series(np.full(mask['a'].shape, False), index=mask['a'].index, name=mask['a'].name)
110
        )
111
        pd.testing.assert_frame_equal(
112
            pd.DataFrame.vbt.signals.empty_like(mask),
113
            pd.DataFrame(np.full(mask.shape, False), index=mask.index, columns=mask.columns)
114
        )
115

116
    def test_generate(self):
117
        @njit
118
        def choice_func_nb(from_i, to_i, col, n):
119
            if col == 0:
120
                return np.arange(from_i, to_i)
121
            elif col == 1:
122
                return np.full(1, from_i)
123
            else:
124
                return np.full(1, to_i - n)
125

126
        pd.testing.assert_series_equal(
127
            pd.Series.vbt.signals.generate(5, choice_func_nb, 1, index=mask['a'].index, name=mask['a'].name),
128
            pd.Series(
129
                np.array([True, True, True, True, True]),
130
                index=mask['a'].index,
131
                name=mask['a'].name
132
            )
133
        )
134
        with pytest.raises(Exception):
135
            _ = pd.Series.vbt.signals.generate((5, 2), choice_func_nb, 1)
136
        pd.testing.assert_frame_equal(
137
            pd.DataFrame.vbt.signals.generate(
138
                (5, 3), choice_func_nb, 1, index=mask.index, columns=mask.columns),
139
            pd.DataFrame(
140
                np.array([
141
                    [True, True, False],
142
                    [True, False, False],
143
                    [True, False, False],
144
                    [True, False, False],
145
                    [True, False, True]
146
                ]),
147
                index=mask.index,
148
                columns=mask.columns
149
            )
150
        )
151
        pd.testing.assert_frame_equal(
152
            pd.DataFrame.vbt.signals.generate(
153
                (5, 3), choice_func_nb, 1, pick_first=True, index=mask.index, columns=mask.columns),
154
            pd.DataFrame(
155
                np.array([
156
                    [True, True, False],
157
                    [False, False, False],
158
                    [False, False, False],
159
                    [False, False, False],
160
                    [False, False, True]
161
                ]),
162
                index=mask.index,
163
                columns=mask.columns
164
            )
165
        )
166

167
    def test_generate_both(self):
168
        @njit
169
        def entry_func_nb(from_i, to_i, col, temp_int):
170
            temp_int[0] = from_i
171
            return temp_int[:1]
172

173
        @njit
174
        def exit_func_nb(from_i, to_i, col, temp_int):
175
            temp_int[0] = from_i
176
            return temp_int[:1]
177

178
        temp_int = np.empty((mask.shape[0],), dtype=np.int64)
179

180
        en, ex = pd.Series.vbt.signals.generate_both(
181
            5, entry_func_nb, (temp_int,), exit_func_nb, (temp_int,),
182
            index=mask['a'].index, name=mask['a'].name)
183
        pd.testing.assert_series_equal(
184
            en,
185
            pd.Series(
186
                np.array([True, False, True, False, True]),
187
                index=mask['a'].index,
188
                name=mask['a'].name
189
            )
190
        )
191
        pd.testing.assert_series_equal(
192
            ex,
193
            pd.Series(
194
                np.array([False, True, False, True, False]),
195
                index=mask['a'].index,
196
                name=mask['a'].name
197
            )
198
        )
199
        en, ex = pd.DataFrame.vbt.signals.generate_both(
200
            (5, 3), entry_func_nb, (temp_int,), exit_func_nb, (temp_int,),
201
            index=mask.index, columns=mask.columns)
202
        pd.testing.assert_frame_equal(
203
            en,
204
            pd.DataFrame(
205
                np.array([
206
                    [True, True, True],
207
                    [False, False, False],
208
                    [True, True, True],
209
                    [False, False, False],
210
                    [True, True, True]
211
                ]),
212
                index=mask.index,
213
                columns=mask.columns
214
            )
215
        )
216
        pd.testing.assert_frame_equal(
217
            ex,
218
            pd.DataFrame(
219
                np.array([
220
                    [False, False, False],
221
                    [True, True, True],
222
                    [False, False, False],
223
                    [True, True, True],
224
                    [False, False, False]
225
                ]),
226
                index=mask.index,
227
                columns=mask.columns
228
            )
229
        )
230
        en, ex = pd.Series.vbt.signals.generate_both(
231
            (5,), entry_func_nb, (temp_int,), exit_func_nb, (temp_int,),
232
            index=mask['a'].index, name=mask['a'].name, entry_wait=1, exit_wait=0)
233
        pd.testing.assert_series_equal(
234
            en,
235
            pd.Series(
236
                np.array([True, True, True, True, True]),
237
                index=mask['a'].index,
238
                name=mask['a'].name
239
            )
240
        )
241
        pd.testing.assert_series_equal(
242
            ex,
243
            pd.Series(
244
                np.array([True, True, True, True, True]),
245
                index=mask['a'].index,
246
                name=mask['a'].name
247
            )
248
        )
249
        en, ex = pd.Series.vbt.signals.generate_both(
250
            (5,), entry_func_nb, (temp_int,), exit_func_nb, (temp_int,),
251
            index=mask['a'].index, name=mask['a'].name, entry_wait=0, exit_wait=1)
252
        pd.testing.assert_series_equal(
253
            en,
254
            pd.Series(
255
                np.array([True, True, True, True, True]),
256
                index=mask['a'].index,
257
                name=mask['a'].name
258
            )
259
        )
260
        pd.testing.assert_series_equal(
261
            ex,
262
            pd.Series(
263
                np.array([False, True, True, True, True]),
264
                index=mask['a'].index,
265
                name=mask['a'].name
266
            )
267
        )
268

269
        @njit
270
        def entry_func2_nb(from_i, to_i, col, temp_int):
271
            temp_int[0] = from_i
272
            if from_i + 1 < to_i:
273
                temp_int[1] = from_i + 1
274
                return temp_int[:2]
275
            return temp_int[:1]
276

277
        @njit
278
        def exit_func2_nb(from_i, to_i, col, temp_int):
279
            temp_int[0] = from_i
280
            if from_i + 1 < to_i:
281
                temp_int[1] = from_i + 1
282
                return temp_int[:2]
283
            return temp_int[:1]
284

285
        en, ex = pd.DataFrame.vbt.signals.generate_both(
286
            (5, 3), entry_func2_nb, (temp_int,), exit_func2_nb, (temp_int,),
287
            entry_pick_first=False, exit_pick_first=False,
288
            index=mask.index, columns=mask.columns)
289
        pd.testing.assert_frame_equal(
290
            en,
291
            pd.DataFrame(
292
                np.array([
293
                    [True, True, True],
294
                    [True, True, True],
295
                    [False, False, False],
296
                    [False, False, False],
297
                    [True, True, True]
298
                ]),
299
                index=mask.index,
300
                columns=mask.columns
301
            )
302
        )
303
        pd.testing.assert_frame_equal(
304
            ex,
305
            pd.DataFrame(
306
                np.array([
307
                    [False, False, False],
308
                    [False, False, False],
309
                    [True, True, True],
310
                    [True, True, True],
311
                    [False, False, False]
312
                ]),
313
                index=mask.index,
314
                columns=mask.columns
315
            )
316
        )
317

318
    def test_generate_exits(self):
319
        @njit
320
        def choice_func_nb(from_i, to_i, col, temp_int):
321
            temp_int[0] = from_i
322
            return temp_int[:1]
323

324
        temp_int = np.empty((mask.shape[0],), dtype=np.int64)
325

326
        pd.testing.assert_series_equal(
327
            mask['a'].vbt.signals.generate_exits(choice_func_nb, temp_int),
328
            pd.Series(
329
                np.array([False, True, False, False, True]),
330
                index=mask['a'].index,
331
                name=mask['a'].name
332
            )
333
        )
334
        pd.testing.assert_frame_equal(
335
            mask.vbt.signals.generate_exits(choice_func_nb, temp_int),
336
            pd.DataFrame(
337
                np.array([
338
                    [False, False, False],
339
                    [True, False, False],
340
                    [False, True, False],
341
                    [False, False, True],
342
                    [True, False, False]
343
                ]),
344
                index=mask.index,
345
                columns=mask.columns
346
            )
347
        )
348
        pd.testing.assert_frame_equal(
349
            mask.vbt.signals.generate_exits(choice_func_nb, temp_int, wait=0),
350
            pd.DataFrame(
351
                np.array([
352
                    [True, False, False],
353
                    [False, True, False],
354
                    [False, False, True],
355
                    [True, False, False],
356
                    [False, True, False]
357
                ]),
358
                index=mask.index,
359
                columns=mask.columns
360
            )
361
        )
362

363
        @njit
364
        def choice_func2_nb(from_i, to_i, col, temp_int):
365
            for i in range(from_i, to_i):
366
                temp_int[i - from_i] = i
367
            return temp_int[:to_i - from_i]
368

369
        pd.testing.assert_frame_equal(
370
            mask.vbt.signals.generate_exits(choice_func2_nb, temp_int, until_next=False, pick_first=False),
371
            pd.DataFrame(
372
                np.array([
373
                    [False, False, False],
374
                    [True, False, False],
375
                    [True, True, False],
376
                    [True, True, True],
377
                    [True, True, True]
378
                ]),
379
                index=mask.index,
380
                columns=mask.columns
381
            )
382
        )
383

384
        mask2 = pd.Series([True, True, True, True, True], index=mask.index)
385
        pd.testing.assert_series_equal(
386
            mask2.vbt.signals.generate_exits(choice_func_nb, temp_int, until_next=False, skip_until_exit=True),
387
            pd.Series(
388
                np.array([False, True, False, True, False]),
389
                index=mask.index
390
            )
391
        )
392

393
    def test_clean(self):
394
        entries = pd.DataFrame([
395
            [True, False, True],
396
            [True, False, False],
397
            [True, True, True],
398
            [False, True, False],
399
            [False, True, True]
400
        ], index=mask.index, columns=mask.columns)
401
        exits = pd.Series([True, False, True, False, True], index=mask.index)
402
        pd.testing.assert_frame_equal(
403
            entries.vbt.signals.clean(),
404
            pd.DataFrame(
405
                np.array([
406
                    [True, False, True],
407
                    [False, False, False],
408
                    [False, True, True],
409
                    [False, False, False],
410
                    [False, False, True]
411
                ]),
412
                index=mask.index,
413
                columns=mask.columns
414
            )
415
        )
416
        pd.testing.assert_frame_equal(
417
            pd.DataFrame.vbt.signals.clean(entries),
418
            pd.DataFrame(
419
                np.array([
420
                    [True, False, True],
421
                    [False, False, False],
422
                    [False, True, True],
423
                    [False, False, False],
424
                    [False, False, True]
425
                ]),
426
                index=mask.index,
427
                columns=mask.columns
428
            )
429
        )
430
        pd.testing.assert_frame_equal(
431
            entries.vbt.signals.clean(exits)[0],
432
            pd.DataFrame(
433
                np.array([
434
                    [False, False, False],
435
                    [True, False, False],
436
                    [False, False, False],
437
                    [False, True, False],
438
                    [False, False, False]
439
                ]),
440
                index=mask.index,
441
                columns=mask.columns
442
            )
443
        )
444
        pd.testing.assert_frame_equal(
445
            entries.vbt.signals.clean(exits)[1],
446
            pd.DataFrame(
447
                np.array([
448
                    [False, False, False],
449
                    [False, False, False],
450
                    [False, False, False],
451
                    [False, False, False],
452
                    [True, False, False]
453
                ]),
454
                index=mask.index,
455
                columns=mask.columns
456
            )
457
        )
458
        pd.testing.assert_frame_equal(
459
            entries.vbt.signals.clean(exits, entry_first=False)[0],
460
            pd.DataFrame(
461
                np.array([
462
                    [False, False, False],
463
                    [True, False, False],
464
                    [False, False, False],
465
                    [False, True, False],
466
                    [False, False, False]
467
                ]),
468
                index=mask.index,
469
                columns=mask.columns
470
            )
471
        )
472
        pd.testing.assert_frame_equal(
473
            entries.vbt.signals.clean(exits, entry_first=False)[1],
474
            pd.DataFrame(
475
                np.array([
476
                    [False, True, False],
477
                    [False, False, False],
478
                    [False, False, False],
479
                    [False, False, False],
480
                    [True, False, False]
481
                ]),
482
                index=mask.index,
483
                columns=mask.columns
484
            )
485
        )
486
        pd.testing.assert_frame_equal(
487
            pd.DataFrame.vbt.signals.clean(entries, exits)[0],
488
            pd.DataFrame(
489
                np.array([
490
                    [False, False, False],
491
                    [True, False, False],
492
                    [False, False, False],
493
                    [False, True, False],
494
                    [False, False, False]
495
                ]),
496
                index=mask.index,
497
                columns=mask.columns
498
            )
499
        )
500
        pd.testing.assert_frame_equal(
501
            pd.DataFrame.vbt.signals.clean(entries, exits)[1],
502
            pd.DataFrame(
503
                np.array([
504
                    [False, False, False],
505
                    [False, False, False],
506
                    [False, False, False],
507
                    [False, False, False],
508
                    [True, False, False]
509
                ]),
510
                index=mask.index,
511
                columns=mask.columns
512
            )
513
        )
514
        with pytest.raises(Exception):
515
            _ = pd.Series.vbt.signals.clean(entries, entries, entries)
516

517
    def test_generate_random(self):
518
        pd.testing.assert_series_equal(
519
            pd.Series.vbt.signals.generate_random(
520
                5, n=3, seed=seed, index=mask['a'].index, name=mask['a'].name),
521
            pd.Series(
522
                np.array([False, True, True, False, True]),
523
                index=mask['a'].index,
524
                name=mask['a'].name
525
            )
526
        )
527
        with pytest.raises(Exception):
528
            _ = pd.Series.vbt.signals.generate_random((5, 2), n=3)
529
        pd.testing.assert_frame_equal(
530
            pd.DataFrame.vbt.signals.generate_random(
531
                (5, 3), n=3, seed=seed, index=mask.index, columns=mask.columns),
532
            pd.DataFrame(
533
                np.array([
534
                    [False, False, True],
535
                    [True, True, True],
536
                    [True, True, False],
537
                    [False, True, True],
538
                    [True, False, False]
539
                ]),
540
                index=mask.index,
541
                columns=mask.columns
542
            )
543
        )
544
        pd.testing.assert_frame_equal(
545
            pd.DataFrame.vbt.signals.generate_random(
546
                (5, 3), n=[0, 1, 2], seed=seed, index=mask.index, columns=mask.columns),
547
            pd.DataFrame(
548
                np.array([
549
                    [False, False, True],
550
                    [False, False, True],
551
                    [False, False, False],
552
                    [False, True, False],
553
                    [False, False, False]
554
                ]),
555
                index=mask.index,
556
                columns=mask.columns
557
            )
558
        )
559
        pd.testing.assert_series_equal(
560
            pd.Series.vbt.signals.generate_random(
561
                5, prob=0.5, seed=seed, index=mask['a'].index, name=mask['a'].name),
562
            pd.Series(
563
                np.array([True, False, False, False, True]),
564
                index=mask['a'].index,
565
                name=mask['a'].name
566
            )
567
        )
568
        with pytest.raises(Exception):
569
            _ = pd.Series.vbt.signals.generate_random((5, 2), prob=3)
570
        pd.testing.assert_frame_equal(
571
            pd.DataFrame.vbt.signals.generate_random(
572
                (5, 3), prob=0.5, seed=seed, index=mask.index, columns=mask.columns),
573
            pd.DataFrame(
574
                np.array([
575
                    [True, True, True],
576
                    [False, True, False],
577
                    [False, False, False],
578
                    [False, False, True],
579
                    [True, False, True]
580
                ]),
581
                index=mask.index,
582
                columns=mask.columns
583
            )
584
        )
585
        pd.testing.assert_frame_equal(
586
            pd.DataFrame.vbt.signals.generate_random(
587
                (5, 3), prob=[0., 0.5, 1.], seed=seed, index=mask.index, columns=mask.columns),
588
            pd.DataFrame(
589
                np.array([
590
                    [False, True, True],
591
                    [False, True, True],
592
                    [False, False, True],
593
                    [False, False, True],
594
                    [False, False, True]
595
                ]),
596
                index=mask.index,
597
                columns=mask.columns
598
            )
599
        )
600
        with pytest.raises(Exception):
601
            pd.DataFrame.vbt.signals.generate_random((5, 3))
602
        pd.testing.assert_frame_equal(
603
            pd.DataFrame.vbt.signals.generate_random(
604
                (5, 3), prob=[0., 0.5, 1.], pick_first=True, seed=seed, index=mask.index, columns=mask.columns),
605
            pd.DataFrame(
606
                np.array([
607
                    [False, True, True],
608
                    [False, False, False],
609
                    [False, False, False],
610
                    [False, False, False],
611
                    [False, False, False]
612
                ]),
613
                index=mask.index,
614
                columns=mask.columns
615
            )
616
        )
617

618
    def test_generate_random_both(self):
619
        # n
620
        en, ex = pd.Series.vbt.signals.generate_random_both(
621
            5, n=2, seed=seed, index=mask['a'].index, name=mask['a'].name)
622
        pd.testing.assert_series_equal(
623
            en,
624
            pd.Series(
625
                np.array([True, False, True, False, False]),
626
                index=mask['a'].index,
627
                name=mask['a'].name
628
            )
629
        )
630
        pd.testing.assert_series_equal(
631
            ex,
632
            pd.Series(
633
                np.array([False, True, False, False, True]),
634
                index=mask['a'].index,
635
                name=mask['a'].name
636
            )
637
        )
638
        en, ex = pd.DataFrame.vbt.signals.generate_random_both(
639
            (5, 3), n=2, seed=seed, index=mask.index, columns=mask.columns)
640
        pd.testing.assert_frame_equal(
641
            en,
642
            pd.DataFrame(
643
                np.array([
644
                    [True, True, True],
645
                    [False, False, False],
646
                    [True, True, False],
647
                    [False, False, True],
648
                    [False, False, False]
649
                ]),
650
                index=mask.index,
651
                columns=mask.columns
652
            )
653
        )
654
        pd.testing.assert_frame_equal(
655
            ex,
656
            pd.DataFrame(
657
                np.array([
658
                    [False, False, False],
659
                    [True, True, True],
660
                    [False, False, False],
661
                    [False, True, False],
662
                    [True, False, True]
663
                ]),
664
                index=mask.index,
665
                columns=mask.columns
666
            )
667
        )
668
        en, ex = pd.DataFrame.vbt.signals.generate_random_both(
669
            (5, 3), n=[0, 1, 2], seed=seed, index=mask.index, columns=mask.columns)
670
        pd.testing.assert_frame_equal(
671
            en,
672
            pd.DataFrame(
673
                np.array([
674
                    [False, False, True],
675
                    [False, True, False],
676
                    [False, False, False],
677
                    [False, False, True],
678
                    [False, False, False]
679
                ]),
680
                index=mask.index,
681
                columns=mask.columns
682
            )
683
        )
684
        pd.testing.assert_frame_equal(
685
            ex,
686
            pd.DataFrame(
687
                np.array([
688
                    [False, False, False],
689
                    [False, False, True],
690
                    [False, False, False],
691
                    [False, True, False],
692
                    [False, False, True]
693
                ]),
694
                index=mask.index,
695
                columns=mask.columns
696
            )
697
        )
698
        en, ex = pd.DataFrame.vbt.signals.generate_random_both((2, 3), n=2, seed=seed, entry_wait=1, exit_wait=0)
699
        pd.testing.assert_frame_equal(
700
            en,
701
            pd.DataFrame(
702
                np.array([
703
                    [True, True, True],
704
                    [True, True, True],
705
                ])
706
            )
707
        )
708
        pd.testing.assert_frame_equal(
709
            ex,
710
            pd.DataFrame(
711
                np.array([
712
                    [True, True, True],
713
                    [True, True, True]
714
                ])
715
            )
716
        )
717
        en, ex = pd.DataFrame.vbt.signals.generate_random_both((3, 3), n=2, seed=seed, entry_wait=0, exit_wait=1)
718
        pd.testing.assert_frame_equal(
719
            en,
720
            pd.DataFrame(
721
                np.array([
722
                    [True, True, True],
723
                    [True, True, True],
724
                    [False, False, False]
725
                ])
726
            )
727
        )
728
        pd.testing.assert_frame_equal(
729
            ex,
730
            pd.DataFrame(
731
                np.array([
732
                    [False, False, False],
733
                    [True, True, True],
734
                    [True, True, True],
735
                ])
736
            )
737
        )
738
        en, ex = pd.DataFrame.vbt.signals.generate_random_both((7, 3), n=2, seed=seed, entry_wait=2, exit_wait=2)
739
        pd.testing.assert_frame_equal(
740
            en,
741
            pd.DataFrame(
742
                np.array([
743
                    [True, True, True],
744
                    [False, False, False],
745
                    [False, False, False],
746
                    [False, False, False],
747
                    [True, True, True],
748
                    [False, False, False],
749
                    [False, False, False]
750
                ])
751
            )
752
        )
753
        pd.testing.assert_frame_equal(
754
            ex,
755
            pd.DataFrame(
756
                np.array([
757
                    [False, False, False],
758
                    [False, False, False],
759
                    [True, True, True],
760
                    [False, False, False],
761
                    [False, False, False],
762
                    [False, False, False],
763
                    [True, True, True]
764
                ])
765
            )
766
        )
767
        n = 10
768
        a = np.full(n * 2, 0.)
769
        for i in range(10000):
770
            en, ex = pd.Series.vbt.signals.generate_random_both(1000, n, entry_wait=2, exit_wait=2)
771
            _a = np.empty((n * 2,), dtype=np.int64)
772
            _a[0::2] = np.flatnonzero(en)
773
            _a[1::2] = np.flatnonzero(ex)
774
            a += _a
775
        greater = a > 10000000 / (2 * n + 1) * np.arange(0, 2 * n)
776
        less = a < 10000000 / (2 * n + 1) * np.arange(2, 2 * n + 2)
777
        assert np.all(greater & less)
778

779
        # probs
780
        en, ex = pd.Series.vbt.signals.generate_random_both(
781
            5, entry_prob=0.5, exit_prob=1., seed=seed, index=mask['a'].index, name=mask['a'].name)
782
        pd.testing.assert_series_equal(
783
            en,
784
            pd.Series(
785
                np.array([True, False, False, False, True]),
786
                index=mask['a'].index,
787
                name=mask['a'].name
788
            )
789
        )
790
        pd.testing.assert_series_equal(
791
            ex,
792
            pd.Series(
793
                np.array([False, True, False, False, False]),
794
                index=mask['a'].index,
795
                name=mask['a'].name
796
            )
797
        )
798
        en, ex = pd.DataFrame.vbt.signals.generate_random_both(
799
            (5, 3), entry_prob=0.5, exit_prob=1., seed=seed, index=mask.index, columns=mask.columns)
800
        pd.testing.assert_frame_equal(
801
            en,
802
            pd.DataFrame(
803
                np.array([
804
                    [True, True, True],
805
                    [False, False, False],
806
                    [False, False, False],
807
                    [False, False, True],
808
                    [True, False, False]
809
                ]),
810
                index=mask.index,
811
                columns=mask.columns
812
            )
813
        )
814
        pd.testing.assert_frame_equal(
815
            ex,
816
            pd.DataFrame(
817
                np.array([
818
                    [False, False, False],
819
                    [True, True, True],
820
                    [False, False, False],
821
                    [False, False, False],
822
                    [False, False, True]
823
                ]),
824
                index=mask.index,
825
                columns=mask.columns
826
            )
827
        )
828
        en, ex = pd.DataFrame.vbt.signals.generate_random_both(
829
            (5, 3), entry_prob=[0., 0.5, 1.], exit_prob=[0., 0.5, 1.],
830
            seed=seed, index=mask.index, columns=mask.columns)
831
        pd.testing.assert_frame_equal(
832
            en,
833
            pd.DataFrame(
834
                np.array([
835
                    [False, True, True],
836
                    [False, False, False],
837
                    [False, False, True],
838
                    [False, False, False],
839
                    [False, False, True]
840
                ]),
841
                index=mask.index,
842
                columns=mask.columns
843
            )
844
        )
845
        pd.testing.assert_frame_equal(
846
            ex,
847
            pd.DataFrame(
848
                np.array([
849
                    [False, False, False],
850
                    [False, True, True],
851
                    [False, False, False],
852
                    [False, False, True],
853
                    [False, False, False]
854
                ]),
855
                index=mask.index,
856
                columns=mask.columns
857
            )
858
        )
859
        en, ex = pd.DataFrame.vbt.signals.generate_random_both(
860
            (5, 3), entry_prob=1., exit_prob=1., exit_wait=0,
861
            seed=seed, index=mask.index, columns=mask.columns)
862
        pd.testing.assert_frame_equal(
863
            en,
864
            pd.DataFrame(
865
                np.array([
866
                    [True, True, True],
867
                    [True, True, True],
868
                    [True, True, True],
869
                    [True, True, True],
870
                    [True, True, True]
871
                ]),
872
                index=mask.index,
873
                columns=mask.columns
874
            )
875
        )
876
        pd.testing.assert_frame_equal(
877
            ex,
878
            pd.DataFrame(
879
                np.array([
880
                    [True, True, True],
881
                    [True, True, True],
882
                    [True, True, True],
883
                    [True, True, True],
884
                    [True, True, True]
885
                ]),
886
                index=mask.index,
887
                columns=mask.columns
888
            )
889
        )
890
        en, ex = pd.DataFrame.vbt.signals.generate_random_both(
891
            (5, 3), entry_prob=1., exit_prob=1., entry_pick_first=False, exit_pick_first=True,
892
            seed=seed, index=mask.index, columns=mask.columns)
893
        pd.testing.assert_frame_equal(
894
            en,
895
            pd.DataFrame(
896
                np.array([
897
                    [True, True, True],
898
                    [True, True, True],
899
                    [True, True, True],
900
                    [True, True, True],
901
                    [True, True, True]
902
                ]),
903
                index=mask.index,
904
                columns=mask.columns
905
            )
906
        )
907
        pd.testing.assert_frame_equal(
908
            ex,
909
            pd.DataFrame(
910
                np.array([
911
                    [False, False, False],
912
                    [False, False, False],
913
                    [False, False, False],
914
                    [False, False, False],
915
                    [False, False, False]
916
                ]),
917
                index=mask.index,
918
                columns=mask.columns
919
            )
920
        )
921
        en, ex = pd.DataFrame.vbt.signals.generate_random_both(
922
            (5, 3), entry_prob=1., exit_prob=1., entry_pick_first=True, exit_pick_first=False,
923
            seed=seed, index=mask.index, columns=mask.columns)
924
        pd.testing.assert_frame_equal(
925
            en,
926
            pd.DataFrame(
927
                np.array([
928
                    [True, True, True],
929
                    [False, False, False],
930
                    [False, False, False],
931
                    [False, False, False],
932
                    [False, False, False]
933
                ]),
934
                index=mask.index,
935
                columns=mask.columns
936
            )
937
        )
938
        pd.testing.assert_frame_equal(
939
            ex,
940
            pd.DataFrame(
941
                np.array([
942
                    [False, False, False],
943
                    [True, True, True],
944
                    [True, True, True],
945
                    [True, True, True],
946
                    [True, True, True]
947
                ]),
948
                index=mask.index,
949
                columns=mask.columns
950
            )
951
        )
952
        # none
953
        with pytest.raises(Exception):
954
            pd.DataFrame.vbt.signals.generate_random((5, 3))
955

956
    def test_generate_random_exits(self):
957
        pd.testing.assert_series_equal(
958
            mask['a'].vbt.signals.generate_random_exits(seed=seed),
959
            pd.Series(
960
                np.array([False, False, True, False, True]),
961
                index=mask['a'].index,
962
                name=mask['a'].name
963
            )
964
        )
965
        pd.testing.assert_frame_equal(
966
            mask.vbt.signals.generate_random_exits(seed=seed),
967
            pd.DataFrame(
968
                np.array([
969
                    [False, False, False],
970
                    [False, False, False],
971
                    [True, True, False],
972
                    [False, False, False],
973
                    [True, False, True]
974
                ]),
975
                index=mask.index,
976
                columns=mask.columns
977
            )
978
        )
979
        pd.testing.assert_frame_equal(
980
            mask.vbt.signals.generate_random_exits(seed=seed, wait=0),
981
            pd.DataFrame(
982
                np.array([
983
                    [True, False, False],
984
                    [False, False, False],
985
                    [False, True, False],
986
                    [False, False, True],
987
                    [True, True, False]
988
                ]),
989
                index=mask.index,
990
                columns=mask.columns
991
            )
992
        )
993
        pd.testing.assert_series_equal(
994
            mask['a'].vbt.signals.generate_random_exits(prob=1., seed=seed),
995
            pd.Series(
996
                np.array([False, True, False, False, True]),
997
                index=mask['a'].index,
998
                name=mask['a'].name
999
            )
1000
        )
1001
        pd.testing.assert_frame_equal(
1002
            mask.vbt.signals.generate_random_exits(prob=1., seed=seed),
1003
            pd.DataFrame(
1004
                np.array([
1005
                    [False, False, False],
1006
                    [True, False, False],
1007
                    [False, True, False],
1008
                    [False, False, True],
1009
                    [True, False, False]
1010
                ]),
1011
                index=mask.index,
1012
                columns=mask.columns
1013
            )
1014
        )
1015
        pd.testing.assert_frame_equal(
1016
            mask.vbt.signals.generate_random_exits(prob=[0., 0.5, 1.], seed=seed),
1017
            pd.DataFrame(
1018
                np.array([
1019
                    [False, False, False],
1020
                    [False, False, False],
1021
                    [False, False, False],
1022
                    [False, True, True],
1023
                    [False, False, False]
1024
                ]),
1025
                index=mask.index,
1026
                columns=mask.columns
1027
            )
1028
        )
1029
        pd.testing.assert_frame_equal(
1030
            mask.vbt.signals.generate_random_exits(prob=1., wait=0, seed=seed),
1031
            pd.DataFrame(
1032
                np.array([
1033
                    [True, False, False],
1034
                    [False, True, False],
1035
                    [False, False, True],
1036
                    [True, False, False],
1037
                    [False, True, False]
1038
                ]),
1039
                index=mask.index,
1040
                columns=mask.columns
1041
            )
1042
        )
1043
        pd.testing.assert_frame_equal(
1044
            mask.vbt.signals.generate_random_exits(prob=1., until_next=False, seed=seed),
1045
            pd.DataFrame(
1046
                np.array([
1047
                    [False, False, False],
1048
                    [True, False, False],
1049
                    [False, True, False],
1050
                    [False, False, True],
1051
                    [True, False, False]
1052
                ]),
1053
                index=mask.index,
1054
                columns=mask.columns
1055
            )
1056
        )
1057

1058
    def test_generate_stop_exits(self):
1059
        e = pd.Series([True, False, False, False, False, False])
1060
        t = pd.Series([2, 3, 4, 3, 2, 1]).astype(np.float64)
1061

1062
        # stop loss
1063
        pd.testing.assert_series_equal(
1064
            e.vbt.signals.generate_stop_exits(t, -0.1),
1065
            pd.Series(np.array([False, False, False, False, False, True]))
1066
        )
1067
        pd.testing.assert_series_equal(
1068
            e.vbt.signals.generate_stop_exits(t, -0.1, trailing=True),
1069
            pd.Series(np.array([False, False, False, True, False, False]))
1070
        )
1071
        pd.testing.assert_series_equal(
1072
            e.vbt.signals.generate_stop_exits(t, -0.1, trailing=True, pick_first=False),
1073
            pd.Series(np.array([False, False, False, True, True, True]))
1074
        )
1075
        pd.testing.assert_frame_equal(
1076
            e.vbt.signals.generate_stop_exits(t.vbt.tile(3), [np.nan, -0.5, -1.], trailing=True, pick_first=False),
1077
            pd.DataFrame(np.array([
1078
                [False, False, False],
1079
                [False, False, False],
1080
                [False, False, False],
1081
                [False, False, False],
1082
                [False, True, False],
1083
                [False, True, False]
1084
            ]))
1085
        )
1086
        pd.testing.assert_series_equal(
1087
            e.vbt.signals.generate_stop_exits(t, -0.1, trailing=True, exit_wait=3),
1088
            pd.Series(np.array([False, False, False, False, True, False]))
1089
        )
1090
        # take profit
1091
        pd.testing.assert_series_equal(
1092
            e.vbt.signals.generate_stop_exits(4 - t, 0.1),
1093
            pd.Series(np.array([False, False, False, False, False, True]))
1094
        )
1095
        pd.testing.assert_series_equal(
1096
            e.vbt.signals.generate_stop_exits(4 - t, 0.1, trailing=True),
1097
            pd.Series(np.array([False, False, False, True, False, False]))
1098
        )
1099
        pd.testing.assert_series_equal(
1100
            e.vbt.signals.generate_stop_exits(4 - t, 0.1, trailing=True, pick_first=False),
1101
            pd.Series(np.array([False, False, False, True, True, True]))
1102
        )
1103
        pd.testing.assert_frame_equal(
1104
            e.vbt.signals.generate_stop_exits((4 - t).vbt.tile(3), [np.nan, 0.5, 1.], trailing=True, pick_first=False),
1105
            pd.DataFrame(np.array([
1106
                [False, False, False],
1107
                [False, False, False],
1108
                [False, False, False],
1109
                [False, True, True],
1110
                [False, True, True],
1111
                [False, True, True]
1112
            ]))
1113
        )
1114
        pd.testing.assert_series_equal(
1115
            e.vbt.signals.generate_stop_exits(4 - t, 0.1, trailing=True, exit_wait=3),
1116
            pd.Series(np.array([False, False, False, False, True, False]))
1117
        )
1118
        # chain
1119
        e = pd.Series([True, True, True, True, True, True])
1120
        en, ex = e.vbt.signals.generate_stop_exits(t, -0.1, trailing=True, chain=True)
1121
        pd.testing.assert_series_equal(
1122
            en,
1123
            pd.Series(np.array([True, False, False, False, True, False]))
1124
        )
1125
        pd.testing.assert_series_equal(
1126
            ex,
1127
            pd.Series(np.array([False, False, False, True, False, True]))
1128
        )
1129
        en, ex = e.vbt.signals.generate_stop_exits(t, -0.1, trailing=True, entry_wait=2, chain=True)
1130
        pd.testing.assert_series_equal(
1131
            en,
1132
            pd.Series(np.array([True, False, False, False, False, True]))
1133
        )
1134
        pd.testing.assert_series_equal(
1135
            ex,
1136
            pd.Series(np.array([False, False, False, True, False, False]))
1137
        )
1138
        en, ex = e.vbt.signals.generate_stop_exits(t, -0.1, trailing=True, exit_wait=2, chain=True)
1139
        pd.testing.assert_series_equal(
1140
            en,
1141
            pd.Series(np.array([True, False, False, False, True, False]))
1142
        )
1143
        pd.testing.assert_series_equal(
1144
            ex,
1145
            pd.Series(np.array([False, False, False, True, False, False]))
1146
        )
1147
        # until_next and pick_first
1148
        e2 = pd.Series([True, True, True, True, True, True])
1149
        t2 = pd.Series([6, 5, 4, 3, 2, 1]).astype(np.float64)
1150
        ex = e2.vbt.signals.generate_stop_exits(t2, -0.1, until_next=False, pick_first=False)
1151
        pd.testing.assert_series_equal(
1152
            ex,
1153
            pd.Series(np.array([False, True, True, True, True, True]))
1154
        )
1155

1156
    def test_generate_ohlc_stop_exits(self):
1157
        with pytest.raises(Exception):
1158
            _ = mask.vbt.signals.generate_ohlc_stop_exits(ts, sl_stop=-0.1)
1159
        with pytest.raises(Exception):
1160
            _ = mask.vbt.signals.generate_ohlc_stop_exits(ts, tp_stop=-0.1)
1161

1162
        pd.testing.assert_frame_equal(
1163
            mask.vbt.signals.generate_stop_exits(ts, -0.1),
1164
            mask.vbt.signals.generate_ohlc_stop_exits(ts, sl_stop=0.1)
1165
        )
1166
        pd.testing.assert_frame_equal(
1167
            mask.vbt.signals.generate_stop_exits(ts, -0.1, trailing=True),
1168
            mask.vbt.signals.generate_ohlc_stop_exits(ts, sl_stop=0.1, sl_trail=True)
1169
        )
1170
        pd.testing.assert_frame_equal(
1171
            mask.vbt.signals.generate_stop_exits(ts, 0.1),
1172
            mask.vbt.signals.generate_ohlc_stop_exits(ts, tp_stop=0.1)
1173
        )
1174
        pd.testing.assert_frame_equal(
1175
            mask.vbt.signals.generate_stop_exits(ts, 0.1),
1176
            mask.vbt.signals.generate_ohlc_stop_exits(ts, sl_stop=0.1, reverse=True)
1177
        )
1178
        pd.testing.assert_frame_equal(
1179
            mask.vbt.signals.generate_stop_exits(ts, 0.1, trailing=True),
1180
            mask.vbt.signals.generate_ohlc_stop_exits(ts, sl_stop=0.1, sl_trail=True, reverse=True)
1181
        )
1182
        pd.testing.assert_frame_equal(
1183
            mask.vbt.signals.generate_stop_exits(ts, -0.1),
1184
            mask.vbt.signals.generate_ohlc_stop_exits(ts, tp_stop=0.1, reverse=True)
1185
        )
1186

1187
        def _test_ohlc_stop_exits(**kwargs):
1188
            out_dict = {'stop_price': np.nan, 'stop_type': -1}
1189
            result = mask.vbt.signals.generate_ohlc_stop_exits(
1190
                price['open'], price['high'], price['low'], price['close'],
1191
                out_dict=out_dict, **kwargs
1192
            )
1193
            if isinstance(result, tuple):
1194
                _, ex = result
1195
            else:
1196
                ex = result
1197
            return result, out_dict['stop_price'], out_dict['stop_type']
1198

1199
        ex, stop_price, stop_type = _test_ohlc_stop_exits()
1200
        pd.testing.assert_frame_equal(
1201
            ex,
1202
            pd.DataFrame(np.array([
1203
                [False, False, False],
1204
                [False, False, False],
1205
                [False, False, False],
1206
                [False, False, False],
1207
                [False, False, False]
1208
            ]), index=mask.index, columns=mask.columns)
1209
        )
1210
        pd.testing.assert_frame_equal(
1211
            stop_price,
1212
            pd.DataFrame(np.array([
1213
                [np.nan, np.nan, np.nan],
1214
                [np.nan, np.nan, np.nan],
1215
                [np.nan, np.nan, np.nan],
1216
                [np.nan, np.nan, np.nan],
1217
                [np.nan, np.nan, np.nan]
1218
            ]), index=mask.index, columns=mask.columns)
1219
        )
1220
        pd.testing.assert_frame_equal(
1221
            stop_type,
1222
            pd.DataFrame(np.array([
1223
                [-1, -1, -1],
1224
                [-1, -1, -1],
1225
                [-1, -1, -1],
1226
                [-1, -1, -1],
1227
                [-1, -1, -1]
1228
            ]), index=mask.index, columns=mask.columns)
1229
        )
1230
        ex, stop_price, stop_type = _test_ohlc_stop_exits(sl_stop=0.1)
1231
        pd.testing.assert_frame_equal(
1232
            ex,
1233
            pd.DataFrame(np.array([
1234
                [False, False, False],
1235
                [False, False, False],
1236
                [False, False, False],
1237
                [False, False, True],
1238
                [True, False, False]
1239
            ]), index=mask.index, columns=mask.columns)
1240
        )
1241
        pd.testing.assert_frame_equal(
1242
            stop_price,
1243
            pd.DataFrame(np.array([
1244
                [np.nan, np.nan, np.nan],
1245
                [np.nan, np.nan, np.nan],
1246
                [np.nan, np.nan, np.nan],
1247
                [np.nan, np.nan, 10.8],
1248
                [9.9, np.nan, np.nan]
1249
            ]), index=mask.index, columns=mask.columns)
1250
        )
1251
        pd.testing.assert_frame_equal(
1252
            stop_type,
1253
            pd.DataFrame(np.array([
1254
                [-1, -1, -1],
1255
                [-1, -1, -1],
1256
                [-1, -1, -1],
1257
                [-1, -1, 0],
1258
                [0, -1, -1]
1259
            ]), index=mask.index, columns=mask.columns)
1260
        )
1261
        ex, stop_price, stop_type = _test_ohlc_stop_exits(sl_stop=0.1, sl_trail=True)
1262
        pd.testing.assert_frame_equal(
1263
            ex,
1264
            pd.DataFrame(np.array([
1265
                [False, False, False],
1266
                [False, False, False],
1267
                [False, False, False],
1268
                [False, True, True],
1269
                [True, False, False]
1270
            ]), index=mask.index, columns=mask.columns)
1271
        )
1272
        pd.testing.assert_frame_equal(
1273
            stop_price,
1274
            pd.DataFrame(np.array([
1275
                [np.nan, np.nan, np.nan],
1276
                [np.nan, np.nan, np.nan],
1277
                [np.nan, np.nan, np.nan],
1278
                [np.nan, 11.7, 10.8],
1279
                [9.9, np.nan, np.nan]
1280
            ]), index=mask.index, columns=mask.columns)
1281
        )
1282
        pd.testing.assert_frame_equal(
1283
            stop_type,
1284
            pd.DataFrame(np.array([
1285
                [-1, -1, -1],
1286
                [-1, -1, -1],
1287
                [-1, -1, -1],
1288
                [-1, 1, 1],
1289
                [1, -1, -1]
1290
            ]), index=mask.index, columns=mask.columns)
1291
        )
1292
        ex, stop_price, stop_type = _test_ohlc_stop_exits(tp_stop=0.1)
1293
        pd.testing.assert_frame_equal(
1294
            ex,
1295
            pd.DataFrame(np.array([
1296
                [False, False, False],
1297
                [True, False, False],
1298
                [False, True, False],
1299
                [False, False, False],
1300
                [False, False, False]
1301
            ]), index=mask.index, columns=mask.columns)
1302
        )
1303
        pd.testing.assert_frame_equal(
1304
            stop_price,
1305
            pd.DataFrame(np.array([
1306
                [np.nan, np.nan, np.nan],
1307
                [11.0, np.nan, np.nan],
1308
                [np.nan, 12.1, np.nan],
1309
                [np.nan, np.nan, np.nan],
1310
                [np.nan, np.nan, np.nan]
1311
            ]), index=mask.index, columns=mask.columns)
1312
        )
1313
        pd.testing.assert_frame_equal(
1314
            stop_type,
1315
            pd.DataFrame(np.array([
1316
                [-1, -1, -1],
1317
                [2, -1, -1],
1318
                [-1, 2, -1],
1319
                [-1, -1, -1],
1320
                [-1, -1, -1]
1321
            ]), index=mask.index, columns=mask.columns)
1322
        )
1323
        ex, stop_price, stop_type = _test_ohlc_stop_exits(sl_stop=0.1, sl_trail=True, tp_stop=0.1)
1324
        pd.testing.assert_frame_equal(
1325
            ex,
1326
            pd.DataFrame(np.array([
1327
                [False, False, False],
1328
                [True, False, False],
1329
                [False, True, False],
1330
                [False, False, True],
1331
                [True, False, False]
1332
            ]), index=mask.index, columns=mask.columns)
1333
        )
1334
        pd.testing.assert_frame_equal(
1335
            stop_price,
1336
            pd.DataFrame(np.array([
1337
                [np.nan, np.nan, np.nan],
1338
                [11.0, np.nan, np.nan],
1339
                [np.nan, 12.1, np.nan],
1340
                [np.nan, np.nan, 10.8],
1341
                [9.9, np.nan, np.nan]
1342
            ]), index=mask.index, columns=mask.columns)
1343
        )
1344
        pd.testing.assert_frame_equal(
1345
            stop_type,
1346
            pd.DataFrame(np.array([
1347
                [-1, -1, -1],
1348
                [2, -1, -1],
1349
                [-1, 2, -1],
1350
                [-1, -1, 1],
1351
                [1, -1, -1]
1352
            ]), index=mask.index, columns=mask.columns)
1353
        )
1354
        ex, stop_price, stop_type = _test_ohlc_stop_exits(
1355
            sl_stop=[np.nan, 0.1, 0.2], sl_trail=True, tp_stop=[np.nan, 0.1, 0.2])
1356
        pd.testing.assert_frame_equal(
1357
            ex,
1358
            pd.DataFrame(np.array([
1359
                [False, False, False],
1360
                [False, False, False],
1361
                [False, True, False],
1362
                [False, False, False],
1363
                [False, False, True]
1364
            ]), index=mask.index, columns=mask.columns)
1365
        )
1366
        pd.testing.assert_frame_equal(
1367
            stop_price,
1368
            pd.DataFrame(np.array([
1369
                [np.nan, np.nan, np.nan],
1370
                [np.nan, np.nan, np.nan],
1371
                [np.nan, 12.1, np.nan],
1372
                [np.nan, np.nan, np.nan],
1373
                [np.nan, np.nan, 9.6]
1374
            ]), index=mask.index, columns=mask.columns)
1375
        )
1376
        pd.testing.assert_frame_equal(
1377
            stop_type,
1378
            pd.DataFrame(np.array([
1379
                [-1, -1, -1],
1380
                [-1, -1, -1],
1381
                [-1, 2, -1],
1382
                [-1, -1, -1],
1383
                [-1, -1, 1]
1384
            ]), index=mask.index, columns=mask.columns)
1385
        )
1386
        ex, stop_price, stop_type = _test_ohlc_stop_exits(sl_stop=0.1, sl_trail=True, tp_stop=0.1, exit_wait=0)
1387
        pd.testing.assert_frame_equal(
1388
            ex,
1389
            pd.DataFrame(np.array([
1390
                [True, False, False],
1391
                [False, False, False],
1392
                [False, True, False],
1393
                [False, False, True],
1394
                [True, True, False]
1395
            ]), index=mask.index, columns=mask.columns)
1396
        )
1397
        pd.testing.assert_frame_equal(
1398
            stop_price,
1399
            pd.DataFrame(np.array([
1400
                [9.0, np.nan, np.nan],
1401
                [np.nan, np.nan, np.nan],
1402
                [np.nan, 12.1, np.nan],
1403
                [np.nan, np.nan, 11.7],
1404
                [10.8, 9.0, np.nan]
1405
            ]), index=mask.index, columns=mask.columns)
1406
        )
1407
        pd.testing.assert_frame_equal(
1408
            stop_type,
1409
            pd.DataFrame(np.array([
1410
                [1, -1, -1],
1411
                [-1, -1, -1],
1412
                [-1, 2, -1],
1413
                [-1, -1, 1],
1414
                [1, 1, -1]
1415
            ]), index=mask.index, columns=mask.columns)
1416
        )
1417
        (en, ex), stop_price, stop_type = _test_ohlc_stop_exits(
1418
            sl_stop=0.1, sl_trail=True, tp_stop=0.1, chain=True)
1419
        pd.testing.assert_frame_equal(
1420
            en,
1421
            pd.DataFrame(np.array([
1422
                [True, False, False],
1423
                [False, True, False],
1424
                [False, False, True],
1425
                [True, False, False],
1426
                [False, True, False]
1427
            ]), index=mask.index, columns=mask.columns)
1428
        )
1429
        pd.testing.assert_frame_equal(
1430
            ex,
1431
            pd.DataFrame(np.array([
1432
                [False, False, False],
1433
                [True, False, False],
1434
                [False, True, False],
1435
                [False, False, True],
1436
                [True, False, False]
1437
            ]), index=mask.index, columns=mask.columns)
1438
        )
1439
        pd.testing.assert_frame_equal(
1440
            stop_price,
1441
            pd.DataFrame(np.array([
1442
                [np.nan, np.nan, np.nan],
1443
                [11.0, np.nan, np.nan],
1444
                [np.nan, 12.1, np.nan],
1445
                [np.nan, np.nan, 10.8],
1446
                [9.9, np.nan, np.nan]
1447
            ]), index=mask.index, columns=mask.columns)
1448
        )
1449
        pd.testing.assert_frame_equal(
1450
            stop_type,
1451
            pd.DataFrame(np.array([
1452
                [-1, -1, -1],
1453
                [2, -1, -1],
1454
                [-1, 2, -1],
1455
                [-1, -1, 1],
1456
                [1, -1, -1]
1457
            ]), index=mask.index, columns=mask.columns)
1458
        )
1459

1460
    def test_between_ranges(self):
1461
        ranges = mask.vbt.signals.between_ranges()
1462
        record_arrays_close(
1463
            ranges.values,
1464
            np.array([
1465
                (0, 0, 0, 3, 1), (1, 1, 1, 4, 1)
1466
            ], dtype=range_dt)
1467
        )
1468
        assert ranges.wrapper == mask.vbt.wrapper
1469

1470
        mask2 = pd.DataFrame([
1471
            [True, True, True],
1472
            [True, True, True],
1473
            [False, False, False],
1474
            [False, False, False],
1475
            [False, False, False]
1476
        ], index=mask.index, columns=mask.columns)
1477

1478
        other_mask = pd.DataFrame([
1479
            [False, False, False],
1480
            [True, False, False],
1481
            [True, True, False],
1482
            [False, True, True],
1483
            [False, False, True]
1484
        ], index=mask.index, columns=mask.columns)
1485

1486
        ranges = mask2.vbt.signals.between_ranges(other=other_mask)
1487
        record_arrays_close(
1488
            ranges.values,
1489
            np.array([
1490
                (0, 0, 0, 1, 1), (1, 0, 1, 1, 1), (2, 1, 0, 2, 1),
1491
                (3, 1, 1, 2, 1), (4, 2, 0, 3, 1), (5, 2, 1, 3, 1)
1492
            ], dtype=range_dt)
1493
        )
1494
        assert ranges.wrapper == mask2.vbt.wrapper
1495

1496
        ranges = mask2.vbt.signals.between_ranges(other=other_mask, from_other=True)
1497
        record_arrays_close(
1498
            ranges.values,
1499
            np.array([
1500
                (0, 0, 1, 1, 1), (1, 0, 1, 2, 1), (2, 1, 1, 2, 1),
1501
                (3, 1, 1, 3, 1), (4, 2, 1, 3, 1), (5, 2, 1, 4, 1)
1502
            ], dtype=range_dt)
1503
        )
1504
        assert ranges.wrapper == mask2.vbt.wrapper
1505

1506
    def test_partition_ranges(self):
1507
        mask2 = pd.DataFrame([
1508
            [False, False, False],
1509
            [True, False, False],
1510
            [True, True, False],
1511
            [False, True, True],
1512
            [True, False, True]
1513
        ], index=mask.index, columns=mask.columns)
1514

1515
        ranges = mask2.vbt.signals.partition_ranges()
1516
        record_arrays_close(
1517
            ranges.values,
1518
            np.array([
1519
                (0, 0, 1, 3, 1), (1, 0, 4, 4, 0), (2, 1, 2, 4, 1), (3, 2, 3, 4, 0)
1520
            ], dtype=range_dt)
1521
        )
1522
        assert ranges.wrapper == mask2.vbt.wrapper
1523

1524
    def test_between_partition_ranges(self):
1525
        mask2 = pd.DataFrame([
1526
            [True, False, False],
1527
            [True, True, False],
1528
            [False, True, True],
1529
            [True, False, True],
1530
            [False, True, False]
1531
        ], index=mask.index, columns=mask.columns)
1532

1533
        ranges = mask2.vbt.signals.between_partition_ranges()
1534
        record_arrays_close(
1535
            ranges.values,
1536
            np.array([
1537
                (0, 0, 1, 3, 1), (1, 1, 2, 4, 1)
1538
            ], dtype=range_dt)
1539
        )
1540
        assert ranges.wrapper == mask2.vbt.wrapper
1541

1542
    def test_pos_rank(self):
1543
        pd.testing.assert_series_equal(
1544
            (~mask['a']).vbt.signals.pos_rank(),
1545
            pd.Series([-1, 0, 1, -1, 0], index=mask['a'].index, name=mask['a'].name)
1546
        )
1547
        pd.testing.assert_frame_equal(
1548
            (~mask).vbt.signals.pos_rank(),
1549
            pd.DataFrame(
1550
                np.array([
1551
                    [-1, 0, 0],
1552
                    [0, -1, 1],
1553
                    [1, 0, -1],
1554
                    [-1, 1, 0],
1555
                    [0, -1, 1]
1556
                ]),
1557
                index=mask.index,
1558
                columns=mask.columns
1559
            )
1560
        )
1561
        pd.testing.assert_frame_equal(
1562
            (~mask).vbt.signals.pos_rank(after_false=True),
1563
            pd.DataFrame(
1564
                np.array([
1565
                    [-1, -1, -1],
1566
                    [0, -1, -1],
1567
                    [1, 0, -1],
1568
                    [-1, 1, 0],
1569
                    [0, -1, 1]
1570
                ]),
1571
                index=mask.index,
1572
                columns=mask.columns
1573
            )
1574
        )
1575
        pd.testing.assert_frame_equal(
1576
            (~mask).vbt.signals.pos_rank(allow_gaps=True),
1577
            pd.DataFrame(
1578
                np.array([
1579
                    [-1, 0, 0],
1580
                    [0, -1, 1],
1581
                    [1, 1, -1],
1582
                    [-1, 2, 2],
1583
                    [2, -1, 3]
1584
                ]),
1585
                index=mask.index,
1586
                columns=mask.columns
1587
            )
1588
        )
1589
        pd.testing.assert_frame_equal(
1590
            (~mask).vbt.signals.pos_rank(reset_by=mask['a'], allow_gaps=True),
1591
            pd.DataFrame(
1592
                np.array([
1593
                    [-1, 0, 0],
1594
                    [0, -1, 1],
1595
                    [1, 1, -1],
1596
                    [-1, 0, 0],
1597
                    [0, -1, 1]
1598
                ]),
1599
                index=mask.index,
1600
                columns=mask.columns
1601
            )
1602
        )
1603
        pd.testing.assert_frame_equal(
1604
            (~mask).vbt.signals.pos_rank(reset_by=mask, allow_gaps=True),
1605
            pd.DataFrame(
1606
                np.array([
1607
                    [-1, 0, 0],
1608
                    [0, -1, 1],
1609
                    [1, 0, -1],
1610
                    [-1, 1, 0],
1611
                    [0, -1, 1]
1612
                ]),
1613
                index=mask.index,
1614
                columns=mask.columns
1615
            )
1616
        )
1617

1618
    def test_partition_pos_rank(self):
1619
        pd.testing.assert_series_equal(
1620
            (~mask['a']).vbt.signals.partition_pos_rank(),
1621
            pd.Series([-1, 0, 0, -1, 1], index=mask['a'].index, name=mask['a'].name)
1622
        )
1623
        pd.testing.assert_frame_equal(
1624
            (~mask).vbt.signals.partition_pos_rank(),
1625
            pd.DataFrame(
1626
                np.array([
1627
                    [-1, 0, 0],
1628
                    [0, -1, 0],
1629
                    [0, 1, -1],
1630
                    [-1, 1, 1],
1631
                    [1, -1, 1]
1632
                ]),
1633
                index=mask.index,
1634
                columns=mask.columns
1635
            )
1636
        )
1637
        pd.testing.assert_frame_equal(
1638
            (~mask).vbt.signals.partition_pos_rank(after_false=True),
1639
            pd.DataFrame(
1640
                np.array([
1641
                    [-1, -1, -1],
1642
                    [0, -1, -1],
1643
                    [0, 0, -1],
1644
                    [-1, 0, 0],
1645
                    [1, -1, 0]
1646
                ]),
1647
                index=mask.index,
1648
                columns=mask.columns
1649
            )
1650
        )
1651
        pd.testing.assert_frame_equal(
1652
            (~mask).vbt.signals.partition_pos_rank(reset_by=mask['a']),
1653
            pd.DataFrame(
1654
                np.array([
1655
                    [-1, 0, 0],
1656
                    [0, -1, 0],
1657
                    [0, 1, -1],
1658
                    [-1, 0, 0],
1659
                    [0, -1, 0]
1660
                ]),
1661
                index=mask.index,
1662
                columns=mask.columns
1663
            )
1664
        )
1665
        pd.testing.assert_frame_equal(
1666
            (~mask).vbt.signals.partition_pos_rank(reset_by=mask),
1667
            pd.DataFrame(
1668
                np.array([
1669
                    [-1, 0, 0],
1670
                    [0, -1, 0],
1671
                    [0, 0, -1],
1672
                    [-1, 0, 0],
1673
                    [0, -1, 0]
1674
                ]),
1675
                index=mask.index,
1676
                columns=mask.columns
1677
            )
1678
        )
1679

1680
    def test_pos_rank_fns(self):
1681
        pd.testing.assert_frame_equal(
1682
            (~mask).vbt.signals.first(),
1683
            pd.DataFrame(
1684
                np.array([
1685
                    [False, True, True],
1686
                    [True, False, False],
1687
                    [False, True, False],
1688
                    [False, False, True],
1689
                    [True, False, False]
1690
                ]),
1691
                index=mask.index,
1692
                columns=mask.columns
1693
            )
1694
        )
1695
        pd.testing.assert_frame_equal(
1696
            (~mask).vbt.signals.nth(1),
1697
            pd.DataFrame(
1698
                np.array([
1699
                    [False, False, False],
1700
                    [False, False, True],
1701
                    [True, False, False],
1702
                    [False, True, False],
1703
                    [False, False, True]
1704
                ]),
1705
                index=mask.index,
1706
                columns=mask.columns
1707
            )
1708
        )
1709
        pd.testing.assert_frame_equal(
1710
            (~mask).vbt.signals.nth(2),
1711
            pd.DataFrame(
1712
                np.array([
1713
                    [False, False, False],
1714
                    [False, False, False],
1715
                    [False, False, False],
1716
                    [False, False, False],
1717
                    [False, False, False]
1718
                ]),
1719
                index=mask.index,
1720
                columns=mask.columns
1721
            )
1722
        )
1723
        pd.testing.assert_frame_equal(
1724
            (~mask).vbt.signals.from_nth(0),
1725
            pd.DataFrame(
1726
                np.array([
1727
                    [False, True, True],
1728
                    [True, False, True],
1729
                    [True, True, False],
1730
                    [False, True, True],
1731
                    [True, False, True]
1732
                ]),
1733
                index=mask.index,
1734
                columns=mask.columns
1735
            )
1736
        )
1737

1738
    def test_pos_rank_mapped(self):
1739
        mask2 = pd.DataFrame([
1740
            [True, False, False],
1741
            [True, True, False],
1742
            [False, True, True],
1743
            [True, False, True],
1744
            [False, True, False]
1745
        ], index=mask.index, columns=mask.columns)
1746

1747
        mapped = mask2.vbt.signals.pos_rank_mapped()
1748
        np.testing.assert_array_equal(
1749
            mapped.values,
1750
            np.array([0, 1, 0, 0, 1, 0, 0, 1])
1751
        )
1752
        np.testing.assert_array_equal(
1753
            mapped.col_arr,
1754
            np.array([0, 0, 0, 1, 1, 1, 2, 2])
1755
        )
1756
        np.testing.assert_array_equal(
1757
            mapped.idx_arr,
1758
            np.array([0, 1, 3, 1, 2, 4, 2, 3])
1759
        )
1760
        assert mapped.wrapper == mask2.vbt.wrapper
1761

1762
    def test_partition_pos_rank_mapped(self):
1763
        mask2 = pd.DataFrame([
1764
            [True, False, False],
1765
            [True, True, False],
1766
            [False, True, True],
1767
            [True, False, True],
1768
            [False, True, False]
1769
        ], index=mask.index, columns=mask.columns)
1770

1771
        mapped = mask2.vbt.signals.partition_pos_rank_mapped()
1772
        np.testing.assert_array_equal(
1773
            mapped.values,
1774
            np.array([0, 0, 1, 0, 0, 1, 0, 0])
1775
        )
1776
        np.testing.assert_array_equal(
1777
            mapped.col_arr,
1778
            np.array([0, 0, 0, 1, 1, 1, 2, 2])
1779
        )
1780
        np.testing.assert_array_equal(
1781
            mapped.idx_arr,
1782
            np.array([0, 1, 3, 1, 2, 4, 2, 3])
1783
        )
1784
        assert mapped.wrapper == mask2.vbt.wrapper
1785

1786
    def test_nth_index(self):
1787
        assert mask['a'].vbt.signals.nth_index(0) == pd.Timestamp('2020-01-01 00:00:00')
1788
        pd.testing.assert_series_equal(
1789
            mask.vbt.signals.nth_index(0),
1790
            pd.Series([
1791
                pd.Timestamp('2020-01-01 00:00:00'),
1792
                pd.Timestamp('2020-01-02 00:00:00'),
1793
                pd.Timestamp('2020-01-03 00:00:00')
1794
            ], index=mask.columns, name='nth_index', dtype='datetime64[ns]')
1795
        )
1796
        pd.testing.assert_series_equal(
1797
            mask.vbt.signals.nth_index(-1),
1798
            pd.Series([
1799
                pd.Timestamp('2020-01-04 00:00:00'),
1800
                pd.Timestamp('2020-01-05 00:00:00'),
1801
                pd.Timestamp('2020-01-03 00:00:00')
1802
            ], index=mask.columns, name='nth_index', dtype='datetime64[ns]')
1803
        )
1804
        pd.testing.assert_series_equal(
1805
            mask.vbt.signals.nth_index(-2),
1806
            pd.Series([
1807
                pd.Timestamp('2020-01-01 00:00:00'),
1808
                pd.Timestamp('2020-01-02 00:00:00'),
1809
                np.nan
1810
            ], index=mask.columns, name='nth_index', dtype='datetime64[ns]')
1811
        )
1812
        pd.testing.assert_series_equal(
1813
            mask.vbt.signals.nth_index(0, group_by=group_by),
1814
            pd.Series([
1815
                pd.Timestamp('2020-01-01 00:00:00'),
1816
                pd.Timestamp('2020-01-03 00:00:00')
1817
            ], index=['g1', 'g2'], name='nth_index', dtype='datetime64[ns]')
1818
        )
1819
        pd.testing.assert_series_equal(
1820
            mask.vbt.signals.nth_index(-1, group_by=group_by),
1821
            pd.Series([
1822
                pd.Timestamp('2020-01-05 00:00:00'),
1823
                pd.Timestamp('2020-01-03 00:00:00')
1824
            ], index=['g1', 'g2'], name='nth_index', dtype='datetime64[ns]')
1825
        )
1826

1827
    def test_norm_avg_index(self):
1828
        assert mask['a'].vbt.signals.norm_avg_index() == -0.25
1829
        pd.testing.assert_series_equal(
1830
            mask.vbt.signals.norm_avg_index(),
1831
            pd.Series([-0.25, 0.25, 0.0], index=mask.columns, name='norm_avg_index')
1832
        )
1833
        pd.testing.assert_series_equal(
1834
            mask.vbt.signals.norm_avg_index(group_by=group_by),
1835
            pd.Series([0.0, 0.0], index=['g1', 'g2'], name='norm_avg_index')
1836
        )
1837

1838
    def test_index_mapped(self):
1839
        mapped = mask.vbt.signals.index_mapped()
1840
        np.testing.assert_array_equal(
1841
            mapped.values,
1842
            np.array([0, 3, 1, 4, 2])
1843
        )
1844
        np.testing.assert_array_equal(
1845
            mapped.col_arr,
1846
            np.array([0, 0, 1, 1, 2])
1847
        )
1848
        np.testing.assert_array_equal(
1849
            mapped.idx_arr,
1850
            np.array([0, 3, 1, 4, 2])
1851
        )
1852
        assert mapped.wrapper == mask.vbt.wrapper
1853

1854
    def test_total(self):
1855
        assert mask['a'].vbt.signals.total() == 2
1856
        pd.testing.assert_series_equal(
1857
            mask.vbt.signals.total(),
1858
            pd.Series([2, 2, 1], index=mask.columns, name='total')
1859
        )
1860
        pd.testing.assert_series_equal(
1861
            mask.vbt.signals.total(group_by=group_by),
1862
            pd.Series([4, 1], index=['g1', 'g2'], name='total')
1863
        )
1864

1865
    def test_rate(self):
1866
        assert mask['a'].vbt.signals.rate() == 0.4
1867
        pd.testing.assert_series_equal(
1868
            mask.vbt.signals.rate(),
1869
            pd.Series([0.4, 0.4, 0.2], index=mask.columns, name='rate')
1870
        )
1871
        pd.testing.assert_series_equal(
1872
            mask.vbt.signals.rate(group_by=group_by),
1873
            pd.Series([0.4, 0.2], index=['g1', 'g2'], name='rate')
1874
        )
1875

1876
    def test_total_partitions(self):
1877
        assert mask['a'].vbt.signals.total_partitions() == 2
1878
        pd.testing.assert_series_equal(
1879
            mask.vbt.signals.total_partitions(),
1880
            pd.Series([2, 2, 1], index=mask.columns, name='total_partitions')
1881
        )
1882
        pd.testing.assert_series_equal(
1883
            mask.vbt.signals.total_partitions(group_by=group_by),
1884
            pd.Series([4, 1], index=['g1', 'g2'], name='total_partitions')
1885
        )
1886

1887
    def test_partition_rate(self):
1888
        assert mask['a'].vbt.signals.partition_rate() == 1.0
1889
        pd.testing.assert_series_equal(
1890
            mask.vbt.signals.partition_rate(),
1891
            pd.Series([1.0, 1.0, 1.0], index=mask.columns, name='partition_rate')
1892
        )
1893
        pd.testing.assert_series_equal(
1894
            mask.vbt.signals.partition_rate(group_by=group_by),
1895
            pd.Series([1.0, 1.0], index=['g1', 'g2'], name='partition_rate')
1896
        )
1897

1898
    def test_stats(self):
1899
        stats_index = pd.Index([
1900
            'Start', 'End', 'Period', 'Total', 'Rate [%]', 'First Index',
1901
            'Last Index', 'Norm Avg Index [-1, 1]', 'Distance: Min',
1902
            'Distance: Max', 'Distance: Mean', 'Distance: Std', 'Total Partitions',
1903
            'Partition Rate [%]', 'Partition Length: Min', 'Partition Length: Max',
1904
            'Partition Length: Mean', 'Partition Length: Std',
1905
            'Partition Distance: Min', 'Partition Distance: Max',
1906
            'Partition Distance: Mean', 'Partition Distance: Std'
1907
        ], dtype='object')
1908
        pd.testing.assert_series_equal(
1909
            mask.vbt.signals.stats(),
1910
            pd.Series([
1911
                pd.Timestamp('2020-01-01 00:00:00'),
1912
                pd.Timestamp('2020-01-05 00:00:00'),
1913
                pd.Timedelta('5 days 00:00:00'),
1914
                1.6666666666666667,
1915
                33.333333333333336,
1916
                pd.Timestamp('2020-01-02 00:00:00'),
1917
                pd.Timestamp('2020-01-04 00:00:00'),
1918
                0.0,
1919
                pd.Timedelta('3 days 00:00:00'),
1920
                pd.Timedelta('3 days 00:00:00'),
1921
                pd.Timedelta('3 days 00:00:00'),
1922
                np.nan,
1923
                1.6666666666666667,
1924
                100.0,
1925
                pd.Timedelta('1 days 00:00:00'),
1926
                pd.Timedelta('1 days 00:00:00'),
1927
                pd.Timedelta('1 days 00:00:00'),
1928
                pd.Timedelta('0 days 00:00:00'),
1929
                pd.Timedelta('3 days 00:00:00'),
1930
                pd.Timedelta('3 days 00:00:00'),
1931
                pd.Timedelta('3 days 00:00:00'),
1932
                np.nan
1933
            ],
1934
                index=stats_index,
1935
                name='agg_func_mean'
1936
            )
1937
        )
1938
        pd.testing.assert_series_equal(
1939
            mask.vbt.signals.stats(column='a'),
1940
            pd.Series([
1941
                pd.Timestamp('2020-01-01 00:00:00'),
1942
                pd.Timestamp('2020-01-05 00:00:00'),
1943
                pd.Timedelta('5 days 00:00:00'),
1944
                2,
1945
                40.0,
1946
                pd.Timestamp('2020-01-01 00:00:00'),
1947
                pd.Timestamp('2020-01-04 00:00:00'),
1948
                -0.25,
1949
                pd.Timedelta('3 days 00:00:00'),
1950
                pd.Timedelta('3 days 00:00:00'),
1951
                pd.Timedelta('3 days 00:00:00'),
1952
                np.nan,
1953
                2,
1954
                100.0,
1955
                pd.Timedelta('1 days 00:00:00'),
1956
                pd.Timedelta('1 days 00:00:00'),
1957
                pd.Timedelta('1 days 00:00:00'),
1958
                pd.Timedelta('0 days 00:00:00'),
1959
                pd.Timedelta('3 days 00:00:00'),
1960
                pd.Timedelta('3 days 00:00:00'),
1961
                pd.Timedelta('3 days 00:00:00'),
1962
                np.nan
1963
            ],
1964
                index=stats_index,
1965
                name='a'
1966
            )
1967
        )
1968
        pd.testing.assert_series_equal(
1969
            mask.vbt.signals.stats(column='a', settings=dict(to_timedelta=False)),
1970
            pd.Series([
1971
                pd.Timestamp('2020-01-01 00:00:00'), pd.Timestamp('2020-01-05 00:00:00'), 5, 2, 40.0,
1972
                pd.Timestamp('2020-01-01 00:00:00'), pd.Timestamp('2020-01-04 00:00:00'), -0.25, 3.0,
1973
                3.0, 3.0, np.nan, 2, 100.0, 1.0, 1.0, 1.0, 0.0, 3.0, 3.0, 3.0, np.nan
1974
            ],
1975
                index=stats_index,
1976
                name='a'
1977
            )
1978
        )
1979
        pd.testing.assert_series_equal(
1980
            mask.vbt.signals.stats(column='a', settings=dict(other=mask['b'], from_other=True)),
1981
            pd.Series([
1982
                pd.Timestamp('2020-01-01 00:00:00'),
1983
                pd.Timestamp('2020-01-05 00:00:00'),
1984
                pd.Timedelta('5 days 00:00:00'),
1985
                2,
1986
                40.0,
1987
                0,
1988
                0.0,
1989
                pd.Timestamp('2020-01-01 00:00:00'),
1990
                pd.Timestamp('2020-01-04 00:00:00'),
1991
                -0.25,
1992
                pd.Timedelta('1 days 00:00:00'),
1993
                pd.Timedelta('1 days 00:00:00'),
1994
                pd.Timedelta('1 days 00:00:00'),
1995
                pd.Timedelta('0 days 00:00:00'),
1996
                2,
1997
                100.0,
1998
                pd.Timedelta('1 days 00:00:00'),
1999
                pd.Timedelta('1 days 00:00:00'),
2000
                pd.Timedelta('1 days 00:00:00'),
2001
                pd.Timedelta('0 days 00:00:00'),
2002
                pd.Timedelta('3 days 00:00:00'),
2003
                pd.Timedelta('3 days 00:00:00'),
2004
                pd.Timedelta('3 days 00:00:00'),
2005
                np.nan
2006
            ],
2007
                index=pd.Index([
2008
                    'Start', 'End', 'Period', 'Total', 'Rate [%]', 'Total Overlapping',
2009
                    'Overlapping Rate [%]', 'First Index', 'Last Index',
2010
                    'Norm Avg Index [-1, 1]', 'Distance <- Other: Min',
2011
                    'Distance <- Other: Max', 'Distance <- Other: Mean',
2012
                    'Distance <- Other: Std', 'Total Partitions', 'Partition Rate [%]',
2013
                    'Partition Length: Min', 'Partition Length: Max',
2014
                    'Partition Length: Mean', 'Partition Length: Std',
2015
                    'Partition Distance: Min', 'Partition Distance: Max',
2016
                    'Partition Distance: Mean', 'Partition Distance: Std'
2017
                ], dtype='object'),
2018
                name='a'
2019
            )
2020
        )
2021
        pd.testing.assert_series_equal(
2022
            mask.vbt.signals.stats(column='g1', group_by=group_by),
2023
            pd.Series([
2024
                pd.Timestamp('2020-01-01 00:00:00'),
2025
                pd.Timestamp('2020-01-05 00:00:00'),
2026
                pd.Timedelta('5 days 00:00:00'),
2027
                4,
2028
                40.0,
2029
                pd.Timestamp('2020-01-01 00:00:00'),
2030
                pd.Timestamp('2020-01-05 00:00:00'),
2031
                0.0,
2032
                pd.Timedelta('3 days 00:00:00'),
2033
                pd.Timedelta('3 days 00:00:00'),
2034
                pd.Timedelta('3 days 00:00:00'),
2035
                pd.Timedelta('0 days 00:00:00'),
2036
                4,
2037
                100.0,
2038
                pd.Timedelta('1 days 00:00:00'),
2039
                pd.Timedelta('1 days 00:00:00'),
2040
                pd.Timedelta('1 days 00:00:00'),
2041
                pd.Timedelta('0 days 00:00:00'),
2042
                pd.Timedelta('3 days 00:00:00'),
2043
                pd.Timedelta('3 days 00:00:00'),
2044
                pd.Timedelta('3 days 00:00:00'),
2045
                pd.Timedelta('0 days 00:00:00')
2046
            ],
2047
                index=stats_index,
2048
                name='g1'
2049
            )
2050
        )
2051
        pd.testing.assert_series_equal(
2052
            mask['c'].vbt.signals.stats(),
2053
            mask.vbt.signals.stats(column='c')
2054
        )
2055
        pd.testing.assert_series_equal(
2056
            mask['c'].vbt.signals.stats(),
2057
            mask.vbt.signals.stats(column='c', group_by=False)
2058
        )
2059
        pd.testing.assert_series_equal(
2060
            mask.vbt.signals(group_by=group_by)['g2'].stats(),
2061
            mask.vbt.signals(group_by=group_by).stats(column='g2')
2062
        )
2063
        pd.testing.assert_series_equal(
2064
            mask.vbt.signals(group_by=group_by)['g2'].stats(),
2065
            mask.vbt.signals.stats(column='g2', group_by=group_by)
2066
        )
2067
        stats_df = mask.vbt.signals.stats(agg_func=None)
2068
        assert stats_df.shape == (3, 22)
2069
        pd.testing.assert_index_equal(stats_df.index, mask.vbt.wrapper.columns)
2070
        pd.testing.assert_index_equal(stats_df.columns, stats_index)
2071

2072
    @pytest.mark.parametrize(
2073
        "test_func,test_func_pd",
2074
        [
2075
            (lambda x, *args, **kwargs: x.AND(args, **kwargs), lambda x, y: x & y),
2076
            (lambda x, *args, **kwargs: x.OR(args, **kwargs), lambda x, y: x | y),
2077
            (lambda x, *args, **kwargs: x.XOR(args, **kwargs), lambda x, y: x ^ y)
2078
        ],
2079
    )
2080
    def test_logical_funcs(self, test_func, test_func_pd):
2081
        pd.testing.assert_series_equal(
2082
            test_func(mask['a'].vbt.signals, True, [True, False, False, False, False]),
2083
            test_func_pd(test_func_pd(mask['a'], True), [True, False, False, False, False])
2084
        )
2085
        pd.testing.assert_frame_equal(
2086
            test_func(mask['a'].vbt.signals, True, [True, False, False, False, False], concat=True),
2087
            pd.concat((
2088
                test_func_pd(mask['a'], True),
2089
                test_func_pd(mask['a'], [True, False, False, False, False])
2090
            ), axis=1, keys=[0, 1], names=['combine_idx'])
2091
        )
2092
        pd.testing.assert_frame_equal(
2093
            test_func(mask.vbt.signals, True, [[True], [False], [False], [False], [False]]),
2094
            test_func_pd(test_func_pd(mask, True),
2095
                         np.broadcast_to([[True], [False], [False], [False], [False]], (5, 3)))
2096
        )
2097
        pd.testing.assert_frame_equal(
2098
            test_func(mask.vbt.signals, True, [[True], [False], [False], [False], [False]], concat=True),
2099
            pd.concat((
2100
                test_func_pd(mask, True),
2101
                test_func_pd(mask, np.broadcast_to([[True], [False], [False], [False], [False]], (5, 3)))
2102
            ), axis=1, keys=[0, 1], names=['combine_idx'])
2103
        )
2104

2105

2106
# ############# factory.py ############# #
2107

2108

2109
class TestFactory:
2110
    def test_entries(self):
2111
        @njit
2112
        def choice_nb(from_i, to_i, col, ts, in_out, n, arg, temp_idx_arr, kw):
2113
            in_out[from_i, col] = ts[from_i, col] * n + arg + kw
2114
            temp_idx_arr[0] = from_i
2115
            return temp_idx_arr[:1]
2116

2117
        MySignals = vbt.SignalFactory(
2118
            mode='entries',
2119
            input_names=['ts2'],
2120
            in_output_names=['in_out2'],
2121
            param_names=['n2']
2122
        ).from_choice_func(
2123
            entry_choice_func=choice_nb,
2124
            entry_settings=dict(
2125
                pass_inputs=['ts2'],
2126
                pass_in_outputs=['in_out2'],
2127
                pass_params=['n2'],
2128
                pass_kwargs=['temp_idx_arr2', ('kw2', 1000)],
2129
                pass_cache=True
2130
            ),
2131
            in_output_settings=dict(
2132
                in_out2=dict(
2133
                    dtype=np.float64
2134
                )
2135
            ),
2136
            in_out2=np.nan,
2137
            var_args=True
2138
        )
2139
        my_sig = MySignals.run((5,), np.arange(5), [1, 0], 100)
2140
        pd.testing.assert_frame_equal(
2141
            my_sig.entries,
2142
            pd.DataFrame(np.array([
2143
                [True, True],
2144
                [False, False],
2145
                [False, False],
2146
                [False, False],
2147
                [False, False]
2148
            ]), columns=pd.Index([1, 0], dtype='int64', name='custom_n2')
2149
            )
2150
        )
2151
        pd.testing.assert_frame_equal(
2152
            my_sig.in_out2,
2153
            pd.DataFrame(np.array([
2154
                [1100.0, 1100.0],
2155
                [np.nan, np.nan],
2156
                [np.nan, np.nan],
2157
                [np.nan, np.nan],
2158
                [np.nan, np.nan],
2159
            ]), columns=pd.Index([1, 0], dtype='int64', name='custom_n2')
2160
            )
2161
        )
2162

2163
    def test_exits(self):
2164
        @njit
2165
        def choice_nb(from_i, to_i, col, ts, in_out, n, arg, temp_idx_arr, kw):
2166
            in_out[from_i, col] = ts[from_i, col] * n + arg + kw
2167
            temp_idx_arr[0] = from_i
2168
            return temp_idx_arr[:1]
2169

2170
        MySignals = vbt.SignalFactory(
2171
            mode='exits',
2172
            input_names=['ts2'],
2173
            in_output_names=['in_out2'],
2174
            param_names=['n2']
2175
        ).from_choice_func(
2176
            exit_choice_func=choice_nb,
2177
            exit_settings=dict(
2178
                pass_inputs=['ts2'],
2179
                pass_in_outputs=['in_out2'],
2180
                pass_params=['n2'],
2181
                pass_kwargs=['temp_idx_arr2', ('kw2', 1000)],
2182
                pass_cache=True
2183
            ),
2184
            in_output_settings=dict(
2185
                in_out2=dict(
2186
                    dtype=np.float64
2187
                )
2188
            ),
2189
            in_out2=np.nan,
2190
            var_args=True
2191
        )
2192
        e = np.array([True, False, True, False, True])
2193
        my_sig = MySignals.run(e, np.arange(5), [1, 0], 100)
2194
        pd.testing.assert_frame_equal(
2195
            my_sig.entries,
2196
            pd.DataFrame(np.array([
2197
                [True, True],
2198
                [False, False],
2199
                [True, True],
2200
                [False, False],
2201
                [True, True]
2202
            ]), columns=pd.Index([1, 0], dtype='int64', name='custom_n2')
2203
            )
2204
        )
2205
        pd.testing.assert_frame_equal(
2206
            my_sig.exits,
2207
            pd.DataFrame(np.array([
2208
                [False, False],
2209
                [True, True],
2210
                [False, False],
2211
                [True, True],
2212
                [False, False]
2213
            ]), columns=pd.Index([1, 0], dtype='int64', name='custom_n2')
2214
            )
2215
        )
2216
        pd.testing.assert_frame_equal(
2217
            my_sig.in_out2,
2218
            pd.DataFrame(np.array([
2219
                [np.nan, np.nan],
2220
                [1101.0, 1100.0],
2221
                [np.nan, np.nan],
2222
                [1103.0, 1100.0],
2223
                [np.nan, np.nan],
2224
            ]), columns=pd.Index([1, 0], dtype='int64', name='custom_n2')
2225
            )
2226
        )
2227
        e = np.array([True, False, False, True, False, False])
2228
        my_sig = MySignals.run(e, np.arange(6), [1, 0], 100, wait=2)
2229
        pd.testing.assert_frame_equal(
2230
            my_sig.entries,
2231
            pd.DataFrame(np.array([
2232
                [True, True],
2233
                [False, False],
2234
                [False, False],
2235
                [True, True],
2236
                [False, False],
2237
                [False, False]
2238
            ]), columns=pd.Index([1, 0], dtype='int64', name='custom_n2')
2239
            )
2240
        )
2241
        pd.testing.assert_frame_equal(
2242
            my_sig.exits,
2243
            pd.DataFrame(np.array([
2244
                [False, False],
2245
                [False, False],
2246
                [True, True],
2247
                [False, False],
2248
                [False, False],
2249
                [True, True]
2250
            ]), columns=pd.Index([1, 0], dtype='int64', name='custom_n2')
2251
            )
2252
        )
2253
        pd.testing.assert_frame_equal(
2254
            my_sig.in_out2,
2255
            pd.DataFrame(np.array([
2256
                [np.nan, np.nan],
2257
                [np.nan, np.nan],
2258
                [1102.0, 1100.0],
2259
                [np.nan, np.nan],
2260
                [np.nan, np.nan],
2261
                [1105.0, 1100.0]
2262
            ]), columns=pd.Index([1, 0], dtype='int64', name='custom_n2')
2263
            )
2264
        )
2265

2266
    def test_chain(self):
2267
        @njit
2268
        def choice_nb(from_i, to_i, col, ts, in_out, n, arg, temp_idx_arr, kw):
2269
            in_out[from_i, col] = ts[from_i, col] * n + arg + kw
2270
            temp_idx_arr[0] = from_i
2271
            return temp_idx_arr[:1]
2272

2273
        MySignals = vbt.SignalFactory(
2274
            mode='chain',
2275
            input_names=['ts2'],
2276
            in_output_names=['in_out2'],
2277
            param_names=['n2']
2278
        ).from_choice_func(
2279
            exit_choice_func=choice_nb,
2280
            exit_settings=dict(
2281
                pass_inputs=['ts2'],
2282
                pass_in_outputs=['in_out2'],
2283
                pass_params=['n2'],
2284
                pass_kwargs=['temp_idx_arr2', ('kw2', 1000)],
2285
                pass_cache=True
2286
            ),
2287
            in_output_settings=dict(
2288
                in_out2=dict(
2289
                    dtype=np.float64
2290
                )
2291
            ),
2292
            in_out2=np.nan,
2293
            var_args=True
2294
        )
2295
        e = np.array([True, True, True, True, True])
2296
        my_sig = MySignals.run(e, np.arange(5), [1, 0], 100)
2297
        pd.testing.assert_frame_equal(
2298
            my_sig.entries,
2299
            pd.DataFrame(np.array([
2300
                [True, True],
2301
                [True, True],
2302
                [True, True],
2303
                [True, True],
2304
                [True, True]
2305
            ]), columns=pd.Index([1, 0], dtype='int64', name='custom_n2')
2306
            )
2307
        )
2308
        pd.testing.assert_frame_equal(
2309
            my_sig.new_entries,
2310
            pd.DataFrame(np.array([
2311
                [True, True],
2312
                [False, False],
2313
                [True, True],
2314
                [False, False],
2315
                [True, True]
2316
            ]), columns=pd.Index([1, 0], dtype='int64', name='custom_n2')
2317
            )
2318
        )
2319
        pd.testing.assert_frame_equal(
2320
            my_sig.exits,
2321
            pd.DataFrame(np.array([
2322
                [False, False],
2323
                [True, True],
2324
                [False, False],
2325
                [True, True],
2326
                [False, False]
2327
            ]), columns=pd.Index([1, 0], dtype='int64', name='custom_n2')
2328
            )
2329
        )
2330
        pd.testing.assert_frame_equal(
2331
            my_sig.in_out2,
2332
            pd.DataFrame(np.array([
2333
                [np.nan, np.nan],
2334
                [1101.0, 1100.0],
2335
                [np.nan, np.nan],
2336
                [1103.0, 1100.0],
2337
                [np.nan, np.nan],
2338
            ]), columns=pd.Index([1, 0], dtype='int64', name='custom_n2')
2339
            )
2340
        )
2341
        e = np.array([True, True, True, True, True, True])
2342
        my_sig = MySignals.run(e, np.arange(6), [1, 0], 100, wait=2)
2343
        pd.testing.assert_frame_equal(
2344
            my_sig.entries,
2345
            pd.DataFrame(np.array([
2346
                [True, True],
2347
                [True, True],
2348
                [True, True],
2349
                [True, True],
2350
                [True, True],
2351
                [True, True]
2352
            ]), columns=pd.Index([1, 0], dtype='int64', name='custom_n2')
2353
            )
2354
        )
2355
        pd.testing.assert_frame_equal(
2356
            my_sig.new_entries,
2357
            pd.DataFrame(np.array([
2358
                [True, True],
2359
                [False, False],
2360
                [False, False],
2361
                [True, True],
2362
                [False, False],
2363
                [False, False]
2364
            ]), columns=pd.Index([1, 0], dtype='int64', name='custom_n2')
2365
            )
2366
        )
2367
        pd.testing.assert_frame_equal(
2368
            my_sig.exits,
2369
            pd.DataFrame(np.array([
2370
                [False, False],
2371
                [False, False],
2372
                [True, True],
2373
                [False, False],
2374
                [False, False],
2375
                [True, True]
2376
            ]), columns=pd.Index([1, 0], dtype='int64', name='custom_n2')
2377
            )
2378
        )
2379
        pd.testing.assert_frame_equal(
2380
            my_sig.in_out2,
2381
            pd.DataFrame(np.array([
2382
                [np.nan, np.nan],
2383
                [np.nan, np.nan],
2384
                [1102.0, 1100.0],
2385
                [np.nan, np.nan],
2386
                [np.nan, np.nan],
2387
                [1105.0, 1100.0]
2388
            ]), columns=pd.Index([1, 0], dtype='int64', name='custom_n2')
2389
            )
2390
        )
2391

2392
    def test_both(self):
2393
        @njit
2394
        def cache_nb(ts1, ts2, in_out1, in_out2, n1, n2, arg0, temp_idx_arr0, kw0):
2395
            return arg0
2396

2397
        @njit
2398
        def choice_nb(from_i, to_i, col, ts, in_out, n, arg, temp_idx_arr, kw, cache):
2399
            in_out[from_i, col] = ts[from_i, col] * n + arg + kw + cache
2400
            temp_idx_arr[0] = from_i
2401
            return temp_idx_arr[:1]
2402

2403
        MySignals = vbt.SignalFactory(
2404
            input_names=['ts1', 'ts2'],
2405
            in_output_names=['in_out1', 'in_out2'],
2406
            param_names=['n1', 'n2']
2407
        ).from_choice_func(
2408
            cache_func=cache_nb,
2409
            cache_settings=dict(
2410
                pass_inputs=['ts1', 'ts2'],
2411
                pass_in_outputs=['in_out1', 'in_out2'],
2412
                pass_params=['n1', 'n2'],
2413
                pass_kwargs=['temp_idx_arr0', ('kw0', 1000)]
2414
            ),
2415
            entry_choice_func=choice_nb,
2416
            entry_settings=dict(
2417
                pass_inputs=['ts1'],
2418
                pass_in_outputs=['in_out1'],
2419
                pass_params=['n1'],
2420
                pass_kwargs=['temp_idx_arr1', ('kw1', 1000)],
2421
                pass_cache=True
2422
            ),
2423
            exit_choice_func=choice_nb,
2424
            exit_settings=dict(
2425
                pass_inputs=['ts2'],
2426
                pass_in_outputs=['in_out2'],
2427
                pass_params=['n2'],
2428
                pass_kwargs=['temp_idx_arr2', ('kw2', 1000)],
2429
                pass_cache=True
2430
            ),
2431
            in_output_settings=dict(
2432
                in_out1=dict(
2433
                    dtype=np.float64
2434
                ),
2435
                in_out2=dict(
2436
                    dtype=np.float64
2437
                )
2438
            ),
2439
            in_out1=np.nan,
2440
            in_out2=np.nan,
2441
            var_args=True,
2442
            require_input_shape=False
2443
        )
2444
        my_sig = MySignals.run(
2445
            np.arange(5), np.arange(5), [0, 1], [1, 0],
2446
            cache_args=(0,), entry_args=(100,), exit_args=(100,)
2447
        )
2448
        pd.testing.assert_frame_equal(
2449
            my_sig.entries,
2450
            pd.DataFrame(np.array([
2451
                [True, True],
2452
                [False, False],
2453
                [True, True],
2454
                [False, False],
2455
                [True, True]
2456
            ]), columns=pd.MultiIndex.from_tuples(
2457
                [(0, 1), (1, 0)],
2458
                names=['custom_n1', 'custom_n2'])
2459
            )
2460
        )
2461
        pd.testing.assert_frame_equal(
2462
            my_sig.exits,
2463
            pd.DataFrame(np.array([
2464
                [False, False],
2465
                [True, True],
2466
                [False, False],
2467
                [True, True],
2468
                [False, False],
2469
            ]), columns=pd.MultiIndex.from_tuples(
2470
                [(0, 1), (1, 0)],
2471
                names=['custom_n1', 'custom_n2'])
2472
            )
2473
        )
2474
        pd.testing.assert_frame_equal(
2475
            my_sig.in_out1,
2476
            pd.DataFrame(np.array([
2477
                [1100.0, 1100.0],
2478
                [np.nan, np.nan],
2479
                [1100.0, 1102.0],
2480
                [np.nan, np.nan],
2481
                [1100.0, 1104.0]
2482
            ]), columns=pd.MultiIndex.from_tuples(
2483
                [(0, 1), (1, 0)],
2484
                names=['custom_n1', 'custom_n2'])
2485
            )
2486
        )
2487
        pd.testing.assert_frame_equal(
2488
            my_sig.in_out2,
2489
            pd.DataFrame(np.array([
2490
                [np.nan, np.nan],
2491
                [1101.0, 1100.0],
2492
                [np.nan, np.nan],
2493
                [1103.0, 1100.0],
2494
                [np.nan, np.nan],
2495
            ]), columns=pd.MultiIndex.from_tuples(
2496
                [(0, 1), (1, 0)],
2497
                names=['custom_n1', 'custom_n2'])
2498
            )
2499
        )
2500
        my_sig = MySignals.run(
2501
            np.arange(7), np.arange(7), [0, 1], [1, 0],
2502
            cache_args=(0,), entry_args=(100,), exit_args=(100,),
2503
            entry_kwargs=dict(wait=2), exit_kwargs=dict(wait=2)
2504
        )
2505
        pd.testing.assert_frame_equal(
2506
            my_sig.entries,
2507
            pd.DataFrame(np.array([
2508
                [True, True],
2509
                [False, False],
2510
                [False, False],
2511
                [False, False],
2512
                [True, True],
2513
                [False, False],
2514
                [False, False]
2515
            ]), columns=pd.MultiIndex.from_tuples(
2516
                [(0, 1), (1, 0)],
2517
                names=['custom_n1', 'custom_n2'])
2518
            )
2519
        )
2520
        pd.testing.assert_frame_equal(
2521
            my_sig.exits,
2522
            pd.DataFrame(np.array([
2523
                [False, False],
2524
                [False, False],
2525
                [True, True],
2526
                [False, False],
2527
                [False, False],
2528
                [False, False],
2529
                [True, True]
2530
            ]), columns=pd.MultiIndex.from_tuples(
2531
                [(0, 1), (1, 0)],
2532
                names=['custom_n1', 'custom_n2'])
2533
            )
2534
        )
2535
        pd.testing.assert_frame_equal(
2536
            my_sig.in_out1,
2537
            pd.DataFrame(np.array([
2538
                [1100.0, 1100.0],
2539
                [np.nan, np.nan],
2540
                [np.nan, np.nan],
2541
                [np.nan, np.nan],
2542
                [1100.0, 1104.0],
2543
                [np.nan, np.nan],
2544
                [np.nan, np.nan]
2545
            ]), columns=pd.MultiIndex.from_tuples(
2546
                [(0, 1), (1, 0)],
2547
                names=['custom_n1', 'custom_n2'])
2548
            )
2549
        )
2550
        pd.testing.assert_frame_equal(
2551
            my_sig.in_out2,
2552
            pd.DataFrame(np.array([
2553
                [np.nan, np.nan],
2554
                [np.nan, np.nan],
2555
                [1102.0, 1100.0],
2556
                [np.nan, np.nan],
2557
                [np.nan, np.nan],
2558
                [np.nan, np.nan],
2559
                [1106.0, 1100.0]
2560
            ]), columns=pd.MultiIndex.from_tuples(
2561
                [(0, 1), (1, 0)],
2562
                names=['custom_n1', 'custom_n2'])
2563
            )
2564
        )
2565

2566

2567
# ############# generators.py ############# #
2568

2569
class TestGenerators:
2570
    def test_RAND(self):
2571
        rand = vbt.RAND.run(n=1, input_shape=(6,), seed=seed)
2572
        pd.testing.assert_series_equal(
2573
            rand.entries,
2574
            pd.Series(np.array([True, False, False, False, False, False]), name=1)
2575
        )
2576
        rand = vbt.RAND.run(n=[1, 2, 3], input_shape=(6,), seed=seed)
2577
        pd.testing.assert_frame_equal(
2578
            rand.entries,
2579
            pd.DataFrame(np.array([
2580
                [True, True, True],
2581
                [False, False, True],
2582
                [False, False, False],
2583
                [False, True, False],
2584
                [False, False, True],
2585
                [False, False, False]
2586
            ]), columns=pd.Index([1, 2, 3], dtype='int64', name='rand_n')
2587
            )
2588
        )
2589
        rand = vbt.RAND.run(n=[np.array([1, 2]), np.array([3, 4])], input_shape=(8, 2), seed=seed)
2590
        pd.testing.assert_frame_equal(
2591
            rand.entries,
2592
            pd.DataFrame(np.array([
2593
                [False, False, True, False],
2594
                [True, False, False, False],
2595
                [False, False, False, True],
2596
                [False, True, True, False],
2597
                [False, False, False, False],
2598
                [False, False, False, True],
2599
                [False, False, True, True],
2600
                [False, True, False, True]
2601
            ]), columns=pd.MultiIndex.from_tuples([
2602
                (1, 0),
2603
                (2, 1),
2604
                (3, 0),
2605
                (4, 1)
2606
            ], names=['rand_n', None])
2607
            )
2608
        )
2609

2610
    def test_RANDX(self):
2611
        randx = vbt.RANDX.run(mask, seed=seed)
2612
        pd.testing.assert_frame_equal(
2613
            randx.exits,
2614
            pd.DataFrame(np.array([
2615
                [False, False, False],
2616
                [False, False, False],
2617
                [True, True, False],
2618
                [False, False, False],
2619
                [True, False, True]
2620
            ]), columns=mask.columns, index=mask.index)
2621
        )
2622

2623
    def test_RANDNX(self):
2624
        randnx = vbt.RANDNX.run(n=1, input_shape=(6,), seed=seed)
2625
        pd.testing.assert_series_equal(
2626
            randnx.entries,
2627
            pd.Series(np.array([True, False, False, False, False, False]), name=1)
2628
        )
2629
        pd.testing.assert_series_equal(
2630
            randnx.exits,
2631
            pd.Series(np.array([False, True, False, False, False, False]), name=1)
2632
        )
2633
        randnx = vbt.RANDNX.run(n=[1, 2, 3], input_shape=(6,), seed=seed)
2634
        pd.testing.assert_frame_equal(
2635
            randnx.entries,
2636
            pd.DataFrame(np.array([
2637
                [True, True, True],
2638
                [False, False, False],
2639
                [False, True, True],
2640
                [False, False, False],
2641
                [False, False, True],
2642
                [False, False, False]
2643
            ]), columns=pd.Index([1, 2, 3], dtype='int64', name='randnx_n')
2644
            )
2645
        )
2646
        pd.testing.assert_frame_equal(
2647
            randnx.exits,
2648
            pd.DataFrame(np.array([
2649
                [False, False, False],
2650
                [True, True, True],
2651
                [False, False, False],
2652
                [False, True, True],
2653
                [False, False, False],
2654
                [False, False, True]
2655
            ]), columns=pd.Index([1, 2, 3], dtype='int64', name='randnx_n')
2656
            )
2657
        )
2658
        randnx = vbt.RANDNX.run(n=[np.array([1, 2]), np.array([3, 4])], input_shape=(8, 2), seed=seed)
2659
        pd.testing.assert_frame_equal(
2660
            randnx.entries,
2661
            pd.DataFrame(np.array([
2662
                [False, True, True, True],
2663
                [True, False, False, False],
2664
                [False, False, False, True],
2665
                [False, False, True, False],
2666
                [False, True, False, True],
2667
                [False, False, True, False],
2668
                [False, False, False, True],
2669
                [False, False, False, False]
2670
            ]), columns=pd.MultiIndex.from_tuples([
2671
                (1, 0),
2672
                (2, 1),
2673
                (3, 0),
2674
                (4, 1)
2675
            ], names=['randnx_n', None])
2676
            )
2677
        )
2678
        pd.testing.assert_frame_equal(
2679
            randnx.exits,
2680
            pd.DataFrame(np.array([
2681
                [False, False, False, False],
2682
                [False, False, True, True],
2683
                [False, False, False, False],
2684
                [False, True, False, True],
2685
                [False, False, True, False],
2686
                [True, False, False, True],
2687
                [False, False, True, False],
2688
                [False, True, False, True]
2689
            ]), columns=pd.MultiIndex.from_tuples([
2690
                (1, 0),
2691
                (2, 1),
2692
                (3, 0),
2693
                (4, 1)
2694
            ], names=['randnx_n', None])
2695
            )
2696
        )
2697

2698
    def test_RPROB(self):
2699
        rprob = vbt.RPROB.run(prob=1, input_shape=(5,), seed=seed)
2700
        pd.testing.assert_series_equal(
2701
            rprob.entries,
2702
            pd.Series(np.array([True, True, True, True, True]), name=1)
2703
        )
2704
        rprob = vbt.RPROB.run(prob=[0, 0.5, 1], input_shape=(5,), seed=seed)
2705
        pd.testing.assert_frame_equal(
2706
            rprob.entries,
2707
            pd.DataFrame(np.array([
2708
                [False, True, True],
2709
                [False, True, True],
2710
                [False, False, True],
2711
                [False, False, True],
2712
                [False, False, True]
2713
            ]), columns=pd.Index([0, 0.5, 1], dtype='float64', name='rprob_prob')
2714
            )
2715
        )
2716
        rprob = vbt.RPROB.run(prob=[np.array([0, 0.25]), np.array([0.75, 1])], input_shape=(5, 2), seed=seed)
2717
        pd.testing.assert_frame_equal(
2718
            rprob.entries,
2719
            pd.DataFrame(np.array([
2720
                [False, True, True, True],
2721
                [False, True, False, True],
2722
                [False, False, False, True],
2723
                [False, False, True, True],
2724
                [False, False, True, True]
2725
            ]), columns=pd.MultiIndex.from_tuples([
2726
                ('array_0', 0),
2727
                ('array_0', 1),
2728
                ('array_1', 0),
2729
                ('array_1', 1)
2730
            ], names=['rprob_prob', None])
2731
            )
2732
        )
2733

2734
    def test_RPROBX(self):
2735
        rprobx = vbt.RPROBX.run(mask, prob=[0., 0.5, 1.], seed=seed)
2736
        pd.testing.assert_frame_equal(
2737
            rprobx.exits,
2738
            pd.DataFrame(np.array([
2739
                [False, False, False, False, False, False, False, False, False],
2740
                [False, False, False, False, False, False, True, False, False],
2741
                [False, False, False, False, True, False, False, True, False],
2742
                [False, False, False, False, False, False, False, False, True],
2743
                [False, False, False, False, False, False, True, False, False]
2744
            ]), index=mask.index, columns=pd.MultiIndex.from_tuples([
2745
                (0.0, 'a'),
2746
                (0.0, 'b'),
2747
                (0.0, 'c'),
2748
                (0.5, 'a'),
2749
                (0.5, 'b'),
2750
                (0.5, 'c'),
2751
                (1.0, 'a'),
2752
                (1.0, 'b'),
2753
                (1.0, 'c')
2754
            ], names=['rprobx_prob', None])
2755
            )
2756
        )
2757

2758
    def test_RPROBCX(self):
2759
        rprobcx = vbt.RPROBCX.run(mask, prob=[0., 0.5, 1.], seed=seed)
2760
        pd.testing.assert_frame_equal(
2761
            rprobcx.new_entries,
2762
            pd.DataFrame(np.array([
2763
                [True, False, False, True, False, False, True, False, False],
2764
                [False, True, False, False, True, False, False, True, False],
2765
                [False, False, True, False, False, True, False, False, True],
2766
                [False, False, False, True, False, False, True, False, False],
2767
                [False, False, False, False, True, False, False, True, False]
2768
            ]), index=mask.index, columns=pd.MultiIndex.from_tuples([
2769
                (0.0, 'a'),
2770
                (0.0, 'b'),
2771
                (0.0, 'c'),
2772
                (0.5, 'a'),
2773
                (0.5, 'b'),
2774
                (0.5, 'c'),
2775
                (1.0, 'a'),
2776
                (1.0, 'b'),
2777
                (1.0, 'c')
2778
            ], names=['rprobcx_prob', None])
2779
            )
2780
        )
2781
        pd.testing.assert_frame_equal(
2782
            rprobcx.exits,
2783
            pd.DataFrame(np.array([
2784
                [False, False, False, False, False, False, False, False, False],
2785
                [False, False, False, False, False, False, True, False, False],
2786
                [False, False, False, True, False, False, False, True, False],
2787
                [False, False, False, False, True, True, False, False, True],
2788
                [False, False, False, False, False, False, True, False, False]
2789
            ]), index=mask.index, columns=pd.MultiIndex.from_tuples([
2790
                (0.0, 'a'),
2791
                (0.0, 'b'),
2792
                (0.0, 'c'),
2793
                (0.5, 'a'),
2794
                (0.5, 'b'),
2795
                (0.5, 'c'),
2796
                (1.0, 'a'),
2797
                (1.0, 'b'),
2798
                (1.0, 'c')
2799
            ], names=['rprobcx_prob', None])
2800
            )
2801
        )
2802

2803
    def test_RPROBNX(self):
2804
        rprobnx = vbt.RPROBNX.run(entry_prob=1., exit_prob=1., input_shape=(5,), seed=seed)
2805
        pd.testing.assert_series_equal(
2806
            rprobnx.entries,
2807
            pd.Series(np.array([True, False, True, False, True]), name=(1.0, 1.0))
2808
        )
2809
        pd.testing.assert_series_equal(
2810
            rprobnx.exits,
2811
            pd.Series(np.array([False, True, False, True, False]), name=(1.0, 1.0))
2812
        )
2813
        rprobnx = vbt.RPROBNX.run(
2814
            entry_prob=np.asarray([1., 0., 1., 0., 1.]),
2815
            exit_prob=np.asarray([0., 1., 0., 1., 0.]),
2816
            input_shape=(5,), seed=seed)
2817
        pd.testing.assert_series_equal(
2818
            rprobnx.entries,
2819
            pd.Series(np.array([True, False, True, False, True]), name=('array_0', 'array_0'))
2820
        )
2821
        pd.testing.assert_series_equal(
2822
            rprobnx.exits,
2823
            pd.Series(np.array([False, True, False, True, False]), name=('array_0', 'array_0'))
2824
        )
2825
        rprobnx = vbt.RPROBNX.run(entry_prob=[0.5, 1.], exit_prob=[1., 0.5], input_shape=(5,), seed=seed)
2826
        pd.testing.assert_frame_equal(
2827
            rprobnx.entries,
2828
            pd.DataFrame(np.array([
2829
                [True, True],
2830
                [False, False],
2831
                [False, True],
2832
                [False, False],
2833
                [True, False]
2834
            ]), columns=pd.MultiIndex.from_tuples(
2835
                [(0.5, 1.0), (1.0, 0.5)],
2836
                names=['rprobnx_entry_prob', 'rprobnx_exit_prob'])
2837
            )
2838
        )
2839
        pd.testing.assert_frame_equal(
2840
            rprobnx.exits,
2841
            pd.DataFrame(np.array([
2842
                [False, False],
2843
                [True, True],
2844
                [False, False],
2845
                [False, False],
2846
                [False, False]
2847
            ]), columns=pd.MultiIndex.from_tuples(
2848
                [(0.5, 1.0), (1.0, 0.5)],
2849
                names=['rprobnx_entry_prob', 'rprobnx_exit_prob'])
2850
            )
2851
        )
2852

2853
    def test_STX(self):
2854
        stx = vbt.STX.run(mask, ts, 0.1)
2855
        pd.testing.assert_frame_equal(
2856
            stx.exits,
2857
            pd.DataFrame(np.array([
2858
                [False, False, False],
2859
                [True, False, False],
2860
                [False, True, False],
2861
                [False, False, False],
2862
                [False, False, False]
2863
            ]), index=mask.index, columns=pd.MultiIndex.from_tuples([
2864
                (0.1, 'a'),
2865
                (0.1, 'b'),
2866
                (0.1, 'c')
2867
            ], names=['stx_stop', None])
2868
            )
2869
        )
2870
        stx = vbt.STX.run(mask, ts, np.asarray([0.1, 0.1, -0.1, -0.1, -0.1])[:, None])
2871
        pd.testing.assert_frame_equal(
2872
            stx.exits,
2873
            pd.DataFrame(np.array([
2874
                [False, False, False],
2875
                [True, False, False],
2876
                [False, True, False],
2877
                [False, False, True],
2878
                [True, False, False]
2879
            ]), index=mask.index, columns=pd.MultiIndex.from_tuples([
2880
                ('array_0', 'a'),
2881
                ('array_0', 'b'),
2882
                ('array_0', 'c')
2883
            ], names=['stx_stop', None])
2884
            )
2885
        )
2886
        stx = vbt.STX.run(mask, ts, [0.1, 0.1, -0.1, -0.1], trailing=[False, True, False, True])
2887
        pd.testing.assert_frame_equal(
2888
            stx.exits,
2889
            pd.DataFrame(np.array([
2890
                [False, False, False, False, False, False, False, False, False, False, False, False],
2891
                [True, False, False, True, False, False, False, False, False, False, False, False],
2892
                [False, True, False, False, True, False, False, False, False, False, False, False],
2893
                [False, False, False, False, False, False, False, False, True, False, True, True],
2894
                [False, False, False, False, False, False, True, False, False, True, False, False]
2895
            ]), index=mask.index, columns=pd.MultiIndex.from_tuples([
2896
                (0.1, False, 'a'),
2897
                (0.1, False, 'b'),
2898
                (0.1, False, 'c'),
2899
                (0.1, True, 'a'),
2900
                (0.1, True, 'b'),
2901
                (0.1, True, 'c'),
2902
                (-0.1, False, 'a'),
2903
                (-0.1, False, 'b'),
2904
                (-0.1, False, 'c'),
2905
                (-0.1, True, 'a'),
2906
                (-0.1, True, 'b'),
2907
                (-0.1, True, 'c')
2908
            ], names=['stx_stop', 'stx_trailing', None])
2909
            )
2910
        )
2911

2912
    def test_STCX(self):
2913
        stcx = vbt.STCX.run(mask, ts, [0.1, 0.1, -0.1, -0.1], trailing=[False, True, False, True])
2914
        pd.testing.assert_frame_equal(
2915
            stcx.new_entries,
2916
            pd.DataFrame(np.array([
2917
                [True, False, False, True, False, False, True, False, False, True, False, False],
2918
                [False, True, False, False, True, False, False, True, False, False, True, False],
2919
                [False, False, True, False, False, True, False, False, True, False, False, True],
2920
                [True, False, False, True, False, False, False, False, False, False, False, False],
2921
                [False, True, False, False, True, False, False, False, False, False, True, False]
2922
            ]), index=mask.index, columns=pd.MultiIndex.from_tuples([
2923
                (0.1, False, 'a'),
2924
                (0.1, False, 'b'),
2925
                (0.1, False, 'c'),
2926
                (0.1, True, 'a'),
2927
                (0.1, True, 'b'),
2928
                (0.1, True, 'c'),
2929
                (-0.1, False, 'a'),
2930
                (-0.1, False, 'b'),
2931
                (-0.1, False, 'c'),
2932
                (-0.1, True, 'a'),
2933
                (-0.1, True, 'b'),
2934
                (-0.1, True, 'c')
2935
            ], names=['stcx_stop', 'stcx_trailing', None])
2936
            )
2937
        )
2938
        pd.testing.assert_frame_equal(
2939
            stcx.exits,
2940
            pd.DataFrame(np.array([
2941
                [False, False, False, False, False, False, False, False, False, False, False, False],
2942
                [True, False, False, True, False, False, False, False, False, False, False, False],
2943
                [False, True, False, False, True, False, False, False, False, False, False, False],
2944
                [False, False, False, False, False, False, False, False, True, True, True, True],
2945
                [False, False, False, False, False, False, False, True, False, False, False, False]
2946
            ]), index=mask.index, columns=pd.MultiIndex.from_tuples([
2947
                (0.1, False, 'a'),
2948
                (0.1, False, 'b'),
2949
                (0.1, False, 'c'),
2950
                (0.1, True, 'a'),
2951
                (0.1, True, 'b'),
2952
                (0.1, True, 'c'),
2953
                (-0.1, False, 'a'),
2954
                (-0.1, False, 'b'),
2955
                (-0.1, False, 'c'),
2956
                (-0.1, True, 'a'),
2957
                (-0.1, True, 'b'),
2958
                (-0.1, True, 'c')
2959
            ], names=['stcx_stop', 'stcx_trailing', None])
2960
            )
2961
        )
2962

2963
    def test_OHLCSTX(self):
2964
        ohlcstx = vbt.OHLCSTX.run(
2965
            mask, price['open'], price['high'], price['low'], price['close'],
2966
            sl_stop=0.1
2967
        )
2968
        pd.testing.assert_frame_equal(
2969
            ohlcstx.exits,
2970
            pd.DataFrame(np.array([
2971
                [False, False, False],
2972
                [False, False, False],
2973
                [False, False, False],
2974
                [False, False, True],
2975
                [True, False, False]
2976
            ]), index=mask.index, columns=pd.MultiIndex.from_tuples([
2977
                (0.1, 'a'),
2978
                (0.1, 'b'),
2979
                (0.1, 'c')
2980
            ], names=['ohlcstx_sl_stop', None])
2981
            )
2982
        )
2983
        pd.testing.assert_frame_equal(
2984
            ohlcstx.stop_price,
2985
            pd.DataFrame(np.array([
2986
                [np.nan, np.nan, np.nan],
2987
                [np.nan, np.nan, np.nan],
2988
                [np.nan, np.nan, np.nan],
2989
                [np.nan, np.nan, 10.8],
2990
                [9.9, np.nan, np.nan]
2991
            ]), index=mask.index, columns=pd.MultiIndex.from_tuples([
2992
                (0.1, 'a'),
2993
                (0.1, 'b'),
2994
                (0.1, 'c')
2995
            ], names=['ohlcstx_sl_stop', None])
2996
            )
2997
        )
2998
        pd.testing.assert_frame_equal(
2999
            ohlcstx.stop_type,
3000
            pd.DataFrame(np.array([
3001
                [-1, -1, -1],
3002
                [-1, -1, -1],
3003
                [-1, -1, -1],
3004
                [-1, -1, 0],
3005
                [0, -1, -1]
3006
            ]), index=mask.index, columns=pd.MultiIndex.from_tuples([
3007
                (0.1, 'a'),
3008
                (0.1, 'b'),
3009
                (0.1, 'c')
3010
            ], names=['ohlcstx_sl_stop', None])
3011
            )
3012
        )
3013
        ohlcstx = vbt.OHLCSTX.run(
3014
            mask, price['open'], price['high'], price['low'], price['close'],
3015
            sl_stop=[0.1, 0.1, np.nan], sl_trail=[False, True, False], tp_stop=[np.nan, np.nan, 0.1]
3016
        )
3017
        pd.testing.assert_frame_equal(
3018
            ohlcstx.exits,
3019
            pd.DataFrame(np.array([
3020
                [False, False, False, False, False, False, False, False, False],
3021
                [False, False, False, False, False, False, True, False, False],
3022
                [False, False, False, False, False, False, False, True, False],
3023
                [False, False, True, False, True, True, False, False, False],
3024
                [True, False, False, True, False, False, False, False, False]
3025
            ]), index=mask.index, columns=pd.MultiIndex.from_tuples([
3026
                (0.1, False, np.nan, 'a'),
3027
                (0.1, False, np.nan, 'b'),
3028
                (0.1, False, np.nan, 'c'),
3029
                (0.1, True, np.nan, 'a'),
3030
                (0.1, True, np.nan, 'b'),
3031
                (0.1, True, np.nan, 'c'),
3032
                (np.nan, False, 0.1, 'a'),
3033
                (np.nan, False, 0.1, 'b'),
3034
                (np.nan, False, 0.1, 'c')
3035
            ], names=['ohlcstx_sl_stop', 'ohlcstx_sl_trail', 'ohlcstx_tp_stop', None])
3036
            )
3037
        )
3038
        pd.testing.assert_frame_equal(
3039
            ohlcstx.stop_price,
3040
            pd.DataFrame(np.array([
3041
                [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan],
3042
                [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, 11., np.nan, np.nan],
3043
                [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, 12.1, np.nan],
3044
                [np.nan, np.nan, 10.8, np.nan, 11.7, 10.8, np.nan, np.nan, np.nan],
3045
                [9.9, np.nan, np.nan, 9.9, np.nan, np.nan, np.nan, np.nan, np.nan]
3046
            ]), index=mask.index, columns=pd.MultiIndex.from_tuples([
3047
                (0.1, False, np.nan, 'a'),
3048
                (0.1, False, np.nan, 'b'),
3049
                (0.1, False, np.nan, 'c'),
3050
                (0.1, True, np.nan, 'a'),
3051
                (0.1, True, np.nan, 'b'),
3052
                (0.1, True, np.nan, 'c'),
3053
                (np.nan, False, 0.1, 'a'),
3054
                (np.nan, False, 0.1, 'b'),
3055
                (np.nan, False, 0.1, 'c')
3056
            ], names=['ohlcstx_sl_stop', 'ohlcstx_sl_trail', 'ohlcstx_tp_stop', None])
3057
            )
3058
        )
3059
        pd.testing.assert_frame_equal(
3060
            ohlcstx.stop_type,
3061
            pd.DataFrame(np.array([
3062
                [-1, -1, -1, -1, -1, -1, -1, -1, -1],
3063
                [-1, -1, -1, -1, -1, -1, 2, -1, -1],
3064
                [-1, -1, -1, -1, -1, -1, -1, 2, -1],
3065
                [-1, -1, 0, -1, 1, 1, -1, -1, -1],
3066
                [0, -1, -1, 1, -1, -1, -1, -1, -1]
3067
            ]), index=mask.index, columns=pd.MultiIndex.from_tuples([
3068
                (0.1, False, np.nan, 'a'),
3069
                (0.1, False, np.nan, 'b'),
3070
                (0.1, False, np.nan, 'c'),
3071
                (0.1, True, np.nan, 'a'),
3072
                (0.1, True, np.nan, 'b'),
3073
                (0.1, True, np.nan, 'c'),
3074
                (np.nan, False, 0.1, 'a'),
3075
                (np.nan, False, 0.1, 'b'),
3076
                (np.nan, False, 0.1, 'c')
3077
            ], names=['ohlcstx_sl_stop', 'ohlcstx_sl_trail', 'ohlcstx_tp_stop', None])
3078
            )
3079
        )
3080
        np.testing.assert_array_equal(
3081
            vbt.OHLCSTX.run(
3082
                mask, price['open'], price['high'], price['low'], price['close'],
3083
                sl_stop=[0.1, np.nan], sl_trail=False, tp_stop=[np.nan, 0.1], reverse=False
3084
            ).exits.values,
3085
            vbt.OHLCSTX.run(
3086
                mask, price['open'], price['high'], price['low'], price['close'],
3087
                sl_stop=[np.nan, 0.1], sl_trail=False, tp_stop=[0.1, np.nan], reverse=True
3088
            ).exits.values
3089
        )
3090

3091
    def test_OHLCSTCX(self):
3092
        ohlcstcx = vbt.OHLCSTCX.run(
3093
            mask, price['open'], price['high'], price['low'], price['close'],
3094
            sl_stop=[0.1, 0.1, np.nan], sl_trail=[False, True, False], tp_stop=[np.nan, np.nan, 0.1]
3095
        )
3096
        pd.testing.assert_frame_equal(
3097
            ohlcstcx.exits,
3098
            pd.DataFrame(np.array([
3099
                [False, False, False, False, False, False, False, False, False],
3100
                [False, False, False, False, False, False, True, False, False],
3101
                [False, False, False, False, False, False, False, True, False],
3102
                [False, False, True, True, True, True, False, False, False],
3103
                [True, True, False, False, False, False, False, False, False]
3104
            ]), index=mask.index, columns=pd.MultiIndex.from_tuples([
3105
                (0.1, False, np.nan, 'a'),
3106
                (0.1, False, np.nan, 'b'),
3107
                (0.1, False, np.nan, 'c'),
3108
                (0.1, True, np.nan, 'a'),
3109
                (0.1, True, np.nan, 'b'),
3110
                (0.1, True, np.nan, 'c'),
3111
                (np.nan, False, 0.1, 'a'),
3112
                (np.nan, False, 0.1, 'b'),
3113
                (np.nan, False, 0.1, 'c')
3114
            ], names=['ohlcstcx_sl_stop', 'ohlcstcx_sl_trail', 'ohlcstcx_tp_stop', None])
3115
            )
3116
        )
3117
        pd.testing.assert_frame_equal(
3118
            ohlcstcx.stop_price,
3119
            pd.DataFrame(np.array([
3120
                [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan],
3121
                [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, 11., np.nan, np.nan],
3122
                [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, 12.1, np.nan],
3123
                [np.nan, np.nan, 10.8, 11.7, 11.7, 10.8, np.nan, np.nan, np.nan],
3124
                [9.0, 9.9, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan]
3125
            ]), index=mask.index, columns=pd.MultiIndex.from_tuples([
3126
                (0.1, False, np.nan, 'a'),
3127
                (0.1, False, np.nan, 'b'),
3128
                (0.1, False, np.nan, 'c'),
3129
                (0.1, True, np.nan, 'a'),
3130
                (0.1, True, np.nan, 'b'),
3131
                (0.1, True, np.nan, 'c'),
3132
                (np.nan, False, 0.1, 'a'),
3133
                (np.nan, False, 0.1, 'b'),
3134
                (np.nan, False, 0.1, 'c')
3135
            ], names=['ohlcstcx_sl_stop', 'ohlcstcx_sl_trail', 'ohlcstcx_tp_stop', None])
3136
            )
3137
        )
3138
        pd.testing.assert_frame_equal(
3139
            ohlcstcx.stop_type,
3140
            pd.DataFrame(np.array([
3141
                [-1, -1, -1, -1, -1, -1, -1, -1, -1],
3142
                [-1, -1, -1, -1, -1, -1, 2, -1, -1],
3143
                [-1, -1, -1, -1, -1, -1, -1, 2, -1],
3144
                [-1, -1, 0, 1, 1, 1, -1, -1, -1],
3145
                [0, 0, -1, -1, -1, -1, -1, -1, -1]
3146
            ]), index=mask.index, columns=pd.MultiIndex.from_tuples([
3147
                (0.1, False, np.nan, 'a'),
3148
                (0.1, False, np.nan, 'b'),
3149
                (0.1, False, np.nan, 'c'),
3150
                (0.1, True, np.nan, 'a'),
3151
                (0.1, True, np.nan, 'b'),
3152
                (0.1, True, np.nan, 'c'),
3153
                (np.nan, False, 0.1, 'a'),
3154
                (np.nan, False, 0.1, 'b'),
3155
                (np.nan, False, 0.1, 'c')
3156
            ], names=['ohlcstcx_sl_stop', 'ohlcstcx_sl_trail', 'ohlcstcx_tp_stop', None])
3157
            )
3158
        )
3159

3160
Product

Resources

Company
