From 59bcd097d5e9395a03c845147bc4b519b1353c3e Mon Sep 17 00:00:00 2001 From: Michael Schatzow Date: Sun, 1 Dec 2013 11:56:31 -0500 Subject: [PATCH] ENH: Add hdf5 and csv source. This creates a data source for csv and hdf5 files, a generator to create a sample csv, and a pytables generator to go from a list of dated gzipped csv's in a directory to a pytables data source. This does not add a unittest yet which we should write for the future. --- zipline/sources/data_source_csv.py | 172 ++++++++++++++++++ zipline/sources/data_source_tables.py | 220 ++++++++++++++++++++++++ zipline/utils/data_source_tables_gen.py | 194 +++++++++++++++++++++ 3 files changed, 586 insertions(+) create mode 100644 zipline/sources/data_source_csv.py create mode 100644 zipline/sources/data_source_tables.py create mode 100644 zipline/utils/data_source_tables_gen.py diff --git a/zipline/sources/data_source_csv.py b/zipline/sources/data_source_csv.py new file mode 100644 index 00000000..519e9eb0 --- /dev/null +++ b/zipline/sources/data_source_csv.py @@ -0,0 +1,172 @@ +# 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. +""" +leverage work of briancappello and quantopian team +(especcially twiecki, eddie, and fawce) +michaelws +""" +import pandas as pd +from zipline.gens.utils import hash_args +from zipline.sources.data_source import DataSource +import datetime +import csv +import numpy as np +import dateutil.parser + + +def gen_ts(date, time): + return pd.Timestamp(datetime.datetime.combine(date, time)) + + +class DatasourceCSVohlc(DataSource): + """ expects dictReader for a csv file + with the following columns in the header + dt, sid, open, high, low, close, volume + dt expected in ISO format and order does not matter""" + def __init__(self, data, **kwargs): + isinstance(data, csv.DictReader) + self.data = data + # Unpack config dictionary with default values. + self.tz_in = kwargs.get('tz_in', "US/Eastern") + self.start = pd.Timestamp(np.datetime64(kwargs.get('start'))) + self.start = self.start.tz_localize('utc') + self.end = pd.Timestamp(np.datetime64(kwargs.get('end'))) + self.end = self.end.tz_localize('utc') + start_time_str = kwargs.get("start_time", "9:30") + end_time_str = kwargs.get("end_time", "16:00") + self.sid_filter = kwargs.get('sid_filter', None) + self.source_id = kwargs.get("source_id", None) + self.sids = kwargs.get('sidsF', None) + self.start_time = dateutil.parser.parse(start_time_str).time() + self.end_time = dateutil.parser.parse(end_time_str).time() + self._raw_data = None + self.arg_string = hash_args(data, **kwargs) + + @property + def instance_hash(self): + return self.arg_string + + def raw_data_gen(self): + previous_ts = None + cols = np.array(["open", "high", "low"]) + for row in self.data: + dt64 = pd.Timestamp(np.datetime64(row["dt"])) + ts = pd.Timestamp(dt64).tz_localize(self.tz_in).tz_convert('utc') + if ts < self.start or ts > self.end: + continue + if previous_ts is None or ts.date() != previous_ts.date(): + start_ts = datetime.date(ts.date(), self.start_time) + end_ts = gen_ts(ts.date(), self.end_time) + volumes = {} + price_volumes = {} + sid = row["sid"] + if self.sid_filter and sid in self.sid_filter: + continue + elif self.sids is None or sid in self.sids: + if sid not in volumes: + volumes[sid] = 0 + price_volumes[sid] = 0 + if ts < start_ts or ts > end_ts: + continue + event = {"sid": sid, "type": "TRADE", "symbol": sid} + event["dt"] = ts + event["price"] = float(row["close"]) + event["close"] = event["price"] + event["volume"] = int(row["volume"]) + volumes[sid] += float(event["volume"]) + price_volumes[sid] += event["price"] * event["volume"] + event["vwap"] = price_volumes[sid] / volumes[sid] + for field in cols: + event[field] = float(row[field]) + yield event + previous_ts = ts + + @property + def raw_data(self): + if not self._raw_data: + self._raw_data = self.raw_data_gen() + return self._raw_data + + @property + def mapping(self): + return { + 'sid': (lambda x: x, 'sid'), + 'dt': (lambda x: x, 'dt'), + 'open': (float, 'open'), + 'high': (float, 'high'), + 'low': (float, 'low'), + 'close': (float, 'close'), + 'price': (float, 'price'), + 'volume': (int, 'volume'), + 'vwap': (lambda x: x, 'vwap') + } + + +class DataSourceCSVSignal(DataSource): + """ expects dictReader for a csv file in form with header + dt, sid, signal + dt expected in ISO format""" + def __init__(self, data, **kwargs): + assert isinstance(data, csv.DictReader) + self.data = data + self.source_id = kwargs.get("source_id", None) + # Unpack config dictionary with default values. + self.start = kwargs.get('start') + self.end = kwargs.get('end') + self.sids = kwargs.get('sids', None) + self.sid_filter = kwargs.get('sid_filter', None) + self.arg_string = hash_args(data, **kwargs) + self._raw_data = None + + @property + def instance_hash(self): + return self.arg_string + + def raw_data_gen(self): + previous_ts = None + for row in self.data: + dt64 = pd.Timestamp(np.datetime64(row["dt"])) + ts = pd.Timestamp(dt64).tz_localize(self.tz_in).tz_convert('utc') + if ts < self.start or ts > self.end: + continue + if previous_ts is None or ts.date() != previous_ts.date(): + start_ts = gen_ts(ts.date(), self.start_time) + end_ts = gen_ts(ts.date(), self.end_time) + volumes = {} + price_volumes = {} + sid = row["sid"] + if self.sid_filter and sid in self.sid_filter: + continue + elif self.sids is None or sid in self.sids: + if sid not in volumes: + volumes[sid] = 0 + price_volumes[sid] = 0 + if ts < start_ts or ts > end_ts: + continue + event = {"sid": sid, "type": "CUSTOM", "dt": ts, + "signal": row["signal"]} + yield event + previous_ts = ts + + @property + def raw_data(self): + if not self._raw_data: + self._raw_data = self.raw_data_gen() + return self._raw_data + + @property + def mapping(self): + return { + 'sid': (lambda x: x, 'symbol'), + 'dt': (lambda x: x, 'dt'), + 'signal': (lambda x: x, 'signal'), + } diff --git a/zipline/sources/data_source_tables.py b/zipline/sources/data_source_tables.py new file mode 100644 index 00000000..fcce1727 --- /dev/null +++ b/zipline/sources/data_source_tables.py @@ -0,0 +1,220 @@ +# 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. +""" +leverage work of briancappello and quantopian team +(especially twiecki, eddie, and fawce) +""" +import pandas as pd +from zipline.gens.utils import hash_args +from zipline.sources.data_source import DataSource +import datetime +import numpy as np +import dateutil.parser +import tables + + +def _iterate_ohlc(date_node, sid_filter, sids, start_ts, end_ts): + last_stamp = None + last_ts = None + volumes = {} + price_volumes = {} + cols = np.array(["open", "high", "low", "close"]) + for row in date_node.iterrows(): + sid = row["sid"] + if sid_filter and sid in sid_filter: + continue + elif sids is None or sid in sids: + if sid not in volumes: + volumes[sid] = 0 + price_volumes[sid] = 0 + if last_stamp and row["dt"] == last_stamp: + ts = last_ts + else: + ts = pd.Timestamp(np.datetime64(row["dt"], "s"), tz='utc') + last_ts = ts + last_stamp = row["dt"] + if (start_ts > ts) or (ts > end_ts): + continue + event = {"sid": sid, "type": "TRADE", "symbol": sid} + event["dt"] = ts + event["price"] = row["close"] + event["volume"] = row["volume"] + volumes[sid] += event["volume"] + price_volumes[sid] += event["price"] * event["volume"] + event["vwap"] = price_volumes[sid] / volumes[sid] + last_ts = ts + for field in cols: + event[field] = row[field] + yield event + + +def _iterate_signal(date_node, sids, sid_filter, start_ts, end_ts): + last_stamp = None + last_ts = None + volumes = {} + price_volumes = {} + for row in date_node.iterrows(): + sid = row["sid"] + if sid_filter and sid in sid_filter: + continue + elif sids is None or sid in sids: + if sid not in volumes: + volumes[sid] = 0 + price_volumes[sid] = 0 + if last_stamp and row["dt"] == last_stamp: + ts = last_ts + else: + ts = pd.Timestamp(np.datetime64(row["dt"], "s"), tz='utc') + last_ts = ts + last_stamp = row["dt"] + if (start_ts > ts) or (ts > end_ts): + continue + event = {"sid": sid, "type": "CUSTOM", + "signal": row["signal"]} + yield event + + +class DataSourceTablesOHLC(DataSource): + """ + Yields all events in event_list that match the given sid_filter. + If no event_list is specified, generates an internal stream of events + to filter. Returns all events if filter is None. + + Configuration options: + + sids : list of values representing simulated internal sids + start : start date + tz_in : timezzone of table + filter : filter to remove the sids + start_time: what time trading should start + end_time: what time trading should end + """ + def __init__(self, data, **kwargs): + assert isinstance(data, tables.file.File) + self.data = data + # Unpack config dictionary with default values. + if 'symbols' in kwargs: + self.sids = kwargs.get('symbols') + else: + self.sids = None + self.tz_in = kwargs.get('tz_in', "US/Eastern") + self.source_id = kwargs.get("source_id", None) + self.sid_filter = kwargs.get("filter", None) + self.start = pd.Timestamp(np.datetime64(kwargs.get('start'))) + self.start = self.start.tz_localize('utc') + self.end = pd.Timestamp(np.datetime64(kwargs.get('end'))) + self.end = self.end.tz_localize('utc') + start_time_str = kwargs.get("start_time", "9:30") + end_time_str = kwargs.get("end_time", "16:00") + self.start_time = dateutil.parser.parse(start_time_str).time() + self.end_time = dateutil.parser.parse(end_time_str).time() + self._raw_data = None + self.arg_string = hash_args(data, **kwargs) + self.root_node = "/" + kwargs.get('root', "TD") + "/" + + @property + def instance_hash(self): + return self.arg_string + + def raw_data_gen(self): + for date_node in self.data.walkNodes(self.root_node): + if isinstance(date_node, tables.group.Group): + continue + date = dateutil.parser.parse(date_node.name.split("_")[1]) + dt64 = np.datetime64(date) + table_dt = pd.Timestamp(dt64).tz_localize("utc") + if table_dt < self.start or table_dt > self.end: + continue + start_ts = pd.Timestamp(datetime.datetime.combine(table_dt.date(), + self.start_time), + tz=self.tz_in) + start_ts = start_ts.tz_convert("utc") + end_ts = pd.Timestamp(datetime.datetime.combine(table_dt.date(), + self.end_time), + tz=self.tz_in) + end_ts = end_ts.tz_convert("utc") + for item in _iterate_ohlc(date_node, self.sids, self.sid_filter, + start_ts, end_ts): + yield item + + @property + def raw_data(self): + if not self._raw_data: + self._raw_data = self.raw_data_gen() + return self._raw_data + + @property + def mapping(self): + return { + 'sid': (lambda x: x, 'sid'), + 'dt': (lambda x: x, 'dt'), + 'open': (lambda x: x, 'open'), + 'high': (lambda x: x, 'high'), + 'low': (lambda x: x, 'low'), + 'close': (lambda x: x, 'close'), + 'price': (lambda x: x, 'price'), + 'volume': (lambda x: x, 'volume'), + 'vwap': (lambda x: x, 'vwap') + } + + +class DataSourceTablesSignal(DataSource): + def __init__(self, data, **kwargs): + assert isinstance(data, tables.file.File) + self.h5file = data + self.sids = kwargs.get('sids', None) + self.start = kwargs.get('start') + self.end = kwargs.get('end') + self.source_id = kwargs.get("source_id", None) + self.arg_string = hash_args(data, **kwargs) + self._raw_data = None + self.root_node = +"/" + kwargs.get('root', "signal") + "/" + + @property + def instance_hash(self): + return self.arg_string + + def raw_data_gen(self): + for date_node in self.data.walkNodes(self.root_node): + if isinstance(date_node, tables.group.Group): + continue + date = dateutil.parser.parse(date_node.name.split("_")[1]) + dt64 = np.datetime64(date) + table_dt = pd.Timestamp(dt64).tz_localize("utc") + if table_dt < self.start or table_dt > self.end: + continue + start_ts = pd.Timestamp(datetime.datetime.combine(table_dt.date(), + self.start_time), + tz=self.tz_in) + start_ts = start_ts.tz_convert("utc") + end_ts = pd.Timestamp(datetime.datetime.combine(table_dt.date(), + self.end_time), + tz=self.tz_in) + end_ts = end_ts.tz_convert("utc") + table = self.data.getNode(date_node) + for row in _iterate_signal(table, self.sids, self.sid_filter, + start_ts, end_ts): + yield row + + @property + def raw_data(self): + if not self._raw_data: + self._raw_data = self.raw_data_gen() + return self._raw_data + + @property + def mapping(self): + return { + 'sid': (lambda x: x, 'symbol'), + 'dt': (lambda x: x, 'dt'), + 'signal': (lambda x: x, 'signal'), + } diff --git a/zipline/utils/data_source_tables_gen.py b/zipline/utils/data_source_tables_gen.py new file mode 100644 index 00000000..f2383613 --- /dev/null +++ b/zipline/utils/data_source_tables_gen.py @@ -0,0 +1,194 @@ +#!/usr/bin/env python +import sys +import getopt +import traceback +import numpy as np +import pandas as pd +import datetime +import logging +import tables +import gzip +import glob +import os +import random +import csv +import time +FORMAT = "%(asctime)-15s -8s %(message)s" + +logging.basicConfig(format=FORMAT, level=logging.INFO) + + +class Usage(Exception): + def __init__(self, msg): + self.msg = msg + + +OHLCTableDescription = {'sid': tables.StringCol(14, pos=2), + 'dt': tables.Int64Col(pos=1), + 'open': tables.Float64Col(dflt=np.NaN, pos=3), + 'high': tables.Float64Col(dflt=np.NaN, pos=4), + 'low': tables.Float64Col(dflt=np.NaN, pos=5), + 'close': tables.Float64Col(dflt=np.NaN, pos=6), + "volume": tables.Int64Col(dflt=0, pos=7)} + + +def process_line(line): + dt = np.datetime64(line["dt"]).astype(np.int64) + sid = line["sid"] + open_p = float(line["open"]) + high_p = float(line["high"]) + low_p = float(line["low"]) + close_p = float(line["close"]) + volume = int(line["volume"]) + return (dt, sid, open_p, high_p, low_p, close_p, volume) + + +def parse_csv(csv_reader): + previous_date = None + data = [] + dtype = [('dt', 'int64'), ('sid', '|S14'), ('open', float), + ('high', float), ('low', float), ('close', float), + ('volume', int)] + for line in csv_reader: + row = process_line(line) + current_date = line["dt"][:10].replace("-", "") + if previous_date and previous_date != current_date: + rows = np.array(data, dtype=dtype).view(np.recarray) + yield current_date, rows + data = [] + data.append(row) + previous_date = current_date + + +def merge_all_files_into_pytables(file_dir, file_out): + """ + process each file into pytables + """ + start = None + start = datetime.datetime.now() + out_h5 = tables.openFile(file_out, + mode="w", + title="bars", + filters=tables.Filters(complevel=9, + complib='zlib')) + table = None + for file_in in glob.glob(file_dir + "/*.gz"): + gzip_file = gzip.open(file_in) + expected_header = ["dt", "sid", "open", "high", "low", "close", + "volume"] + csv_reader = csv.DictReader(gzip_file) + header = csv_reader.fieldnames + if header != expected_header: + logging.warn("expected header %s\n" % (expected_header)) + logging.warn("header_found %s" % (header)) + return + + for current_date, rows in parse_csv(csv_reader): + table = out_h5.createTable("/TD", "date_" + current_date, + OHLCTableDescription, + expectedrows=len(rows), + createparents=True) + table.append(rows) + table.flush() + if not table is None: + table.flush() + end = datetime.datetime.now() + diff = (end - start).seconds + logging.debug("finished it took %d." % (diff)) + + +def create_fake_csv(file_in): + fields = ["dt", "sid", "open", "high", "low", "close", "volume"] + gzip_file = gzip.open(file_in, "w") + dict_writer = csv.DictWriter(gzip_file, fieldnames=fields) + current_dt = datetime.date.today() - datetime.timedelta(days=2) + current_dt = pd.Timestamp(current_dt).replace(hour=9) + current_dt = current_dt.replace(minute=30) + end_time = pd.Timestamp(datetime.date.today()) + end_time = end_time.replace(hour=16) + last_price = 10.0 + while current_dt < end_time: + row = {} + row["dt"] = current_dt + row["sid"] = "test" + last_price += random.randint(-20, 100) / 10000.0 + row["close"] = last_price + row["open"] = last_price - 0.01 + row["low"] = last_price - 0.02 + row["high"] = last_price + 0.02 + row["volume"] = random.randint(10, 1000) * 10 + dict_writer.writerow(row) + current_dt += datetime.timedelta(minutes=1) + if current_dt.hour > 16: + current_dt += datetime.timedelta(days=1) + current_dt = current_dt.replace(hour=9) + current_dt = current_dt.replace(minute=30) + gzip_file.close() + + +def main(argv=None): + """ + This script cleans minute bars into pytables file + data_source_tables_gen.py + [--tz_in] sets time zone of data only reasonably fast way to use + time.tzset() + [--dir_in] iterates through directory provided of csv files in gzip form + in form: + dt, sid, open, high, low, close, volume + 2012-01-01T12:30:30,1234HT,1, 2,3,4.0 + [--fake_csv] creates a fake sample csv to iterate through + [--file_out] determines output file + """ + if argv is None: + argv = sys.argv + try: + dir_in = None + file_out = "./all.h5" + fake_csv = None + try: + opts, args = getopt.getopt(argv[1:], "hdft", + ["help", + "dir_in=", + "debug", + "tz_in=", + "fake_csv=", + "file_out="]) + except getopt.error, msg: + raise Usage(msg) + for opt, value in opts: + if opt in ("--help", "-h"): + print main.__doc__ + if opt in ("-d", "--debug"): + logging.basicConfig(format=FORMAT, + level=logging.DEBUG) + if opt in ("-d", "--dir_in"): + dir_in = value + if opt in ("-o", "--file_out"): + file_out = value + if opt in ("--fake_csv"): + fake_csv = value + if opt in ("--tz_in"): + os.environ['TZ'] = value + time.tzset() + try: + if dir_in: + merge_all_files_into_pytables(dir_in, file_out) + if fake_csv: + create_fake_csv(fake_csv) + except Exception: + error = "An unhandled error occured in the" + error += "data_source_tables_gen.py script." + error += "\n\nTraceback:\n" + error += '-' * 70 + "\n" + error += "".join(traceback.format_tb(sys.exc_info()[2])) + error += repr(sys.exc_info()[1]) + "\n" + error += str(sys.exc_info()[1]) + "\n" + error += '-' * 70 + "\n" + print error + except Usage, err: + print >>sys.stderr, err.msg + print >>sys.stderr, "for help use --help" + return 2 + +if __name__ == "__main__": + sys.exit(main())