-
-
Notifications
You must be signed in to change notification settings - Fork 3.9k
/
ols.py
128 lines (102 loc) · 3.89 KB
/
ols.py
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
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
#!/usr/bin/env python
# -*- coding: utf-8; py-indent-offset:4 -*-
###############################################################################
#
# Copyright (C) 2015-2023 Daniel Rodriguez
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
#
###############################################################################
from __future__ import (absolute_import, division, print_function,
unicode_literals)
import backtrader as bt
from . import PeriodN
__all__ = ['OLS_Slope_InterceptN', 'OLS_TransformationN', 'OLS_BetaN',
'CointN']
class OLS_Slope_InterceptN(PeriodN):
'''
Calculates a linear regression using ``statsmodel.OLS`` (Ordinary least
squares) of data1 on data0
Uses ``pandas`` and ``statsmodels``
'''
_mindatas = 2 # ensure at least 2 data feeds are passed
packages = (
('pandas', 'pd'),
('statsmodels.api', 'sm'),
)
lines = ('slope', 'intercept',)
params = (
('period', 10),
)
def next(self):
p0 = pd.Series(self.data0.get(size=self.p.period))
p1 = pd.Series(self.data1.get(size=self.p.period))
p1 = sm.add_constant(p1)
intercept, slope = sm.OLS(p0, p1).fit().params
self.lines.slope[0] = slope
self.lines.intercept[0] = intercept
class OLS_TransformationN(PeriodN):
'''
Calculates the ``zscore`` for data0 and data1. Although it doesn't directly
uses any external package it relies on ``OLS_SlopeInterceptN`` which uses
``pandas`` and ``statsmodels``
'''
_mindatas = 2 # ensure at least 2 data feeds are passed
lines = ('spread', 'spread_mean', 'spread_std', 'zscore',)
params = (('period', 10),)
def __init__(self):
slint = OLS_Slope_InterceptN(*self.datas)
spread = self.data0 - (slint.slope * self.data1 + slint.intercept)
self.l.spread = spread
self.l.spread_mean = bt.ind.SMA(spread, period=self.p.period)
self.l.spread_std = bt.ind.StdDev(spread, period=self.p.period)
self.l.zscore = (spread - self.l.spread_mean) / self.l.spread_std
class OLS_BetaN(PeriodN):
'''
Calculates a regression of data1 on data0 using ``pandas.ols``
Uses ``pandas``
'''
_mindatas = 2 # ensure at least 2 data feeds are passed
packages = (
('pandas', 'pd'),
)
lines = ('beta',)
params = (('period', 10),)
def next(self):
y, x = (pd.Series(d.get(size=self.p.period)) for d in self.datas)
r_beta = pd.ols(y=y, x=x, window_type='full_sample')
self.lines.beta[0] = r_beta.beta['x']
class CointN(PeriodN):
'''
Calculates the score (coint_t) and pvalue for a given ``period`` for the
data feeds
Uses ``pandas`` and ``statsmodels`` (for ``coint``)
'''
_mindatas = 2 # ensure at least 2 data feeds are passed
packages = (
('pandas', 'pd'), # import pandas as pd
)
frompackages = (
('statsmodels.tsa.stattools', 'coint'), # from st... import coint
)
lines = ('score', 'pvalue',)
params = (
('period', 10),
('trend', 'c'), # see statsmodel.tsa.statttools
)
def next(self):
x, y = (pd.Series(d.get(size=self.p.period)) for d in self.datas)
score, pvalue, _ = coint(x, y, trend=self.p.trend)
self.lines.score[0] = score
self.lines.pvalue[0] = pvalue