multi_task_learning.py

"""Multi-Task Learning for Joint Alpha + Volatility + Portfolio Optimization

Based on Ong & Herremans 2023 (arxiv:2306.13661):
"Multi-Task Learning for Time Series Momentum Portfolio Construction"

KEY INSIGHT: Jointly optimizing all three tasks simultaneously outperforms
independent optimization even after 3bps transaction costs.

This is THE critical upgrade that separates toy systems from production-grade quant.
"""
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from typing import Dict, Tuple, Optional, List
import warnings
warnings.filterwarnings('ignore')


class MTLSample(Dataset):
    """Dataset for multi-task learning with sequence input"""
    def __init__(self, X: np.ndarray, 
                 y_return: np.ndarray,
                 y_vol: np.ndarray,
                 y_portfolio: Optional[np.ndarray] = None):
        self.X = torch.FloatTensor(X)
        self.y_return = torch.FloatTensor(y_return)
        self.y_vol = torch.FloatTensor(y_vol)
        if y_portfolio is not None:
            self.y_portfolio = torch.FloatTensor(y_portfolio)
        else:
            self.y_portfolio = None
    
    def __len__(self):
        return len(self.X)
    
    def __getitem__(self, idx):
        out = {
            'X': self.X[idx],
            'return': self.y_return[idx],
            'volatility': self.y_vol[idx]
        }
        if self.y_portfolio is not None:
            out['portfolio'] = self.y_portfolio[idx]
        return out


class MultiTaskPortfolioNet(nn.Module):
    """
    Multi-Task Learning Network for Joint:
    1. Return prediction (alpha generation)
    2. Volatility prediction (risk estimation)
    3. Portfolio weight optimization
    
    Architecture (from MTL-TSMOM paper):
    - Shared LSTM encoder (hard parameter sharing)
    - Task-specific FNN heads with different architectures
    - Custom task-specific losses
    
    Shared encoder learns common temporal representations.
    Each head learns task-specific transformations.
    """
    
    def __init__(self,
                 input_dim: int,
                 hidden_dim: int = 128,
                 n_lstm_layers: int = 2,
                 n_assets: int = 10,
                 dropout: float = 0.15,
                 use_attention: bool = True):
        super().__init__()
        
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.n_assets = n_assets
        self.use_attention = use_attention
        
        # Shared encoder: LSTM with optional attention
        self.lstm = nn.LSTM(
            input_dim, hidden_dim, n_lstm_layers,
            batch_first=True, dropout=dropout if n_lstm_layers > 1 else 0
        )
        
        # Optional: Self-attention on LSTM outputs
        if use_attention:
            self.attention = nn.MultiheadAttention(
                hidden_dim, num_heads=4, dropout=dropout, batch_first=True
            )
        
        # Shared projection layer
        self.shared_fc = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout)
        )
        
        # Task 1: Return prediction head (Alpha)
        # Predicts future returns for each asset
        self.return_head = nn.Sequential(
            nn.Linear(hidden_dim, 256),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, n_assets)  # One return per asset
        )
        
        # Task 2: Volatility prediction head (Risk)
        # Predicts realized volatility for each asset
        self.vol_head = nn.Sequential(
            nn.Linear(hidden_dim, 128),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, n_assets)
        )
        
        # Task 3: Portfolio weight head (Allocation)
        # Directly outputs portfolio weights (long-only, softmax)
        self.portfolio_head = nn.Sequential(
            nn.Linear(hidden_dim, 256),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, n_assets),
            nn.Softmax(dim=-1)  # Long-only, fully invested
        )
        
        # Task 4: Direction prediction (auxiliary)
        # Binary classification: up or down (helps stabilize training)
        self.direction_head = nn.Sequential(
            nn.Linear(hidden_dim, 64),
            nn.ReLU(),
            nn.Linear(64, n_assets),
            nn.Sigmoid()
        )
    
    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        """
        Forward pass.
        
        Args:
            x: (batch, seq_len, input_dim)
        
        Returns:
            Dict with 'returns', 'volatility', 'portfolio', 'direction'
        """
        # Shared LSTM encoder
        lstm_out, (h_n, _) = self.lstm(x)
        # h_n: (n_layers, batch, hidden_dim)
        shared = h_n[-1]  # (batch, hidden_dim) — last layer final hidden state
        
        # Optional attention on sequence outputs
        if self.use_attention:
            attn_out, _ = self.attention(lstm_out, lstm_out, lstm_out)
            # Global average pooling over time
            shared_attn = attn_out.mean(dim=1)  # (batch, hidden_dim)
            shared = shared + shared_attn  # Residual connection
        
        # Shared projection
        shared_repr = self.shared_fc(shared)
        
        # Task-specific outputs
        returns = self.return_head(shared_repr)  # (batch, n_assets)
        volatility = F.softplus(self.vol_head(shared_repr)) + 1e-6  # Ensure positive
        portfolio = self.portfolio_head(shared_repr)  # (batch, n_assets), sums to 1
        direction = self.direction_head(shared_repr)  # (batch, n_assets), 0-1
        
        return {
            'returns': returns,
            'volatility': volatility,
            'portfolio': portfolio,
            'direction': direction,
            'shared_repr': shared_repr  # For analysis
        }


