CoCalc -- test_labels.py

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

3
import numpy as np
4
import pandas as pd
5

6
import vectorbt as vbt
7

8
close_ts = pd.DataFrame({
9
    'a': [1, 2, 1, 2, 3, 2],
10
    'b': [3, 2, 3, 2, 1, 2]
11
}, index=pd.Index([
12
    datetime(2020, 1, 1),
13
    datetime(2020, 1, 2),
14
    datetime(2020, 1, 3),
15
    datetime(2020, 1, 4),
16
    datetime(2020, 1, 5),
17
    datetime(2020, 1, 6)
18
]))
19

20
pos_ths = [np.array([1, 1 / 2]), np.array([2, 1 / 2]), np.array([3, 1 / 2])]
21
neg_ths = [np.array([1 / 2, 1 / 3]), np.array([1 / 2, 2 / 3]), np.array([1 / 2, 3 / 4])]
22

23

24
# ############# Global ############# #
25

26
def setup_module():
27
    vbt.settings.numba['check_func_suffix'] = True
28
    vbt.settings.caching.enabled = False
29
    vbt.settings.caching.whitelist = []
30
    vbt.settings.caching.blacklist = []
31

32

33
def teardown_module():
34
    vbt.settings.reset()
35

36

37
# ############# generators.py ############# #
38

39
class TestGenerators:
40
    def test_FMEAN(self):
41
        pd.testing.assert_frame_equal(
42
            vbt.FMEAN.run(close_ts, window=(2, 3), ewm=False).fmean,
43
            pd.DataFrame(
44
                np.array([
45
                    [1.5, 2.5, 1.6666666666666667, 2.3333333333333335],
46
                    [1.5, 2.5, 2.0, 2.0],
47
                    [2.5, 1.5, 2.3333333333333335, 1.6666666666666667],
48
                    [2.5, 1.5, np.nan, np.nan],
49
                    [np.nan, np.nan, np.nan, np.nan],
50
                    [np.nan, np.nan, np.nan, np.nan]
51
                ]),
52
                index=close_ts.index,
53
                columns=pd.MultiIndex.from_tuples([
54
                    (2, False, 'a'),
55
                    (2, False, 'b'),
56
                    (3, False, 'a'),
57
                    (3, False, 'b'),
58
                ], names=['fmean_window', 'fmean_ewm', None])
59
            )
60
        )
61
        pd.testing.assert_frame_equal(
62
            vbt.FMEAN.run(close_ts, window=(2, 3), ewm=True).fmean,
63
            pd.DataFrame(
64
                np.array([
65
                    [1.8024691358024691, 2.197530864197531, 1.8125, 2.1875],
66
                    [1.4074074074074074, 2.5925925925925926, 1.625, 2.375],
67
                    [2.2222222222222223, 1.7777777777777777, 2.25, 1.75],
68
                    [2.666666666666667, 1.3333333333333335, np.nan, np.nan],
69
                    [np.nan, np.nan, np.nan, np.nan],
70
                    [np.nan, np.nan, np.nan, np.nan]
71
                ]),
72
                index=close_ts.index,
73
                columns=pd.MultiIndex.from_tuples([
74
                    (2, True, 'a'),
75
                    (2, True, 'b'),
76
                    (3, True, 'a'),
77
                    (3, True, 'b'),
78
                ], names=['fmean_window', 'fmean_ewm', None])
79
            )
80
        )
81

82
    def test_FSTD(self):
83
        pd.testing.assert_frame_equal(
84
            vbt.FSTD.run(close_ts, window=(2, 3), ewm=False).fstd,
85
            pd.DataFrame(
86
                np.array([
87
                    [0.5, 0.5, 0.4714045207910384, 0.4714045207910183],
88
                    [0.5, 0.5, 0.816496580927726, 0.816496580927726],
89
                    [0.5, 0.5, 0.4714045207910183, 0.4714045207910384],
90
                    [0.5, 0.5, np.nan, np.nan],
91
                    [np.nan, np.nan, np.nan, np.nan],
92
                    [np.nan, np.nan, np.nan, np.nan]
93
                ]),
94
                index=close_ts.index,
95
                columns=pd.MultiIndex.from_tuples([
96
                    (2, False, 'a'),
97
                    (2, False, 'b'),
98
                    (3, False, 'a'),
99
                    (3, False, 'b'),
100
                ], names=['fstd_window', 'fstd_ewm', None])
101
            )
102
        )
