Artifex/gpu_feature_engine_v1_2.py

79 lines
4.7 KiB
Python
Raw Permalink Normal View History

"""V1.2 GPU feature engine.
This module intentionally leaves ``gpu_feature_engine_v1`` unchanged. Only
Supertrend uses the V1.2 ordered, on-device RMA seed implementation.
"""
from __future__ import annotations
from typing import Any
import torch
from gpu_feature_engine_v1 import _empty, _true_range, evaluate_batch as _v1_evaluate_batch
def rma_ordered_seed(values: torch.Tensor, period: int) -> torch.Tensor:
"""RMA with the historical oracle's left-to-right seed association."""
out = _empty(values)
valid = torch.nonzero(~torch.isnan(values), as_tuple=False).flatten()
if valid.numel() < period:
return out
seed_bar = int(valid[period - 1].item())
total = torch.zeros((), dtype=values.dtype, device=values.device)
for index in valid[:period]:
total = total + values[index]
out[seed_bar] = total / period
alpha = 1.0 / period
for index in range(seed_bar + 1, values.numel()):
out[index] = out[index - 1] if torch.isnan(values[index]) else alpha * values[index] + (1.0 - alpha) * out[index - 1]
return out
def supertrend_trace(close: torch.Tensor, high: torch.Tensor, low: torch.Tensor, period: int, multiplier: float) -> dict[str, torch.Tensor]:
"""Return complete V1.2 Supertrend state, including branch predicates."""
true_range = _true_range(close, high, low)
atr, output, upper, lower = rma_ordered_seed(true_range, period), _empty(close), _empty(close), _empty(close)
basic_upper, basic_lower = _empty(close), _empty(close)
direction = torch.zeros_like(close, dtype=torch.int8)
initialized = torch.zeros_like(close, dtype=torch.bool)
prior_close_le_upper = torch.zeros_like(close, dtype=torch.bool)
prior_close_ge_lower = torch.zeros_like(close, dtype=torch.bool)
active_long = torch.zeros_like(close, dtype=torch.bool)
close_ge_lower = torch.zeros_like(close, dtype=torch.bool)
close_le_upper = torch.zeros_like(close, dtype=torch.bool)
output_uses_lower = torch.zeros_like(close, dtype=torch.bool)
direction_transition = torch.zeros_like(close, dtype=torch.bool)
for index in range(period, close.numel()):
midpoint = (high[index] + low[index]) / 2.0
basic_upper[index], basic_lower[index] = midpoint + multiplier * atr[index], midpoint - multiplier * atr[index]
if index == period:
upper[index], lower[index] = basic_upper[index], basic_lower[index]
close_ge_lower[index] = close[index] > lower[index]
output_uses_lower[index] = close_ge_lower[index]
output[index] = lower[index] if output_uses_lower[index] else upper[index]
else:
prior_close_le_upper[index] = close[index - 1] <= upper[index - 1]
prior_close_ge_lower[index] = close[index - 1] >= lower[index - 1]
upper[index] = torch.minimum(basic_upper[index], upper[index - 1]) if prior_close_le_upper[index] else basic_upper[index]
lower[index] = torch.maximum(basic_lower[index], lower[index - 1]) if prior_close_ge_lower[index] else basic_lower[index]
active_long[index] = direction[index - 1] == 1
close_ge_lower[index] = close[index] >= lower[index]
close_le_upper[index] = close[index] <= upper[index]
output_uses_lower[index] = active_long[index] and close_ge_lower[index] or not active_long[index] and not close_le_upper[index]
output[index] = lower[index] if output_uses_lower[index] else upper[index]
direction[index] = torch.where(close[index] > output[index], 1, -1)
initialized[index] = True
if index > period:
direction_transition[index] = direction[index] != direction[index - 1]
return {"output": output, "true_range": true_range, "atr": atr, "basic_upper": basic_upper, "basic_lower": basic_lower, "upper": upper, "lower": lower, "direction": direction, "initialized": initialized, "prior_close_le_upper": prior_close_le_upper, "prior_close_ge_lower": prior_close_ge_lower, "active_long": active_long, "close_ge_lower": close_ge_lower, "close_le_upper": close_le_upper, "output_uses_lower": output_uses_lower, "direction_transition": direction_transition}
def evaluate_batch(request: dict[str, Any], close: torch.Tensor, high: torch.Tensor, low: torch.Tensor, volume: torch.Tensor) -> dict[str, torch.Tensor]:
"""Evaluate Batch01, replacing only Supertrend requests with V1.2 logic."""
other = {"requests": [item for item in request["requests"] if int(item["indicator_id"]) != 19]}
outputs = _v1_evaluate_batch(other, close, high, low, volume) if other["requests"] else {}
for item in request["requests"]:
if int(item["indicator_id"]) == 19:
outputs[str(item["request_id"])] = supertrend_trace(close, high, low, int(item["period"]), float(item["p1"]))["output"]
return outputs