Fix stochastic transpose error and add tests

.gitignore

@@ -14,3 +14,4 @@ Icon
 /sql-aggregate.txt
 /tests
 .vscode
+nosy.py

qtpylib/indicators.py

@@ -536,32 +536,22 @@ def stoch(df, window=14, d=3, k=3, fast=False):
     compute the n period relative strength indicator
     http://excelta.blogspot.co.il/2013/09/stochastic-oscillator-technical.html
     """
-    highs_ma = pd.concat([df['high'].shift(i)
-                          for i in np.arange(window)], 1, sort=True).apply(list, 1)
-    highs_ma = highs_ma.T.max().T
 
-    lows_ma = pd.concat([df['low'].shift(i)
-                         for i in np.arange(window)], 1, sort=True).apply(list, 1)
-    lows_ma = lows_ma.T.min().T
+    my_df = pd.DataFrame(index=df.index)
 
-    fast_k = ((df['close'] - lows_ma) / (highs_ma - lows_ma)) * 100
-    fast_d = numpy_rolling_mean(fast_k, d)
+    my_df['rolling_max'] = df['high'].rolling(window).max()
+    my_df['rolling_min'] = df['low'].rolling(window).min()
+
+    my_df['fast_k'] = 100 * (df['close'] - my_df['rolling_min'])/(my_df['rolling_max'] - my_df['rolling_min'])
+    my_df['fast_d'] = my_df['fast_k'].rolling(d).mean()
 
     if fast:
-        data = {
-            'k': fast_k,
-            'd': fast_d
-        }
+        return my_df.loc[:, ['fast_k', 'fast_d']]
 
-    else:
-        slow_k = numpy_rolling_mean(fast_k, k)
-        slow_d = numpy_rolling_mean(slow_k, d)
-        data = {
-            'k': slow_k,
-            'd': slow_d
-        }
+    my_df['slow_k'] = my_df['fast_k'].rolling(k).mean()
+    my_df['slow_d'] = my_df['slow_k'].rolling(d).mean()
 
-    return pd.DataFrame(index=df.index, data=data)
+    return my_df.loc[:, ['slow_k', 'slow_d']]
 
 # ---------------------------------------------
 

qtpylib/tests/test_indicators.py

@@ -0,0 +1,38 @@
+from nose.tools import eq_
+import pandas as pd
+import numpy as np
+from qtpylib import indicators as qtind
+
+def test_indicator_stoch_slow():
+    """Test the stochastic indicator logic"""
+
+    data = {'open': range(15, 150, 15),
+            'high': range(20, 200, 20),
+            'low': range(10, 100, 10),
+            'close': range(10, 100, 10),
+        }
+
+    df = pd.DataFrame(data=data)
+    my_stoch = qtind.stoch(df, window=5, d=3, k=3, fast=False)
+
+    last_stoch_slow_k = int(my_stoch['slow_k'].tail(1)*1000)
+    last_stoch_slow_d = int(my_stoch['slow_d'].tail(1)*1000)
+    eq_(last_stoch_slow_k, 33488)
+    eq_(last_stoch_slow_d, 36774)
+
+def test_indicator_stoch_fast():
+    """Test the stochastic indicator logic"""
+
+    data = {'open': range(15, 150, 15),
+            'high': range(20, 200, 20),
+            'low': range(10, 100, 10),
+            'close': range(10, 100, 10),
+        }
+
+    df = pd.DataFrame(data=data)
+    my_stoch = qtind.stoch(df, window=5, d=3, k=3, fast=True)
+
+    last_stoch_fast_k = int(my_stoch['fast_k'].tail(1)*1000)
+    last_stoch_fast_d = int(my_stoch['fast_d'].tail(1)*1000)
+    eq_(last_stoch_fast_k, 30769)
+    eq_(last_stoch_fast_d, 33488)