103
        pd.testing.assert_frame_equal(
104
            vbt.FSTD.run(close_ts, window=(2, 3), ewm=True).fstd,
105
            pd.DataFrame(
106
                np.array([
107
                    [0.64486716348143, 0.6448671634814303, 0.6462561866810479, 0.6462561866810479],
108
                    [0.8833005039168617, 0.8833005039168604, 0.8591246929842246, 0.8591246929842246],
109
                    [0.5916079783099623, 0.5916079783099623, 0.5477225575051662, 0.5477225575051662],
110
                    [0.7071067811865476, 0.7071067811865476, np.nan, np.nan],
111
                    [np.nan, np.nan, np.nan, np.nan],
112
                    [np.nan, np.nan, np.nan, np.nan]
113
                ]),
114
                index=close_ts.index,
115
                columns=pd.MultiIndex.from_tuples([
116
                    (2, True, 'a'),
117
                    (2, True, 'b'),
118
                    (3, True, 'a'),
119
                    (3, True, 'b'),
120
                ], names=['fstd_window', 'fstd_ewm', None])
121
            )
122
        )
123

124
    def test_FMIN(self):
125
        pd.testing.assert_frame_equal(
126
            vbt.FMIN.run(close_ts, window=(2, 3)).fmin,
127
            pd.DataFrame(
128
                np.array([
129
                    [1.0, 2.0, 1.0, 2.0],
130
                    [1.0, 2.0, 1.0, 1.0],
131
                    [2.0, 1.0, 2.0, 1.0],
132
                    [2.0, 1.0, np.nan, np.nan],
133
                    [np.nan, np.nan, np.nan, np.nan],
134
                    [np.nan, np.nan, np.nan, np.nan]
135
                ]),
136
                index=close_ts.index,
137
                columns=pd.MultiIndex.from_tuples([
138
                    (2, 'a'),
139
                    (2, 'b'),
140
                    (3, 'a'),
141
                    (3, 'b'),
142
                ], names=['fmin_window', None])
143
            )
144
        )
145

146
    def test_FMAX(self):
147
        pd.testing.assert_frame_equal(
148
            vbt.FMAX.run(close_ts, window=(2, 3)).fmax,
149
            pd.DataFrame(
150
                np.array([
151
                    [2.0, 3.0, 2.0, 3.0],
152
                    [2.0, 3.0, 3.0, 3.0],
153
                    [3.0, 2.0, 3.0, 2.0],
154
                    [3.0, 2.0, np.nan, np.nan],
155
                    [np.nan, np.nan, np.nan, np.nan],
156
                    [np.nan, np.nan, np.nan, np.nan]
157
                ]),
158
                index=close_ts.index,
159
                columns=pd.MultiIndex.from_tuples([
160
                    (2, 'a'),
161
                    (2, 'b'),
162
                    (3, 'a'),
163
                    (3, 'b'),
164
                ], names=['fmax_window', None])
165
            )
166
        )
167

168
    def test_FIXLB(self):
169
        pd.testing.assert_frame_equal(
170
            vbt.FIXLB.run(close_ts, n=(2, 3)).labels,
171
            pd.DataFrame(
172
                np.array([
173
                    [0.0, 0.0, 1.0, -0.3333333333333333],
174
                    [0.0, 0.0, 0.5, -0.5],
175
                    [2.0, -0.6666666666666666, 1.0, -0.3333333333333333],
176
                    [0.0, 0.0, np.nan, np.nan],
177
                    [np.nan, np.nan, np.nan, np.nan],
178
                    [np.nan, np.nan, np.nan, np.nan]
179
                ]),
180
                index=close_ts.index,
181
                columns=pd.MultiIndex.from_tuples([
182
                    (2, 'a'),
183
                    (2, 'b'),
184
                    (3, 'a'),
185
                    (3, 'b'),
186
                ], names=['fixlb_n', None])
187
            )
188
        )
189

190
    def test_MEANLB(self):
