zipline.data.BcolzMinuteBarWriter.full_minutes_for_days() - Code Metrics - Inspection of "Q 2.0" - quantopian/zipline - Measure and Improve Code Quality continuously with Scrutinizer

Completed

Pull Request — master (#858)

by Eddie

created 2015-12-30 18:07 UTC

full_minutes_for_days() A

↳ Parent: zipline.data.BcolzMinuteBarWriter

Complexity

Conditions

Size

Total Lines

Duplication

Lines	0
Ratio	0 %

Metric	Value
cc	1
dl	0
loc	6
rs	9.4286

# Copyright 2015 Quantopian, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from abc import (
    ABCMeta,
    abstractmethod,
)
import bcolz
import json
import os
from bcolz import ctable
from datetime import datetime
import numpy as np
from numpy import float64
from os.path import join
import pandas as pd
from pandas import read_csv
from six import with_metaclass

from zipline.finance.trading import TradingEnvironment
from zipline.utils import tradingcalendar

MINUTES_PER_DAY = 390

_writer_env = TradingEnvironment()

_NANOS_IN_MINUTE = 60000000000

METADATA_FILENAME = 'metadata.json'


def write_metadata(directory, first_trading_day):
    metadata_path = os.path.join(directory, METADATA_FILENAME)

    metadata = {
        'first_trading_day': str(first_trading_day.date())
    }

    with open(metadata_path, 'w') as fp:
        json.dump(metadata, fp)


def _bcolz_minute_index(trading_days):
    minutes = np.zeros(len(trading_days) * MINUTES_PER_DAY,
                       dtype='datetime64[ns]')
    market_opens = tradingcalendar.open_and_closes.market_open
    mask = market_opens.index.slice_indexer(start=trading_days[0],
                                            end=trading_days[-1])
    opens = market_opens[mask]

    deltas = np.arange(0, MINUTES_PER_DAY, dtype='timedelta64[m]')
    for i, market_open in enumerate(opens):
        start = market_open.asm8
        minute_values = start + deltas
        start_ix = MINUTES_PER_DAY * i
        end_ix = start_ix + MINUTES_PER_DAY
        minutes[start_ix:end_ix] = minute_values
    return pd.to_datetime(minutes, utc=True, box=True)


class BcolzMinuteBarWriter(with_metaclass(ABCMeta)):
    """
    Class capable of writing minute OHLCV data to disk into bcolz format.
    """
    @property
    def first_trading_day(self):
        return self._first_trading_day

    @abstractmethod
    def gen_frames(self, assets):
        """
        Return an iterator of pairs of (asset_id, pd.dataframe).
        """
        raise NotImplementedError()

    def write(self, directory, assets, sid_path_func=None):
        _iterator = self.gen_frames(assets)

        return self._write_internal(directory, _iterator,
                                    sid_path_func=sid_path_func)

    def full_minutes_for_days(self, dt1, dt2):
        start_date = _writer_env.normalize_date(dt1)
        end_date = _writer_env.normalize_date(dt2)

        trading_days = _writer_env.days_in_range(start_date, end_date)
        return _bcolz_minute_index(trading_days)

    def _write_internal(self, directory, iterator, sid_path_func=None):
        first_trading_day = self.first_trading_day

        write_metadata(directory, first_trading_day)

        first_open = pd.Timestamp(
            datetime(
                year=first_trading_day.year,
                month=first_trading_day.month,
                day=first_trading_day.day,
                hour=9,
                minute=31
            ), tz='US/Eastern').tz_convert('UTC')

        all_minutes = None

        for asset_id, df in iterator:
            if sid_path_func is None:
                path = join(directory, "{0}.bcolz".format(asset_id))
            else:
                path = sid_path_func(directory, asset_id)

            os.makedirs(path)

            last_dt = df.index[-1]

            if all_minutes is None:
                all_minutes = \
                    self.full_minutes_for_days(first_open, last_dt)
                minutes = all_minutes
            else:
                if df.index[-1] in all_minutes:
                    mask = all_minutes.slice_indexer(end=last_dt)
                    minutes = all_minutes[mask]
                else:
                    # Need to extend all minutes from open after last value
                    # in all_minutes to the last_dt.
                    next_open, _ = _writer_env.next_open_and_close(
                        all_minutes[-1])
                    to_append = self.full_minutes_for_days(next_open, last_dt)
                    all_minutes = all_minutes.append(to_append)
                    minutes = all_minutes

            minutes_count = len(minutes)

            open_col = np.zeros(minutes_count, dtype=np.uint32)
            high_col = np.zeros(minutes_count, dtype=np.uint32)
            low_col = np.zeros(minutes_count, dtype=np.uint32)
            close_col = np.zeros(minutes_count, dtype=np.uint32)
            vol_col = np.zeros(minutes_count, dtype=np.uint32)

            opens = df.open.values.astype(np.uint32)
            highs = df.high.values.astype(np.uint32)
            lows = df.low.values.astype(np.uint32)
            closes = df.close.values.astype(np.uint32)
            volumes = df.volume.values.astype(np.uint32)

            dt_ixs = np.searchsorted(minutes.values, df.index.values)

            for i, dt_ix in enumerate(dt_ixs):
                # Each day has 390 slots, where 9:31 is the first
                # slot (ix=0) of the day, and each slot represents a
                # minute's data.
                #
                # Get the difference in seconds between the current
                # minute and market open and then divide by 60 to get
                # the index into which to write, while still writing
                    # from the same row aligned with the dt in the from the
                # CSV.
                open_col[dt_ix] = opens[i]
                high_col[dt_ix] = highs[i]
                low_col[dt_ix] = lows[i]
                close_col[dt_ix] = closes[i]
                vol_col[dt_ix] = volumes[i]

            ctable(
                columns=[
                    open_col,
                    high_col,
                    low_col,
                    close_col,
                    vol_col,
                ],
                names=[
                    "open",
                    "high",
                    "low",
                    "close",
                    "volume",
                ],
                rootdir=path,
                mode='w'
            )


