rustytorch/demos/rtx-portfolio-demo/src/optimizer.rs

//! Portfolio optimizer using Modern Portfolio Theory

use crate::constraints::{project_to_constraints, validate_weights};
use crate::covariance::{estimate_covariance_matrix, portfolio_variance, risk_contributions};
use crate::{
    EfficientFrontier, FrontierPoint, OptimizationObjective, OptimizationResult, OptimizerStatus,
    PortfolioConfig, PortfolioError, Result,
};
use nalgebra::DMatrix;
use std::time::Instant;

/// Portfolio optimizer using mean-variance optimization
pub struct PortfolioOptimizer {
    config: Option<PortfolioConfig>,
    covariance_matrix: Option<DMatrix<f64>>,
    optimization_count: u64,
    total_optimization_time_ms: f64,
}

impl PortfolioOptimizer {
    /// Create a new portfolio optimizer
    #[must_use]
    pub fn new() -> Self {
        Self {
            config: None,
            covariance_matrix: None,
            optimization_count: 0,
            total_optimization_time_ms: 0.0,
        }
    }

    /// Initialize with portfolio configuration
    pub fn initialize(&mut self, config: PortfolioConfig) -> Result<()> {
        if config.assets.len() < 2 {
            return Err(PortfolioError::InsufficientData(
                "Need at least 2 assets for optimization".to_string(),
            ));
        }

        let covariance = estimate_covariance_matrix(&config.assets)?;
        self.covariance_matrix = Some(covariance);
        self.config = Some(config);
        Ok(())
    }

    /// Optimize portfolio according to configured objective
    pub fn optimize(&mut self) -> Result<OptimizationResult> {
        let start = Instant::now();

        let config = self.config.as_ref().ok_or_else(|| {
            PortfolioError::InvalidConfig("Optimizer not initialized".to_string())
        })?;

        let covariance = self.covariance_matrix.as_ref().ok_or_else(|| {
            PortfolioError::InvalidConfig("Covariance matrix not computed".to_string())
        })?;

        let result = match config.objective {
            OptimizationObjective::MaxSharpe => self.optimize_max_sharpe(config, covariance)?,
            OptimizationObjective::MinVariance => self.optimize_min_variance(config, covariance)?,
            OptimizationObjective::RiskParity => self.optimize_risk_parity(config, covariance)?,
            OptimizationObjective::MaxReturn => self.optimize_max_return(config, covariance)?,
        };

        let processing_time_ms = start.elapsed().as_secs_f64() * 1000.0;
        self.optimization_count += 1;
        self.total_optimization_time_ms += processing_time_ms;

        Ok(OptimizationResult {
            processing_time_ms,
            ..result
        })
    }

    /// Compute efficient frontier
    pub fn compute_efficient_frontier(&self, num_points: usize) -> Result<EfficientFrontier> {
        let config = self.config.as_ref().ok_or_else(|| {
            PortfolioError::InvalidConfig("Optimizer not initialized".to_string())
        })?;

        let covariance = self.covariance_matrix.as_ref().ok_or_else(|| {
            PortfolioError::InvalidConfig("Covariance matrix not computed".to_string())
        })?;

        let min_return = config
            .assets
            .iter()
            .map(|a| a.expected_return)
            .min_by(|a, b| a.partial_cmp(b).unwrap())
            .unwrap_or(0.0);
        let max_return = config
            .assets
            .iter()
            .map(|a| a.expected_return)
            .max_by(|a, b| a.partial_cmp(b).unwrap())
            .unwrap_or(0.10);

        let mut points = Vec::with_capacity(num_points);
        let mut max_sharpe_index = 0;
        let mut max_sharpe = f64::NEG_INFINITY;
        let mut min_variance_index = 0;
        let mut min_variance = f64::INFINITY;

        for i in 0..num_points {
            let target_return =
                min_return + (max_return - min_return) * (i as f64) / ((num_points - 1) as f64);

            let mut temp_config = config.clone();
            temp_config.objective = OptimizationObjective::MinVariance;
            temp_config.constraints.target_return = Some(target_return);

            if let Ok(result) = self.optimize_min_variance(&temp_config, covariance) {
                let point = FrontierPoint {
                    expected_return: result.expected_return,
                    volatility: result.volatility,
                    sharpe_ratio: result.sharpe_ratio,
                    weights: result.weights,
                };

                if result.sharpe_ratio > max_sharpe {
                    max_sharpe = result.sharpe_ratio;
                    max_sharpe_index = i;
                }

                if result.volatility * result.volatility < min_variance {
                    min_variance = result.volatility * result.volatility;
                    min_variance_index = i;
                }

                points.push(point);
            }
        }

        if points.is_empty() {
            return Err(PortfolioError::OptimizationFailed(
                "Could not compute efficient frontier".to_string(),
            ));
        }

        Ok(EfficientFrontier {
            points,
            max_sharpe_index,
            min_variance_index,
        })
    }

