wassname/catalyst
1
0
Fork 0
mirror of https://github.com/wassname/catalyst.git synced 2026-07-09 21:49:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

ENH: Write dividend payouts to adjustments db.

Browse Source 
To prepare for querying for payouts from SQLite, write the dividend
payouts to a new table `dividend_payouts`.

Change the expected columns of the passed dividend frame to contain the
payout data, and use that data to calculate the ratios (this moves
internal code that was calcualting the ratios into Zipline.)

The end result is that instead of just a `dividends` table with the
backward looking adjustment ratios, also write a `dividend_payouts`
table and a `stock_dividend_payout` table.
This commit is contained in:
Eddie Hebert
2015-10-10 06:11:51 -04:00
parent 752a2c3962
commit ccdc815526
4 changed files with 391 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/pipeline/test_pipeline_algo.py
+21 -2
View File
@@ -13,7 +13,9 @@ from numpy import (
array,
arange,
full_like,
float64,
nan,
uint32,
)
from numpy.testing import assert_almost_equal
from pandas import (
@@ -304,6 +306,14 @@ class ClosesOnly(TestCase):
algo.run(source=self.closes.iloc[10:17])
class MockDailyBarSpotReader(object):
"""
A BcolzDailyBarReader which returns a constant value for spot price.
"""
def spot_price(self, sid, day, column):
return 100.0
class PipelineAlgorithmTestCase(TestCase):
@classmethod
@@ -364,7 +374,8 @@ class PipelineAlgorithmTestCase(TestCase):
@classmethod
def create_adjustment_reader(cls, tempdir):
dbpath = tempdir.getpath('adjustments.sqlite')
writer = SQLiteAdjustmentWriter(dbpath)
writer = SQLiteAdjustmentWriter(dbpath, cls.env.trading_days,
MockDailyBarSpotReader())
splits = DataFrame.from_records([
{
'effective_date': str_to_seconds('2014-06-09'),
@@ -372,7 +383,7 @@ class PipelineAlgorithmTestCase(TestCase):
'sid': cls.AAPL,
}
])
mergers = dividends = DataFrame(
mergers = DataFrame(
{
# Hackery to make the dtypes correct on an empty frame.
'effective_date': array([], dtype=int),
@@ -382,6 +393,14 @@ class PipelineAlgorithmTestCase(TestCase):
index=DatetimeIndex([], tz='UTC'),
columns=['effective_date', 'ratio', 'sid'],
)
dividends = DataFrame({
'sid': array([], dtype=uint32),
'amount': array([], dtype=float64),
'record_date': array([], dtype='datetime64[ns]'),
'ex_date': array([], dtype='datetime64[ns]'),
'declared_date': array([], dtype='datetime64[ns]'),
'pay_date': array([], dtype='datetime64[ns]'),
})
writer.write(splits, mergers, dividends)
return SQLiteAdjustmentReader(dbpath)
tests/pipeline/test_us_equity_pricing_loader.py
+94 -15
View File
@@ -172,51 +172,114 @@ MERGERS = DataFrame(
DIVIDENDS = DataFrame(
[
# Before query range, should be excluded.
{'declared_date': Timestamp('2015-05-01', tz='UTC').to_datetime64(),
'ex_date': Timestamp('2015-06-01', tz='UTC').to_datetime64(),
'record_date': Timestamp('2015-06-03', tz='UTC').to_datetime64(),
'pay_date': Timestamp('2015-06-05', tz='UTC').to_datetime64(),
'amount': 90.0,
'sid': 1},
# First day of query range, should be excluded.
{'declared_date': Timestamp('2015-06-01', tz='UTC').to_datetime64(),
'ex_date': Timestamp('2015-06-10', tz='UTC').to_datetime64(),
'record_date': Timestamp('2015-06-15', tz='UTC').to_datetime64(),
'pay_date': Timestamp('2015-06-17', tz='UTC').to_datetime64(),
'amount': 80.0,
'sid': 3},
# Third day of query range, should have last_row of 2
{'declared_date': Timestamp('2015-06-01', tz='UTC').to_datetime64(),
'ex_date': Timestamp('2015-06-12', tz='UTC').to_datetime64(),
'record_date': Timestamp('2015-06-15', tz='UTC').to_datetime64(),
'pay_date': Timestamp('2015-06-17', tz='UTC').to_datetime64(),
'amount': 70.0,
'sid': 3},
# After query range, should be excluded.
{'declared_date': Timestamp('2015-06-01', tz='UTC').to_datetime64(),
'ex_date': Timestamp('2015-06-25', tz='UTC').to_datetime64(),
'record_date': Timestamp('2015-06-28', tz='UTC').to_datetime64(),
'pay_date': Timestamp('2015-06-30', tz='UTC').to_datetime64(),
'amount': 60.0,
'sid': 6},
# Another action in query range, should have last_row of 3
{'declared_date': Timestamp('2015-06-01', tz='UTC').to_datetime64(),
'ex_date': Timestamp('2015-06-15', tz='UTC').to_datetime64(),
'record_date': Timestamp('2015-06-18', tz='UTC').to_datetime64(),
'pay_date': Timestamp('2015-06-20', tz='UTC').to_datetime64(),
'amount': 50.0,
'sid': 3},
# Last day of range. Should have last_row of 7
{'declared_date': Timestamp('2015-06-01', tz='UTC').to_datetime64(),
'ex_date': Timestamp('2015-06-19', tz='UTC').to_datetime64(),
'record_date': Timestamp('2015-06-22', tz='UTC').to_datetime64(),
'pay_date': Timestamp('2015-06-30', tz='UTC').to_datetime64(),
'amount': 40.0,
'sid': 3},
],
columns=['declared_date',
'ex_date',
'record_date',
'pay_date',
'amount',
'sid'],
)
DIVIDENDS_EXPECTED = DataFrame(
[
# Before query range, should be excluded.
{'effective_date': str_to_seconds('2015-06-01'),
'ratio': 1.301,
'ratio': 0.1,
'sid': 1},
# First day of query range, should be excluded.
{'effective_date': str_to_seconds('2015-06-10'),
'ratio': 3.310,
'ratio': 0.20,
'sid': 3},
# Third day of query range, should have last_row of 2
{'effective_date': str_to_seconds('2015-06-12'),
'ratio': 3.312,
'ratio': 0.30,
'sid': 3},
# After query range, should be excluded.
{'effective_date': str_to_seconds('2015-06-25'),
'ratio': 6.325,
'ratio': 0.40,
'sid': 6},
# Another action in query range, should have last_row of 3
{'effective_date': str_to_seconds('2015-06-15'),
'ratio': 3.315,
'ratio': 0.50,
'sid': 3},
# Last day of range. Should have last_row of 7
{'effective_date': str_to_seconds('2015-06-19'),
'ratio': 3.319,
'ratio': 0.60,
'sid': 3},
],
columns=['effective_date', 'ratio', 'sid'],
)
class MockDailyBarSpotReader(object):
"""
A BcolzDailyBarReader which returns a constant value for spot price.
"""
def spot_price(self, sid, day, column):
return 100.0
class USEquityPricingLoaderTestCase(TestCase):
@classmethod
def setUpClass(cls):
cls.test_data_dir = TempDirectory()
cls.db_path = cls.test_data_dir.getpath('adjustments.db')
writer = SQLiteAdjustmentWriter(cls.db_path)
writer.write(SPLITS, MERGERS, DIVIDENDS)
cls.assets = TEST_QUERY_ASSETS
all_days = TradingEnvironment().trading_days
cls.calendar_days = all_days[
all_days.slice_indexer(TEST_CALENDAR_START, TEST_CALENDAR_STOP)
]
daily_bar_reader = MockDailyBarSpotReader()
writer = SQLiteAdjustmentWriter(cls.db_path, cls.calendar_days,
daily_bar_reader)
writer.write(SPLITS, MERGERS, DIVIDENDS)
cls.assets = TEST_QUERY_ASSETS
cls.asset_info = EQUITY_INFO
cls.bcolz_writer = SyntheticDailyBarWriter(
cls.asset_info,
@@ -232,7 +295,7 @@ class USEquityPricingLoaderTestCase(TestCase):
def test_input_sanity(self):
# Ensure that the input data doesn't contain adjustments during periods
# where the corresponding asset didn't exist.
for table in SPLITS, MERGERS, DIVIDENDS:
for table in SPLITS, MERGERS:
for eff_date_secs, _, sid in table.itertuples(index=False):
eff_date = Timestamp(eff_date_secs, unit='s')
asset_start, asset_end = EQUITY_INFO.ix[
@@ -256,7 +319,7 @@ class USEquityPricingLoaderTestCase(TestCase):
query_days = self.calendar_days_between(start_date, end_date)
start_loc = query_days.get_loc(start_date)
for table in SPLITS, MERGERS, DIVIDENDS:
for table in SPLITS, MERGERS, DIVIDENDS_EXPECTED:
for eff_date_secs, ratio, sid in table.itertuples(index=False):
eff_date = Timestamp(eff_date_secs, unit='s', tz='UTC')
@@ -309,8 +372,23 @@ class USEquityPricingLoaderTestCase(TestCase):
expected_close_adjustments, expected_volume_adjustments = \
self.expected_adjustments(TEST_QUERY_START, TEST_QUERY_STOP)
self.assertEqual(close_adjustments, expected_close_adjustments)
self.assertEqual(volume_adjustments, expected_volume_adjustments)
for key in expected_close_adjustments:
close_adjustment = close_adjustments[key]
for j, adj in enumerate(close_adjustment):
expected = expected_close_adjustments[key][j]
self.assertEqual(adj.first_row, expected.first_row)
self.assertEqual(adj.last_row, expected.last_row)
self.assertEqual(adj.col, expected.col)
assert_allclose(adj.value, expected.value)
for key in expected_volume_adjustments:
volume_adjustment = volume_adjustments[key]
for j, adj in enumerate(volume_adjustment):
expected = expected_volume_adjustments[key][j]
self.assertEqual(adj.first_row, expected.first_row)
self.assertEqual(adj.last_row, expected.last_row)
self.assertEqual(adj.col, expected.col)
assert_allclose(adj.value, expected.value)
def test_read_no_adjustments(self):
adjustment_reader = NullAdjustmentReader()
@@ -447,7 +525,8 @@ class USEquityPricingLoaderTestCase(TestCase):
self.assets,
baseline,
# Apply all adjustments.
concat([SPLITS, MERGERS, DIVIDENDS], ignore_index=True),
concat([SPLITS, MERGERS, DIVIDENDS_EXPECTED],
ignore_index=True),
)
assert_allclose(expected_adjusted_highs, window)
zipline/data/us_equity_pricing.py
+266 -5
View File
@@ -33,9 +33,11 @@ from numpy import (
iinfo,
integer,
issubdtype,
nan,
uint32,
)
from pandas import (
DataFrame,
DatetimeIndex,
read_csv,
Timestamp,
@@ -49,6 +51,9 @@ from six import (
from ._equities import _compute_row_slices, _read_bcolz_data
from ._adjustments import load_adjustments_from_sqlite
import logbook
logger = logbook.Logger('UsEquityPricing')
OHLC = frozenset(['open', 'high', 'low', 'close'])
US_EQUITY_PRICING_BCOLZ_COLUMNS = [
'open', 'high', 'low', 'close', 'volume', 'day', 'id'
@@ -61,6 +66,41 @@ SQLITE_ADJUSTMENT_COLUMN_DTYPES = {
}
SQLITE_ADJUSTMENT_TABLENAMES = frozenset(['splits', 'dividends', 'mergers'])
SQLITE_DIVIDEND_PAYOUT_COLUMNS = frozenset(
['sid',
'ex_date',
'declared_date',
'pay_date',
'record_date',
'amount'])
SQLITE_DIVIDEND_PAYOUT_COLUMN_DTYPES = {
'sid': integer,
'ex_date': integer,
'declared_date': integer,
'record_date': integer,
'pay_date': integer,
'amount': float,
}
SQLITE_STOCK_DIVIDEND_PAYOUT_COLUMNS = frozenset(
['sid',
'ex_date',
'declared_date',
'record_date',
'pay_date',
'payment_sid',
'ratio'])
SQLITE_STOCK_DIVIDEND_PAYOUT_COLUMN_DTYPES = {
'sid': integer,
'ex_date': integer,
'declared_date': integer,
'record_date': integer,
'pay_date': integer,
'payment_sid': integer,
'ratio': float,
}
UINT32_MAX = iinfo(uint32).max
@@ -477,7 +517,8 @@ class SQLiteAdjustmentWriter(object):
SQLiteAdjustmentReader
"""
def __init__(self, conn_or_path, overwrite=False):
def __init__(self, conn_or_path, calendar, daily_bar_reader,
overwrite=False):
if isinstance(conn_or_path, sqlite3.Connection):
self.conn = conn_or_path
elif isinstance(conn_or_path, str):
@@ -491,6 +532,9 @@ class SQLiteAdjustmentWriter(object):
else:
raise TypeError("Unknown connection type %s" % type(conn_or_path))
self._daily_bar_reader = daily_bar_reader
self._calendar = calendar
def write_frame(self, tablename, frame):
if frozenset(frame.columns) != SQLITE_ADJUSTMENT_COLUMNS:
raise ValueError(
@@ -523,7 +567,167 @@ class SQLiteAdjustmentWriter(object):
)
return frame.to_sql(tablename, self.conn)
def write(self, splits, mergers, dividends):
def write_dividend_payouts(self, frame):
"""
Write dividend payout data to SQLite table `dividend_payouts`.
"""
if frozenset(frame.columns) != SQLITE_DIVIDEND_PAYOUT_COLUMNS:
raise ValueError(
"Unexpected frame columns:\n"
"Expected Columns: %s\n"
"Received Columns: %s" % (
sorted(SQLITE_DIVIDEND_PAYOUT_COLUMNS),
sorted(frame.columns.tolist()),
)
)
expected_dtypes = SQLITE_DIVIDEND_PAYOUT_COLUMN_DTYPES
actual_dtypes = frame.dtypes
for colname, expected in iteritems(expected_dtypes):
actual = actual_dtypes[colname]
if not issubdtype(actual, expected):
raise TypeError(
"Expected data of type {expected} for column '{colname}', "
"but got {actual}.".format(
expected=expected,
colname=colname,
actual=actual,
)
)
return frame.to_sql('dividend_payouts', self.conn)
def write_stock_dividend_payouts(self, frame):
if frozenset(frame.columns) != SQLITE_STOCK_DIVIDEND_PAYOUT_COLUMNS:
raise ValueError(
"Unexpected frame columns:\n"
"Expected Columns: %s\n"
"Received Columns: %s" % (
sorted(SQLITE_STOCK_DIVIDEND_PAYOUT_COLUMNS),
sorted(frame.columns.tolist()),
)
)
expected_dtypes = SQLITE_STOCK_DIVIDEND_PAYOUT_COLUMN_DTYPES
actual_dtypes = frame.dtypes
for colname, expected in iteritems(expected_dtypes):
actual = actual_dtypes[colname]
if not issubdtype(actual, expected):
raise TypeError(
"Expected data of type {expected} for column '{colname}', "
"but got {actual}.".format(
expected=expected,
colname=colname,
actual=actual,
)
)
return frame.to_sql('stock_dividend_payouts', self.conn)
def calc_dividend_ratios(self, dividends):
"""
Calculate the ratios to apply to equities when looking back at pricing
history so that the price is smoothed over the ex_date, when the market
adjusts to the change in equity value due to upcoming dividend.
Returns
-------
DataFrame
A frame in the same format as splits and mergers, with keys
- sid, the id of the equity
- effective_date, the date in seconds on which to apply the ratio.
- ratio, the ratio to apply to backwards looking pricing data.
"""
ex_dates = dividends.ex_date.values
sids = dividends.sid.values
amounts = dividends.amount.values
ratios = full(len(amounts), nan)
daily_bar_reader = self._daily_bar_reader
calendar = self._calendar
for i, amount in enumerate(amounts):
sid = sids[i]
ex_date = ex_dates[i]
day_loc = calendar.get_loc(ex_date)
div_adj_date = calendar[day_loc - 1]
try:
prev_close = daily_bar_reader.spot_price(
sid, div_adj_date, 'close')
ratio = 1.0 - amount / (prev_close)
ratios[i] = ratio
except NoDataOnDate:
logger.warn("Couldn't compute ratio for dividend %s" % {
'sid': sid,
'ex_date': ex_date,
'amount': amount,
})
continue
effective_dates = ex_dates.astype('datetime64[s]').astype(uint32)
return DataFrame({
'sid': sids,
'effective_date': effective_dates,
'ratio': ratios,
})
def write_dividend_data(self, dividends, stock_dividends=None):
"""
Write both dividend payouts and the derived price adjustment ratios.
"""
# First write the dividend payouts.
dividend_payouts = dividends.copy()
dividend_payouts['ex_date'] = dividend_payouts['ex_date'].values.\
astype('datetime64[s]').astype(integer)
dividend_payouts['record_date'] = \
dividend_payouts['record_date'].values.astype('datetime64[s]').\
astype(integer)
dividend_payouts['declared_date'] = \
dividend_payouts['declared_date'].values.astype('datetime64[s]').\
astype(integer)
dividend_payouts['pay_date'] = \
dividend_payouts['pay_date'].values.astype('datetime64[s]').\
astype(integer)
self.write_dividend_payouts(dividend_payouts)
if stock_dividends is not None:
stock_dividend_payouts = stock_dividends.copy()
stock_dividend_payouts['ex_date'] = \
stock_dividend_payouts['ex_date'].values.\
astype('datetime64[s]').astype(integer)
stock_dividend_payouts['record_date'] = \
stock_dividend_payouts['record_date'].values.\
astype('datetime64[s]').astype(integer)
stock_dividend_payouts['declared_date'] = \
stock_dividend_payouts['declared_date'].\
values.astype('datetime64[s]').astype(integer)
stock_dividend_payouts['pay_date'] = \
stock_dividend_payouts['pay_date'].\
values.astype('datetime64[s]').astype(integer)
else:
stock_dividend_payouts = DataFrame({
'sid': array([], dtype=uint32),
'record_date': array([], dtype=uint32),
'ex_date': array([], dtype=uint32),
'declared_date': array([], dtype=uint32),
'pay_date': array([], dtype=uint32),
'payment_sid': array([], dtype=uint32),
'ratio': array([], dtype=float),
})
self.write_stock_dividend_payouts(stock_dividend_payouts)
# Second from the dividend payouts, calculate ratios.
dividend_ratios = self.calc_dividend_ratios(dividends)
self.write_frame('dividends', dividend_ratios)
def write(self, splits, mergers, dividends, stock_dividends=None):
"""
Writes data to a SQLite file to be read by SQLiteAdjustmentReader.
@@ -538,7 +742,7 @@ class SQLiteAdjustmentWriter(object):
Notes
-----
DataFrame input (`splits`, `mergers`, and `dividends`) should all have
DataFrame input (`splits`, `mergers`) should all have
the following columns:
effective_date : int
@@ -554,9 +758,50 @@ class SQLiteAdjustmentWriter(object):
'low', and 'close') by the ratio.
- For **splits only**, **divide** volume by the adjustment ratio.
Dividend ratios should be calculated as
DataFrame input, 'dividends' should have the following columns:
sid : int
The asset id associated with this adjustment.
ex_date : datetime64
The date on which an equity must be held to be eligible to receive
payment.
declared_date : datetime64
The date on which the dividend is announced to the public.
pay_date : datetime64
The date on which the dividend is distributed.
record_date : datetime64
The date on which the stock ownership is checked to determine
distribution of dividends.
amount : float
The cash amount paid for each share.
Dividend ratios are calculated as
1.0 - (dividend_value / "close on day prior to dividend ex_date").
DataFrame input, 'stock_dividends' should have the following columns:
sid : int
The asset id associated with this adjustment.
ex_date : datetime64
The date on which an equity must be held to be eligible to receive
payment.
declared_date : datetime64
The date on which the dividend is announced to the public.
pay_date : datetime64
The date on which the dividend is distributed.
record_date : datetime64
The date on which the stock ownership is checked to determine
distribution of dividends.
payment_sid : int
The asset id of the shares that should be paid instead of cash.
ratio: float
The ratio of currently held shares in the held sid that should
be paid with new shares of the payment_sid.
stock_dividends is optional.
Returns
-------
None
@@ -567,7 +812,7 @@ class SQLiteAdjustmentWriter(object):
"""
self.write_frame('splits', splits)
self.write_frame('mergers', mergers)
self.write_frame('dividends', dividends)
self.write_dividend_data(dividends, stock_dividends)
self.conn.execute(
"CREATE INDEX splits_sids "
"ON splits(sid)"
@@ -592,6 +837,22 @@ class SQLiteAdjustmentWriter(object):
"CREATE INDEX dividends_effective_date "
"ON dividends(effective_date)"
)
self.conn.execute(
"CREATE INDEX dividend_payouts_sid "
"ON dividend_payouts(sid)"
)
self.conn.execute(
"CREATE INDEX dividends_payouts_ex_date "
"ON dividend_payouts(ex_date)"
)
self.conn.execute(
"CREATE INDEX stock_dividend_payouts_sid "
"ON stock_dividend_payouts(sid)"
)
self.conn.execute(
"CREATE INDEX stock_dividends_payouts_ex_date "
"ON stock_dividend_payouts(ex_date)"
)
def close(self):
self.conn.close()
zipline/pipeline/loaders/synthetic.py
+10 -2
View File
@@ -253,11 +253,19 @@ class NullAdjustmentReader(SQLiteAdjustmentReader):
def __init__(self):
conn = sqlite3_connect(':memory:')
writer = SQLiteAdjustmentWriter(conn)
writer = SQLiteAdjustmentWriter(conn, None, None)
empty = DataFrame({
'sid': array([], dtype=uint32),
'effective_date': array([], dtype=uint32),
'ratio': array([], dtype=float),
})
writer.write(splits=empty, mergers=empty, dividends=empty)
empty_dividends = DataFrame({
'sid': array([], dtype=uint32),
'amount': array([], dtype=float64),
'record_date': array([], dtype='datetime64[ns]'),
'ex_date': array([], dtype='datetime64[ns]'),
'declared_date': array([], dtype='datetime64[ns]'),
'pay_date': array([], dtype='datetime64[ns]'),
})
writer.write(splits=empty, mergers=empty, dividends=empty_dividends)
super(NullAdjustmentReader, self).__init__(conn)