"""Lightweight neural network for trend prediction.""" from __future__ import annotations import math import os from typing import List, Optional from core.trend_analysis import Candle try: import torch import torch.nn as nn TORCH_AVAILABLE = True except ImportError: TORCH_AVAILABLE = False class TrendPredictor: """ Lightweight trend predictor using a small convolutional network. Architecture (~10k parameters): - Conv1d: 5 features -> 16 channels, kernel=3 - Conv1d: 16 channels -> 8 channels, kernel=3 - Flatten + Linear -> 32 units - Linear -> 3 units (uptrend/neutral/downtrend) Input: sequence of candles (normalized OHLCV) Output: trend classification with confidence """ def __init__(self, sequence_length: int = 24, model_path: Optional[str] = None): self.sequence_length = sequence_length self._model = None self._initialized = False self._classes = ["downtrend", "neutral", "uptrend"] self.model_path = model_path or os.path.join( os.path.dirname(os.path.dirname(__file__)), "models", "trend_predictor.pt" ) if TORCH_AVAILABLE: self._model = _TrendNet(sequence_length) self._initialized = True # Try to load saved weights if os.path.exists(self.model_path): try: self._model.load_state_dict(torch.load(self.model_path, weights_only=True, map_location="cpu")) self._model.eval() except Exception: pass def is_available(self) -> bool: """Check if PyTorch is available.""" return TORCH_AVAILABLE and self._initialized def predict(self, candles: List[Candle]) -> dict: """ Predict trend from candle sequence. Args: candles: List of candles (oldest to newest), at least sequence_length Returns: dict with 'direction', 'confidence', and 'raw_output' """ if not self.is_available(): return {"direction": "neutral", "confidence": 0.0, "raw_output": None} if len(candles) < self.sequence_length: return {"direction": "neutral", "confidence": 0.0, "raw_output": None} # Use the most recent sequence_length candles sequence = candles[-self.sequence_length:] features = self._extract_features(sequence) self._model.eval() with torch.no_grad(): features_tensor = torch.FloatTensor(features).reshape(5, self.sequence_length).unsqueeze(0) output = self._model(features_tensor) probabilities = torch.softmax(output, dim=1) predicted = torch.argmax(probabilities, dim=1).item() direction = self._classes[predicted] confidence = probabilities[0][predicted].item() return { "direction": direction, "confidence": round(confidence, 4), "raw_output": [round(p.item(), 4) for p in probabilities[0]], } def train(self, training_data: List[List[Candle]], labels: List[str], epochs: int = 50, lr: float = 0.001, save: bool = True) -> dict: """ Train the model on historical candle data. Args: training_data: List of candle sequences labels: Corresponding trend labels ('uptrend', 'neutral', 'downtrend') epochs: Number of training epochs lr: Learning rate save: Whether to save model weights after training Returns: Training history with final loss and accuracy """ if not self.is_available(): return {"error": "PyTorch not available"} if len(training_data) != len(labels): return {"error": "Data and labels length mismatch"} label_map = {"downtrend": 0, "neutral": 1, "uptrend": 2} targets = torch.LongTensor([label_map.get(l, 1) for l in labels]) features_list = [] for candles in training_data: if len(candles) >= self.sequence_length: seq = candles[-self.sequence_length:] features_list.append(self._extract_features(seq)) if not features_list: return {"error": "No valid training sequences"} features_tensor = torch.FloatTensor(features_list) targets = targets[:len(features_list)] self._model.train() optimizer = torch.optim.Adam(self._model.parameters(), lr=lr) criterion = nn.CrossEntropyLoss() history = {"loss": [], "accuracy": []} for epoch in range(epochs): optimizer.zero_grad() # Reshape features from (N, seq_len*5) to (N, 5, seq_len) for Conv1d batch_input = features_tensor.reshape(len(features_list), 5, self.sequence_length) outputs = self._model(batch_input) loss = criterion(outputs, targets) loss.backward() optimizer.step() with torch.no_grad(): predicted = torch.argmax(outputs, dim=1) accuracy = (predicted == targets).float().mean().item() history["loss"].append(round(loss.item(), 6)) history["accuracy"].append(round(accuracy, 4)) result = { "final_loss": round(history["loss"][-1], 6), "final_accuracy": round(history["accuracy"][-1], 4), "epochs_trained": epochs, "total_samples": len(features_list), } if save: self.save_model() return result def save_model(self) -> bool: """Save model weights to disk.""" if not self.is_available(): return False try: os.makedirs(os.path.dirname(self.model_path), exist_ok=True) torch.save(self._model.state_dict(), self.model_path) return True except Exception: return False def load_model(self) -> bool: """Load model weights from disk.""" if not self.is_available() or not os.path.exists(self.model_path): return False try: self._model.load_state_dict(torch.load(self.model_path, weights_only=True, map_location="cpu")) self._model.eval() return True except Exception: return False @staticmethod def _extract_features(candles: List[Candle]) -> List[float]: """ Extract normalized features from candle sequence. Features per candle: open, high, low, close, volume (normalized) Returns flat list of length sequence_length * 5. """ if not candles: return [] # Normalize using the range of the sequence closes = [c.close for c in candles] price_min = min(closes) price_max = max(closes) price_range = price_max - price_min if price_max != price_min else 1.0 volumes = [c.volume for c in candles] vol_min = min(volumes) vol_max = max(volumes) vol_range = vol_max - vol_min if vol_max != vol_min else 1.0 features = [] for c in candles: features.append((c.open - price_min) / price_range) features.append((c.high - price_min) / price_range) features.append((c.low - price_min) / price_range) features.append((c.close - price_min) / price_range) features.append((c.volume - vol_min) / vol_range) return features class _TrendNet(nn.Module): """Internal PyTorch model for trend prediction.""" def __init__(self, sequence_length: int): super().__init__() # Input: (batch, 5 features, sequence_length) self.conv1 = nn.Conv1d(5, 16, kernel_size=3, padding=1) self.conv2 = nn.Conv1d(16, 8, kernel_size=3, padding=1) self.pool = nn.AdaptiveAvgPool1d(1) self.fc1 = nn.Linear(8, 16) self.fc2 = nn.Linear(16, 3) # 3 classes: downtrend, neutral, uptrend self.relu = nn.ReLU() self.dropout = nn.Dropout(0.2) def forward(self, x): # x shape: (batch, 5, seq_len) x = self.relu(self.conv1(x)) x = self.relu(self.conv2(x)) x = self.pool(x) # (batch, 8, 1) x = x.squeeze(-1) # (batch, 8) x = self.dropout(x) x = self.relu(self.fc1(x)) x = self.fc2(x) return x