191
        pd.testing.assert_frame_equal(
192
            vbt.MEANLB.run(close_ts, window=(2, 3), ewm=False).labels,
193
            pd.DataFrame(
194
                np.array([
195
                    [0.5, -0.16666666666666666, 0.6666666666666667, -0.22222222222222218],
196
                    [-0.25, 0.25, 0.0, 0.0],
197
                    [1.5, -0.5, 1.3333333333333335, -0.4444444444444444],
198
                    [0.25, -0.25, np.nan, np.nan],
199
                    [np.nan, np.nan, np.nan, np.nan],
200
                    [np.nan, np.nan, np.nan, np.nan]
201
                ]),
202
                index=close_ts.index,
203
                columns=pd.MultiIndex.from_tuples([
204
                    (2, False, 'a'),
205
                    (2, False, 'b'),
206
                    (3, False, 'a'),
207
                    (3, False, 'b'),
208
                ], names=['meanlb_window', 'meanlb_ewm', None])
209
            )
210
        )
211
        pd.testing.assert_frame_equal(
212
            vbt.MEANLB.run(close_ts, window=(2, 3), ewm=True).labels,
213
            pd.DataFrame(
214
                np.array([
215
                    [0.8024691358024691, -0.2674897119341564, 0.8125, -0.2708333333333333],
216
                    [-0.2962962962962963, 0.2962962962962963, -0.1875, 0.1875],
217
                    [1.2222222222222223, -0.40740740740740744, 1.25, -0.4166666666666667],
218
                    [0.3333333333333335, -0.33333333333333326, np.nan, np.nan],
219
                    [np.nan, np.nan, np.nan, np.nan],
220
                    [np.nan, np.nan, np.nan, np.nan]
221
                ]),
222
                index=close_ts.index,
223
                columns=pd.MultiIndex.from_tuples([
224
                    (2, True, 'a'),
225
                    (2, True, 'b'),
226
                    (3, True, 'a'),
227
                    (3, True, 'b'),
228
                ], names=['meanlb_window', 'meanlb_ewm', None])
229
            )
230
        )
231

232
    def test_LEXLB(self):
233
        pd.testing.assert_frame_equal(
234
            vbt.LEXLB.run(close_ts, pos_th=pos_ths, neg_th=neg_ths).labels,
235
            pd.DataFrame(
236
                np.array([
237
                    [-1, 1, -1, 1, 0, 0],
238
                    [1, -1, 0, 0, 0, 0],
239
                    [-1, 1, 0, 0, 0, 0],
240
                    [0, 0, 0, 0, 0, 0],
241
                    [1, -1, 1, -1, 0, 0],
242
                    [0, 1, 0, 1, 0, 0]
243
                ]),
244
                index=close_ts.index,
245
                columns=pd.MultiIndex.from_tuples([
246
                    ('array_0', 'array_0', 'a'),
247
                    ('array_0', 'array_0', 'b'),
248
                    ('array_1', 'array_1', 'a'),
249
                    ('array_1', 'array_1', 'b'),
250
                    ('array_2', 'array_2', 'a'),
251
                    ('array_2', 'array_2', 'b')
252
                ], names=['lexlb_pos_th', 'lexlb_neg_th', None])
253
            )
254
        )
255

256
    def test_TRENDLB(self):
257
        pd.testing.assert_frame_equal(
258
            vbt.TRENDLB.run(close_ts, pos_th=pos_ths, neg_th=neg_ths, mode='Binary').labels,
259
            pd.DataFrame(
260
                np.array([
261
                    [1.0, 0.0, 1.0, 0.0, np.nan, np.nan],
262
                    [0.0, 1.0, 1.0, 0.0, np.nan, np.nan],
263
                    [1.0, 0.0, 1.0, 0.0, np.nan, np.nan],
264
                    [1.0, 0.0, 1.0, 0.0, np.nan, np.nan],
265
                    [np.nan, 1.0, np.nan, 1.0, np.nan, np.nan],
266
                    [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan]
267
                ]),
268
                index=close_ts.index,
269
                columns=pd.MultiIndex.from_tuples([
270
                    ('array_0', 'array_0', 0, 'a'),
271
                    ('array_0', 'array_0', 0, 'b'),
272
                    ('array_1', 'array_1', 0, 'a'),
273
                    ('array_1', 'array_1', 0, 'b'),
274
                    ('array_2', 'array_2', 0, 'a'),
275
                    ('array_2', 'array_2', 0, 'b')
276
                ], names=['trendlb_pos_th', 'trendlb_neg_th', 'trendlb_mode', None])
277
            )
278
        )
