From 283c959cc4bc51a90ee8cd9e805004cad1cc9406 Mon Sep 17 00:00:00 2001 From: Stewart Douglas Date: Wed, 9 Sep 2015 13:33:05 -0400 Subject: [PATCH] MAINT: Move analyze methods into algorithm files --- tests/test_examples.py | 1 + zipline/examples/buyapple.py | 30 +++++---- zipline/examples/buyapple_analyze.py | 10 --- zipline/examples/dual_ema_talib.py | 61 +++++++++++++----- zipline/examples/dual_moving_average.py | 64 +++++++++++++------ .../examples/dual_moving_average_analyze.py | 24 ------- zipline/examples/olmar.py | 28 ++++++-- zipline/examples/pairtrade.py | 48 +++++++++----- zipline/utils/cli.py | 6 -- 9 files changed, 161 insertions(+), 111 deletions(-) delete mode 100644 zipline/examples/buyapple_analyze.py delete mode 100644 zipline/examples/dual_moving_average_analyze.py diff --git a/tests/test_examples.py b/tests/test_examples.py index 38474471..67d83975 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -39,6 +39,7 @@ def example_dir(): class ExamplesTests(TestCase): + # Test algorithms as if they being executed directly from the command line. @parameterized.expand(((os.path.basename(f).replace('.', '_'), f) for f in glob.glob(os.path.join(example_dir(), '*.py')))) def test_example(self, name, example): diff --git a/zipline/examples/buyapple.py b/zipline/examples/buyapple.py index 9ef53fce..371b7a60 100755 --- a/zipline/examples/buyapple.py +++ b/zipline/examples/buyapple.py @@ -26,11 +26,27 @@ def handle_data(context, data): record(AAPL=data[symbol('AAPL')].price) +# Note: this function can be removed if running +# this algorithm on quantopian.com +def analyze(context=None, results=None): + import matplotlib.pyplot as plt + # Plot the portfolio and asset data. + ax1 = plt.subplot(211) + results.portfolio_value.plot(ax=ax1) + ax1.set_ylabel('Portfolio value (USD)') + ax2 = plt.subplot(212, sharex=ax1) + results.AAPL.plot(ax=ax2) + ax2.set_ylabel('AAPL price (USD)') + + # Show the plot. + plt.gcf().set_size_inches(18, 8) + plt.show() + + # Note: this if-block should be removed if running # this algorithm on quantopian.com if __name__ == '__main__': from datetime import datetime - import matplotlib.pyplot as plt import pytz from zipline.algorithm import TradingAlgorithm from zipline.utils.factory import load_from_yahoo @@ -48,14 +64,4 @@ if __name__ == '__main__': identifiers=['AAPL']) results = algo.run(data) - # Plot the portfolio and asset data. - ax1 = plt.subplot(211) - results.portfolio_value.plot(ax=ax1) - ax1.set_ylabel('Portfolio value (USD)') - ax2 = plt.subplot(212, sharex=ax1) - results.AAPL.plot(ax=ax2) - ax2.set_ylabel('AAPL price (USD)') - - # Show the plot. - plt.gcf().set_size_inches(18, 8) - plt.show() + analyze(results=results) diff --git a/zipline/examples/buyapple_analyze.py b/zipline/examples/buyapple_analyze.py deleted file mode 100644 index d0029be6..00000000 --- a/zipline/examples/buyapple_analyze.py +++ /dev/null @@ -1,10 +0,0 @@ -import matplotlib.pyplot as plt - - -def analyze(context, perf): - ax1 = plt.subplot(211) - perf.portfolio_value.plot(ax=ax1) - ax2 = plt.subplot(212, sharex=ax1) - perf.AAPL.plot(ax=ax2) - plt.gcf().set_size_inches(18, 8) - plt.show() diff --git a/zipline/examples/dual_ema_talib.py b/zipline/examples/dual_ema_talib.py index 25e71e15..f8706ba6 100755 --- a/zipline/examples/dual_ema_talib.py +++ b/zipline/examples/dual_ema_talib.py @@ -65,35 +65,62 @@ def handle_data(context, data): sell=sell) +# Note: this function can be removed if running +# this algorithm on quantopian.com +def analyze(context=None, results=None): + import matplotlib.pyplot as plt + import logbook + logbook.StderrHandler().push_application() + log = logbook.Logger('Algorithm') + + fig = plt.figure() + ax1 = fig.add_subplot(211) + results.portfolio_value.plot(ax=ax1) + ax1.set_ylabel('Portfolio value (USD)') + + ax2 = fig.add_subplot(212) + ax2.set_ylabel('Price (USD)') + + # If data has been record()ed, then plot it. + # Otherwise, log the fact that no data has been recorded. + if 'AAPL' in results and 'short_ema' in results and 'long_ema' in results: + results[['AAPL', 'short_ema', 'long_ema']].plot(ax=ax2) + + ax2.plot(results.ix[results.buy].index, results.short_ema[results.buy], + '^', markersize=10, color='m') + ax2.plot(results.ix[results.sell].index, + results.short_ema[results.sell], + 'v', markersize=10, color='k') + plt.legend(loc=0) + plt.gcf().set_size_inches(18, 8) + else: + msg = 'AAPL, short_ema and long_ema data not captured using record().' + ax2.annotate(msg, xy=(0.1, 0.5)) + log.info(msg) + + plt.show() + + +# Note: this if-block should be removed if running +# this algorithm on quantopian.com if __name__ == '__main__': from datetime import datetime - import logbook - import matplotlib.pyplot as plt import pytz from zipline.algorithm import TradingAlgorithm from zipline.utils.factory import load_from_yahoo - logbook.StderrHandler().push_application() + # Set the simulation start and end dates. start = datetime(2014, 1, 1, 0, 0, 0, 0, pytz.utc) end = datetime(2014, 11, 1, 0, 0, 0, 0, pytz.utc) + + # Load price data from yahoo. data = load_from_yahoo(stocks=['AAPL'], indexes={}, start=start, end=end) + # Create and run the algorithm. algo = TradingAlgorithm(initialize=initialize, handle_data=handle_data, identifiers=['AAPL']) results = algo.run(data).dropna() - fig = plt.figure() - ax1 = fig.add_subplot(211, ylabel='portfolio value') - results.portfolio_value.plot(ax=ax1) - - ax2 = fig.add_subplot(212) - results[['AAPL', 'short_ema', 'long_ema']].plot(ax=ax2) - - ax2.plot(results.ix[results.buy].index, results.short_ema[results.buy], - '^', markersize=10, color='m') - ax2.plot(results.ix[results.sell].index, results.short_ema[results.sell], - 'v', markersize=10, color='k') - plt.legend(loc=0) - plt.gcf().set_size_inches(18, 8) - plt.show() + # Plot the portfolio and asset data. + analyze(results=results) diff --git a/zipline/examples/dual_moving_average.py b/zipline/examples/dual_moving_average.py index c1204d02..1b3dfa0e 100755 --- a/zipline/examples/dual_moving_average.py +++ b/zipline/examples/dual_moving_average.py @@ -62,16 +62,56 @@ def handle_data(context, data): long_mavg=long_mavg[context.sym]) +# Note: this function can be removed if running +# this algorithm on quantopian.com +def analyze(context=None, results=None): + import matplotlib.pyplot as plt + import logbook + logbook.StderrHandler().push_application() + log = logbook.Logger('Algorithm') + + fig = plt.figure() + ax1 = fig.add_subplot(211) + results.portfolio_value.plot(ax=ax1) + ax1.set_ylabel('Portfolio value (USD)') + + ax2 = fig.add_subplot(212) + ax2.set_ylabel('Price (USD)') + + # If data has been record()ed, then plot it. + # Otherwise, log the fact that no data has been recorded. + if ('AAPL' in results and 'short_mavg' in results and + 'long_mavg' in results): + results['AAPL'].plot(ax=ax2) + results[['short_mavg', 'long_mavg']].plot(ax=ax2) + + trans = results.ix[[t != [] for t in results.transactions]] + buys = trans.ix[[t[0]['amount'] > 0 for t in + trans.transactions]] + sells = trans.ix[ + [t[0]['amount'] < 0 for t in trans.transactions]] + ax2.plot(buys.index, results.short_mavg.ix[buys.index], + '^', markersize=10, color='m') + ax2.plot(sells.index, results.short_mavg.ix[sells.index], + 'v', markersize=10, color='k') + plt.legend(loc=0) + else: + msg = 'AAPL, short_mavg & long_mavg data not captured using record().' + ax2.annotate(msg, xy=(0.1, 0.5)) + log.info(msg) + + plt.show() + + # Note: this if-block should be removed if running # this algorithm on quantopian.com if __name__ == '__main__': from datetime import datetime - import matplotlib.pyplot as plt import pytz from zipline.algorithm import TradingAlgorithm from zipline.utils.factory import load_from_yahoo - # Set the simulation start and end dates + # Set the simulation start and end dates. start = datetime(2011, 1, 1, 0, 0, 0, 0, pytz.utc) end = datetime(2013, 1, 1, 0, 0, 0, 0, pytz.utc) @@ -85,22 +125,4 @@ if __name__ == '__main__': results = algo.run(data) # Plot the portfolio and asset data. - fig = plt.figure() - ax1 = fig.add_subplot(211) - results.portfolio_value.plot(ax=ax1) - ax1.set_ylabel('Portfolio value (USD)') - ax2 = fig.add_subplot(212) - ax2.set_ylabel('Price in (USD)') - results[['AAPL', 'short_mavg', 'long_mavg']].plot(ax=ax2) - trans = results.ix[[t != [] for t in results.transactions]] - buys = trans.ix[[t[0]['amount'] > 0 for t in - trans.transactions]] - sells = trans.ix[[t[0]['amount'] < 0 for t in trans.transactions]] - ax2.plot(buys.index, results.short_mavg.ix[buys.index], - '^', markersize=10, color='m') - ax2.plot(sells.index, results.short_mavg.ix[sells.index], - 'v', markersize=10, color='k') - plt.legend(loc=0) - - # Show the plot. - plt.show() + analyze(results=results) diff --git a/zipline/examples/dual_moving_average_analyze.py b/zipline/examples/dual_moving_average_analyze.py deleted file mode 100644 index cec76cac..00000000 --- a/zipline/examples/dual_moving_average_analyze.py +++ /dev/null @@ -1,24 +0,0 @@ -import matplotlib.pyplot as plt - - -def analyze(context, perf): - fig = plt.figure() - ax1 = fig.add_subplot(211) - perf.portfolio_value.plot(ax=ax1) - ax1.set_ylabel('portfolio value in $') - - ax2 = fig.add_subplot(212) - perf['AAPL'].plot(ax=ax2) - perf[['short_mavg', 'long_mavg']].plot(ax=ax2) - - perf_trans = perf.ix[[t != [] for t in perf.transactions]] - buys = perf_trans.ix[[t[0]['amount'] > 0 for t in perf_trans.transactions]] - sells = perf_trans.ix[ - [t[0]['amount'] < 0 for t in perf_trans.transactions]] - ax2.plot(buys.index, perf.short_mavg.ix[buys.index], - '^', markersize=10, color='m') - ax2.plot(sells.index, perf.short_mavg.ix[sells.index], - 'v', markersize=10, color='k') - ax2.set_ylabel('price in $') - plt.legend(loc=0) - plt.show() diff --git a/zipline/examples/olmar.py b/zipline/examples/olmar.py index 3a1704a2..73950523 100644 --- a/zipline/examples/olmar.py +++ b/zipline/examples/olmar.py @@ -149,18 +149,34 @@ def simplex_projection(v, b=1): w[w < 0] = 0 return w -if __name__ == '__main__': + +# Note: this function can be removed if running +# this algorithm on quantopian.com +def analyze(context=None, results=None): import matplotlib.pyplot as plt + fig = plt.figure() + ax = fig.add_subplot(111) + results.portfolio_value.plot(ax=ax) + ax.set_ylabel('Portfolio value (USD)') + plt.show() + + +# Note: this if-block should be removed if running +# this algorithm on quantopian.com +if __name__ == '__main__': + # Set the simulation start and end dates. start = datetime(2004, 1, 1, 0, 0, 0, 0, pytz.utc) end = datetime(2008, 1, 1, 0, 0, 0, 0, pytz.utc) + + # Load price data from yahoo. data = load_from_yahoo(stocks=STOCKS, indexes={}, start=start, end=end) data = data.dropna() + + # Create and run the algorithm. olmar = TradingAlgorithm(handle_data=handle_data, initialize=initialize, identifiers=STOCKS) results = olmar.run(data) - fig = plt.figure() - ax = fig.add_subplot(111) - results.portfolio_value.plot(ax=ax) - ax.set_ylabel('portfolio value in $') - plt.show() + + # Plot the portfolio data. + analyze(results=results) diff --git a/zipline/examples/pairtrade.py b/zipline/examples/pairtrade.py index f78a20b4..d04e8e58 100755 --- a/zipline/examples/pairtrade.py +++ b/zipline/examples/pairtrade.py @@ -24,6 +24,7 @@ import pytz from zipline.algorithm import TradingAlgorithm from zipline.transforms import batch_transform from zipline.utils.factory import load_from_yahoo +from zipline.api import symbol @batch_transform @@ -74,7 +75,9 @@ class Pairtrade(TradingAlgorithm): ###################################################### # 2. Compute spread and zscore zscore = self.compute_zscore(data, slope, intercept) - self.record(zscores=zscore) + self.record(zscores=zscore, + PEP=data[symbol('PEP')].price, + KO=data[symbol('KO')].price) ###################################################### # 3. Place orders @@ -116,25 +119,40 @@ class Pairtrade(TradingAlgorithm): pep_amount = self.portfolio.positions[self.PEP].amount self.order(self.PEP, -1 * pep_amount) -if __name__ == '__main__': - logbook.StderrHandler().push_application() - start = datetime(2000, 1, 1, 0, 0, 0, 0, pytz.utc) - end = datetime(2002, 1, 1, 0, 0, 0, 0, pytz.utc) - data = load_from_yahoo(stocks=['PEP', 'KO'], indexes={}, - start=start, end=end) - - pairtrade = Pairtrade() - results = pairtrade.run(data) - data['spreads'] = np.nan +# Note: this function can be removed if running +# this algorithm on quantopian.com +def analyze(context=None, results=None): ax1 = plt.subplot(211) - # TODO Bugged - indices are out of bounds - # data[[pairtrade.PEPsid, pairtrade.KOsid]].plot(ax=ax1) - plt.ylabel('price') + plt.title('PepsiCo & Coca-Cola Co. share prices') + results[['PEP', 'KO']].plot(ax=ax1) + plt.ylabel('Price (USD)') plt.setp(ax1.get_xticklabels(), visible=False) ax2 = plt.subplot(212, sharex=ax1) results.zscores.plot(ax=ax2, color='r') - plt.ylabel('zscored spread') + plt.ylabel('Z-scored spread') plt.gcf().set_size_inches(18, 8) + plt.show() + + +# Note: this if-block should be removed if running +# this algorithm on quantopian.com +if __name__ == '__main__': + logbook.StderrHandler().push_application() + + # Set the simulation start and end dates. + start = datetime(2000, 1, 1, 0, 0, 0, 0, pytz.utc) + end = datetime(2002, 1, 1, 0, 0, 0, 0, pytz.utc) + + # Load price data from yahoo. + data = load_from_yahoo(stocks=['PEP', 'KO'], indexes={}, + start=start, end=end) + + # Create and run the algorithm. + pairtrade = Pairtrade() + results = pairtrade.run(data) + + # Plot the portfolio data. + analyze(results=results) diff --git a/zipline/utils/cli.py b/zipline/utils/cli.py index 70753828..4bd16c61 100644 --- a/zipline/utils/cli.py +++ b/zipline/utils/cli.py @@ -229,12 +229,6 @@ def run_pipeline(print_algo=True, **kwargs): with open(algo_fname, 'r') as fd: algo_text = fd.read() - analyze_fname = os.path.splitext(algo_fname)[0] + '_analyze.py' - if os.path.exists(analyze_fname): - with open(analyze_fname, 'r') as fd: - # Simply append - algo_text += fd.read() - if print_algo: if PYGMENTS: highlight(algo_text, PythonLexer(), TerminalFormatter(),