ENH: Adds a perf tracker method to handle SIDs leaving the universe

This commit is contained in:
jfkirk
2015-07-09 13:28:40 -04:00
parent 7cde3939bf
commit efa6d8dbce
2 changed files with 56 additions and 0 deletions
+40
View File
@@ -2020,6 +2020,46 @@ class TestPerformanceTracker(unittest.TestCase):
# Test not-owned SID
self.assertIsNone(txn)
def test_handle_sid_removed_from_universe(self):
# post some trades in the market
sim_params, _, _ = create_random_simulation_parameters()
events = factory.create_trade_history(
1,
[10, 10, 10, 10, 10],
[100, 100, 100, 100, 100],
oneday,
sim_params
)
# Create a tracker and a dividend
perf_tracker = perf.PerformanceTracker(sim_params)
dividend = factory.create_dividend(
1,
10.00,
# declared date, when the algorithm finds out about
# the dividend
events[0].dt,
# ex_date, the date before which the algorithm must hold stock
# to receive the dividend
events[1].dt,
# pay date, when the algorithm receives the dividend.
events[2].dt
)
dividend_frame = pd.DataFrame(
[dividend.to_series(index=zp.DIVIDEND_FIELDS)],
)
perf_tracker.update_dividends(dividend_frame)
# Ensure that the dividend is in the tracker
self.assertIn(1, perf_tracker.dividend_frame['sid'].values)
# Inform the tracker that sid 1 has been removed from the universe
perf_tracker.handle_sid_removed_from_universe(1)
# Ensure that the dividend for sid 1 has been removed from dividend
# frame
self.assertNotIn(1, perf_tracker.dividend_frame['sid'].values)
def test_serialization(self):
start_dt = datetime(year=2008,
month=10,
+16
View File
@@ -236,6 +236,22 @@ class PerformanceTracker(object):
self.dividend_frame = other.dividend_frame
self._dividend_count = other._dividend_count
def handle_sid_removed_from_universe(self, sid):
"""
This method handles any behaviors that must occur when a SID leaves the
universe of the TradingAlgorithm.
Parameters
__________
sid : int
The sid of the Asset being removed from the universe.
"""
# Drop any dividends for the sid from the dividends frame
self.dividend_frame = self.dividend_frame[
self.dividend_frame.sid != sid
]
def update_performance(self):
# calculate performance as of last trade
for perf_period in self.perf_periods: