trade_pool/features.py

"""Causal technical indicators over an OHLCV frame.

Every function takes a 1-D array/series of closes (or OHLCV) and returns a series
aligned to the SAME index, where value[i] uses ONLY data at indices <= i. This is
the structural anti-lookahead guarantee: a strategy can request features at bar t
and physically cannot see t+1.
"""

from __future__ import annotations

import numpy as np


def sma(close: np.ndarray, window: int) -> np.ndarray:
    out = np.full_like(close, np.nan, dtype=float)
    if window <= 0:
        return out
    csum = np.cumsum(np.insert(close, 0, 0.0))
    out[window - 1:] = (csum[window:] - csum[:-window]) / window
    return out


def ema(close: np.ndarray, window: int) -> np.ndarray:
    out = np.full_like(close, np.nan, dtype=float)
    if len(close) == 0 or window <= 0:
        return out
    alpha = 2.0 / (window + 1.0)
    out[0] = close[0]
    for i in range(1, len(close)):
        out[i] = alpha * close[i] + (1 - alpha) * out[i - 1]
    return out


def rsi(close: np.ndarray, window: int = 14) -> np.ndarray:
    out = np.full_like(close, np.nan, dtype=float)
    if len(close) < window + 1:
        return out
    delta = np.diff(close)
    gain = np.where(delta > 0, delta, 0.0)
    loss = np.where(delta < 0, -delta, 0.0)
    avg_gain = np.mean(gain[:window])
    avg_loss = np.mean(loss[:window])
    for i in range(window, len(close)):
        if i > window:
            avg_gain = (avg_gain * (window - 1) + gain[i - 1]) / window
            avg_loss = (avg_loss * (window - 1) + loss[i - 1]) / window
        rs = avg_gain / avg_loss if avg_loss > 1e-12 else np.inf
        out[i] = 100.0 - (100.0 / (1.0 + rs))
    return out


def macd(close: np.ndarray, fast: int = 12, slow: int = 26, signal: int = 9):
    macd_line = ema(close, fast) - ema(close, slow)
    signal_line = ema(np.nan_to_num(macd_line), signal)
    return macd_line, signal_line, macd_line - signal_line


def zscore(close: np.ndarray, window: int = 20) -> np.ndarray:
    out = np.full_like(close, np.nan, dtype=float)
    for i in range(window - 1, len(close)):
        w = close[i - window + 1: i + 1]
        mu, sd = w.mean(), w.std()
        out[i] = (close[i] - mu) / sd if sd > 1e-12 else 0.0
    return out


def bollinger(close: np.ndarray, window: int = 20, k: float = 2.0):
    mid = sma(close, window)
    out_w = np.full_like(close, np.nan, dtype=float)
    for i in range(window - 1, len(close)):
        out_w[i] = close[i - window + 1: i + 1].std()
    return mid - k * out_w, mid, mid + k * out_w


def returns(close: np.ndarray) -> np.ndarray:
    out = np.zeros_like(close, dtype=float)
    out[1:] = np.diff(close) / np.where(close[:-1] == 0, np.nan, close[:-1])
    return np.nan_to_num(out)


def volatility(close: np.ndarray, window: int = 20) -> np.ndarray:
    r = returns(close)
    out = np.full_like(close, np.nan, dtype=float)
    for i in range(window - 1, len(close)):
        out[i] = r[i - window + 1: i + 1].std()
    return out


# The full feature surface exposed to a strategy at bar t (all causal).
def feature_frame(close: np.ndarray) -> dict[str, np.ndarray]:
    macd_line, macd_signal, macd_hist = macd(close)
    bb_lo, bb_mid, bb_hi = bollinger(close)
    return {
        "close": close,
        "sma_10": sma(close, 10),
        "sma_20": sma(close, 20),
        "sma_50": sma(close, 50),
        "ema_12": ema(close, 12),
        "ema_26": ema(close, 26),
        "rsi_14": rsi(close, 14),
        "macd": macd_line,
        "macd_signal": macd_signal,
        "macd_hist": macd_hist,
        "zscore_20": zscore(close, 20),
        "bb_lo": bb_lo,
        "bb_mid": bb_mid,
        "bb_hi": bb_hi,
        "ret_1": returns(close),
        "vol_20": volatility(close, 20),
    }