benchmark_test.py

"""
Monte Carlo VaR Benchmark Suite
Comprehensive tests to measure performance improvements of Rust implementation
"""

import numpy as np
import time
import os
from datetime import datetime
from typing import Dict, List
from dataclasses import dataclass
from concurrent.futures import ProcessPoolExecutor

try:
    import psutil
    PSUTIL_AVAILABLE = True
except ImportError:
    PSUTIL_AVAILABLE = False

# numba wheels often lag new CPython releases, so treat it as optional and skip
# that comparison when it isn't installed.
try:
    import numba
    NUMBA_AVAILABLE = True
except ImportError:
    numba = None
    NUMBA_AVAILABLE = False

try:
    import var_mc  # Our Rust library
    RUST_AVAILABLE = True
except ImportError:
    RUST_AVAILABLE = False


def _numba_jit(**kwargs):
    """Apply numba.jit when available, otherwise return the function untouched."""
    def wrap(func):
        return numba.jit(**kwargs)(func) if NUMBA_AVAILABLE else func
    return wrap


def _rss_mb() -> float:
    """Current resident-set size in MB (0.0 if psutil is unavailable)."""
    if not PSUTIL_AVAILABLE:
        return 0.0
    return psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024


def _mp_simulate_chunk(args):
    """Worker for the multiprocessing benchmark. Must be a module-level function
    (not a nested closure) so it is picklable under the 'spawn' start method
    used on macOS and Windows."""
    chunk_size, daily_returns, L, weights, portfolio_value, time_horizon = args
    randn = np.random.standard_normal((chunk_size, len(weights)))
    correlated = randn @ L.T
    asset_returns = daily_returns[None, :] * time_horizon + correlated * np.sqrt(time_horizon)
    portfolio_returns = asset_returns @ weights
    return portfolio_value - portfolio_value * (1 + portfolio_returns)


@dataclass
class BenchmarkResult:
    """Store benchmark results for analysis"""
    method: str
    n_assets: int
    n_simulations: int
    time_horizon: int
    execution_time: float
    memory_usage: float
    var_95: float
    var_99: float
    cvar_95: float
    cpu_cores_used: int
    timestamp: datetime