class MinuteBarWriterFromDataFrames(BcolzMinuteBarWriter):
    _csv_dtypes = {
        'open': float64,
        'high': float64,
        'low': float64,
        'close': float64,
        'volume': float64,
    }

    def __init__(self, first_trading_day):
        self._first_trading_day = first_trading_day

    def gen_frames(self, assets):
        for asset in assets:
            df = assets[asset]
            yield asset, df.set_index("minute")


class MinuteBarWriterFromCSVs(BcolzMinuteBarWriter):
    """
    BcolzMinuteBarWriter constructed from a map of CSVs to assets.

    Parameters
    ----------
    asset_map: dict
        A map from asset_id -> path to csv with data for that asset.

    CSVs should have the following columns:
        minute : datetime64
        open : float64
        high : float64
        low : float64
        close : float64
        volume : int64
    """
    _csv_dtypes = {
        'open': float64,
        'high': float64,
        'low': float64,
        'close': float64,
        'volume': float64,
    }

    def __init__(self, asset_map, first_trading_day):
        self._asset_map = asset_map
        self._first_trading_day = first_trading_day

    def gen_frames(self, assets):
        """
        Read CSVs as DataFrames from our asset map.
        """
        dtypes = self._csv_dtypes

        for asset in assets:
            path = self._asset_map.get(asset)
            if path is None:
                raise KeyError("No path supplied for asset %s" % asset)
            df = read_csv(path, parse_dates=['minute'], dtype=dtypes)
            df = df.set_index("minute").tz_localize("UTC")

            yield asset, df


