Pratik Bhadane
225 lines
7.0 KiB
Python
225 lines
7.0 KiB
Python
"""
|
|
ferro_ta.gpu — Optional GPU-accelerated indicator backend via PyTorch.
|
|
|
|
When the caller passes a PyTorch Tensor as input, the GPU path is used and the
|
|
result is returned as a PyTorch Tensor. When a NumPy array (or plain Python
|
|
sequence) is passed, the standard CPU path is used — there is **no behaviour
|
|
change** for existing CPU-only code.
|
|
|
|
Install the optional GPU extra to enable this feature:
|
|
|
|
pip install "ferro-ta[gpu]"
|
|
|
|
Or install PyTorch manually:
|
|
|
|
pip install torch
|
|
|
|
Usage
|
|
-----
|
|
>>> import torch
|
|
>>> from ferro_ta.tools.gpu import sma, ema, rsi
|
|
>>>
|
|
>>> close_gpu = torch.tensor([44.34, 44.09, 44.15, 43.61, 44.33, 44.83, 45.10], device='cuda') # or 'mps'
|
|
>>> result = sma(close_gpu, timeperiod=3)
|
|
>>> type(result) # torch.Tensor
|
|
>>> result_cpu = result.cpu().numpy()
|
|
|
|
See ``docs/gpu-backend.md`` for design notes, limitations, and benchmark data.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, cast
|
|
|
|
import numpy as np
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# PyTorch detection
|
|
# ---------------------------------------------------------------------------
|
|
|
|
try:
|
|
import torch as _torch
|
|
|
|
_TORCH_AVAILABLE = True
|
|
except ImportError:
|
|
_torch = None # type: ignore[assignment]
|
|
_TORCH_AVAILABLE = False
|
|
|
|
|
|
def _is_torch(arr: object) -> bool:
|
|
"""Return True when *arr* is a PyTorch Tensor."""
|
|
return (
|
|
_TORCH_AVAILABLE is True
|
|
and _torch is not None
|
|
and isinstance(arr, _torch.Tensor)
|
|
)
|
|
|
|
|
|
def _to_cpu(arr: object) -> np.ndarray:
|
|
"""Convert a PyTorch Tensor to a NumPy array; pass NumPy arrays through."""
|
|
if _is_torch(arr):
|
|
return cast(Any, arr).cpu().numpy()
|
|
return np.asarray(arr, dtype=np.float64)
|
|
|
|
|
|
def _to_gpu(arr: np.ndarray, device: Any = None) -> Any:
|
|
"""Move a NumPy array to the GPU (returns torch.Tensor)."""
|
|
assert _torch is not None
|
|
return _torch.tensor(arr, device=device)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# GPU implementations
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _sma_gpu(close, timeperiod: int):
|
|
"""SMA on a PyTorch Tensor using cumsum-based rolling mean."""
|
|
if _torch is None:
|
|
raise RuntimeError("PyTorch is not installed")
|
|
torch = _torch
|
|
n = close.shape[0]
|
|
result = torch.full((n,), float("nan"), dtype=close.dtype, device=close.device)
|
|
if timeperiod < 1 or n < timeperiod:
|
|
return result
|
|
# cumsum-based O(n) rolling sum
|
|
cs = torch.cumsum(close, dim=0)
|
|
# window sum for index i: cs[i] - cs[i - timeperiod] (i >= timeperiod-1)
|
|
win = cs[timeperiod - 1 :]
|
|
win = win.clone()
|
|
win[1:] -= cs[: len(win) - 1]
|
|
result[timeperiod - 1 :] = win / timeperiod
|
|
return result
|
|
|
|
|
|
def _ema_gpu(close, timeperiod: int):
|
|
"""EMA on a PyTorch Tensor — SMA-seeded, element-wise loop in Python/PyTorch."""
|
|
if _torch is None:
|
|
raise RuntimeError("PyTorch is not installed")
|
|
torch = _torch
|
|
n = close.shape[0]
|
|
result = torch.full((n,), float("nan"), dtype=close.dtype, device=close.device)
|
|
if timeperiod < 1 or n < timeperiod:
|
|
return result
|
|
k = 2.0 / (timeperiod + 1.0)
|
|
# Seed with SMA of first window (already on GPU)
|
|
seed = float(torch.mean(close[:timeperiod]).item())
|
|
result[timeperiod - 1] = seed
|
|
# Recurrence on CPU for numerical correctness then move back
|
|
close_cpu = close.cpu().numpy()
|
|
res_cpu = np.full(n, np.nan)
|
|
res_cpu[timeperiod - 1] = seed
|
|
prev = seed
|
|
for i in range(timeperiod, n):
|
|
val = float(close_cpu[i]) * k + prev * (1.0 - k)
|
|
res_cpu[i] = val
|
|
prev = val
|
|
return torch.tensor(res_cpu, dtype=close.dtype, device=close.device)
|
|
|
|
|
|
def _rsi_gpu(close, timeperiod: int):
|
|
"""RSI on a PyTorch Tensor — compute diffs on GPU, finish on CPU."""
|
|
if _torch is None:
|
|
raise RuntimeError("PyTorch is not installed")
|
|
torch = _torch
|
|
n = close.shape[0]
|
|
result = torch.full((n,), float("nan"), dtype=close.dtype, device=close.device)
|
|
if timeperiod < 1 or n <= timeperiod:
|
|
return result
|
|
# Compute price diffs on GPU
|
|
diffs = torch.diff(close).cpu().numpy() # (n-1,) numpy array
|
|
# CPU recurrence (Wilder smoothing)
|
|
res_cpu = np.full(n, np.nan)
|
|
avg_gain = np.mean(np.maximum(diffs[:timeperiod], 0.0))
|
|
avg_loss = np.mean(np.maximum(-diffs[:timeperiod], 0.0))
|
|
rs = avg_gain / avg_loss if avg_loss != 0.0 else np.inf
|
|
res_cpu[timeperiod] = 100.0 - 100.0 / (1.0 + rs)
|
|
for i in range(timeperiod + 1, n):
|
|
d = diffs[i - 1]
|
|
gain = d if d > 0.0 else 0.0
|
|
loss = -d if d < 0.0 else 0.0
|
|
avg_gain = (avg_gain * (timeperiod - 1) + gain) / timeperiod
|
|
avg_loss = (avg_loss * (timeperiod - 1) + loss) / timeperiod
|
|
rs = avg_gain / avg_loss if avg_loss != 0.0 else np.inf
|
|
res_cpu[i] = 100.0 - 100.0 / (1.0 + rs)
|
|
return torch.tensor(res_cpu, dtype=close.dtype, device=close.device)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Public API — PyTorch in → PyTorch out; NumPy in → NumPy out
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def sma(close, timeperiod: int = 30):
|
|
"""Simple Moving Average — GPU-accelerated when *close* is a PyTorch Tensor.
|
|
|
|
Parameters
|
|
----------
|
|
close : numpy.ndarray or torch.Tensor
|
|
Close price array.
|
|
timeperiod : int, default 30
|
|
Look-back window.
|
|
|
|
Returns
|
|
-------
|
|
numpy.ndarray or torch.Tensor
|
|
Same type as *close*. First ``timeperiod - 1`` values are NaN.
|
|
"""
|
|
if _is_torch(close):
|
|
if not close.is_floating_point():
|
|
close = close.float()
|
|
return _sma_gpu(close, timeperiod)
|
|
# CPU fallback
|
|
from ferro_ta import SMA # noqa: PLC0415
|
|
|
|
return SMA(np.asarray(close, dtype=np.float64), timeperiod=timeperiod)
|
|
|
|
|
|
def ema(close, timeperiod: int = 30):
|
|
"""Exponential Moving Average — GPU-accelerated when *close* is a PyTorch Tensor.
|
|
|
|
Parameters
|
|
----------
|
|
close : numpy.ndarray or torch.Tensor
|
|
timeperiod : int, default 30
|
|
|
|
Returns
|
|
-------
|
|
numpy.ndarray or torch.Tensor — same type as *close*.
|
|
"""
|
|
if _is_torch(close):
|
|
if not close.is_floating_point():
|
|
close = close.float()
|
|
return _ema_gpu(close, timeperiod)
|
|
from ferro_ta import EMA # noqa: PLC0415
|
|
|
|
return EMA(np.asarray(close, dtype=np.float64), timeperiod=timeperiod)
|
|
|
|
|
|
def rsi(close, timeperiod: int = 14):
|
|
"""Relative Strength Index — GPU-accelerated when *close* is a PyTorch Tensor.
|
|
|
|
Parameters
|
|
----------
|
|
close : numpy.ndarray or torch.Tensor
|
|
timeperiod : int, default 14
|
|
|
|
Returns
|
|
-------
|
|
numpy.ndarray or torch.Tensor — same type as *close*. Values in [0, 100].
|
|
"""
|
|
if _is_torch(close):
|
|
if not close.is_floating_point():
|
|
close = close.float()
|
|
return _rsi_gpu(close, timeperiod)
|
|
from ferro_ta import RSI # noqa: PLC0415
|
|
|
|
return RSI(np.asarray(close, dtype=np.float64), timeperiod=timeperiod)
|
|
|
|
|
|
__all__ = [
|
|
"sma",
|
|
"ema",
|
|
"rsi",
|
|
]
|