From 27f94f83fa05f6a83c9585179438eea4d5b44808 Mon Sep 17 00:00:00 2001 From: Eddie Hebert Date: Tue, 2 Feb 2016 10:26:37 -0500 Subject: [PATCH] ENH: Allow passing of numpy arrays to writer. For faster parsing and writing workflows, do not require a DataFrame. --- tests/data/test_minute_bars.py | 56 ++++++++++++++++++++++++++++++++- zipline/data/minute_bars.py | 57 +++++++++++++++++++++++++++------- 2 files changed, 101 insertions(+), 12 deletions(-) diff --git a/tests/data/test_minute_bars.py b/tests/data/test_minute_bars.py index 1496f605..c218555d 100644 --- a/tests/data/test_minute_bars.py +++ b/tests/data/test_minute_bars.py @@ -17,7 +17,7 @@ import os from unittest import TestCase -from numpy import nan +from numpy import nan, array from numpy.testing import assert_almost_equal from pandas import ( DataFrame, @@ -446,3 +446,57 @@ class BcolzMinuteBarTestCase(TestCase): volume_price = self.reader.get_value(sid, minute, 'volume') self.assertEquals(100.0, volume_price) + + def test_write_cols(self): + minute_0 = self.market_opens[self.test_calendar_start] + minute_1 = minute_0 + timedelta(minutes=1) + sid = 1 + cols = { + 'open': array([10.0, 11.0]), + 'high': array([20.0, 21.0]), + 'low': array([30.0, 31.0]), + 'close': array([40.0, 41.0]), + 'volume': array([50.0, 51.0]) + } + dts = array([minute_0, minute_1], dtype='datetime64[s]') + self.writer.write_cols(sid, dts, cols) + + open_price = self.reader.get_value(sid, minute_0, 'open') + + self.assertEquals(10.0, open_price) + + high_price = self.reader.get_value(sid, minute_0, 'high') + + self.assertEquals(20.0, high_price) + + low_price = self.reader.get_value(sid, minute_0, 'low') + + self.assertEquals(30.0, low_price) + + close_price = self.reader.get_value(sid, minute_0, 'close') + + self.assertEquals(40.0, close_price) + + volume_price = self.reader.get_value(sid, minute_0, 'volume') + + self.assertEquals(50.0, volume_price) + + open_price = self.reader.get_value(sid, minute_1, 'open') + + self.assertEquals(11.0, open_price) + + high_price = self.reader.get_value(sid, minute_1, 'high') + + self.assertEquals(21.0, high_price) + + low_price = self.reader.get_value(sid, minute_1, 'low') + + self.assertEquals(31.0, low_price) + + close_price = self.reader.get_value(sid, minute_1, 'close') + + self.assertEquals(41.0, close_price) + + volume_price = self.reader.get_value(sid, minute_1, 'volume') + + self.assertEquals(51.0, volume_price) diff --git a/zipline/data/minute_bars.py b/zipline/data/minute_bars.py index dc604dcd..b459851e 100644 --- a/zipline/data/minute_bars.py +++ b/zipline/data/minute_bars.py @@ -20,7 +20,6 @@ from os.path import join import json import os import pandas as pd -from pandas.core.datetools import normalize_date US_EQUITIES_MINUTES_PER_DAY = 390 @@ -403,8 +402,6 @@ class BcolzMinuteBarWriter(object): ----------- sid : int The asset identifer for the data being written. - days : pd.DatetimeIndex - The days for which to write data from the given df. df : pd.DataFrame DataFrame of market data with the following characteristics. columns : ('open', 'high', 'low', 'close', 'volume') @@ -415,11 +412,49 @@ class BcolzMinuteBarWriter(object): volume : float64|int64 index : DatetimeIndex of market minutes. """ + cols = { + 'open': df.open.values, + 'high': df.high.values, + 'low': df.low.values, + 'close': df.close.values, + 'volume': df.volume.values, + } + dts = df.index.values + self.write_cols(sid, dts, cols) + + def write_cols(self, sid, dts, cols): + """ + Write the OHLCV data for the given sid. + + If there is no bcolz ctable yet created for the sid, create it. + + If the length of the bcolz ctable is not exactly to the date before + the first day provided, fill the ctable with 0s up to that date. + + Writes in blocks of the size of the days times minutes per day. + + Parameters: + ----------- + sid : int + The asset identifer for the data being written. + dts : datetime64 array + The dts corresponding to values in cols. + cols : dict of str -> np.array + dict of market data with the following characteristics. + keys are ('open', 'high', 'low', 'close', 'volume') + open : float64 + high : float64 + low : float64 + close : float64 + volume : float64|int64 + """ table = self._ensure_ctable(sid) tds = self._trading_days - input_first_day = normalize_date(df.index[0]) - input_last_day = normalize_date(df.index[-1]) + input_first_day = pd.Timestamp(dts[0].astype('datetime64[D]'), + tz='UTC') + input_last_day = pd.Timestamp(dts[-1].astype('datetime64[D]'), + tz='UTC') last_date = self.last_date_in_output_for_sid(sid) @@ -449,15 +484,15 @@ class BcolzMinuteBarWriter(object): vol_col = np.zeros(minutes_count, dtype=np.uint32) dt_ixs = np.searchsorted(all_minutes_in_window.values, - df.index.values) + dts.astype('datetime64[ns]')) ohlc_ratio = self._ohlc_ratio - open_col[dt_ixs] = (df.open.values * ohlc_ratio).astype(np.uint32) - high_col[dt_ixs] = (df.high.values * ohlc_ratio).astype(np.uint32) - low_col[dt_ixs] = (df.low.values * ohlc_ratio).astype(np.uint32) - close_col[dt_ixs] = (df.close.values * ohlc_ratio).astype( + open_col[dt_ixs] = (cols['open'] * ohlc_ratio).astype(np.uint32) + high_col[dt_ixs] = (cols['high'] * ohlc_ratio).astype(np.uint32) + low_col[dt_ixs] = (cols['low'] * ohlc_ratio).astype(np.uint32) + close_col[dt_ixs] = (cols['close'] * ohlc_ratio).astype( np.uint32) - vol_col[dt_ixs] = df.volume.values.astype(np.uint32) + vol_col[dt_ixs] = cols['volume'].astype(np.uint32) table.append([ open_col,