catalyst/tests/test_benchmark.py

#
# Copyright 2015 Quantopian, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
from unittest import TestCase
from datetime import timedelta
import numpy as np
import pandas as pd
from testfixtures import TempDirectory
from zipline.data.us_equity_pricing import SQLiteAdjustmentWriter, \
    SQLiteAdjustmentReader
from zipline.errors import (
    BenchmarkAssetNotAvailableTooEarly,
    BenchmarkAssetNotAvailableTooLate,
    InvalidBenchmarkAsset)

from zipline.finance.trading import TradingEnvironment
from zipline.sources.benchmark_source import BenchmarkSource
from zipline.utils import factory
from zipline.testing.core import create_data_portal, write_minute_data, \
    create_empty_splits_mergers_frame
from .test_perf_tracking import MockDailyBarSpotReader


class TestBenchmark(TestCase):
    @classmethod
    def setUpClass(cls):
        cls.env = TradingEnvironment()
        cls.tempdir = TempDirectory()

        cls.sim_params = factory.create_simulation_parameters()

        cls.env.write_data(equities_data={
            1: {
                "start_date": cls.sim_params.trading_days[0],
                "end_date": cls.sim_params.trading_days[-1] + timedelta(days=1)
            },
            2: {
                "start_date": cls.sim_params.trading_days[0],
                "end_date": cls.sim_params.trading_days[-1] + timedelta(days=1)
            },
            3: {
                "start_date": cls.sim_params.trading_days[100],
                "end_date": cls.sim_params.trading_days[-100]
            },
            4: {
                "start_date": cls.sim_params.trading_days[0],
                "end_date": cls.sim_params.trading_days[-1] + timedelta(days=1)
            }

        })

        dbpath = os.path.join(cls.tempdir.path, "adjustments.db")

        writer = SQLiteAdjustmentWriter(dbpath, cls.env.trading_days,
                                        MockDailyBarSpotReader())
        splits = mergers = create_empty_splits_mergers_frame()
        dividends = pd.DataFrame({
            'sid': np.array([], dtype=np.uint32),
            'amount': np.array([], dtype=np.float64),
            'declared_date': np.array([], dtype='datetime64[ns]'),
            'ex_date': np.array([], dtype='datetime64[ns]'),
            'pay_date': np.array([], dtype='datetime64[ns]'),
            'record_date': np.array([], dtype='datetime64[ns]'),
        })
        declared_date = cls.sim_params.trading_days[45]
        ex_date = cls.sim_params.trading_days[50]
        record_date = pay_date = cls.sim_params.trading_days[55]

        stock_dividends = pd.DataFrame({
            'sid': np.array([4], dtype=np.uint32),
            'payment_sid': np.array([5], dtype=np.uint32),
            'ratio': np.array([2], dtype=np.float64),
            'declared_date': np.array([declared_date], dtype='datetime64[ns]'),
            'ex_date': np.array([ex_date], dtype='datetime64[ns]'),
            'record_date': np.array([record_date], dtype='datetime64[ns]'),
            'pay_date': np.array([pay_date], dtype='datetime64[ns]'),
        })
        writer.write(splits, mergers, dividends,
                     stock_dividends=stock_dividends)

        cls.data_portal = create_data_portal(
            cls.env,
            cls.tempdir,
            cls.sim_params,
            [1, 2, 3, 4],
            adjustment_reader=SQLiteAdjustmentReader(dbpath)
        )

    @classmethod
    def tearDownClass(cls):
        del cls.env
        cls.tempdir.cleanup()

    def test_normal(self):
        days_to_use = self.sim_params.trading_days[1:]

        source = BenchmarkSource(
            1, self.env, days_to_use, self.data_portal
        )

        # should be the equivalent of getting the price history, then doing
        # a pct_change on it
        manually_calculated = self.data_portal.get_history_window(
            [1], days_to_use[-1], len(days_to_use), "1d", "close"
        )[1].pct_change()

        # compare all the fields except the first one, for which we don't have
        # data in manually_calculated
        for idx, day in enumerate(days_to_use[1:]):
            self.assertEqual(
                source.get_value(day),
                manually_calculated[idx + 1]
            )

    def test_asset_not_trading(self):
        with self.assertRaises(BenchmarkAssetNotAvailableTooEarly) as exc:
            BenchmarkSource(
                3,
                self.env,
                self.sim_params.trading_days[1:],
                self.data_portal
            )

        self.assertEqual(
            '3 does not exist on 2006-01-04 00:00:00+00:00. '
            'It started trading on 2006-05-26 00:00:00+00:00.',
            exc.exception.message
        )

        with self.assertRaises(BenchmarkAssetNotAvailableTooLate) as exc2:
            BenchmarkSource(
                3,
                self.env,
                self.sim_params.trading_days[120:],
                self.data_portal
            )

        self.assertEqual(
            '3 does not exist on 2006-06-26 00:00:00+00:00. '
            'It stopped trading on 2006-08-09 00:00:00+00:00.',
            exc2.exception.message
        )

    def test_asset_IPOed_same_day(self):
        # gotta get some minute data up in here.
        # add sid 4 for a couple of days
        minutes = self.env.minutes_for_days_in_range(
            self.sim_params.trading_days[0],
            self.sim_params.trading_days[5]
        )

        path = write_minute_data(
            self.env,
            self.tempdir,
            minutes,
            [2]
        )

        self.data_portal._minutes_equities_path = path

        source = BenchmarkSource(
            2,
            self.env,
            self.sim_params.trading_days,
            self.data_portal
        )

        days_to_use = self.sim_params.trading_days

        # first value should be 0.0, coming from daily data
        self.assertAlmostEquals(0.0, source.get_value(days_to_use[0]))

        manually_calculated = self.data_portal.get_history_window(
            [2], days_to_use[-1], len(days_to_use), "1d", "close"
        )[2].pct_change()

        for idx, day in enumerate(days_to_use[1:]):
            self.assertEqual(
                source.get_value(day),
                manually_calculated[idx + 1]
            )

    def test_no_stock_dividends_allowed(self):
        # try to use sid(4) as benchmark, should blow up due to the presence
        # of a stock dividend

        with self.assertRaises(InvalidBenchmarkAsset) as exc:
            BenchmarkSource(
                4, self.env, self.sim_params.trading_days, self.data_portal
            )

        self.assertEqual("4 cannot be used as the benchmark because it has a "
                         "stock dividend on 2006-03-16 00:00:00.  Choose "
                         "another asset to use as the benchmark.",
                         exc.exception.message)