"""Model evaluation metrics and utilities.""" from typing import Dict, Literal import numpy as np import torch from sklearn.metrics import ( accuracy_score, precision_score, recall_score, f1_score, confusion_matrix, mean_squared_error, mean_absolute_error ) class ModelEvaluator: """Evaluate model performance with task-specific metrics.""" def __init__(self, task: Literal['classification', 'regression']): """Initialize evaluator. Args: task: Task type ('classification' or 'regression') """ self.task = task def evaluate( self, y_true: np.ndarray, y_pred: np.ndarray, y_prob: np.ndarray = None ) -> Dict[str, float]: """Compute evaluation metrics. Args: y_true: True labels y_pred: Predicted labels/values y_prob: Prediction probabilities (for classification) Returns: Dictionary of metrics """ if self.task == 'classification': return self._evaluate_classification(y_true, y_pred, y_prob) else: return self._evaluate_regression(y_true, y_pred) def _evaluate_classification( self, y_true: np.ndarray, y_pred: np.ndarray, y_prob: np.ndarray = None ) -> Dict[str, float]: """Compute classification metrics. Args: y_true: True labels (0 or 1) y_pred: Predicted labels (0 or 1) y_prob: Prediction probabilities Returns: Dictionary of classification metrics """ metrics = { 'accuracy': accuracy_score(y_true, y_pred), 'precision': precision_score(y_true, y_pred, zero_division=0), 'recall': recall_score(y_true, y_pred, zero_division=0), 'f1': f1_score(y_true, y_pred, zero_division=0) } # Confusion matrix cm = confusion_matrix(y_true, y_pred) if cm.shape == (2, 2): tn, fp, fn, tp = cm.ravel() metrics['true_negatives'] = int(tn) metrics['false_positives'] = int(fp) metrics['false_negatives'] = int(fn) metrics['true_positives'] = int(tp) # Specificity if (tn + fp) > 0: metrics['specificity'] = tn / (tn + fp) # Directional accuracy (if probabilities provided) if y_prob is not None: metrics['mean_confidence'] = float(np.mean(y_prob)) return metrics def _evaluate_regression( self, y_true: np.ndarray, y_pred: np.ndarray ) -> Dict[str, float]: """Compute regression metrics. Args: y_true: True values y_pred: Predicted values Returns: Dictionary of regression metrics """ # Basic regression metrics mse = mean_squared_error(y_true, y_pred) mae = mean_absolute_error(y_true, y_pred) rmse = np.sqrt(mse) # MAPE (Mean Absolute Percentage Error) # Avoid division by zero epsilon = 1e-8 mape = np.mean(np.abs((y_true - y_pred) / (np.abs(y_true) + epsilon))) * 100 # Directional accuracy (sign agreement) directional_accuracy = np.mean(np.sign(y_true) == np.sign(y_pred)) metrics = { 'mse': float(mse), 'rmse': float(rmse), 'mae': float(mae), 'mape': float(mape), 'directional_accuracy': float(directional_accuracy) } return metrics def backtest_strategy( self, y_true: np.ndarray, y_pred: np.ndarray, initial_capital: float = 10000.0, transaction_cost: float = 0.001 ) -> Dict[str, float]: """Simple backtesting with long/short strategy. Args: y_true: True returns y_pred: Predicted labels (classification) or returns (regression) initial_capital: Starting capital transaction_cost: Transaction cost per trade (e.g., 0.001 = 0.1%) Returns: Dictionary of backtest metrics """ capital = initial_capital positions = [] returns_list = [] for i in range(len(y_pred)): if self.task == 'classification': # Classification: y_pred已经是0/1的标签（不是概率） # 1 = buy, 0 = sell/short position = 1 if y_pred[i] == 1 else -1 else: # Regression: positive prediction = buy, negative = sell position = 1 if y_pred[i] > 0 else -1 # Calculate return based on position actual_return = y_true[i] strategy_return = position * actual_return # Apply transaction cost if position changes if i > 0 and positions[-1] != position: strategy_return -= transaction_cost # Update capital capital *= (1 + strategy_return) positions.append(position) returns_list.append(strategy_return) # Calculate metrics returns_array = np.array(returns_list) total_return = (capital - initial_capital) / initial_capital # Sharpe ratio (annualized, assuming hourly data) # For hourly data: sqrt(24 * 365) ≈ 93 if len(returns_array) > 0 and returns_array.std() > 0: sharpe_ratio = (returns_array.mean() / returns_array.std()) * np.sqrt(93) else: sharpe_ratio = 0.0 # Maximum drawdown cumulative = np.cumprod(1 + returns_array) running_max = np.maximum.accumulate(cumulative) drawdown = (cumulative - running_max) / running_max max_drawdown = float(drawdown.min()) # Win rate win_rate = np.mean(returns_array > 0) return { 'total_return': float(total_return), 'final_capital': float(capital), 'sharpe_ratio': float(sharpe_ratio), 'max_drawdown': float(max_drawdown), 'win_rate': float(win_rate), 'num_trades': int(np.sum(np.diff(positions) != 0)) } def evaluate_model( model: torch.nn.Module, dataloader: torch.utils.data.DataLoader, task: Literal['classification', 'regression'], device: str = 'cpu' ) -> Dict[str, float]: """Evaluate a model on a dataset. Args: model: PyTorch model dataloader: DataLoader with evaluation data task: Task type device: Device to run evaluation on Returns: Dictionary of evaluation metrics """ model.eval() model = model.to(device) all_preds = [] all_probs = [] all_labels = [] with torch.no_grad(): for batch_X, batch_y in dataloader: batch_X = batch_X.to(device) batch_y = batch_y.to(device) # Forward pass outputs = model(batch_X).squeeze() if task == 'classification': # 注意：修复后的模型输出logits，需要应用sigmoid得到概率 probs = torch.sigmoid(outputs).cpu().numpy() preds = (probs > 0.5).astype(int) # 确保是1维数组 if probs.ndim == 0: probs = np.array([probs]) if preds.ndim == 0: preds = np.array([preds]) all_probs.extend(probs.flatten()) all_preds.extend(preds.flatten()) else: # Regression predictions preds = outputs.cpu().numpy() if preds.ndim == 0: preds = np.array([preds]) all_preds.extend(preds.flatten()) # Labels labels = batch_y.cpu().numpy() if labels.ndim == 0: labels = np.array([labels]) all_labels.extend(labels.flatten()) # Convert to arrays y_true = np.array(all_labels) y_pred = np.array(all_preds) y_prob = np.array(all_probs) if task == 'classification' else None # Evaluate evaluator = ModelEvaluator(task) metrics = evaluator.evaluate(y_true, y_pred, y_prob) # Add backtesting metrics backtest = evaluator.backtest_strategy(y_true, y_pred) metrics.update({f'backtest_{k}': v for k, v in backtest.items()}) return metrics if __name__ == "__main__": # Test classification evaluation y_true = np.array([0, 1, 1, 0, 1, 0, 1, 1, 0, 0]) y_pred = np.array([0, 1, 1, 0, 0, 1, 1, 1, 0, 1]) y_prob = np.array([0.3, 0.8, 0.9, 0.2, 0.4, 0.6, 0.7, 0.85, 0.1, 0.55]) evaluator = ModelEvaluator('classification') metrics = evaluator.evaluate(y_true, y_pred, y_prob) print("Classification metrics:") for k, v in metrics.items(): print(f" {k}: {v}") # Test regression evaluation y_true_reg = np.random.randn(100) * 0.02 y_pred_reg = y_true_reg + np.random.randn(100) * 0.01 evaluator_reg = ModelEvaluator('regression') metrics_reg = evaluator_reg.evaluate(y_true_reg, y_pred_reg) print("\nRegression metrics:") for k, v in metrics_reg.items(): print(f" {k}: {v}")