Python zipline.TradingAlgorithm类代码示例

    def test_minutely_fetcher(self):
        self.responses.add(
            self.responses.GET,
            'https://fake.urls.com/aapl_minute_csv_data.csv',
            body=AAPL_MINUTE_CSV_DATA,
            content_type='text/csv',
        )

        sim_params = factory.create_simulation_parameters(
            start=pd.Timestamp("2006-01-03", tz='UTC'),
            end=pd.Timestamp("2006-01-10", tz='UTC'),
            emission_rate="minute",
            data_frequency="minute"
        )

        test_algo = TradingAlgorithm(
            script="""
from zipline.api import fetch_csv, record, sid

def initialize(context):
    fetch_csv('https://fake.urls.com/aapl_minute_csv_data.csv')

def handle_data(context, data):
    record(aapl_signal=data.current(sid(24), "signal"))
""", sim_params=sim_params, data_frequency="minute", env=self.env)

        # manually setting data portal and getting generator because we need
        # the minutely emission packets here.  TradingAlgorithm.run() only
        # returns daily packets.
        test_algo.data_portal = FetcherDataPortal(self.env,
                                                  self.trading_calendar)
        gen = test_algo.get_generator()
        perf_packets = list(gen)

        signal = [result["minute_perf"]["recorded_vars"]["aapl_signal"] for
                  result in perf_packets if "minute_perf" in result]

        self.assertEqual(6 * 390, len(signal))

        # csv data is:
        # symbol,date,signal
        # aapl,1/4/06 5:31AM, 1
        # aapl,1/4/06 11:30AM, 2
        # aapl,1/5/06 5:31AM, 1
        # aapl,1/5/06 11:30AM, 3
        # aapl,1/9/06 5:31AM, 1
        # aapl,1/9/06 11:30AM, 4 for dates 1/3 to 1/10

        # 2 signals per day, only last signal is taken. So we expect
        # 390 bars of signal NaN on 1/3
        # 390 bars of signal 2 on 1/4
        # 390 bars of signal 3 on 1/5
        # 390 bars of signal 3 on 1/6 (forward filled)
        # 390 bars of signal 4 on 1/9
        # 390 bars of signal 4 on 1/9 (forward filled)

        np.testing.assert_array_equal([np.NaN] * 390, signal[0:390])
        np.testing.assert_array_equal([2] * 390, signal[390:780])
        np.testing.assert_array_equal([3] * 780, signal[780:1560])
        np.testing.assert_array_equal([4] * 780, signal[1560:])

    def test_bts_simulation_dt(self):
        code = """
def initialize(context):
    pass
"""
        algo = TradingAlgorithm(
            script=code,
            sim_params=self.sim_params,
            env=self.env,
            metrics=metrics.load('none'),
        )

        algo.metrics_tracker = algo._create_metrics_tracker()
        benchmark_source = algo._create_benchmark_source()
        algo.metrics_tracker.handle_start_of_simulation(benchmark_source)

        dt = pd.Timestamp("2016-08-04 9:13:14", tz='US/Eastern')
        algo_simulator = AlgorithmSimulator(
            algo,
            self.sim_params,
            self.data_portal,
            BeforeTradingStartsOnlyClock(dt),
            benchmark_source,
            NoRestrictions(),
            None
        )

        # run through the algo's simulation
        list(algo_simulator.transform())

        # since the clock only ever emitted a single before_trading_start
        # event, we can check that the simulation_dt was properly set
        self.assertEqual(dt, algo_simulator.simulation_dt)

    def test_bts_simulation_dt(self):
        code = """
def initialize(context):
    pass
"""
        algo = TradingAlgorithm(script=code,
                                sim_params=self.sim_params,
                                env=self.env)

        algo.perf_tracker = PerformanceTracker(
            sim_params=self.sim_params,
            trading_calendar=self.trading_calendar,
            asset_finder=self.asset_finder,
        )

        dt = pd.Timestamp("2016-08-04 9:13:14", tz='US/Eastern')
        algo_simulator = AlgorithmSimulator(
            algo,
            self.sim_params,
            self.data_portal,
            BeforeTradingStartsOnlyClock(dt),
            algo._create_benchmark_source(),
            NoRestrictions(),
            None
        )

        # run through the algo's simulation
        list(algo_simulator.transform())

        # since the clock only ever emitted a single before_trading_start
        # event, we can check that the simulation_dt was properly set
        self.assertEqual(dt, algo_simulator.simulation_dt)

    def test_history_daily_data_1m_window(self):
        algo_text = """
from zipline.api import history, add_history

def initialize(context):
    add_history(bar_count=1, frequency='1m', field='price')

def handle_data(context, data):
    prices = history(bar_count=3, frequency='1d', field='price')
""".strip()

        start = pd.Timestamp('2006-03-20', tz='UTC')
        end = pd.Timestamp('2006-03-30', tz='UTC')

        sim_params = factory.create_simulation_parameters(
            start=start, end=end)

        with self.assertRaises(IncompatibleHistoryFrequency):
            algo = TradingAlgorithm(
                script=algo_text,
                data_frequency='daily',
                sim_params=sim_params
            )
            source = RandomWalkSource(start=start, end=end)
            algo.run(source)

    def test_history_passed_to_talib(self):
        """
        Had an issue where MagicMock was causing errors during validation
        with talib.

        We don't officially support a talib integration, yet.
        But using talib directly should work.
        """
        algo_text = """
import talib
import numpy as np

from zipline.api import history, add_history, record

def initialize(context):
    add_history(2, '1d', 'price')

def handle_data(context, data):
    prices = history(2, '1d', 'price')

    ma_result = talib.MA(np.asarray(prices[0]), timeperiod=2)
    record(ma=ma_result[-1])
""".strip()

        #      April 2007
        # Su Mo Tu We Th Fr Sa
        #  1  2  3  4  5  6  7
        #  8  9 10 11 12 13 14
        # 15 16 17 18 19 20 21
        # 22 23 24 25 26 27 28
        # 29 30

        # Eddie: this was set to 04-10 but I don't see how that makes
        # sense as it does not generate enough data to get at -2 index
        # below.
        start = pd.Timestamp("2007-04-05", tz="UTC")
        end = pd.Timestamp("2007-04-10", tz="UTC")

        sim_params = SimulationParameters(
            period_start=start,
            period_end=end,
            capital_base=float("1.0e5"),
            data_frequency="minute",
            emission_rate="daily",
        )

        test_algo = TradingAlgorithm(
            script=algo_text, data_frequency="minute", sim_params=sim_params, env=TestHistoryAlgo.env
        )

        source = RandomWalkSource(start=start, end=end)
        output = test_algo.run(source)
        # At this point, just ensure that there is no crash.
        self.assertIsNotNone(output)

        recorded_ma = output.ix[-2, "ma"]

        self.assertFalse(pd.isnull(recorded_ma))
        # Depends on seed
        np.testing.assert_almost_equal(recorded_ma, 159.76304468946876)

    def test_history_container_constructed_at_runtime(self):
        algo_text = dedent(
            """\
            from zipline.api import history
            def handle_data(context, data):
                context.prices = history(2, '1d', 'price')
            """
        )
        start = pd.Timestamp('2007-04-05', tz='UTC')
        end = pd.Timestamp('2007-04-10', tz='UTC')

        sim_params = SimulationParameters(
            period_start=start,
            period_end=end,
            capital_base=float("1.0e5"),
            data_frequency='minute',
            emission_rate='daily'
        )

        test_algo = TradingAlgorithm(
            script=algo_text,
            data_frequency='minute',
            sim_params=sim_params
        )

        source = RandomWalkSource(start=start, end=end)

        self.assertIsNone(test_algo.history_container)
        test_algo.run(source)
        self.assertIsNotNone(
            test_algo.history_container,
            msg='HistoryContainer was not constructed at runtime',
        )

        container = test_algo.history_container
        self.assertEqual(
            container.buffer_panel.window_length,
            Frequency.MAX_MINUTES['d'],
            msg='HistoryContainer.buffer_panel was not large enough to service'
            ' the given HistorySpec',
        )

        self.assertEqual(
            len(container.digest_panels),
            1,
            msg='The HistoryContainer created too many digest panels',
        )

        freq, digest = list(container.digest_panels.items())[0]
        self.assertEqual(
            freq.unit_str,
            'd',
        )

        self.assertEqual(
            digest.window_length,
            1,
            msg='The digest panel is not large enough to service the given'
            ' HistorySpec',
        )

    def test_history_in_initialize(self):
        algo_text = dedent(
            """\
            from zipline.api import history

            def initialize(context):
                history(10, '1d', 'price')

            def handle_data(context, data):
                pass
            """
        )

        start = pd.Timestamp('2007-04-05', tz='UTC')
        end = pd.Timestamp('2007-04-10', tz='UTC')

        sim_params = SimulationParameters(
            period_start=start,
            period_end=end,
            capital_base=float("1.0e5"),
            data_frequency='minute',
            emission_rate='daily',
            env=self.env,
        )

        test_algo = TradingAlgorithm(
            script=algo_text,
            data_frequency='minute',
            sim_params=sim_params,
            env=self.env,
        )

        with self.assertRaises(HistoryInInitialize):
            test_algo.initialize()

    def get_results(self, algo_code):
        algo = TradingAlgorithm(
            script=algo_code,
            env=self.env,
            sim_params=self.sim_params
        )

        return algo.run(self.data_portal)

    def run_algo(self, code, sim_params=None, data_frequency="daily"):
        if sim_params is None:
            sim_params = self.sim_params

        test_algo = TradingAlgorithm(script=code, sim_params=sim_params, env=self.env, data_frequency=data_frequency)

        results = test_algo.run(FetcherDataPortal(self.env))

        return results

    def test_basic_history_positional_args(self):
        """
        Ensure that positional args work.
        """
        algo_text = """
import copy
from zipline.api import history, add_history

def initialize(context):
    add_history(2, '1d', 'price')

def handle_data(context, data):

    prices = history(2, '1d', 'price')
    context.last_prices = copy.deepcopy(prices)
""".strip()

        #      March 2006
        # Su Mo Tu We Th Fr Sa
        #          1  2  3  4
        #  5  6  7  8  9 10 11
        # 12 13 14 15 16 17 18
        # 19 20 21 22 23 24 25
        # 26 27 28 29 30 31

        start = pd.Timestamp('2006-03-20', tz='UTC')
        end = pd.Timestamp('2006-03-21', tz='UTC')

        sim_params = factory.create_simulation_parameters(
            start=start, end=end)

        test_algo = TradingAlgorithm(
            script=algo_text,
            data_frequency='minute',
            sim_params=sim_params
        )

        source = RandomWalkSource(start=start,
                                  end=end)
        output = test_algo.run(source)
        self.assertIsNotNone(output)

        last_prices = test_algo.last_prices[0]
        oldest_dt = pd.Timestamp(
            '2006-03-20 4:00 PM', tz='US/Eastern').tz_convert('UTC')
        newest_dt = pd.Timestamp(
            '2006-03-21 4:00 PM', tz='US/Eastern').tz_convert('UTC')

        self.assertEquals(oldest_dt, last_prices.index[0])
        self.assertEquals(newest_dt, last_prices.index[-1])

        self.assertEquals(139.36946942498648, last_prices[oldest_dt])
        self.assertEquals(180.15661995395106, last_prices[newest_dt])

    def test_basic_history_one_day(self):
        algo_text = """
from zipline.api import history, add_history

def initialize(context):
    add_history(bar_count=1, frequency='1d', field='price')

def handle_data(context, data):
    prices = history(bar_count=1, frequency='1d', field='price')
    context.last_prices = prices
""".strip()

        #      March 2006
        # Su Mo Tu We Th Fr Sa
        #          1  2  3  4
        #  5  6  7  8  9 10 11
        # 12 13 14 15 16 17 18
        # 19 20 21 22 23 24 25
        # 26 27 28 29 30 31

        start = pd.Timestamp('2006-03-20', tz='UTC')
        end = pd.Timestamp('2006-03-21', tz='UTC')

        sim_params = factory.create_simulation_parameters(
            start=start, end=end)

        test_algo = TradingAlgorithm(
            script=algo_text,
            data_frequency='minute',
            sim_params=sim_params,
            env=TestHistoryAlgo.env,
        )

        source = RandomWalkSource(start=start,
                                  end=end)
        output = test_algo.run(source)

        self.assertIsNotNone(output)

        last_prices = test_algo.last_prices[0]
        # oldest and newest should be the same if there is only 1 bar
        oldest_dt = pd.Timestamp(
            '2006-03-21 4:00 PM', tz='US/Eastern').tz_convert('UTC')
        newest_dt = pd.Timestamp(
            '2006-03-21 4:00 PM', tz='US/Eastern').tz_convert('UTC')

        self.assertEquals(oldest_dt, last_prices.index[0])
        self.assertEquals(newest_dt, last_prices.index[-1])

        # Random, depends on seed
        self.assertEquals(180.15661995395106, last_prices[oldest_dt])
        self.assertEquals(180.15661995395106, last_prices[newest_dt])

    def test_current_contract_in_algo(self):
        code = dedent("""
from zipline.api import (
    record,
    continuous_future,
    schedule_function,
    get_datetime,
)

def initialize(algo):
    algo.primary_cl = continuous_future('FO', 0, 'calendar')
    algo.secondary_cl = continuous_future('FO', 1, 'calendar')
    schedule_function(record_current_contract)

def record_current_contract(algo, data):
    record(datetime=get_datetime())
    record(primary=data.current(algo.primary_cl, 'contract'))
    record(secondary=data.current(algo.secondary_cl, 'contract'))
""")
        algo = TradingAlgorithm(script=code,
                                sim_params=self.sim_params,
                                trading_calendar=self.trading_calendar,
                                env=self.env)
        results = algo.run(self.data_portal)
        result = results.iloc[0]

        self.assertEqual(result.primary.symbol,
                         'FOF16',
                         'Primary should be FOF16 on first session.')
        self.assertEqual(result.secondary.symbol,
                         'FOG16',
                         'Secondary should be FOG16 on first session.')

        result = results.iloc[1]
        # Second day, primary should switch to FOG
        self.assertEqual(result.primary.symbol,
                         'FOG16',
                         'Primary should be FOG16 on second session, auto '
                         'close is at beginning of the session.')
        self.assertEqual(result.secondary.symbol,
                         'FOH16',
                         'Secondary should be FOH16 on second session, auto '
                         'close is at beginning of the session.')

        result = results.iloc[2]
        # Second day, primary should switch to FOG
        self.assertEqual(result.primary.symbol,
                         'FOG16',
                         'Primary should remain as FOG16 on third session.')
        self.assertEqual(result.secondary.symbol,
                         'FOH16',
                         'Secondary should remain as FOH16 on third session.')

    def test_history_passed_to_func(self):
        """
        Had an issue where MagicMock was causing errors during validation
        with rolling mean.
        """
        algo_text = """
from zipline.api import history, add_history
import pandas as pd

def initialize(context):
    add_history(2, '1d', 'price')

def handle_data(context, data):
    prices = history(2, '1d', 'price')

    pd.rolling_mean(prices, 2)
""".strip()

        #      April 2007
        # Su Mo Tu We Th Fr Sa
        #  1  2  3  4  5  6  7
        #  8  9 10 11 12 13 14
        # 15 16 17 18 19 20 21
        # 22 23 24 25 26 27 28
        # 29 30

        start = pd.Timestamp('2007-04-10', tz='UTC')
        end = pd.Timestamp('2007-04-10', tz='UTC')

        sim_params = SimulationParameters(
            period_start=start,
            period_end=end,
            capital_base=float("1.0e5"),
            data_frequency='minute',
            emission_rate='minute'
        )

        test_algo = TradingAlgorithm(
            script=algo_text,
            data_frequency='minute',
            sim_params=sim_params,
            env=TestHistoryAlgo.env,
        )

        source = RandomWalkSource(start=start,
                                  end=end)
        output = test_algo.run(source)

        # At this point, just ensure that there is no crash.
        self.assertIsNotNone(output)

    def test_history_daily(self):
        bar_count = 3
        algo_text = """
from zipline.api import history, add_history

def initialize(context):
    add_history(bar_count={bar_count}, frequency='1d', field='price')
    context.history_trace = []

def handle_data(context, data):
    prices = history(bar_count={bar_count}, frequency='1d', field='price')
    context.history_trace.append(prices)
""".format(bar_count=bar_count).strip()

        #      March 2006
        # Su Mo Tu We Th Fr Sa
        #          1  2  3  4
        #  5  6  7  8  9 10 11
        # 12 13 14 15 16 17 18
        # 19 20 21 22 23 24 25
        # 26 27 28 29 30 31

        start = pd.Timestamp('2006-03-20', tz='UTC')
        end = pd.Timestamp('2006-03-30', tz='UTC')

        sim_params = factory.create_simulation_parameters(
            start=start, end=end, data_frequency='daily', env=self.env,
        )

        _, df = factory.create_test_df_source(sim_params, self.env)
        df = df.astype(np.float64)
        source = DataFrameSource(df)

        test_algo = TradingAlgorithm(
            script=algo_text,
            data_frequency='daily',
            sim_params=sim_params,
            env=TestHistoryAlgo.env,
        )

        output = test_algo.run(source)
        self.assertIsNotNone(output)

        df.columns = self.env.asset_finder.retrieve_all(df.columns)

        for i, received in enumerate(test_algo.history_trace[bar_count - 1:]):
            expected = df.iloc[i:i + bar_count]
            assert_frame_equal(expected, received)

def main():
    with open(api.__file__.rstrip('c') + 'i', 'w') as stub:
        # Imports so that Asset et al can be resolved.
        # "from MOD import *" will re-export the imports from the stub, so
        # explicitly importing.
        stub.write(dedent("""\
        from zipline.assets import Asset, Equity, Future
        from zipline.assets.futures import FutureChain
        from zipline.finance.cancel_policy import CancelPolicy
        from zipline.pipeline import Pipeline
        from zipline.protocol import Order
        from zipline.utils.events import EventRule


        """))

        # Sort to generate consistent stub file:
        for api_func in sorted(TradingAlgorithm.all_api_methods(),
                               key=attrgetter('__name__')):
            sig = inspect._signature_bound_method(inspect.signature(api_func))

            indent = ' ' * 4
            stub.write(dedent('''\
                def {func_name}{func_sig}:
                    """'''.format(func_name=api_func.__name__,
                                  func_sig=sig)))
            stub.write(dedent('{indent}{func_doc}'.format(
                func_doc=api_func.__doc__ or '\n',  # handle None docstring
                indent=indent,
            )))
            stub.write('{indent}"""\n\n'.format(indent=indent))