279
        pd.testing.assert_frame_equal(
280
            vbt.TRENDLB.run(close_ts, pos_th=pos_ths, neg_th=neg_ths, mode='BinaryCont').labels,
281
            pd.DataFrame(
282
                np.array([
283
                    [1.0, 0.0, 1.0, 0.0, np.nan, np.nan],
284
                    [0.0, 1.0, 0.5, 0.5, np.nan, np.nan],
285
                    [1.0, 0.0, 1.0, 0.0, np.nan, np.nan],
286
                    [0.5, 0.5, 0.5, 0.5, np.nan, np.nan],
287
                    [np.nan, 1.0, np.nan, 1.0, np.nan, np.nan],
288
                    [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan]
289
                ]),
290
                index=close_ts.index,
291
                columns=pd.MultiIndex.from_tuples([
292
                    ('array_0', 'array_0', 1, 'a'),
293
                    ('array_0', 'array_0', 1, 'b'),
294
                    ('array_1', 'array_1', 1, 'a'),
295
                    ('array_1', 'array_1', 1, 'b'),
296
                    ('array_2', 'array_2', 1, 'a'),
297
                    ('array_2', 'array_2', 1, 'b')
298
                ], names=['trendlb_pos_th', 'trendlb_neg_th', 'trendlb_mode', None])
299
            )
300
        )
301
        pd.testing.assert_frame_equal(
302
            vbt.TRENDLB.run(close_ts, pos_th=pos_ths, neg_th=neg_ths, mode='BinaryContSat').labels,
303
            pd.DataFrame(
304
                np.array([
305
                    [1.0, 0.0, 1.0, 0.0, np.nan, np.nan],
306
                    [0.0, 1.0, 0.5, 0.4999999999999999, np.nan, np.nan],
307
                    [1.0, 0.0, 1.0, 0.0, np.nan, np.nan],
308
                    [0.6666666666666667, 0.0, 0.5, 0.4999999999999999, np.nan, np.nan],
309
                    [np.nan, 1.0, np.nan, 1.0, np.nan, np.nan],
310
                    [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan]
311
                ]),
312
                index=close_ts.index,
313
                columns=pd.MultiIndex.from_tuples([
314
                    ('array_0', 'array_0', 2, 'a'),
315
                    ('array_0', 'array_0', 2, 'b'),
316
                    ('array_1', 'array_1', 2, 'a'),
317
                    ('array_1', 'array_1', 2, 'b'),
318
                    ('array_2', 'array_2', 2, 'a'),
319
                    ('array_2', 'array_2', 2, 'b')
320
                ], names=['trendlb_pos_th', 'trendlb_neg_th', 'trendlb_mode', None])
321
            )
322
        )
323
        pd.testing.assert_frame_equal(
324
            vbt.TRENDLB.run(close_ts, pos_th=pos_ths, neg_th=neg_ths, mode='PctChange').labels,
325
            pd.DataFrame(
326
                np.array([
327
                    [1.0, -0.3333333333333333, 2.0, -0.6666666666666666, np.nan, np.nan],
328
                    [-0.5, 0.5, 0.5, -0.5, np.nan, np.nan],
329
                    [2.0, -0.6666666666666666, 2.0, -0.6666666666666666, np.nan, np.nan],
330
                    [0.5, -0.5, 0.5, -0.5, np.nan, np.nan],
331
                    [np.nan, 1.0, np.nan, 1.0, np.nan, np.nan],
332
                    [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan]
333
                ]),
334
                index=close_ts.index,
335
                columns=pd.MultiIndex.from_tuples([
336
                    ('array_0', 'array_0', 3, 'a'),
337
                    ('array_0', 'array_0', 3, 'b'),
338
                    ('array_1', 'array_1', 3, 'a'),
339
                    ('array_1', 'array_1', 3, 'b'),
340
                    ('array_2', 'array_2', 3, 'a'),
341
                    ('array_2', 'array_2', 3, 'b')
342
                ], names=['trendlb_pos_th', 'trendlb_neg_th', 'trendlb_mode', None])
343
            )
344
        )
