wassname/catalyst
1
0
Fork 0
mirror of https://github.com/wassname/catalyst.git synced 2026-07-02 01:01:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

MAINT: Make TALib zipline_wrapper a module level function.

Browse Source 
Prepare for making the zipline_wrapper operate on multiple sids,
as the needed nested logic will get cramped within the nested function.

Also, should help clearly define the inputs of the zipline_wrapper
function that are needed before it is passed to the BatchTransform
constructor.
This commit is contained in:
Eddie Hebert
2013-07-10 13:52:15 -04:00
parent 15323a8179
commit 37352210c0
1 changed files with 40 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
zipline/transforms/ta.py
+40 -35
View File
@@ -12,6 +12,7 @@
# 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
@@ -20,6 +21,42 @@ import copy
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 = []
# build talib_data from zipline data
talib_data = dict()
for talib_key, zipline_key in key_map.iteritems():
# if zipline_key is found, add it to talib_data
if zipline_key in data:
talib_data[talib_key] = data[zipline_key].values[:, 0]
# 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
result = talib_fn(talib_data)
# keep only the most recent result
if isinstance(result, (list, tuple)):
return tuple([r[-1] for r in result])
else:
return result[-1]
def make_transform(talib_fn, name):
"""
A factory for BatchTransforms based on TALIB abstract functions.
@@ -125,43 +162,11 @@ def make_transform(talib_fn, name):
# Ensure that window_length is at least 1 day's worth of data.
window_length = max(lookback, 1)
def zipline_wrapper(data):
# get required TA-Lib input names
if 'price' in self.talib_fn.input_names:
req_inputs = [self.talib_fn.input_names['price']]
elif 'prices' in self.talib_fn.input_names:
req_inputs = self.talib_fn.input_names['prices']
else:
req_inputs = []
# build talib_data from zipline data
talib_data = dict()
for talib_key, zipline_key in key_map.iteritems():
# if zipline_key is found, add it to talib_data
if zipline_key in data:
talib_data[talib_key] = data[zipline_key].values[:, 0]
# 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
result = self.talib_fn(talib_data)
# keep only the most recent result
if isinstance(result, (list, tuple)):
return tuple([r[-1] for r in result])
else:
return result[-1]
transform_func = functools.partial(
zipline_wrapper, self.talib_fn, key_map)
super(TALibTransform, self).__init__(
func=zipline_wrapper,
func=transform_func,
refresh_period=refresh_period,
window_length=window_length,
compute_only_full=False,