CoCalc -- test_records.py

GitHub Repository: polakowo/vectorbt
Path: blob/master/tests/test_records.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.enums import range_dt, drawdown_dt
11
from vectorbt.portfolio.enums import order_dt, trade_dt, log_dt
12

13
day_dt = np.timedelta64(86400000000000)
14

15
example_dt = np.dtype([
16
    ('id', np.int64),
17
    ('col', np.int64),
18
    ('idx', np.int64),
19
    ('some_field1', np.float64),
20
    ('some_field2', np.float64)
21
], align=True)
22

23
records_arr = np.asarray([
24
    (0, 0, 0, 10, 21),
25
    (1, 0, 1, 11, 20),
26
    (2, 0, 2, 12, 19),
27
    (3, 1, 0, 13, 18),
28
    (4, 1, 1, 14, 17),
29
    (5, 1, 2, 13, 18),
30
    (6, 2, 0, 12, 19),
31
    (7, 2, 1, 11, 20),
32
    (8, 2, 2, 10, 21)
33
], dtype=example_dt)
34
records_nosort_arr = np.concatenate((
35
    records_arr[0::3],
36
    records_arr[1::3],
37
    records_arr[2::3]
38
))
39

40
group_by = pd.Index(['g1', 'g1', 'g2', 'g2'])
41

42
wrapper = vbt.ArrayWrapper(
43
    index=['x', 'y', 'z'],
44
    columns=['a', 'b', 'c', 'd'],
45
    ndim=2,
46
    freq='1 days'
47
)
48
wrapper_grouped = wrapper.replace(group_by=group_by)
49

50
records = vbt.records.Records(wrapper, records_arr)
51
records_grouped = vbt.records.Records(wrapper_grouped, records_arr)
52
records_nosort = records.replace(records_arr=records_nosort_arr)
53
records_nosort_grouped = vbt.records.Records(wrapper_grouped, records_nosort_arr)
54

55

56
# ############# Global ############# #
57

58
def setup_module():
59
    vbt.settings.numba['check_func_suffix'] = True
60
    vbt.settings.caching.enabled = False
61
    vbt.settings.caching.whitelist = []
62
    vbt.settings.caching.blacklist = []
63

64

65
def teardown_module():
66
    vbt.settings.reset()
67

68

69
# ############# col_mapper.py ############# #
70

71

72
class TestColumnMapper:
73
    def test_col_arr(self):
74
        np.testing.assert_array_equal(
75
            records['a'].col_mapper.col_arr,
76
            np.array([0, 0, 0])
77
        )
78
        np.testing.assert_array_equal(
79
            records.col_mapper.col_arr,
80
            np.array([0, 0, 0, 1, 1, 1, 2, 2, 2])
81
        )
82

83
    def test_get_col_arr(self):
84
        np.testing.assert_array_equal(
85
            records.col_mapper.get_col_arr(),
86
            records.col_mapper.col_arr
87
        )
88
        np.testing.assert_array_equal(
89
            records_grouped['g1'].col_mapper.get_col_arr(),
90
            np.array([0, 0, 0, 0, 0, 0])
91
        )
92
        np.testing.assert_array_equal(
93
            records_grouped.col_mapper.get_col_arr(),
94
            np.array([0, 0, 0, 0, 0, 0, 1, 1, 1])
95
        )
96

97
    def test_col_range(self):
98
        np.testing.assert_array_equal(
99
            records['a'].col_mapper.col_range,
100
            np.array([
101
                [0, 3]
102
            ])
103
        )
104
        np.testing.assert_array_equal(
105
            records.col_mapper.col_range,
106
            np.array([
107
                [0, 3],
108
                [3, 6],
109
                [6, 9],
110
                [-1, -1]
111
            ])
112
        )
113

114
    def test_get_col_range(self):
115
        np.testing.assert_array_equal(
116
            records.col_mapper.get_col_range(),
117
            np.array([
118
                [0, 3],
119
                [3, 6],
120
                [6, 9],
121
                [-1, -1]
122
            ])
123
        )
124
        np.testing.assert_array_equal(
125
            records_grouped['g1'].col_mapper.get_col_range(),
126
            np.array([[0, 6]])
127
        )
128
        np.testing.assert_array_equal(
129
            records_grouped.col_mapper.get_col_range(),
130
            np.array([[0, 6], [6, 9]])
131
        )
132

133
    def test_col_map(self):
134
        np.testing.assert_array_equal(
135
            records['a'].col_mapper.col_map[0],
136
            np.array([0, 1, 2])
137
        )
138
        np.testing.assert_array_equal(
139
            records['a'].col_mapper.col_map[1],
140
            np.array([3])
141
        )
142
        np.testing.assert_array_equal(
143
            records.col_mapper.col_map[0],
144
            np.array([0, 1, 2, 3, 4, 5, 6, 7, 8])
145
        )
146
        np.testing.assert_array_equal(
147
            records.col_mapper.col_map[1],
148
            np.array([3, 3, 3, 0])
149
        )
150

151
    def test_get_col_map(self):
152
        np.testing.assert_array_equal(
153
            records.col_mapper.get_col_map()[0],
154
            records.col_mapper.col_map[0]
155
        )
156
        np.testing.assert_array_equal(
157
            records.col_mapper.get_col_map()[1],
158
            records.col_mapper.col_map[1]
159
        )
160
        np.testing.assert_array_equal(
161
            records_grouped['g1'].col_mapper.get_col_map()[0],
162
            np.array([0, 1, 2, 3, 4, 5])
163
        )
164
        np.testing.assert_array_equal(
165
            records_grouped['g1'].col_mapper.get_col_map()[1],
166
            np.array([6])
167
        )
168
        np.testing.assert_array_equal(
169
            records_grouped.col_mapper.get_col_map()[0],
170
            np.array([0, 1, 2, 3, 4, 5, 6, 7, 8])
171
        )
172
        np.testing.assert_array_equal(
173
            records_grouped.col_mapper.get_col_map()[1],
174
            np.array([6, 3])
175
        )
176

177
    def test_is_sorted(self):
178
        assert records.col_mapper.is_sorted()
179
        assert not records_nosort.col_mapper.is_sorted()
180

181

182
# ############# mapped_array.py ############# #
183

184
mapped_array = records.map_field('some_field1')
185
mapped_array_grouped = records_grouped.map_field('some_field1')
186
mapped_array_nosort = records_nosort.map_field('some_field1')
187
mapped_array_nosort_grouped = records_nosort_grouped.map_field('some_field1')
188
mapping = {x: 'test_' + str(x) for x in pd.unique(mapped_array.values)}
189
mp_mapped_array = mapped_array.replace(mapping=mapping)
190
mp_mapped_array_grouped = mapped_array_grouped.replace(mapping=mapping)
191

192

193
class TestMappedArray:
194
    def test_config(self, tmp_path):
195
        assert vbt.MappedArray.loads(mapped_array.dumps()) == mapped_array
196
        mapped_array.save(tmp_path / 'mapped_array')
197
        assert vbt.MappedArray.load(tmp_path / 'mapped_array') == mapped_array
198

199
    def test_mapped_arr(self):
200
        np.testing.assert_array_equal(
201
            mapped_array['a'].values,
202
            np.array([10., 11., 12.])
203
        )
204
        np.testing.assert_array_equal(
205
            mapped_array.values,
206
            np.array([10., 11., 12., 13., 14., 13., 12., 11., 10.])
207
        )
208

209
    def test_id_arr(self):
210
        np.testing.assert_array_equal(
211
            mapped_array['a'].id_arr,
212
            np.array([0, 1, 2])
213
        )
214
        np.testing.assert_array_equal(
215
            mapped_array.id_arr,
216
            np.array([0, 1, 2, 3, 4, 5, 6, 7, 8])
217
        )
218

219
    def test_col_arr(self):
220
        np.testing.assert_array_equal(
221
            mapped_array['a'].col_arr,
222
            np.array([0, 0, 0])
223
        )
224
        np.testing.assert_array_equal(
225
            mapped_array.col_arr,
226
            np.array([0, 0, 0, 1, 1, 1, 2, 2, 2])
227
        )
228

229
    def test_idx_arr(self):
230
        np.testing.assert_array_equal(
231
            mapped_array['a'].idx_arr,
232
            np.array([0, 1, 2])
233
        )
234
        np.testing.assert_array_equal(
235
            mapped_array.idx_arr,
236
            np.array([0, 1, 2, 0, 1, 2, 0, 1, 2])
237
        )
238

239
    def test_is_sorted(self):
240
        assert mapped_array.is_sorted()
241
        assert mapped_array.is_sorted(incl_id=True)
242
        assert not mapped_array_nosort.is_sorted()
243
        assert not mapped_array_nosort.is_sorted(incl_id=True)
244

245
    def test_sort(self):
246
        assert mapped_array.sort().is_sorted()
247
        assert mapped_array.sort().is_sorted(incl_id=True)
248
        assert mapped_array.sort(incl_id=True).is_sorted(incl_id=True)
249
        assert mapped_array_nosort.sort().is_sorted()
250
        assert mapped_array_nosort.sort().is_sorted(incl_id=True)
251
        assert mapped_array_nosort.sort(incl_id=True).is_sorted(incl_id=True)
252

253
    def test_apply_mask(self):
254
        mask_a = mapped_array['a'].values >= mapped_array['a'].values.mean()
255
        np.testing.assert_array_equal(
256
            mapped_array['a'].apply_mask(mask_a).id_arr,
257
            np.array([1, 2])
258
        )
259
        mask = mapped_array.values >= mapped_array.values.mean()
260
        filtered = mapped_array.apply_mask(mask)
261
        np.testing.assert_array_equal(
262
            filtered.id_arr,
263
            np.array([2, 3, 4, 5, 6])
264
        )
265
        np.testing.assert_array_equal(filtered.col_arr, mapped_array.col_arr[mask])
266
        np.testing.assert_array_equal(filtered.idx_arr, mapped_array.idx_arr[mask])
267
        assert mapped_array_grouped.apply_mask(mask).wrapper == mapped_array_grouped.wrapper
268
        assert mapped_array_grouped.apply_mask(mask, group_by=False).wrapper.grouper.group_by is None
269

270
    def test_map_to_mask(self):
271
        @njit
272
        def every_2_nb(inout, idxs, col, mapped_arr):
273
            inout[idxs[::2]] = True
274

275
        np.testing.assert_array_equal(
276
            mapped_array.map_to_mask(every_2_nb),
277
            np.array([True, False, True, True, False, True, True, False, True])
278
        )
279

280
    def test_top_n_mask(self):
281
        np.testing.assert_array_equal(
282
            mapped_array.top_n_mask(1),
283
            np.array([False, False, True, False, True, False, True, False, False])
284
        )
285

286
    def test_bottom_n_mask(self):
287
        np.testing.assert_array_equal(
288
            mapped_array.bottom_n_mask(1),
289
            np.array([True, False, False, True, False, False, False, False, True])
290
        )
291

292
    def test_top_n(self):
293
        np.testing.assert_array_equal(
294
            mapped_array.top_n(1).id_arr,
295
            np.array([2, 4, 6])
296
        )
297

298
    def test_bottom_n(self):
299
        np.testing.assert_array_equal(
300
            mapped_array.bottom_n(1).id_arr,
301
            np.array([0, 3, 8])
302
        )
303

304
    def test_to_pd(self):
305
        target = pd.DataFrame(
306
            np.array([
307
                [10., 13., 12., np.nan],
308
                [11., 14., 11., np.nan],
309
                [12., 13., 10., np.nan]
310
            ]),
311
            index=wrapper.index,
312
            columns=wrapper.columns
313
        )
314
        pd.testing.assert_series_equal(
315
            mapped_array['a'].to_pd(),
316
            target['a']
317
        )
318
        pd.testing.assert_frame_equal(
319
            mapped_array.to_pd(),
320
            target
321
        )
322
        pd.testing.assert_frame_equal(
323
            mapped_array.to_pd(fill_value=0.),
324
            target.fillna(0.)
325
        )
326
        mapped_array2 = vbt.MappedArray(
327
            wrapper,
328
            records_arr['some_field1'].tolist() + [1],
329
            records_arr['col'].tolist() + [2],
330
            idx_arr=records_arr['idx'].tolist() + [2]
331
        )
332
        with pytest.raises(Exception):
333
            _ = mapped_array2.to_pd()
334
        pd.testing.assert_series_equal(
335
            mapped_array['a'].to_pd(ignore_index=True),
336
            pd.Series(np.array([10., 11., 12.]), name='a')
337
        )
338
        pd.testing.assert_frame_equal(
339
            mapped_array.to_pd(ignore_index=True),
340
            pd.DataFrame(
341
                np.array([
342
                    [10., 13., 12., np.nan],
343
                    [11., 14., 11., np.nan],
344
                    [12., 13., 10., np.nan]
345
                ]),
346
                columns=wrapper.columns
347
            )
348
        )
349
        pd.testing.assert_frame_equal(
350
            mapped_array.to_pd(fill_value=0, ignore_index=True),
351
            pd.DataFrame(
352
                np.array([
353
                    [10., 13., 12., 0.],
354
                    [11., 14., 11., 0.],
355
                    [12., 13., 10., 0.]
356
                ]),
357
                columns=wrapper.columns
358
            )
359
        )
360
        pd.testing.assert_frame_equal(
361
            mapped_array_grouped.to_pd(ignore_index=True),
362
            pd.DataFrame(
363
                np.array([
364
                    [10., 12.],
365
                    [11., 11.],
366
                    [12., 10.],
367
                    [13., np.nan],
368
                    [14., np.nan],
369
                    [13., np.nan],
370
                ]),
371
                columns=pd.Index(['g1', 'g2'], dtype='object')
372
            )
373
        )
374

375
    def test_apply(self):
376
        @njit
377
        def cumsum_apply_nb(idxs, col, a):
378
            return np.cumsum(a)
379

380
        np.testing.assert_array_equal(
381
            mapped_array['a'].apply(cumsum_apply_nb).values,
382
            np.array([10., 21., 33.])
383
        )
384
        np.testing.assert_array_equal(
385
            mapped_array.apply(cumsum_apply_nb).values,
386
            np.array([10., 21., 33., 13., 27., 40., 12., 23., 33.])
387
        )
388
        np.testing.assert_array_equal(
389
            mapped_array_grouped.apply(cumsum_apply_nb, apply_per_group=False).values,
390
            np.array([10., 21., 33., 13., 27., 40., 12., 23., 33.])
391
        )
392
        np.testing.assert_array_equal(
393
            mapped_array_grouped.apply(cumsum_apply_nb, apply_per_group=True).values,
394
            np.array([10., 21., 33., 46., 60., 73., 12., 23., 33.])
395
        )
396
        assert mapped_array_grouped.apply(cumsum_apply_nb).wrapper == \
397
               mapped_array.apply(cumsum_apply_nb, group_by=group_by).wrapper
398
        assert mapped_array.apply(cumsum_apply_nb, group_by=False).wrapper.grouper.group_by is None
399

400
    def test_reduce(self):
401
        @njit
402
        def mean_reduce_nb(col, a):
403
            return np.mean(a)
404

405
        assert mapped_array['a'].reduce(mean_reduce_nb) == 11.
406
        pd.testing.assert_series_equal(
407
            mapped_array.reduce(mean_reduce_nb),
408
            pd.Series(np.array([11., 13.333333333333334, 11., np.nan]), index=wrapper.columns).rename('reduce')
409
        )
410
        pd.testing.assert_series_equal(
411
            mapped_array.reduce(mean_reduce_nb, fill_value=0.),
412
            pd.Series(np.array([11., 13.333333333333334, 11., 0.]), index=wrapper.columns).rename('reduce')
413
        )
414
        pd.testing.assert_series_equal(
415
            mapped_array.reduce(mean_reduce_nb, wrap_kwargs=dict(to_timedelta=True)),
416
            pd.Series(np.array([11., 13.333333333333334, 11., np.nan]), index=wrapper.columns).rename('reduce') * day_dt
417
        )
418
        pd.testing.assert_series_equal(
419
            mapped_array_grouped.reduce(mean_reduce_nb),
420
            pd.Series([12.166666666666666, 11.0], index=pd.Index(['g1', 'g2'], dtype='object')).rename('reduce')
421
        )
422
        assert mapped_array_grouped['g1'].reduce(mean_reduce_nb) == 12.166666666666666
423
        pd.testing.assert_series_equal(
424
            mapped_array_grouped[['g1']].reduce(mean_reduce_nb),
425
            pd.Series([12.166666666666666], index=pd.Index(['g1'], dtype='object')).rename('reduce')
426
        )
427
        pd.testing.assert_series_equal(
428
            mapped_array.reduce(mean_reduce_nb),
429
            mapped_array_grouped.reduce(mean_reduce_nb, group_by=False)
430
        )
431
        pd.testing.assert_series_equal(
432
            mapped_array.reduce(mean_reduce_nb, group_by=group_by),
433
            mapped_array_grouped.reduce(mean_reduce_nb)
434
        )
435

436
    def test_reduce_to_idx(self):
437
        @njit
438
        def argmin_reduce_nb(col, a):
439
            return np.argmin(a)
440

441
        assert mapped_array['a'].reduce(argmin_reduce_nb, returns_idx=True) == 'x'
442
        pd.testing.assert_series_equal(
443
            mapped_array.reduce(argmin_reduce_nb, returns_idx=True),
444
            pd.Series(np.array(['x', 'x', 'z', np.nan], dtype=object), index=wrapper.columns).rename('reduce')
445
        )
446
        pd.testing.assert_series_equal(
447
            mapped_array.reduce(argmin_reduce_nb, returns_idx=True, to_index=False),
448
            pd.Series(np.array([0, 0, 2, -1], dtype=int), index=wrapper.columns).rename('reduce')
449
        )
450
        pd.testing.assert_series_equal(
451
            mapped_array_grouped.reduce(argmin_reduce_nb, returns_idx=True, to_index=False),
452
            pd.Series(np.array([0, 2], dtype=int), index=pd.Index(['g1', 'g2'], dtype='object')).rename('reduce')
453
        )
454

455
    def test_reduce_to_array(self):
456
        @njit
457
        def min_max_reduce_nb(col, a):
458
            return np.array([np.min(a), np.max(a)])
459

460
        pd.testing.assert_series_equal(
461
            mapped_array['a'].reduce(min_max_reduce_nb, returns_array=True,
462
                                     wrap_kwargs=dict(name_or_index=['min', 'max'])),
463
            pd.Series([10., 12.], index=pd.Index(['min', 'max'], dtype='object'), name='a')
464
        )
465
        pd.testing.assert_frame_equal(
466
            mapped_array.reduce(min_max_reduce_nb, returns_array=True, wrap_kwargs=dict(name_or_index=['min', 'max'])),
467
            pd.DataFrame(
468
                np.array([
469
                    [10., 13., 10., np.nan],
470
                    [12., 14., 12., np.nan]
471
                ]),
472
                index=pd.Index(['min', 'max'], dtype='object'),
473
                columns=wrapper.columns
474
            )
475
        )
476
        pd.testing.assert_frame_equal(
477
            mapped_array.reduce(min_max_reduce_nb, returns_array=True, fill_value=0.),
478
            pd.DataFrame(
479
                np.array([
480
                    [10., 13., 10., 0.],
481
                    [12., 14., 12., 0.]
482
                ]),
483
                columns=wrapper.columns
484
            )
485
        )
486
        pd.testing.assert_frame_equal(
487
            mapped_array.reduce(min_max_reduce_nb, returns_array=True, wrap_kwargs=dict(to_timedelta=True)),
488
            pd.DataFrame(
489
                np.array([
490
                    [10., 13., 10., np.nan],
491
                    [12., 14., 12., np.nan]
492
                ]),
493
                columns=wrapper.columns
494
            ) * day_dt
495
        )
496
        pd.testing.assert_frame_equal(
497
            mapped_array_grouped.reduce(min_max_reduce_nb, returns_array=True),
498
            pd.DataFrame(
499
                np.array([
500
                    [10., 10.],
501
                    [14., 12.]
502
                ]),
503
                columns=pd.Index(['g1', 'g2'], dtype='object')
504
            )
505
        )
506
        pd.testing.assert_frame_equal(
507
            mapped_array.reduce(min_max_reduce_nb, returns_array=True),
508
            mapped_array_grouped.reduce(min_max_reduce_nb, returns_array=True, group_by=False)
509
        )
510
        pd.testing.assert_frame_equal(
511
            mapped_array.reduce(min_max_reduce_nb, returns_array=True, group_by=group_by),
512
            mapped_array_grouped.reduce(min_max_reduce_nb, returns_array=True)
513
        )
514
        pd.testing.assert_series_equal(
515
            mapped_array_grouped['g1'].reduce(min_max_reduce_nb, returns_array=True),
516
            pd.Series([10., 14.], name='g1')
517
        )
518
        pd.testing.assert_frame_equal(
519
            mapped_array_grouped[['g1']].reduce(min_max_reduce_nb, returns_array=True),
520
            pd.DataFrame([[10.], [14.]], columns=pd.Index(['g1'], dtype='object'))
521
        )
522

523
    def test_reduce_to_idx_array(self):
524
        @njit
525
        def idxmin_idxmax_reduce_nb(col, a):
526
            return np.array([np.argmin(a), np.argmax(a)])
527