345
        pd.testing.assert_frame_equal(
346
            vbt.TRENDLB.run(close_ts, pos_th=pos_ths, neg_th=neg_ths, mode='PctChangeNorm').labels,
347
            pd.DataFrame(
348
                np.array([
349
                    [0.5, -0.3333333333333333, 0.6666666666666666, -0.6666666666666666, np.nan, np.nan],
350
                    [-0.5, 0.3333333333333333, 0.3333333333333333, -0.5, np.nan, np.nan],
351
                    [0.6666666666666666, -0.6666666666666666, 0.6666666666666666,
352
                     -0.6666666666666666, np.nan, np.nan],
353
                    [0.3333333333333333, -0.5, 0.3333333333333333, -0.5, np.nan, np.nan],
354
                    [np.nan, 0.5, np.nan, 0.5, np.nan, np.nan],
355
                    [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan]
356
                ]),
357
                index=close_ts.index,
358
                columns=pd.MultiIndex.from_tuples([
359
                    ('array_0', 'array_0', 4, 'a'),
360
                    ('array_0', 'array_0', 4, 'b'),
361
                    ('array_1', 'array_1', 4, 'a'),
362
                    ('array_1', 'array_1', 4, 'b'),
363
                    ('array_2', 'array_2', 4, 'a'),
364
                    ('array_2', 'array_2', 4, 'b')
365
                ], names=['trendlb_pos_th', 'trendlb_neg_th', 'trendlb_mode', None])
366
            )
367
        )
368

369
    def test_BOLB(self):
370
        pd.testing.assert_frame_equal(
371
            vbt.BOLB.run(close_ts, window=1, pos_th=pos_ths, neg_th=neg_ths).labels,
372
            pd.DataFrame(
373
                np.array([
374
                    [1.0, -1.0, 0.0, 0.0, 0.0, 0.0],
375
                    [-1.0, 1.0, -1.0, 1.0, -1.0, 1.0],
376
                    [1.0, -1.0, 0.0, 0.0, 0.0, 0.0],
377
                    [0.0, -1.0, 0.0, 0.0, 0.0, 0.0],
378
                    [0.0, 1.0, 0.0, 1.0, 0.0, 1.0],
379
                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
380
                ]),
381
                index=close_ts.index,
382
                columns=pd.MultiIndex.from_tuples([
383
                    (1, 'array_0', 'array_0', 'a'),
384
                    (1, 'array_0', 'array_0', 'b'),
385
                    (1, 'array_1', 'array_1', 'a'),
386
                    (1, 'array_1', 'array_1', 'b'),
387
                    (1, 'array_2', 'array_2', 'a'),
388
                    (1, 'array_2', 'array_2', 'b')
389
                ], names=['bolb_window', 'bolb_pos_th', 'bolb_neg_th', None])
390
            )
391
        )
392
        pd.testing.assert_frame_equal(
393
            vbt.BOLB.run(close_ts, window=2, pos_th=pos_ths, neg_th=neg_ths).labels,
394
            pd.DataFrame(
395
                np.array([
396
                    [1.0, -1.0, 0.0, 0.0, 0.0, 0.0],
397
                    [-1.0, 1.0, -1.0, 1.0, -1.0, 1.0],
398
                    [1.0, -1.0, 1.0, -1.0, 0.0, 0.0],
399
                    [0.0, -1.0, 0.0, 0.0, 0.0, 0.0],
400
                    [0.0, 1.0, 0.0, 1.0, 0.0, 1.0],
401
                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
402
                ]),
403
                index=close_ts.index,
404
                columns=pd.MultiIndex.from_tuples([
405
                    (2, 'array_0', 'array_0', 'a'),
406
                    (2, 'array_0', 'array_0', 'b'),
407
                    (2, 'array_1', 'array_1', 'a'),
408
                    (2, 'array_1', 'array_1', 'b'),
409
                    (2, 'array_2', 'array_2', 'a'),
410
                    (2, 'array_2', 'array_2', 'b')
411
                ], names=['bolb_window', 'bolb_pos_th', 'bolb_neg_th', None])
412
            )
413
        )
414

415
Product

Resources

Company
