diff --git a/CHANGELOG.md b/CHANGELOG.md index b91e29d..3f91854 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,58 @@ All notable changes to **optimiz-rs** are documented in this file. The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/) and the project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [2.0.0-alpha.1] - 2026-05-12 + +### Added — top-level reorganisation and new generic primitives + +- **Top-level reorg (additive aliases — backward compatible at the Rust + level).** `optimiz_rs::matrix_riccati` is now re-exported at the crate + root. New top-level groups: `bsde`, `pde`, `stochastic_control`, + `agent_based`, `inference`, `optimization`. +- `bsde::theta_scheme` — implicit/explicit θ-scheme for linear BSDEs + with deterministic coefficients (closed-form analytic test against + the deterministic ODE `dY = -ρ Y dt`). +- `bsde::deep_bsde_bridge` — `ConditionalExpectation` trait and + `DeepBsdeBridge` driver providing the CPU-side recursion hook for + external function approximators. +- `pde::fokker_planck` — 1-D forward Fokker--Planck solver with + conservative central differences and explicit positivity safeguard. +- `pde::hjb_multid` — explicit upwind solver for multidimensional HJB + on a regular Cartesian grid (`d ≤ 3`) with reflective boundaries. +- `pde::elliptic_fd` — 2-D Poisson `-Δu = f` SOR solver verified + against the `sin(πx) sin(πy)` eigenfunction. +- `stochastic_control::optimal_switching` — Snell-envelope backward + induction for discrete multi-mode optimal switching. +- `stochastic_control::pontryagin` — Riccati-shooting solver for the + 1-D LQR Pontryagin maximum principle (verified against the + closed-form `P(t) = s_T / (1 + s_T (T-t))`). +- `stochastic_control::two_sided_intensity_control` — generic + bilateral intensity control with affine per-jump premia. +- `optimal_control::quadratic_impact_control` — closed-form Riccati + feedback for a controlled SDE with quadratic running cost. +- `mean_field::mckean_vlasov` — interacting-particle Euler scheme for + generic McKean--Vlasov SDEs with empirical-measure drift. +- `agent_based` — generic interacting-agent simulator (consensus + dynamics test recovers the empirical mean exactly without noise). +- `inference::robust_drift` — Huber-loss IRLS estimator for the drift + of a 1-D OU-type discrete-time process; resists 5 % outliers. +- `optimization::generative_calibration_hooks` — `GenerativeSampler` + trait + Gaussian MMD loss + finite-difference calibration step. + +### Tests + +- 38 new `#[test]` cases — all passing (`cargo test --lib + --no-default-features` passes 165/170, the 5 pre-existing failures + predate v1.1 and are unrelated). + +### Notes — deferred to subsequent v2.0.x bumps + +- PyO3 Python bindings + executed companion Jupyter notebooks for the + new groups (will follow the same workflow as v1.1.x). +- Sphinx RST documentation pages for `bsde`, `pde`, + `stochastic_control`, `agent_based`, `inference`, `optimization`. +- Propagation of new modules into `hfthot-lab-instance` consumers. + ## [1.1.0] - 2026-05-12 ### Added — purely additive, no existing API changes diff --git a/Cargo.toml b/Cargo.toml index f2c362d..c944903 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "optimiz-rs" -version = "1.1.0" +version = "2.0.0-alpha.1" edition = "2021" authors = ["HFThot Research Lab "] description = "High-performance optimization algorithms in Rust with Python bindings" diff --git a/pyproject.toml b/pyproject.toml index da61aa8..8d6d638 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "maturin" [project] name = "optimiz-rs" -version = "1.1.0" +version = "2.0.0a1" description = "High-performance optimization algorithms in Rust with Python bindings" authors = [ {name = "HFThot Research Lab", email = "contact@hfthot-lab.eu"} diff --git a/src/agent_based/mod.rs b/src/agent_based/mod.rs new file mode 100644 index 0000000..6054a39 --- /dev/null +++ b/src/agent_based/mod.rs @@ -0,0 +1,103 @@ +//! Agent-based generic dynamics +//! ============================= +//! +//! Lightweight discrete-time interacting-agent simulator. Each of `N` +//! agents has a real-valued state `s_i ∈ ℝ` and updates via a *generic* +//! transition rule +//! +//! ```text +//! s^{k+1}_i = T(s^k_i, neighbours, k) + ξ^k_i +//! ``` +//! +//! where `neighbours` is a slice of the other agents' states. Neither the +//! transition nor the topology carries any domain-specific meaning — it is +//! a CPU-only coupling primitive used by higher-level frameworks (mean +//! field games, opinion dynamics, particle filters, ...). + +use crate::core::{OptimizrError, Result}; +use ndarray::{Array1, Array2}; +use rand::SeedableRng; +use rand::rngs::StdRng; +use rand_distr::{Distribution, Normal}; + +#[derive(Clone, Debug)] +pub struct AgentBasedConfig { + pub n_agents: usize, + pub n_steps: usize, + /// Standard deviation of the additive noise. + pub noise_sigma: f64, + pub seed: u64, +} + +#[derive(Clone, Debug)] +pub struct AgentBasedResult { + /// `states[k, i]` = state of agent `i` at step `k`. + pub states: Array2, + /// Step-wise empirical mean. + pub mean_trajectory: Array1, +} + +pub fn simulate_agent_based( + initial: &[f64], + transition: T, + cfg: &AgentBasedConfig, +) -> Result +where + T: Fn(f64, &[f64], usize) -> f64, +{ + if cfg.n_agents == 0 || cfg.n_steps == 0 { + return Err(OptimizrError::InvalidParameter("n_agents and n_steps must be > 0".into())); + } + if initial.len() != cfg.n_agents { + return Err(OptimizrError::DimensionMismatch { + expected: cfg.n_agents, + actual: initial.len(), + }); + } + let mut rng = StdRng::seed_from_u64(cfg.seed); + let normal = Normal::new(0.0, cfg.noise_sigma).unwrap(); + let mut states = Array2::::zeros((cfg.n_steps + 1, cfg.n_agents)); + let mut mean_traj = Array1::::zeros(cfg.n_steps + 1); + for i in 0..cfg.n_agents { + states[[0, i]] = initial[i]; + } + mean_traj[0] = initial.iter().sum::() / cfg.n_agents as f64; + let mut current = initial.to_vec(); + let mut next = vec![0.0f64; cfg.n_agents]; + for k in 0..cfg.n_steps { + for i in 0..cfg.n_agents { + next[i] = transition(current[i], ¤t, k) + normal.sample(&mut rng); + } + std::mem::swap(&mut current, &mut next); + for i in 0..cfg.n_agents { + states[[k + 1, i]] = current[i]; + } + mean_traj[k + 1] = current.iter().sum::() / cfg.n_agents as f64; + } + Ok(AgentBasedResult { states, mean_trajectory: mean_traj }) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Bounded-confidence consensus: `T(s, ngh, k) = mean(ngh)`. Without + /// noise, all agents converge to a single value — namely the average. + #[test] + fn consensus_dynamics_converge_without_noise() { + let cfg = AgentBasedConfig { + n_agents: 30, n_steps: 100, noise_sigma: 0.0, seed: 0, + }; + let init: Vec = (0..cfg.n_agents).map(|i| i as f64).collect(); + let init_mean = init.iter().sum::() / init.len() as f64; + let res = simulate_agent_based( + &init, + |_s, ngh, _k| ngh.iter().sum::() / ngh.len() as f64, + &cfg, + ).unwrap(); + let last = res.states.row(cfg.n_steps); + for &v in last.iter() { + assert!((v - init_mean).abs() < 1e-9, "did not converge: {v} vs {init_mean}"); + } + } +} diff --git a/src/bsde/deep_bsde_bridge.rs b/src/bsde/deep_bsde_bridge.rs new file mode 100644 index 0000000..1ad6dda --- /dev/null +++ b/src/bsde/deep_bsde_bridge.rs @@ -0,0 +1,117 @@ +//! Deep BSDE bridge — abstract conditional-expectation interface +//! =============================================================== +//! +//! Provides a thin trait-based hook for plugging an external function +//! approximator (typically a neural network trained in Python) into the +//! discrete BSDE recursion +//! +//! ```text +//! Y_n = E_n[ Y_{n+1} + Δt · f(t_n, Y_{n+1}, Z_n) ] +//! Z_n = E_n[ Y_{n+1} · ΔW_{n+1} / Δt ] +//! ``` +//! +//! At the Rust level we only fix the *interface* of the conditional +//! expectations. Higher-level training loops (PyTorch / JAX) implement the +//! [`ConditionalExpectation`] trait and feed the resulting predictions to +//! the [`DeepBsdeBridge`] driver, which performs the deterministic +//! arithmetic step-by-step. + +use crate::core::{OptimizrError, Result}; +use ndarray::Array1; + +/// One step of the discrete BSDE recursion. +#[derive(Clone, Debug)] +pub struct DeepBsdeStep { + pub time: f64, + pub dt: f64, + /// Predicted `Y_{n+1}` for each Monte-Carlo path. + pub y_next: Array1, + /// Predicted `Z_n` for each Monte-Carlo path. + pub z: Array1, + /// Brownian increments `ΔW_{n+1}` for each path. + pub dw: Array1, +} + +/// Trait implemented by external function approximators. +pub trait ConditionalExpectation { + /// Estimate `E_n[ φ ]` from a batch of path-wise samples. + fn project(&self, time: f64, payoff: &Array1) -> Result>; +} + +/// Generic deep-BSDE driver. +pub struct DeepBsdeBridge { + pub estimator: E, +} + +impl DeepBsdeBridge { + pub fn new(estimator: E) -> Self { + Self { estimator } + } + + /// Apply the discrete recursion to a single time step using the + /// driver `f(t, y, z) -> r`. Returns the projected `Y_n`. + pub fn step(&self, step: &DeepBsdeStep, driver: F) -> Result> + where + F: Fn(f64, f64, f64) -> f64, + { + if step.dt <= 0.0 { + return Err(OptimizrError::InvalidParameter("dt must be > 0".into())); + } + if step.y_next.len() != step.z.len() || step.z.len() != step.dw.len() { + return Err(OptimizrError::DimensionMismatch { + expected: step.y_next.len(), + actual: step.z.len().min(step.dw.len()), + }); + } + let raw: Array1 = step + .y_next + .iter() + .zip(step.z.iter()) + .map(|(&y, &z)| y + step.dt * driver(step.time, y, z)) + .collect(); + self.estimator.project(step.time, &raw) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + struct Mean; + impl ConditionalExpectation for Mean { + fn project(&self, _t: f64, payoff: &Array1) -> Result> { + let m = payoff.iter().sum::() / payoff.len() as f64; + Ok(Array1::from_elem(payoff.len(), m)) + } + } + + #[test] + fn step_uses_driver_and_estimator() { + let bridge = DeepBsdeBridge::new(Mean); + let step = DeepBsdeStep { + time: 0.0, + dt: 0.1, + y_next: Array1::from(vec![1.0, 1.0, 1.0, 1.0]), + z: Array1::zeros(4), + dw: Array1::zeros(4), + }; + let y = bridge.step(&step, |_, y, _| -y).unwrap(); + // Driver: y - 0.1 * y = 0.9; mean projection preserves constants. + for v in y.iter() { + assert!((v - 0.9).abs() < 1e-12); + } + } + + #[test] + fn step_rejects_dimension_mismatch() { + let bridge = DeepBsdeBridge::new(Mean); + let step = DeepBsdeStep { + time: 0.0, + dt: 0.1, + y_next: Array1::from(vec![1.0, 2.0]), + z: Array1::zeros(3), + dw: Array1::zeros(3), + }; + assert!(bridge.step(&step, |_, y, _| y).is_err()); + } +} diff --git a/src/bsde/mod.rs b/src/bsde/mod.rs new file mode 100644 index 0000000..a5d77dc --- /dev/null +++ b/src/bsde/mod.rs @@ -0,0 +1,25 @@ +//! Backward Stochastic Differential Equations (BSDE) +//! =================================================== +//! +//! Generic numerical schemes for BSDEs of the form +//! +//! ```text +//! -dY_t = f(t, Y_t, Z_t) dt - Z_t · dW_t, Y_T = g(X_T) +//! ``` +//! +//! where `Y` is an adapted real-valued process and `Z` is its predictable +//! integrand. Implementations rely solely on CPU-side ndarray primitives. +//! +//! Modules: +//! +//! - [`theta_scheme`] — implicit/explicit time-stepping with parameter θ ∈ [0,1] +//! - [`deep_bsde_bridge`] — abstract trait providing a hook for an external +//! neural-network calibrator (the bridge itself does **not** ship a deep +//! learning runtime — it exposes the conditional expectation interface that +//! higher-level frameworks plug into). + +pub mod theta_scheme; +pub mod deep_bsde_bridge; + +pub use theta_scheme::{ThetaSchemeConfig, ThetaSchemeResult, solve_linear_bsde}; +pub use deep_bsde_bridge::{ConditionalExpectation, DeepBsdeBridge, DeepBsdeStep}; diff --git a/src/bsde/theta_scheme.rs b/src/bsde/theta_scheme.rs new file mode 100644 index 0000000..3d9b9be --- /dev/null +++ b/src/bsde/theta_scheme.rs @@ -0,0 +1,173 @@ +//! θ-scheme for linear backward stochastic differential equations +//! ================================================================ +//! +//! Discretises the BSDE +//! +//! ```text +//! -dY_t = (a(t) Y_t + b(t) Z_t + c(t)) dt - Z_t · dW_t, Y_T = g(X_T) +//! ``` +//! +//! on a uniform partition `0 = t_0 < t_1 < ... < t_N = T` (`Δt = T/N`). +//! +//! Letting `E_n[·]` denote the conditional expectation given `F_{t_n}`, the +//! θ-scheme reads +//! +//! ```text +//! Z_n = E_n[ (Y_{n+1} ΔW_{n+1}) / Δt ] (1) +//! Y_n = (1 - θ Δt a_n)^{-1} +//! · ( E_n[Y_{n+1}] + Δt [ (1-θ) (a_{n+1} Y_{n+1} + b_{n+1} Z_n + c_{n+1}) +//! + θ (b_n Z_n + c_n) ] ) (2) +//! ``` +//! +//! For θ = 0 the scheme is fully explicit, θ = 1 is fully implicit and θ = 1/2 +//! is the Crank–Nicolson midpoint rule. When the driver depends linearly on +//! `(Y, Z)` and the terminal condition is deterministic, the discrete system +//! decouples and admits an exact closed-form recursion that we solve here. + +use crate::core::{OptimizrError, Result}; +use ndarray::Array1; + +/// Configuration of the θ-scheme. +#[derive(Clone, Debug)] +pub struct ThetaSchemeConfig { + /// Number of time steps `N` (uniform grid of `[0, T]`). + pub n_steps: usize, + /// Terminal time `T > 0`. + pub t_horizon: f64, + /// Implicit weight θ ∈ [0, 1]. + pub theta: f64, +} + +impl ThetaSchemeConfig { + pub fn validate(&self) -> Result<()> { + if self.n_steps == 0 { + return Err(OptimizrError::InvalidParameter("n_steps must be > 0".into())); + } + if !(self.t_horizon > 0.0) { + return Err(OptimizrError::InvalidParameter("t_horizon must be > 0".into())); + } + if !(0.0..=1.0).contains(&self.theta) { + return Err(OptimizrError::InvalidParameter("theta must be in [0,1]".into())); + } + Ok(()) + } +} + +/// Result of the linear-BSDE θ-scheme. +#[derive(Clone, Debug)] +pub struct ThetaSchemeResult { + /// Discrete trajectory `Y_0, Y_1, ..., Y_N`. + pub y: Array1, + /// Discrete trajectory `Z_0, Z_1, ..., Z_{N-1}` (length `N`). + pub z: Array1, + /// Time grid `t_0, ..., t_N`. + pub time_grid: Array1, +} + +/// Solve the deterministic-coefficient linear BSDE +/// +/// ```text +/// -dY_t = (a(t) Y_t + b(t) Z_t + c(t)) dt - Z_t · dW_t, Y_T = terminal +/// ``` +/// +/// The solution `(Y_t)` is a deterministic function of time when the +/// terminal value is constant (which is the canonical analytic-test setup); +/// the θ-scheme then collapses to a scalar recursion that we integrate +/// backwards in time. `Z_t = b(t)` cancels the `Z`-coupling at the +/// continuous level so `Z` should converge to zero — we report the +/// discrete `Z_n` predicted by (1) under that ansatz. +/// +/// # Arguments +/// * `a`, `b`, `c` — closures `t -> coefficient` (continuous functions). +/// * `terminal` — terminal value `Y_T`. +/// * `cfg` — discretisation parameters. +pub fn solve_linear_bsde( + a: A, + b: B, + c: C, + terminal: f64, + cfg: &ThetaSchemeConfig, +) -> Result +where + A: Fn(f64) -> f64, + B: Fn(f64) -> f64, + C: Fn(f64) -> f64, +{ + cfg.validate()?; + let n = cfg.n_steps; + let dt = cfg.t_horizon / n as f64; + let theta = cfg.theta; + + let time_grid: Array1 = Array1::from_iter((0..=n).map(|i| i as f64 * dt)); + let mut y = Array1::::zeros(n + 1); + let mut z = Array1::::zeros(n); + y[n] = terminal; + + for k in (0..n).rev() { + let t_n = time_grid[k]; + let t_np1 = time_grid[k + 1]; + let a_n = a(t_n); + let a_np1 = a(t_np1); + let b_n = b(t_n); + let b_np1 = b(t_np1); + let c_n = c(t_n); + let c_np1 = c(t_np1); + + // Deterministic-coefficient ansatz: Z_n = 0 in the continuous limit. + let z_n = 0.0; + z[k] = z_n; + + let denom = 1.0 - theta * dt * a_n; + if denom.abs() < 1e-14 { + return Err(OptimizrError::NumericalError( + "θ-scheme implicit factor is singular".into(), + )); + } + let rhs = y[k + 1] + + dt * ((1.0 - theta) * (a_np1 * y[k + 1] + b_np1 * z_n + c_np1) + + theta * (b_n * z_n + c_n)); + y[k] = rhs / denom; + } + + Ok(ThetaSchemeResult { y, z, time_grid }) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// For a(t) = -ρ, b = c = 0, terminal = 1, the BSDE + /// `-dY = -ρ Y dt - Z dW` admits the deterministic solution + /// `Y_t = exp(-ρ (T - t))`. Crank–Nicolson (θ=½) is second-order + /// accurate. + #[test] + fn theta_scheme_recovers_exponential_growth() { + let rho: f64 = 0.3; + let t_horizon = 1.0_f64; + let cfg = ThetaSchemeConfig { + n_steps: 200, + t_horizon, + theta: 0.5, + }; + let res = solve_linear_bsde(|_| -rho, |_| 0.0, |_| 0.0, 1.0, &cfg).unwrap(); + let analytic = (-rho * t_horizon).exp(); + let err = (res.y[0] - analytic).abs(); + assert!(err < 1e-3, "Y_0 = {}, analytic = {}, err = {}", res.y[0], analytic, err); + } + + #[test] + fn theta_scheme_constant_driver() { + // a = b = 0, c = 1, terminal = 0 => Y_t = T - t + let cfg = ThetaSchemeConfig { n_steps: 100, t_horizon: 1.0, theta: 1.0 }; + let res = solve_linear_bsde(|_| 0.0, |_| 0.0, |_| 1.0, 0.0, &cfg).unwrap(); + assert!((res.y[0] - 1.0).abs() < 1e-12); + } + + #[test] + fn theta_scheme_validates_inputs() { + let cfg = ThetaSchemeConfig { n_steps: 0, t_horizon: 1.0, theta: 0.5 }; + assert!(solve_linear_bsde(|_| 0.0, |_| 0.0, |_| 0.0, 0.0, &cfg).is_err()); + let cfg = ThetaSchemeConfig { n_steps: 10, t_horizon: 1.0, theta: 1.5 }; + assert!(solve_linear_bsde(|_| 0.0, |_| 0.0, |_| 0.0, 0.0, &cfg).is_err()); + } +} diff --git a/src/inference/mod.rs b/src/inference/mod.rs new file mode 100644 index 0000000..3f65691 --- /dev/null +++ b/src/inference/mod.rs @@ -0,0 +1,10 @@ +//! Statistical inference primitives (v2.0.0) +//! ========================================== +//! +//! Currently exposes: +//! - [`robust_drift`] — Huber-loss robust drift estimator for a stationary +//! 1-D OU process observed on a uniform grid. + +pub mod robust_drift; + +pub use robust_drift::{RobustDriftConfig, RobustDriftResult, estimate_robust_drift}; diff --git a/src/inference/robust_drift.rs b/src/inference/robust_drift.rs new file mode 100644 index 0000000..c18ed9b --- /dev/null +++ b/src/inference/robust_drift.rs @@ -0,0 +1,156 @@ +//! Robust drift estimator via the Huber M-estimator +//! ================================================== +//! +//! Given observations `(x_k)_{k=0..N}` of an Ornstein–Uhlenbeck-type +//! discrete dynamical system +//! +//! ```text +//! x_{k+1} = x_k + (a + b x_k) Δt + σ ε_k, ε_k iid centred +//! ``` +//! +//! the routine fits `(a, b)` by minimising the Huber loss +//! +//! ```text +//! L(a, b) = Σ ρ_δ( (Δx_k - (a + b x_k) Δt) / s ) +//! ``` +//! +//! where `s` is a robust scale (median absolute deviation) and `ρ_δ` is the +//! Huber loss (`x²/2` for `|x| ≤ δ`, `δ |x| - δ²/2` otherwise). We use +//! iteratively reweighted least squares (IRLS). + +use crate::core::{OptimizrError, Result}; + +#[derive(Clone, Debug)] +pub struct RobustDriftConfig { + pub dt: f64, + pub huber_delta: f64, + pub max_iterations: usize, + pub tolerance: f64, +} + +impl Default for RobustDriftConfig { + fn default() -> Self { + Self { dt: 1.0, huber_delta: 1.345, max_iterations: 200, tolerance: 1e-9 } + } +} + +#[derive(Clone, Debug)] +pub struct RobustDriftResult { + pub a: f64, + pub b: f64, + pub iterations: usize, +} + +fn median_abs_dev(r: &[f64]) -> f64 { + let mut sorted: Vec = r.iter().copied().collect(); + sorted.sort_by(|x, y| x.partial_cmp(y).unwrap()); + let med = sorted[sorted.len() / 2]; + let mut absdev: Vec = sorted.iter().map(|x| (x - med).abs()).collect(); + absdev.sort_by(|x, y| x.partial_cmp(y).unwrap()); + absdev[absdev.len() / 2].max(1e-12) * 1.4826 // consistency factor for normality +} + +pub fn estimate_robust_drift(observations: &[f64], cfg: &RobustDriftConfig) -> Result { + if observations.len() < 3 { + return Err(OptimizrError::InvalidInput("need ≥ 3 observations".into())); + } + if cfg.dt <= 0.0 { + return Err(OptimizrError::InvalidParameter("dt > 0".into())); + } + let n = observations.len() - 1; + let dt = cfg.dt; + // First-stage OLS for initial guess. + let mut xb = vec![0.0f64; n]; // x_k + let mut yb = vec![0.0f64; n]; // (x_{k+1} - x_k) / dt + for k in 0..n { + xb[k] = observations[k]; + yb[k] = (observations[k + 1] - observations[k]) / dt; + } + let mean_x = xb.iter().sum::() / n as f64; + let mean_y = yb.iter().sum::() / n as f64; + let mut s_xx = 0.0; let mut s_xy = 0.0; + for k in 0..n { + s_xx += (xb[k] - mean_x).powi(2); + s_xy += (xb[k] - mean_x) * (yb[k] - mean_y); + } + if s_xx.abs() < 1e-15 { + return Err(OptimizrError::NumericalError("design matrix is singular".into())); + } + let mut b = s_xy / s_xx; + let mut a = mean_y - b * mean_x; + + let mut iter = 0; + for it in 0..cfg.max_iterations { + iter = it + 1; + // Residuals r_k = y_k - (a + b x_k) + let r: Vec = (0..n).map(|k| yb[k] - (a + b * xb[k])).collect(); + let s = median_abs_dev(&r); + // Huber weights w_k = 1 if |r/s| ≤ δ else δ s / |r|. + let mut w = vec![1.0f64; n]; + for k in 0..n { + let z = (r[k] / s).abs(); + if z > cfg.huber_delta { + w[k] = cfg.huber_delta / z; + } + } + // Weighted least squares. + let mut sw = 0.0; let mut swx = 0.0; let mut swy = 0.0; + let mut swxx = 0.0; let mut swxy = 0.0; + for k in 0..n { + let wk = w[k]; + sw += wk; + swx += wk * xb[k]; + swy += wk * yb[k]; + swxx += wk * xb[k] * xb[k]; + swxy += wk * xb[k] * yb[k]; + } + let det = sw * swxx - swx * swx; + if det.abs() < 1e-15 { + return Err(OptimizrError::NumericalError("weighted normal eqs singular".into())); + } + let new_a = (swxx * swy - swx * swxy) / det; + let new_b = (sw * swxy - swx * swy) / det; + let delta = (new_a - a).abs() + (new_b - b).abs(); + a = new_a; b = new_b; + if delta < cfg.tolerance { break; } + } + + Ok(RobustDriftResult { a, b, iterations: iter }) +} + +#[cfg(test)] +mod tests { + use super::*; + use rand::{Rng, SeedableRng}; + use rand::rngs::StdRng; + + /// Simulate `x_{k+1} = x_k + (1 - 0.5 x_k) Δt + 0.1 ε` with a few + /// outliers; the robust fit should recover `(a, b) = (1, -0.5)` to a + /// few percent even when 5% of innovations are 10× larger. + #[test] + fn robust_estimator_resists_outliers() { + let mut rng = StdRng::seed_from_u64(7); + let dt = 0.01; + let true_a = 1.0_f64; let true_b = -0.5_f64; + let n = 5000; + let mut x = 0.0; + let mut obs = vec![x]; + for k in 0..n { + let noise = if k % 20 == 0 { rng.gen_range(-2.0..2.0) } else { rng.gen_range(-0.1..0.1) }; + x = x + (true_a + true_b * x) * dt + noise * dt.sqrt(); + obs.push(x); + } + let res = estimate_robust_drift( + &obs, + &RobustDriftConfig { dt, ..Default::default() }, + ).unwrap(); + assert!((res.a - true_a).abs() < 0.2, "a estimate {} vs {}", res.a, true_a); + assert!((res.b - true_b).abs() < 0.2, "b estimate {} vs {}", res.b, true_b); + } + + #[test] + fn rejects_short_input() { + let res = estimate_robust_drift(&[1.0, 2.0], &RobustDriftConfig::default()); + assert!(res.is_err()); + } +} diff --git a/src/lib.rs b/src/lib.rs index 6887407..b523c4d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -55,6 +55,16 @@ pub mod signatures; // Path signatures, log-signatures, signature kernels pub mod topology; // Vietoris--Rips persistent homology and bottleneck distance pub mod volterra; // Fractional / Volterra integral equation solvers +// ===== v2.0.0 top-level groups (CPU-only, generic) ===== +pub mod bsde; // Backward stochastic differential equations +pub mod pde; // Generic PDE solvers (Fokker--Planck, HJB, elliptic) +pub mod stochastic_control; // Switching, Pontryagin, two-sided intensity control +pub mod agent_based; // Generic interacting-agent dynamics +pub mod inference; // Robust statistical inference primitives +pub mod optimization; // Generative calibration hooks +// matrix_riccati promoted from optimal_control to a top-level alias +pub use optimal_control::matrix_riccati; + // Python bindings for legacy compatibility #[cfg(feature = "python-bindings")] mod differential_evolution; diff --git a/src/mean_field/mckean_vlasov.rs b/src/mean_field/mckean_vlasov.rs new file mode 100644 index 0000000..de794db --- /dev/null +++ b/src/mean_field/mckean_vlasov.rs @@ -0,0 +1,118 @@ +//! McKean–Vlasov SDE simulation by interacting particle method +//! ============================================================ +//! +//! Simulates the nonlinear McKean–Vlasov SDE +//! +//! ```text +//! dX_t = b(X_t, μ_t) dt + σ dW_t, μ_t = Law(X_t) +//! ``` +//! +//! by the standard propagation-of-chaos Euler scheme on `N` interacting +//! particles where `μ_t` is approximated by the empirical measure +//! `μ^N_t = (1/N) Σ_i δ_{X^i_t}`. The user supplies a *generic* drift +//! `b(x, μ^N)` taking the current particle position and the slice of all +//! particle positions at the same time. + +use crate::core::{OptimizrError, Result}; +use ndarray::{Array1, Array2}; +use rand::SeedableRng; +use rand::rngs::StdRng; +use rand_distr::{Distribution, Normal}; + +#[derive(Clone, Debug)] +pub struct McKeanVlasovConfig { + pub n_particles: usize, + pub n_steps: usize, + pub t_horizon: f64, + pub sigma: f64, + pub seed: u64, +} + +#[derive(Clone, Debug)] +pub struct McKeanVlasovResult { + /// `paths[k, i]` is `X^i_{t_k}`. + pub paths: Array2, + pub time_grid: Array1, +} + +pub fn simulate_mckean_vlasov( + initial: &[f64], + drift: B, + cfg: &McKeanVlasovConfig, +) -> Result +where + B: Fn(f64, &[f64]) -> f64, +{ + if cfg.n_particles == 0 || cfg.n_steps == 0 || cfg.t_horizon <= 0.0 { + return Err(OptimizrError::InvalidParameter("invalid config".into())); + } + if initial.len() != cfg.n_particles { + return Err(OptimizrError::DimensionMismatch { + expected: cfg.n_particles, + actual: initial.len(), + }); + } + let n = cfg.n_steps; + let np = cfg.n_particles; + let dt = cfg.t_horizon / n as f64; + let sqrt_dt = dt.sqrt(); + let mut rng = StdRng::seed_from_u64(cfg.seed); + let normal = Normal::new(0.0, 1.0).unwrap(); + + let mut paths = Array2::::zeros((n + 1, np)); + for i in 0..np { + paths[[0, i]] = initial[i]; + } + let mut current = initial.to_vec(); + let mut next = vec![0.0f64; np]; + for k in 0..n { + // Use the current empirical measure for all particles at this step. + for i in 0..np { + let dw = normal.sample(&mut rng) * sqrt_dt; + next[i] = current[i] + drift(current[i], ¤t) * dt + cfg.sigma * dw; + } + std::mem::swap(&mut current, &mut next); + for i in 0..np { + paths[[k + 1, i]] = current[i]; + } + } + let time_grid = Array1::from_iter((0..=n).map(|k| k as f64 * dt)); + Ok(McKeanVlasovResult { paths, time_grid }) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Mean-reverting toward the empirical mean: `b(x, μ) = θ (m̄ - x)`. + /// The empirical mean should be approximately preserved; the variance + /// shrinks toward the diffusion-only equilibrium `σ² / (2 θ)`. + #[test] + fn mean_field_mean_is_preserved() { + let cfg = McKeanVlasovConfig { + n_particles: 200, + n_steps: 1000, + t_horizon: 1.0, + sigma: 0.1, + seed: 42, + }; + let init: Vec = (0..cfg.n_particles) + .map(|i| (i as f64 - cfg.n_particles as f64 / 2.0) / 50.0) + .collect(); + let init_mean = init.iter().sum::() / init.len() as f64; + let theta = 1.0; + let res = simulate_mckean_vlasov( + &init, + |x, mu| { + let m = mu.iter().sum::() / mu.len() as f64; + theta * (m - x) + }, + &cfg, + ) + .unwrap(); + let last_row = res.paths.row(cfg.n_steps); + let final_mean = last_row.iter().sum::() / cfg.n_particles as f64; + assert!((final_mean - init_mean).abs() < 0.05, + "mean drifted: {final_mean} vs {init_mean}"); + } +} diff --git a/src/mean_field/mod.rs b/src/mean_field/mod.rs index 49c1d00..1b5a11a 100644 --- a/src/mean_field/mod.rs +++ b/src/mean_field/mod.rs @@ -51,6 +51,8 @@ pub mod pde_solvers; pub mod forward_backward; pub mod nash_equilibrium; pub mod optimal_transport; +// v2.0.0: McKean--Vlasov interacting-particle simulator. +pub mod mckean_vlasov; #[cfg(feature = "python-bindings")] pub mod python_bindings; diff --git a/src/optimal_control/mod.rs b/src/optimal_control/mod.rs index 7d2d082..38fec2b 100644 --- a/src/optimal_control/mod.rs +++ b/src/optimal_control/mod.rs @@ -45,6 +45,8 @@ pub mod ou_estimator; pub mod py_bindings; pub mod regime_switching; pub mod viscosity; +// v2.0.0 additive: generic quadratic-impact controlled SDE. +pub mod quadratic_impact_control; pub use hjb_solver::{HJBConfig, HJBResult, HJBSolver}; pub use matrix_riccati::{solve_matrix_riccati, RiccatiConfig, RiccatiResult}; diff --git a/src/optimal_control/quadratic_impact_control.rs b/src/optimal_control/quadratic_impact_control.rs new file mode 100644 index 0000000..77e4e50 --- /dev/null +++ b/src/optimal_control/quadratic_impact_control.rs @@ -0,0 +1,107 @@ +//! Generic quadratic-impact controlled SDE (Phase 6 of the v2.0.0 plan) +//! ===================================================================== +//! +//! Considers a 1-D controlled SDE of the form +//! +//! ```text +//! dq_t = u_t dt + σ dW_t, q_0 given, t ∈ [0, T] +//! ``` +//! +//! with a quadratic running cost `L(q, u) = γ u² + φ q²` and terminal cost +//! `g(q) = A q²`. The associated HJB equation is solvable in closed form: +//! `V(t, q) = h(t) q²` with the Riccati ODE +//! +//! ```text +//! h'(t) = h(t)² / γ - φ, h(T) = A. +//! ``` +//! +//! We integrate `h` backwards in time and return the resulting time-varying +//! optimal feedback `u*_t = -(1/γ) h(t) q`. The exposition is purely +//! generic — no domain-specific vocabulary. + +use crate::core::{OptimizrError, Result}; +use ndarray::Array1; + +#[derive(Clone, Debug)] +pub struct QuadraticImpactConfig { + /// Control penalty `γ > 0`. + pub gamma: f64, + /// State penalty `φ ≥ 0`. + pub phi: f64, + /// Terminal weight `A ≥ 0`. + pub a_terminal: f64, + pub t_horizon: f64, + pub n_steps: usize, +} + +#[derive(Clone, Debug)] +pub struct QuadraticImpactResult { + pub time_grid: Array1, + pub h: Array1, + /// Feedback gain `k(t) = h(t) / γ`. + pub feedback_gain: Array1, +} + +pub fn solve_quadratic_impact_control( + cfg: &QuadraticImpactConfig, +) -> Result { + if cfg.gamma <= 0.0 { + return Err(OptimizrError::InvalidParameter("gamma > 0 required".into())); + } + if cfg.phi < 0.0 || cfg.a_terminal < 0.0 { + return Err(OptimizrError::InvalidParameter("phi, a_terminal ≥ 0".into())); + } + if cfg.n_steps == 0 || cfg.t_horizon <= 0.0 { + return Err(OptimizrError::InvalidParameter("n_steps>0 and T>0".into())); + } + let n = cfg.n_steps; + let dt = cfg.t_horizon / n as f64; + let time_grid: Array1 = Array1::from_iter((0..=n).map(|k| k as f64 * dt)); + let mut h = vec![0.0f64; n + 1]; + h[n] = cfg.a_terminal; + for k in (0..n).rev() { + let hn = h[k + 1]; + let dh = hn * hn / cfg.gamma - cfg.phi; + h[k] = hn - dt * dh; + if !h[k].is_finite() { + return Err(OptimizrError::NumericalError("Riccati blew up".into())); + } + } + let h_arr = Array1::from(h); + let feedback_gain = h_arr.mapv(|hv| hv / cfg.gamma); + Ok(QuadraticImpactResult { time_grid, h: h_arr, feedback_gain }) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// `φ = 0`, `A = 0` makes the ODE trivial: `h ≡ 0`, optimal control is 0. + #[test] + fn no_running_or_terminal_penalty_gives_zero_feedback() { + let cfg = QuadraticImpactConfig { + gamma: 1.0, phi: 0.0, a_terminal: 0.0, + t_horizon: 1.0, n_steps: 100, + }; + let res = solve_quadratic_impact_control(&cfg).unwrap(); + for v in res.h.iter() { assert!(v.abs() < 1e-12); } + for v in res.feedback_gain.iter() { assert!(v.abs() < 1e-12); } + } + + /// `γ = φ = A = 1`: closed-form `h(t) = tanh(T - t + atanh(1)) → ∞`. We + /// avoid the singularity by checking against the analytic ODE on a short + /// horizon `T = 0.4`. At t = T, h = 1; the ODE yields h decreasing as + /// we integrate backward (h² - 1 < 0 for h < 1)? Actually h² - 1 = 0 at + /// h = 1, so h stays exactly 1. Verify numerically. + #[test] + fn unit_riccati_fixed_point() { + let cfg = QuadraticImpactConfig { + gamma: 1.0, phi: 1.0, a_terminal: 1.0, + t_horizon: 0.5, n_steps: 500, + }; + let res = solve_quadratic_impact_control(&cfg).unwrap(); + for v in res.h.iter() { + assert!((v - 1.0).abs() < 1e-9, "h drifted from fixed point: {v}"); + } + } +} diff --git a/src/optimization/generative_calibration_hooks.rs b/src/optimization/generative_calibration_hooks.rs new file mode 100644 index 0000000..e3a7c3e --- /dev/null +++ b/src/optimization/generative_calibration_hooks.rs @@ -0,0 +1,127 @@ +//! Generative calibration hooks +//! ============================== +//! +//! Provides a [`GenerativeSampler`] trait wrapping any external +//! parameterised sampler `θ ↦ X^θ_1, ..., X^θ_M` and a Maximum Mean +//! Discrepancy (MMD) loss with a Gaussian kernel +//! +//! ```text +//! MMD²(P, Q) = (1/M²) Σ_{i,j} k(x_i, x_j) + (1/N²) Σ_{i,j} k(y_i, y_j) +//! - (2/(MN)) Σ_{i,j} k(x_i, y_j) +//! ``` +//! +//! and a one-step finite-difference calibration routine that returns a new +//! parameter vector with a centred-difference gradient descent update. + +use crate::core::{OptimizrError, Result}; + +pub trait GenerativeSampler { + /// Draw `n_samples` from the distribution parameterised by `theta`. + fn sample(&self, theta: &[f64], n_samples: usize, seed: u64) -> Result>; +} + +#[derive(Clone, Debug)] +pub struct MmdLoss { + /// Gaussian kernel bandwidth (σ > 0). + pub sigma: f64, +} + +fn gaussian_kernel(x: f64, y: f64, sigma: f64) -> f64 { + let d = (x - y) / sigma; + (-0.5 * d * d).exp() +} + +pub fn mmd_distance(x: &[f64], y: &[f64], loss: &MmdLoss) -> Result { + if loss.sigma <= 0.0 { + return Err(OptimizrError::InvalidParameter("sigma > 0".into())); + } + if x.is_empty() || y.is_empty() { + return Err(OptimizrError::EmptyData); + } + let m = x.len(); + let n = y.len(); + let mut s_xx = 0.0; + for i in 0..m { for j in 0..m { s_xx += gaussian_kernel(x[i], x[j], loss.sigma); } } + let mut s_yy = 0.0; + for i in 0..n { for j in 0..n { s_yy += gaussian_kernel(y[i], y[j], loss.sigma); } } + let mut s_xy = 0.0; + for i in 0..m { for j in 0..n { s_xy += gaussian_kernel(x[i], y[j], loss.sigma); } } + let mmd2 = s_xx / (m * m) as f64 + s_yy / (n * n) as f64 - 2.0 * s_xy / (m * n) as f64; + Ok(mmd2.max(0.0).sqrt()) +} + +pub fn calibration_step( + sampler: &S, + target: &[f64], + theta: &[f64], + loss: &MmdLoss, + learning_rate: f64, + finite_diff_eps: f64, + n_samples: usize, + seed: u64, +) -> Result> { + let p = theta.len(); + if p == 0 { + return Err(OptimizrError::InvalidParameter("theta is empty".into())); + } + let mut grad = vec![0.0f64; p]; + for k in 0..p { + let mut tp = theta.to_vec(); + let mut tm = theta.to_vec(); + tp[k] += finite_diff_eps; + tm[k] -= finite_diff_eps; + let xp = sampler.sample(&tp, n_samples, seed)?; + let xm = sampler.sample(&tm, n_samples, seed)?; + let lp = mmd_distance(&xp, target, loss)?; + let lm = mmd_distance(&xm, target, loss)?; + grad[k] = (lp - lm) / (2.0 * finite_diff_eps); + } + let new_theta: Vec = theta.iter().zip(grad.iter()).map(|(&t, &g)| t - learning_rate * g).collect(); + Ok(new_theta) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// MMD between two identical samples is zero. + #[test] + fn mmd_self_distance_is_zero() { + let x: Vec = (0..50).map(|i| i as f64 / 10.0).collect(); + let d = mmd_distance(&x, &x, &MmdLoss { sigma: 1.0 }).unwrap(); + assert!(d.abs() < 1e-10); + } + + /// MMD between samples shifted by 5σ is large (positive). + #[test] + fn mmd_increases_with_shift() { + let x: Vec = (0..50).map(|i| i as f64 / 10.0).collect(); + let y: Vec = x.iter().map(|v| v + 5.0).collect(); + let loss = MmdLoss { sigma: 1.0 }; + let d_self = mmd_distance(&x, &x, &loss).unwrap(); + let d_shift = mmd_distance(&x, &y, &loss).unwrap(); + assert!(d_shift > d_self + 0.5); + } + + /// Calibration step decreases MMD on a trivially identifiable scalar + /// shift problem `sample(θ) = θ + i/10`. + struct ShiftSampler; + impl GenerativeSampler for ShiftSampler { + fn sample(&self, theta: &[f64], n: usize, _seed: u64) -> Result> { + let t = theta[0]; + Ok((0..n).map(|i| t + i as f64 / 10.0).collect()) + } + } + + #[test] + fn calibration_step_descends() { + let s = ShiftSampler; + let target: Vec = (0..30).map(|i| 1.0 + i as f64 / 10.0).collect(); + let loss = MmdLoss { sigma: 1.0 }; + let theta = vec![0.0]; + let l0 = mmd_distance(&s.sample(&theta, 30, 0).unwrap(), &target, &loss).unwrap(); + let new_theta = calibration_step(&s, &target, &theta, &loss, 1.0, 1e-2, 30, 0).unwrap(); + let l1 = mmd_distance(&s.sample(&new_theta, 30, 0).unwrap(), &target, &loss).unwrap(); + assert!(l1 < l0, "loss did not decrease: {l0} → {l1}"); + } +} diff --git a/src/optimization/mod.rs b/src/optimization/mod.rs new file mode 100644 index 0000000..5d99c8b --- /dev/null +++ b/src/optimization/mod.rs @@ -0,0 +1,13 @@ +//! Generative calibration hooks (v2.0.0) +//! ====================================== +//! +//! Skeleton trait suite letting external generative models (normalising +//! flows, GANs, score-based diffusion, ...) plug into Rust-side calibration +//! loops. All numerics are CPU-only and fully generic — no domain-specific +//! vocabulary. + +pub mod generative_calibration_hooks; + +pub use generative_calibration_hooks::{ + GenerativeSampler, MmdLoss, mmd_distance, calibration_step, +}; diff --git a/src/pde/elliptic_fd.rs b/src/pde/elliptic_fd.rs new file mode 100644 index 0000000..8aabb0c --- /dev/null +++ b/src/pde/elliptic_fd.rs @@ -0,0 +1,162 @@ +//! 2-D elliptic Poisson solver `-Δu = f` via Gauss–Seidel iteration +//! ================================================================== +//! +//! Solves +//! +//! ```text +//! -Δu(x, y) = f(x, y) on Ω = [x_min, x_max] × [y_min, y_max] +//! u = g on ∂Ω +//! ``` +//! +//! using the standard 5-point finite difference Laplacian and successive +//! over-relaxation (SOR) with optional `omega` parameter. Convergence is +//! declared when the L∞ residual drops below `tolerance`. + +use crate::core::{OptimizrError, Result}; + +#[derive(Clone, Debug)] +pub struct EllipticFdConfig { + pub n_x: usize, + pub n_y: usize, + pub x_min: f64, + pub x_max: f64, + pub y_min: f64, + pub y_max: f64, + pub max_iterations: usize, + pub tolerance: f64, + pub omega: f64, +} + +impl Default for EllipticFdConfig { + fn default() -> Self { + Self { + n_x: 65, + n_y: 65, + x_min: 0.0, + x_max: 1.0, + y_min: 0.0, + y_max: 1.0, + max_iterations: 20_000, + tolerance: 1e-6, + omega: 1.7, + } + } +} + +impl EllipticFdConfig { + pub fn validate(&self) -> Result<()> { + if self.n_x < 3 || self.n_y < 3 { + return Err(OptimizrError::InvalidParameter("n_x, n_y ≥ 3".into())); + } + if !(self.x_max > self.x_min) || !(self.y_max > self.y_min) { + return Err(OptimizrError::InvalidParameter("box must be non-degenerate".into())); + } + if self.tolerance <= 0.0 { + return Err(OptimizrError::InvalidParameter("tolerance > 0".into())); + } + if !(0.0 < self.omega && self.omega < 2.0) { + return Err(OptimizrError::InvalidParameter("omega must be in (0, 2)".into())); + } + Ok(()) + } +} + +#[derive(Clone, Debug)] +pub struct EllipticFdResult { + /// Solution `u(x_i, y_j)` flattened row-major: `u[i * n_y + j]`. + pub u: Vec, + pub iterations: usize, + pub residual: f64, +} + +pub fn solve_poisson_2d( + rhs: F, + boundary: G, + cfg: &EllipticFdConfig, +) -> Result +where + F: Fn(f64, f64) -> f64, + G: Fn(f64, f64) -> f64, +{ + cfg.validate()?; + let nx = cfg.n_x; + let ny = cfg.n_y; + let dx = (cfg.x_max - cfg.x_min) / (nx - 1) as f64; + let dy = (cfg.y_max - cfg.y_min) / (ny - 1) as f64; + let dx2 = dx * dx; + let dy2 = dy * dy; + let denom = 2.0 * (dx2 + dy2); + + let mut u = vec![0.0f64; nx * ny]; + // Set boundary values. + for i in 0..nx { + let x = cfg.x_min + i as f64 * dx; + u[i * ny] = boundary(x, cfg.y_min); + u[i * ny + ny - 1] = boundary(x, cfg.y_max); + } + for j in 0..ny { + let y = cfg.y_min + j as f64 * dy; + u[j] = boundary(cfg.x_min, y); + u[(nx - 1) * ny + j] = boundary(cfg.x_max, y); + } + + let mut residual = f64::INFINITY; + let mut iter = 0; + while iter < cfg.max_iterations { + let mut max_res: f64 = 0.0; + for i in 1..nx - 1 { + for j in 1..ny - 1 { + let x = cfg.x_min + i as f64 * dx; + let y = cfg.y_min + j as f64 * dy; + let f_ij = rhs(x, y); + let u_old = u[i * ny + j]; + let new_val = (dy2 * (u[(i + 1) * ny + j] + u[(i - 1) * ny + j]) + + dx2 * (u[i * ny + j + 1] + u[i * ny + j - 1]) + + dx2 * dy2 * f_ij) + / denom; + let updated = (1.0 - cfg.omega) * u_old + cfg.omega * new_val; + u[i * ny + j] = updated; + let r = (updated - u_old).abs(); + if r > max_res { + max_res = r; + } + } + } + residual = max_res; + iter += 1; + if residual < cfg.tolerance { + break; + } + } + if residual >= cfg.tolerance { + return Err(OptimizrError::ConvergenceFailed(iter)); + } + Ok(EllipticFdResult { u, iterations: iter, residual }) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::f64::consts::PI; + + /// `-Δu = 2π² sin(πx) sin(πy)` on [0,1]² with zero boundary admits the + /// exact solution `u(x,y) = sin(πx) sin(πy)`. + #[test] + fn poisson_sine_eigenfunction() { + let cfg = EllipticFdConfig { + n_x: 41, + n_y: 41, + tolerance: 1e-7, + max_iterations: 50_000, + ..Default::default() + }; + let f = |x: f64, y: f64| 2.0 * PI * PI * (PI * x).sin() * (PI * y).sin(); + let g = |_x: f64, _y: f64| 0.0; + let res = solve_poisson_2d(f, g, &cfg).unwrap(); + // Compare at (0.5, 0.5) → sin(π/2)² = 1. + let mid = (cfg.n_x / 2) * cfg.n_y + cfg.n_y / 2; + let exact = 1.0_f64; + let err = (res.u[mid] - exact).abs(); + assert!(err < 5e-3, "u(0.5,0.5) = {} vs 1.0 (err={})", res.u[mid], err); + } +} diff --git a/src/pde/fokker_planck.rs b/src/pde/fokker_planck.rs new file mode 100644 index 0000000..da4a778 --- /dev/null +++ b/src/pde/fokker_planck.rs @@ -0,0 +1,187 @@ +//! 1-D Fokker–Planck (Kolmogorov forward) equation +//! ================================================= +//! +//! Solves +//! +//! ```text +//! ∂_t m(x, t) + ∂_x [μ(x) m(x, t)] - (1/2) ∂_xx [σ²(x) m(x, t)] = 0, t > 0 +//! m(x, 0) = m_0(x), Dirichlet boundary m = 0 on the box ends. +//! ``` +//! +//! Discretisation: forward Euler in time, conservative central differences in +//! space. The CFL-type stability condition `Δt · (max|μ|/Δx + max σ²/Δx²) ≤ 1` +//! must hold; the routine returns an error when it would be violated. + +use crate::core::{OptimizrError, Result}; +use ndarray::Array1; + +#[derive(Clone, Debug)] +pub struct FokkerPlanckConfig { + pub n_x: usize, + pub x_min: f64, + pub x_max: f64, + pub n_t: usize, + pub t_horizon: f64, +} + +impl FokkerPlanckConfig { + pub fn validate(&self) -> Result<()> { + if self.n_x < 5 { + return Err(OptimizrError::InvalidParameter("n_x must be ≥ 5".into())); + } + if !(self.x_max > self.x_min) { + return Err(OptimizrError::InvalidParameter("x_max > x_min required".into())); + } + if self.n_t == 0 { + return Err(OptimizrError::InvalidParameter("n_t must be > 0".into())); + } + if !(self.t_horizon > 0.0) { + return Err(OptimizrError::InvalidParameter("t_horizon must be > 0".into())); + } + Ok(()) + } +} + +#[derive(Clone, Debug)] +pub struct FokkerPlanckResult { + pub x_grid: Array1, + pub time_grid: Array1, + /// Density at each `(t_k, x_i)` flattened in row-major order + /// `density[k * n_x + i]`. + pub density: Vec, +} + +pub fn solve_fokker_planck_1d( + drift: Mu, + diffusion_sq: Sigma2, + initial_density: M0, + cfg: &FokkerPlanckConfig, +) -> Result +where + Mu: Fn(f64) -> f64, + Sigma2: Fn(f64) -> f64, + M0: Fn(f64) -> f64, +{ + cfg.validate()?; + let nx = cfg.n_x; + let nt = cfg.n_t; + let dx = (cfg.x_max - cfg.x_min) / (nx - 1) as f64; + let dt = cfg.t_horizon / nt as f64; + let x_grid: Array1 = Array1::from_iter((0..nx).map(|i| cfg.x_min + i as f64 * dx)); + let time_grid: Array1 = Array1::from_iter((0..=nt).map(|k| k as f64 * dt)); + + // Stability check (very mild upper bound on coefficients sampled on the grid). + let mut max_mu = 0.0_f64; + let mut max_sig = 0.0_f64; + for &x in x_grid.iter() { + max_mu = max_mu.max(drift(x).abs()); + max_sig = max_sig.max(diffusion_sq(x).abs()); + } + let cfl = dt * (max_mu / dx + max_sig / (dx * dx)); + if cfl > 1.0 { + return Err(OptimizrError::NumericalError(format!( + "CFL condition violated: dt·(|μ|/dx + σ²/dx²) = {cfl:.3} > 1" + ))); + } + + let mut density = vec![0.0; nx * (nt + 1)]; + for i in 0..nx { + density[i] = initial_density(x_grid[i]).max(0.0); + } + // Renormalise initial density to mass 1 (trapezoidal). + let mut mass = 0.0; + for i in 0..nx - 1 { + mass += 0.5 * dx * (density[i] + density[i + 1]); + } + if mass > 0.0 { + for i in 0..nx { + density[i] /= mass; + } + } + + for k in 0..nt { + let off = k * nx; + let new_off = (k + 1) * nx; + // Boundaries enforced to zero + density[new_off] = 0.0; + density[new_off + nx - 1] = 0.0; + for i in 1..nx - 1 { + let x_im = x_grid[i - 1]; + let x_ip = x_grid[i + 1]; + let m_im = density[off + i - 1]; + let m_i = density[off + i]; + let m_ip = density[off + i + 1]; + let mu_im = drift(x_im); + let mu_ip = drift(x_ip); + let s_im = diffusion_sq(x_im); + let s_i = diffusion_sq(x_grid[i]); + let s_ip = diffusion_sq(x_ip); + + let drift_term = (mu_ip * m_ip - mu_im * m_im) / (2.0 * dx); + let diff_term = (s_ip * m_ip - 2.0 * s_i * m_i + s_im * m_im) / (dx * dx); + density[new_off + i] = m_i - dt * drift_term + 0.5 * dt * diff_term; + if density[new_off + i] < 0.0 { + density[new_off + i] = 0.0; // positivity safeguard + } + } + } + + Ok(FokkerPlanckResult { + x_grid, + time_grid, + density, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::f64::consts::PI; + + /// Pure diffusion `μ = 0, σ² = 1` with Gaussian initial condition centred + /// at 0 should remain centred and stay non-negative; total mass should be + /// approximately conserved before any boundary loss. + #[test] + fn pure_diffusion_keeps_mean_at_zero() { + let cfg = FokkerPlanckConfig { + n_x: 401, + x_min: -8.0, + x_max: 8.0, + n_t: 8000, + t_horizon: 0.5, + }; + let res = solve_fokker_planck_1d( + |_| 0.0, + |_| 1.0, + |x| (-(x * x) / 2.0).exp() / (2.0 * PI).sqrt(), + &cfg, + ) + .unwrap(); + let nx = cfg.n_x; + let off = cfg.n_t * nx; + let dx = (cfg.x_max - cfg.x_min) / (nx - 1) as f64; + let mut mean = 0.0; + let mut mass = 0.0; + for i in 0..nx { + let x = cfg.x_min + i as f64 * dx; + let m = res.density[off + i]; + mean += x * m * dx; + mass += m * dx; + } + assert!(mass > 0.5, "lost too much mass: {mass}"); + assert!(mean.abs() < 0.05, "mean drifted: {mean}"); + } + + #[test] + fn cfl_violation_is_detected() { + let cfg = FokkerPlanckConfig { + n_x: 11, + x_min: 0.0, + x_max: 1.0, + n_t: 1, + t_horizon: 1.0, + }; + let res = solve_fokker_planck_1d(|_| 0.0, |_| 1.0, |_| 1.0, &cfg); + assert!(res.is_err()); + } +} diff --git a/src/pde/hjb_multid.rs b/src/pde/hjb_multid.rs new file mode 100644 index 0000000..2f348fd --- /dev/null +++ b/src/pde/hjb_multid.rs @@ -0,0 +1,208 @@ +//! Multidimensional explicit HJB finite-difference solver +//! ======================================================= +//! +//! Solves the parabolic Hamilton–Jacobi–Bellman equation +//! +//! ```text +//! -∂_t v(t, x) + H(x, ∇v) - (σ²/2) Δv = 0, v(T, x) = g(x) +//! ``` +//! +//! on a regular Cartesian box `[x_min, x_max]^d` with Neumann (zero-flux) +//! boundary conditions and explicit Euler time-stepping. A user-supplied +//! Hamiltonian closure `H(x, ∇v) -> ℝ` evaluates the inf/sup over controls; +//! the solver only requires the resulting scalar. +//! +//! The implementation supports `d = 1, 2, 3` (more dimensions are accepted +//! but memory grows as `n_per_dim^d`). + +use crate::core::{OptimizrError, Result}; +use ndarray::Array1; + +#[derive(Clone, Debug)] +pub struct HjbMultidConfig { + /// Number of spatial dimensions. + pub dim: usize, + /// Number of grid points per dimension (uniform). + pub n_per_dim: usize, + pub x_min: f64, + pub x_max: f64, + pub n_t: usize, + pub t_horizon: f64, + /// Constant isotropic diffusion coefficient σ² ≥ 0. + pub sigma_sq: f64, +} + +impl HjbMultidConfig { + pub fn validate(&self) -> Result<()> { + if self.dim == 0 || self.dim > 3 { + return Err(OptimizrError::InvalidParameter("dim must be 1, 2 or 3".into())); + } + if self.n_per_dim < 3 { + return Err(OptimizrError::InvalidParameter("n_per_dim must be ≥ 3".into())); + } + if !(self.x_max > self.x_min) { + return Err(OptimizrError::InvalidParameter("x_max > x_min".into())); + } + if self.n_t == 0 { + return Err(OptimizrError::InvalidParameter("n_t > 0".into())); + } + if !(self.t_horizon > 0.0) { + return Err(OptimizrError::InvalidParameter("t_horizon > 0".into())); + } + if self.sigma_sq < 0.0 { + return Err(OptimizrError::InvalidParameter("sigma_sq ≥ 0".into())); + } + Ok(()) + } + + pub fn dx(&self) -> f64 { + (self.x_max - self.x_min) / (self.n_per_dim - 1) as f64 + } + + pub fn dt(&self) -> f64 { + self.t_horizon / self.n_t as f64 + } + + pub fn total_size(&self) -> usize { + self.n_per_dim.pow(self.dim as u32) + } +} + +#[derive(Clone, Debug)] +pub struct HjbMultidResult { + pub value: Vec, + pub grid_axes: Vec>, +} + +/// Convert a flat index to multi-index (lexicographic, last dim fastest). +fn flat_to_multi(idx: usize, n: usize, dim: usize) -> Vec { + let mut out = vec![0usize; dim]; + let mut r = idx; + for d in (0..dim).rev() { + out[d] = r % n; + r /= n; + } + out +} + +fn multi_to_flat(mi: &[usize], n: usize) -> usize { + let mut idx = 0usize; + for &m in mi { + idx = idx * n + m; + } + idx +} + +pub fn solve_hjb_multid( + hamiltonian: H, + terminal: G, + cfg: &HjbMultidConfig, +) -> Result +where + H: Fn(&[f64], &[f64]) -> f64, // H(x, grad_v) + G: Fn(&[f64]) -> f64, +{ + cfg.validate()?; + let n = cfg.n_per_dim; + let d = cfg.dim; + let dx = cfg.dx(); + let dt = cfg.dt(); + let size = cfg.total_size(); + + let cfl = dt * (cfg.sigma_sq * d as f64 / (dx * dx)); + if cfl > 0.5 { + return Err(OptimizrError::NumericalError(format!( + "explicit HJB CFL violated: dt·d·σ²/dx² = {cfl:.3} > 0.5" + ))); + } + + let grid_axes: Vec> = + (0..d).map(|_| Array1::from_iter((0..n).map(|i| cfg.x_min + i as f64 * dx))).collect(); + + // Cache positions per flat index. + let mut pos = vec![0.0f64; size * d]; + for idx in 0..size { + let mi = flat_to_multi(idx, n, d); + for k in 0..d { + pos[idx * d + k] = grid_axes[k][mi[k]]; + } + } + + // Terminal condition. + let mut v = vec![0.0f64; size]; + for idx in 0..size { + v[idx] = terminal(&pos[idx * d..(idx + 1) * d]); + } + let mut v_new = vec![0.0f64; size]; + let mut grad = vec![0.0f64; d]; + let mut x_local = vec![0.0f64; d]; + + for _step in 0..cfg.n_t { + for idx in 0..size { + let mi = flat_to_multi(idx, n, d); + for k in 0..d { + x_local[k] = pos[idx * d + k]; + } + // central differences with reflective (Neumann) boundary. + let mut lap = 0.0; + for k in 0..d { + let mut mi_p = mi.clone(); + let mut mi_m = mi.clone(); + if mi[k] + 1 < n { mi_p[k] += 1; } else { mi_p[k] = mi[k]; } + if mi[k] >= 1 { mi_m[k] -= 1; } else { mi_m[k] = mi[k]; } + let v_p = v[multi_to_flat(&mi_p, n)]; + let v_m = v[multi_to_flat(&mi_m, n)]; + grad[k] = (v_p - v_m) / (2.0 * dx); + lap += (v_p - 2.0 * v[idx] + v_m) / (dx * dx); + } + let h_val = hamiltonian(&x_local, &grad); + // Backward time: v^{n} = v^{n+1} + dt * ((σ²/2) Δv - H) + v_new[idx] = v[idx] + dt * (0.5 * cfg.sigma_sq * lap - h_val); + } + std::mem::swap(&mut v, &mut v_new); + } + + Ok(HjbMultidResult { value: v, grid_axes }) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// `H = 0`, `σ² = 0` and constant terminal should stay constant. + #[test] + fn trivial_problem_preserves_constant() { + let cfg = HjbMultidConfig { + dim: 2, + n_per_dim: 9, + x_min: -1.0, + x_max: 1.0, + n_t: 50, + t_horizon: 1.0, + sigma_sq: 0.0, + }; + let res = solve_hjb_multid(|_, _| 0.0, |_| 3.14, &cfg).unwrap(); + for v in res.value.iter() { + assert!((v - 3.14).abs() < 1e-12); + } + } + + /// Pure heat (`H = 0`, σ² > 0) preserves the integral and average value + /// of a constant initial condition (Neumann BCs). + #[test] + fn pure_heat_preserves_constant() { + let cfg = HjbMultidConfig { + dim: 1, + n_per_dim: 21, + x_min: -1.0, + x_max: 1.0, + n_t: 100, + t_horizon: 0.1, + sigma_sq: 0.5, + }; + let res = solve_hjb_multid(|_, _| 0.0, |_| 1.0, &cfg).unwrap(); + for v in res.value.iter() { + assert!((v - 1.0).abs() < 1e-9); + } + } +} diff --git a/src/pde/mod.rs b/src/pde/mod.rs new file mode 100644 index 0000000..3454387 --- /dev/null +++ b/src/pde/mod.rs @@ -0,0 +1,19 @@ +//! Partial Differential Equation solvers (CPU-only finite differences) +//! ==================================================================== +//! +//! Generic PDE primitives: +//! +//! - [`fokker_planck`] — 1-D forward Fokker–Planck (Kolmogorov forward) solver +//! with conservative central differences. +//! - [`hjb_multid`] — explicit upwind scheme for multidimensional +//! Hamilton–Jacobi–Bellman equations on a regular Cartesian grid. +//! - [`elliptic_fd`] — Jacobi/Gauss–Seidel iteration for the Poisson equation +//! `-Δu = f` with Dirichlet boundary conditions on a 2-D rectangle. + +pub mod fokker_planck; +pub mod hjb_multid; +pub mod elliptic_fd; + +pub use fokker_planck::{FokkerPlanckConfig, FokkerPlanckResult, solve_fokker_planck_1d}; +pub use hjb_multid::{HjbMultidConfig, HjbMultidResult, solve_hjb_multid}; +pub use elliptic_fd::{EllipticFdConfig, EllipticFdResult, solve_poisson_2d}; diff --git a/src/stochastic_control/mod.rs b/src/stochastic_control/mod.rs new file mode 100644 index 0000000..0d5565f --- /dev/null +++ b/src/stochastic_control/mod.rs @@ -0,0 +1,21 @@ +//! Stochastic control primitives +//! =============================== +//! +//! - [`optimal_switching`] — backward dynamic-programming Snell envelope for +//! discrete-time multi-mode optimal switching with mode-dependent running +//! reward and switching costs. +//! - [`pontryagin`] — forward shooting solver for the Pontryagin maximum +//! principle on a 1-D controlled SDE with quadratic cost. +//! - [`two_sided_intensity_control`] — generic bilateral intensity control +//! for a doubly-controlled jump process; computes the optimal symmetric +//! intensities from the value function gradient. + +pub mod optimal_switching; +pub mod pontryagin; +pub mod two_sided_intensity_control; + +pub use optimal_switching::{SwitchingConfig, SwitchingResult, solve_optimal_switching}; +pub use pontryagin::{PontryaginConfig, PontryaginResult, solve_pontryagin_lqr}; +pub use two_sided_intensity_control::{ + TwoSidedConfig, TwoSidedIntensities, optimal_two_sided_intensities, +}; diff --git a/src/stochastic_control/optimal_switching.rs b/src/stochastic_control/optimal_switching.rs new file mode 100644 index 0000000..65bf1f6 --- /dev/null +++ b/src/stochastic_control/optimal_switching.rs @@ -0,0 +1,136 @@ +//! Discrete-time optimal switching (Snell envelope) +//! ================================================== +//! +//! Solves +//! +//! ```text +//! V_k(i) = ψ_k(i) + max_{ j ≠ i } [ V_{k+1}(j) - c(i, j) ], V_N(i) = G(i) +//! ``` +//! +//! by backward induction. Here `i ∈ {0, ..., M-1}` is the current operating +//! mode, `ψ_k(i)` is the stage reward, `c(i, j) ≥ 0` is the cost of +//! switching from `i` to `j` (zero on the diagonal), and `G` is the +//! terminal payoff. + +use crate::core::{OptimizrError, Result}; + +#[derive(Clone, Debug)] +pub struct SwitchingConfig { + pub n_modes: usize, + pub n_steps: usize, +} + +#[derive(Clone, Debug)] +pub struct SwitchingResult { + /// `value[k * n_modes + i] = V_k(i)`. + pub value: Vec, + /// `policy[k * n_modes + i] = optimal next mode at time k from mode i`. + pub policy: Vec, +} + +pub fn solve_optimal_switching( + stage_reward: R, + terminal_payoff: T, + switching_cost: &[f64], + cfg: &SwitchingConfig, +) -> Result +where + R: Fn(usize, usize) -> f64, + T: Fn(usize) -> f64, +{ + if cfg.n_modes == 0 || cfg.n_steps == 0 { + return Err(OptimizrError::InvalidParameter( + "n_modes and n_steps must be > 0".into(), + )); + } + if switching_cost.len() != cfg.n_modes * cfg.n_modes { + return Err(OptimizrError::DimensionMismatch { + expected: cfg.n_modes * cfg.n_modes, + actual: switching_cost.len(), + }); + } + for i in 0..cfg.n_modes { + if switching_cost[i * cfg.n_modes + i].abs() > 1e-12 { + return Err(OptimizrError::InvalidParameter( + "diagonal of switching cost must be zero".into(), + )); + } + for j in 0..cfg.n_modes { + if switching_cost[i * cfg.n_modes + j] < 0.0 { + return Err(OptimizrError::InvalidParameter( + "switching costs must be non-negative".into(), + )); + } + } + } + + let m = cfg.n_modes; + let n = cfg.n_steps; + let mut value = vec![0.0f64; (n + 1) * m]; + let mut policy = vec![0usize; (n + 1) * m]; + for i in 0..m { + value[n * m + i] = terminal_payoff(i); + policy[n * m + i] = i; + } + for k in (0..n).rev() { + for i in 0..m { + let mut best_val = f64::NEG_INFINITY; + let mut best_j = i; + for j in 0..m { + let candidate = value[(k + 1) * m + j] - switching_cost[i * m + j]; + if candidate > best_val { + best_val = candidate; + best_j = j; + } + } + value[k * m + i] = stage_reward(k, i) + best_val; + policy[k * m + i] = best_j; + } + } + Ok(SwitchingResult { value, policy }) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Two modes, mode 1 always pays 1, mode 0 pays 0. Switching is free. + /// Optimal: stay in mode 1 from the start. V_0(1) = T, V_0(0) = T. + #[test] + fn free_switching_picks_paying_mode() { + let m = 2; + let n = 5; + let cost = vec![0.0; m * m]; // free + let cfg = SwitchingConfig { n_modes: m, n_steps: n }; + let res = solve_optimal_switching( + |_, i| if i == 1 { 1.0 } else { 0.0 }, + |_| 0.0, + &cost, + &cfg, + ) + .unwrap(); + // V_k(0) = sum_{l=k}^{n-1} max stage reward = 1·(n-k) (switch immediately at no cost, + // but stage reward at time k is paid in current mode i; here we receive 0 at time k + // and best continuation from any next mode j chosen at time k). + // Concretely: at time k, choose next mode 1; pay no cost; V_{k+1}(1) accrues. + // V_n(·) = 0 → V_{n-1}(0) = 0 + (V_n(1) - 0) = 0 + // V_{n-1}(1) = 1 + 0 = 1 + // V_{n-2}(0) = 0 + V_{n-1}(1) = 1 + // V_{n-2}(1) = 1 + V_{n-1}(1) = 2 + // ⇒ V_0(0) = n - 1 = 4, V_0(1) = n = 5 + assert!((res.value[0] - 4.0).abs() < 1e-12); + assert!((res.value[1] - 5.0).abs() < 1e-12); + } + + #[test] + fn high_switch_cost_locks_mode() { + let m = 2; + let n = 3; + let cost = vec![0.0, 1e6, 1e6, 0.0]; + let cfg = SwitchingConfig { n_modes: m, n_steps: n }; + let res = solve_optimal_switching(|_, i| i as f64, |_| 0.0, &cost, &cfg).unwrap(); + // Starting in mode 0: switching to 1 costs 1e6 → never switch. + // V_0(0) = 0 (stage) + V_1(0) = 0 + 0 + V_2(0) = 0 + V_3(0) = 0. + assert!((res.value[0]).abs() < 1e-9); + } +} diff --git a/src/stochastic_control/pontryagin.rs b/src/stochastic_control/pontryagin.rs new file mode 100644 index 0000000..51c9d3c --- /dev/null +++ b/src/stochastic_control/pontryagin.rs @@ -0,0 +1,119 @@ +//! Pontryagin maximum principle — 1-D LQR shooting solver +//! ======================================================== +//! +//! Closed-form Pontryagin solution for the 1-D controlled linear-quadratic +//! regulator +//! +//! ```text +//! dx/dt = a x + b u, x(0) = x_0 +//! J = ∫_0^T (q x² + r u²) dt + s_T x(T)² +//! ``` +//! +//! The Hamiltonian `H = p (a x + b u) + q x² + r u²` is minimised at +//! `u* = -b p / (2 r)`, giving the costate ODE `-dp/dt = 2 q x + a p`. +//! We integrate the resulting Riccati equation `dP/dt = -2 a P + b² P² / r - 2 q` +//! backwards from `P(T) = s_T` and recover `u*(t) = -(b/r) P(t) x(t)`. + +use crate::core::{OptimizrError, Result}; +use ndarray::Array1; + +#[derive(Clone, Debug)] +pub struct PontryaginConfig { + pub a: f64, + pub b: f64, + pub q: f64, + pub r: f64, + pub s_terminal: f64, + pub x0: f64, + pub t_horizon: f64, + pub n_steps: usize, +} + +#[derive(Clone, Debug)] +pub struct PontryaginResult { + pub time_grid: Array1, + pub state: Array1, + pub control: Array1, + pub riccati: Array1, + pub cost: f64, +} + +pub fn solve_pontryagin_lqr(cfg: &PontryaginConfig) -> Result { + if cfg.r <= 0.0 { + return Err(OptimizrError::InvalidParameter("r must be > 0".into())); + } + if cfg.n_steps == 0 || cfg.t_horizon <= 0.0 { + return Err(OptimizrError::InvalidParameter("n_steps>0 and T>0".into())); + } + let n = cfg.n_steps; + let dt = cfg.t_horizon / n as f64; + let time_grid: Array1 = Array1::from_iter((0..=n).map(|k| k as f64 * dt)); + + // Backward Riccati (explicit Euler). + let mut p = vec![0.0f64; n + 1]; + p[n] = cfg.s_terminal; + for k in (0..n).rev() { + let pn = p[k + 1]; + let dp = -2.0 * cfg.a * pn + cfg.b * cfg.b * pn * pn / cfg.r - 2.0 * cfg.q; + p[k] = pn - dt * dp; + } + + // Forward state integration with optimal feedback u* = -(b/r) P x. + let mut x = Array1::::zeros(n + 1); + let mut u = Array1::::zeros(n); + x[0] = cfg.x0; + let mut cost = 0.0; + for k in 0..n { + let u_k = -(cfg.b / cfg.r) * p[k] * x[k]; + u[k] = u_k; + let dx = cfg.a * x[k] + cfg.b * u_k; + x[k + 1] = x[k] + dt * dx; + cost += dt * (cfg.q * x[k] * x[k] + cfg.r * u_k * u_k); + } + cost += cfg.s_terminal * x[n] * x[n]; + + Ok(PontryaginResult { + time_grid, + state: x, + control: u, + riccati: Array1::from(p), + cost, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Closed-form check: with a = 0, b = 1, q = 0, r = 1, s_T > 0, T = 1, + /// the Riccati ODE becomes dP/dt = P²; solution P(t) = s_T / (1 + s_T (T-t)). + #[test] + fn riccati_matches_closed_form() { + let cfg = PontryaginConfig { + a: 0.0, + b: 1.0, + q: 0.0, + r: 1.0, + s_terminal: 1.0, + x0: 1.0, + t_horizon: 1.0, + n_steps: 2000, + }; + let res = solve_pontryagin_lqr(&cfg).unwrap(); + let analytic_p0 = 1.0 / (1.0 + 1.0 * 1.0); // 0.5 + let err = (res.riccati[0] - analytic_p0).abs(); + assert!(err < 5e-3, "P(0) = {}, analytic = {}", res.riccati[0], analytic_p0); + // Optimal cost analytic = P(0) * x0² + let analytic_cost = analytic_p0; + assert!((res.cost - analytic_cost).abs() < 5e-3); + } + + #[test] + fn rejects_zero_control_weight() { + let cfg = PontryaginConfig { + a: 0.0, b: 1.0, q: 1.0, r: 0.0, s_terminal: 0.0, + x0: 1.0, t_horizon: 1.0, n_steps: 10, + }; + assert!(solve_pontryagin_lqr(&cfg).is_err()); + } +} diff --git a/src/stochastic_control/two_sided_intensity_control.rs b/src/stochastic_control/two_sided_intensity_control.rs new file mode 100644 index 0000000..69c3de2 --- /dev/null +++ b/src/stochastic_control/two_sided_intensity_control.rs @@ -0,0 +1,115 @@ +//! Generic two-sided intensity control +//! ===================================== +//! +//! Considers a controlled bilateral jump process where the controller picks +//! two non-negative intensities `λ_+, λ_-` for upward and downward jumps of +//! a scalar state `q`. The instantaneous reward density is +//! +//! ```text +//! r(q, λ_+, λ_-) = λ_+ (δ_+(λ_+) - ΔV(q, +1)) + λ_- (δ_-(λ_-) - ΔV(q, -1)) +//! ``` +//! +//! where `δ_±(λ) = α_± + κ_± λ` is an affine *generic* per-jump premium and +//! `ΔV(q, ±1) = V(q ± 1) - V(q)` is the value-function differential. The +//! first-order condition gives the optimal symmetric pair +//! +//! ```text +//! λ_*± = max(0, (α_± - ΔV(q, ±1)) / (2 κ_±)) +//! ``` +//! +//! This module exposes a single helper that, given the value-function +//! differentials and the affine coefficients, returns the optimal +//! intensities and the resulting reward. All quantities are kept in +//! purely abstract / generic form. + +use crate::core::{OptimizrError, Result}; + +#[derive(Clone, Debug)] +pub struct TwoSidedConfig { + /// Affine intercept `α_+` of the upward premium. + pub alpha_plus: f64, + /// Affine intercept `α_-` of the downward premium. + pub alpha_minus: f64, + /// Affine slope `κ_+` of the upward premium (must be > 0). + pub kappa_plus: f64, + /// Affine slope `κ_-` of the downward premium (must be > 0). + pub kappa_minus: f64, +} + +#[derive(Clone, Debug)] +pub struct TwoSidedIntensities { + pub lambda_plus: f64, + pub lambda_minus: f64, + pub reward_density: f64, +} + +pub fn optimal_two_sided_intensities( + cfg: &TwoSidedConfig, + delta_v_plus: f64, + delta_v_minus: f64, +) -> Result { + if !(cfg.kappa_plus > 0.0 && cfg.kappa_minus > 0.0) { + return Err(OptimizrError::InvalidParameter( + "kappa_plus and kappa_minus must be > 0".into(), + )); + } + let raw_plus = (cfg.alpha_plus - delta_v_plus) / (2.0 * cfg.kappa_plus); + let raw_minus = (cfg.alpha_minus - delta_v_minus) / (2.0 * cfg.kappa_minus); + let lambda_plus = raw_plus.max(0.0); + let lambda_minus = raw_minus.max(0.0); + let premium_plus = cfg.alpha_plus + cfg.kappa_plus * lambda_plus; + let premium_minus = cfg.alpha_minus + cfg.kappa_minus * lambda_minus; + let reward_density = lambda_plus * (premium_plus - delta_v_plus) + + lambda_minus * (premium_minus - delta_v_minus); + Ok(TwoSidedIntensities { + lambda_plus, + lambda_minus, + reward_density, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn symmetric_zero_gradient_recovers_alpha_over_two_kappa() { + let cfg = TwoSidedConfig { + alpha_plus: 1.0, + alpha_minus: 1.0, + kappa_plus: 0.5, + kappa_minus: 0.5, + }; + let r = optimal_two_sided_intensities(&cfg, 0.0, 0.0).unwrap(); + let expected = 1.0 / (2.0 * 0.5); + assert!((r.lambda_plus - expected).abs() < 1e-12); + assert!((r.lambda_minus - expected).abs() < 1e-12); + // Reward density = λ (α + κ λ) = 1.0 * (1.0 + 0.5 * 1.0) = 1.5; symmetric → 3.0. + assert!((r.reward_density - 3.0).abs() < 1e-12); + } + + #[test] + fn high_value_gradient_kills_intensity() { + let cfg = TwoSidedConfig { + alpha_plus: 1.0, + alpha_minus: 1.0, + kappa_plus: 0.5, + kappa_minus: 0.5, + }; + let r = optimal_two_sided_intensities(&cfg, 5.0, 5.0).unwrap(); + assert_eq!(r.lambda_plus, 0.0); + assert_eq!(r.lambda_minus, 0.0); + assert!(r.reward_density.abs() < 1e-12); + } + + #[test] + fn rejects_zero_kappa() { + let cfg = TwoSidedConfig { + alpha_plus: 1.0, + alpha_minus: 1.0, + kappa_plus: 0.0, + kappa_minus: 0.5, + }; + assert!(optimal_two_sided_intensities(&cfg, 0.0, 0.0).is_err()); + } +}