mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-14 11:15:09 +08:00
Richard Frank
30847a10a7
but MultiColumnLoader was returning a list of lists of arrays in some cases.
268 lines
8.1 KiB
Python
268 lines
8.1 KiB
Python
"""
|
|
Synthetic data loaders for testing.
|
|
"""
|
|
from bcolz import ctable
|
|
|
|
from numpy import (
|
|
arange,
|
|
array,
|
|
float64,
|
|
full,
|
|
iinfo,
|
|
uint32,
|
|
)
|
|
from pandas import (
|
|
DataFrame,
|
|
Timestamp,
|
|
)
|
|
from sqlite3 import connect as sqlite3_connect
|
|
|
|
from six import iteritems
|
|
|
|
from zipline.data.ffc.base import FFCLoader
|
|
from zipline.data.ffc.frame import DataFrameFFCLoader
|
|
from zipline.data.ffc.loaders.us_equity_pricing import (
|
|
BcolzDailyBarWriter,
|
|
SQLiteAdjustmentReader,
|
|
SQLiteAdjustmentWriter,
|
|
US_EQUITY_PRICING_BCOLZ_COLUMNS,
|
|
)
|
|
|
|
|
|
UINT_32_MAX = iinfo(uint32).max
|
|
|
|
|
|
def nanos_to_seconds(nanos):
|
|
return nanos / (1000 * 1000 * 1000)
|
|
|
|
|
|
class MultiColumnLoader(FFCLoader):
|
|
"""
|
|
FFCLoader that can delegate to sub-loaders.
|
|
|
|
Parameters
|
|
----------
|
|
loaders : dict
|
|
Dictionary mapping columns -> loader
|
|
"""
|
|
def __init__(self, loaders):
|
|
self._loaders = loaders
|
|
|
|
def load_adjusted_array(self, columns, mask):
|
|
"""
|
|
Load by delegating to sub-loaders.
|
|
"""
|
|
out = []
|
|
for column in columns:
|
|
try:
|
|
loader = self._loaders[column]
|
|
except KeyError:
|
|
raise ValueError("Couldn't find loader for %s" % column)
|
|
out.extend(loader.load_adjusted_array([column], mask))
|
|
return out
|
|
|
|
|
|
class ConstantLoader(MultiColumnLoader):
|
|
"""
|
|
Synthetic FFCLoader that returns a constant value for each column.
|
|
|
|
Parameters
|
|
----------
|
|
constants : dict
|
|
Map from column to value(s) to use for that column.
|
|
Values can be anything that can be passed as the first positional
|
|
argument to a DataFrame of the same shape as `mask`.
|
|
mask : pandas.DataFrame
|
|
Mask indicating when assets existed.
|
|
Indices of this frame are used to align input queries.
|
|
|
|
Notes
|
|
-----
|
|
Adjustments are unsupported with ConstantLoader.
|
|
"""
|
|
def __init__(self, constants, dates, assets):
|
|
loaders = {}
|
|
for column, const in iteritems(constants):
|
|
frame = DataFrame(
|
|
const,
|
|
index=dates,
|
|
columns=assets,
|
|
dtype=column.dtype,
|
|
)
|
|
loaders[column] = DataFrameFFCLoader(
|
|
column=column,
|
|
baseline=frame,
|
|
adjustments=None,
|
|
)
|
|
|
|
super(ConstantLoader, self).__init__(loaders)
|
|
|
|
|
|
class SyntheticDailyBarWriter(BcolzDailyBarWriter):
|
|
"""
|
|
Bcolz writer that creates synthetic data based on asset lifetime metadata.
|
|
|
|
For a given asset/date/column combination, we generate a corresponding raw
|
|
value using the following formula for OHLCV columns:
|
|
|
|
data(asset, date, column) = (100,000 * asset_id)
|
|
+ (10,000 * column_num)
|
|
+ (date - Jan 1 2000).days # ~6000 for 2015
|
|
where:
|
|
column_num('open') = 0
|
|
column_num('high') = 1
|
|
column_num('low') = 2
|
|
column_num('close') = 3
|
|
column_num('volume') = 4
|
|
|
|
We use days since Jan 1, 2000 to guarantee that there are no collisions
|
|
while also the produced values smaller than UINT32_MAX / 1000.
|
|
|
|
For 'day' and 'id', we use the standard format expected by the base class.
|
|
|
|
Parameters
|
|
----------
|
|
asset_info : DataFrame
|
|
DataFrame with asset_id as index and 'start_date'/'end_date' columns.
|
|
calendar : DatetimeIndex
|
|
Calendar to use for constructing asset lifetimes.
|
|
"""
|
|
OHLCV = ('open', 'high', 'low', 'close', 'volume')
|
|
OHLC = ('open', 'high', 'low', 'close')
|
|
PSEUDO_EPOCH = Timestamp('2000-01-01', tz='UTC')
|
|
|
|
def __init__(self, asset_info, calendar):
|
|
super(SyntheticDailyBarWriter, self).__init__()
|
|
assert (
|
|
# Using .value here to avoid having to care about UTC-aware dates.
|
|
self.PSEUDO_EPOCH.value <
|
|
calendar.min().value <=
|
|
asset_info['start_date'].min().value
|
|
)
|
|
assert (asset_info['start_date'] < asset_info['end_date']).all()
|
|
self._asset_info = asset_info
|
|
self._calendar = calendar
|
|
|
|
def _raw_data_for_asset(self, asset_id):
|
|
"""
|
|
Generate 'raw' data that encodes information about the asset.
|
|
|
|
See class docstring for a description of the data format.
|
|
"""
|
|
# Get the dates for which this asset existed according to our asset
|
|
# info.
|
|
dates = self._calendar[
|
|
self._calendar.slice_indexer(
|
|
self.asset_start(asset_id), self.asset_end(asset_id)
|
|
)
|
|
]
|
|
|
|
data = full(
|
|
(len(dates), len(US_EQUITY_PRICING_BCOLZ_COLUMNS)),
|
|
asset_id * (100 * 1000),
|
|
dtype=uint32,
|
|
)
|
|
|
|
# Add 10,000 * column-index to OHLCV columns
|
|
data[:, :5] += arange(5) * (10 * 1000)
|
|
|
|
# Add days since Jan 1 2001 for OHLCV columns.
|
|
data[:, :5] += (dates - self.PSEUDO_EPOCH).days[:, None]
|
|
|
|
frame = DataFrame(
|
|
data,
|
|
index=dates,
|
|
columns=US_EQUITY_PRICING_BCOLZ_COLUMNS,
|
|
)
|
|
|
|
frame['day'] = nanos_to_seconds(dates.asi8)
|
|
frame['id'] = asset_id
|
|
|
|
return ctable.fromdataframe(frame)
|
|
|
|
def asset_start(self, asset):
|
|
ret = self._asset_info.loc[asset]['start_date']
|
|
if ret.tz is None:
|
|
ret = ret.tz_localize('UTC')
|
|
assert ret.tzname() == 'UTC', "Unexpected non-UTC timestamp"
|
|
return ret
|
|
|
|
def asset_end(self, asset):
|
|
ret = self._asset_info.loc[asset]['end_date']
|
|
if ret.tz is None:
|
|
ret = ret.tz_localize('UTC')
|
|
assert ret.tzname() == 'UTC', "Unexpected non-UTC timestamp"
|
|
return ret
|
|
|
|
@classmethod
|
|
def expected_value(cls, asset_id, date, colname):
|
|
"""
|
|
Check that the raw value for an asset/date/column triple is as
|
|
expected.
|
|
|
|
Used by tests to verify data written by a writer.
|
|
"""
|
|
from_asset = asset_id * 100 * 1000
|
|
from_colname = cls.OHLCV.index(colname) * (10 * 1000)
|
|
from_date = (date - cls.PSEUDO_EPOCH).days
|
|
return from_asset + from_colname + from_date
|
|
|
|
def expected_values_2d(self, dates, assets, colname):
|
|
"""
|
|
Return an 2D array containing cls.expected_value(asset_id, date,
|
|
colname) for each date/asset pair in the inputs.
|
|
|
|
Values before/after an assets lifetime are filled with 0 for volume and
|
|
NaN for price columns.
|
|
"""
|
|
if colname == 'volume':
|
|
dtype = uint32
|
|
missing = 0
|
|
else:
|
|
dtype = float64
|
|
missing = float('nan')
|
|
|
|
data = full((len(dates), len(assets)), missing, dtype=dtype)
|
|
for j, asset in enumerate(assets):
|
|
start, end = self.asset_start(asset), self.asset_end(asset)
|
|
for i, date in enumerate(dates):
|
|
# No value expected for dates outside the asset's start/end
|
|
# date.
|
|
if not (start <= date <= end):
|
|
continue
|
|
data[i, j] = self.expected_value(asset, date, colname)
|
|
return data
|
|
|
|
# BEGIN SUPERCLASS INTERFACE
|
|
def gen_tables(self, assets):
|
|
for asset in assets:
|
|
yield asset, self._raw_data_for_asset(asset)
|
|
|
|
def to_uint32(self, array, colname):
|
|
if colname in {'open', 'high', 'low', 'close'}:
|
|
# Data is stored as 1000 * raw value.
|
|
assert array.max() < (UINT_32_MAX / 1000), "Test data overflow!"
|
|
return array * 1000
|
|
else:
|
|
assert colname in ('volume', 'day'), "Unknown column: %s" % colname
|
|
return array
|
|
# END SUPERCLASS INTERFACE
|
|
|
|
|
|
class NullAdjustmentReader(SQLiteAdjustmentReader):
|
|
"""
|
|
A SQLiteAdjustmentReader that stores no adjustments and uses in-memory
|
|
SQLite.
|
|
"""
|
|
|
|
def __init__(self):
|
|
conn = sqlite3_connect(':memory:')
|
|
writer = SQLiteAdjustmentWriter(conn)
|
|
empty = DataFrame({
|
|
'sid': array([], dtype=uint32),
|
|
'effective_date': array([], dtype=uint32),
|
|
'ratio': array([], dtype=float),
|
|
})
|
|
writer.write(splits=empty, mergers=empty, dividends=empty)
|
|
super(NullAdjustmentReader, self).__init__(conn)
|