    /// Get optimizer status
    #[must_use]
    pub fn status(&self) -> OptimizerStatus {
        let num_assets = self.config.as_ref().map_or(0, |c| c.assets.len());

        let avg_time = if self.optimization_count > 0 {
            self.total_optimization_time_ms / (self.optimization_count as f64)
        } else {
            0.0
        };

        OptimizerStatus {
            initialized: self.config.is_some(),
            num_assets,
            device: "CPU".to_string(),
            optimization_count: self.optimization_count,
            avg_optimization_time_ms: avg_time,
        }
    }

    /// Reset optimizer state
    pub fn reset(&mut self) {
        self.config = None;
        self.covariance_matrix = None;
        self.optimization_count = 0;
        self.total_optimization_time_ms = 0.0;
    }

    fn optimize_max_sharpe(
        &self,
        config: &PortfolioConfig,
        covariance: &DMatrix<f64>,
    ) -> Result<OptimizationResult> {
        let n = config.assets.len();
        let mut best_weights = vec![1.0 / (n as f64); n];
        let mut best_sharpe = f64::NEG_INFINITY;

        let max_iterations = 1000;
        let learning_rate = 0.01;

        for _ in 0..max_iterations {
            let expected_return: f64 = config
                .assets
                .iter()
                .enumerate()
                .map(|(i, a)| best_weights[i] * a.expected_return)
                .sum();

            let variance = portfolio_variance(&best_weights, covariance)?;
            let volatility = variance.sqrt();

            if volatility == 0.0 {
                break;
            }

            let sharpe_ratio = (expected_return - config.risk_free_rate) / volatility;

            if sharpe_ratio > best_sharpe {
                best_sharpe = sharpe_ratio;
            }

            for i in 0..n {
                let mut marginal_contribution = 0.0;
                for j in 0..n {
                    marginal_contribution += best_weights[j] * covariance[(i, j)];
                }

                let gradient = (config.assets[i].expected_return - config.risk_free_rate)
                    / volatility
                    - sharpe_ratio * marginal_contribution / volatility;

                best_weights[i] += learning_rate * gradient;
            }

            project_to_constraints(&mut best_weights, &config.assets, &config.constraints)?;
        }

        self.build_result(config, &best_weights, covariance)
    }

    fn optimize_min_variance(
        &self,
        config: &PortfolioConfig,
        covariance: &DMatrix<f64>,
    ) -> Result<OptimizationResult> {
        let n = config.assets.len();
        let mut weights = vec![1.0 / (n as f64); n];

        let max_iterations = 500;
        let learning_rate = 0.05;

        for _ in 0..max_iterations {
            let variance = portfolio_variance(&weights, covariance)?;
            let volatility = variance.sqrt();

            if volatility == 0.0 {
                break;
            }

            for i in 0..n {
                let mut gradient = 0.0;
                for j in 0..n {
                    gradient += 2.0 * weights[j] * covariance[(i, j)];
                }
                weights[i] -= learning_rate * gradient / volatility;
            }

            project_to_constraints(&mut weights, &config.assets, &config.constraints)?;
        }

        self.build_result(config, &weights, covariance)
    }

    fn optimize_risk_parity(
        &self,
        config: &PortfolioConfig,
        covariance: &DMatrix<f64>,
    ) -> Result<OptimizationResult> {
        let n = config.assets.len();
        let mut weights = vec![1.0 / (n as f64); n];

        let max_iterations = 500;
        let learning_rate = 0.01;

        for _ in 0..max_iterations {
            let contributions = risk_contributions(&weights, covariance)?;
            let target_contribution = contributions.iter().sum::<f64>() / (n as f64);

            for i in 0..n {
                let error = contributions[i] - target_contribution;
                weights[i] -= learning_rate * error;
            }

            project_to_constraints(&mut weights, &config.assets, &config.constraints)?;
        }

        self.build_result(config, &weights, covariance)
    }

    fn optimize_max_return(
        &self,
        config: &PortfolioConfig,
        covariance: &DMatrix<f64>,
    ) -> Result<OptimizationResult> {
        let n = config.assets.len();
        let mut weights = vec![1.0 / (n as f64); n];

        let max_iterations = 500;
        let learning_rate = 0.02;

        for _ in 0..max_iterations {
            for i in 0..n {
                weights[i] += learning_rate * config.assets[i].expected_return;
            }

            project_to_constraints(&mut weights, &config.assets, &config.constraints)?;
        }

        self.build_result(config, &weights, covariance)
    }

    fn build_result(
        &self,
        config: &PortfolioConfig,
        weights: &[f64],
        covariance: &DMatrix<f64>,
    ) -> Result<OptimizationResult> {
        validate_weights(weights, &config.assets, &config.constraints)?;

        let expected_return: f64 = config
            .assets
            .iter()
            .enumerate()
            .map(|(i, a)| weights[i] * a.expected_return)
            .sum();

        let variance = portfolio_variance(weights, covariance)?;
        let volatility = variance.sqrt();

        let sharpe_ratio = if volatility > 0.0 {
            (expected_return - config.risk_free_rate) / volatility
        } else {
            0.0
        };

        let risk_contribs = risk_contributions(weights, covariance)?;

        Ok(OptimizationResult {
            weights: weights.to_vec(),
            expected_return,
            volatility,
            sharpe_ratio,
            risk_contributions: risk_contribs,
            status: "Converged".to_string(),
            success: true,
            processing_time_ms: 0.0,
        })
    }
}

