Return NaN when a rolling window is longer than the series

technical/vendor/qtpylib/indicators.py

@@ -32,7 +32,11 @@
 
 
 def numpy_rolling_window(data, window):
-    shape = data.shape[:-1] + (data.shape[-1] - window + 1, window)
+    # A series shorter than the window has no complete observation. An empty
+    # view lets the rolling reducers return NaN instead of raising.
+    # n - window + 1 stays 1 when the series is exactly one window long.
+    count = max(data.shape[-1] - window + 1, 0)
+    shape = data.shape[:-1] + (count, window)
     strides = data.strides + (data.strides[-1],)
     return np.lib.stride_tricks.as_strided(data, shape=shape, strides=strides)
 
@@ -43,7 +47,9 @@ def func_wrapper(data, window, as_source=False):
 
         new_series = np.empty(len(series)) * np.nan
         calculated = func(series, window)
-        new_series[-len(calculated) :] = calculated
+        # len == 0 must not use new_series[-0:], which is the whole array.
+        if len(calculated):
+            new_series[-len(calculated) :] = calculated
 
         if as_source and isinstance(data, pd.Series):
             return pd.Series(index=data.index, data=new_series)

tests/test_rolling_window.py

@@ -0,0 +1,58 @@
+import numpy as np
+import pandas as pd
+
+from technical.qtpylib import awesome_oscillator, numpy_rolling_mean, numpy_rolling_window
+
+
+def test_rolling_window_shorter_than_period_is_empty():
+    data = np.array([1.0, 2.0, 3.0])
+
+    windows = numpy_rolling_window(data, 5)
+
+    assert windows.shape == (0, 5)
+    mean = numpy_rolling_mean(data, 5)
+    assert mean.shape == (3,)
+    assert np.isnan(mean).all()
+
+
+def test_rolling_mean_keeps_warmup_nan():
+    data = np.array([1.0, 2.0, 3.0, 4.0])
+
+    mean = numpy_rolling_mean(data, 3)
+
+    assert np.isnan(mean[:2]).all()
+    np.testing.assert_allclose(mean[2:], [2.0, 3.0])
+
+
+def test_rolling_window_length_equal_to_window_has_one_row():
+    data = np.array([1.0, 2.0, 3.0, 4.0])
+
+    windows = numpy_rolling_window(data, len(data))
+
+    assert windows.shape == (1, 4)
+    np.testing.assert_array_equal(windows[0], data)
+    mean = numpy_rolling_mean(data, len(data))
+    assert mean.shape == (4,)
+    assert np.isnan(mean[:-1]).all()
+    np.testing.assert_allclose(mean[-1], 2.5)
+
+
+def test_rolling_window_full_period_is_unchanged():
+    data = np.array([1.0, 2.0, 3.0, 4.0])
+
+    windows = numpy_rolling_window(data, 3)
+
+    assert windows.shape == (2, 3)
+    np.testing.assert_array_equal(windows[0], [1.0, 2.0, 3.0])
+    np.testing.assert_array_equal(windows[1], [2.0, 3.0, 4.0])
+
+
+def test_awesome_oscillator_short_frame_is_nan():
+    frame = pd.DataFrame(
+        {"high": [1.0, 2.0, 3.0], "low": [0.5, 1.5, 2.5]},
+    )
+
+    result = awesome_oscillator(frame)
+
+    assert len(result) == 3
+    assert result.isna().all()