rustytorch/demos/rtx-marketsim-demo/src/world_model.rs

//! Market world model for dynamics simulation.
//!
//! Implements a simplified world model for market state evolution.

use marketsim_shared::{EconomicFactors, MarketIndicators, ScenarioType};

/// Market dynamics parameters.
#[derive(Debug, Clone)]
pub struct DynamicsParams {
    /// Base drift (annualized).
    pub drift: f64,
    /// Base volatility (annualized).
    pub volatility: f64,
    /// Mean reversion speed.
    pub mean_reversion: f64,
    /// Correlation with market.
    pub market_correlation: f64,
    /// Jump intensity (probability per day).
    pub jump_intensity: f64,
    /// Jump size (mean).
    pub jump_size_mean: f64,
    /// Jump size (std).
    pub jump_size_std: f64,
}

impl Default for DynamicsParams {
    fn default() -> Self {
        Self {
            drift: 0.08,
            volatility: 0.20,
            mean_reversion: 0.1,
            market_correlation: 0.6,
            jump_intensity: 0.01,
            jump_size_mean: -0.03,
            jump_size_std: 0.02,
        }
    }
}

/// Market world model.
#[derive(Debug)]
pub struct MarketWorldModel {
    /// Default dynamics parameters.
    default_params: DynamicsParams,
    /// Risk-free rate.
    risk_free_rate: f64,
}

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

impl MarketWorldModel {
    /// Create a new world model.
    #[must_use]
    pub fn new() -> Self {
        Self {
            default_params: DynamicsParams::default(),
            risk_free_rate: 0.04,
        }
    }

    /// Get parameters adjusted for scenario type.
    #[must_use]
    pub fn get_scenario_params(
        &self,
        scenario_type: ScenarioType,
        severity: f64,
    ) -> DynamicsParams {
        let base = self.default_params.clone();

        match scenario_type {
            ScenarioType::Crisis => DynamicsParams {
                drift: base.drift - 0.30 * severity,
                volatility: base.volatility * (1.0 + 2.0 * severity),
                mean_reversion: base.mean_reversion * 0.5,
                market_correlation: 0.9, // Correlations spike in crisis
                jump_intensity: base.jump_intensity * (1.0 + 5.0 * severity),
                jump_size_mean: -0.05 - 0.05 * severity,
                jump_size_std: base.jump_size_std * 2.0,
            },
            ScenarioType::Rally => DynamicsParams {
                drift: base.drift + 0.20 * severity,
                volatility: base.volatility * (1.0 - 0.3 * severity),
                mean_reversion: base.mean_reversion,
                market_correlation: 0.7,
                jump_intensity: base.jump_intensity * (1.0 + severity),
                jump_size_mean: 0.02 + 0.02 * severity,
                jump_size_std: base.jump_size_std,
            },
            ScenarioType::HighVolatility => DynamicsParams {
                drift: base.drift,
                volatility: base.volatility * (1.0 + 1.5 * severity),
                mean_reversion: base.mean_reversion * 1.5,
                market_correlation: 0.5,
                jump_intensity: base.jump_intensity * (1.0 + 3.0 * severity),
                jump_size_mean: 0.0,
                jump_size_std: base.jump_size_std * 2.0,
            },
            ScenarioType::LowVolatility => DynamicsParams {
                drift: base.drift * 0.8,
                volatility: base.volatility * (1.0 - 0.5 * severity),
                mean_reversion: base.mean_reversion * 2.0,
                market_correlation: 0.4,
                jump_intensity: base.jump_intensity * (1.0 - 0.5 * severity),
                jump_size_mean: 0.0,
                jump_size_std: base.jump_size_std * 0.5,
            },
            ScenarioType::SectorRotation => DynamicsParams {
                drift: base.drift,
                volatility: base.volatility * 1.2,
                mean_reversion: base.mean_reversion,
                market_correlation: 0.3, // Lower correlations
                jump_intensity: base.jump_intensity,
                jump_size_mean: 0.0,
                jump_size_std: base.jump_size_std,
            },
            ScenarioType::RateShock => DynamicsParams {
                drift: base.drift - 0.10 * severity,
                volatility: base.volatility * (1.0 + 0.5 * severity),
                mean_reversion: base.mean_reversion,
                market_correlation: 0.7,
                jump_intensity: 0.0,
                jump_size_mean: 0.0,
                jump_size_std: 0.0,
            },
            ScenarioType::InflationSpike => DynamicsParams {
                drift: base.drift - 0.15 * severity,
                volatility: base.volatility * (1.0 + 0.8 * severity),
                mean_reversion: base.mean_reversion * 0.8,
                market_correlation: 0.6,
                jump_intensity: base.jump_intensity,
                jump_size_mean: -0.02 * severity,
                jump_size_std: base.jump_size_std,
            },
            ScenarioType::Custom => base,
        }
    }

