diff --git a/technical/vendor/qtpylib/indicators.py b/technical/vendor/qtpylib/indicators.py index 103ab59d..6f1a2fdf 100644 --- a/technical/vendor/qtpylib/indicators.py +++ b/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) diff --git a/tests/test_rolling_window.py b/tests/test_rolling_window.py new file mode 100644 index 00000000..f015a56d --- /dev/null +++ b/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()