Helper for nan index array and reinsert.

jesse/helpers.py

@@ -263,6 +263,10 @@ def get_config(keys: str, default=None):
     return CACHED_CONFIG[keys]
 
 
+def get_nan_indices(array):
+    return np.where(np.isnan(array))[0]
+
+
 def get_strategy_class(strategy_name):
     from pydoc import locate
 
@@ -530,6 +534,12 @@ def readable_duration(seconds, granularity=2):
     return ', '.join(result[:granularity])
 
 
+def reinsert_nan(array, nan_indices):
+    for i in range(nan_indices.shape[0]):
+        array = np.concatenate((array[:nan_indices[i]], [np.nan], array[nan_indices[i]:]))
+    return array
+
+
 def relative_to_absolute(path: str) -> str:
     return os.path.abspath(path)
 

tests/test_helpers.py

@@ -217,6 +217,12 @@ def test_get_config():
     assert jh.get_config('env.logging.order_submission', 2020) is True
 
 
+def test_get_nan_indices():
+    arr = np.array([0, 11, 22, np.nan, 33, 44, 54, 55, np.nan, np.nan, 20])
+    ind = jh.get_nan_indices(arr)
+    assert (ind == np.array([3, 8, 9])).all()
+
+
 def test_get_strategy_class():
     from jesse.strategies import Strategy
     assert issubclass(jh.get_strategy_class("Test01"), Strategy)
@@ -410,6 +416,14 @@ def test_readable_duration():
     assert jh.readable_duration(604312) == "6 days, 23 hours"
 
 
+def test_reinsert_nan():
+    arr = np.array([0, 11, 22, np.nan, 33, 44, 54, 55, np.nan, np.nan, 20])
+    arr_without_nan = arr[~np.isnan(arr)]
+    ind = jh.get_nan_indices(arr)
+    arr_with_nan = jh.reinsert_nan(arr_without_nan, ind)
+    assert ((arr == arr_with_nan) | (np.isnan(arr) & np.isnan(arr_with_nan))).all()
+
+
 def test_relative_to_absolute():
     from pathlib import Path
     assert jh.relative_to_absolute("tests/test_helpers.py") == str(Path(__file__).absolute())