    /// Simulate one step of market dynamics.
    #[must_use]
    pub fn step(
        &self,
        current_price: f64,
        params: &DynamicsParams,
        dt: f64,
        z: f64,      // Standard normal random
        z_jump: f64, // For jump size
        u_jump: f64, // Uniform for jump occurrence
    ) -> f64 {
        // Daily drift
        let drift = (params.drift - 0.5 * params.volatility.powi(2)) * dt;

        // Diffusion
        let diffusion = params.volatility * dt.sqrt() * z;

        // Jump component
        let jump = if u_jump < params.jump_intensity * dt {
            params.jump_size_mean + params.jump_size_std * z_jump
        } else {
            0.0
        };

        // New price (geometric Brownian motion with jumps)
        current_price * (drift + diffusion + jump).exp()
    }

    /// Evolve market indicators based on scenario.
    #[must_use]
    pub fn evolve_indicators(
        &self,
        current: &MarketIndicators,
        scenario_type: ScenarioType,
        severity: f64,
        dt: f64,
    ) -> MarketIndicators {
        let vix_target = match scenario_type {
            ScenarioType::Crisis => 40.0 + 30.0 * severity,
            ScenarioType::Rally => 12.0 - 2.0 * severity,
            ScenarioType::HighVolatility => 30.0 + 15.0 * severity,
            ScenarioType::LowVolatility => 10.0 - 3.0 * severity,
            _ => current.vix,
        };

        // Mean reversion towards target
        let vix_speed = 0.1;
        let new_vix = current.vix + vix_speed * (vix_target - current.vix) * dt * 252.0;

        MarketIndicators {
            vix: new_vix.clamp(8.0, 80.0),
            sp500: current.sp500, // Updated separately
            treasury_10y: current.treasury_10y,
            credit_spread: current.credit_spread
                + if scenario_type == ScenarioType::Crisis {
                    0.01 * severity
                } else {
                    0.0
                },
            put_call_ratio: current.put_call_ratio,
        }
    }

    /// Evolve economic factors.
    #[must_use]
    pub fn evolve_economic_factors(
        &self,
        current: &EconomicFactors,
        scenario_type: ScenarioType,
        severity: f64,
        _dt: f64,
    ) -> EconomicFactors {
        match scenario_type {
            ScenarioType::Crisis => EconomicFactors {
                gdp_growth: current.gdp_growth - 0.5 * severity,
                inflation: current.inflation - 0.2 * severity,
                unemployment: current.unemployment + 0.3 * severity,
                fed_funds_rate: current.fed_funds_rate - 0.25 * severity,
                consumer_sentiment: current.consumer_sentiment - 10.0 * severity,
            },
            ScenarioType::InflationSpike => EconomicFactors {
                gdp_growth: current.gdp_growth - 0.3 * severity,
                inflation: current.inflation + 2.0 * severity,
                unemployment: current.unemployment + 0.1 * severity,
                fed_funds_rate: current.fed_funds_rate + 0.5 * severity,
                consumer_sentiment: current.consumer_sentiment - 5.0 * severity,
            },
            _ => current.clone(),
        }
    }

    /// Get the risk-free rate.
    #[must_use]
    pub fn risk_free_rate(&self) -> f64 {
        self.risk_free_rate
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_world_model_creation() {
        let model = MarketWorldModel::new();
        assert!(model.risk_free_rate() > 0.0);
    }

    #[test]
    fn test_crisis_params() {
        let model = MarketWorldModel::new();
        let params = model.get_scenario_params(ScenarioType::Crisis, 0.8);

        // Crisis should have negative drift
        assert!(params.drift < 0.0);
        // Crisis should have high volatility
        assert!(params.volatility > model.default_params.volatility);
        // Crisis should have high jump intensity
        assert!(params.jump_intensity > model.default_params.jump_intensity);
    }

    #[test]
    fn test_rally_params() {
        let model = MarketWorldModel::new();
        let params = model.get_scenario_params(ScenarioType::Rally, 0.6);

        // Rally should have positive drift
        assert!(params.drift > model.default_params.drift);
        // Rally should have lower volatility
        assert!(params.volatility < model.default_params.volatility);
    }

    #[test]
    fn test_step() {
        let model = MarketWorldModel::new();
        let params = model.default_params.clone();

        let new_price = model.step(100.0, &params, 1.0 / 252.0, 0.0, 0.0, 1.0);

        // With z=0 and no jump, price should change slightly due to drift
        assert!((new_price - 100.0).abs() < 1.0);
    }

    #[test]
    fn test_evolve_indicators() {
        let model = MarketWorldModel::new();
        let current = MarketIndicators::default();

        let evolved = model.evolve_indicators(&current, ScenarioType::Crisis, 0.8, 1.0 / 252.0);

        // VIX should increase towards crisis target
        assert!(evolved.vix > current.vix);
    }
}