Fix DataSource id/get_id mixup.

This commit is contained in:
Stephen Diehl
2012-05-16 16:19:08 -04:00
parent 3ad1f250e6
commit 8d864462ff
8 changed files with 322 additions and 109 deletions
+2 -6
View File
@@ -9,10 +9,6 @@ LOGGER = logging.getLogger('ZiplineLogger')
class TestClient(Component):
def __init__(self):
Component.__init__(self)
self.init()
def init(self):
self.received_count = 0
self.prev_dt = None
@@ -48,10 +44,10 @@ class TestClient(Component):
def do_work(self):
socks = dict(self.poll.poll(self.heartbeat_timeout))
if self.control_in in socks and socks[self.control_in] == self.zmq.POLLIN:
if socks.get(self.control_in) == self.zmq.POLLIN:
msg = self.control_in.recv()
if self.data_feed in socks and socks[self.data_feed] == self.zmq.POLLIN:
if socks.get(self.data_feed) == self.zmq.POLLIN:
msg = self.data_feed.recv()
#logger.info('msg:' + str(msg))
+74 -77
View File
@@ -26,16 +26,16 @@ EXTENDED_TIMEOUT = 90
allocator = AddressAllocator(1000)
class FinanceTestCase(TestCase):
leased_sockets = defaultdict(list)
def setUp(self):
#qutil.configure_logging()
self.zipline_test_config = {
'allocator':allocator,
'sid':133
}
@timed(DEFAULT_TIMEOUT)
def test_factory_daily(self):
trading_environment = factory.create_trading_environment()
@@ -49,12 +49,12 @@ class FinanceTestCase(TestCase):
if prev:
self.assertTrue(trade.dt > prev.dt)
prev = trade
@timed(DEFAULT_TIMEOUT)
def test_trading_environment(self):
benchmark_returns, treasury_curves = \
factory.load_market_data()
env = TradingEnvironment(
benchmark_returns,
treasury_curves,
@@ -88,90 +88,90 @@ class FinanceTestCase(TestCase):
a_saturday,
a_sunday
]
for holiday in holidays:
self.assertTrue(not env.is_trading_day(holiday))
first_trading_day = datetime(2008, 1, 2, tzinfo = pytz.utc)
last_trading_day = datetime(2008, 12, 31, tzinfo = pytz.utc)
workdays = [first_trading_day, last_trading_day]
for workday in workdays:
self.assertTrue(env.is_trading_day(workday))
self.assertTrue(env.last_close.month == 12)
self.assertTrue(env.last_close.day == 31)
# The following two tests appear broken no that the order source is
# non blocking. HUNCH: The trades are streaming through before the orders
# are placed.
@timed(DEFAULT_TIMEOUT)
def test_orders(self):
# Simulation
# ----------
self.zipline_test_config['simulation_style'] = \
SIMULATION_STYLE.FIXED_SLIPPAGE
zipline = SimulatedTrading.create_test_zipline(
**self.zipline_test_config
)
zipline.simulate(blocking=True)
self.assertTrue(zipline.sim.ready())
self.assertFalse(zipline.sim.exception)
# TODO: Make more assertions about the final state of the components.
self.assertEqual(zipline.sim.feed.pending_messages(), 0, \
"The feed should be drained of all messages, found {n} remaining." \
.format(n=zipline.sim.feed.pending_messages()))
# the trading client should receive one transaction for every
# order placed.
self.assertEqual(
zipline.trading_client.txn_count,
zipline.trading_client.order_count
)
@timed(DEFAULT_TIMEOUT)
def test_aggressive_buying(self):
# Simulation
# ----------
# TODO: for some reason the orders aren't filled without an extra
# trade.
trade_count = 5
self.zipline_test_config['order_count'] = trade_count - 1
self.zipline_test_config['trade_count'] = trade_count
self.zipline_test_config['order_amount'] = 1
# tell the simulator to fill the orders in individual transactions
# matching the order volume exactly.
self.zipline_test_config['simulation_style'] = \
SIMULATION_STYLE.FIXED_SLIPPAGE
self.zipline_test_config['environment'] = factory.create_trading_environment()
sid_list = [self.zipline_test_config['sid']]
self.zipline_test_config['trade_source'] = factory.create_minutely_trade_source(
sid_list,
trade_count,
self.zipline_test_config['environment']
)
zipline = SimulatedTrading.create_test_zipline(**self.zipline_test_config)
zipline.simulate(blocking=True)
self.assertTrue(zipline.sim.ready())
self.assertFalse(zipline.sim.exception)
self.assertEqual(zipline.sim.feed.pending_messages(), 0, \
"The feed should be drained of all messages, found {n} remaining." \
.format(n=zipline.sim.feed.pending_messages()))
#
# the trading client should receive one transaction for every
# order placed.
@@ -179,9 +179,9 @@ class FinanceTestCase(TestCase):
zipline.trading_client.txn_count,
zipline.trading_client.order_count
)
@timed(DEFAULT_TIMEOUT)
def test_performance(self):
#provide enough trades to ensure all orders are filled.
@@ -189,27 +189,27 @@ class FinanceTestCase(TestCase):
self.zipline_test_config['trade_count'] = 200
zipline = SimulatedTrading.create_test_zipline(**self.zipline_test_config)
zipline.simulate(blocking=True)
self.assertEqual(
zipline.sim.feed.pending_messages(),
0,
"The feed should be drained of all messages, found {n} remaining." \
.format(n=zipline.sim.feed.pending_messages())
)
self.assertEqual(
zipline.sim.merge.pending_messages(),
0,
"The merge should be drained of all messages, found {n} remaining." \
.format(n=zipline.sim.merge.pending_messages())
)
self.assertEqual(
zipline.algorithm.count,
zipline.algorithm.incr,
"The test algorithm should send as many orders as specified.")
transaction_sim = zipline.trading_client.txn_sim
self.assertEqual(
transaction_sim.txn_count,
@@ -217,32 +217,32 @@ class FinanceTestCase(TestCase):
"The perf tracker should handle the same number of transactions \
as the simulator emits."
)
self.assertEqual(
len(zipline.get_positions()),
1,
"Portfolio should have one position."
)
SID = self.zipline_test_config['sid']
self.assertEqual(
zipline.get_positions()[SID]['sid'],
SID,
"Portfolio should have one position in " + str(SID)
)
self.assertEqual(
zipline.sources['flat'].count,
self.zipline_test_config['trade_count'],
"The simulated trade source should send all trades."
)
self.assertEqual(
zipline.algorithm.frame_count,
self.zipline_test_config['trade_count'],
"The algorithm should receive all trades."
)
@timed(DEFAULT_TIMEOUT)
def test_sid_filter(self):
"""Ensure the algorithm's filter prevents events from arriving."""
@@ -256,14 +256,14 @@ class FinanceTestCase(TestCase):
order_amount,
order_count
)
self.zipline_test_config['trade_count'] = 200
self.zipline_test_config['algorithm'] = test_algo
zipline = SimulatedTrading.create_test_zipline(
**self.zipline_test_config
)
zipline.simulate(blocking=True)
#check that the algorithm received no events
self.assertEqual(
@@ -271,14 +271,14 @@ class FinanceTestCase(TestCase):
test_algo.frame_count,
"The algorithm should not receive any events due to filtering."
)
# TODO: write tests for short sales
# TODO: write a test to do massive buying or shorting.
@timed(DEFAULT_TIMEOUT)
def test_partially_filled_orders(self):
# create a scenario where order size and trade size are equal
# so that orders must be spread out over several trades.
params ={
@@ -294,9 +294,9 @@ class FinanceTestCase(TestCase):
'expected_txn_count':8,
'expected_txn_volume':2 * 100
}
self.transaction_sim(**params)
# same scenario, but with short sales
params2 ={
'trade_count':360,
@@ -308,9 +308,9 @@ class FinanceTestCase(TestCase):
'expected_txn_count':8,
'expected_txn_volume':2 * -100
}
self.transaction_sim(**params2)
@timed(DEFAULT_TIMEOUT)
def test_collapsing_orders(self):
# create a scenario where order.amount <<< trade.volume
@@ -328,7 +328,7 @@ class FinanceTestCase(TestCase):
'expected_txn_volume':24 * 1
}
self.transaction_sim(**params1)
# second verse, same as the first. except short!
params2 ={
'trade_count':6,
@@ -341,7 +341,7 @@ class FinanceTestCase(TestCase):
'expected_txn_volume':24 * -1
}
self.transaction_sim(**params2)
@timed(DEFAULT_TIMEOUT)
def test_partial_expiration_orders(self):
# create a scenario where orders expire without being filled
@@ -360,7 +360,7 @@ class FinanceTestCase(TestCase):
'expected_txn_volume' : 25
}
self.transaction_sim(**params1)
# same scenario, but short sales.
params2 = {
'trade_count' : 100,
@@ -376,7 +376,7 @@ class FinanceTestCase(TestCase):
'expected_txn_volume' : -25
}
self.transaction_sim(**params2)
@timed(DEFAULT_TIMEOUT)
def test_alternating_long_short(self):
# create a scenario where we alternate buys and sells
@@ -393,9 +393,9 @@ class FinanceTestCase(TestCase):
'expected_txn_volume' : 0 #equal buys and sells
}
self.transaction_sim(**params1)
def transaction_sim(self, **params):
trade_count = params['trade_count']
trade_amount = params['trade_amount']
trade_interval = params['trade_interval']
@@ -411,14 +411,14 @@ class FinanceTestCase(TestCase):
alternate = params.get('alternate')
# if present, expect transaction amounts to match orders exactly.
complete_fill = params.get('complete_fill')
trading_environment = factory.create_trading_environment()
trade_sim = TransactionSimulator()
price = [10.1] * trade_count
volume = [100] * trade_count
start_date = trading_environment.first_open
sid = 1
generated_trades = factory.create_trade_history(
sid,
price,
@@ -426,12 +426,12 @@ class FinanceTestCase(TestCase):
trade_interval,
trading_environment
)
if alternate:
alternator = -1
else:
alternator = 1
order_date = start_date
for i in xrange(order_count):
order = ndict(
@@ -441,9 +441,9 @@ class FinanceTestCase(TestCase):
'type' : zp.DATASOURCE_TYPE.ORDER,
'dt' : order_date
})
trade_sim.add_open_order(order)
order_date = order_date + order_interval
# move after market orders to just after market next
# market open.
@@ -451,40 +451,40 @@ class FinanceTestCase(TestCase):
if order_date.minute >= 00:
order_date = order_date + timedelta(days=1)
order_date = order_date.replace(hour=14, minute=30)
# there should now be one open order list stored under the sid
oo = trade_sim.open_orders
self.assertEqual(len(oo), 1)
self.assertTrue(oo.has_key(sid))
order_list = oo[sid]
self.assertEqual(order_count, len(order_list))
for i in xrange(order_count):
order = order_list[i]
self.assertEqual(order.sid, sid)
self.assertEqual(order.amount, order_amount * alternator**i)
tracker = PerformanceTracker(trading_environment)
# this approximates the loop inside TradingSimulationClient
transactions = []
for trade in generated_trades:
if trade_delay:
trade.dt = trade.dt + trade_delay
txn = trade_sim.apply_trade_to_open_orders(trade)
if txn:
transactions.append(txn)
trade.TRANSACTION = txn
else:
trade.TRANSACTION = None
tracker.process_event(trade)
if complete_fill:
self.assertEqual(len(transactions), len(order_list))
total_volume = 0
for i in xrange(len(transactions)):
txn = transactions[i]
@@ -492,18 +492,15 @@ class FinanceTestCase(TestCase):
if complete_fill:
order = order_list[i]
self.assertEqual(order.amount, txn.amount)
self.assertEqual(total_volume, expected_txn_volume)
self.assertEqual(len(transactions), expected_txn_count)
cumulative_pos = tracker.cumulative_performance.positions[sid]
self.assertEqual(total_volume, cumulative_pos.amount)
# the open orders should now be empty
oo = trade_sim.open_orders
self.assertTrue(oo.has_key(sid))
order_list = oo[sid]
self.assertEqual(0, len(order_list))
+210
View File
@@ -0,0 +1,210 @@
"""
Component
|
Aggregate
|
/ \
Feed Merge
"""
import logging
from collections import Counter
import zipline.protocol as zp
from zipline.core.component import Component
from zipline.protocol import CONTROL_PROTOCOL, COMPONENT_TYPE, \
CONTROL_FRAME, CONTROL_UNFRAME
LOGGER = logging.getLogger('ZiplineLogger')
class Aggregate(Component):
"""
Abstract superclass to Merge & Feed. Acts on two sockets
- pull_socket
- feed_socket
Feed and Merge define these differently.
"""
def init(self):
self.sent_count = 0
self.received_count = 0
self.draining = False
self.ds_finished_counter = 0
# Depending on the size of this, might want to use a data
# structure with better asymptotics.
self.data_buffer = {}
# source_id -> integer count
self.sent_counters = Counter()
self.recv_counters = Counter()
@property
def get_id(self):
raise NotImplementedError
@property
def get_type(self):
return COMPONENT_TYPE.CONDUIT
# -------------
# Core Methods
# -------------
def open(self):
self.pull_socket = self.bind_data()
self.feed_socket = self.bind_feed()
def do_work(self):
# wait for synchronization reply from the host
socks = dict(self.poll.poll(self.heartbeat_timeout))
# TODO: Abstract this out, maybe on base component
if socks.get(self.control_in) == self.zmq.POLLIN:
msg = self.control_in.recv()
event, payload = CONTROL_UNFRAME(msg)
# -- Heartbeat --
if event == CONTROL_PROTOCOL.HEARTBEAT:
# Heart outgoing
heartbeat_frame = CONTROL_FRAME(
CONTROL_PROTOCOL.OK,
payload
)
self.control_out.send(heartbeat_frame)
# -- Soft Kill --
elif event == CONTROL_PROTOCOL.SHUTDOWN:
self.signal_done()
self.shutdown()
# -- Hard Kill --
elif event == CONTROL_PROTOCOL.KILL:
self.kill()
if socks.get(self.pull_socket) == self.zmq.POLLIN:
message = self.pull_socket.recv()
if message == str(CONTROL_PROTOCOL.DONE):
self.ds_finished_counter += 1
if len(self.data_buffer) == self.ds_finished_counter:
#drain any remaining messages in the buffer
LOGGER.debug("draining feed")
self.drain()
self.signal_done()
else:
try:
event = self.unframe(message)
# deserialization error
except zp.INVALID_DATASOURCE_FRAME as exc:
return self.signal_exception(exc)
try:
self.append(event)
self.send_next()
# Invalid message
except zp.INVALID_DATASOURCE_FRAME as exc:
return self.signal_exception(exc)
# -------------
# Flow Control
# -------------
def drain(self):
"""
Send all messages in the buffer.
"""
self.draining = True
while self.pending_messages() > 0:
self.send_next()
def send_next(self):
"""
Send the (chronologically) next message in the buffer.
"""
if not (self.is_full() or self.draining):
return
# TODO: implement this in __iter__
event = self.next()
if event is not None:
self.feed_socket.send(self.frame(event), self.zmq.NOBLOCK)
self.sent_counters[event.source_id] += 1
self.sent_count += 1
def append(self, event):
"""
Add an event to the buffer for the source specified by
source_id.
"""
self.data_buffer[event.source_id].append(event)
self.recv_counters[event.source_id] += 1
self.received_count += 1
def next(self):
"""
Get the next message in chronological order.
"""
if not(self.is_full() or self.draining):
return
cur_source = None
earliest_source = None
earliest_event = None
#iterate over the queues of events from all sources
#(1 queue per datasource)
for events in self.data_buffer.itervalues():
if len(events) == 0:
continue
cur_source = events
first_in_list = events[0]
if first_in_list.dt == None:
#this is a filler event, discard
events.pop(0)
continue
if (earliest_event == None) or (first_in_list.dt <= earliest_event.dt):
earliest_event = first_in_list
earliest_source = cur_source
if earliest_event != None:
return earliest_source.pop(0)
def is_full(self):
"""
Indicates whether the buffer has messages in buffer for
all un-DONE, blocking sources.
"""
for source_id, events in self.data_buffer.iteritems():
if len(events) == 0:
return False
return True
def pending_messages(self):
"""
Returns the count of all events from all sources in the
buffer.
"""
total = 0
for events in self.data_buffer.itervalues():
total += len(events)
return total
def add_source(self, source_id):
"""
Add a data source to the buffer.
"""
self.data_buffer[source_id] = []
def __len__(self):
"""
Buffer's length is same as internal map holding separate
sorted arrays of events keyed by source id.
"""
return len(self.data_buffer)
+15 -9
View File
@@ -12,9 +12,9 @@ LOGGER = logging.getLogger('ZiplineLogger')
class DataSource(Component):
"""
Baseclass for data sources. Subclass and implement send_all - usually this
means looping through all records in a store, converting to a dict, and
calling send(map).
Abstract baseclass for data sources. Subclass and implement send_all -
usually this means looping through all records in a store, converting
to a dict, and calling send(map).
Every datasource has a dict property to hold filters::
- key -- name of the filter, e.g. SID
@@ -22,14 +22,10 @@ class DataSource(Component):
Modify the datasource's filters via the set_filter(name, value)
"""
def __init__(self, source_id):
Component.__init__(self)
def init(self, source_id):
self.id = source_id
self.init()
self.filter = {}
def init(self):
self.cur_event = None
def set_filter(self, name, value):
@@ -37,7 +33,17 @@ class DataSource(Component):
@property
def get_id(self):
return self.id
"""
Returns this component id, this is fixed at a class
level. This should not and cannot be contingent on
arguments to the init function. Examples:
- "TradeDataSource"
- "RandomEquityTrades"
- "SpecificEquityTrades"
"""
raise NotImplementedError
@property
def get_type(self):
+2 -2
View File
@@ -49,7 +49,7 @@ class Feed(Component):
def do_work(self):
# wait for synchronization reply from the host
socks = dict(self.poll.poll(self.heartbeat_timeout))
socks = dict(self.poll.poll(self.heartbeat_timeout))
# TODO: Abstract this out, maybe on base component
if self.control_in in socks and socks[self.control_in] == self.zmq.POLLIN:
@@ -151,7 +151,7 @@ class Feed(Component):
cur_source = None
earliest_source = None
earliest_event = None
#iterate over the queues of events from all sources
#iterate over the queues of events from all sources
#(1 queue per datasource)
for events in self.data_buffer.values():
if len(events) == 0:
+1 -2
View File
@@ -61,7 +61,7 @@ class Merge(Feed):
def append(self, event):
"""
:param event: a ndict with one entry. key is the name of the
:param event: a ndict with one entry. key is the name of the
transform, value is the transformed value.
Add an event to the buffer for the source specified by
source_id.
@@ -69,4 +69,3 @@ class Merge(Feed):
self.data_buffer[event.keys()[0]].append(event)
self.received_count += 1
+9 -6
View File
@@ -13,6 +13,10 @@ import zipline.protocol as zp
class TradeDataSource(DataSource):
@property
def get_id(self):
return 'TradeDataSource'
def send(self, event):
"""
Sends the event iff it matches the internal SID filter.
@@ -48,7 +52,11 @@ class RandomEquityTrades(TradeDataSource):
self.day = datetime.timedelta(days=1)
self.price = random.uniform(5.0, 50.0)
@property
def get_id(self):
return 'RandomEquityTrades'
@property
def get_type(self):
zp.COMPONENT_TYPE.SOURCE
@@ -100,12 +108,7 @@ class SpecificEquityTrades(TradeDataSource):
@property
def get_id(self):
"""
The descriptive name of the component.
"""
# Prevents the bug that Thomas ran into
return "Unique ID"
return "SpecificEquityTrades"
def do_work(self):
if(len(self.event_list) == 0):
+9 -7
View File
@@ -191,7 +191,7 @@ class SimulatedTrading(object):
- algorithm - optional parameter providing an algorithm. defaults
to :py:class:`zipline.test.algorithms.TestAlgorithm`
- trade_source - optional parameter to specify trades, if present.
If not present :py:class:`ziplien.sources.SpecificEquityTrades`
If not present :py:class:`zipline.sources.SpecificEquityTrades`
is the source, with daily frequency in trades.
- simulation_style: optional parameter that configures the
:py:class:`zipline.finance.trading.TransactionSimulator`. Expects
@@ -264,11 +264,11 @@ class SimulatedTrading(object):
# Simulation
#-------------------
zipline = SimulatedTrading(**{
'algorithm':test_algo,
'trading_environment':trading_environment,
'allocator':allocator,
'simulator_class':simulator_class,
'simulation_style':simulation_style
'algorithm' : test_algo,
'trading_environment' : trading_environment,
'allocator' : allocator,
'simulator_class' : simulator_class,
'simulation_style' : simulation_style
})
#-------------------
@@ -285,7 +285,9 @@ class SimulatedTrading(object):
self.check_started()
source.set_filter('SID', self.algorithm.get_sid_filter())
self.sim.register_components([source])
self.sources[source.get_id] = source
# ``id`` is name of source_id, ``get_id`` is the class name
self.sources[source.id] = source
def add_transform(self, transform):
assert isinstance(transform, BaseTransform)