Completed
Pull Request — master (#858)
by Eddie
01:35
created

zipline.sources.BenchmarkSource   A

Complexity

Total Complexity 12

Size/Duplication

Total Lines 163
Duplicated Lines 0 %
Metric Value
dl 0
loc 163
rs 10
wmc 12

4 Methods

Rating   Name   Duplication   Size   Complexity  
B _validate_benchmark() 0 28 4
B _initialize_precalculated_series() 0 107 5
A get_value() 0 2 1
A __init__() 0 20 2
1
#
2
# Copyright 2015 Quantopian, Inc.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15
16
from zipline.errors import (
17
    InvalidBenchmarkAsset,
18
    BenchmarkAssetNotAvailableTooEarly,
19
    BenchmarkAssetNotAvailableTooLate
20
)
21
from zipline.data.us_equity_pricing import NoDataOnDate
22
23
24
class BenchmarkSource(object):
25
    def __init__(self, benchmark_sid, env, trading_days, data_portal,
26
                 emission_rate="daily"):
27
        self.benchmark_sid = benchmark_sid
28
        self.env = env
29
        self.trading_days = trading_days
30
        self.emission_rate = emission_rate
31
        self.data_portal = data_portal
32
33
        if self.benchmark_sid:
34
            self.benchmark_asset = self.env.asset_finder.retrieve_asset(
35
                self.benchmark_sid)
36
37
            self._validate_benchmark()
38
39
        self.precalculated_series = \
40
            self._initialize_precalculated_series(
41
                self.benchmark_sid,
42
                self.env,
43
                self.trading_days,
44
                self.data_portal
45
            )
46
47
    def get_value(self, dt):
48
        return self.precalculated_series.loc[dt]
49
50
    def _validate_benchmark(self):
51
        # check if this security has a stock dividend.  if so, raise an
52
        # error suggesting that the user pick a different asset to use
53
        # as benchmark.
54
        stock_dividends = \
55
            self.data_portal.get_stock_dividends(self.benchmark_sid,
56
                                                 self.trading_days)
57
58
        if len(stock_dividends) > 0:
59
            raise InvalidBenchmarkAsset(
60
                sid=str(self.benchmark_sid),
61
                dt=stock_dividends[0]["ex_date"]
62
            )
63
64
        if self.benchmark_asset.start_date > self.trading_days[0]:
65
            # the asset started trading after the first simulation day
66
            raise BenchmarkAssetNotAvailableTooEarly(
67
                sid=str(self.benchmark_sid),
68
                dt=self.trading_days[0],
69
                start_dt=self.benchmark_asset.start_date
70
            )
71
72
        if self.benchmark_asset.end_date < self.trading_days[-1]:
73
            # the asset stopped trading before the last simulation day
74
            raise BenchmarkAssetNotAvailableTooLate(
75
                sid=str(self.benchmark_sid),
76
                dt=self.trading_days[0],
77
                end_dt=self.benchmark_asset.end_date
78
            )
79
80
    def _initialize_precalculated_series(self, sid, env, trading_days,
81
                                         data_portal):
82
        """
83
        Internal method that precalculates the benchmark return series for
84
        use in the simulation.
85
86
        Parameters
87
        ----------
88
        sid: (int) Asset to use
89
90
        env: TradingEnvironment
91
92
        trading_days: pd.DateTimeIndex
93
94
        data_portal: DataPortal
95
96
        Notes
97
        -----
98
        If the benchmark asset started trading after the simulation start,
99
        or finished trading before the simulation end, exceptions are raised.
100
101
        If the benchmark asset started trading the same day as the simulation
102
        start, the first available minute price on that day is used instead
103
        of the previous close.
104
105
        We use history to get an adjusted price history for each day's close,
106
        as of the look-back date (the last day of the simulation).  Prices are
107
        fully adjusted for dividends, splits, and mergers.
108
109
        Returns
110
        -------
111
        A pd.Series, indexed by trading day, whose values represent the %
112
        change from close to close.
113
        """
114
        if sid is None:
115
            # get benchmark info from trading environment, which defaults to
116
            # downloading data from Yahoo.
117
            daily_series = \
118
                env.benchmark_returns[trading_days[0]:trading_days[-1]]
119
120
            if self.emission_rate == "minute":
121
                # we need to take the env's benchmark returns, which are daily,
122
                # and resample them to minute
123
                minutes = env.minutes_for_days_in_range(
124
                    start=trading_days[0],
125
                    end=trading_days[-1]
126
                )
127
128
                minute_series = daily_series.reindex(
129
                    index=minutes,
130
                    method="ffill"
131
                )
132
133
                return minute_series
134
            else:
135
                return daily_series
136
        elif self.emission_rate == "minute":
137
            minutes = env.minutes_for_days_in_range(self.trading_days[0],
138
                                                    self.trading_days[-1])
139
            benchmark_series = data_portal.get_history_window(
140
                [sid],
141
                minutes[-1],
142
                bar_count=len(minutes) + 1,
143
                frequency="1m",
144
                field="price",
145
                ffill=True
146
            )
147
148
            return benchmark_series.pct_change()[1:]
149
        else:
150
            try:
151
                # get the window of close prices for benchmark_sid from the
152
                # last trading day of the simulation, going up to one day
153
                # before the simulation start day (so that we can get the %
154
                # change on day 1)
155
                benchmark_series = data_portal.get_history_window(
156
                    [sid],
157
                    trading_days[-1],
158
                    bar_count=len(trading_days) + 1,
159
                    frequency="1d",
160
                    field="price",
161
                    ffill=True
162
                )[sid]
163
                return benchmark_series.pct_change()[1:]
164
            except NoDataOnDate:
165
                # Attempt to handle case where stock data starts on first
166
                # day, in this case use the open to close return.
167
                benchmark_series = data_portal.get_history_window(
168
                    [sid],
169
                    trading_days[-1],
170
                    bar_count=len(trading_days),
171
                    frequency="1d",
172
                    field="price",
173
                    ffill=True
174
                )[sid]
175
176
                # get a minute history window of the first day
177
                first_open = data_portal.get_spot_value(
178
                    sid, 'open', trading_days[0])
179
                first_close = data_portal.get_spot_value(
180
                    sid, 'close', trading_days[0])
181
182
                first_day_return = (first_close - first_open) / first_open
183
184
                returns = benchmark_series.pct_change()[:]
185
                returns[0] = first_day_return
186
                return returns
187