diff --git a/zipline/algorithm.py b/zipline/algorithm.py index c55fc38d..736b0425 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -41,6 +41,8 @@ from zipline.gens.composites import ( from zipline.gens.tradesimulation import TradeSimulationClient as tsc from zipline import MESSAGES +DEFAULT_CAPITAL_BASE = float("1.0e5") + class TradingAlgorithm(object): """Base class for trading algorithms. Inherit and overload @@ -82,6 +84,9 @@ class TradingAlgorithm(object): self.slippage = VolumeShareSlippage() self.commission = PerShare() + # set the capital base + self.capital_base = kwargs.get('capital_base', DEFAULT_CAPITAL_BASE) + # an algorithm subclass needs to set initialized to True # when it is fully initialized. self.initialized = False @@ -175,7 +180,11 @@ class TradingAlgorithm(object): self.transforms.append(sf) - environment = create_trading_environment(start=start, end=end) + environment = create_trading_environment( + start=start, + end=end, + capital_base=self.capital_base + ) # create transforms and zipline self.gen = self._create_generator(environment) diff --git a/zipline/utils/factory.py b/zipline/utils/factory.py index ac2e94b2..6d4ade15 100644 --- a/zipline/utils/factory.py +++ b/zipline/utils/factory.py @@ -92,7 +92,8 @@ Fetching data from data.treasury.gov return bm_returns, tr_curves -def create_trading_environment(year=2006, start=None, end=None): +def create_trading_environment(year=2006, start=None, end=None, + capital_base=float("1.0e5")): """Construct a complete environment with reasonable defaults""" benchmark_returns, treasury_curves = load_market_data() @@ -106,7 +107,7 @@ def create_trading_environment(year=2006, start=None, end=None): treasury_curves, period_start=start, period_end=end, - capital_base=100000.0 + capital_base=capital_base ) return trading_environment