TabNet Financial Predictive Models

import numpy as np
import pandas as pd
import torch
from pytorch_tabnet.tab_model import TabNetRegressor, TabNetClassifier
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score
import yfinance as yf
import ta
import matplotlib.pyplot as plt
import seaborn as sns
import warnings
warnings.filterwarnings('ignore')


class FinancialDataProcessor:
    """Process and prepare financial data for TabNet models"""
    
    def __init__(self):
        self.scaler = StandardScaler()
        self.label_encoder = LabelEncoder()
        
    def fetch_stock_data(self, symbols, period='2y'):
        """Fetch stock data from Yahoo Finance"""
        data = {}
        for symbol in symbols:
            ticker = yf.Ticker(symbol)
            df = ticker.history(period=period)
            df['Symbol'] = symbol
            data[symbol] = df
        return pd.concat(data.values(), ignore_index=True)
    
    def create_technical_indicators(self, df):
        """Create technical indicators for stock data"""
        # Price-based indicators
        df['SMA_10'] = ta.trend.sma_indicator(df['Close'], window=10)
        df['SMA_20'] = ta.trend.sma_indicator(df['Close'], window=20)
        df['EMA_12'] = ta.trend.ema_indicator(df['Close'], window=12)
        df['EMA_26'] = ta.trend.ema_indicator(df['Close'], window=26)
        
        # Momentum indicators
        df['RSI'] = ta.momentum.rsi(df['Close'])
        df['MACD'] = ta.trend.macd_diff(df['Close'])
        df['MACD_signal'] = ta.trend.macd_signal(df['Close'])
        df['Stoch_K'] = ta.momentum.stoch(df['High'], df['Low'], df['Close'])
        df['Stoch_D'] = ta.momentum.stoch_signal(df['High'], df['Low'], df['Close'])
        
        # Volatility indicators
        df['BB_upper'] = ta.volatility.bollinger_hband(df['Close'])
        df['BB_lower'] = ta.volatility.bollinger_lband(df['Close'])
        df['BB_width'] = df['BB_upper'] - df['BB_lower']
        df['ATR'] = ta.volatility.average_true_range(df['High'], df['Low'], df['Close'])
        
        # Volume indicators
        df['Volume_SMA'] = ta.volume.volume_sma(df['Close'], df['Volume'])
        df['OBV'] = ta.volume.on_balance_volume(df['Close'], df['Volume'])
        
        # Price features
        df['Price_Change'] = df['Close'].pct_change()
        df['High_Low_Ratio'] = df['High'] / df['Low']
        df['Close_Open_Ratio'] = df['Close'] / df['Open']
        
        return df
    
    def create_target_variables(self, df, target_type='next_day_return'):
        """Create target variables for prediction"""
        if target_type == 'next_day_return':
            df['Target'] = df['Close'].pct_change().shift(-1)
        elif target_type == 'price_direction':
            df['Target'] = (df['Close'].shift(-1) > df['Close']).astype(int)
        elif target_type == 'volatility':
            df['Target'] = df['Close'].rolling(window=10).std().shift(-1)
        
        return df
    
    def prepare_features(self, df):
        """Prepare features for model training"""
        # Select numeric columns (excluding target and date-related columns)
        feature_cols = [col for col in df.columns if df[col].dtype in ['float64', 'int64'] 
                       and col not in ['Target']]
        
        # Handle missing values
        df[feature_cols] = df[feature_cols].fillna(method='ffill').fillna(method='bfill')
        
        # Encode categorical variables if any
        categorical_cols = [col for col in df.columns if df[col].dtype == 'object' 
                           and col not in ['Date']]
        
        for col in categorical_cols:
            df[col] = self.label_encoder.fit_transform(df[col].astype(str))
            feature_cols.append(col)
        
        return df, feature_cols