class PortfolioGenerator:
    """Generate realistic portfolio parameters for testing"""
    
    @staticmethod
    def generate_portfolio(n_assets: int, seed: int = 42) -> Dict:
        """Generate portfolio with realistic parameters"""
        np.random.seed(seed)
        
        # Generate returns (annual)
        base_returns = np.random.uniform(0.02, 0.15, n_assets)
        
        # Generate correlation matrix using a random factor model. Each asset
        # gets factor exposure plus idiosyncratic variance so the matrix is
        # full-rank (positive definite) — a pure low-rank factor product is
        # singular and would make every Cholesky-based method fail.
        n_factors = max(1, min(5, n_assets // 3))
        factor_loadings = np.random.randn(n_assets, n_factors) * 0.6
        factor_cov = factor_loadings @ factor_loadings.T
        idiosyncratic = np.maximum(1.0 - np.diag(factor_cov), 0.05)
        correlation = factor_cov + np.diag(idiosyncratic)
        # Normalise to unit diagonal → a valid, positive-definite correlation matrix.
        inv_std = 1.0 / np.sqrt(np.diag(correlation))
        correlation = correlation * inv_std[:, None] * inv_std[None, :]
        
        # Generate volatilities (annual)
        volatilities = np.random.uniform(0.10, 0.40, n_assets)
        
        # Create covariance matrix
        covariance = correlation * volatilities[:, None] * volatilities[None, :]
        
        # Generate weights (random then normalized)
        weights = np.random.dirichlet(np.ones(n_assets) * 2)
        
        return {
            'weights': weights,
            'returns': base_returns,
            'covariance': covariance,
            'correlation': correlation,
            'volatilities': volatilities
        }


class VaRCalculators:
    """Different implementations of Monte Carlo VaR for benchmarking"""
    
    @staticmethod
    def numpy_basic(portfolio_value: float, weights: np.ndarray, 
                   returns: np.ndarray, covariance: np.ndarray,
                   time_horizon: int, confidence_levels: List[float],
                   n_simulations: int) -> Dict:
        """Basic NumPy implementation"""
        start_time = time.time()
        start_memory = _rss_mb()
        
        # Convert to daily
        daily_returns = returns / 252
        daily_cov = covariance / 252
        
        # Portfolio parameters
        portfolio_return = np.dot(weights, daily_returns)
        portfolio_variance = np.dot(weights, np.dot(daily_cov, weights))
        portfolio_std = np.sqrt(portfolio_variance)
        
        # Generate random returns
        random_returns = np.random.normal(
            portfolio_return * time_horizon,
            portfolio_std * np.sqrt(time_horizon),
            n_simulations
        )
        
        # Calculate portfolio values
        portfolio_values = portfolio_value * (1 + random_returns)
        losses = portfolio_value - portfolio_values
        
        # Calculate VaR and CVaR
        results = {}
        for conf in confidence_levels:
            var = np.percentile(losses, conf * 100)
            cvar = losses[losses >= var].mean()
            results[f'var_{int(conf*100)}'] = var
            results[f'cvar_{int(conf*100)}'] = cvar
        
        end_memory = _rss_mb()
        results['execution_time'] = time.time() - start_time
        results['memory_usage'] = end_memory - start_memory
        
        return results
    
    @staticmethod
    def numpy_optimized(portfolio_value: float, weights: np.ndarray, 
                       returns: np.ndarray, covariance: np.ndarray,
                       time_horizon: int, confidence_levels: List[float],
                       n_simulations: int) -> Dict:
        """Optimized NumPy with Cholesky decomposition"""
        start_time = time.time()
        start_memory = _rss_mb()
        
        # Convert to daily
        daily_returns = returns / 252
        daily_cov = covariance / 252
        
        # Cholesky decomposition for correlated returns
        L = np.linalg.cholesky(daily_cov)
        
        # Generate uncorrelated random numbers
        randn = np.random.standard_normal((n_simulations, len(weights)))
        
        # Create correlated returns
        correlated_randn = randn @ L.T
        
        # Calculate returns for time horizon
        asset_returns = (daily_returns[None, :] * time_horizon + 
                        correlated_randn * np.sqrt(time_horizon))
        
        # Portfolio returns
        portfolio_returns = asset_returns @ weights
        portfolio_values = portfolio_value * (1 + portfolio_returns)
        losses = portfolio_value - portfolio_values
        
        # Sort once for all percentiles
        losses.sort()
        
        results = {}
        for conf in confidence_levels:
            idx = int(conf * n_simulations)
            var = losses[idx]
            cvar = losses[idx:].mean()
            results[f'var_{int(conf*100)}'] = var
            results[f'cvar_{int(conf*100)}'] = cvar
        
        end_memory = _rss_mb()
        results['execution_time'] = time.time() - start_time
        results['memory_usage'] = end_memory - start_memory
        
        return results
    
    @staticmethod
    @_numba_jit(nopython=True, parallel=True)
    def _numba_simulation(daily_returns: np.ndarray, L: np.ndarray,
                         weights: np.ndarray, portfolio_value: float,
                         time_horizon: int, n_simulations: int) -> np.ndarray:
        """Numba-accelerated simulation kernel"""
        losses = np.empty(n_simulations)
        sqrt_time = np.sqrt(time_horizon)
        
        for i in numba.prange(n_simulations):
            # Generate random numbers
            randn = np.random.standard_normal(len(weights))
            
            # Correlated returns
            correlated_randn = L @ randn
            
            # Asset returns
            asset_returns = daily_returns * time_horizon + correlated_randn * sqrt_time
            
            # Portfolio return
            portfolio_return = np.dot(weights, asset_returns)
            portfolio_value_end = portfolio_value * (1 + portfolio_return)
            losses[i] = portfolio_value - portfolio_value_end
        
        return losses
    
    @staticmethod
    def numba_parallel(portfolio_value: float, weights: np.ndarray, 
                      returns: np.ndarray, covariance: np.ndarray,
                      time_horizon: int, confidence_levels: List[float],
                      n_simulations: int) -> Dict:
        """Numba JIT-compiled with parallel execution"""
        start_time = time.time()
        start_memory = _rss_mb()
        
        # Convert to daily
        daily_returns = returns / 252
        daily_cov = covariance / 252
        
        # Cholesky decomposition
        L = np.linalg.cholesky(daily_cov)
        
        # Run simulation
        losses = VaRCalculators._numba_simulation(
            daily_returns, L, weights, portfolio_value, 
            time_horizon, n_simulations
        )
        
        # Sort for percentiles
        losses.sort()
        
        results = {}
        for conf in confidence_levels:
            idx = int(conf * n_simulations)
            var = losses[idx]
            cvar = losses[idx:].mean()
            results[f'var_{int(conf*100)}'] = var
            results[f'cvar_{int(conf*100)}'] = cvar
        
        end_memory = _rss_mb()
        results['execution_time'] = time.time() - start_time
        results['memory_usage'] = end_memory - start_memory
        
        return results
    
    @staticmethod
    def multiprocessing_chunked(portfolio_value: float, weights: np.ndarray, 
                               returns: np.ndarray, covariance: np.ndarray,
                               time_horizon: int, confidence_levels: List[float],
                               n_simulations: int, n_processes: int = None) -> Dict:
        """Multiprocessing with chunked simulations"""
        start_time = time.time()
        start_memory = _rss_mb()
        
        if n_processes is None:
            n_processes = os.cpu_count()
        
        # Prepare parameters
        daily_returns = returns / 252
        daily_cov = covariance / 252
        L = np.linalg.cholesky(daily_cov)

        # Split work across processes. The worker is a module-level function
        # (`_mp_simulate_chunk`) rather than a nested closure so it can be
        # pickled under the 'spawn' start method used on macOS and Windows.
        chunk_sizes = [n_simulations // n_processes] * n_processes
        chunk_sizes[-1] += n_simulations % n_processes
        tasks = [(cs, daily_returns, L, weights, portfolio_value, time_horizon)
                 for cs in chunk_sizes]

        with ProcessPoolExecutor(max_workers=n_processes) as executor:
            loss_chunks = list(executor.map(_mp_simulate_chunk, tasks))
        
        losses = np.concatenate(loss_chunks)
        losses.sort()
        
        results = {}
        for conf in confidence_levels:
            idx = int(conf * n_simulations)
            var = losses[idx]
            cvar = losses[idx:].mean()
            results[f'var_{int(conf*100)}'] = var
            results[f'cvar_{int(conf*100)}'] = cvar
        
        end_memory = _rss_mb()
        results['execution_time'] = time.time() - start_time
        results['memory_usage'] = end_memory - start_memory
        
        return results


class VaRBenchmarkSuite:
    """Main benchmark suite for comparing implementations"""
    
    def __init__(self):
        self.results: List[BenchmarkResult] = []
        self.portfolio_configs = [
            {'n_assets': 10, 'name': 'Small Portfolio'},
            {'n_assets': 50, 'name': 'Medium Portfolio'},
            {'n_assets': 100, 'name': 'Large Portfolio'},
            {'n_assets': 500, 'name': 'Very Large Portfolio'},
        ]
        self.simulation_configs = [
            {'n_simulations': 10_000, 'name': '10K simulations'},
            {'n_simulations': 100_000, 'name': '100K simulations'},
            {'n_simulations': 1_000_000, 'name': '1M simulations'},
            {'n_simulations': 10_000_000, 'name': '10M simulations'},
        ]
        
    def run_benchmark(self, method_name: str, method_func, **kwargs):
        """Run a single benchmark configuration"""
        portfolio = kwargs['portfolio']
        n_simulations = kwargs['n_simulations']
        n_assets = kwargs['n_assets']
        
        print(f"Running {method_name}: {n_assets} assets, {n_simulations:,} simulations...")
        
        try:
            results = method_func(
                portfolio_value=1_000_000,
                weights=portfolio['weights'],
                returns=portfolio['returns'],
                covariance=portfolio['covariance'],
                time_horizon=10,
                confidence_levels=[0.95, 0.99],
                n_simulations=n_simulations
            )
            
            benchmark_result = BenchmarkResult(
                method=method_name,
                n_assets=n_assets,
                n_simulations=n_simulations,
                time_horizon=10,
                execution_time=results['execution_time'],
                memory_usage=results['memory_usage'],
                var_95=results['var_95'],
                var_99=results['var_99'],
                cvar_95=results['cvar_95'],
                cpu_cores_used=1 if 'parallel' not in method_name else os.cpu_count(),
                timestamp=datetime.now()
            )
            
            self.results.append(benchmark_result)
            print(f"  Time: {results['execution_time']:.3f}s, Memory: {results['memory_usage']:.1f}MB")
            print(f"  VaR 95%: ${results['var_95']:,.2f}, CVaR 95%: ${results['cvar_95']:,.2f}")
            
        except Exception as e:
            print(f"  Failed: {str(e)}")
    
    def run_all_benchmarks(self):
        """Run complete benchmark suite"""
        methods = [
            ('NumPy Basic', VaRCalculators.numpy_basic),
            ('NumPy Optimized', VaRCalculators.numpy_optimized),
            ('Multiprocessing', VaRCalculators.multiprocessing_chunked),
        ]

        if NUMBA_AVAILABLE:
            methods.insert(2, ('Numba Parallel', VaRCalculators.numba_parallel))
        else:
            print("(numba not installed — skipping the Numba Parallel comparison)")

        if RUST_AVAILABLE:
            methods.append(('Rust', var_mc.calculate_var))
        else:
            print("(var_mc not importable — build it with `maturin develop --release`)")
        
        # Warm up JIT compilation
        print("Warming up JIT compilation...")
        small_portfolio = PortfolioGenerator.generate_portfolio(5)
        for name, func in methods:
            if 'Numba' in name:
                try:
                    func(100_000, small_portfolio['weights'], small_portfolio['returns'],
                         small_portfolio['covariance'], 1, [0.95], 100)
                except:
                    pass
        
        print("\nStarting benchmarks...\n")
        
        for portfolio_config in self.portfolio_configs:
            n_assets = portfolio_config['n_assets']
            portfolio = PortfolioGenerator.generate_portfolio(n_assets)
            
            for sim_config in self.simulation_configs:
                n_simulations = sim_config['n_simulations']
                
                # Skip very large combinations that would take too long
                if n_assets * n_simulations > 500 * 1_000_000:
                    continue
                
                print(f"\n{portfolio_config['name']} - {sim_config['name']}")
                print("="*60)
                
                for method_name, method_func in methods:
                    self.run_benchmark(
                        method_name=method_name,
                        method_func=method_func,
                        portfolio=portfolio,
                        n_simulations=n_simulations,
                        n_assets=n_assets
                    )
    
    def analyze_results(self):
        """Analyze and visualize benchmark results"""
        # pandas + matplotlib are only needed for reporting/plotting, so import
        # them lazily — the benchmarks themselves run with just NumPy.
        import pandas as pd
        import matplotlib.pyplot as plt

        df = pd.DataFrame([
            {
                'method': r.method,
                'n_assets': r.n_assets,
                'n_simulations': r.n_simulations,
                'execution_time': r.execution_time,
                'memory_usage': r.memory_usage,
                'var_95': r.var_95,
                'simulations_per_second': r.n_simulations / r.execution_time
            }
            for r in self.results
        ])
        
        # Reference slices. Each subplot fixes one dimension so the pivot has a
        # unique index (the grid varies BOTH n_assets and n_simulations, so a
        # raw pivot on either alone would collide).
        REF_SIM = 1_000_000            # for the "by portfolio size" plots
        REF_ASSETS = 10                # for the "by simulation count" plot
        BASELINE = 'NumPy Optimized'   # fair same-model baseline (see README)
        by_size = df[df['n_simulations'] == REF_SIM]
        by_sims = df[df['n_assets'] == REF_ASSETS]

        # Create visualizations
        fig, axes = plt.subplots(2, 2, figsize=(15, 12))

        # 1. Execution time by portfolio size (at REF_SIM simulations)
        ax = axes[0, 0]
        by_size.pivot(index='n_assets', columns='method', values='execution_time').plot(ax=ax, marker='o')
        ax.set_xlabel('Number of Assets')
        ax.set_ylabel('Execution Time (seconds)')
        ax.set_title(f'Execution Time by Portfolio Size ({REF_SIM:,} sims)')
        ax.set_yscale('log')
        ax.grid(True, alpha=0.3)

        # 2. Throughput by simulation count (at REF_ASSETS assets)
        ax = axes[0, 1]
        by_sims.pivot(index='n_simulations', columns='method', values='simulations_per_second').plot(ax=ax, marker='o')
        ax.set_xlabel('Number of Simulations')
        ax.set_ylabel('Simulations per Second')
        ax.set_title(f'Throughput by Simulation Count ({REF_ASSETS} assets)')
        ax.set_xscale('log')
        ax.set_yscale('log')
        ax.grid(True, alpha=0.3)

        # 3. Memory usage by portfolio size (at REF_SIM simulations)
        ax = axes[1, 0]
        by_size.pivot(index='n_assets', columns='method', values='memory_usage').plot(ax=ax, marker='o')
        ax.set_xlabel('Number of Assets')
        ax.set_ylabel('Memory Usage (MB)')
        ax.set_title(f'Memory Usage by Portfolio Size ({REF_SIM:,} sims)')
        ax.grid(True, alpha=0.3)

        # 4. Speedup relative to the fair NumPy baseline (at REF_SIM simulations)
        ax = axes[1, 1]
        base_by_size = by_size[by_size['method'] == BASELINE].set_index('n_assets')['execution_time']
        speedups = []
        for method in by_size['method'].unique():
            if method == BASELINE:
                continue
            method_times = by_size[by_size['method'] == method].set_index('n_assets')['execution_time']
            for n_assets, val in (base_by_size / method_times).items():
                speedups.append((method, n_assets, val))
        speedup_df = pd.DataFrame(speedups, columns=['method', 'n_assets', 'speedup'])
        speedup_df.pivot(index='n_assets', columns='method', values='speedup').plot(ax=ax, kind='bar')
        ax.set_xlabel('Number of Assets')
        ax.set_ylabel(f'Speedup vs {BASELINE}')
        ax.set_title(f'Speedup vs {BASELINE} ({REF_SIM:,} sims)')
        ax.axhline(y=1, color='black', linestyle='--', alpha=0.5)
        ax.grid(True, alpha=0.3)

        plt.tight_layout()
        plt.savefig('var_benchmark_results.png', dpi=300)

        # Print summary statistics
        print("\n" + "=" * 80)
        print("BENCHMARK SUMMARY")
        print("=" * 80)

        baseline = df[df['method'] == BASELINE].set_index(['n_assets', 'n_simulations'])['execution_time']
        for method in df['method'].unique():
            method_df = df[df['method'] == method]
            avg_speedup = (baseline / method_df.set_index(['n_assets', 'n_simulations'])['execution_time']).mean()

            print(f"\n{method}:")
            print(f"  Average execution time: {method_df['execution_time'].mean():.3f}s")
            print(f"  Average throughput: {method_df['simulations_per_second'].mean():,.0f} sims/sec")
            if method != BASELINE:
                print(f"  Average speedup vs {BASELINE}: {avg_speedup:.1f}x")

        return df


# Example usage
if __name__ == "__main__":
    benchmark = VaRBenchmarkSuite()
    
    # Run all benchmarks
    benchmark.run_all_benchmarks()
    
    # Analyze and visualize results
    results_df = benchmark.analyze_results()
    
    # Save results for future reference
    results_df.to_csv('var_benchmark_results.csv', index=False)
    print(f"\nResults saved to var_benchmark_results.csv")