Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 33 additions & 27 deletions src/ml4t/engineer/features/momentum/stochrsi.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,37 +17,14 @@


@jit(nopython=True, cache=True, fastmath=False) # type: ignore[misc]
def stochrsi_fastk_numba(
close: npt.NDArray[np.float64],
timeperiod: int = 14,
def _stochrsi_fastk_from_rsi_numba(
rsi_values: npt.NDArray[np.float64],
fastk_period: int = 5,
) -> npt.NDArray[np.float64]:
"""
Calculate STOCHRSI %K using Numba for performance.

Parameters
----------
close : npt.NDArray
Price data (typically closing close)
timeperiod : int, default 14
Period for RSI calculation
fastk_period : int, default 5
Period for Stochastic calculation on RSI

Returns
-------
npt.NDArray
The STOCHRSI %K close (0-100 scale)
"""
n = len(close)

# Calculate RSI first
rsi_values = rsi_numba(close, timeperiod)

# Initialize output array
"""Calculate STOCHRSI %K from precomputed RSI values."""
n = len(rsi_values)
fastk = np.full(n, np.nan)

# Find first valid RSI
first_rsi = -1
for i in range(n):
if not np.isnan(rsi_values[i]):
Expand Down Expand Up @@ -82,6 +59,35 @@ def stochrsi_fastk_numba(
return fastk


def stochrsi_fastk_numba(
close: npt.NDArray[np.float64],
timeperiod: int = 14,
fastk_period: int = 5,
) -> npt.NDArray[np.float64]:
"""
Calculate STOCHRSI %K using Numba for performance.

RSI is computed outside the cached StochRSI kernel so a cached caller cannot
retain an older compiled RSI implementation after a package upgrade.

Parameters
----------
close : npt.NDArray
Price data (typically closing close)
timeperiod : int, default 14
Period for RSI calculation
fastk_period : int, default 5
Period for Stochastic calculation on RSI

Returns
-------
npt.NDArray
The STOCHRSI %K close (0-100 scale)
"""
rsi_values = rsi_numba(close, timeperiod)
return _stochrsi_fastk_from_rsi_numba(rsi_values, fastk_period)


def stochrsi_numba(
close: npt.NDArray[np.float64],
timeperiod: int = 14,
Expand Down
Loading