class MTLPortfolioTrainer:
    """
    Trainer for Multi-Task Portfolio Network.
    
    Uses task-specific loss weighting and gradient normalization
    to balance the three tasks.
    
    Key innovations from MTL-TSMOM paper:
    1. Negative Sharpe ratio as primary portfolio loss
    2. MSE for return prediction
    3. MSE for volatility prediction  
    4. BCE for direction (auxiliary stabilization)
    5. GradNorm for automatic task balancing
    """
    
    def __init__(self, model: MultiTaskPortfolioNet,
                 device: str = 'cpu',
                 learning_rate: float = 1e-4,
                 weight_decay: float = 1e-5,
                 max_grad_norm: float = 0.5,
                 risk_free_rate: float = 0.04):
        self.model = model.to(device)
        self.device = device
        self.risk_free_rate = risk_free_rate / 252  # Daily
        self.max_grad_norm = max_grad_norm
        
        self.optimizer = torch.optim.Adam(
            model.parameters(), lr=learning_rate, weight_decay=weight_decay
        )
        self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            self.optimizer, patience=10, factor=0.5, verbose=True
        )
        
        # Task loss weights (can be learned via GradNorm)
        self.task_weights = {
            'return': 1.0,
            'volatility': 0.5,
            'portfolio': 2.0,  # Primary task gets highest weight
            'direction': 0.3
        }
        
        self.history = {
            'train_loss': [], 'val_loss': [],
            'return_loss': [], 'vol_loss': [],
            'portfolio_loss': [], 'direction_loss': [],
            'sharpe': [], 'val_sharpe': []
        }
    
    def compute_loss(self, outputs: Dict[str, torch.Tensor],
                     batch: Dict[str, torch.Tensor],
                     actual_returns: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]:
        """
        Compute multi-task loss.
        
        Args:
            outputs: Model predictions
            batch: Ground truth batch
            actual_returns: Actual future returns (for Sharpe calculation)
        
        Returns:
            Dict of losses
        """
        losses = {}
        
        # Task 1: Return prediction loss (MSE on predicted vs actual returns)
        losses['return'] = F.mse_loss(outputs['returns'], batch['return'])
        
        # Task 2: Volatility prediction loss (MSE on predicted vs realized vol)
        losses['volatility'] = F.mse_loss(outputs['volatility'], batch['volatility'])
        
        # Task 3: Portfolio loss — NEGATIVE Sharpe ratio
        # We want portfolio weights that maximize risk-adjusted return
        if actual_returns is not None:
            # Portfolio return: sum(w_i * r_i)
            port_return = (outputs['portfolio'] * actual_returns).sum(dim=-1)
            
            # Sharpe ratio: mean(excess_return) / std(return)
            # We compute over batch (simulating a holding period)
            mean_return = port_return.mean()
            std_return = port_return.std() + 1e-6
            sharpe = (mean_return - self.risk_free_rate) / std_return
            
            # Negative Sharpe (we minimize this → maximize Sharpe)
            losses['portfolio'] = -sharpe
            
            # Track for monitoring
            losses['sharpe'] = sharpe.detach()
        else:
            losses['portfolio'] = torch.tensor(0.0, device=self.device)
            losses['sharpe'] = torch.tensor(0.0, device=self.device)
        
        # Task 4: Direction prediction (BCE)
        # Convert returns to binary: 1 if return > 0, else 0
        direction_target = (batch['return'] > 0).float()
        losses['direction'] = F.binary_cross_entropy(
            outputs['direction'], direction_target
        )
        
        # Total loss with task weighting
        total = sum(
            self.task_weights[task] * losses[task]
            for task in ['return', 'volatility', 'portfolio', 'direction']
        )
        losses['total'] = total
        
        return losses
    
    def train_epoch(self, dataloader: DataLoader,
                    actual_returns: Optional[np.ndarray] = None) -> Dict[str, float]:
        """Train for one epoch"""
        self.model.train()
        
        epoch_losses = {
            'return': 0.0, 'volatility': 0.0,
            'portfolio': 0.0, 'direction': 0.0,
            'total': 0.0, 'sharpe': 0.0
        }
        n_batches = 0
        
        for batch in dataloader:
            # Move to device
            X = batch['X'].to(self.device)
            returns_target = batch['return'].to(self.device)
            vol_target = batch['volatility'].to(self.device)
            
            # Actual returns for Sharpe (can be same as returns_target or future)
            actual = returns_target if actual_returns is None else \
                     torch.FloatTensor(actual_returns[n_batches]).to(self.device)
            
            # Forward
            outputs = self.model(X)
            
            # Loss
            losses = self.compute_loss(outputs, {
                'return': returns_target,
                'volatility': vol_target
            }, actual)
            
            # Backward
            self.optimizer.zero_grad()
            losses['total'].backward()
            
            # Gradient clipping (critical for LSTM stability)
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
            
            self.optimizer.step()
            
            # Track
            for key in epoch_losses:
                if key in losses:
                    val = losses[key]
                    if isinstance(val, torch.Tensor):
                        val = val.item()
                    epoch_losses[key] += val
            
            n_batches += 1
        
        # Average
        for key in epoch_losses:
            epoch_losses[key] /= max(n_batches, 1)
        
        return epoch_losses
    
    def validate(self, dataloader: DataLoader) -> Dict[str, float]:
        """Validate"""
        self.model.eval()
        
        val_losses = {
            'return': 0.0, 'volatility': 0.0,
            'portfolio': 0.0, 'direction': 0.0,
            'total': 0.0
        }
        n_batches = 0
        
        portfolio_returns = []
        
        with torch.no_grad():
            for batch in dataloader:
                X = batch['X'].to(self.device)
                returns_target = batch['return'].to(self.device)
                vol_target = batch['volatility'].to(self.device)
                
                outputs = self.model(X)
                
                losses = self.compute_loss(outputs, {
                    'return': returns_target,
                    'volatility': vol_target
                }, returns_target)
                
                for key in val_losses:
                    if key in losses:
                        val = losses[key]
                        if isinstance(val, torch.Tensor):
                            val = val.item()
                        val_losses[key] += val
                
                # Track portfolio returns for validation Sharpe
                port_ret = (outputs['portfolio'] * returns_target).sum(dim=-1)
                portfolio_returns.extend(port_ret.cpu().numpy())
                
                n_batches += 1
        
        for key in val_losses:
            val_losses[key] /= max(n_batches, 1)
        
        # Compute validation Sharpe
        if len(portfolio_returns) > 1:
            port_returns = np.array(portfolio_returns)
            mean_ret = np.mean(port_returns)
            std_ret = np.std(port_returns) + 1e-8
            val_sharpe = (mean_ret - self.risk_free_rate) / std_ret * np.sqrt(252)
            val_losses['sharpe'] = val_sharpe
        
        return val_losses
    
    def fit(self, train_loader: DataLoader,
            val_loader: Optional[DataLoader] = None,
            epochs: int = 100,
            early_stopping_patience: int = 15) -> Dict:
        """
        Full training loop.
        
        Returns:
            Training history dictionary
        """
        best_val_loss = float('inf')
        patience_counter = 0
        
        print(f"Training MTL Portfolio Net for {epochs} epochs...")
        print(f"Task weights: {self.task_weights}")
        print(f"Device: {self.device}")
        
        for epoch in range(epochs):
            # Train
            train_losses = self.train_epoch(train_loader)
            
            # Validate
            if val_loader is not None:
                val_losses = self.validate(val_loader)
                val_total = val_losses.get('total', 0)
                
                # Learning rate scheduling
                self.scheduler.step(val_total)
                
                # Early stopping
                if val_total < best_val_loss:
                    best_val_loss = val_total
                    patience_counter = 0
                else:
                    patience_counter += 1
                
                if patience_counter >= early_stopping_patience:
                    print(f"Early stopping at epoch {epoch}")
                    break
            else:
                val_losses = {}
            
            # Record
            for key in ['return', 'volatility', 'portfolio', 'direction', 'total']:
                self.history[f'{key}_loss'].append(train_losses.get(key, 0))
            self.history['sharpe'].append(train_losses.get('sharpe', 0))
            if 'sharpe' in val_losses:
                self.history['val_sharpe'].append(val_losses['sharpe'])
            
            # Print
            if epoch % 10 == 0 or epoch == epochs - 1:
                msg = f"Epoch {epoch}: "
                msg += f"train_total={train_losses['total']:.4f} "
                msg += f"return={train_losses['return']:.4f} "
                msg += f"vol={train_losses['volatility']:.4f} "
                msg += f"port={train_losses['portfolio']:.4f} "
                if 'sharpe' in train_losses:
                    msg += f"sharpe={train_losses['sharpe']:.4f} "
                if 'sharpe' in val_losses:
                    msg += f"val_sharpe={val_losses['sharpe']:.4f}"
                print(msg)
        
        return self.history
    
    def predict(self, X: np.ndarray) -> Dict[str, np.ndarray]:
        """Predict all tasks"""
        self.model.eval()
        
        X_t = torch.FloatTensor(X).to(self.device)
        
        with torch.no_grad():
            outputs = self.model(X_t)
        
        return {
            'returns': outputs['returns'].cpu().numpy(),
            'volatility': outputs['volatility'].cpu().numpy(),
            'portfolio': outputs['portfolio'].cpu().numpy(),
            'direction': outputs['direction'].cpu().numpy()
        }