class BcolzMinuteBarReader(object):

    def __init__(self, rootdir, sid_path_func=None):
        self.rootdir = rootdir

        metadata = self._get_metadata()

        self.first_trading_day = pd.Timestamp(
            metadata['first_trading_day'], tz='UTC')
        mask = tradingcalendar.trading_days.slice_indexer(
            self.first_trading_day)
        # TODO: Read/write calendar to match daily, so that calendar is not
        # 'hardcoded'.
        self.trading_days = tradingcalendar.trading_days[mask]
        self._sid_path_func = sid_path_func

        self._carrays = {
            'open': {},
            'high': {},
            'low': {},
            'close': {},
            'volume': {},
            'sid': {},
            'dt': {},
        }

        self._minute_index = _bcolz_minute_index(self.trading_days)

    def _get_metadata(self):
        with open(os.path.join(self.rootdir, METADATA_FILENAME)) as fp:
            return json.load(fp)

    def _get_ctable(self, asset):
        sid = int(asset)
        if self._sid_path_func is not None:
            path = self._sid_path_func(self.rootdir, sid)
        else:
            path = "{0}/{1}.bcolz".format(self.rootdir, sid)

        return bcolz.open(path, mode='r')

    def get_last_traded_dt(self, asset, dt):
        return self._minute_index[
            self._find_last_traded_position(asset, dt)]

    def _find_last_traded_position(self, asset, dt):
        volumes = self._open_minute_file('volume', asset)
        minute_pos = self._find_position_of_minute(dt)
        while True:
            if minute_pos == 0 or volumes[minute_pos] != 0:
                return minute_pos
            minute_pos -= 1

    def _find_position_of_minute(self, minute_dt):
        """
        Internal method that returns the position of the given minute in the
        list of every trading minute since market open of the first trading
        day.

        IMPORTANT: This method assumes every day is 390 minutes long, even
        early closes.  Our minute bcolz files are generated like this to
        support fast lookup.

        ex. this method would return 2 for 1/2/2002 9:32 AM Eastern, if
        1/2/2002 is the first trading day of the dataset.

        Parameters
        ----------
        minute_dt: pd.Timestamp
            The minute whose position should be calculated.

        Returns
        -------
        The position of the given minute in the list of all trading minutes
        since market open on the first trading day.
        """
        return self._minute_index.get_loc(minute_dt)

    def _open_minute_file(self, field, asset):
        sid_str = str(int(asset))

        try:
            carray = self._carrays[field][sid_str]
        except KeyError:
            carray = self._carrays[field][sid_str] = \
                self._get_ctable(asset)[field]

        return carray


