From 6979ae8d6aeb4e4cd825810259c182e0f855e987 Mon Sep 17 00:00:00 2001
From: Jonny Elliott <jonnybazookatone@users.noreply.github.com>
Date: Mon, 6 Jun 2016 17:06:34 -0400
Subject: [PATCH] ENH: fast stochastic oscillator added (#1255)

ENH: fast stochastic oscillator added.

A fast stochastic oscillator has been added to the technical
factors. This is the simplest of the stochastic oscillators,
and can be used to build the others.

Tests have been added that compare against the values expected
from that of ta-lib STOCHF.

FastStochasticOscillator is marked as window_safe=True to allow taking
moving averages for smoothing.
---
 tests/pipeline/test_technical.py      | 65 ++++++++++++++++++++++++++-
 zipline/pipeline/factors/__init__.py  |  2 +
 zipline/pipeline/factors/technical.py | 42 +++++++++++++++++
 3 files changed, 108 insertions(+), 1 deletion(-)

diff --git a/tests/pipeline/test_technical.py b/tests/pipeline/test_technical.py
index 110d01ee..97734477 100644
--- a/tests/pipeline/test_technical.py
+++ b/tests/pipeline/test_technical.py
@@ -10,7 +10,11 @@ from zipline.pipeline import TermGraph
 from zipline.pipeline.data import USEquityPricing
 from zipline.pipeline.engine import SimplePipelineEngine
 from zipline.pipeline.term import AssetExists
-from zipline.pipeline.factors import BollingerBands, Aroon
+from zipline.pipeline.factors import (
+    BollingerBands,
+    Aroon,
+    FastStochasticOscillator
+)
 from zipline.testing import ExplodingObject, parameter_space
 from zipline.testing.fixtures import WithAssetFinder, ZiplineTestCase
 from zipline.testing.predicates import assert_equal
@@ -174,3 +178,62 @@ class AroonTestCase(ZiplineTestCase):
         aroon.compute(today, assets, out, lows, highs)
 
         assert_equal(out, expected_out)
+
+
+class TestFastStochasticOscillator(WithTechnicalFactor, ZiplineTestCase):
+    """
+    Test the Fast Stochastic Oscillator
+    """
+
+    def test_fso_expected_basic(self):
+        """
+        Simple test of expected output from fast stochastic oscillator
+        """
+        fso = FastStochasticOscillator()
+
+        today = pd.Timestamp('2015')
+        assets = np.arange(3, dtype=np.float)
+        out = np.empty(shape=(3,), dtype=np.float)
+
+        highs = np.full((50, 3), 3)
+        lows = np.full((50, 3), 2)
+        closes = np.full((50, 3), 4)
+
+        fso.compute(today, assets, out, closes, lows, highs)
+
+        # Expected %K
+        assert_equal(out, np.full((3,), 200))
+
+    def test_fso_expected_with_talib(self):
+        """
+        Test the output that is returned from the fast stochastic oscillator
+        is the same as that from the ta-lib STOCHF function.
+        """
+        window_length = 14
+        nassets = 6
+        closes = np.random.random_integers(1, 6, size=(50, nassets))*1.0
+        highs = np.random.random_integers(4, 6, size=(50, nassets))*1.0
+        lows = np.random.random_integers(1, 3, size=(50, nassets))*1.0
+
+        expected_out_k = []
+        for i in range(nassets):
+            e = talib.STOCHF(
+                high=highs[:, i],
+                low=lows[:, i],
+                close=closes[:, i],
+                fastk_period=window_length,
+            )
+
+            expected_out_k.append(e[0][-1])
+        expected_out_k = np.array(expected_out_k)
+
+        today = pd.Timestamp('2015')
+        out = np.empty(shape=(nassets,), dtype=np.float)
+        assets = np.arange(nassets, dtype=np.float)
+
+        fso = FastStochasticOscillator()
+        fso.compute(
+            today, assets, out, closes, lows, highs
+        )
+
+        assert_equal(out, expected_out_k)
diff --git a/zipline/pipeline/factors/__init__.py b/zipline/pipeline/factors/__init__.py
index 53eecf03..7858920c 100644
--- a/zipline/pipeline/factors/__init__.py
+++ b/zipline/pipeline/factors/__init__.py
@@ -21,6 +21,7 @@ from .technical import (
     EWMSTD,
     ExponentialWeightedMovingAverage,
     ExponentialWeightedMovingStdDev,
+    FastStochasticOscillator,
     MaxDrawdown,
     Returns,
     RollingLinearRegressionOfReturns,
@@ -49,6 +50,7 @@ __all__ = [
     'ExponentialWeightedMovingAverage',
     'ExponentialWeightedMovingStdDev',
     'Factor',
+    'FastStochasticOscillator',
     'Latest',
     'MaxDrawdown',
     'RecarrayField',
diff --git a/zipline/pipeline/factors/technical.py b/zipline/pipeline/factors/technical.py
index 8eb9b85b..c54a6f1d 100644
--- a/zipline/pipeline/factors/technical.py
+++ b/zipline/pipeline/factors/technical.py
@@ -38,6 +38,7 @@ from zipline.utils.math_utils import (
     nanmean,
     nanstd,
     nansum,
+    nanmin,
 )
 from .factor import CustomFactor
 
@@ -782,3 +783,44 @@ class Aroon(CustomFactor):
             },
             out=out.down,
         )
+
+
+class FastStochasticOscillator(CustomFactor):
+    """
+    Fast Stochastic Oscillator Indicator [%K, Momentum Indicator]
+    https://wiki.timetotrade.eu/Stochastic
+
+    This stochastic is considered volatile, and varies a lot when used in
+    market analysis. It is recommended to use the slow stochastic oscillator
+    or a moving average of the %K [%D].
+
+    **Default Inputs:** :data: `zipline.pipeline.data.USEquityPricing.close`
+                        :data: `zipline.pipeline.data.USEquityPricing.low`
+                        :data: `zipline.pipeline.data.USEquityPricing.high`
+
+    **Default Window Length:** 14
+
+    Returns
+    -------
+    out: %K oscillator
+    """
+    inputs = (USEquityPricing.close, USEquityPricing.low, USEquityPricing.high)
+    window_safe = True
+    window_length = 14
+
+    def compute(self, today, assets, out, closes, lows, highs):
+
+        highest_highs = nanmax(highs, axis=0)
+        lowest_lows = nanmin(lows, axis=0)
+        today_closes = closes[-1]
+
+        evaluate(
+            '((tc - ll) / (hh - ll)) * 100',
+            local_dict={
+                'tc': today_closes,
+                'll': lowest_lows,
+                'hh': highest_highs,
+            },
+            global_dict={},
+            out=out,
+        )