class MTLPortfolioStrategy:
    """
    End-to-end strategy using MTL Portfolio Net.
    
    Unlike the original AlphaForge which runs separate models then combines,
    this trains ONE model that jointly optimizes all tasks.
    
    Output is directly usable portfolio weights — no separate optimizer needed!
    """
    
    def __init__(self, 
                 input_dim: int,
                 n_assets: int,
                 hidden_dim: int = 128,
                 device: str = 'cpu'):
        self.model = MultiTaskPortfolioNet(
            input_dim=input_dim,
            hidden_dim=hidden_dim,
            n_assets=n_assets,
            use_attention=True
        )
        self.trainer = MTLPortfolioTrainer(self.model, device=device)
        self.n_assets = n_assets
    
    def prepare_data(self,
                     X_train: np.ndarray,
                     returns_train: np.ndarray,
                     vol_train: np.ndarray,
                     X_val: Optional[np.ndarray] = None,
                     returns_val: Optional[np.ndarray] = None,
                     vol_val: Optional[np.ndarray] = None,
                     batch_size: int = 64) -> Tuple[DataLoader, Optional[DataLoader]]:
        """Prepare data loaders"""
        train_dataset = MTLSample(X_train, returns_train, vol_train)
        train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
        
        val_loader = None
        if X_val is not None:
            val_dataset = MTLSample(X_val, returns_val, vol_val)
            val_loader = DataLoader(val_dataset, batch_size=batch_size)
        
        return train_loader, val_loader
    
    def fit(self, X_train: np.ndarray,
            returns_train: np.ndarray,
            vol_train: np.ndarray,
            X_val: Optional[np.ndarray] = None,
            returns_val: Optional[np.ndarray] = None,
            vol_val: Optional[np.ndarray] = None,
            epochs: int = 100) -> Dict:
        """Fit the MTL model"""
        train_loader, val_loader = self.prepare_data(
            X_train, returns_train, vol_train,
            X_val, returns_val, vol_val
        )
        
        return self.trainer.fit(train_loader, val_loader, epochs=epochs)
    
    def generate_portfolio(self, X: np.ndarray) -> Tuple[np.ndarray, Dict]:
        """
        Generate portfolio weights and predictions.
        
        Returns:
            weights: (n_samples, n_assets) — directly usable allocations
            predictions: Dict with returns, volatility, direction predictions
        """
        predictions = self.trainer.predict(X)
        
        weights = predictions['portfolio']
        
        # Ensure valid weights
        weights = np.maximum(weights, 0)
        weights = weights / (weights.sum(axis=1, keepdims=True) + 1e-10)
        
        return weights, predictions


