def init_db(self, txn=None):
"""Connect to database and create tables.
Parameters
----------
txn : sa.engine.Connection, optional
The transaction to execute in. If this is not provided, a new
transaction will be started with the engine provided.
Returns
-------
metadata : sa.MetaData
The metadata that describes the new assets db.
"""
with ExitStack() as stack:
if txn is None:
txn = stack.enter_context(self.engine.begin())
tables_already_exist = self._all_tables_present(txn)
# Create the SQL tables if they do not already exist.
metadata.create_all(txn, checkfirst=True)
if tables_already_exist:
check_version_info(txn, version_info, ASSET_DB_VERSION)
else:
write_version_info(txn, version_info, ASSET_DB_VERSION)
def transform(self):
"""
Main generator work loop.
"""
algo = self.algo
metrics_tracker = algo.metrics_tracker
emission_rate = metrics_tracker.emission_rate
def every_bar(
dt_to_use,
current_data=self.current_data,
handle_data=algo.event_manager.handle_data,
):
for capital_change in calculate_minute_capital_changes(dt_to_use):
yield capital_change
self.simulation_dt = dt_to_use
# called every tick (minute or day).
algo.on_dt_changed(dt_to_use)
blotter = algo.blotter
# handle any transactions and commissions coming out new orders
# placed in the last bar
new_transactions, new_commissions, closed_orders = blotter.get_transactions(
current_data)
blotter.prune_orders(closed_orders)
for transaction in new_transactions:
metrics_tracker.process_transaction(transaction)
# since this order was modified, record it
order = blotter.orders[transaction.order_id]
metrics_tracker.process_order(order)
for commission in new_commissions:
metrics_tracker.process_commission(commission)
handle_data(algo, current_data, dt_to_use)
# grab any new orders from the blotter, then clear the list.
# this includes cancelled orders.
new_orders = blotter.new_orders
blotter.new_orders = []
# if we have any new orders, record them so that we know
# in what perf period they were placed.
for new_order in new_orders:
metrics_tracker.process_order(new_order)
def once_a_day(midnight_dt,
current_data=self.current_data,
data_portal=self.data_portal):
# process any capital changes that came overnight
for capital_change in algo.calculate_capital_changes(
midnight_dt, emission_rate=emission_rate,
is_interday=True):
yield capital_change
# set all the timestamps
self.simulation_dt = midnight_dt
algo.on_dt_changed(midnight_dt)
metrics_tracker.handle_market_open(
midnight_dt,
algo.data_portal,
)
# handle any splits that impact any positions or any open orders.
assets_we_care_about = (metrics_tracker.positions.keys()
| algo.blotter.open_orders.keys())
if assets_we_care_about:
splits = data_portal.get_splits(assets_we_care_about,
midnight_dt)
if splits:
algo.blotter.process_splits(splits)
metrics_tracker.handle_splits(splits)
def on_exit():
# Remove references to algo, data portal, et al to break cycles
# and ensure deterministic cleanup of these objects when the
# simulation finishes.
self.algo = None
self.benchmark_source = self.current_data = self.data_portal = None
with ExitStack() as stack:
stack.callback(on_exit)
stack.enter_context(self.processor)
stack.enter_context(ZiplineAPI(self.algo))
if algo.data_frequency == "minute":
def execute_order_cancellation_policy():
algo.blotter.execute_cancel_policy(SESSION_END)
def calculate_minute_capital_changes(dt):
# process any capital changes that came between the last
# and current minutes
return algo.calculate_capital_changes(
dt, emission_rate=emission_rate, is_interday=False)
else:
def execute_order_cancellation_policy():
pass
def calculate_minute_capital_changes(dt):
return []
for dt, action in self.clock:
if action == BAR:
for capital_change_packet in every_bar(dt):
yield capital_change_packet
elif action == SESSION_START:
for capital_change_packet in once_a_day(dt):
yield capital_change_packet
elif action == SESSION_END:
# End of the session.
positions = metrics_tracker.positions
position_assets = algo.asset_finder.retrieve_all(positions)
self._cleanup_expired_assets(dt, position_assets)
execute_order_cancellation_policy()
algo.validate_account_controls()
yield self._get_daily_message(dt, algo, metrics_tracker)
elif action == BEFORE_TRADING_START_BAR:
self.simulation_dt = dt
algo.on_dt_changed(dt)
algo.before_trading_start(self.current_data)
elif action == MINUTE_END:
minute_msg = self._get_minute_message(
dt,
algo,
metrics_tracker,
)
yield minute_msg
risk_message = metrics_tracker.handle_simulation_end(
self.data_portal, )
yield risk_message
def ingest(name,
environ=os.environ,
timestamp=None,
assets_versions=(),
show_progress=False):
"""Ingest data for a given bundle.
Parameters
----------
name : str
The name of the bundle.
environ : mapping, optional
The environment variables. By default this is os.environ.
timestamp : datetime, optional
The timestamp to use for the load.
By default this is the current time.
assets_versions : Iterable[int], optional
Versions of the assets db to which to downgrade.
show_progress : bool, optional
Tell the ingest function to display the progress where possible.
"""
try:
bundle = bundles[name]
except KeyError:
raise UnknownBundle(name)
calendar = get_calendar(bundle.calendar_name)
start_session = bundle.start_session
end_session = bundle.end_session
if start_session is None or start_session < calendar.first_session:
start_session = calendar.first_session
if end_session is None or end_session > calendar.last_session:
end_session = calendar.last_session
if timestamp is None:
timestamp = pd.Timestamp.utcnow()
timestamp = timestamp.tz_convert('utc').tz_localize(None)
timestr = to_bundle_ingest_dirname(timestamp)
cachepath = cache_path(name, environ=environ)
pth.ensure_directory(pth.data_path([name, timestr], environ=environ))
pth.ensure_directory(cachepath)
with dataframe_cache(cachepath, clean_on_failure=False) as cache, \
ExitStack() as stack:
# we use `cleanup_on_failure=False` so that we don't purge the
# cache directory if the load fails in the middle
if bundle.create_writers:
wd = stack.enter_context(
working_dir(pth.data_path([], environ=environ)))
daily_bars_path = wd.ensure_dir(
*daily_equity_relative(name, timestr))
daily_bar_writer = BcolzDailyBarWriter(
daily_bars_path,
calendar,
start_session,
end_session,
)
# Do an empty write to ensure that the daily ctables exist
# when we create the SQLiteAdjustmentWriter below. The
# SQLiteAdjustmentWriter needs to open the daily ctables so
# that it can compute the adjustment ratios for the dividends.
daily_bar_writer.write(())
minute_bar_writer = BcolzMinuteBarWriter(
wd.ensure_dir(*minute_equity_relative(name, timestr)),
calendar,
start_session,
end_session,
minutes_per_day=bundle.minutes_per_day,
)
assets_db_path = wd.getpath(*asset_db_relative(name, timestr))
asset_db_writer = AssetDBWriter(assets_db_path)
adjustment_db_writer = stack.enter_context(
SQLiteAdjustmentWriter(
wd.getpath(*adjustment_db_relative(name, timestr)),
BcolzDailyBarReader(daily_bars_path),
overwrite=True,
))
else:
daily_bar_writer = None
minute_bar_writer = None
asset_db_writer = None
adjustment_db_writer = None
if assets_versions:
raise ValueError('Need to ingest a bundle that creates '
'writers in order to downgrade the assets'
' db.')
log.info("Ingesting {}.", name)
bundle.ingest(
environ,
asset_db_writer,
minute_bar_writer,
daily_bar_writer,
adjustment_db_writer,
calendar,
start_session,
end_session,
cache,
show_progress,
pth.data_path([name, timestr], environ=environ),
)
for version in sorted(set(assets_versions), reverse=True):
version_path = wd.getpath(*asset_db_relative(
name,
timestr,
db_version=version,
))
with working_file(version_path) as wf:
shutil.copy2(assets_db_path, wf.path)
downgrade(wf.path, version)
def ctx(self, *args, **kwargs):
with ExitStack() as stack:
for hook in self._hooks:
sub_ctx = getattr(hook, method_name)(*args, **kwargs)
stack.enter_context(sub_ctx)
yield stack