123 lines
3.6 KiB
Python
123 lines
3.6 KiB
Python
from __future__ import annotations
|
|
|
|
import math
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from crypto_spot_bot.models import Candle
|
|
from crypto_spot_bot.time_series import _torch_head_outputs
|
|
from tools.train_torch_recurrent_forecaster import (
|
|
OUTPUT_LAYOUT,
|
|
RecurrentReturnModel,
|
|
_barrier_outcome,
|
|
_export_head_state,
|
|
_prepare_data,
|
|
)
|
|
|
|
|
|
def _candle(index: int, *, open_: float, high: float, low: float, close: float) -> Candle:
|
|
return Candle(index, open_, high, low, close, 100.0)
|
|
|
|
|
|
def test_barrier_target_uses_next_open_and_marks_take_profit_first() -> None:
|
|
candles = [
|
|
_candle(0, open_=90.0, high=101.0, low=89.0, close=100.0),
|
|
_candle(1, open_=100.0, high=102.0, low=99.0, close=101.0),
|
|
_candle(2, open_=101.0, high=104.0, low=100.0, close=103.0),
|
|
]
|
|
|
|
net_return, event = _barrier_outcome(
|
|
candles,
|
|
end_index=0,
|
|
horizon=2,
|
|
stop_loss_percent=0.02,
|
|
take_profit_percent=0.03,
|
|
round_trip_cost=0.002,
|
|
) or (math.nan, math.nan)
|
|
|
|
assert event == 1.0
|
|
assert net_return == pytest.approx(math.log(1.03) - 0.002)
|
|
|
|
|
|
def test_barrier_target_resolves_same_candle_tie_as_stop_loss() -> None:
|
|
candles = [
|
|
_candle(0, open_=100.0, high=101.0, low=99.0, close=100.0),
|
|
_candle(1, open_=100.0, high=104.0, low=97.0, close=101.0),
|
|
]
|
|
|
|
net_return, event = _barrier_outcome(
|
|
candles,
|
|
end_index=0,
|
|
horizon=1,
|
|
stop_loss_percent=0.02,
|
|
take_profit_percent=0.03,
|
|
round_trip_cost=0.002,
|
|
) or (math.nan, math.nan)
|
|
|
|
assert event == 0.0
|
|
assert net_return == pytest.approx(math.log(0.98) - 0.002)
|
|
|
|
|
|
def test_multitask_head_export_matches_runtime_inference() -> None:
|
|
torch.manual_seed(7)
|
|
model = RecurrentReturnModel(
|
|
architecture="gru",
|
|
input_size=2,
|
|
hidden_size=4,
|
|
num_layers=1,
|
|
dropout=0.0,
|
|
output_size=2 * len(OUTPUT_LAYOUT),
|
|
attention_pooling=False,
|
|
context_norm=False,
|
|
multitask_head=True,
|
|
head_hidden_size=6,
|
|
)
|
|
model.eval()
|
|
context = torch.tensor([[0.2, -0.1, 0.4, 0.3]], dtype=torch.float32)
|
|
with torch.no_grad():
|
|
shared = model.head_activation(model.head_hidden(context))
|
|
returns = model.return_head(shared).view(1, 2, 4)
|
|
events = model.event_head(shared).view(1, 2, 1)
|
|
expected = torch.cat((returns, events), dim=2).reshape(-1).tolist()
|
|
|
|
entry = {"multitask_head": True, **_export_head_state(model)}
|
|
actual = _torch_head_outputs(context[0].tolist(), entry, hidden_size=4)
|
|
|
|
assert actual == pytest.approx(expected, abs=2e-6)
|
|
|
|
|
|
def test_pooled_training_populates_symbol_identity_feature() -> None:
|
|
candles = [
|
|
_candle(
|
|
index,
|
|
open_=100.0 + index * 0.01,
|
|
high=100.2 + index * 0.01,
|
|
low=99.8 + index * 0.01,
|
|
close=100.0 + index * 0.01,
|
|
)
|
|
for index in range(180)
|
|
]
|
|
|
|
prepared = _prepare_data(
|
|
symbol="BTCUSDT",
|
|
candles=candles,
|
|
feature_names=["return_1", "symbol_is_BTCUSDT", "symbol_is_ETHUSDT"],
|
|
lookback=8,
|
|
target_horizons=[3],
|
|
decision_horizon=3,
|
|
round_trip_cost=0.002,
|
|
stop_loss_percent=0.04,
|
|
take_profit_percent=0.035,
|
|
market_candles={"BTCUSDT": candles},
|
|
trend_candles=candles,
|
|
validation_window=24,
|
|
holdout_window=32,
|
|
clip=8.0,
|
|
device=torch.device("cpu"),
|
|
)
|
|
|
|
assert prepared is not None
|
|
assert torch.all(prepared.train_x[:, :, 1] == 1.0)
|
|
assert torch.all(prepared.train_x[:, :, 2] == 0.0)
|