Add PyTorch recurrent forecaster

2026-06-20 21:28:05 +03:00
parent bac55f22b7
commit 92538850ad
7 changed files with 781 additions and 13 deletions
@@ -122,3 +122,50 @@ def test_time_series_forecaster_reads_lstm_artifact(make_settings, tmp_path) ->

    assert forecast.usable is True
    assert any(candidate["model"] == "lstm" for candidate in forecast.candidates)
+
+
+def test_time_series_forecaster_reads_torch_gru_artifact(make_settings, tmp_path) -> None:
+    artifact_path = tmp_path / "lstm_forecaster.json"
+    hidden_size = 2
+    artifact_path.write_text(
+        json.dumps(
+            {
+                "version": 2,
+                "type": "pytorch_recurrent_forecaster",
+                "symbols": {
+                    "BTCUSDT": {
+                        "model": "torch_gru",
+                        "architecture": "gru",
+                        "lookback": 8,
+                        "hidden_size": hidden_size,
+                        "num_layers": 1,
+                        "mean": 0.0,
+                        "scale": 0.001,
+                        "clip": 8.0,
+                        "state_dict": {
+                            "weight_ih_l0": [[0.0] for _ in range(3 * hidden_size)],
+                            "weight_hh_l0": [[0.0, 0.0] for _ in range(3 * hidden_size)],
+                            "bias_ih_l0": [0.0 for _ in range(3 * hidden_size)],
+                            "bias_hh_l0": [0.0 for _ in range(3 * hidden_size)],
+                        },
+                        "head_weight": [0.0, 0.0],
+                        "head_bias": 0.2,
+                    },
+                },
+            }
+        ),
+        encoding="utf-8",
+    )
+    settings = make_settings(
+        tmp_path,
+        time_series_min_candles=80,
+        time_series_validation_window=20,
+        time_series_lstm_enabled=True,
+        time_series_lstm_model_path=artifact_path,
+    )
+    returns = [0.00015 if index % 4 else -0.00005 for index in range(140)]
+
+    forecast = TimeSeriesForecaster(settings).forecast(_candles_from_returns(returns), symbol="BTCUSDT")
+
+    assert forecast.usable is True
+    assert any(candidate["model"] == "torch_gru" for candidate in forecast.candidates)