diff --git a/tests/pipeline/test_pipeline_algo.py b/tests/pipeline/test_pipeline_algo.py index 40fef5f5..a9115322 100644 --- a/tests/pipeline/test_pipeline_algo.py +++ b/tests/pipeline/test_pipeline_algo.py @@ -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) diff --git a/tests/pipeline/test_us_equity_pricing_loader.py b/tests/pipeline/test_us_equity_pricing_loader.py index 7a768d2b..a10ea0fb 100644 --- a/tests/pipeline/test_us_equity_pricing_loader.py +++ b/tests/pipeline/test_us_equity_pricing_loader.py @@ -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) diff --git a/zipline/data/us_equity_pricing.py b/zipline/data/us_equity_pricing.py index 9fb4d32c..59ab952b 100644 --- a/zipline/data/us_equity_pricing.py +++ b/zipline/data/us_equity_pricing.py @@ -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() diff --git a/zipline/pipeline/loaders/synthetic.py b/zipline/pipeline/loaders/synthetic.py index 2aff90de..513b0ec8 100644 --- a/zipline/pipeline/loaders/synthetic.py +++ b/zipline/pipeline/loaders/synthetic.py @@ -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)