wassname/catalyst
1
0
Fork 0
mirror of https://github.com/wassname/catalyst.git synced 2026-06-30 17:22:40 +08:00
Code Issues Packages Projects Releases Wiki Activity

ENH: Enable TALib transforms to perform on multiple stocks.

Browse Source 
The TALib transform only supported operating on the first value
of a given batch transform panel row.

Instead of returning the one value, even if an panel with multiple
sids was provided, return a dictionary that maps stock to TALib
result.
This commit is contained in:
Eddie Hebert
2013-07-10 14:33:14 -04:00
parent 37352210c0
commit eac882b773
2 changed files with 34 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_transforms.py
+5 -2
View File
@@ -350,11 +350,14 @@ class TestTALIB(TestCase):
algo = TALIBAlgorithm(talib=zipline_transforms)
algo.run(self.source)
# Test if computed values match those computed by pandas rolling mean.
talib_values = np.array(algo.talib_results[zipline_transforms[0]])
sid = 0
talib_values = np.array([x[sid] for x in
algo.talib_results[zipline_transforms[0]]])
np.testing.assert_array_equal(talib_values,
pd.rolling_mean(self.panel[0]['price'],
10).values)
talib_values = np.array(algo.talib_results[zipline_transforms[1]])
talib_values = np.array([x[sid] for x in
algo.talib_results[zipline_transforms[1]]])
np.testing.assert_array_equal(talib_values,
pd.rolling_mean(self.panel[0]['price'],
25).values)
zipline/transforms/ta.py
+29 -22
View File
@@ -30,31 +30,38 @@ def zipline_wrapper(talib_fn, key_map, data):
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
all_results = {}
for sid in data.minor_axis:
# 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
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:
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))
sid_result = talib_result[-1]
# call talib
result = talib_fn(talib_data)
all_results[sid] = sid_result
# keep only the most recent result
if isinstance(result, (list, tuple)):
return tuple([r[-1] for r in result])
else:
return result[-1]
return all_results
def make_transform(talib_fn, name):