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.
This commit is contained in:
Michael Schatzow
2013-12-01 11:56:31 -05:00
committed by twiecki
parent 07e25ae018
commit 59bcd097d5
3 changed files with 586 additions and 0 deletions
+172
View File
@@ -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'),
}
+220
View File
@@ -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'),
}
+194
View File
@@ -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())