use ndarray::Array1;
use scilib::math::basic::gamma;
use stochastic_rs_core::simd_rng::Deterministic;
use stochastic_rs_core::simd_rng::SeedExt;
use stochastic_rs_core::simd_rng::Unseeded;
use stochastic_rs_distributions::exp::SimdExp;
use stochastic_rs_distributions::non_central_chi_squared;
use stochastic_rs_distributions::uniform::SimdUniform;
use crate::process::poisson::Poisson;
use crate::traits::FloatExt;
use crate::traits::ProcessExt;
pub struct Svcgmy<T: FloatExt, S: SeedExt = Unseeded> {
pub lambda_plus: T,
pub lambda_minus: T,
pub alpha: T,
pub kappa: T,
pub eta: T,
pub zeta: T,
pub rho: T,
pub n: usize,
pub j: usize,
pub x0: Option<T>,
pub v0: Option<T>,
pub t: Option<T>,
pub seed: S,
}
impl<T: FloatExt> Svcgmy<T> {
pub fn new(
lambda_plus: T,
lambda_minus: T,
alpha: T,
kappa: T,
eta: T,
zeta: T,
rho: T,
n: usize,
j: usize,
x0: Option<T>,
v0: Option<T>,
t: Option<T>,
) -> Self {
assert!(lambda_plus > T::zero(), "lambda_plus must be positive");
assert!(lambda_minus > T::zero(), "lambda_minus must be positive");
assert!(
alpha > T::zero() && alpha < T::from_usize_(2),
"alpha must be in (0, 2)"
);
assert!(kappa > T::zero(), "kappa must be positive");
assert!(eta >= T::zero(), "eta must be non-negative");
assert!(zeta > T::zero(), "zeta must be positive");
assert!(n >= 2, "n must be >= 2");
if let Some(v0) = v0 {
assert!(v0 >= T::zero(), "v0 must be non-negative");
}
Self {
lambda_plus,
lambda_minus,
alpha,
kappa,
eta,
zeta,
rho,
n,
j,
x0,
v0,
t,
seed: Unseeded,
}
}
}
impl<T: FloatExt> Svcgmy<T, Deterministic> {
pub fn seeded(
lambda_plus: T,
lambda_minus: T,
alpha: T,
kappa: T,
eta: T,
zeta: T,
rho: T,
n: usize,
j: usize,
x0: Option<T>,
v0: Option<T>,
t: Option<T>,
seed: u64,
) -> Self {
assert!(lambda_plus > T::zero(), "lambda_plus must be positive");
assert!(lambda_minus > T::zero(), "lambda_minus must be positive");
assert!(
alpha > T::zero() && alpha < T::from_usize_(2),
"alpha must be in (0, 2)"
);
assert!(kappa > T::zero(), "kappa must be positive");
assert!(eta >= T::zero(), "eta must be non-negative");
assert!(zeta > T::zero(), "zeta must be positive");
assert!(n >= 2, "n must be >= 2");
if let Some(v0) = v0 {
assert!(v0 >= T::zero(), "v0 must be non-negative");
}
Self {
lambda_plus,
lambda_minus,
alpha,
kappa,
eta,
zeta,
rho,
n,
j,
x0,
v0,
t,
seed: Deterministic::new(seed),
}
}
}
impl<T: FloatExt, S: SeedExt> ProcessExt<T> for Svcgmy<T, S> {
type Output = [Array1<T>; 2];
fn sample(&self) -> Self::Output {
let t_max = self.t.unwrap_or(T::one());
let dt = t_max / T::from_usize_(self.n - 1);
let mut x = Array1::<T>::zeros(self.n);
let mut v = Array1::<T>::zeros(self.n);
let mut y = Array1::<T>::zeros(self.n);
x[0] = self.x0.unwrap_or(T::zero());
v[0] = self.v0.unwrap_or(T::zero());
y[0] = x[0] - self.rho * v[0];
let f2 = T::from_usize_(2);
let g = gamma(2.0 - self.alpha.to_f64().unwrap());
let C = T::one()
/ (T::from_f64_fast(g)
* (self.lambda_plus.powf(self.alpha - f2) + self.lambda_minus.powf(self.alpha - f2)));
let c = (f2 * self.kappa) / ((T::one() - (-self.kappa * dt).exp()) * self.zeta.powi(2));
let df = T::from_usize_(4) * self.kappa * self.eta / self.zeta.powi(2);
for i in 1..self.n {
let ncp = f2 * c * v[i - 1] * (-self.kappa * dt).exp();
let xi = non_central_chi_squared::sample(df, ncp);
v[i] = xi / (f2 * c);
}
let J = self.j;
let size = J + 1;
let uniform = SimdUniform::from_seed_source(T::zero(), T::one(), &self.seed);
let exp = SimdExp::from_seed_source(T::one(), &self.seed);
let mut U = Array1::<T>::zeros(size);
uniform.fill_slice_fast(U.as_slice_mut().unwrap());
let E = Array1::from_shape_fn(size, |_| exp.sample_fast());
let mut tau_raw = Array1::<T>::zeros(size);
uniform.fill_slice_fast(tau_raw.as_slice_mut().unwrap());
let tau = tau_raw * t_max;
let P = Poisson::new(T::one(), Some(size), None).sample();
let mut c_tau = Array1::<T>::zeros(size);
for j in 1..size {
let tau_j = tau[j];
let k = ((tau_j / dt).ceil()).min(T::from_usize_(self.n - 1));
let v_km1 = if k == T::zero() {
v[0]
} else {
v[k.to_usize().unwrap() - 1]
};
c_tau[j] = C * v_km1;
}
for i in 1..self.n {
let numerator = v[i - 1]
* (self.lambda_plus.powf(self.alpha - T::one())
- self.lambda_minus.powf(self.alpha - T::one()));
let denominator = (T::one() - self.alpha)
* (self.lambda_plus.powf(self.alpha - f2) + self.lambda_minus.powf(self.alpha - f2));
let b = -numerator / denominator;
let mut jump_component = T::zero();
let t_1 = T::from_usize_(i - 1) * dt;
let t = T::from_usize_(i) * dt;
for j in 1..=J {
if tau[j] > t_1 && tau[j] <= t {
let v_j = if uniform.sample_fast() < T::from_f64_fast(0.5) {
self.lambda_plus
} else {
-self.lambda_minus
};
let num = self.alpha * P[j];
let den = f2 * c_tau[j] * t_max;
let term1 = (num / den).powf(-T::one() / self.alpha);
let term2 = E[j] * U[j].powf(T::one() / self.alpha) / v_j.abs();
let min_term = term1.min(term2);
jump_component += min_term * (v_j / v_j.abs());
}
}
y[i] = y[i - 1] + jump_component + b * dt;
}
for i in 1..self.n {
x[i] = y[i] + self.rho * v[i];
}
[x, v]
}
}
py_process_2x1d!(PySvcgmy, Svcgmy,
sig: (lambda_plus, lambda_minus, alpha, kappa, eta, zeta, rho, n, j, x0=None, v0=None, t=None, seed=None, dtype=None),
params: (lambda_plus: f64, lambda_minus: f64, alpha: f64, kappa: f64, eta: f64, zeta: f64, rho: f64, n: usize, j: usize, x0: Option<f64>, v0: Option<f64>, t: Option<f64>)
);