528
        pd.testing.assert_series_equal(
529
            mapped_array['a'].reduce(
530
                idxmin_idxmax_reduce_nb,
531
                returns_array=True,
532
                returns_idx=True,
533
                wrap_kwargs=dict(name_or_index=['min', 'max'])
534
            ),
535
            pd.Series(
536
                np.array(['x', 'z'], dtype=object),
537
                index=pd.Index(['min', 'max'], dtype='object'),
538
                name='a'
539
            )
540
        )
541
        pd.testing.assert_frame_equal(
542
            mapped_array.reduce(
543
                idxmin_idxmax_reduce_nb,
544
                returns_array=True,
545
                returns_idx=True,
546
                wrap_kwargs=dict(name_or_index=['min', 'max'])
547
            ),
548
            pd.DataFrame(
549
                {
550
                    'a': ['x', 'z'],
551
                    'b': ['x', 'y'],
552
                    'c': ['z', 'x'],
553
                    'd': [np.nan, np.nan]
554
                },
555
                index=pd.Index(['min', 'max'], dtype='object')
556
            )
557
        )
558
        pd.testing.assert_frame_equal(
559
            mapped_array.reduce(
560
                idxmin_idxmax_reduce_nb,
561
                returns_array=True,
562
                returns_idx=True,
563
                to_index=False
564
            ),
565
            pd.DataFrame(
566
                np.array([
567
                    [0, 0, 2, -1],
568
                    [2, 1, 0, -1]
569
                ]),
570
                columns=wrapper.columns
571
            )
572
        )
573
        pd.testing.assert_frame_equal(
574
            mapped_array_grouped.reduce(
575
                idxmin_idxmax_reduce_nb,
576
                returns_array=True,
577
                returns_idx=True,
578
                to_index=False
579
            ),
580
            pd.DataFrame(
581
                np.array([
582
                    [0, 2],
583
                    [1, 0]
584
                ]),
585
                columns=pd.Index(['g1', 'g2'], dtype='object')
586
            )
587
        )
588

589
    def test_nth(self):
590
        assert mapped_array['a'].nth(0) == 10.
591
        pd.testing.assert_series_equal(
592
            mapped_array.nth(0),
593
            pd.Series(np.array([10., 13., 12., np.nan]), index=wrapper.columns).rename('nth')
594
        )
595
        assert mapped_array['a'].nth(-1) == 12.
596
        pd.testing.assert_series_equal(
597
            mapped_array.nth(-1),
598
            pd.Series(np.array([12., 13., 10., np.nan]), index=wrapper.columns).rename('nth')
599
        )
600
        with pytest.raises(Exception):
601
            _ = mapped_array.nth(10)
602
        pd.testing.assert_series_equal(
603
            mapped_array_grouped.nth(0),
604
            pd.Series(np.array([10., 12.]), index=pd.Index(['g1', 'g2'], dtype='object')).rename('nth')
605
        )
606

607
    def test_nth_index(self):
608
        assert mapped_array['a'].nth(0) == 10.
609
        pd.testing.assert_series_equal(
610
            mapped_array.nth_index(0),
611
            pd.Series(
612
                np.array(['x', 'x', 'x', np.nan], dtype='object'),
613
                index=wrapper.columns
614
            ).rename('nth_index')
615
        )
616
        assert mapped_array['a'].nth(-1) == 12.
617
        pd.testing.assert_series_equal(
618
            mapped_array.nth_index(-1),
619
            pd.Series(
620
                np.array(['z', 'z', 'z', np.nan], dtype='object'),
621
                index=wrapper.columns
622
            ).rename('nth_index')
623
        )
624
        with pytest.raises(Exception):
625
            _ = mapped_array.nth_index(10)
626
        pd.testing.assert_series_equal(
627
            mapped_array_grouped.nth_index(0),
628
            pd.Series(
629
                np.array(['x', 'x'], dtype='object'),
630
                index=pd.Index(['g1', 'g2'], dtype='object')
631
            ).rename('nth_index')
632
        )
633

634
    def test_min(self):
635
        assert mapped_array['a'].min() == mapped_array['a'].to_pd().min()
636
        pd.testing.assert_series_equal(
637
            mapped_array.min(),
638
            mapped_array.to_pd().min().rename('min')
639
        )
640
        pd.testing.assert_series_equal(
641
            mapped_array_grouped.min(),
642
            pd.Series([10., 10.], index=pd.Index(['g1', 'g2'], dtype='object')).rename('min')
643
        )
644

645
    def test_max(self):
646
        assert mapped_array['a'].max() == mapped_array['a'].to_pd().max()
647
        pd.testing.assert_series_equal(
648
            mapped_array.max(),
649
            mapped_array.to_pd().max().rename('max')
650
        )
651
        pd.testing.assert_series_equal(
652
            mapped_array_grouped.max(),
653
            pd.Series([14., 12.], index=pd.Index(['g1', 'g2'], dtype='object')).rename('max')
654
        )
655

656
    def test_mean(self):
657
        assert mapped_array['a'].mean() == mapped_array['a'].to_pd().mean()
658
        pd.testing.assert_series_equal(
659
            mapped_array.mean(),
660
            mapped_array.to_pd().mean().rename('mean')
661
        )
662
        pd.testing.assert_series_equal(
663
            mapped_array_grouped.mean(),
664
            pd.Series([12.166667, 11.], index=pd.Index(['g1', 'g2'], dtype='object')).rename('mean')
665
        )
666

667
    def test_median(self):
668
        assert mapped_array['a'].median() == mapped_array['a'].to_pd().median()
669
        pd.testing.assert_series_equal(
670
            mapped_array.median(),
671
            mapped_array.to_pd().median().rename('median')
672
        )
673
        pd.testing.assert_series_equal(
674
            mapped_array_grouped.median(),
675
            pd.Series([12.5, 11.], index=pd.Index(['g1', 'g2'], dtype='object')).rename('median')
676
        )
677

678
    def test_std(self):
679
        assert mapped_array['a'].std() == mapped_array['a'].to_pd().std()
680
        pd.testing.assert_series_equal(
681
            mapped_array.std(),
682
            mapped_array.to_pd().std().rename('std')
683
        )
684
        pd.testing.assert_series_equal(
685
            mapped_array.std(ddof=0),
686
            mapped_array.to_pd().std(ddof=0).rename('std')
687
        )
688
        pd.testing.assert_series_equal(
689
            mapped_array_grouped.std(),
690
            pd.Series([1.4719601443879746, 1.0], index=pd.Index(['g1', 'g2'], dtype='object')).rename('std')
691
        )
692

693
    def test_sum(self):
694
        assert mapped_array['a'].sum() == mapped_array['a'].to_pd().sum()
695
        pd.testing.assert_series_equal(
696
            mapped_array.sum(),
697
            mapped_array.to_pd().sum().rename('sum')
698
        )
699
        pd.testing.assert_series_equal(
700
            mapped_array_grouped.sum(),
701
            pd.Series([73.0, 33.0], index=pd.Index(['g1', 'g2'], dtype='object')).rename('sum')
702
        )
703

704
    def test_count(self):
705
        assert mapped_array['a'].count() == mapped_array['a'].to_pd().count()
706
        pd.testing.assert_series_equal(
707
            mapped_array.count(),
708
            mapped_array.to_pd().count().rename('count')
709
        )
710
        pd.testing.assert_series_equal(
711
            mapped_array_grouped.count(),
712
            pd.Series([6, 3], index=pd.Index(['g1', 'g2'], dtype='object')).rename('count')
713
        )
714

715
    def test_idxmin(self):
716
        assert mapped_array['a'].idxmin() == mapped_array['a'].to_pd().idxmin()
717
        pd.testing.assert_series_equal(
718
            mapped_array.idxmin(),
719
            mapped_array.to_pd().idxmin().rename('idxmin')
720
        )
721
        pd.testing.assert_series_equal(
722
            mapped_array_grouped.idxmin(),
723
            pd.Series(
724
                np.array(['x', 'z'], dtype=object),
725
                index=pd.Index(['g1', 'g2'], dtype='object')
726
            ).rename('idxmin')
727
        )
728

729
    def test_idxmax(self):
730
        assert mapped_array['a'].idxmax() == mapped_array['a'].to_pd().idxmax()
731
        pd.testing.assert_series_equal(
732
            mapped_array.idxmax(),
733
            mapped_array.to_pd().idxmax().rename('idxmax')
734
        )
735
        pd.testing.assert_series_equal(
736
            mapped_array_grouped.idxmax(),
737
            pd.Series(
738
                np.array(['y', 'x'], dtype=object),
739
                index=pd.Index(['g1', 'g2'], dtype='object')
740
            ).rename('idxmax')
741
        )
742

743
    def test_describe(self):
744
        pd.testing.assert_series_equal(
745
            mapped_array['a'].describe(),
746
            mapped_array['a'].to_pd().describe()
747
        )
748
        pd.testing.assert_frame_equal(
749
            mapped_array.describe(percentiles=None),
750
            mapped_array.to_pd().describe(percentiles=None)
751
        )
752
        pd.testing.assert_frame_equal(
753
            mapped_array.describe(percentiles=[]),
754
            mapped_array.to_pd().describe(percentiles=[])
755
        )
756
        pd.testing.assert_frame_equal(
757
            mapped_array.describe(percentiles=np.arange(0, 1, 0.1)),
758
            mapped_array.to_pd().describe(percentiles=np.arange(0, 1, 0.1))
759
        )
760
        pd.testing.assert_frame_equal(
761
            mapped_array_grouped.describe(),
762
            pd.DataFrame(
763
                np.array([
764
                    [6., 3.],
765
                    [12.16666667, 11.],
766
                    [1.47196014, 1.],
767
                    [10., 10.],
768
                    [11.25, 10.5],
769
                    [12.5, 11.],
770
                    [13., 11.5],
771
                    [14., 12.]
772
                ]),
773
                columns=pd.Index(['g1', 'g2'], dtype='object'),
774
                index=mapped_array.describe().index
775
            )
776
        )
777

778
    def test_value_counts(self):
779
        pd.testing.assert_series_equal(
780
            mapped_array['a'].value_counts(),
781
            pd.Series(
782
                np.array([1, 1, 1]),
783
                index=pd.Index([10.0, 11.0, 12.0], dtype='float64'),
784
                name='a'
785
            )
786
        )
787
        pd.testing.assert_series_equal(
788
            mapped_array['a'].value_counts(mapping=mapping),
789
            pd.Series(
790
                np.array([1, 1, 1]),
791
                index=pd.Index(['test_10.0', 'test_11.0', 'test_12.0'], dtype='object'),
792
                name='a'
793
            )
794
        )
795
        pd.testing.assert_frame_equal(
796
            mapped_array.value_counts(),
797
            pd.DataFrame(
798
                np.array([
799
                    [1, 0, 1, 0],
800
                    [1, 0, 1, 0],
801
                    [1, 0, 1, 0],
802
                    [0, 2, 0, 0],
803
                    [0, 1, 0, 0]
804
                ]),
805
                index=pd.Index([10.0, 11.0, 12.0, 13.0, 14.0], dtype='float64'),
806
                columns=wrapper.columns
807
            )
808
        )
809
        pd.testing.assert_frame_equal(
810
            mapped_array_grouped.value_counts(),
811
            pd.DataFrame(
812
                np.array([
813
                    [1, 1],
814
                    [1, 1],
815
                    [1, 1],
816
                    [2, 0],
817
                    [1, 0]
818
                ]),
819
                index=pd.Index([10.0, 11.0, 12.0, 13.0, 14.0], dtype='float64'),
820
                columns=pd.Index(['g1', 'g2'], dtype='object')
821
            )
822
        )
823
        mapped_array2 = mapped_array.replace(mapped_arr=[4, 4, 3, 2, np.nan, 4, 3, 2, 1])
824
        pd.testing.assert_frame_equal(
825
            mapped_array2.value_counts(sort_uniques=True),
826
            pd.DataFrame(
827
                np.array([
828
                    [0, 0, 1, 0],
829
                    [0, 1, 1, 0],
830
                    [1, 0, 1, 0],
831
                    [2, 1, 0, 0],
832
                    [0, 1, 0, 0]
833
                ]),
834
                index=pd.Index([1.0, 2.0, 3.0, 4.0, None], dtype='float64'),
835
                columns=wrapper.columns
836
            )
837
        )
838
        pd.testing.assert_frame_equal(
839
            mapped_array2.value_counts(sort=True),
840
            pd.DataFrame(
841
                np.array([
842
                    [2, 1, 0, 0],
843
                    [0, 1, 1, 0],
844
                    [1, 0, 1, 0],
845
                    [0, 0, 1, 0],
846
                    [0, 1, 0, 0]
847
                ]),
848
                index=pd.Index([4.0, 2.0, 3.0, 1.0, np.nan], dtype='float64'),
849
                columns=wrapper.columns
850
            )
851
        )
852
        pd.testing.assert_frame_equal(
853
            mapped_array2.value_counts(sort=True, ascending=True),
854
            pd.DataFrame(
855
                np.array([
856
                    [0, 0, 1, 0],
857
                    [0, 1, 0, 0],
858
                    [0, 1, 1, 0],
859
                    [1, 0, 1, 0],
860
                    [2, 1, 0, 0]
861
                ]),
862
                index=pd.Index([1.0, np.nan, 2.0, 3.0, 4.0], dtype='float64'),
863
                columns=wrapper.columns
864
            )
865
        )
866
        pd.testing.assert_frame_equal(
867
            mapped_array2.value_counts(sort=True, normalize=True),
868
            pd.DataFrame(
869
                np.array([
870
                    [0.2222222222222222, 0.1111111111111111, 0.0, 0.0],
871
                    [0.0, 0.1111111111111111, 0.1111111111111111, 0.0],
872
                    [0.1111111111111111, 0.0, 0.1111111111111111, 0.0],
873
                    [0.0, 0.0, 0.1111111111111111, 0.0],
874
                    [0.0, 0.1111111111111111, 0.0, 0.0]
875
                ]),
876
                index=pd.Index([4.0, 2.0, 3.0, 1.0, np.nan], dtype='float64'),
877
                columns=wrapper.columns
878
            )
879
        )
880
        pd.testing.assert_frame_equal(
881
            mapped_array2.value_counts(sort=True, normalize=True, dropna=True),
882
            pd.DataFrame(
883
                np.array([
884
                    [0.25, 0.125, 0.0, 0.0],
885
                    [0.0, 0.125, 0.125, 0.0],
886
                    [0.125, 0.0, 0.125, 0.0],
887
                    [0.0, 0.0, 0.125, 0.0]
888
                ]),
889
                index=pd.Index([4.0, 2.0, 3.0, 1.0], dtype='float64'),
890
                columns=wrapper.columns
891
            )
892
        )
893

894
    @pytest.mark.parametrize(
895
        "test_nosort",
896
        [False, True],
897
    )
898
    def test_indexing(self, test_nosort):
899
        if test_nosort:
900
            ma = mapped_array_nosort
901
            ma_grouped = mapped_array_nosort_grouped
902
        else:
903
            ma = mapped_array
904
            ma_grouped = mapped_array_grouped
905
        np.testing.assert_array_equal(
906
            ma['a'].id_arr,
907
            np.array([0, 1, 2])
908
        )
909
        np.testing.assert_array_equal(
910
            ma['a'].col_arr,
911
            np.array([0, 0, 0])
912
        )
913
        pd.testing.assert_index_equal(
914
            ma['a'].wrapper.columns,
915
            pd.Index(['a'], dtype='object')
916
        )
917
        np.testing.assert_array_equal(
918
            ma['b'].id_arr,
919
            np.array([3, 4, 5])
920
        )
921
        np.testing.assert_array_equal(
922
            ma['b'].col_arr,
923
            np.array([0, 0, 0])
924
        )
925
        pd.testing.assert_index_equal(
926
            ma['b'].wrapper.columns,
927
            pd.Index(['b'], dtype='object')
928
        )
929
        np.testing.assert_array_equal(
930
            ma[['a', 'a']].id_arr,
931
            np.array([0, 1, 2, 0, 1, 2])
932
        )
933
        np.testing.assert_array_equal(
934
            ma[['a', 'a']].col_arr,
935
            np.array([0, 0, 0, 1, 1, 1])
936
        )
937
        pd.testing.assert_index_equal(
938
            ma[['a', 'a']].wrapper.columns,
939
            pd.Index(['a', 'a'], dtype='object')
940
        )
941
        np.testing.assert_array_equal(
942
            ma[['a', 'b']].id_arr,
943
            np.array([0, 1, 2, 3, 4, 5])
944
        )
945
        np.testing.assert_array_equal(
946
            ma[['a', 'b']].col_arr,
947
            np.array([0, 0, 0, 1, 1, 1])
948
        )
949
        pd.testing.assert_index_equal(
950
            ma[['a', 'b']].wrapper.columns,
951
            pd.Index(['a', 'b'], dtype='object')
952
        )
953
        with pytest.raises(Exception):
954
            _ = ma.iloc[::2, :]  # changing time not supported
955
        pd.testing.assert_index_equal(
956
            ma_grouped['g1'].wrapper.columns,
957
            pd.Index(['a', 'b'], dtype='object')
958
        )
959
        assert ma_grouped['g1'].wrapper.ndim == 2
960
        assert ma_grouped['g1'].wrapper.grouped_ndim == 1
961
        pd.testing.assert_index_equal(
962
            ma_grouped['g1'].wrapper.grouper.group_by,
963
            pd.Index(['g1', 'g1'], dtype='object')
964
        )
965
        pd.testing.assert_index_equal(
966
            ma_grouped['g2'].wrapper.columns,
967
            pd.Index(['c', 'd'], dtype='object')
968
        )
969
        assert ma_grouped['g2'].wrapper.ndim == 2
970
        assert ma_grouped['g2'].wrapper.grouped_ndim == 1
971
        pd.testing.assert_index_equal(
972
            ma_grouped['g2'].wrapper.grouper.group_by,
973
            pd.Index(['g2', 'g2'], dtype='object')
974
        )
975
        pd.testing.assert_index_equal(
976
            ma_grouped[['g1']].wrapper.columns,
977
            pd.Index(['a', 'b'], dtype='object')
978
        )
979
        assert ma_grouped[['g1']].wrapper.ndim == 2
980
        assert ma_grouped[['g1']].wrapper.grouped_ndim == 2
981
        pd.testing.assert_index_equal(
982
            ma_grouped[['g1']].wrapper.grouper.group_by,
983
            pd.Index(['g1', 'g1'], dtype='object')
984
        )
985
        pd.testing.assert_index_equal(
986
            ma_grouped[['g1', 'g2']].wrapper.columns,
987
            pd.Index(['a', 'b', 'c', 'd'], dtype='object')
988
        )
989
        assert ma_grouped[['g1', 'g2']].wrapper.ndim == 2
990
        assert ma_grouped[['g1', 'g2']].wrapper.grouped_ndim == 2
991
        pd.testing.assert_index_equal(
992
            ma_grouped[['g1', 'g2']].wrapper.grouper.group_by,
993
            pd.Index(['g1', 'g1', 'g2', 'g2'], dtype='object')
994
        )
995

996
    def test_magic(self):
997
        a = vbt.MappedArray(
998
            wrapper,
999
            records_arr['some_field1'],
1000
            records_arr['col'],
1001
            id_arr=records_arr['id'],
1002
            idx_arr=records_arr['idx']
1003
        )
1004
        a_inv = vbt.MappedArray(
1005
            wrapper,
1006
            records_arr['some_field1'][::-1],
1007
            records_arr['col'][::-1],
1008
            id_arr=records_arr['id'][::-1],
1009
            idx_arr=records_arr['idx'][::-1]
1010
        )
1011
        b = records_arr['some_field2']
1012
        a_bool = vbt.MappedArray(
1013
            wrapper,
1014
            records_arr['some_field1'] > np.mean(records_arr['some_field1']),
1015
            records_arr['col'],
1016
            id_arr=records_arr['id'],
1017
            idx_arr=records_arr['idx']
1018
        )
1019
        b_bool = records_arr['some_field2'] > np.mean(records_arr['some_field2'])
1020
        assert a ** a == a ** 2
1021
        with pytest.raises(Exception):
1022
            _ = a * a_inv
1023

1024
        # binary ops
1025
        # comparison ops
1026
        np.testing.assert_array_equal((a == b).values, a.values == b)
1027
        np.testing.assert_array_equal((a != b).values, a.values != b)
1028
        np.testing.assert_array_equal((a < b).values, a.values < b)
1029
        np.testing.assert_array_equal((a > b).values, a.values > b)
1030
        np.testing.assert_array_equal((a <= b).values, a.values <= b)
1031
        np.testing.assert_array_equal((a >= b).values, a.values >= b)
1032
        # arithmetic ops
1033
        np.testing.assert_array_equal((a + b).values, a.values + b)
1034
        np.testing.assert_array_equal((a - b).values, a.values - b)
1035
        np.testing.assert_array_equal((a * b).values, a.values * b)
1036
        np.testing.assert_array_equal((a ** b).values, a.values ** b)
1037
        np.testing.assert_array_equal((a % b).values, a.values % b)