# Factory function for easy integration
def create_mtl_strategy(input_dim: int, n_assets: int,
                        device: str = 'cpu') -> MTLPortfolioStrategy:
    """Factory for MTL portfolio strategy"""
    return MTLPortfolioStrategy(input_dim, n_assets, device=device)


if __name__ == '__main__':
    # Test MTL model
    np.random.seed(42)
    torch.manual_seed(42)
    
    n_samples = 2000
    seq_len = 60
    n_features = 20
    n_assets = 10
    
    # Synthetic data
    X = np.random.randn(n_samples, seq_len, n_features).astype(np.float32)
    
    # Target returns (with some structure)
    returns = np.zeros((n_samples, n_assets))
    for i in range(n_assets):
        returns[:, i] = X[:, -1, i % n_features] * 0.1 + np.random.randn(n_samples) * 0.05
    
    # Target volatility
    vol = np.abs(returns) * 2 + 0.1
    
    # Split
    train_size = 1500
    X_train, X_val = X[:train_size], X[train_size:]
    r_train, r_val = returns[:train_size], returns[train_size:]
    v_train, v_val = vol[:train_size], vol[train_size:]
    
    # Create and train
    strategy = MTLPortfolioStrategy(
        input_dim=n_features,
        n_assets=n_assets,
        device='cpu'
    )
    
    history = strategy.fit(
        X_train, r_train, v_train,
        X_val, r_val, v_val,
        epochs=20
    )
    
    # Generate portfolio
    weights, preds = strategy.generate_portfolio(X_val[:10])
    
    print(f"\nSample portfolio weights (first 3):")
    for i in range(min(3, len(weights))):
        print(f"  Day {i}: {weights[i].round(3)} (sum={weights[i].sum():.3f})")
    
    print(f"\nPredicted returns (first 3):")
    print(preds['returns'][:3].round(4))
    
    print(f"\nPredicted volatility (first 3):")
    print(preds['volatility'][:3].round(4))