class TabNetStockPredictor:
    """TabNet model for stock price prediction"""
    
    def __init__(self, model_params=None):
        self.model_params = model_params or {
            'n_d': 64, 'n_a': 64, 'n_steps': 5,
            'gamma': 1.5, 'n_independent': 2, 'n_shared': 2,
            'lambda_sparse': 1e-4, 'optimizer_fn': torch.optim.Adam,
            'optimizer_params': dict(lr=0.02),
            'scheduler_params': dict(step_size=50, gamma=0.9),
            'scheduler_fn': torch.optim.lr_scheduler.StepLR,
            'verbose': 1
        }
        self.model = TabNetRegressor(**self.model_params)
        self.processor = FinancialDataProcessor()
        
    def prepare_data(self, symbols, period='2y', target_type='next_day_return'):
        """Prepare data for training"""
        # Fetch and process data
        df = self.processor.fetch_stock_data(symbols, period)
        df = self.processor.create_technical_indicators(df)
        df = self.processor.create_target_variables(df, target_type)
        
        # Prepare features
        df, feature_cols = self.processor.prepare_features(df)
        
        # Remove rows with NaN targets
        df = df.dropna(subset=['Target'])
        
        X = df[feature_cols].values
        y = df['Target'].values
        
        return X, y, feature_cols
    
    def train(self, X, y, test_size=0.2, validation_split=0.2):
        """Train the TabNet model"""
        # Split data
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=test_size, random_state=42, shuffle=False
        )
        
        # Further split training data for validation
        X_train, X_val, y_train, y_val = train_test_split(
            X_train, y_train, test_size=validation_split, random_state=42
        )
        
        # Train model
        self.model.fit(
            X_train, y_train,
            eval_set=[(X_val, y_val)],
            max_epochs=200,
            patience=20,
            batch_size=1024,
            virtual_batch_size=128,
            drop_last=False
        )
        
        return X_test, y_test
    
    def evaluate(self, X_test, y_test):
        """Evaluate model performance"""
        predictions = self.model.predict(X_test)
        
        metrics = {
            'MAE': mean_absolute_error(y_test, predictions),
            'MSE': mean_squared_error(y_test, predictions),
            'RMSE': np.sqrt(mean_squared_error(y_test, predictions)),
            'R2': r2_score(y_test, predictions),
            'Directional_Accuracy': np.mean(np.sign(y_test) == np.sign(predictions))
        }
        
        return metrics, predictions
    
    def plot_feature_importance(self, feature_names):
        """Plot feature importance from TabNet"""
        importance = self.model.feature_importances_
        
        plt.figure(figsize=(12, 8))
        indices = np.argsort(importance)[-20:]  # Top 20 features
        plt.barh(range(len(indices)), importance[indices])
        plt.yticks(range(len(indices)), [feature_names[i] for i in indices])
        plt.xlabel('Feature Importance')
        plt.title('TabNet Feature Importance - Top 20 Features')
        plt.tight_layout()
        plt.show()


class TabNetCreditRiskClassifier:
    """TabNet model for credit risk classification"""
    
    def __init__(self, model_params=None):
        self.model_params = model_params or {
            'n_d': 32, 'n_a': 32, 'n_steps': 4,
            'gamma': 1.3, 'n_independent': 2, 'n_shared': 2,
            'lambda_sparse': 1e-3, 'optimizer_fn': torch.optim.Adam,
            'optimizer_params': dict(lr=0.02),
            'scheduler_params': dict(step_size=50, gamma=0.9),
            'scheduler_fn': torch.optim.lr_scheduler.StepLR,
            'verbose': 1
        }
        self.model = TabNetClassifier(**self.model_params)
        
    def create_synthetic_credit_data(self, n_samples=10000):
        """Create synthetic credit data for demonstration"""
        np.random.seed(42)
        
        # Generate features
        data = {
            'age': np.random.normal(40, 12, n_samples),
            'income': np.random.exponential(50000, n_samples),
            'credit_score': np.random.normal(650, 100, n_samples),
            'debt_to_income': np.random.beta(2, 5, n_samples),
            'employment_length': np.random.exponential(5, n_samples),
            'loan_amount': np.random.exponential(25000, n_samples),
            'loan_term': np.random.choice([12, 24, 36, 48, 60], n_samples),
            'num_credit_lines': np.random.poisson(3, n_samples),
            'num_delinquencies': np.random.poisson(0.5, n_samples),
            'home_ownership': np.random.choice([0, 1, 2], n_samples, p=[0.3, 0.6, 0.1])
        }
        
        df = pd.DataFrame(data)
        
        # Create target based on logical rules
        risk_score = (
            -0.01 * df['credit_score'] +
            0.02 * df['debt_to_income'] * 100 +
            0.001 * df['loan_amount'] / df['income'] * 100 +
            0.1 * df['num_delinquencies'] +
            np.random.normal(0, 0.5, n_samples)
        )
        
        df['default'] = (risk_score > np.percentile(risk_score, 80)).astype(int)
        
        return df
    
    def train(self, df, target_col='default', test_size=0.2):
        """Train the credit risk model"""
        feature_cols = [col for col in df.columns if col != target_col]
        
        X = df[feature_cols].values
        y = df[target_col].values
        
        # Split data
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=test_size, random_state=42, stratify=y
        )
        
        # Further split for validation
        X_train, X_val, y_train, y_val = train_test_split(
            X_train, y_train, test_size=0.2, random_state=42, stratify=y_train
        )
        
        # Train model
        self.model.fit(
            X_train, y_train,
            eval_set=[(X_val, y_val)],
            max_epochs=200,
            patience=20,
            batch_size=1024,
            virtual_batch_size=128,
            drop_last=False
        )
        
        return X_test, y_test, feature_cols
    
    def evaluate(self, X_test, y_test):
        """Evaluate classification performance"""
        predictions = self.model.predict(X_test)
        probabilities = self.model.predict_proba(X_test)[:, 1]
        
        metrics = {
            'Accuracy': accuracy_score(y_test, predictions),
            'Precision': precision_score(y_test, predictions),
            'Recall': recall_score(y_test, predictions),
            'F1': f1_score(y_test, predictions),
            'AUC_ROC': roc_auc_score(y_test, probabilities)
        }
        
        return metrics, predictions, probabilities