1038
        np.testing.assert_array_equal((a // b).values, a.values // b)
1039
        np.testing.assert_array_equal((a / b).values, a.values / b)
1040
        # __r*__ is only called if the left object does not have an __*__ method
1041
        np.testing.assert_array_equal((10 + a).values, 10 + a.values)
1042
        np.testing.assert_array_equal((10 - a).values, 10 - a.values)
1043
        np.testing.assert_array_equal((10 * a).values, 10 * a.values)
1044
        np.testing.assert_array_equal((10 ** a).values, 10 ** a.values)
1045
        np.testing.assert_array_equal((10 % a).values, 10 % a.values)
1046
        np.testing.assert_array_equal((10 // a).values, 10 // a.values)
1047
        np.testing.assert_array_equal((10 / a).values, 10 / a.values)
1048
        # mask ops
1049
        np.testing.assert_array_equal((a_bool & b_bool).values, a_bool.values & b_bool)
1050
        np.testing.assert_array_equal((a_bool | b_bool).values, a_bool.values | b_bool)
1051
        np.testing.assert_array_equal((a_bool ^ b_bool).values, a_bool.values ^ b_bool)
1052
        np.testing.assert_array_equal((True & a_bool).values, True & a_bool.values)
1053
        np.testing.assert_array_equal((True | a_bool).values, True | a_bool.values)
1054
        np.testing.assert_array_equal((True ^ a_bool).values, True ^ a_bool.values)
1055
        # unary ops
1056
        np.testing.assert_array_equal((-a).values, -a.values)
1057
        np.testing.assert_array_equal((+a).values, +a.values)
1058
        np.testing.assert_array_equal((abs(-a)).values, abs((-a.values)))
1059

1060
    def test_stats(self):
1061
        stats_index = pd.Index([
1062
            'Start', 'End', 'Period', 'Count', 'Mean', 'Std', 'Min', 'Median', 'Max', 'Min Index', 'Max Index'
1063
        ], dtype='object')
1064
        pd.testing.assert_series_equal(
1065
            mapped_array.stats(),
1066
            pd.Series([
1067
                'x', 'z', pd.Timedelta('3 days 00:00:00'),
1068
                2.25, 11.777777777777779, 0.859116756396542, 11.0, 11.666666666666666, 12.666666666666666
1069
            ],
1070
                index=stats_index[:-2],
1071
                name='agg_func_mean'
1072
            )
1073
        )
1074
        pd.testing.assert_series_equal(
1075
            mapped_array.stats(column='a'),
1076
            pd.Series([
1077
                'x', 'z', pd.Timedelta('3 days 00:00:00'),
1078
                3, 11.0, 1.0, 10.0, 11.0, 12.0, 'x', 'z'
1079
            ],
1080
                index=stats_index,
1081
                name='a'
1082
            )
1083
        )
1084
        pd.testing.assert_series_equal(
1085
            mapped_array.stats(column='g1', group_by=group_by),
1086
            pd.Series([
1087
                'x', 'z', pd.Timedelta('3 days 00:00:00'),
1088
                6, 12.166666666666666, 1.4719601443879746, 10.0, 12.5, 14.0, 'x', 'y'
1089
            ],
1090
                index=stats_index,
1091
                name='g1'
1092
            )
1093
        )
1094
        pd.testing.assert_series_equal(
1095
            mapped_array['c'].stats(),
1096
            mapped_array.stats(column='c')
1097
        )
1098
        pd.testing.assert_series_equal(
1099
            mapped_array['c'].stats(),
1100
            mapped_array.stats(column='c', group_by=False)
1101
        )
1102
        pd.testing.assert_series_equal(
1103
            mapped_array_grouped['g2'].stats(),
1104
            mapped_array_grouped.stats(column='g2')
1105
        )
1106
        pd.testing.assert_series_equal(
1107
            mapped_array_grouped['g2'].stats(),
1108
            mapped_array.stats(column='g2', group_by=group_by)
1109
        )
1110
        stats_df = mapped_array.stats(agg_func=None)
1111
        assert stats_df.shape == (4, 11)
1112
        pd.testing.assert_index_equal(stats_df.index, mapped_array.wrapper.columns)
1113
        pd.testing.assert_index_equal(stats_df.columns, stats_index)
1114

1115
    def test_stats_mapping(self):
1116
        stats_index = pd.Index([
1117
            'Start', 'End', 'Period', 'Count', 'Value Counts: test_10.0',
1118
            'Value Counts: test_11.0', 'Value Counts: test_12.0',
1119
            'Value Counts: test_13.0', 'Value Counts: test_14.0'
1120
        ], dtype='object')
1121
        pd.testing.assert_series_equal(
1122
            mp_mapped_array.stats(),
1123
            pd.Series([
1124
                'x',
1125
                'z',
1126
                pd.Timedelta('3 days 00:00:00'),
1127
                2.25, 0.5, 0.5, 0.5, 0.5, 0.25
1128
            ],
1129
                index=stats_index,
1130
                name='agg_func_mean'
1131
            )
1132
        )
1133
        pd.testing.assert_series_equal(
1134
            mp_mapped_array.stats(column='a'),
1135
            pd.Series([
1136
                'x',
1137
                'z',
1138
                pd.Timedelta('3 days 00:00:00'),
1139
                3, 1, 1, 1, 0, 0
1140
            ],
1141
                index=stats_index,
1142
                name='a'
1143
            )
1144
        )
1145
        pd.testing.assert_series_equal(
1146
            mp_mapped_array.stats(column='g1', group_by=group_by),
1147
            pd.Series([
1148
                'x',
1149
                'z',
1150
                pd.Timedelta('3 days 00:00:00'),
1151
                6, 1, 1, 1, 2, 1
1152
            ],
1153
                index=stats_index,
1154
                name='g1'
1155
            )
1156
        )
1157
        pd.testing.assert_series_equal(
1158
            mp_mapped_array.stats(),
1159
            mapped_array.stats(settings=dict(mapping=mapping))
1160
        )
1161
        pd.testing.assert_series_equal(
1162
            mp_mapped_array['c'].stats(settings=dict(incl_all_keys=True)),
1163
            mp_mapped_array.stats(column='c')
1164
        )
1165
        pd.testing.assert_series_equal(
1166
            mp_mapped_array['c'].stats(settings=dict(incl_all_keys=True)),
1167
            mp_mapped_array.stats(column='c', group_by=False)
1168
        )
1169
        pd.testing.assert_series_equal(
1170
            mp_mapped_array_grouped['g2'].stats(settings=dict(incl_all_keys=True)),
1171
            mp_mapped_array_grouped.stats(column='g2')
1172
        )
1173
        pd.testing.assert_series_equal(
1174
            mp_mapped_array_grouped['g2'].stats(settings=dict(incl_all_keys=True)),
1175
            mp_mapped_array.stats(column='g2', group_by=group_by)
1176
        )
1177
        stats_df = mp_mapped_array.stats(agg_func=None)
1178
        assert stats_df.shape == (4, 9)
1179
        pd.testing.assert_index_equal(stats_df.index, mp_mapped_array.wrapper.columns)
1180
        pd.testing.assert_index_equal(stats_df.columns, stats_index)
1181

1182

1183
# ############# base.py ############# #
1184

1185
class TestRecords:
1186
    def test_config(self, tmp_path):
1187
        assert vbt.Records.loads(records['a'].dumps()) == records['a']
1188
        assert vbt.Records.loads(records.dumps()) == records
1189
        records.save(tmp_path / 'records')
1190
        assert vbt.Records.load(tmp_path / 'records') == records
1191

1192
    def test_records(self):
1193
        pd.testing.assert_frame_equal(
1194
            records.records,
1195
            pd.DataFrame.from_records(records_arr)
1196
        )
1197

1198
    def test_recarray(self):
1199
        np.testing.assert_array_equal(records['a'].recarray.some_field1, records['a'].values['some_field1'])
1200
        np.testing.assert_array_equal(records.recarray.some_field1, records.values['some_field1'])
1201

1202
    def test_records_readable(self):
1203
        pd.testing.assert_frame_equal(
1204
            records.records_readable,
1205
            pd.DataFrame([
1206
                [0, 'a', 'x', 10.0, 21.0], [1, 'a', 'y', 11.0, 20.0], [2, 'a', 'z', 12.0, 19.0],
1207
                [3, 'b', 'x', 13.0, 18.0], [4, 'b', 'y', 14.0, 17.0], [5, 'b', 'z', 13.0, 18.0],
1208
                [6, 'c', 'x', 12.0, 19.0], [7, 'c', 'y', 11.0, 20.0], [8, 'c', 'z', 10.0, 21.0]
1209
            ], columns=pd.Index(['Id', 'Column', 'Timestamp', 'some_field1', 'some_field2'], dtype='object'))
1210
        )
1211

1212
    def test_is_sorted(self):
1213
        assert records.is_sorted()
1214
        assert records.is_sorted(incl_id=True)
1215
        assert not records_nosort.is_sorted()
1216
        assert not records_nosort.is_sorted(incl_id=True)
1217

1218
    def test_sort(self):
1219
        assert records.sort().is_sorted()
1220
        assert records.sort().is_sorted(incl_id=True)
1221
        assert records.sort(incl_id=True).is_sorted(incl_id=True)
1222
        assert records_nosort.sort().is_sorted()
1223
        assert records_nosort.sort().is_sorted(incl_id=True)
1224
        assert records_nosort.sort(incl_id=True).is_sorted(incl_id=True)
1225

1226
    def test_apply_mask(self):
1227
        mask_a = records['a'].values['some_field1'] >= records['a'].values['some_field1'].mean()
1228
        record_arrays_close(
1229
            records['a'].apply_mask(mask_a).values,
1230
            np.array([
1231
                (1, 0, 1, 11., 20.), (2, 0, 2, 12., 19.)
1232
            ], dtype=example_dt)
1233
        )
1234
        mask = records.values['some_field1'] >= records.values['some_field1'].mean()
1235
        filtered = records.apply_mask(mask)
1236
        record_arrays_close(
1237
            filtered.values,
1238
            np.array([
1239
                (2, 0, 2, 12., 19.), (3, 1, 0, 13., 18.), (4, 1, 1, 14., 17.),
1240
                (5, 1, 2, 13., 18.), (6, 2, 0, 12., 19.)
1241
            ], dtype=example_dt)
1242
        )
1243
        assert records_grouped.apply_mask(mask).wrapper == records_grouped.wrapper
1244

1245
    def test_map_field(self):
1246
        np.testing.assert_array_equal(
1247
            records['a'].map_field('some_field1').values,
1248
            np.array([10., 11., 12.])
1249
        )
1250
        np.testing.assert_array_equal(
1251
            records.map_field('some_field1').values,
1252
            np.array([10., 11., 12., 13., 14., 13., 12., 11., 10.])
1253
        )
1254
        assert records_grouped.map_field('some_field1').wrapper == \
1255
               records.map_field('some_field1', group_by=group_by).wrapper
1256
        assert records_grouped.map_field('some_field1', group_by=False).wrapper.grouper.group_by is None
1257

1258
    def test_map(self):
1259
        @njit
1260
        def map_func_nb(record):
1261
            return record['some_field1'] + record['some_field2']
1262

1263
        np.testing.assert_array_equal(
1264
            records['a'].map(map_func_nb).values,
1265
            np.array([31., 31., 31.])
1266
        )
1267
        np.testing.assert_array_equal(
1268
            records.map(map_func_nb).values,
1269
            np.array([31., 31., 31., 31., 31., 31., 31., 31., 31.])
1270
        )
1271
        assert records_grouped.map(map_func_nb).wrapper == \
1272
               records.map(map_func_nb, group_by=group_by).wrapper
1273
        assert records_grouped.map(map_func_nb, group_by=False).wrapper.grouper.group_by is None
1274

1275
    def test_map_array(self):
1276
        arr = records_arr['some_field1'] + records_arr['some_field2']
1277
        np.testing.assert_array_equal(
1278
            records['a'].map_array(arr[:3]).values,
1279
            np.array([31., 31., 31.])
1280
        )
1281
        np.testing.assert_array_equal(
1282
            records.map_array(arr).values,
1283
            np.array([31., 31., 31., 31., 31., 31., 31., 31., 31.])
1284
        )
1285
        assert records_grouped.map_array(arr).wrapper == \
1286
               records.map_array(arr, group_by=group_by).wrapper
1287
        assert records_grouped.map_array(arr, group_by=False).wrapper.grouper.group_by is None
1288

1289
    def test_apply(self):
1290
        @njit
1291
        def cumsum_apply_nb(records):
1292
            return np.cumsum(records['some_field1'])
1293

1294
        np.testing.assert_array_equal(
1295
            records['a'].apply(cumsum_apply_nb).values,
1296
            np.array([10., 21., 33.])
1297
        )
1298
        np.testing.assert_array_equal(
1299
            records.apply(cumsum_apply_nb).values,
1300
            np.array([10., 21., 33., 13., 27., 40., 12., 23., 33.])
1301
        )
1302
        np.testing.assert_array_equal(
1303
            records_grouped.apply(cumsum_apply_nb, apply_per_group=False).values,
1304
            np.array([10., 21., 33., 13., 27., 40., 12., 23., 33.])
1305
        )
1306
        np.testing.assert_array_equal(
1307
            records_grouped.apply(cumsum_apply_nb, apply_per_group=True).values,
1308
            np.array([10., 21., 33., 46., 60., 73., 12., 23., 33.])
1309
        )
1310
        assert records_grouped.apply(cumsum_apply_nb).wrapper == \
1311
               records.apply(cumsum_apply_nb, group_by=group_by).wrapper
1312
        assert records_grouped.apply(cumsum_apply_nb, group_by=False).wrapper.grouper.group_by is None
1313

1314
    def test_count(self):
1315
        assert records['a'].count() == 3
1316
        pd.testing.assert_series_equal(
1317
            records.count(),
1318
            pd.Series(
1319
                np.array([3, 3, 3, 0]),
1320
                index=wrapper.columns
1321
            ).rename('count')
1322
        )
1323
        assert records_grouped['g1'].count() == 6
1324
        pd.testing.assert_series_equal(
1325
            records_grouped.count(),
1326
            pd.Series(
1327
                np.array([6, 3]),
1328
                index=pd.Index(['g1', 'g2'], dtype='object')
1329
            ).rename('count')
1330
        )
1331

1332
    @pytest.mark.parametrize(
1333
        "test_nosort",
1334
        [False, True],
1335
    )
1336
    def test_indexing(self, test_nosort):
1337
        if test_nosort:
1338
            r = records_nosort
1339
            r_grouped = records_nosort_grouped
1340
        else:
1341
            r = records
1342
            r_grouped = records_grouped
1343
        record_arrays_close(
1344
            r['a'].values,
1345
            np.array([
1346
                (0, 0, 0, 10., 21.), (1, 0, 1, 11., 20.), (2, 0, 2, 12., 19.)
1347
            ], dtype=example_dt)
1348
        )
1349
        pd.testing.assert_index_equal(
1350
            r['a'].wrapper.columns,
1351
            pd.Index(['a'], dtype='object')
1352
        )
1353
        pd.testing.assert_index_equal(
1354
            r['b'].wrapper.columns,
1355
            pd.Index(['b'], dtype='object')
1356
        )
1357
        record_arrays_close(
1358
            r[['a', 'a']].values,
1359
            np.array([
1360
                (0, 0, 0, 10., 21.), (1, 0, 1, 11., 20.), (2, 0, 2, 12., 19.),
1361
                (0, 1, 0, 10., 21.), (1, 1, 1, 11., 20.), (2, 1, 2, 12., 19.)
1362
            ], dtype=example_dt)
1363
        )
1364
        pd.testing.assert_index_equal(
1365
            r[['a', 'a']].wrapper.columns,
1366
            pd.Index(['a', 'a'], dtype='object')
1367
        )
1368
        record_arrays_close(
1369
            r[['a', 'b']].values,
1370
            np.array([
1371
                (0, 0, 0, 10., 21.), (1, 0, 1, 11., 20.), (2, 0, 2, 12., 19.),
1372
                (3, 1, 0, 13., 18.), (4, 1, 1, 14., 17.), (5, 1, 2, 13., 18.)
1373
            ], dtype=example_dt)
1374
        )
1375
        pd.testing.assert_index_equal(
1376
            r[['a', 'b']].wrapper.columns,
1377
            pd.Index(['a', 'b'], dtype='object')
1378
        )
1379
        with pytest.raises(Exception):
1380
            _ = r.iloc[::2, :]  # changing time not supported
1381
        pd.testing.assert_index_equal(
1382
            r_grouped['g1'].wrapper.columns,
1383
            pd.Index(['a', 'b'], dtype='object')
1384
        )
1385
        assert r_grouped['g1'].wrapper.ndim == 2
1386
        assert r_grouped['g1'].wrapper.grouped_ndim == 1
1387
        pd.testing.assert_index_equal(
1388
            r_grouped['g1'].wrapper.grouper.group_by,
1389
            pd.Index(['g1', 'g1'], dtype='object')
1390
        )
1391
        pd.testing.assert_index_equal(
1392
            r_grouped['g2'].wrapper.columns,
1393
            pd.Index(['c', 'd'], dtype='object')
1394
        )
1395
        assert r_grouped['g2'].wrapper.ndim == 2
1396
        assert r_grouped['g2'].wrapper.grouped_ndim == 1
1397
        pd.testing.assert_index_equal(
1398
            r_grouped['g2'].wrapper.grouper.group_by,
1399
            pd.Index(['g2', 'g2'], dtype='object')
1400
        )
1401
        pd.testing.assert_index_equal(
1402
            r_grouped[['g1']].wrapper.columns,
1403
            pd.Index(['a', 'b'], dtype='object')
1404
        )
1405
        assert r_grouped[['g1']].wrapper.ndim == 2
1406
        assert r_grouped[['g1']].wrapper.grouped_ndim == 2
1407
        pd.testing.assert_index_equal(
1408
            r_grouped[['g1']].wrapper.grouper.group_by,
1409
            pd.Index(['g1', 'g1'], dtype='object')
1410
        )
1411
        pd.testing.assert_index_equal(
1412
            r_grouped[['g1', 'g2']].wrapper.columns,
1413
            pd.Index(['a', 'b', 'c', 'd'], dtype='object')
1414
        )
1415
        assert r_grouped[['g1', 'g2']].wrapper.ndim == 2
1416
        assert r_grouped[['g1', 'g2']].wrapper.grouped_ndim == 2
1417
        pd.testing.assert_index_equal(
1418
            r_grouped[['g1', 'g2']].wrapper.grouper.group_by,
1419
            pd.Index(['g1', 'g1', 'g2', 'g2'], dtype='object')
1420
        )
1421

1422
    def test_filtering(self):
1423
        filtered_records = vbt.Records(wrapper, records_arr[[0, -1]])
1424
        record_arrays_close(
1425
            filtered_records.values,
1426
            np.array([(0, 0, 0, 10., 21.), (8, 2, 2, 10., 21.)], dtype=example_dt)
1427
        )
1428
        # a
1429
        record_arrays_close(
1430
            filtered_records['a'].values,
1431
            np.array([(0, 0, 0, 10., 21.)], dtype=example_dt)
1432
        )
1433
        np.testing.assert_array_equal(
1434
            filtered_records['a'].map_field('some_field1').id_arr,
1435
            np.array([0])
1436
        )
1437
        assert filtered_records['a'].map_field('some_field1').min() == 10.
1438
        assert filtered_records['a'].count() == 1.
1439
        # b
1440
        record_arrays_close(
1441
            filtered_records['b'].values,
1442
            np.array([], dtype=example_dt)
1443
        )
1444
        np.testing.assert_array_equal(
1445
            filtered_records['b'].map_field('some_field1').id_arr,
1446
            np.array([])
1447
        )
1448
        assert np.isnan(filtered_records['b'].map_field('some_field1').min())
1449
        assert filtered_records['b'].count() == 0.
1450
        # c
1451
        record_arrays_close(
1452
            filtered_records['c'].values,
1453
            np.array([(8, 0, 2, 10., 21.)], dtype=example_dt)
1454
        )
1455
        np.testing.assert_array_equal(
1456
            filtered_records['c'].map_field('some_field1').id_arr,
1457
            np.array([8])
1458
        )
1459
        assert filtered_records['c'].map_field('some_field1').min() == 10.
1460
        assert filtered_records['c'].count() == 1.
1461
        # d
1462
        record_arrays_close(
1463
            filtered_records['d'].values,
1464
            np.array([], dtype=example_dt)
1465
        )
1466
        np.testing.assert_array_equal(
1467
            filtered_records['d'].map_field('some_field1').id_arr,
1468
            np.array([])
1469
        )
1470
        assert np.isnan(filtered_records['d'].map_field('some_field1').min())
1471
        assert filtered_records['d'].count() == 0.
1472

1473
    def test_stats(self):
1474
        stats_index = pd.Index([
1475
            'Start', 'End', 'Period', 'Count'
1476
        ], dtype='object')
1477
        pd.testing.assert_series_equal(
1478
            records.stats(),
1479
            pd.Series([
1480
                'x', 'z', pd.Timedelta('3 days 00:00:00'), 2.25
1481
            ],
1482
                index=stats_index,
1483
                name='agg_func_mean'
1484
            )
1485
        )
1486
        pd.testing.assert_series_equal(
1487
            records.stats(column='a'),
1488
            pd.Series([
1489
                'x', 'z', pd.Timedelta('3 days 00:00:00'), 3
1490
            ],
1491
                index=stats_index,
1492
                name='a'
1493
            )
1494
        )
1495
        pd.testing.assert_series_equal(
1496
            records.stats(column='g1', group_by=group_by),
1497
            pd.Series([
1498
                'x', 'z', pd.Timedelta('3 days 00:00:00'), 6
1499
            ],
1500
                index=stats_index,
1501
                name='g1'
1502
            )
1503
        )
1504
        pd.testing.assert_series_equal(
1505
            records['c'].stats(),
1506
            records.stats(column='c')
1507
        )
1508
        pd.testing.assert_series_equal(
1509
            records['c'].stats(),
1510
            records.stats(column='c', group_by=False)
1511
        )
1512
        pd.testing.assert_series_equal(
1513
            records_grouped['g2'].stats(),
1514
            records_grouped.stats(column='g2')
1515
        )
1516
        pd.testing.assert_series_equal(
1517
            records_grouped['g2'].stats(),
1518
            records.stats(column='g2', group_by=group_by)
1519
        )
1520
        stats_df = records.stats(agg_func=None)
1521
        assert stats_df.shape == (4, 4)
1522
        pd.testing.assert_index_equal(stats_df.index, records.wrapper.columns)
1523
        pd.testing.assert_index_equal(stats_df.columns, stats_index)
1524

1525

1526
# ############# ranges.py ############# #
1527

1528
ts = pd.DataFrame({
1529
    'a': [1, -1, 3, -1, 5, -1],
1530
    'b': [-1, -1, -1, 4, 5, 6],
1531
    'c': [1, 2, 3, -1, -1, -1],
1532
    'd': [-1, -1, -1, -1, -1, -1]
1533
}, index=[
1534
    datetime(2020, 1, 1),
1535
    datetime(2020, 1, 2),
1536
    datetime(2020, 1, 3),
1537
    datetime(2020, 1, 4),
1538
    datetime(2020, 1, 5),
1539
    datetime(2020, 1, 6)
1540
])
1541

1542
ranges = vbt.Ranges.from_ts(ts, wrapper_kwargs=dict(freq='1 days'))
1543
ranges_grouped = vbt.Ranges.from_ts(ts, wrapper_kwargs=dict(freq='1 days', group_by=group_by))
1544

1545

1546
class TestRanges:
1547
    def test_mapped_fields(self):
1548
        for name in range_dt.names:
1549
            np.testing.assert_array_equal(
1550
                getattr(ranges, name).values,
1551
                ranges.values[name]
1552
            )
1553

1554
    def test_from_ts(self):
1555
        record_arrays_close(
1556
            ranges.values,
1557
            np.array([
1558
                (0, 0, 0, 1, 1), (1, 0, 2, 3, 1), (2, 0, 4, 5, 1), (3, 1, 3, 5, 0), (4, 2, 0, 3, 1)
1559
            ], dtype=range_dt)
1560
        )
1561
        assert ranges.wrapper.freq == day_dt
1562
        pd.testing.assert_index_equal(
1563
            ranges_grouped.wrapper.grouper.group_by,
1564
            group_by
1565
        )
1566

1567
    def test_records_readable(self):
1568
        records_readable = ranges.records_readable
1569

1570
        np.testing.assert_array_equal(
1571
            records_readable['Range Id'].values,
1572
            np.array([
1573
                0, 1, 2, 3, 4
1574
            ])
1575
        )
1576
        np.testing.assert_array_equal(
1577
            records_readable['Column'].values,
1578
            np.array([
1579
                'a', 'a', 'a', 'b', 'c'
1580
            ])
1581
        )
1582
        np.testing.assert_array_equal(
1583
            records_readable['Start Timestamp'].values,
1584
            np.array([
1585
                '2020-01-01T00:00:00.000000000', '2020-01-03T00:00:00.000000000',
1586
                '2020-01-05T00:00:00.000000000', '2020-01-04T00:00:00.000000000',
1587
                '2020-01-01T00:00:00.000000000'
1588
            ], dtype='datetime64[ns]')
1589
        )
1590
        np.testing.assert_array_equal(
1591
            records_readable['End Timestamp'].values,
1592
            np.array([
1593
                '2020-01-02T00:00:00.000000000', '2020-01-04T00:00:00.000000000',
1594
                '2020-01-06T00:00:00.000000000', '2020-01-06T00:00:00.000000000',
1595
                '2020-01-04T00:00:00.000000000'
1596
            ], dtype='datetime64[ns]')
1597
        )
1598
        np.testing.assert_array_equal(
1599
            records_readable['Status'].values,
1600
            np.array([
1601
                'Closed', 'Closed', 'Closed', 'Open', 'Closed'
1602
            ])
1603
        )
1604

1605
    def test_to_mask(self):
1606
        pd.testing.assert_series_equal(
1607
            ranges['a'].to_mask(),
1608
            ts['a'] != -1
1609
        )
1610
        pd.testing.assert_frame_equal(
1611
            ranges.to_mask(),
1612
            ts != -1
1613
        )
1614
        pd.testing.assert_frame_equal(
1615
            ranges_grouped.to_mask(),
1616
            pd.DataFrame(
1617
                [
1618
                    [True, True],
1619
                    [False, True],
1620
                    [True, True],
1621
                    [True, False],
1622
                    [True, False],
1623
                    [True, False]
1624
                ],
1625
                index=ts.index,
1626
                columns=pd.Index(['g1', 'g2'], dtype='object')
1627
            )
1628
        )
1629

1630
    def test_duration(self):
1631
        np.testing.assert_array_equal(
1632
            ranges['a'].duration.values,
1633
            np.array([1, 1, 1])
1634
        )
1635
        np.testing.assert_array_equal(
1636
            ranges.duration.values,
1637
            np.array([1, 1, 1, 3, 3])
1638
        )
1639

1640
    def test_avg_duration(self):
1641
        assert ranges['a'].avg_duration() == pd.Timedelta('1 days 00:00:00')
1642
        pd.testing.assert_series_equal(
1643
            ranges.avg_duration(),
1644
            pd.Series(
1645
                np.array([86400000000000, 259200000000000, 259200000000000, 'NaT'], dtype='timedelta64[ns]'),
1646
                index=wrapper.columns
1647
            ).rename('avg_duration')
1648
        )
1649
        pd.testing.assert_series_equal(
1650
            ranges_grouped.avg_duration(),
1651
            pd.Series(
1652
                np.array([129600000000000, 259200000000000], dtype='timedelta64[ns]'),
1653
                index=pd.Index(['g1', 'g2'], dtype='object')
1654
            ).rename('avg_duration')
1655
        )
1656

1657
    def test_max_duration(self):
1658
        assert ranges['a'].max_duration() == pd.Timedelta('1 days 00:00:00')
1659
        pd.testing.assert_series_equal(
1660
            ranges.max_duration(),
1661
            pd.Series(
1662
                np.array([86400000000000, 259200000000000, 259200000000000, 'NaT'], dtype='timedelta64[ns]'),
1663
                index=wrapper.columns
1664
            ).rename('max_duration')
1665
        )
1666
        pd.testing.assert_series_equal(
1667
            ranges_grouped.max_duration(),
1668
            pd.Series(
1669
                np.array([259200000000000, 259200000000000], dtype='timedelta64[ns]'),
1670
                index=pd.Index(['g1', 'g2'], dtype='object')
1671
            ).rename('max_duration')
1672
        )
1673

1674
    def test_coverage(self):
1675
        assert ranges['a'].coverage() == 0.5
1676
        pd.testing.assert_series_equal(
1677
            ranges.coverage(),
1678
            pd.Series(
1679
                np.array([0.5, 0.5, 0.5, np.nan]),
1680
                index=ts2.columns
1681
            ).rename('coverage')
1682
        )
1683
        pd.testing.assert_series_equal(
1684
            ranges.coverage(),
1685
            ranges.replace(records_arr=np.repeat(ranges.values, 2)).coverage()
1686
        )
1687
        pd.testing.assert_series_equal(
1688
            ranges.replace(records_arr=np.repeat(ranges.values, 2)).coverage(overlapping=True),
1689
            pd.Series(
1690
                np.array([1.0, 1.0, 1.0, np.nan]),
1691
                index=ts2.columns
1692
            ).rename('coverage')
1693
        )
1694
        pd.testing.assert_series_equal(
1695
            ranges.coverage(normalize=False),
1696
            pd.Series(
1697
                np.array([3.0, 3.0, 3.0, np.nan]),
1698
                index=ts2.columns
1699
            ).rename('coverage')
1700
        )
1701
        pd.testing.assert_series_equal(
1702
            ranges.replace(records_arr=np.repeat(ranges.values, 2)).coverage(overlapping=True, normalize=False),
1703
            pd.Series(
1704
                np.array([3.0, 3.0, 3.0, np.nan]),
1705
                index=ts2.columns
1706
            ).rename('coverage')
1707
        )
1708
        pd.testing.assert_series_equal(
1709
            ranges_grouped.coverage(),
1710
            pd.Series(
1711
                np.array([0.4166666666666667, 0.25]),
1712
                index=pd.Index(['g1', 'g2'], dtype='object')
1713
            ).rename('coverage')
1714
        )
1715
        pd.testing.assert_series_equal(
1716
            ranges_grouped.coverage(),
1717
            ranges_grouped.replace(records_arr=np.repeat(ranges_grouped.values, 2)).coverage()
1718
        )
1719

1720
    def test_stats(self):
1721
        stats_index = pd.Index([
1722
            'Start', 'End', 'Period', 'Coverage', 'Overlap Coverage',
1723
            'Total Records', 'Duration: Min', 'Duration: Median', 'Duration: Max',
1724
            'Duration: Mean', 'Duration: Std'
1725
        ], dtype='object')
1726
        pd.testing.assert_series_equal(
1727
            ranges.stats(),
1728
            pd.Series([
1729
                pd.Timestamp('2020-01-01 00:00:00'), pd.Timestamp('2020-01-06 00:00:00'),
1730
                pd.Timedelta('6 days 00:00:00'), pd.Timedelta('3 days 00:00:00'),
1731
                pd.Timedelta('0 days 00:00:00'), 1.25, pd.Timedelta('2 days 08:00:00'),
1732
                pd.Timedelta('2 days 08:00:00'), pd.Timedelta('2 days 08:00:00'),
1733
                pd.Timedelta('2 days 08:00:00'), pd.Timedelta('0 days 00:00:00')
1734
            ],
1735
                index=stats_index,
1736
                name='agg_func_mean'
1737
            )
1738
        )
1739
        pd.testing.assert_series_equal(
1740
            ranges.stats(column='a'),
1741
            pd.Series([
1742
                pd.Timestamp('2020-01-01 00:00:00'), pd.Timestamp('2020-01-06 00:00:00'),
1743
                pd.Timedelta('6 days 00:00:00'), pd.Timedelta('3 days 00:00:00'),
1744
                pd.Timedelta('0 days 00:00:00'), 3, pd.Timedelta('1 days 00:00:00'),
1745
                pd.Timedelta('1 days 00:00:00'), pd.Timedelta('1 days 00:00:00'),
1746
                pd.Timedelta('1 days 00:00:00'), pd.Timedelta('0 days 00:00:00')
1747
            ],
1748
                index=stats_index,
1749
                name='a'
1750
            )
1751
        )
1752
        pd.testing.assert_series_equal(
1753
            ranges.stats(column='g1', group_by=group_by),
1754
            pd.Series([
1755
                pd.Timestamp('2020-01-01 00:00:00'), pd.Timestamp('2020-01-06 00:00:00'),
1756
                pd.Timedelta('6 days 00:00:00'), pd.Timedelta('5 days 00:00:00'),
1757
                pd.Timedelta('1 days 00:00:00'), 4, pd.Timedelta('1 days 00:00:00'),
1758
                pd.Timedelta('1 days 00:00:00'), pd.Timedelta('3 days 00:00:00'),
1759
                pd.Timedelta('1 days 12:00:00'), pd.Timedelta('1 days 00:00:00')
1760
            ],
1761
                index=stats_index,
1762
                name='g1'
1763
            )
1764
        )
1765
        pd.testing.assert_series_equal(
1766
            ranges['c'].stats(),
1767
            ranges.stats(column='c')
1768
        )
1769
        pd.testing.assert_series_equal(
1770
            ranges['c'].stats(),
1771
            ranges.stats(column='c', group_by=False)
1772
        )
1773
        pd.testing.assert_series_equal(
1774
            ranges_grouped['g2'].stats(),
1775
            ranges_grouped.stats(column='g2')
1776
        )
1777
        pd.testing.assert_series_equal(
1778
            ranges_grouped['g2'].stats(),
1779
            ranges.stats(column='g2', group_by=group_by)
1780
        )
1781
        stats_df = ranges.stats(agg_func=None)
1782
        assert stats_df.shape == (4, 11)
1783
        pd.testing.assert_index_equal(stats_df.index, ranges.wrapper.columns)
1784
        pd.testing.assert_index_equal(stats_df.columns, stats_index)
1785

1786

1787
# ############# drawdowns.py ############# #
1788

1789

1790
ts2 = pd.DataFrame({
1791
    'a': [2, 1, 3, 1, 4, 1],
1792
    'b': [1, 2, 1, 3, 1, 4],
1793
    'c': [1, 2, 3, 2, 1, 2],
1794
    'd': [1, 2, 3, 4, 5, 6]
1795
}, index=[
1796
    datetime(2020, 1, 1),
1797
    datetime(2020, 1, 2),
1798
    datetime(2020, 1, 3),
1799
    datetime(2020, 1, 4),
1800
    datetime(2020, 1, 5),
1801
    datetime(2020, 1, 6)
1802
])
1803

1804
drawdowns = vbt.Drawdowns.from_ts(ts2, wrapper_kwargs=dict(freq='1 days'))
1805
drawdowns_grouped = vbt.Drawdowns.from_ts(ts2, wrapper_kwargs=dict(freq='1 days', group_by=group_by))
1806

1807

1808
class TestDrawdowns:
1809
    def test_mapped_fields(self):
1810
        for name in drawdown_dt.names:
1811
            np.testing.assert_array_equal(
1812
                getattr(drawdowns, name).values,
1813
                drawdowns.values[name]
1814
            )
1815

1816
    def test_ts(self):
1817
        pd.testing.assert_frame_equal(
1818
            drawdowns.ts,
1819
            ts2
1820
        )
1821
        pd.testing.assert_series_equal(
1822
            drawdowns['a'].ts,
1823
            ts2['a']
1824
        )
1825
        pd.testing.assert_frame_equal(
1826
            drawdowns_grouped['g1'].ts,
1827
            ts2[['a', 'b']]
1828
        )
1829
        assert drawdowns.replace(ts=None)['a'].ts is None
1830

1831
    def test_from_ts(self):
1832
        record_arrays_close(
1833
            drawdowns.values,
1834
            np.array([
1835
                (0, 0, 0, 1, 1, 2, 2.0, 1.0, 3.0, 1), (1, 0, 2, 3, 3, 4, 3.0, 1.0, 4.0, 1),
1836
                (2, 0, 4, 5, 5, 5, 4.0, 1.0, 1.0, 0), (3, 1, 1, 2, 2, 3, 2.0, 1.0, 3.0, 1),
1837
                (4, 1, 3, 4, 4, 5, 3.0, 1.0, 4.0, 1), (5, 2, 2, 3, 4, 5, 3.0, 1.0, 2.0, 0)
1838
            ], dtype=drawdown_dt)
1839
        )
1840
        assert drawdowns.wrapper.freq == day_dt
1841
        pd.testing.assert_index_equal(
1842
            drawdowns_grouped.wrapper.grouper.group_by,
1843
            group_by
1844
        )
1845

1846
    def test_records_readable(self):
1847
        records_readable = drawdowns.records_readable
1848

1849
        np.testing.assert_array_equal(
1850
            records_readable['Drawdown Id'].values,
1851
            np.array([
1852
                0, 1, 2, 3, 4, 5
1853
            ])
1854
        )
1855
        np.testing.assert_array_equal(
1856
            records_readable['Column'].values,
1857
            np.array([
1858
                'a', 'a', 'a', 'b', 'b', 'c'
1859
            ])
1860
        )
1861
        np.testing.assert_array_equal(
1862
            records_readable['Peak Timestamp'].values,
1863
            np.array([
1864
                '2020-01-01T00:00:00.000000000', '2020-01-03T00:00:00.000000000',
1865
                '2020-01-05T00:00:00.000000000', '2020-01-02T00:00:00.000000000',
1866
                '2020-01-04T00:00:00.000000000', '2020-01-03T00:00:00.000000000'
1867
            ], dtype='datetime64[ns]')
1868
        )
1869
        np.testing.assert_array_equal(
1870
            records_readable['Start Timestamp'].values,
1871
            np.array([
1872
                '2020-01-02T00:00:00.000000000', '2020-01-04T00:00:00.000000000',
1873
                '2020-01-06T00:00:00.000000000', '2020-01-03T00:00:00.000000000',
1874
                '2020-01-05T00:00:00.000000000', '2020-01-04T00:00:00.000000000'
1875
            ], dtype='datetime64[ns]')
1876
        )
1877
        np.testing.assert_array_equal(
1878
            records_readable['Valley Timestamp'].values,
1879
            np.array([
1880
                '2020-01-02T00:00:00.000000000', '2020-01-04T00:00:00.000000000',
1881
                '2020-01-06T00:00:00.000000000', '2020-01-03T00:00:00.000000000',
1882
                '2020-01-05T00:00:00.000000000', '2020-01-05T00:00:00.000000000'
1883
            ], dtype='datetime64[ns]')
1884
        )
1885
        np.testing.assert_array_equal(
1886
            records_readable['End Timestamp'].values,
1887
            np.array([
1888
                '2020-01-03T00:00:00.000000000', '2020-01-05T00:00:00.000000000',
1889
                '2020-01-06T00:00:00.000000000', '2020-01-04T00:00:00.000000000',
1890
                '2020-01-06T00:00:00.000000000', '2020-01-06T00:00:00.000000000'
1891
            ], dtype='datetime64[ns]')
1892
        )
1893
        np.testing.assert_array_equal(
1894
            records_readable['Peak Value'].values,
1895
            np.array([
1896
                2., 3., 4., 2., 3., 3.
1897
            ])
1898
        )
1899
        np.testing.assert_array_equal(
1900
            records_readable['Valley Value'].values,
1901
            np.array([
1902
                1., 1., 1., 1., 1., 1.
1903
            ])
1904
        )
1905
        np.testing.assert_array_equal(
1906
            records_readable['End Value'].values,
1907
            np.array([
1908
                3., 4., 1., 3., 4., 2.
1909
            ])
1910
        )
1911
        np.testing.assert_array_equal(
1912
            records_readable['Status'].values,
1913
            np.array([
1914
                'Recovered', 'Recovered', 'Active', 'Recovered', 'Recovered', 'Active'
1915
            ])
1916
        )
1917

1918
    def test_drawdown(self):
1919
        np.testing.assert_array_almost_equal(
1920
            drawdowns['a'].drawdown.values,
1921
            np.array([-0.5, -0.66666667, -0.75])
1922
        )
1923
        np.testing.assert_array_almost_equal(
1924
            drawdowns.drawdown.values,
1925
            np.array([-0.5, -0.66666667, -0.75, -0.5, -0.66666667, -0.66666667])
1926
        )
1927
        pd.testing.assert_frame_equal(
1928
            drawdowns.drawdown.to_pd(),
1929
            pd.DataFrame(
1930
                np.array([
1931
                    [np.nan, np.nan, np.nan, np.nan],
1932
                    [np.nan, np.nan, np.nan, np.nan],
1933
                    [-0.5, np.nan, np.nan, np.nan],
1934
                    [np.nan, -0.5, np.nan, np.nan],
1935
                    [-0.66666669, np.nan, np.nan, np.nan],
1936
                    [-0.75, -0.66666669, -0.66666669, np.nan]
1937
                ]),
1938
                index=ts2.index,
1939
                columns=ts2.columns
1940
            )
1941
        )
1942

1943
    def test_avg_drawdown(self):
1944
        assert drawdowns['a'].avg_drawdown() == -0.6388888888888888
1945
        pd.testing.assert_series_equal(
1946
            drawdowns.avg_drawdown(),
1947
            pd.Series(
1948
                np.array([-0.63888889, -0.58333333, -0.66666667, np.nan]),
1949
                index=wrapper.columns
1950
            ).rename('avg_drawdown')
1951
        )
1952
        pd.testing.assert_series_equal(
1953
            drawdowns_grouped.avg_drawdown(),
1954
            pd.Series(
1955
                np.array([-0.6166666666666666, -0.6666666666666666]),
1956
                index=pd.Index(['g1', 'g2'], dtype='object')
1957
            ).rename('avg_drawdown')
1958
        )
1959

1960
    def test_max_drawdown(self):
1961
        assert drawdowns['a'].max_drawdown() == -0.75
1962
        pd.testing.assert_series_equal(
1963
            drawdowns.max_drawdown(),
1964
            pd.Series(
1965
                np.array([-0.75, -0.66666667, -0.66666667, np.nan]),
1966
                index=wrapper.columns
1967
            ).rename('max_drawdown')
1968
        )
1969
        pd.testing.assert_series_equal(
1970
            drawdowns_grouped.max_drawdown(),
1971
            pd.Series(
1972
                np.array([-0.75, -0.6666666666666666]),
1973
                index=pd.Index(['g1', 'g2'], dtype='object')
1974
            ).rename('max_drawdown')
1975
        )
1976

1977
    def test_recovery_return(self):
1978
        np.testing.assert_array_almost_equal(
1979
            drawdowns['a'].recovery_return.values,
1980
            np.array([2., 3., 0.])
1981
        )
1982
        np.testing.assert_array_almost_equal(
1983
            drawdowns.recovery_return.values,
1984
            np.array([2., 3., 0., 2., 3., 1.])
1985
        )
1986
        pd.testing.assert_frame_equal(
1987
            drawdowns.recovery_return.to_pd(),
1988
            pd.DataFrame(
1989
                np.array([
1990
                    [np.nan, np.nan, np.nan, np.nan],
1991
                    [np.nan, np.nan, np.nan, np.nan],
1992
                    [2.0, np.nan, np.nan, np.nan],
1993
                    [np.nan, 2.0, np.nan, np.nan],
1994
                    [3.0, np.nan, np.nan, np.nan],
1995
                    [0.0, 3.0, 1.0, np.nan]
1996
                ]),
1997
                index=ts2.index,
1998
                columns=ts2.columns
1999
            )
2000
        )
2001

2002
    def test_avg_recovery_return(self):
2003
        assert drawdowns['a'].avg_recovery_return() == 1.6666666666666667
2004
        pd.testing.assert_series_equal(
2005
            drawdowns.avg_recovery_return(),
2006
            pd.Series(
2007
                np.array([1.6666666666666667, 2.5, 1.0, np.nan]),
2008
                index=wrapper.columns
2009
            ).rename('avg_recovery_return')
2010
        )
2011
        pd.testing.assert_series_equal(
2012
            drawdowns_grouped.avg_recovery_return(),
2013
            pd.Series(
2014
                np.array([2.0, 1.0]),
2015
                index=pd.Index(['g1', 'g2'], dtype='object')
2016
            ).rename('avg_recovery_return')
2017
        )
2018

2019
    def test_max_recovery_return(self):
2020
        assert drawdowns['a'].max_recovery_return() == 3.0
2021
        pd.testing.assert_series_equal(
2022
            drawdowns.max_recovery_return(),
2023
            pd.Series(
2024
                np.array([3.0, 3.0, 1.0, np.nan]),
2025
                index=wrapper.columns
2026
            ).rename('max_recovery_return')
2027
        )
2028
        pd.testing.assert_series_equal(
2029
            drawdowns_grouped.max_recovery_return(),
2030
            pd.Series(
2031
                np.array([3.0, 1.0]),
2032
                index=pd.Index(['g1', 'g2'], dtype='object')
2033
            ).rename('max_recovery_return')
2034
        )
2035

2036
    def test_duration(self):
2037
        np.testing.assert_array_almost_equal(
2038
            drawdowns['a'].duration.values,
2039
            np.array([1, 1, 1])
2040
        )
2041
        np.testing.assert_array_almost_equal(
2042
            drawdowns.duration.values,
2043
            np.array([1, 1, 1, 1, 1, 3])
2044
        )
2045

2046
    def test_avg_duration(self):
2047
        assert drawdowns['a'].avg_duration() == pd.Timedelta('1 days 00:00:00')
2048
        pd.testing.assert_series_equal(
2049
            drawdowns.avg_duration(),
2050
            pd.Series(
2051
                np.array([86400000000000, 86400000000000, 259200000000000, 'NaT'], dtype='timedelta64[ns]'),
2052
                index=wrapper.columns
2053
            ).rename('avg_duration')
2054
        )
2055
        pd.testing.assert_series_equal(
2056
            drawdowns_grouped.avg_duration(),
2057
            pd.Series(
2058
                np.array([86400000000000, 259200000000000], dtype='timedelta64[ns]'),
2059
                index=pd.Index(['g1', 'g2'], dtype='object')
2060
            ).rename('avg_duration')
2061
        )
2062

2063
    def test_max_duration(self):
2064
        assert drawdowns['a'].max_duration() == pd.Timedelta('1 days 00:00:00')
2065
        pd.testing.assert_series_equal(
2066
            drawdowns.max_duration(),
2067
            pd.Series(
2068
                np.array([86400000000000, 86400000000000, 259200000000000, 'NaT'], dtype='timedelta64[ns]'),
2069
                index=wrapper.columns
2070
            ).rename('max_duration')
2071
        )
2072
        pd.testing.assert_series_equal(
2073
            drawdowns_grouped.max_duration(),
2074
            pd.Series(
2075
                np.array([86400000000000, 259200000000000], dtype='timedelta64[ns]'),
2076
                index=pd.Index(['g1', 'g2'], dtype='object')
2077
            ).rename('max_duration')
2078
        )
2079

2080
    def test_coverage(self):
2081
        assert drawdowns['a'].coverage() == 0.5
2082
        pd.testing.assert_series_equal(
2083
            drawdowns.coverage(),
2084
            pd.Series(
2085
                np.array([0.5, 0.3333333333333333, 0.5, np.nan]),
2086
                index=ts2.columns
2087
            ).rename('coverage')
2088
        )
2089
        pd.testing.assert_series_equal(
2090
            drawdowns_grouped.coverage(),
2091
            pd.Series(
2092
                np.array([0.4166666666666667, 0.25]),
2093
                index=pd.Index(['g1', 'g2'], dtype='object')
2094
            ).rename('coverage')
2095
        )
2096

2097
    def test_decline_duration(self):
2098
        np.testing.assert_array_almost_equal(
2099
            drawdowns['a'].decline_duration.values,
2100
            np.array([1., 1., 1.])
2101
        )
2102
        np.testing.assert_array_almost_equal(
2103
            drawdowns.decline_duration.values,
2104
            np.array([1., 1., 1., 1., 1., 2.])
2105
        )
2106

2107
    def test_recovery_duration(self):
2108
        np.testing.assert_array_almost_equal(
2109
            drawdowns['a'].recovery_duration.values,
2110
            np.array([1, 1, 0])
2111
        )
2112
        np.testing.assert_array_almost_equal(
2113
            drawdowns.recovery_duration.values,
2114
            np.array([1, 1, 0, 1, 1, 1])
2115
        )
2116

2117
    def test_recovery_duration_ratio(self):
2118
        np.testing.assert_array_almost_equal(
2119
            drawdowns['a'].recovery_duration_ratio.values,
2120
            np.array([1., 1., 0.])
2121
        )
2122
        np.testing.assert_array_almost_equal(
2123
            drawdowns.recovery_duration_ratio.values,
2124
            np.array([1., 1., 0., 1., 1., 0.5])
2125
        )
2126

2127
    def test_active_records(self):
2128
        assert isinstance(drawdowns.active, vbt.Drawdowns)
2129
        assert drawdowns.active.wrapper == drawdowns.wrapper
2130
        record_arrays_close(
2131
            drawdowns['a'].active.values,
2132
            np.array([
2133
                (2, 0, 4, 5, 5, 5, 4., 1., 1., 0)
2134
            ], dtype=drawdown_dt)
2135
        )
2136
        record_arrays_close(
2137
            drawdowns['a'].active.values,
2138
            drawdowns.active['a'].values
2139
        )
2140
        record_arrays_close(
2141
            drawdowns.active.values,
2142
            np.array([
2143
                (2, 0, 4, 5, 5, 5, 4.0, 1.0, 1.0, 0), (5, 2, 2, 3, 4, 5, 3.0, 1.0, 2.0, 0)
2144
            ], dtype=drawdown_dt)
2145
        )
2146

2147
    def test_recovered_records(self):
2148
        assert isinstance(drawdowns.recovered, vbt.Drawdowns)
2149
        assert drawdowns.recovered.wrapper == drawdowns.wrapper
2150
        record_arrays_close(
2151
            drawdowns['a'].recovered.values,
2152
            np.array([
2153
                (0, 0, 0, 1, 1, 2, 2.0, 1.0, 3.0, 1), (1, 0, 2, 3, 3, 4, 3.0, 1.0, 4.0, 1)
2154
            ], dtype=drawdown_dt)
2155
        )
2156
        record_arrays_close(
2157
            drawdowns['a'].recovered.values,
2158
            drawdowns.recovered['a'].values
2159
        )
2160
        record_arrays_close(
2161
            drawdowns.recovered.values,
2162
            np.array([
2163
                (0, 0, 0, 1, 1, 2, 2.0, 1.0, 3.0, 1), (1, 0, 2, 3, 3, 4, 3.0, 1.0, 4.0, 1),
2164
                (3, 1, 1, 2, 2, 3, 2.0, 1.0, 3.0, 1), (4, 1, 3, 4, 4, 5, 3.0, 1.0, 4.0, 1)
2165
            ], dtype=drawdown_dt)
2166
        )
2167

2168
    def test_active_drawdown(self):
2169
        assert drawdowns['a'].active_drawdown() == -0.75
2170
        pd.testing.assert_series_equal(
2171
            drawdowns.active_drawdown(),
2172
            pd.Series(
2173
                np.array([-0.75, np.nan, -0.3333333333333333, np.nan]),
2174
                index=wrapper.columns
2175
            ).rename('active_drawdown')
2176
        )
2177
        with pytest.raises(Exception):
2178
            drawdowns_grouped.active_drawdown()
2179

2180
    def test_active_duration(self):
2181
        assert drawdowns['a'].active_duration() == np.timedelta64(86400000000000)
2182
        pd.testing.assert_series_equal(
2183
            drawdowns.active_duration(),
2184
            pd.Series(
2185
                np.array([86400000000000, 'NaT', 259200000000000, 'NaT'], dtype='timedelta64[ns]'),
2186
                index=wrapper.columns
2187
            ).rename('active_duration')
2188
        )
2189
        with pytest.raises(Exception):
2190
            drawdowns_grouped.active_duration()
2191

2192
    def test_active_recovery(self):
2193
        assert drawdowns['a'].active_recovery() == 0.
2194
        pd.testing.assert_series_equal(
2195
            drawdowns.active_recovery(),
2196
            pd.Series(
2197
                np.array([0., np.nan, 0.5, np.nan]),
2198
                index=wrapper.columns
2199
            ).rename('active_recovery')
2200
        )
2201
        with pytest.raises(Exception):
2202
            drawdowns_grouped.active_recovery()
2203

2204
    def test_active_recovery_return(self):
2205
        assert drawdowns['a'].active_recovery_return() == 0.
2206
        pd.testing.assert_series_equal(
2207
            drawdowns.active_recovery_return(),
2208
            pd.Series(
2209
                np.array([0., np.nan, 1., np.nan]),
2210
                index=wrapper.columns
2211
            ).rename('active_recovery_return')
2212
        )
2213
        with pytest.raises(Exception):
2214
            drawdowns_grouped.active_recovery_return()
2215

2216
    def test_active_recovery_duration(self):
2217
        assert drawdowns['a'].active_recovery_duration() == pd.Timedelta('0 days 00:00:00')
2218
        pd.testing.assert_series_equal(
2219
            drawdowns.active_recovery_duration(),
2220
            pd.Series(
2221
                np.array([0, 'NaT', 86400000000000, 'NaT'], dtype='timedelta64[ns]'),
2222
                index=wrapper.columns
2223
            ).rename('active_recovery_duration')
2224
        )
2225
        with pytest.raises(Exception):
2226
            drawdowns_grouped.active_recovery_duration()
2227

2228
    def test_stats(self):
2229
        stats_index = pd.Index([
2230
            'Start', 'End', 'Period', 'Coverage [%]', 'Total Records',
2231
            'Total Recovered Drawdowns', 'Total Active Drawdowns',
2232
            'Active Drawdown [%]', 'Active Duration', 'Active Recovery [%]',
2233
            'Active Recovery Return [%]', 'Active Recovery Duration',
2234
            'Max Drawdown [%]', 'Avg Drawdown [%]', 'Max Drawdown Duration',
2235
            'Avg Drawdown Duration', 'Max Recovery Return [%]',
2236
            'Avg Recovery Return [%]', 'Max Recovery Duration',
2237
            'Avg Recovery Duration', 'Avg Recovery Duration Ratio'
2238
        ], dtype='object')
2239
        pd.testing.assert_series_equal(
2240
            drawdowns.stats(),
2241
            pd.Series([
2242
                pd.Timestamp('2020-01-01 00:00:00'), pd.Timestamp('2020-01-06 00:00:00'),
2243
                pd.Timedelta('6 days 00:00:00'), 44.444444444444436, 1.5, 1.0, 0.5,
2244
                54.166666666666664, pd.Timedelta('2 days 00:00:00'), 25.0, 50.0,
2245
                pd.Timedelta('0 days 12:00:00'), 66.66666666666666, 58.33333333333333,
2246
                pd.Timedelta('1 days 00:00:00'), pd.Timedelta('1 days 00:00:00'), 300.0, 250.0,
2247
                pd.Timedelta('1 days 00:00:00'), pd.Timedelta('1 days 00:00:00'), 1.0
2248
            ],
2249
                index=stats_index,
2250
                name='agg_func_mean'
2251
            )
2252
        )
2253
        pd.testing.assert_series_equal(
2254
            drawdowns.stats(settings=dict(incl_active=True)),
2255
            pd.Series([
2256
                pd.Timestamp('2020-01-01 00:00:00'), pd.Timestamp('2020-01-06 00:00:00'),
2257
                pd.Timedelta('6 days 00:00:00'), 44.444444444444436, 1.5, 1.0, 0.5,
2258
                54.166666666666664, pd.Timedelta('2 days 00:00:00'), 25.0, 50.0,
2259
                pd.Timedelta('0 days 12:00:00'), 69.44444444444444, 62.962962962962955,
2260
                pd.Timedelta('1 days 16:00:00'), pd.Timedelta('1 days 16:00:00'), 300.0, 250.0,
2261
                pd.Timedelta('1 days 00:00:00'), pd.Timedelta('1 days 00:00:00'), 1.0
2262
            ],
2263
                index=stats_index,
2264
                name='agg_func_mean'
2265
            )
2266
        )
2267
        pd.testing.assert_series_equal(
2268
            drawdowns.stats(column='a'),
2269
            pd.Series([
2270
                pd.Timestamp('2020-01-01 00:00:00'), pd.Timestamp('2020-01-06 00:00:00'),
2271
                pd.Timedelta('6 days 00:00:00'), 50.0, 3, 2, 1, 75.0, pd.Timedelta('1 days 00:00:00'),
2272
                0.0, 0.0, pd.Timedelta('0 days 00:00:00'), 66.66666666666666, 58.33333333333333,
2273
                pd.Timedelta('1 days 00:00:00'), pd.Timedelta('1 days 00:00:00'), 300.0, 250.0,
2274
                pd.Timedelta('1 days 00:00:00'), pd.Timedelta('1 days 00:00:00'), 1.0
2275
            ],
2276
                index=stats_index,
2277
                name='a'
2278
            )
2279
        )
2280
        pd.testing.assert_series_equal(
2281
            drawdowns.stats(column='g1', group_by=group_by),
2282
            pd.Series([
2283
                pd.Timestamp('2020-01-01 00:00:00'), pd.Timestamp('2020-01-06 00:00:00'),
2284
                pd.Timedelta('6 days 00:00:00'), 41.66666666666667, 5, 4, 1, 66.66666666666666,
2285
                58.33333333333333, pd.Timedelta('1 days 00:00:00'), pd.Timedelta('1 days 00:00:00'),
2286
                300.0, 250.0, pd.Timedelta('1 days 00:00:00'), pd.Timedelta('1 days 00:00:00'), 1.0
2287
            ],
2288
                index=pd.Index([
2289
                    'Start', 'End', 'Period', 'Coverage [%]', 'Total Records',
2290
                    'Total Recovered Drawdowns', 'Total Active Drawdowns',
2291
                    'Max Drawdown [%]', 'Avg Drawdown [%]', 'Max Drawdown Duration',
2292
                    'Avg Drawdown Duration', 'Max Recovery Return [%]',
2293
                    'Avg Recovery Return [%]', 'Max Recovery Duration',
2294
                    'Avg Recovery Duration', 'Avg Recovery Duration Ratio'
2295
                ], dtype='object'),
2296
                name='g1'
2297
            )
2298
        )
2299
        pd.testing.assert_series_equal(
2300
            drawdowns['c'].stats(),
2301
            drawdowns.stats(column='c')
2302
        )
2303
        pd.testing.assert_series_equal(
2304
            drawdowns['c'].stats(),
2305
            drawdowns.stats(column='c', group_by=False)
2306
        )
2307
        pd.testing.assert_series_equal(
2308
            drawdowns_grouped['g2'].stats(),
2309
            drawdowns_grouped.stats(column='g2')
2310
        )
2311
        pd.testing.assert_series_equal(
2312
            drawdowns_grouped['g2'].stats(),
2313
            drawdowns.stats(column='g2', group_by=group_by)
2314
        )
2315
        stats_df = drawdowns.stats(agg_func=None)
2316
        assert stats_df.shape == (4, 21)
2317
        pd.testing.assert_index_equal(stats_df.index, drawdowns.wrapper.columns)
2318
        pd.testing.assert_index_equal(stats_df.columns, stats_index)
2319

2320

2321
# ############# orders.py ############# #
2322

2323
close = pd.Series([1, 2, 3, 4, 5, 6, 7, 8], index=[
2324
    datetime(2020, 1, 1),
2325
    datetime(2020, 1, 2),
2326
    datetime(2020, 1, 3),
2327
    datetime(2020, 1, 4),
2328
    datetime(2020, 1, 5),
2329
    datetime(2020, 1, 6),
2330
    datetime(2020, 1, 7),
2331
    datetime(2020, 1, 8)
2332
]).vbt.tile(4, keys=['a', 'b', 'c', 'd'])
2333

2334
size = np.full(close.shape, np.nan, dtype=np.float64)
2335
size[:, 0] = [1, 0.1, -1, -0.1, np.nan, 1, -1, 2]
2336
size[:, 1] = [-1, -0.1, 1, 0.1, np.nan, -1, 1, -2]
2337
size[:, 2] = [1, 0.1, -1, -0.1, np.nan, 1, -2, 2]
2338
orders = vbt.Portfolio.from_orders(close, size, fees=0.01, freq='1 days').orders
2339
orders_grouped = orders.regroup(group_by)
2340

2341

2342
class TestOrders:
2343
    def test_mapped_fields(self):
2344
        for name in order_dt.names:
2345
            np.testing.assert_array_equal(
2346
                getattr(orders, name).values,
2347
                orders.values[name]
2348
            )
2349

2350
    def test_close(self):
2351
        pd.testing.assert_frame_equal(
2352
            orders.close,
2353
            close
2354
        )
2355
        pd.testing.assert_series_equal(
2356
            orders['a'].close,
2357
            close['a']
2358
        )
2359
        pd.testing.assert_frame_equal(
2360
            orders_grouped['g1'].close,
2361
            close[['a', 'b']]
2362
        )
2363
        assert orders.replace(close=None)['a'].close is None
2364

2365
    def test_records_readable(self):
2366
        records_readable = orders.records_readable
2367

2368
        np.testing.assert_array_equal(
2369
            records_readable['Order Id'].values,
2370
            np.array([
2371
                0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20
2372
            ])
2373
        )
2374
        np.testing.assert_array_equal(
2375
            records_readable['Timestamp'].values,
2376
            np.array([
2377
                '2020-01-01T00:00:00.000000000', '2020-01-02T00:00:00.000000000',
2378
                '2020-01-03T00:00:00.000000000', '2020-01-04T00:00:00.000000000',
2379
                '2020-01-06T00:00:00.000000000', '2020-01-07T00:00:00.000000000',
2380
                '2020-01-08T00:00:00.000000000', '2020-01-01T00:00:00.000000000',
2381
                '2020-01-02T00:00:00.000000000', '2020-01-03T00:00:00.000000000',
2382
                '2020-01-04T00:00:00.000000000', '2020-01-06T00:00:00.000000000',
2383
                '2020-01-07T00:00:00.000000000', '2020-01-08T00:00:00.000000000',
2384
                '2020-01-01T00:00:00.000000000', '2020-01-02T00:00:00.000000000',
2385
                '2020-01-03T00:00:00.000000000', '2020-01-04T00:00:00.000000000',
2386
                '2020-01-06T00:00:00.000000000', '2020-01-07T00:00:00.000000000',
2387
                '2020-01-08T00:00:00.000000000'
2388
            ], dtype='datetime64[ns]')
2389
        )
2390
        np.testing.assert_array_equal(
2391
            records_readable['Column'].values,
2392
            np.array([
2393
                'a', 'a', 'a', 'a', 'a', 'a', 'a', 'b', 'b', 'b', 'b', 'b', 'b',
2394
                'b', 'c', 'c', 'c', 'c', 'c', 'c', 'c'
2395
            ])
2396
        )
2397
        np.testing.assert_array_equal(
2398
            records_readable['Size'].values,
2399
            np.array([
2400
                1.0, 0.1, 1.0, 0.1, 1.0, 1.0, 2.0, 1.0, 0.1, 1.0, 0.1, 1.0, 1.0,
2401
                2.0, 1.0, 0.1, 1.0, 0.1, 1.0, 2.0, 2.0
2402
            ])
2403
        )
2404
        np.testing.assert_array_equal(
2405
            records_readable['Price'].values,
2406
            np.array([
2407
                1.0, 2.0, 3.0, 4.0, 6.0, 7.0, 8.0, 1.0, 2.0, 3.0, 4.0, 6.0, 7.0,
2408
                8.0, 1.0, 2.0, 3.0, 4.0, 6.0, 7.0, 8.0
2409
            ])
2410
        )
2411
        np.testing.assert_array_equal(
2412
            records_readable['Fees'].values,
2413
            np.array([
2414
                0.01, 0.002, 0.03, 0.004, 0.06, 0.07, 0.16, 0.01, 0.002, 0.03,
2415
                0.004, 0.06, 0.07, 0.16, 0.01, 0.002, 0.03, 0.004, 0.06, 0.14,
2416
                0.16
2417
            ])
2418
        )
2419
        np.testing.assert_array_equal(
2420
            records_readable['Side'].values,
2421
            np.array([
2422
                'Buy', 'Buy', 'Sell', 'Sell', 'Buy', 'Sell', 'Buy', 'Sell', 'Sell',
2423
                'Buy', 'Buy', 'Sell', 'Buy', 'Sell', 'Buy', 'Buy', 'Sell', 'Sell',
2424
                'Buy', 'Sell', 'Buy'
2425
            ])
2426
        )
2427

2428
    def test_buy_records(self):
2429
        assert isinstance(orders.buy, vbt.Orders)
2430
        assert orders.buy.wrapper == orders.wrapper
2431
        record_arrays_close(
2432
            orders['a'].buy.values,
2433
            np.array([
2434
                (0, 0, 0, 1., 1., 0.01, 0), (1, 0, 1, 0.1, 2., 0.002, 0),
2435
                (4, 0, 5, 1., 6., 0.06, 0), (6, 0, 7, 2., 8., 0.16, 0)
2436
            ], dtype=order_dt)
2437
        )
2438
        record_arrays_close(
2439
            orders['a'].buy.values,
2440
            orders.buy['a'].values
2441
        )
2442
        record_arrays_close(
2443
            orders.buy.values,
2444
            np.array([
2445
                (0, 0, 0, 1., 1., 0.01, 0), (1, 0, 1, 0.1, 2., 0.002, 0),
2446
                (4, 0, 5, 1., 6., 0.06, 0), (6, 0, 7, 2., 8., 0.16, 0),
2447
                (9, 1, 2, 1., 3., 0.03, 0), (10, 1, 3, 0.1, 4., 0.004, 0),
2448
                (12, 1, 6, 1., 7., 0.07, 0), (14, 2, 0, 1., 1., 0.01, 0),
2449
                (15, 2, 1, 0.1, 2., 0.002, 0), (18, 2, 5, 1., 6., 0.06, 0),
2450
                (20, 2, 7, 2., 8., 0.16, 0)
2451
            ], dtype=order_dt)
2452
        )
2453

2454
    def test_sell_records(self):
2455
        assert isinstance(orders.sell, vbt.Orders)
2456
        assert orders.sell.wrapper == orders.wrapper
2457
        record_arrays_close(
2458
            orders['a'].sell.values,
2459
            np.array([
2460
                (2, 0, 2, 1., 3., 0.03, 1), (3, 0, 3, 0.1, 4., 0.004, 1),
2461
                (5, 0, 6, 1., 7., 0.07, 1)
2462
            ], dtype=order_dt)
2463
        )
2464
        record_arrays_close(
2465
            orders['a'].sell.values,
2466
            orders.sell['a'].values
2467
        )
2468
        record_arrays_close(
2469
            orders.sell.values,
2470
            np.array([
2471
                (2, 0, 2, 1., 3., 0.03, 1), (3, 0, 3, 0.1, 4., 0.004, 1),
2472
                (5, 0, 6, 1., 7., 0.07, 1), (7, 1, 0, 1., 1., 0.01, 1),
2473
                (8, 1, 1, 0.1, 2., 0.002, 1), (11, 1, 5, 1., 6., 0.06, 1),
2474
                (13, 1, 7, 2., 8., 0.16, 1), (16, 2, 2, 1., 3., 0.03, 1),
2475
                (17, 2, 3, 0.1, 4., 0.004, 1), (19, 2, 6, 2., 7., 0.14, 1)
2476
            ], dtype=order_dt)
2477
        )
2478

2479
    def test_stats(self):
2480
        stats_index = pd.Index([
2481
            'Start', 'End', 'Period', 'Total Records', 'Total Buy Orders', 'Total Sell Orders',
2482
            'Min Size', 'Max Size', 'Avg Size', 'Avg Buy Size', 'Avg Sell Size',
2483
            'Avg Buy Price', 'Avg Sell Price', 'Total Fees', 'Min Fees', 'Max Fees',
2484
            'Avg Fees', 'Avg Buy Fees', 'Avg Sell Fees'
2485
        ], dtype='object')
2486
        pd.testing.assert_series_equal(
2487
            orders.stats(),
2488
            pd.Series([
2489
                pd.Timestamp('2020-01-01 00:00:00'), pd.Timestamp('2020-01-08 00:00:00'),
2490
                pd.Timedelta('8 days 00:00:00'), 5.25, 2.75, 2.5, 0.10000000000000002, 2.0,
2491
                0.9333333333333335, 0.9166666666666666, 0.9194444444444446, 4.388888888888889,
2492
                4.527777777777779, 0.26949999999999996, 0.002, 0.16, 0.051333333333333335,
2493
                0.050222222222222224, 0.050222222222222224
2494
            ],
2495
                index=stats_index,
2496
                name='agg_func_mean'
2497
            )
2498
        )
2499
        pd.testing.assert_series_equal(
2500
            orders.stats(column='a'),
2501
            pd.Series([
2502
                pd.Timestamp('2020-01-01 00:00:00'), pd.Timestamp('2020-01-08 00:00:00'),
2503
                pd.Timedelta('8 days 00:00:00'), 7, 4, 3, 0.1, 2.0, 0.8857142857142858,
2504
                1.025, 0.7000000000000001, 4.25, 4.666666666666667, 0.33599999999999997,
2505
                0.002, 0.16, 0.047999999999999994, 0.057999999999999996, 0.03466666666666667
2506
            ],
2507
                index=stats_index,
2508
                name='a'
2509
            )
2510
        )
2511
        pd.testing.assert_series_equal(
2512
            orders.stats(column='g1', group_by=group_by),
2513
            pd.Series([
2514
                pd.Timestamp('2020-01-01 00:00:00'), pd.Timestamp('2020-01-08 00:00:00'),
2515
                pd.Timedelta('8 days 00:00:00'), 14, 7, 7, 0.1, 2.0, 0.8857142857142858,
2516
                0.8857142857142856, 0.8857142857142858, 4.428571428571429, 4.428571428571429,
2517
                0.672, 0.002, 0.16, 0.048, 0.048, 0.047999999999999994
2518
            ],
2519
                index=stats_index,
2520
                name='g1'
2521
            )
2522
        )
2523
        pd.testing.assert_series_equal(
2524
            orders['c'].stats(),
2525
            orders.stats(column='c')
2526
        )
2527
        pd.testing.assert_series_equal(
2528
            orders['c'].stats(),
2529
            orders.stats(column='c', group_by=False)
2530
        )
2531
        pd.testing.assert_series_equal(
2532
            orders_grouped['g2'].stats(),
2533
            orders_grouped.stats(column='g2')
2534
        )
2535
        pd.testing.assert_series_equal(
2536
            orders_grouped['g2'].stats(),
2537
            orders.stats(column='g2', group_by=group_by)
2538
        )
2539
        stats_df = orders.stats(agg_func=None)
2540
        assert stats_df.shape == (4, 19)
2541
        pd.testing.assert_index_equal(stats_df.index, orders.wrapper.columns)
2542
        pd.testing.assert_index_equal(stats_df.columns, stats_index)
2543

2544

2545
# ############# trades.py ############# #
2546

2547
exit_trades = vbt.ExitTrades.from_orders(orders)
2548
exit_trades_grouped = vbt.ExitTrades.from_orders(orders_grouped)
2549

2550

2551
class TestExitTrades:
2552
    def test_mapped_fields(self):
2553
        for name in trade_dt.names:
2554
            if name == 'return':
2555
                np.testing.assert_array_equal(
2556
                    getattr(exit_trades, 'returns').values,
2557
                    exit_trades.values[name]
2558
                )
2559
            else:
2560
                np.testing.assert_array_equal(
2561
                    getattr(exit_trades, name).values,
2562
                    exit_trades.values[name]
2563
                )
2564

2565
    def test_close(self):
2566
        pd.testing.assert_frame_equal(
2567
            exit_trades.close,
2568
            close
2569
        )
2570
        pd.testing.assert_series_equal(
2571
            exit_trades['a'].close,
2572
            close['a']
2573
        )
2574
        pd.testing.assert_frame_equal(
2575
            exit_trades_grouped['g1'].close,
2576
            close[['a', 'b']]
2577
        )
2578
        assert exit_trades.replace(close=None)['a'].close is None
2579

2580
    def test_records_arr(self):
2581
        record_arrays_close(
2582
            exit_trades.values,
2583
            np.array([
2584
                (0, 0, 1., 0, 1.09090909, 0.01090909, 2, 3., 0.03, 1.86818182, 1.7125, 0, 1, 0),
2585
                (1, 0, 0.1, 0, 1.09090909, 0.00109091, 3, 4., 0.004, 0.28581818, 2.62, 0, 1, 0),
2586
                (2, 0, 1., 5, 6., 0.06, 6, 7., 0.07, 0.87, 0.145, 0, 1, 1),
2587
                (3, 0, 2., 7, 8., 0.16, 7, 8., 0., -0.16, -0.01, 0, 0, 2),
2588
                (4, 1, 1., 0, 1.09090909, 0.01090909, 2, 3., 0.03, -1.95, -1.7875, 1, 1, 3),
2589
                (5, 1, 0.1, 0, 1.09090909, 0.00109091, 3, 4., 0.004, -0.296, -2.71333333, 1, 1, 3),
2590
                (6, 1, 1., 5, 6., 0.06, 6, 7., 0.07, -1.13, -0.18833333, 1, 1, 4),
2591
                (7, 1, 2., 7, 8., 0.16, 7, 8., 0., -0.16, -0.01, 1, 0, 5),
2592
                (8, 2, 1., 0, 1.09090909, 0.01090909, 2, 3., 0.03, 1.86818182, 1.7125, 0, 1, 6),
2593
                (9, 2, 0.1, 0, 1.09090909, 0.00109091, 3, 4., 0.004, 0.28581818, 2.62, 0, 1, 6),
2594
                (10, 2, 1., 5, 6., 0.06, 6, 7., 0.07, 0.87, 0.145, 0, 1, 7),
2595
                (11, 2, 1., 6, 7., 0.07, 7, 8., 0.08, -1.15, -0.16428571, 1, 1, 8),
2596
                (12, 2, 1., 7, 8., 0.08, 7, 8., 0., -0.08, -0.01, 0, 0, 9)
2597
            ], dtype=trade_dt)
2598
        )
2599
        reversed_col_orders = orders.replace(records_arr=np.concatenate((
2600
            orders.values[orders.values['col'] == 2],
2601
            orders.values[orders.values['col'] == 1],
2602
            orders.values[orders.values['col'] == 0]
2603
        )))
2604
        record_arrays_close(
2605
            vbt.ExitTrades.from_orders(reversed_col_orders).values,
2606
            exit_trades.values
2607
        )
2608

2609
    def test_records_readable(self):
2610
        records_readable = exit_trades.records_readable
2611

2612
        np.testing.assert_array_equal(
2613
            records_readable['Exit Trade Id'].values,
2614
            np.array([
2615
                0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12
2616
            ])
2617
        )
2618
        np.testing.assert_array_equal(
2619
            records_readable['Column'].values,
2620
            np.array([
2621
                'a', 'a', 'a', 'a', 'b', 'b', 'b', 'b', 'c', 'c', 'c', 'c', 'c'
2622
            ])
2623
        )
2624
        np.testing.assert_array_equal(
2625
            records_readable['Size'].values,
2626
            np.array([
2627
                1.0, 0.10000000000000009, 1.0, 2.0, 1.0, 0.10000000000000009, 1.0,
2628
                2.0, 1.0, 0.10000000000000009, 1.0, 1.0, 1.0
2629
            ])
2630
        )
2631
        np.testing.assert_array_equal(
2632
            records_readable['Entry Timestamp'].values,
2633
            np.array([
2634
                '2020-01-01T00:00:00.000000000', '2020-01-01T00:00:00.000000000',
2635
                '2020-01-06T00:00:00.000000000', '2020-01-08T00:00:00.000000000',
2636
                '2020-01-01T00:00:00.000000000', '2020-01-01T00:00:00.000000000',
2637
                '2020-01-06T00:00:00.000000000', '2020-01-08T00:00:00.000000000',
2638
                '2020-01-01T00:00:00.000000000', '2020-01-01T00:00:00.000000000',
2639
                '2020-01-06T00:00:00.000000000', '2020-01-07T00:00:00.000000000',
2640
                '2020-01-08T00:00:00.000000000'
2641
            ], dtype='datetime64[ns]')
2642
        )
2643
        np.testing.assert_array_equal(
2644
            records_readable['Avg Entry Price'].values,
2645
            np.array([
2646
                1.0909090909090908, 1.0909090909090908, 6.0, 8.0,
2647
                1.0909090909090908, 1.0909090909090908, 6.0, 8.0,
2648
                1.0909090909090908, 1.0909090909090908, 6.0, 7.0, 8.0
2649
            ])
2650
        )
2651
        np.testing.assert_array_equal(
2652
            records_readable['Entry Fees'].values,
2653
            np.array([
2654
                0.010909090909090908, 0.0010909090909090918, 0.06, 0.16,
2655
                0.010909090909090908, 0.0010909090909090918, 0.06, 0.16,
2656
                0.010909090909090908, 0.0010909090909090918, 0.06, 0.07, 0.08
2657
            ])
2658
        )
2659
        np.testing.assert_array_equal(
2660
            records_readable['Exit Timestamp'].values,
2661
            np.array([
2662
                '2020-01-03T00:00:00.000000000', '2020-01-04T00:00:00.000000000',
2663
                '2020-01-07T00:00:00.000000000', '2020-01-08T00:00:00.000000000',
2664
                '2020-01-03T00:00:00.000000000', '2020-01-04T00:00:00.000000000',
2665
                '2020-01-07T00:00:00.000000000', '2020-01-08T00:00:00.000000000',
2666
                '2020-01-03T00:00:00.000000000', '2020-01-04T00:00:00.000000000',
2667
                '2020-01-07T00:00:00.000000000', '2020-01-08T00:00:00.000000000',
2668
                '2020-01-08T00:00:00.000000000'
2669
            ], dtype='datetime64[ns]')
2670
        )
2671
        np.testing.assert_array_equal(
2672
            records_readable['Avg Exit Price'].values,
2673
            np.array([
2674
                3.0, 4.0, 7.0, 8.0, 3.0, 4.0, 7.0, 8.0, 3.0, 4.0, 7.0, 8.0, 8.0
2675
            ])
2676
        )
2677
        np.testing.assert_array_equal(
2678
            records_readable['Exit Fees'].values,
2679
            np.array([
2680
                0.03, 0.004, 0.07, 0.0, 0.03, 0.004, 0.07, 0.0, 0.03, 0.004, 0.07, 0.08, 0.0
2681
            ])
2682
        )
2683
        np.testing.assert_array_equal(
2684
            records_readable['PnL'].values,
2685
            np.array([
2686
                1.8681818181818182, 0.2858181818181821, 0.8699999999999999, -0.16,
2687
                -1.9500000000000002, -0.29600000000000026, -1.1300000000000001,
2688
                -0.16, 1.8681818181818182, 0.2858181818181821, 0.8699999999999999,
2689
                -1.1500000000000001, -0.08
2690
            ])
2691
        )
2692
        np.testing.assert_array_equal(
2693
            records_readable['Return'].values,
2694
            np.array([
2695
                1.7125000000000001, 2.62, 0.145, -0.01, -1.7875000000000003,
2696
                -2.7133333333333334, -0.18833333333333335, -0.01,
2697
                1.7125000000000001, 2.62, 0.145, -0.1642857142857143, -0.01
2698
            ])
2699
        )
2700
        np.testing.assert_array_equal(
2701
            records_readable['Direction'].values,
2702
            np.array([
2703
                'Long', 'Long', 'Long', 'Long', 'Short', 'Short', 'Short',
2704
                'Short', 'Long', 'Long', 'Long', 'Short', 'Long'
2705
            ])
2706
        )
2707
        np.testing.assert_array_equal(
2708
            records_readable['Status'].values,
2709
            np.array([
2710
                'Closed', 'Closed', 'Closed', 'Open', 'Closed', 'Closed', 'Closed',
2711
                'Open', 'Closed', 'Closed', 'Closed', 'Closed', 'Open'
2712
            ])
2713
        )
2714
        np.testing.assert_array_equal(
2715
            records_readable['Position Id'].values,
2716
            np.array([
2717
                0, 0, 1, 2, 3, 3, 4, 5, 6, 6, 7, 8, 9
2718
            ])
2719
        )
2720

2721
    def test_duration(self):
2722
        np.testing.assert_array_almost_equal(
2723
            exit_trades['a'].duration.values,
2724
            np.array([2, 3, 1, 1])
2725
        )
2726
        np.testing.assert_array_almost_equal(
2727
            exit_trades.duration.values,
2728
            np.array([2, 3, 1, 1, 2, 3, 1, 1, 2, 3, 1, 1, 1])
2729
        )
2730

2731
    def test_winning_records(self):
2732
        assert isinstance(exit_trades.winning, vbt.ExitTrades)
2733
        assert exit_trades.winning.wrapper == exit_trades.wrapper
2734
        record_arrays_close(
2735
            exit_trades['a'].winning.values,
2736
            np.array([
2737
                (0, 0, 1., 0, 1.09090909, 0.01090909, 2, 3., 0.03, 1.86818182, 1.7125, 0, 1, 0),
2738
                (1, 0, 0.1, 0, 1.09090909, 0.00109091, 3, 4., 0.004, 0.28581818, 2.62, 0, 1, 0),
2739
                (2, 0, 1., 5, 6., 0.06, 6, 7., 0.07, 0.87, 0.145, 0, 1, 1)
2740
            ], dtype=trade_dt)
2741
        )
2742
        record_arrays_close(
2743
            exit_trades['a'].winning.values,
2744
            exit_trades.winning['a'].values
2745
        )
2746
        record_arrays_close(
2747
            exit_trades.winning.values,
2748
            np.array([
2749
                (0, 0, 1., 0, 1.09090909, 0.01090909, 2, 3., 0.03, 1.86818182, 1.7125, 0, 1, 0),
2750
                (1, 0, 0.1, 0, 1.09090909, 0.00109091, 3, 4., 0.004, 0.28581818, 2.62, 0, 1, 0),
2751
                (2, 0, 1., 5, 6., 0.06, 6, 7., 0.07, 0.87, 0.145, 0, 1, 1),
2752
                (8, 2, 1., 0, 1.09090909, 0.01090909, 2, 3., 0.03, 1.86818182, 1.7125, 0, 1, 6),
2753
                (9, 2, 0.1, 0, 1.09090909, 0.00109091, 3, 4., 0.004, 0.28581818, 2.62, 0, 1, 6),
2754
                (10, 2, 1., 5, 6., 0.06, 6, 7., 0.07, 0.87, 0.145, 0, 1, 7)
2755
            ], dtype=trade_dt)
2756
        )
2757

2758
    def test_losing_records(self):
2759
        assert isinstance(exit_trades.losing, vbt.ExitTrades)
2760
        assert exit_trades.losing.wrapper == exit_trades.wrapper
2761
        record_arrays_close(
2762
            exit_trades['a'].losing.values,
2763
            np.array([
2764
                (3, 0, 2., 7, 8., 0.16, 7, 8., 0., -0.16, -0.01, 0, 0, 2)
2765
            ], dtype=trade_dt)
2766
        )
2767
        record_arrays_close(
2768
            exit_trades['a'].losing.values,
2769
            exit_trades.losing['a'].values
2770
        )
2771
        record_arrays_close(
2772
            exit_trades.losing.values,
2773
            np.array([
2774
                (3, 0, 2., 7, 8., 0.16, 7, 8., 0., -0.16, -0.01, 0, 0, 2),
2775
                (4, 1, 1., 0, 1.09090909, 0.01090909, 2, 3., 0.03, -1.95, -1.7875, 1, 1, 3),
2776
                (5, 1, 0.1, 0, 1.09090909, 0.00109091, 3, 4., 0.004, -0.296, -2.71333333, 1, 1, 3),
2777
                (6, 1, 1., 5, 6., 0.06, 6, 7., 0.07, -1.13, -0.18833333, 1, 1, 4),
2778
                (7, 1, 2., 7, 8., 0.16, 7, 8., 0., -0.16, -0.01, 1, 0, 5),
2779
                (11, 2, 1., 6, 7., 0.07, 7, 8., 0.08, -1.15, -0.16428571, 1, 1, 8),
2780
                (12, 2, 1., 7, 8., 0.08, 7, 8., 0., -0.08, -0.01, 0, 0, 9)
2781
            ], dtype=trade_dt)
2782
        )
2783

2784
    def test_win_rate(self):
2785
        assert exit_trades['a'].win_rate() == 0.75
2786
        pd.testing.assert_series_equal(
2787
            exit_trades.win_rate(),
2788
            pd.Series(
2789
                np.array([0.75, 0., 0.6, np.nan]),
2790
                index=close.columns
2791
            ).rename('win_rate')
2792
        )
2793
        pd.testing.assert_series_equal(
2794
            exit_trades_grouped.win_rate(),
2795
            pd.Series(
2796
                np.array([0.375, 0.6]),
2797
                index=pd.Index(['g1', 'g2'], dtype='object')
2798
            ).rename('win_rate')
2799
        )
2800

2801
    def test_winning_streak(self):
2802
        np.testing.assert_array_almost_equal(
2803
            exit_trades['a'].winning_streak.values,
2804
            np.array([1, 2, 3, 0])
2805
        )
2806
        np.testing.assert_array_almost_equal(
2807
            exit_trades.winning_streak.values,
2808
            np.array([1, 2, 3, 0, 0, 0, 0, 0, 1, 2, 3, 0, 0])
2809
        )
2810

2811
    def test_losing_streak(self):
2812
        np.testing.assert_array_almost_equal(
2813
            exit_trades['a'].losing_streak.values,
2814
            np.array([0, 0, 0, 1])
2815
        )
2816
        np.testing.assert_array_almost_equal(
2817
            exit_trades.losing_streak.values,
2818
            np.array([0, 0, 0, 1, 1, 2, 3, 4, 0, 0, 0, 1, 2])
2819
        )
2820

2821
    def test_profit_factor(self):
2822
        assert exit_trades['a'].profit_factor() == 18.9
2823
        pd.testing.assert_series_equal(
2824
            exit_trades.profit_factor(),
2825
            pd.Series(
2826
                np.array([18.9, 0., 2.45853659, np.nan]),
2827
                index=ts2.columns
2828
            ).rename('profit_factor')
2829
        )
2830
        pd.testing.assert_series_equal(
2831
            exit_trades_grouped.profit_factor(),
2832
            pd.Series(
2833
                np.array([0.81818182, 2.45853659]),
2834
                index=pd.Index(['g1', 'g2'], dtype='object')
2835
            ).rename('profit_factor')
2836
        )
2837

2838
    def test_expectancy(self):
2839
        assert exit_trades['a'].expectancy() == 0.716
2840
        pd.testing.assert_series_equal(
2841
            exit_trades.expectancy(),
2842
            pd.Series(
2843
                np.array([0.716, -0.884, 0.3588, np.nan]),
2844
                index=ts2.columns
2845
            ).rename('expectancy')
2846
        )
2847
        pd.testing.assert_series_equal(
2848
            exit_trades_grouped.expectancy(),
2849
            pd.Series(
2850
                np.array([-0.084, 0.3588]),
2851
                index=pd.Index(['g1', 'g2'], dtype='object')
2852
            ).rename('expectancy')
2853
        )
2854

2855
    def test_sqn(self):
2856
        assert exit_trades['a'].sqn() == 1.634155521947584
2857
        pd.testing.assert_series_equal(
2858
            exit_trades.sqn(),
2859
            pd.Series(
2860
                np.array([1.63415552, -2.13007307, 0.71660403, np.nan]),
2861
                index=ts2.columns
2862
            ).rename('sqn')
2863
        )
2864
        pd.testing.assert_series_equal(
2865
            exit_trades_grouped.sqn(),
2866
            pd.Series(
2867
                np.array([-0.20404671, 0.71660403]),
2868
                index=pd.Index(['g1', 'g2'], dtype='object')
2869
            ).rename('sqn')
2870
        )
2871

2872
    def test_long_records(self):
2873
        assert isinstance(exit_trades.long, vbt.ExitTrades)
2874
        assert exit_trades.long.wrapper == exit_trades.wrapper
2875
        record_arrays_close(
2876
            exit_trades['a'].long.values,
2877
            np.array([
2878
                (0, 0, 1., 0, 1.09090909, 0.01090909, 2, 3., 0.03, 1.86818182, 1.7125, 0, 1, 0),
2879
                (1, 0, 0.1, 0, 1.09090909, 0.00109091, 3, 4., 0.004, 0.28581818, 2.62, 0, 1, 0),
2880
                (2, 0, 1., 5, 6., 0.06, 6, 7., 0.07, 0.87, 0.145, 0, 1, 1),
2881
                (3, 0, 2., 7, 8., 0.16, 7, 8., 0., -0.16, -0.01, 0, 0, 2)
2882
            ], dtype=trade_dt)
2883
        )
2884
        record_arrays_close(
2885
            exit_trades['a'].long.values,
2886
            exit_trades.long['a'].values
2887
        )
2888
        record_arrays_close(
2889
            exit_trades.long.values,
2890
            np.array([
2891
                (0, 0, 1., 0, 1.09090909, 0.01090909, 2, 3., 0.03, 1.86818182, 1.7125, 0, 1, 0),
2892
                (1, 0, 0.1, 0, 1.09090909, 0.00109091, 3, 4., 0.004, 0.28581818, 2.62, 0, 1, 0),
2893
                (2, 0, 1., 5, 6., 0.06, 6, 7., 0.07, 0.87, 0.145, 0, 1, 1),
2894
                (3, 0, 2., 7, 8., 0.16, 7, 8., 0., -0.16, -0.01, 0, 0, 2),
2895
                (8, 2, 1., 0, 1.09090909, 0.01090909, 2, 3., 0.03, 1.86818182, 1.7125, 0, 1, 6),
2896
                (9, 2, 0.1, 0, 1.09090909, 0.00109091, 3, 4., 0.004, 0.28581818, 2.62, 0, 1, 6),
2897
                (10, 2, 1., 5, 6., 0.06, 6, 7., 0.07, 0.87, 0.145, 0, 1, 7),
2898
                (12, 2, 1., 7, 8., 0.08, 7, 8., 0., -0.08, -0.01, 0, 0, 9)
2899
            ], dtype=trade_dt)
2900
        )
2901

2902
    def test_short_records(self):
2903
        assert isinstance(exit_trades.short, vbt.ExitTrades)
2904
        assert exit_trades.short.wrapper == exit_trades.wrapper
2905
        record_arrays_close(
2906
            exit_trades['a'].short.values,
2907
            np.array([], dtype=trade_dt)
2908
        )
2909
        record_arrays_close(
2910
            exit_trades['a'].short.values,
2911
            exit_trades.short['a'].values
2912
        )
2913
        record_arrays_close(
2914
            exit_trades.short.values,
2915
            np.array([
2916
                (4, 1, 1., 0, 1.09090909, 0.01090909, 2, 3., 0.03, -1.95, -1.7875, 1, 1, 3),
2917
                (5, 1, 0.1, 0, 1.09090909, 0.00109091, 3, 4., 0.004, -0.296, -2.71333333, 1, 1, 3),
2918
                (6, 1, 1., 5, 6., 0.06, 6, 7., 0.07, -1.13, -0.18833333, 1, 1, 4),
2919
                (7, 1, 2., 7, 8., 0.16, 7, 8., 0., -0.16, -0.01, 1, 0, 5),
2920
                (11, 2, 1., 6, 7., 0.07, 7, 8., 0.08, -1.15, -0.16428571, 1, 1, 8)
2921
            ], dtype=trade_dt)
2922
        )
2923

2924
    def test_open_records(self):
2925
        assert isinstance(exit_trades.open, vbt.ExitTrades)
2926
        assert exit_trades.open.wrapper == exit_trades.wrapper
2927
        record_arrays_close(
2928
            exit_trades['a'].open.values,
2929
            np.array([
2930
                (3, 0, 2., 7, 8., 0.16, 7, 8., 0., -0.16, -0.01, 0, 0, 2)
2931
            ], dtype=trade_dt)
2932
        )
2933
        record_arrays_close(
2934
            exit_trades['a'].open.values,
2935
            exit_trades.open['a'].values
2936
        )
2937
        record_arrays_close(
2938
            exit_trades.open.values,
2939
            np.array([
2940
                (3, 0, 2., 7, 8., 0.16, 7, 8., 0., -0.16, -0.01, 0, 0, 2),
2941
                (7, 1, 2., 7, 8., 0.16, 7, 8., 0., -0.16, -0.01, 1, 0, 5),
2942
                (12, 2, 1., 7, 8., 0.08, 7, 8., 0., -0.08, -0.01, 0, 0, 9)
2943
            ], dtype=trade_dt)
2944
        )
2945

2946
    def test_closed_records(self):
2947
        assert isinstance(exit_trades.closed, vbt.ExitTrades)
2948
        assert exit_trades.closed.wrapper == exit_trades.wrapper
2949
        record_arrays_close(
2950
            exit_trades['a'].closed.values,
2951
            np.array([
2952
                (0, 0, 1., 0, 1.09090909, 0.01090909, 2, 3., 0.03, 1.86818182, 1.7125, 0, 1, 0),
2953
                (1, 0, 0.1, 0, 1.09090909, 0.00109091, 3, 4., 0.004, 0.28581818, 2.62, 0, 1, 0),
2954
                (2, 0, 1., 5, 6., 0.06, 6, 7., 0.07, 0.87, 0.145, 0, 1, 1)
2955
            ], dtype=trade_dt)
2956
        )
2957
        record_arrays_close(
2958
            exit_trades['a'].closed.values,
2959
            exit_trades.closed['a'].values
2960
        )
2961
        record_arrays_close(
2962
            exit_trades.closed.values,
2963
            np.array([
2964
                (0, 0, 1., 0, 1.09090909, 0.01090909, 2, 3., 0.03, 1.86818182, 1.7125, 0, 1, 0),
2965
                (1, 0, 0.1, 0, 1.09090909, 0.00109091, 3, 4., 0.004, 0.28581818, 2.62, 0, 1, 0),
2966
                (2, 0, 1., 5, 6., 0.06, 6, 7., 0.07, 0.87, 0.145, 0, 1, 1),
2967
                (4, 1, 1., 0, 1.09090909, 0.01090909, 2, 3., 0.03, -1.95, -1.7875, 1, 1, 3),
2968
                (5, 1, 0.1, 0, 1.09090909, 0.00109091, 3, 4., 0.004, -0.296, -2.71333333, 1, 1, 3),
2969
                (6, 1, 1., 5, 6., 0.06, 6, 7., 0.07, -1.13, -0.18833333, 1, 1, 4),
2970
                (8, 2, 1., 0, 1.09090909, 0.01090909, 2, 3., 0.03, 1.86818182, 1.7125, 0, 1, 6),
2971
                (9, 2, 0.1, 0, 1.09090909, 0.00109091, 3, 4., 0.004, 0.28581818, 2.62, 0, 1, 6),
2972
                (10, 2, 1., 5, 6., 0.06, 6, 7., 0.07, 0.87, 0.145, 0, 1, 7),
2973
                (11, 2, 1., 6, 7., 0.07, 7, 8., 0.08, -1.15, -0.16428571, 1, 1, 8)
2974
            ], dtype=trade_dt)
2975
        )
2976

2977
    def test_stats(self):
2978
        stats_index = pd.Index([
2979
            'Start', 'End', 'Period', 'First Trade Start', 'Last Trade End',
2980
            'Coverage', 'Overlap Coverage', 'Total Records', 'Total Long Trades',
2981
            'Total Short Trades', 'Total Closed Trades', 'Total Open Trades',
2982
            'Open Trade PnL', 'Win Rate [%]', 'Max Win Streak', 'Max Loss Streak',
2983
            'Best Trade [%]', 'Worst Trade [%]', 'Avg Winning Trade [%]',
2984
            'Avg Losing Trade [%]', 'Avg Winning Trade Duration',
2985
            'Avg Losing Trade Duration', 'Profit Factor', 'Expectancy', 'SQN'
2986
        ], dtype='object')
2987
        pd.testing.assert_series_equal(
2988
            exit_trades.stats(),
2989
            pd.Series([
2990
                pd.Timestamp('2020-01-01 00:00:00'), pd.Timestamp('2020-01-08 00:00:00'),
2991
                pd.Timedelta('8 days 00:00:00'), pd.Timestamp('2020-01-01 00:00:00'),
2992
                pd.Timestamp('2020-01-08 00:00:00'), pd.Timedelta('5 days 08:00:00'),
2993
                pd.Timedelta('2 days 00:00:00'), 3.25, 2.0, 1.25, 2.5, 0.75, -0.1,
2994
                58.333333333333336, 2.0, 1.3333333333333333, 168.38888888888889,
2995
                -91.08730158730158, 149.25, -86.3670634920635, pd.Timedelta('2 days 00:00:00'),
2996
                pd.Timedelta('1 days 12:00:00'), np.inf, 0.11705555555555548, 0.18931590012681135
2997
            ],
2998
                index=stats_index,
2999
                name='agg_func_mean'
3000
            )
3001
        )
3002
        pd.testing.assert_series_equal(
3003
            exit_trades.stats(settings=dict(incl_open=True)),
3004
            pd.Series([
3005
                pd.Timestamp('2020-01-01 00:00:00'), pd.Timestamp('2020-01-08 00:00:00'),
3006
                pd.Timedelta('8 days 00:00:00'), pd.Timestamp('2020-01-01 00:00:00'),
3007
                pd.Timestamp('2020-01-08 00:00:00'), pd.Timedelta('5 days 08:00:00'),
3008
                pd.Timedelta('2 days 00:00:00'), 3.25, 2.0, 1.25, 2.5, 0.75, -0.1,
3009
                58.333333333333336, 2.0, 2.3333333333333335, 174.33333333333334,
3010
                -96.25396825396825, 149.25, -42.39781746031746, pd.Timedelta('2 days 00:00:00'),
3011
                pd.Timedelta('1 days 06:00:00'), 7.11951219512195, 0.06359999999999993, 0.07356215977397455
3012
            ],
3013
                index=stats_index,
3014
                name='agg_func_mean'
3015
            )
3016
        )
3017
        pd.testing.assert_series_equal(
3018
            exit_trades.stats(column='a'),
3019
            pd.Series([
3020
                pd.Timestamp('2020-01-01 00:00:00'), pd.Timestamp('2020-01-08 00:00:00'),
3021
                pd.Timedelta('8 days 00:00:00'), pd.Timestamp('2020-01-01 00:00:00'),
3022
                pd.Timestamp('2020-01-08 00:00:00'), pd.Timedelta('5 days 00:00:00'),
3023
                pd.Timedelta('2 days 00:00:00'), 4, 4, 0, 3, 1, -0.16, 100.0, 3, 0,
3024
                262.0, 14.499999999999998, 149.25, np.nan, pd.Timedelta('2 days 00:00:00'),
3025
                pd.NaT, np.inf, 1.008, 2.181955050824476
3026
            ],
3027
                index=stats_index,
3028
                name='a'
3029
            )
3030
        )
3031
        pd.testing.assert_series_equal(
3032
            exit_trades.stats(column='g1', group_by=group_by),
3033
            pd.Series([
3034
                pd.Timestamp('2020-01-01 00:00:00'), pd.Timestamp('2020-01-08 00:00:00'),
3035
                pd.Timedelta('8 days 00:00:00'), pd.Timestamp('2020-01-01 00:00:00'),
3036
                pd.Timestamp('2020-01-08 00:00:00'), pd.Timedelta('5 days 00:00:00'),
3037
                pd.Timedelta('5 days 00:00:00'), 8, 4, 4, 6, 2, -0.32, 50.0, 3, 3, 262.0,
3038
                -271.3333333333333, 149.25, -156.30555555555557, pd.Timedelta('2 days 00:00:00'),
3039
                pd.Timedelta('2 days 00:00:00'), 0.895734597156398, -0.058666666666666756, -0.10439051512510047
3040
            ],
3041
                index=stats_index,
3042
                name='g1'
3043
            )
3044
        )
3045
        pd.testing.assert_index_equal(
3046
            exit_trades.stats(tags='trades').index,
3047
            pd.Index([
3048
                'First Trade Start', 'Last Trade End', 'Total Long Trades',
3049
                'Total Short Trades', 'Total Closed Trades', 'Total Open Trades',
3050
                'Open Trade PnL', 'Win Rate [%]', 'Max Win Streak', 'Max Loss Streak',
3051
                'Best Trade [%]', 'Worst Trade [%]', 'Avg Winning Trade [%]',
3052
                'Avg Losing Trade [%]', 'Avg Winning Trade Duration',
3053
                'Avg Losing Trade Duration', 'Profit Factor', 'Expectancy', 'SQN'
3054
            ], dtype='object')
3055
        )
3056
        pd.testing.assert_series_equal(
3057
            exit_trades['c'].stats(),
3058
            exit_trades.stats(column='c')
3059
        )
3060
        pd.testing.assert_series_equal(
3061
            exit_trades['c'].stats(),
3062
            exit_trades.stats(column='c', group_by=False)
3063
        )
3064
        pd.testing.assert_series_equal(
3065
            exit_trades_grouped['g2'].stats(),
3066
            exit_trades_grouped.stats(column='g2')
3067
        )
3068
        pd.testing.assert_series_equal(
3069
            exit_trades_grouped['g2'].stats(),
3070
            exit_trades.stats(column='g2', group_by=group_by)
3071
        )
3072
        stats_df = exit_trades.stats(agg_func=None)
3073
        assert stats_df.shape == (4, 25)
3074
        pd.testing.assert_index_equal(stats_df.index, exit_trades.wrapper.columns)
3075
        pd.testing.assert_index_equal(stats_df.columns, stats_index)
3076

3077

3078
entry_trades = vbt.EntryTrades.from_orders(orders)
3079
entry_trades_grouped = vbt.EntryTrades.from_orders(orders_grouped)
3080

3081

3082
class TestEntryTrades:
3083
    def test_records_arr(self):
3084
        record_arrays_close(
3085
            entry_trades.values,
3086
            np.array([
3087
                (0, 0, 1.0, 0, 1.0, 0.01, 3, 3.0909090909090904, 0.03090909090909091, 2.05, 2.05, 0, 1, 0),
3088
                (1, 0, 0.1, 1, 2.0, 0.002, 3, 3.0909090909090904, 0.003090909090909091,
3089
                 0.10399999999999998, 0.5199999999999999, 0, 1, 0),
3090
                (2, 0, 1.0, 5, 6.0, 0.06, 6, 7.0, 0.07, 0.8699999999999999, 0.145, 0, 1, 1),
3091
                (3, 0, 2.0, 7, 8.0, 0.16, 7, 8.0, 0.0, -0.16, -0.01, 0, 0, 2),
3092
                (4, 1, 1.0, 0, 1.0, 0.01, 3, 3.0909090909090904, 0.03090909090909091,
3093
                 -2.131818181818181, -2.131818181818181, 1, 1, 3),
3094
                (5, 1, 0.1, 1, 2.0, 0.002, 3, 3.0909090909090904, 0.003090909090909091,
3095
                 -0.11418181818181816, -0.5709090909090908, 1, 1, 3),
3096
                (6, 1, 1.0, 5, 6.0, 0.06, 6, 7.0, 0.07, -1.1300000000000001, -0.18833333333333335, 1, 1, 4),
3097
                (7, 1, 2.0, 7, 8.0, 0.16, 7, 8.0, 0.0, -0.16, -0.01, 1, 0, 5),
3098
                (8, 2, 1.0, 0, 1.0, 0.01, 3, 3.0909090909090904, 0.03090909090909091, 2.05, 2.05, 0, 1, 6),
3099
                (9, 2, 0.1, 1, 2.0, 0.002, 3, 3.0909090909090904, 0.003090909090909091,
3100
                 0.10399999999999998, 0.5199999999999999, 0, 1, 6),
3101
                (10, 2, 1.0, 5, 6.0, 0.06, 6, 7.0, 0.07, 0.8699999999999999, 0.145, 0, 1, 7),
3102
                (11, 2, 1.0, 6, 7.0, 0.07, 7, 8.0, 0.08, -1.1500000000000001, -0.1642857142857143, 1, 1, 8),
3103
                (12, 2, 1.0, 7, 8.0, 0.08, 7, 8.0, 0.0, -0.08, -0.01, 0, 0, 9)
3104
            ], dtype=trade_dt)
3105
        )
3106
        reversed_col_orders = orders.replace(records_arr=np.concatenate((
3107
            orders.values[orders.values['col'] == 2],
3108
            orders.values[orders.values['col'] == 1],
3109
            orders.values[orders.values['col'] == 0]
3110
        )))
3111
        record_arrays_close(
3112
            vbt.EntryTrades.from_orders(reversed_col_orders).values,
3113
            entry_trades.values
3114
        )
3115

3116
    def test_records_readable(self):
3117
        records_readable = entry_trades.records_readable
3118

3119
        np.testing.assert_array_equal(
3120
            records_readable['Entry Trade Id'].values,
3121
            np.array([
3122
                0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12
3123
            ])
3124
        )
3125
        np.testing.assert_array_equal(
3126
            records_readable['Position Id'].values,
3127
            np.array([
3128
                0, 0, 1, 2, 3, 3, 4, 5, 6, 6, 7, 8, 9
3129
            ])
3130
        )
3131

3132

3133
positions = vbt.Positions.from_trades(exit_trades)
3134
positions_grouped = vbt.Positions.from_trades(exit_trades_grouped)
3135

3136

3137
class TestPositions:
3138
    def test_records_arr(self):
3139
        record_arrays_close(
3140
            positions.values,
3141
            np.array([
3142
                (0, 0, 1.1, 0, 1.09090909, 0.012, 3, 3.09090909, 0.034, 2.154, 1.795, 0, 1, 0),
3143
                (1, 0, 1., 5, 6., 0.06, 6, 7., 0.07, 0.87, 0.145, 0, 1, 1),
3144
                (2, 0, 2., 7, 8., 0.16, 7, 8., 0., -0.16, -0.01, 0, 0, 2),
3145
                (3, 1, 1.1, 0, 1.09090909, 0.012, 3, 3.09090909, 0.034, -2.246, -1.87166667, 1, 1, 3),
3146
                (4, 1, 1., 5, 6., 0.06, 6, 7., 0.07, -1.13, -0.18833333, 1, 1, 4),
3147
                (5, 1, 2., 7, 8., 0.16, 7, 8., 0., -0.16, -0.01, 1, 0, 5),
3148
                (6, 2, 1.1, 0, 1.09090909, 0.012, 3, 3.09090909, 0.034, 2.154, 1.795, 0, 1, 6),
3149
                (7, 2, 1., 5, 6., 0.06, 6, 7., 0.07, 0.87, 0.145, 0, 1, 7),
3150
                (8, 2, 1., 6, 7., 0.07, 7, 8., 0.08, -1.15, -0.16428571, 1, 1, 8),
3151
                (9, 2, 1., 7, 8., 0.08, 7, 8., 0., -0.08, -0.01, 0, 0, 9)
3152
            ], dtype=trade_dt)
3153
        )
3154
        reversed_col_trades = exit_trades.replace(records_arr=np.concatenate((
3155
            exit_trades.values[exit_trades.values['col'] == 2],
3156
            exit_trades.values[exit_trades.values['col'] == 1],
3157
            exit_trades.values[exit_trades.values['col'] == 0]
3158
        )))
3159
        record_arrays_close(
3160
            vbt.Positions.from_trades(reversed_col_trades).values,
3161
            positions.values
3162
        )
3163

3164
    def test_records_readable(self):
3165
        records_readable = positions.records_readable
3166

3167
        np.testing.assert_array_equal(
3168
            records_readable['Position Id'].values,
3169
            np.array([
3170
                0, 1, 2, 3, 4, 5, 6, 7, 8, 9
3171
            ])
3172
        )
3173
        assert 'Parent Id' not in records_readable.columns
3174

3175

3176
# ############# logs.py ############# #
3177

3178
logs = vbt.Portfolio.from_orders(close, size, fees=0.01, log=True, freq='1 days').logs
3179
logs_grouped = logs.regroup(group_by)
3180

3181

3182
class TestLogs:
3183
    def test_mapped_fields(self):
3184
        for name in log_dt.names:
3185
            np.testing.assert_array_equal(
3186
                getattr(logs, name).values,
3187
                logs.values[name]
3188
            )
3189

3190
    def test_records_readable(self):
3191
        records_readable = logs.records_readable
3192

3193
        np.testing.assert_array_equal(
3194
            records_readable['Log Id'].values,
3195
            np.array([
3196
                0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27,
3197
                28, 29, 30, 31
3198
            ])
3199
        )
3200
        np.testing.assert_array_equal(
3201
            records_readable['Timestamp'].values,
3202
            np.array([
3203
                '2020-01-01T00:00:00.000000000', '2020-01-02T00:00:00.000000000',
3204
                '2020-01-03T00:00:00.000000000', '2020-01-04T00:00:00.000000000',
3205
                '2020-01-05T00:00:00.000000000', '2020-01-06T00:00:00.000000000',
3206
                '2020-01-07T00:00:00.000000000', '2020-01-08T00:00:00.000000000',
3207
                '2020-01-01T00:00:00.000000000', '2020-01-02T00:00:00.000000000',
3208
                '2020-01-03T00:00:00.000000000', '2020-01-04T00:00:00.000000000',
3209
                '2020-01-05T00:00:00.000000000', '2020-01-06T00:00:00.000000000',
3210
                '2020-01-07T00:00:00.000000000', '2020-01-08T00:00:00.000000000',
3211
                '2020-01-01T00:00:00.000000000', '2020-01-02T00:00:00.000000000',
3212
                '2020-01-03T00:00:00.000000000', '2020-01-04T00:00:00.000000000',
3213
                '2020-01-05T00:00:00.000000000', '2020-01-06T00:00:00.000000000',
3214
                '2020-01-07T00:00:00.000000000', '2020-01-08T00:00:00.000000000',
3215
                '2020-01-01T00:00:00.000000000', '2020-01-02T00:00:00.000000000',
3216
                '2020-01-03T00:00:00.000000000', '2020-01-04T00:00:00.000000000',
3217
                '2020-01-05T00:00:00.000000000', '2020-01-06T00:00:00.000000000',
3218
                '2020-01-07T00:00:00.000000000', '2020-01-08T00:00:00.000000000'
3219
            ], dtype='datetime64[ns]')
3220
        )
3221
        np.testing.assert_array_equal(
3222
            records_readable['Column'].values,
3223
            np.array([
3224
                'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'b', 'b', 'b', 'b', 'b', 'b', 'b', 'b', 'c', 'c', 'c', 'c', 'c',
3225
                'c', 'c', 'c', 'd', 'd', 'd', 'd', 'd', 'd', 'd', 'd'
3226
            ])
3227
        )
3228
        np.testing.assert_array_equal(
3229
            records_readable['Group'].values,
3230
            np.array([
3231
                0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3
3232
            ])
3233
        )
3234
        np.testing.assert_array_equal(
3235
            records_readable['Cash'].values,
3236
            np.array([
3237
                100.0, 98.99, 98.788, 101.758, 102.154, 102.154, 96.094, 103.024, 100.0, 100.99, 101.18799999999999,
3238
                98.15799999999999, 97.75399999999999, 97.75399999999999, 103.69399999999999, 96.624, 100.0, 98.99,
3239
                98.788, 101.758, 102.154, 102.154, 96.094, 109.954, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0,
3240
                100.0
3241
            ])
3242
        )
3243
        np.testing.assert_array_equal(
3244
            records_readable['Position'].values,
3245
            np.array([
3246
                0.0, 1.0, 1.1, 0.10000000000000009, 0.0, 0.0, 1.0, 0.0, 0.0, -1.0, -1.1, -0.10000000000000009, 0.0, 0.0,
3247
                -1.0, 0.0, 0.0, 1.0, 1.1, 0.10000000000000009, 0.0, 0.0, 1.0, -1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
3248
                0.0
3249
            ])
3250
        )
3251
        np.testing.assert_array_equal(
3252
            records_readable['Debt'].values,
3253
            np.array([
3254
                0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.2, 0.10909090909090913, 0.0, 0.0, 6.0, 0.0, 0.0,
3255
                0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 7.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
3256
            ])
3257
        )
3258
        np.testing.assert_array_equal(
3259
            records_readable['Free Cash'].values,
3260
            np.array([
3261
                100.0, 98.99, 98.788, 101.758, 102.154, 102.154, 96.094, 103.024, 100.0, 98.99, 98.788,
3262
                97.93981818181818, 97.754, 97.754, 91.694, 96.624, 100.0, 98.99, 98.788, 101.758, 102.154, 102.154,
3263
                96.094, 95.954, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0
3264
            ])
3265
        )
3266
        np.testing.assert_array_equal(
3267
            records_readable['Val Price'].values,
3268
            np.array([
3269
                1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 1.0, 2.0, 3.0, 4.0, 5.0,
3270
                6.0, 7.0, 8.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0
3271
            ])
3272
        )
3273
        np.testing.assert_array_equal(
3274
            records_readable['Value'].values,
3275
            np.array([
3276
                100.0, 100.99, 102.088, 102.158, 102.154, 102.154, 103.094, 103.024, 100.0, 98.99, 97.88799999999999,
3277
                97.75799999999998, 97.75399999999999, 97.75399999999999, 96.69399999999999, 96.624, 100.0, 100.99,
3278
                102.088, 102.158, 102.154, 102.154, 103.094, 101.954, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0,
3279
                100.0
3280
            ])
3281
        )
3282
        np.testing.assert_array_equal(
3283
            records_readable['Request Size'].values,
3284
            np.array([
3285
                1.0, 0.1, -1.0, -0.1, np.nan, 1.0, -1.0, 2.0, -1.0, -0.1, 1.0, 0.1, np.nan, -1.0, 1.0, -2.0, 1.0, 0.1,
3286
                -1.0, -0.1, np.nan, 1.0, -2.0, 2.0, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan
3287
            ])
3288
        )
3289
        np.testing.assert_array_equal(
3290
            records_readable['Request Price'].values,
3291
            np.array([
3292
                1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 1.0, 2.0, 3.0, 4.0, 5.0,
3293
                6.0, 7.0, 8.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0
3294
            ])
3295
        )
3296
        np.testing.assert_array_equal(
3297
            records_readable['Request Size Type'].values,
3298
            np.array([
3299
                'Amount', 'Amount', 'Amount', 'Amount', 'Amount', 'Amount', 'Amount', 'Amount', 'Amount', 'Amount',
3300
                'Amount', 'Amount', 'Amount', 'Amount', 'Amount', 'Amount', 'Amount', 'Amount', 'Amount', 'Amount',
3301
                'Amount', 'Amount', 'Amount', 'Amount', 'Amount', 'Amount', 'Amount', 'Amount', 'Amount', 'Amount',
3302
                'Amount', 'Amount'
3303
            ])
3304
        )
3305
        np.testing.assert_array_equal(
3306
            records_readable['Request Direction'].values,
3307
            np.array([
3308
                'Both', 'Both', 'Both', 'Both', 'Both', 'Both', 'Both', 'Both', 'Both', 'Both', 'Both', 'Both',
3309
                'Both', 'Both', 'Both', 'Both', 'Both', 'Both', 'Both', 'Both', 'Both', 'Both', 'Both', 'Both',
3310
                'Both', 'Both', 'Both', 'Both', 'Both', 'Both', 'Both', 'Both'
3311
            ])
3312
        )
3313
        np.testing.assert_array_equal(
3314
            records_readable['Request Fees'].values,
3315
            np.array([
3316
                0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01,
3317
                0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01
3318
            ])
3319
        )
3320
        np.testing.assert_array_equal(
3321
            records_readable['Request Fixed Fees'].values,
3322
            np.array([
3323
                0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
3324
                0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
3325
            ])
3326
        )
3327
        np.testing.assert_array_equal(
3328
            records_readable['Request Slippage'].values,
3329
            np.array([
3330
                0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
3331
                0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
3332
            ])
3333
        )
3334
        np.testing.assert_array_equal(
3335
            records_readable['Request Min Size'].values,
3336
            np.array([
3337
                1e-08, 1e-08, 1e-08, 1e-08, 1e-08, 1e-08, 1e-08, 1e-08, 1e-08, 1e-08, 1e-08, 1e-08, 1e-08, 1e-08, 1e-08,
3338
                1e-08, 1e-08, 1e-08, 1e-08, 1e-08, 1e-08, 1e-08, 1e-08, 1e-08, 1e-08, 1e-08, 1e-08, 1e-08, 1e-08, 1e-08,
3339
                1e-08, 1e-08
3340
            ])
3341
        )
3342
        np.testing.assert_array_equal(
3343
            records_readable['Request Max Size'].values,
3344
            np.array([
3345
                np.inf, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf,
3346
                np.inf, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf,
3347
                np.inf, np.inf, np.inf, np.inf, np.inf, np.inf
3348
            ])
3349
        )
3350
        np.testing.assert_array_equal(
3351
            records_readable['Request Size Granularity'].values,
3352
            np.array([
3353
                np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan,
3354
                np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan,
3355
                np.nan, np.nan, np.nan, np.nan, np.nan, np.nan
3356
            ])
3357
        )
3358
        np.testing.assert_array_equal(
3359
            records_readable['Request Rejection Prob'].values,
3360
            np.array([
3361
                0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
3362
                0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
3363
            ])
3364
        )
3365
        np.testing.assert_array_equal(
3366
            records_readable['Request Lock Cash'].values,
3367
            np.array([
3368
                False, False, False, False, False, False, False, False, False, False, False, False, False, False, False,
3369
                False, False, False, False, False, False, False, False, False, False, False, False, False, False, False,
3370
                False, False
3371
            ])
3372
        )
3373
        np.testing.assert_array_equal(
3374
            records_readable['Request Allow Partial'].values,
3375
            np.array([
3376
                True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True,
3377
                True, True, True, True, True, True, True, True, True, True, True, True, True, True, True
3378
            ])
3379
        )
3380
        np.testing.assert_array_equal(
3381
            records_readable['Request Raise Rejection'].values,
3382
            np.array([
3383
                False, False, False, False, False, False, False, False, False, False, False, False, False, False, False,
3384
                False, False, False, False, False, False, False, False, False, False, False, False, False, False, False,
3385
                False, False
3386
            ])
3387
        )
3388
        np.testing.assert_array_equal(
3389
            records_readable['Request Log'].values,
3390
            np.array([
3391
                True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True,
3392
                True, True, True, True, True, True, True, True, True, True, True, True, True, True, True
3393
            ])
3394
        )
3395
        np.testing.assert_array_equal(
3396
            records_readable['New Cash'].values,
3397
            np.array([
3398
                98.99, 98.788, 101.758, 102.154, 102.154, 96.094, 103.024, 86.864, 100.99, 101.18799999999999,
3399
                98.15799999999999, 97.75399999999999, 97.75399999999999, 103.69399999999999, 96.624, 112.464, 98.99,
3400
                98.788, 101.758, 102.154, 102.154, 96.094, 109.954, 93.794, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0,
3401
                100.0, 100.0
3402
            ])
3403
        )
3404
        np.testing.assert_array_equal(
3405
            records_readable['New Position'].values,
3406
            np.array([
3407
                1.0, 1.1, 0.10000000000000009, 0.0, 0.0, 1.0, 0.0, 2.0, -1.0, -1.1, -0.10000000000000009, 0.0, 0.0,
3408
                -1.0, 0.0, -2.0, 1.0, 1.1, 0.10000000000000009, 0.0, 0.0, 1.0, -1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
3409
                0.0, 0.0
3410
            ])
3411
        )
3412
        np.testing.assert_array_equal(
3413
            records_readable['New Debt'].values,
3414
            np.array([
3415
                0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.2, 0.10909090909090913, 0.0, 0.0, 6.0, 0.0, 16.0, 0.0,
3416
                0.0, 0.0, 0.0, 0.0, 0.0, 7.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
3417
            ])
3418
        )
3419
        np.testing.assert_array_equal(
3420
            records_readable['New Free Cash'].values,
3421
            np.array([
3422
                98.99, 98.788, 101.758, 102.154, 102.154, 96.094, 103.024, 86.864, 98.99, 98.788, 97.93981818181818,
3423
                97.754, 97.754, 91.694, 96.624, 80.464, 98.99, 98.788, 101.758, 102.154, 102.154, 96.094, 95.954,
3424
                93.794, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0
3425
            ])
3426
        )
3427
        np.testing.assert_array_equal(
3428
            records_readable['New Val Price'].values,
3429
            np.array([
3430
                1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 1.0, 2.0, 3.0, 4.0, 5.0,
3431
                6.0, 7.0, 8.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0
3432
            ])
3433
        )
3434
        np.testing.assert_array_equal(
3435
            records_readable['New Value'].values,
3436
            np.array([
3437
                100.0, 100.99, 102.088, 102.158, 102.154, 102.154, 103.094, 103.024, 100.0, 98.99, 97.88799999999999,
3438
                97.75799999999998, 97.75399999999999, 97.75399999999999, 96.69399999999999, 96.624, 100.0, 100.99,
3439
                102.088, 102.158, 102.154, 102.154, 103.094, 101.954, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0,
3440
                100.0
3441
            ])
3442
        )
3443
        np.testing.assert_array_equal(
3444
            records_readable['Result Size'].values,
3445
            np.array([
3446
                1.0, 0.1, 1.0, 0.1, np.nan, 1.0, 1.0, 2.0, 1.0, 0.1, 1.0, 0.1, np.nan, 1.0, 1.0, 2.0, 1.0, 0.1, 1.0,
3447
                0.1, np.nan, 1.0, 2.0, 2.0, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan
3448
            ])
3449
        )
3450
        np.testing.assert_array_equal(
3451
            records_readable['Result Price'].values,
3452
            np.array([
3453
                1.0, 2.0, 3.0, 4.0, np.nan, 6.0, 7.0, 8.0, 1.0, 2.0, 3.0, 4.0, np.nan, 6.0, 7.0, 8.0, 1.0, 2.0, 3.0,
3454
                4.0, np.nan, 6.0, 7.0, 8.0, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan
3455
            ])
3456
        )
3457
        np.testing.assert_array_equal(
3458
            records_readable['Result Fees'].values,
3459
            np.array([
3460
                0.01, 0.002, 0.03, 0.004, np.nan, 0.06, 0.07, 0.16, 0.01, 0.002, 0.03, 0.004, np.nan, 0.06, 0.07, 0.16,
3461
                0.01, 0.002, 0.03, 0.004, np.nan, 0.06, 0.14, 0.16, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan,
3462
                np.nan, np.nan
3463
            ])
3464
        )
3465
        np.testing.assert_array_equal(
3466
            records_readable['Result Side'].values,
3467
            np.array([
3468
                'Buy', 'Buy', 'Sell', 'Sell', None, 'Buy', 'Sell', 'Buy', 'Sell', 'Sell', 'Buy', 'Buy', None, 'Sell',
3469
                'Buy', 'Sell', 'Buy', 'Buy', 'Sell', 'Sell', None, 'Buy', 'Sell', 'Buy', None, None, None, None, None,
3470
                None, None, None
3471
            ])
3472
        )
3473
        np.testing.assert_array_equal(
3474
            records_readable['Result Status'].values,
3475
            np.array([
3476
                'Filled', 'Filled', 'Filled', 'Filled', 'Ignored', 'Filled', 'Filled', 'Filled', 'Filled', 'Filled',
3477
                'Filled', 'Filled', 'Ignored', 'Filled', 'Filled', 'Filled', 'Filled', 'Filled', 'Filled', 'Filled',
3478
                'Ignored', 'Filled', 'Filled', 'Filled', 'Ignored', 'Ignored', 'Ignored', 'Ignored', 'Ignored',
3479
                'Ignored', 'Ignored', 'Ignored'
3480
            ])
3481
        )
3482
        np.testing.assert_array_equal(
3483
            records_readable['Result Status Info'].values,
3484
            np.array([
3485
                None, None, None, None, 'SizeNaN', None, None, None, None, None, None, None, 'SizeNaN', None, None,
3486
                None, None, None, None, None, 'SizeNaN', None, None, None, 'SizeNaN', 'SizeNaN', 'SizeNaN', 'SizeNaN',
3487
                'SizeNaN', 'SizeNaN', 'SizeNaN', 'SizeNaN'
3488
            ])
3489
        )
3490
        np.testing.assert_array_equal(
3491
            records_readable['Order Id'].values,
3492
            np.array([
3493
                0, 1, 2, 3, -1, 4, 5, 6, 7, 8, 9, 10, -1, 11, 12, 13, 14, 15, 16, 17, -1, 18, 19, 20, -1, -1, -1, -1,
3494
                -1, -1, -1, -1
3495
            ])
3496
        )
3497

3498
    def test_stats(self):
3499
        stats_index = pd.Index([
3500
            'Start', 'End', 'Period', 'Total Records', 'Status Counts: None',
3501
            'Status Counts: Filled', 'Status Counts: Ignored',
3502
            'Status Counts: Rejected', 'Status Info Counts: None',
3503
            'Status Info Counts: SizeNaN'
3504
        ], dtype='object')
3505
        pd.testing.assert_series_equal(
3506
            logs.stats(),
3507
            pd.Series([
3508
                pd.Timestamp('2020-01-01 00:00:00'), pd.Timestamp('2020-01-08 00:00:00'),
3509
                pd.Timedelta('8 days 00:00:00'), 8.0, 0.0, 5.25, 2.75, 0.0, 5.25, 2.75
3510
            ],
3511
                index=stats_index,
3512
                name='agg_func_mean'
3513
            )
3514
        )
3515
        pd.testing.assert_series_equal(
3516
            logs.stats(column='a'),
3517
            pd.Series([
3518
                pd.Timestamp('2020-01-01 00:00:00'), pd.Timestamp('2020-01-08 00:00:00'),
3519
                pd.Timedelta('8 days 00:00:00'), 8, 0, 7, 1, 0, 7, 1
3520
            ],
3521
                index=stats_index,
3522
                name='a'
3523
            )
3524
        )
3525
        pd.testing.assert_series_equal(
3526
            logs.stats(column='g1', group_by=group_by),
3527
            pd.Series([
3528
                pd.Timestamp('2020-01-01 00:00:00'), pd.Timestamp('2020-01-08 00:00:00'),
3529
                pd.Timedelta('8 days 00:00:00'), 16, 0, 14, 2, 0, 14, 2
3530
            ],
3531
                index=stats_index,
3532
                name='g1'
3533
            )
3534
        )
3535
        pd.testing.assert_series_equal(
3536
            logs['c'].stats(),
3537
            logs.stats(column='c')
3538
        )
3539
        pd.testing.assert_series_equal(
3540
            logs['c'].stats(),
3541
            logs.stats(column='c', group_by=False)
3542
        )
3543
        pd.testing.assert_series_equal(
3544
            logs_grouped['g2'].stats(),
3545
            logs_grouped.stats(column='g2')
3546
        )
3547
        pd.testing.assert_series_equal(
3548
            logs_grouped['g2'].stats(),
3549
            logs.stats(column='g2', group_by=group_by)
3550
        )
3551
        stats_df = logs.stats(agg_func=None)
3552
        assert stats_df.shape == (4, 10)
3553
        pd.testing.assert_index_equal(stats_df.index, logs.wrapper.columns)
3554
        pd.testing.assert_index_equal(stats_df.columns, stats_index)
3555

3556
    def test_count(self):
3557
        assert logs['a'].count() == 8
3558
        pd.testing.assert_series_equal(
3559
            logs.count(),
3560
            pd.Series(
3561
                np.array([8, 8, 8, 8]),
3562
                index=pd.Index(['a', 'b', 'c', 'd'], dtype='object')
3563
            ).rename('count')
3564
        )
3565
        pd.testing.assert_series_equal(
3566
            logs_grouped.count(),
3567
            pd.Series(
3568
                np.array([16, 16]),
3569
                index=pd.Index(['g1', 'g2'], dtype='object')
3570
            ).rename('count')
3571
        )
3572

3573
Product

Resources

Company