1			# Copyright 2015 Quantopian, Inc.
2			#
3			# Licensed under the Apache License, Version 2.0 (the "License");
4			# you may not use this file except in compliance with the License.
5			# You may obtain a copy of the License at
6			#
7			# http://www.apache.org/licenses/LICENSE-2.0
8			#
9			# Unless required by applicable law or agreed to in writing, software
10			# distributed under the License is distributed on an "AS IS" BASIS,
11			# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12			# See the License for the specific language governing permissions and
13			# limitations under the License.
14
15			from abc import (
16			ABCMeta,
17			abstractmethod,
18			)
19			import bcolz
20			import json
21			import os
22			from bcolz import ctable
23			from datetime import datetime
24			import numpy as np
25			from numpy import float64
26			from os.path import join
27			import pandas as pd
28			from pandas import read_csv
29			from six import with_metaclass
30
31			from zipline.finance.trading import TradingEnvironment
32			from zipline.utils import tradingcalendar
33
34			MINUTES_PER_DAY = 390
35
36			_writer_env = TradingEnvironment()
37
38			_NANOS_IN_MINUTE = 60000000000
39
40			METADATA_FILENAME = 'metadata.json'
41
42
43			def write_metadata(directory, first_trading_day):
44			metadata_path = os.path.join(directory, METADATA_FILENAME)
45
46			metadata = {
47			'first_trading_day': str(first_trading_day.date())
48			}
49
50			with open(metadata_path, 'w') as fp:
51			json.dump(metadata, fp)
52
53
54			def _bcolz_minute_index(trading_days):
55			minutes = np.zeros(len(trading_days) * MINUTES_PER_DAY,
56			dtype='datetime64[ns]')
57			market_opens = tradingcalendar.open_and_closes.market_open
58			mask = market_opens.index.slice_indexer(start=trading_days[0],
59			end=trading_days[-1])
60			opens = market_opens[mask]
61
62			deltas = np.arange(0, MINUTES_PER_DAY, dtype='timedelta64[m]')
63			for i, market_open in enumerate(opens):
64			start = market_open.asm8
65			minute_values = start + deltas
66			start_ix = MINUTES_PER_DAY * i
67			end_ix = start_ix + MINUTES_PER_DAY
68			minutes[start_ix:end_ix] = minute_values
69			return pd.to_datetime(minutes, utc=True, box=True)
70
71
72			class BcolzMinuteBarWriter(with_metaclass(ABCMeta)):
73			"""
74			Class capable of writing minute OHLCV data to disk into bcolz format.
75			"""
76			@property
77			def first_trading_day(self):
78			return self._first_trading_day
79
80			@abstractmethod
81			def gen_frames(self, assets):
82			"""
83			Return an iterator of pairs of (asset_id, pd.dataframe).
84			"""
85			raise NotImplementedError()
86
87			def write(self, directory, assets, sid_path_func=None):
88			_iterator = self.gen_frames(assets)
89
90			return self._write_internal(directory, _iterator,
91			sid_path_func=sid_path_func)
92
93			def full_minutes_for_days(self, dt1, dt2):
94			start_date = _writer_env.normalize_date(dt1)
95			end_date = _writer_env.normalize_date(dt2)
96
97			trading_days = _writer_env.days_in_range(start_date, end_date)
98			return _bcolz_minute_index(trading_days)
99
100			def _write_internal(self, directory, iterator, sid_path_func=None):
101			first_trading_day = self.first_trading_day
102
103			write_metadata(directory, first_trading_day)
104
105			first_open = pd.Timestamp(
106			datetime(
107			year=first_trading_day.year,
108			month=first_trading_day.month,
109			day=first_trading_day.day,
110			hour=9,
111			minute=31
112			), tz='US/Eastern').tz_convert('UTC')
113
114			all_minutes = None
115
116			for asset_id, df in iterator:
117			if sid_path_func is None:
118			path = join(directory, "{0}.bcolz".format(asset_id))
119			else:
120			path = sid_path_func(directory, asset_id)
121
122			os.makedirs(path)
123
124			last_dt = df.index[-1]
125
126			if all_minutes is None:
127			all_minutes = \
128			self.full_minutes_for_days(first_open, last_dt)
129			minutes = all_minutes
130			else:
131			if df.index[-1] in all_minutes:
132			mask = all_minutes.slice_indexer(end=last_dt)
133			minutes = all_minutes[mask]
134			else:
135			# Need to extend all minutes from open after last value
136			# in all_minutes to the last_dt.
137			next_open, _ = _writer_env.next_open_and_close(
138			all_minutes[-1])
139			to_append = self.full_minutes_for_days(next_open, last_dt)
140			all_minutes = all_minutes.append(to_append)
141			minutes = all_minutes
142
143			minutes_count = len(minutes)
144
145			open_col = np.zeros(minutes_count, dtype=np.uint32)
146			high_col = np.zeros(minutes_count, dtype=np.uint32)
147			low_col = np.zeros(minutes_count, dtype=np.uint32)
148			close_col = np.zeros(minutes_count, dtype=np.uint32)
149			vol_col = np.zeros(minutes_count, dtype=np.uint32)
150
151			opens = df.open.values.astype(np.uint32)
152			highs = df.high.values.astype(np.uint32)
153			lows = df.low.values.astype(np.uint32)
154			closes = df.close.values.astype(np.uint32)
155			volumes = df.volume.values.astype(np.uint32)
156
157			dt_ixs = np.searchsorted(minutes.values, df.index.values)
158
159			for i, dt_ix in enumerate(dt_ixs):
160			# Each day has 390 slots, where 9:31 is the first
161			# slot (ix=0) of the day, and each slot represents a
162			# minute's data.
163			#
164			# Get the difference in seconds between the current
165			# minute and market open and then divide by 60 to get
166			# the index into which to write, while still writing
167			# from the same row aligned with the dt in the from the
168			# CSV.
169			open_col[dt_ix] = opens[i]
170			high_col[dt_ix] = highs[i]
171			low_col[dt_ix] = lows[i]
172			close_col[dt_ix] = closes[i]
173			vol_col[dt_ix] = volumes[i]
174
175			ctable(
176			columns=[
177			open_col,
178			high_col,
179			low_col,
180			close_col,
181			vol_col,
182			],
183			names=[
184			"open",
185			"high",
186			"low",
187			"close",
188			"volume",
189			],
190			rootdir=path,
191			mode='w'
192			)
193
194
195			class MinuteBarWriterFromDataFrames(BcolzMinuteBarWriter):
196			_csv_dtypes = {
197			'open': float64,
198			'high': float64,
199			'low': float64,
200			'close': float64,
201			'volume': float64,
202			}
203
204			def __init__(self, first_trading_day):
205			self._first_trading_day = first_trading_day
206
207			def gen_frames(self, assets):
208			for asset in assets:
209			df = assets[asset]
210			yield asset, df.set_index("minute")
211
212
213			class MinuteBarWriterFromCSVs(BcolzMinuteBarWriter):
214			"""
215			BcolzMinuteBarWriter constructed from a map of CSVs to assets.
216
217			Parameters
218			----------
219			asset_map: dict
220			A map from asset_id -> path to csv with data for that asset.
221
222			CSVs should have the following columns:
223			minute : datetime64
224			open : float64
225			high : float64
226			low : float64
227			close : float64
228			volume : int64
229			"""
230			_csv_dtypes = {
231			'open': float64,
232			'high': float64,
233			'low': float64,
234			'close': float64,
235			'volume': float64,
236			}
237
238			def __init__(self, asset_map, first_trading_day):
239			self._asset_map = asset_map
240			self._first_trading_day = first_trading_day
241
242			def gen_frames(self, assets):
243			"""
244			Read CSVs as DataFrames from our asset map.
245			"""
246			dtypes = self._csv_dtypes
247
248			for asset in assets:
249			path = self._asset_map.get(asset)
250			if path is None:
251			raise KeyError("No path supplied for asset %s" % asset)
252			df = read_csv(path, parse_dates=['minute'], dtype=dtypes)
253			df = df.set_index("minute").tz_localize("UTC")
254
255			yield asset, df
256
257
258			class BcolzMinuteBarReader(object):
259
260			def __init__(self, rootdir, sid_path_func=None):
261			self.rootdir = rootdir
262
263			metadata = self._get_metadata()
264
265			self.first_trading_day = pd.Timestamp(
266			metadata['first_trading_day'], tz='UTC')
267			mask = tradingcalendar.trading_days.slice_indexer(
268			self.first_trading_day)
269			# TODO: Read/write calendar to match daily, so that calendar is not
270			# 'hardcoded'.
271			self.trading_days = tradingcalendar.trading_days[mask]
272			self._sid_path_func = sid_path_func
273
274			self._carrays = {
275			'open': {},
276			'high': {},
277			'low': {},
278			'close': {},
279			'volume': {},
280			'sid': {},
281			'dt': {},
282			}
283
284			self._minute_index = _bcolz_minute_index(self.trading_days)
285
286			def _get_metadata(self):
287			with open(os.path.join(self.rootdir, METADATA_FILENAME)) as fp:
288			return json.load(fp)
289
290			def _get_ctable(self, asset):
291			sid = int(asset)
292			if self._sid_path_func is not None:
293			path = self._sid_path_func(self.rootdir, sid)
294			else:
295			path = "{0}/{1}.bcolz".format(self.rootdir, sid)
296
297			return bcolz.open(path, mode='r')
298
299			def get_last_traded_dt(self, asset, dt):
300			return self._minute_index[
301			self._find_last_traded_position(asset, dt)]
302
303			def _find_last_traded_position(self, asset, dt):
304			volumes = self._open_minute_file('volume', asset)
305			minute_pos = self._find_position_of_minute(dt)
306			while True:
307			if minute_pos == 0 or volumes[minute_pos] != 0:
308			return minute_pos
309			minute_pos -= 1
310
311			def _find_position_of_minute(self, minute_dt):
312			"""
313			Internal method that returns the position of the given minute in the
314			list of every trading minute since market open of the first trading
315			day.
316
317			IMPORTANT: This method assumes every day is 390 minutes long, even
318			early closes. Our minute bcolz files are generated like this to
319			support fast lookup.
320
321			ex. this method would return 2 for 1/2/2002 9:32 AM Eastern, if
322			1/2/2002 is the first trading day of the dataset.
323
324			Parameters
325			----------
326			minute_dt: pd.Timestamp
327			The minute whose position should be calculated.
328
329			Returns
330			-------
331			The position of the given minute in the list of all trading minutes
332			since market open on the first trading day.
333			"""
334			return self._minute_index.get_loc(minute_dt)
335
336			def _open_minute_file(self, field, asset):
337			sid_str = str(int(asset))
338
339			try:
340			carray = self._carrays[field][sid_str]
341			except KeyError:
342			carray = self._carrays[field][sid_str] = \
343			self._get_ctable(asset)[field]
344
345			return carray
346

quantopian / zipline

Pull Request — master (#858)

full_minutes_for_days() A

Complexity

Size

Duplication

Duplication Side-by-Side

Filter issues like