class TabNetPortfolioAnalyzer:
    """TabNet model for portfolio risk analysis"""
    
    def __init__(self):
        self.processor = FinancialDataProcessor()
        
    def create_portfolio_features(self, symbols, weights=None, period='2y'):
        """Create portfolio-level features"""
        if weights is None:
            weights = np.ones(len(symbols)) / len(symbols)
        
        # Fetch individual stock data
        portfolio_data = []
        for symbol in symbols:
            ticker = yf.Ticker(symbol)
            df = ticker.history(period=period)
            df['Symbol'] = symbol
            df['Weight'] = weights[symbols.index(symbol)]
            portfolio_data.append(df)
        
        # Combine data
        combined_df = pd.concat(portfolio_data, ignore_index=True)
        
        # Calculate portfolio returns
        portfolio_returns = []
        dates = combined_df.index.get_level_values(0).unique()
        
        for date in dates:
            day_data = combined_df.loc[date]
            if len(day_data) == len(symbols):
                returns = day_data['Close'].pct_change().fillna(0)
                portfolio_return = np.sum(returns * weights)
                portfolio_returns.append(portfolio_return)
        
        return np.array(portfolio_returns)


def main():
    """Main function to demonstrate TabNet financial models"""
    
    print("🚀 TabNet Financial Predictive Models Demo")
    print("=" * 50)
    
    # 1. Stock Price Prediction
    print("\n📈 Stock Price Prediction")
    print("-" * 30)
    
    stock_predictor = TabNetStockPredictor()
    symbols = ['AAPL', 'GOOGL', 'MSFT', 'AMZN']
    
    try:
        X, y, feature_names = stock_predictor.prepare_data(symbols, period='1y')
        X_test, y_test = stock_predictor.train(X, y)
        metrics, predictions = stock_predictor.evaluate(X_test, y_test)
        
        print("Stock Prediction Metrics:")
        for metric, value in metrics.items():
            print(f"  {metric}: {value:.4f}")
        
        # Plot feature importance
        stock_predictor.plot_feature_importance(feature_names)
        
    except Exception as e:
        print(f"Stock prediction demo failed: {e}")
    
    # 2. Credit Risk Classification
    print("\n🏦 Credit Risk Assessment")
    print("-" * 30)
    
    credit_classifier = TabNetCreditRiskClassifier()
    
    try:
        # Generate synthetic credit data
        credit_data = credit_classifier.create_synthetic_credit_data()
        X_test, y_test, feature_cols = credit_classifier.train(credit_data)
        metrics, predictions, probabilities = credit_classifier.evaluate(X_test, y_test)
        
        print("Credit Risk Metrics:")
        for metric, value in metrics.items():
            print(f"  {metric}: {value:.4f}")
        
    except Exception as e:
        print(f"Credit risk demo failed: {e}")
    
    print("\n✅ Demo completed successfully!")
    print("\nNext steps:")
    print("1. Customize model parameters for your specific use case")
    print("2. Implement additional evaluation metrics")
    print("3. Add model interpretability analysis")
    print("4. Set up automated model retraining pipeline")


if __name__ == "__main__":
    main()