catalyst/zipline/transforms/ta.py

#
# Copyright 2013 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 functools
import math

import numpy as np
import pandas as pd
import talib
import copy

from six import iteritems

from zipline.transforms import BatchTransform


def zipline_wrapper(talib_fn, key_map, data):
    # get required TA-Lib input names
    if 'price' in talib_fn.input_names:
        req_inputs = [talib_fn.input_names['price']]
    elif 'prices' in talib_fn.input_names:
        req_inputs = talib_fn.input_names['prices']
    else:
        req_inputs = []

    # If there are multiple output names then the results are named,
    # if there is only one output name, it usually 'real' is best represented
    # by a float.
    # Use a DataFrame to map sid to named values, and a Series map sid
    # to floats.
    if len(talib_fn.output_names) > 1:
        all_results = pd.DataFrame(index=talib_fn.output_names,
                                   columns=data.minor_axis)
    else:
        all_results = pd.Series(index=data.minor_axis)

    for sid in data.minor_axis:
        # build talib_data from zipline data
        talib_data = dict()
        for talib_key, zipline_key in iteritems(key_map):
            # if zipline_key is found, add it to talib_data
            if zipline_key in data:
                values = data[zipline_key][sid].values
                # Do not include sids that have only nans, passing only nans
                # is incompatible with many of the underlying TALib functions.
                if pd.isnull(values).all():
                    break
                else:
                    talib_data[talib_key] = data[zipline_key][sid].values
            # if zipline_key is not found and not required, add zeros
            elif talib_key not in req_inputs:
                talib_data[talib_key] = np.zeros(data.shape[1])
            # if zipline key is not found and required, raise error
            else:
                raise KeyError(
                    'Tried to set required TA-Lib data with key '
                    '\'{0}\' but no Zipline data is available under '
                    'expected key \'{1}\'.'.format(
                        talib_key, zipline_key))

        # call talib
        if talib_data:
            talib_result = talib_fn(talib_data)

            # keep only the most recent result
            if isinstance(talib_result, (list, tuple)):
                sid_result = tuple([r[-1] for r in talib_result])
            else:
                sid_result = talib_result[-1]

            all_results[sid] = sid_result

    return all_results


def make_transform(talib_fn, name):
    """
    A factory for BatchTransforms based on TALIB abstract functions.
    """
    # make class docstring
    header = '\n#---- TA-Lib docs\n\n'
    talib_docs = getattr(talib, talib_fn.info['name']).__doc__
    divider1 = '\n#---- Default mapping (TA-Lib : Zipline)\n\n'
    mappings = '\n'.join('        {0} : {1}'.format(k, v)
                         for k, v in talib_fn.input_names.items())
    divider2 = '\n\n#---- Zipline docs\n'
    help_str = header + talib_docs + divider1 + mappings + divider2

    class TALibTransform(BatchTransform):
        __doc__ = help_str + """
        TA-Lib keyword arguments must be passed at initialization. For
        example, to construct a moving average with timeperiod of 5, pass
        "timeperiod=5" during initialization.

        All abstract TA-Lib functions accept a data dictionary containing
        'open', 'high', 'low', 'close', and 'volume' keys, even if they do
        not require those keys to run. For example, talib.MA (moving
        average) is always computed using the data under the 'close'
        key. By default, Zipline constructs this data dictionary with the
        appropriate sid data, but users may overwrite this by passing
        mappings as keyword arguments. For example, to compute the moving
        average of the sid's high, provide "close = 'high'" and Zipline's
        'high' data will be used as TA-Lib's 'close' data. Similarly, if a
        user had a data column named 'Oil', they could compute its moving
        average by passing "close='Oil'".


        **Example**

        A moving average of a data column called 'Oil' with timeperiod 5,
            talib.transforms.ta.MA(close='Oil', timeperiod=5)

        The user could find the default arguments and mappings by calling:
            help(zipline.transforms.ta.MA)


        **Arguments**

        open   : string, default 'open'
        high   : string, default 'high'
        low    : string, default 'low'
        close  : string, default 'price'
        volume : string, default 'volume'

        refresh_period : int, default 0
            The refresh_period of the BatchTransform determines the number
            of iterations that pass before the BatchTransform updates its
            internal data.

        \*\*kwargs : any arguments to be passed to the TA-Lib function.
        """

        def __init__(self,
                     close='price',
                     open='open',
                     high='high',
                     low='low',
                     volume='volume',
                     refresh_period=0,
                     bars='daily',
                     **kwargs):

            key_map = {'high': high,
                       'low': low,
                       'open': open,
                       'volume': volume,
                       'close': close}

            self.call_kwargs = kwargs

            # Make deepcopy of talib abstract function.
            # This is necessary because talib abstract functions remember
            # state, including parameters, and we need to set the parameters
            # in order to compute the lookback period that will determine the
            # BatchTransform window_length. TALIB has no way to restore default
            # parameters, so the deepcopy lets us change this function's
            # parameters without affecting other TALibTransforms of the same
            # function.
            self.talib_fn = copy.deepcopy(talib_fn)

            # set the parameters
            for param in self.talib_fn.get_parameters().keys():
                if param in kwargs:
                    self.talib_fn.set_parameters({param: kwargs[param]})

            # get the lookback
            self.lookback = self.talib_fn.lookback

            self.bars = bars
            if bars == 'daily':
                lookback = self.lookback + 1
            elif bars == 'minute':
                lookback = int(math.ceil(self.lookback / (6.5 * 60)))

            # Ensure that window_length is at least 1 day's worth of data.
            window_length = max(lookback, 1)

            transform_func = functools.partial(
                zipline_wrapper, self.talib_fn, key_map)

            super(TALibTransform, self).__init__(
                func=transform_func,
                refresh_period=refresh_period,
                window_length=window_length,
                compute_only_full=False,
                bars=bars)

        def __repr__(self):
            return 'Zipline BatchTransform: {0}'.format(
                self.talib_fn.info['name'])

    TALibTransform.__name__ = name
    # return class
    return TALibTransform


# add all TA-Lib functions to locals
for name in talib.abstract.__FUNCTION_NAMES:
    fn = getattr(talib.abstract, name)
    locals()[name] = make_transform(fn, name)