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::uniform::SimdUniform;
use crate::process::poisson::Poisson;
use crate::traits::FloatExt;
use crate::traits::ProcessExt;
pub struct Cts<T: FloatExt, S: SeedExt = Unseeded> {
pub lambda_plus: T, pub lambda_minus: T, pub alpha: T,
pub n: usize,
pub j: usize,
pub x0: Option<T>,
pub t: Option<T>,
pub seed: S,
}
impl<T: FloatExt> Cts<T> {
pub fn new(
lambda_plus: T,
lambda_minus: T,
alpha: T,
n: usize,
j: usize,
x0: 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)"
);
Self {
lambda_plus,
lambda_minus,
alpha,
n,
j,
x0,
t,
seed: Unseeded,
}
}
}
impl<T: FloatExt> Cts<T, Deterministic> {
pub fn seeded(
lambda_plus: T,
lambda_minus: T,
alpha: T,
n: usize,
j: usize,
x0: 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)"
);
Self {
lambda_plus,
lambda_minus,
alpha,
n,
j,
x0,
t,
seed: Deterministic::new(seed),
}
}
}
impl<T: FloatExt, S: SeedExt> ProcessExt<T> for Cts<T, S> {
type Output = Array1<T>;
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 g = gamma(2.0 - self.alpha.to_f64().unwrap());
let C = (T::from_f64_fast(g)
* (self.lambda_plus.powf(self.alpha - T::from_usize_(2))
+ self.lambda_minus.powf(self.alpha - T::from_usize_(2))))
.powi(-1);
let g = gamma(1.0 - self.alpha.to_f64().unwrap());
let b_t = -C
* T::from_f64_fast(g)
* (self.lambda_plus.powf(self.alpha - T::one())
- self.lambda_minus.powf(self.alpha - T::one()));
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 P = Poisson::new(T::one(), Some(size), None).sample();
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 mut jump_size = Array1::<T>::zeros(size);
for j in 1..size {
let v_j = if uniform.sample_fast() < T::from_f64_fast(0.5) {
self.lambda_plus
} else {
-self.lambda_minus
};
let numerator = self.alpha * P[j];
let term1 = (numerator / C).powf(-T::one() / self.alpha);
let term2 = E[j] * U[j].powf(T::one() / self.alpha) / v_j.abs();
jump_size[j] = term1.min(term2) * (v_j / v_j.abs());
}
let mut idx = (1..size).collect::<Vec<_>>();
idx.sort_by(|&a, &b| {
tau[a]
.to_f64()
.unwrap()
.partial_cmp(&tau[b].to_f64().unwrap())
.unwrap()
});
let mut x = Array1::<T>::zeros(self.n);
x[0] = self.x0.unwrap_or(T::zero());
let mut k: usize = 0;
let mut cum_jumps = T::zero();
for i in 1..self.n {
let t_i = T::from_usize_(i) * dt;
while k < idx.len() && tau[idx[k]] <= t_i {
cum_jumps += jump_size[idx[k]];
k += 1;
}
x[i] = x[0] + cum_jumps + b_t * t_i;
}
x
}
}
py_process_1d!(PyCts, Cts,
sig: (lambda_plus, lambda_minus, alpha, n, j, x0=None, t=None, seed=None, dtype=None),
params: (lambda_plus: f64, lambda_minus: f64, alpha: f64, n: usize, j: usize, x0: Option<f64>, t: Option<f64>)
);