impl Default for PortfolioOptimizer {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{Asset, PortfolioConstraints};

    fn create_test_config() -> PortfolioConfig {
        PortfolioConfig {
            assets: vec![
                Asset {
                    ticker: "STOCK".to_string(),
                    name: "Stock ETF".to_string(),
                    asset_class: "equity".to_string(),
                    sector: Some("Broad Market".to_string()),
                    expected_return: 0.10,
                    volatility: 0.18,
                },
                Asset {
                    ticker: "BOND".to_string(),
                    name: "Bond ETF".to_string(),
                    asset_class: "bond".to_string(),
                    sector: Some("Fixed Income".to_string()),
                    expected_return: 0.03,
                    volatility: 0.05,
                },
            ],
            objective: OptimizationObjective::MaxSharpe,
            constraints: PortfolioConstraints::default(),
            risk_free_rate: 0.02,
            use_gpu: false,
        }
    }

    #[test]
    fn test_new_optimizer() {
        let optimizer = PortfolioOptimizer::new();
        assert!(!optimizer.status().initialized);
        assert_eq!(optimizer.status().optimization_count, 0);
    }

    #[test]
    fn test_initialize_success() {
        let mut optimizer = PortfolioOptimizer::new();
        let config = create_test_config();

        let result = optimizer.initialize(config);
        assert!(result.is_ok());
        assert!(optimizer.status().initialized);
        assert_eq!(optimizer.status().num_assets, 2);
    }

    #[test]
    fn test_initialize_insufficient_assets() {
        let mut optimizer = PortfolioOptimizer::new();
        let mut config = create_test_config();
        config.assets = vec![config.assets[0].clone()];

        let result = optimizer.initialize(config);
        assert!(result.is_err());
    }

    #[test]
    fn test_optimize_max_sharpe() {
        let mut optimizer = PortfolioOptimizer::new();
        let config = create_test_config();
        optimizer.initialize(config).unwrap();

        let result = optimizer.optimize();
        assert!(result.is_ok());

        let optimization_result = result.unwrap();
        assert!(optimization_result.success);
        assert_eq!(optimization_result.weights.len(), 2);
        assert!((optimization_result.weights.iter().sum::<f64>() - 1.0).abs() < 1e-5);
        assert!(optimization_result.sharpe_ratio > 0.0);
    }

    #[test]
    fn test_optimize_min_variance() {
        let mut optimizer = PortfolioOptimizer::new();
        let mut config = create_test_config();
        config.objective = OptimizationObjective::MinVariance;
        optimizer.initialize(config).unwrap();

        let result = optimizer.optimize();
        assert!(result.is_ok());

        let optimization_result = result.unwrap();
        assert!(optimization_result.success);
        assert!(optimization_result.volatility > 0.0);
    }

    #[test]
    fn test_optimize_risk_parity() {
        let mut optimizer = PortfolioOptimizer::new();
        let mut config = create_test_config();
        config.objective = OptimizationObjective::RiskParity;
        optimizer.initialize(config).unwrap();

        let result = optimizer.optimize();
        assert!(result.is_ok());

        let optimization_result = result.unwrap();
        assert!(optimization_result.success);
        assert_eq!(optimization_result.risk_contributions.len(), 2);
    }

    #[test]
    fn test_optimize_without_initialization() {
        let mut optimizer = PortfolioOptimizer::new();
        let result = optimizer.optimize();
        assert!(result.is_err());
    }

    #[test]
    fn test_compute_efficient_frontier() {
        let mut optimizer = PortfolioOptimizer::new();
        let config = create_test_config();
        optimizer.initialize(config).unwrap();

        let result = optimizer.compute_efficient_frontier(10);
        assert!(result.is_ok());

        let frontier = result.unwrap();
        assert!(!frontier.points.is_empty());
        assert!(frontier.max_sharpe_index < frontier.points.len());
        assert!(frontier.min_variance_index < frontier.points.len());
    }

    #[test]
    fn test_reset() {
        let mut optimizer = PortfolioOptimizer::new();
        let config = create_test_config();
        optimizer.initialize(config).unwrap();
        optimizer.optimize().unwrap();

        assert!(optimizer.status().initialized);
        assert!(optimizer.status().optimization_count > 0);

        optimizer.reset();

        assert!(!optimizer.status().initialized);
        assert_eq!(optimizer.status().optimization_count, 0);
    }

    #[test]
    fn test_optimization_metrics() {
        let mut optimizer = PortfolioOptimizer::new();
        let config = create_test_config();
        optimizer.initialize(config).unwrap();

        optimizer.optimize().unwrap();
        optimizer.optimize().unwrap();

        let status = optimizer.status();
        assert_eq!(status.optimization_count, 2);
        assert!(status.avg_optimization_time_ms > 0.0);
    }
}