mirror of
https://github.com/wassname/catalyst.git
synced 2026-06-28 05:41:19 +08:00
fawce
379 lines
10 KiB
Python
379 lines
10 KiB
Python
import zmq
|
|
import pytz
|
|
from pprint import pformat as pf
|
|
from datetime import datetime, timedelta
|
|
|
|
from unittest2 import TestCase
|
|
from collections import defaultdict
|
|
from zipline.gens.composites import date_sorted_sources, merged_transforms
|
|
|
|
from zipline.core.devsimulator import AddressAllocator
|
|
from zipline.gens.transform import Passthrough, StatefulTransform
|
|
from zipline.gens.mavg import MovingAverage
|
|
from zipline.gens.tradesimulation import TradeSimulationClient as tsc
|
|
|
|
from zipline.utils.factory import create_trading_environment
|
|
from zipline.test_algorithms import TestAlgorithm
|
|
|
|
|
|
from zipline.utils.test_utils import (
|
|
setup_logger,
|
|
teardown_logger,
|
|
create_monitor,
|
|
launch_monitor
|
|
)
|
|
|
|
from zipline.core import Component
|
|
from zipline.protocol import (
|
|
DATASOURCE_FRAME,
|
|
DATASOURCE_UNFRAME,
|
|
FEED_FRAME,
|
|
FEED_UNFRAME,
|
|
MERGE_FRAME,
|
|
MERGE_UNFRAME,
|
|
SIMULATION_STYLE,
|
|
PERF_FRAME,
|
|
BT_UPDATE_UNFRAME
|
|
)
|
|
|
|
from zipline.gens.tradegens import SpecificEquityTrades
|
|
|
|
import logbook
|
|
log = logbook.Logger('ComponentTestCase')
|
|
|
|
allocator = AddressAllocator(1000)
|
|
|
|
|
|
class ComponentTestCase(TestCase):
|
|
|
|
leased_sockets = defaultdict(list)
|
|
|
|
def setUp(self):
|
|
self.zipline_test_config = {
|
|
'allocator' : allocator,
|
|
'sid' : 133,
|
|
'devel' : False,
|
|
'results_socket' : allocator.lease(1)[0],
|
|
'simulation_style' : SIMULATION_STYLE.FIXED_SLIPPAGE
|
|
}
|
|
self.ctx = zmq.Context()
|
|
setup_logger(self)
|
|
|
|
count = 250
|
|
filter = [2,3]
|
|
#Set up source a. One minute between events.
|
|
args_a = tuple()
|
|
kwargs_a = {
|
|
'count' : 2*count,
|
|
'sids' : [1,2,3],
|
|
'start' : datetime(2002,1,3,15, tzinfo = pytz.utc),
|
|
'delta' : timedelta(hours = 6),
|
|
'filter' : filter
|
|
}
|
|
self.source_a = SpecificEquityTrades(*args_a, **kwargs_a)
|
|
|
|
#Set up source b. Two minutes between events.
|
|
args_b = tuple()
|
|
kwargs_b = {
|
|
'count' : count,
|
|
'sids' : [2,3,4],
|
|
'start' : datetime(2002,1,3,14, tzinfo = pytz.utc),
|
|
'delta' : timedelta(minutes = 5),
|
|
'filter' : filter
|
|
}
|
|
self.source_b = SpecificEquityTrades(*args_b, **kwargs_b)
|
|
|
|
self.environment = create_trading_environment(year = 2002)
|
|
|
|
|
|
|
|
def tearDown(self):
|
|
teardown_logger(self)
|
|
|
|
def test_source(self):
|
|
monitor = create_monitor(allocator)
|
|
socket_uri = allocator.lease(1)[0]
|
|
count = 100
|
|
|
|
filter = [1,2,3,4]
|
|
#Set up source a. One minute between events.
|
|
args_a = tuple()
|
|
kwargs_a = {
|
|
'sids' : [1,2],
|
|
'start' : datetime(2012,6,6,0,tzinfo=pytz.utc),
|
|
'delta' : timedelta(minutes = 1),
|
|
'filter' : filter,
|
|
'count' : count
|
|
}
|
|
|
|
trade_gen = SpecificEquityTrades(*args_a, **kwargs_a)
|
|
|
|
|
|
comp_a = Component(
|
|
trade_gen,
|
|
monitor,
|
|
socket_uri,
|
|
DATASOURCE_FRAME,
|
|
DATASOURCE_UNFRAME,
|
|
"source_a"
|
|
)
|
|
|
|
launch_monitor(monitor)
|
|
|
|
for event in comp_a:
|
|
log.info(event)
|
|
|
|
# wait for the sending process to exit
|
|
comp_a.proc.join()
|
|
|
|
|
|
def test_sort(self):
|
|
monitor = create_monitor(allocator)
|
|
socket_uris = allocator.lease(3)
|
|
count = 100
|
|
|
|
filter = [1,2,3,4]
|
|
#Set up source a. One minute between events.
|
|
args_a = tuple()
|
|
kwargs_a = {
|
|
'sids' : [1,2],
|
|
'start' : datetime(2012,6,6,0,tzinfo=pytz.utc),
|
|
'delta' : timedelta(minutes = 1),
|
|
'filter' : filter,
|
|
'count' : count
|
|
}
|
|
trade_gen_a = SpecificEquityTrades(*args_a, **kwargs_a)
|
|
|
|
#Set up source b. Two minutes between events.
|
|
args_b = tuple()
|
|
kwargs_b = {
|
|
'sids' : [2],
|
|
'start' : datetime(2012,1,3,15, tzinfo = pytz.utc),
|
|
'delta' : timedelta(minutes = 1),
|
|
'filter' : filter,
|
|
'count' : count
|
|
}
|
|
trade_gen_b = SpecificEquityTrades(*args_b, **kwargs_b)
|
|
|
|
#Set up source c. Three minutes between events.
|
|
args_c = tuple()
|
|
kwargs_c = {
|
|
'sids' : [3],
|
|
'start' : datetime(2012,1,3,15, tzinfo = pytz.utc),
|
|
'delta' : timedelta(minutes = 1),
|
|
'filter' : filter,
|
|
'count' : count
|
|
}
|
|
|
|
trade_gen_c = SpecificEquityTrades(*args_c, **kwargs_c)
|
|
|
|
|
|
comp_a = Component(
|
|
trade_gen_a,
|
|
monitor,
|
|
socket_uris[0],
|
|
DATASOURCE_FRAME,
|
|
DATASOURCE_UNFRAME,
|
|
trade_gen_a.get_hash()
|
|
)
|
|
|
|
comp_b = Component(
|
|
trade_gen_b,
|
|
monitor,
|
|
socket_uris[1],
|
|
DATASOURCE_FRAME,
|
|
DATASOURCE_UNFRAME,
|
|
trade_gen_b.get_hash()
|
|
)
|
|
|
|
comp_c = Component(
|
|
trade_gen_c,
|
|
monitor,
|
|
socket_uris[2],
|
|
DATASOURCE_FRAME,
|
|
DATASOURCE_UNFRAME,
|
|
trade_gen_c.get_hash()
|
|
)
|
|
|
|
sources = [comp_a, comp_b, comp_c]
|
|
|
|
sorted_out = date_sorted_sources(*sources)
|
|
|
|
launch_monitor(monitor)
|
|
|
|
prev = None
|
|
sort_count = 0
|
|
for msg in sorted_out:
|
|
if prev:
|
|
self.assertTrue(msg.dt >= prev.dt, \
|
|
"Messages should be in date ascending order")
|
|
prev = msg
|
|
sort_count += 1
|
|
|
|
self.assertEqual(count*3, sort_count)
|
|
|
|
# wait for processes to finish
|
|
comp_a.proc.join()
|
|
comp_b.proc.join()
|
|
comp_c.proc.join()
|
|
|
|
|
|
def test_full(self):
|
|
monitor = create_monitor(allocator)
|
|
|
|
# ------------------------
|
|
# Run sources in dedicated processes
|
|
comp_a = Component(
|
|
self.source_a,
|
|
monitor,
|
|
allocator.lease(1)[0],
|
|
DATASOURCE_FRAME,
|
|
DATASOURCE_UNFRAME,
|
|
self.source_a.get_hash()
|
|
)
|
|
|
|
comp_b = Component(
|
|
self.source_b,
|
|
monitor,
|
|
allocator.lease(1)[0],
|
|
DATASOURCE_FRAME,
|
|
DATASOURCE_UNFRAME,
|
|
self.source_b.get_hash()
|
|
)
|
|
|
|
# Date sort the sources, and run the sort in a dedicated
|
|
# process
|
|
sources = [comp_a, comp_b]
|
|
|
|
sorted_out = date_sorted_sources(*sources)
|
|
|
|
#launch_monitor(monitor)
|
|
#import nose.tools; nose.tools.set_trace()
|
|
#for feed_msg in sorted_out:
|
|
# log.info(pf(feed_msg))
|
|
|
|
#return
|
|
|
|
sorted = Component(
|
|
sorted_out,
|
|
monitor,
|
|
allocator.lease(1)[0],
|
|
FEED_FRAME,
|
|
FEED_UNFRAME,
|
|
"sort"
|
|
)
|
|
|
|
|
|
passthrough = StatefulTransform(Passthrough)
|
|
mavg_price = StatefulTransform(
|
|
MovingAverage,
|
|
timedelta(minutes = 20),
|
|
['price']
|
|
)
|
|
|
|
merged_gen = merged_transforms(sorted, passthrough, mavg_price)
|
|
|
|
merged = Component(
|
|
merged_gen,
|
|
monitor,
|
|
allocator.lease(1)[0],
|
|
MERGE_FRAME,
|
|
MERGE_UNFRAME,
|
|
"merge"
|
|
)
|
|
|
|
algo = TestAlgorithm(2, 10, 100, sid_filter = [2,3])
|
|
|
|
style = SIMULATION_STYLE.FIXED_SLIPPAGE
|
|
|
|
trading_client = tsc(algo, self.environment, style)
|
|
tsc_gen = trading_client.simulate(merged)
|
|
|
|
tsc_comp = Component(
|
|
tsc_gen,
|
|
monitor,
|
|
allocator.lease(1)[0],
|
|
PERF_FRAME,
|
|
BT_UPDATE_UNFRAME,
|
|
"tsc"
|
|
)
|
|
|
|
launch_monitor(monitor)
|
|
for message in tsc_comp:
|
|
log.info(pf(message))
|
|
|
|
|
|
# wait for processes to finish
|
|
comp_a.proc.join()
|
|
comp_b.proc.join()
|
|
sorted.proc.join()
|
|
merged.proc.join()
|
|
return
|
|
|
|
def test_single_thread(self):
|
|
|
|
#Set up source c. Three minutes between events.
|
|
|
|
sorted = date_sorted_sources(self.source_a, self.source_b)
|
|
|
|
passthrough = StatefulTransform(Passthrough)
|
|
mavg_price = StatefulTransform(MovingAverage, timedelta(minutes = 20), ['price'])
|
|
|
|
merged = merged_transforms(sorted, passthrough, mavg_price)
|
|
|
|
algo = TestAlgorithm(2, 10, 100, sid_filter = [2,3])
|
|
style = SIMULATION_STYLE.FIXED_SLIPPAGE
|
|
|
|
trading_client = tsc(algo, self.environment, style)
|
|
for message in trading_client.simulate(merged):
|
|
log.info(pf(message))
|
|
|
|
def test_compound(self):
|
|
monitor = create_monitor(allocator)
|
|
|
|
sorted_out = date_sorted_sources(self.source_a, self.source_b)
|
|
|
|
sorted = Component(
|
|
sorted_out,
|
|
monitor,
|
|
allocator.lease(1)[0],
|
|
FEED_FRAME,
|
|
FEED_UNFRAME,
|
|
"feed"
|
|
)
|
|
|
|
passthrough = StatefulTransform(Passthrough)
|
|
mavg_price = StatefulTransform(
|
|
MovingAverage,
|
|
timedelta(minutes = 20),
|
|
['price']
|
|
)
|
|
|
|
merged_gen = merged_transforms(sorted, passthrough, mavg_price)
|
|
|
|
merged = Component(
|
|
merged_gen,
|
|
monitor,
|
|
allocator.lease(1)[0],
|
|
MERGE_FRAME,
|
|
MERGE_UNFRAME,
|
|
"merge"
|
|
)
|
|
|
|
algo = TestAlgorithm(2, 10, 100, sid_filter = [2,3])
|
|
style = SIMULATION_STYLE.FIXED_SLIPPAGE
|
|
|
|
trading_client = tsc(algo, self.environment, style)
|
|
tsc_gen = trading_client.simulate(merged)
|
|
|
|
|
|
launch_monitor(monitor)
|
|
for message in tsc_gen:
|
|
log.info(pf(message))
|
|
|
|
|
|
# wait for processes to finish
|
|
sorted.proc.join()
|
|
merged.proc.join()
|
|
return
|