feat: rework simd acceleration

Author: dancixx

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "stochastic-rs"
-version = "0.15.1"
+version = "0.15.2"
 edition = "2021"
 license = "MIT"
 description = "A Rust library for quant finance and simulating stochastic processes."
@@ -82,7 +82,7 @@ f64 = []
 jemalloc = ["dep:tikv-jemallocator"]
 malliavin = []
 mimalloc = ["dep:mimalloc"]
-simd = []
+simd = ["f32"]
 yahoo = ["dep:time", "dep:yahoo_finance_api", "dep:polars"]
 
 [lib]

src/stochastic/diffusion/cev.rs

@@ -61,25 +61,6 @@ impl SamplingExt<f64> for CEV<f64> {
     cev
   }
 
-  #[cfg(feature = "simd")]
-  fn sample_simd(&self) -> Array1<f64> {
-    use crate::stats::distr::normal::SimdNormal;
-
-    let dt = self.t.unwrap_or(1.0) / (self.n - 1) as f64;
-    let gn = Array1::random(self.n - 1, SimdNormal::new(0.0, dt.sqrt() as f32));
-
-    let mut cev = Array1::<f64>::zeros(self.n);
-    cev[0] = self.x0.unwrap_or(0.0);
-
-    for i in 1..self.n {
-      cev[i] = cev[i - 1]
-        + self.mu * cev[i - 1] * dt
-        + self.sigma * cev[i - 1].powf(self.gamma) * gn[i - 1] as f64
-    }
-
-    cev
-  }
-
   /// Number of time steps
   fn n(&self) -> usize {
     self.n

src/stochastic/diffusion/cir.rs

@@ -48,34 +48,6 @@ impl SamplingExt<f64> for CIR<f64> {
     cir
   }
 
-  #[cfg(feature = "simd")]
-  fn sample_simd(&self) -> Array1<f64> {
-    use crate::stats::distr::normal::SimdNormal;
-
-    assert!(
-      2.0 * self.theta * self.mu >= self.sigma.powi(2),
-      "2 * theta * mu < sigma^2"
-    );
-
-    let dt = self.t.unwrap_or(1.0) / (self.n - 1) as f64;
-    let gn = Array1::random(self.n - 1, SimdNormal::new(0.0, dt.sqrt() as f32));
-
-    let mut cir = Array1::<f64>::zeros(self.n);
-    cir[0] = self.x0.unwrap_or(0.0);
-
-    for i in 1..self.n {
-      let dcir = self.theta * (self.mu - cir[i - 1]) * dt
-        + self.sigma * (cir[i - 1]).abs().sqrt() * gn[i - 1] as f64;
-
-      cir[i] = match self.use_sym.unwrap_or(false) {
-        true => (cir[i - 1] + dcir).abs(),
-        false => (cir[i - 1] + dcir).max(0.0),
-      };
-    }
-
-    cir
-  }
-
   /// Number of time steps
   fn n(&self) -> usize {
     self.n

src/stochastic/diffusion/feller.rs

@@ -75,6 +75,30 @@ impl SamplingExt<f32> for FellerLogistic<f32> {
     x
   }
 
+  #[cfg(feature = "simd")]
+  fn sample_simd(&self) -> Array1<f32> {
+    use crate::stats::distr::normal::SimdNormal;
+
+    let dt = self.t.unwrap_or(1.0) / (self.n - 1) as f32;
+    let gn = Array1::random(self.n - 1, SimdNormal::new(0.0, dt.sqrt()));
+
+    let mut x = Array1::<f32>::zeros(self.n);
+    x[0] = self.x0.unwrap_or(0.0);
+
+    for i in 1..self.n {
+      let xi = x[i - 1].max(0.0);
+      let drift = self.kappa * (self.theta - xi) * xi * dt;
+      let diff = self.sigma * xi.sqrt() * gn[i - 1];
+      let next = xi + drift + diff;
+      x[i] = match self.use_sym.unwrap_or(false) {
+        true => next.abs(),
+        false => next.max(0.0),
+      };
+    }
+
+    x
+  }
+
   fn n(&self) -> usize {
     self.n
   }

src/stochastic/diffusion/gbm.rs

@@ -59,23 +59,6 @@ impl SamplingExt<f64> for GBM<f64> {
     gbm
   }
 
-  #[cfg(feature = "simd")]
-  fn sample_simd(&self) -> Array1<f64> {
-    use crate::stats::distr::normal::SimdNormal;
-
-    let dt = self.t.unwrap_or(1.0) / (self.n - 1) as f64;
-    let gn = Array1::random(self.n - 1, SimdNormal::new(0.0, dt.sqrt() as f32));
-
-    let mut gbm = Array1::<f64>::zeros(self.n);
-    gbm[0] = self.x0.unwrap_or(0.0);
-
-    for i in 1..self.n {
-      gbm[i] = gbm[i - 1] + self.mu * gbm[i - 1] * dt + self.sigma * gbm[i - 1] * gn[i - 1] as f64
-    }
-
-    gbm
-  }
-
   /// Number of time steps
   fn n(&self) -> usize {
     self.n

src/stochastic/diffusion/gbm_ih.rs

@@ -71,6 +71,28 @@ impl SamplingExt<f32> for GBMIH<f32> {
     x
   }
 
+  #[cfg(feature = "simd")]
+  fn sample_simd(&self) -> Array1<f32> {
+    use crate::stats::distr::normal::SimdNormal;
+
+    if let Some(s) = &self.sigmas {
+      assert_eq!(s.len(), self.n - 1, "sigmas length must be n - 1");
+    }
+
+    let dt = self.t.unwrap_or(1.0) / (self.n - 1) as f32;
+    let gn = Array1::random(self.n - 1, SimdNormal::new(0.0, dt.sqrt()));
+
+    let mut x = Array1::<f32>::zeros(self.n);
+    x[0] = self.x0.unwrap_or(0.0);
+
+    for i in 1..self.n {
+      let sigma_i = self.sigmas.as_ref().map(|s| s[i - 1]).unwrap_or(self.sigma);
+      x[i] = x[i - 1] + self.mu * x[i - 1] * dt + sigma_i * x[i - 1] * gn[i - 1];
+    }
+
+    x
+  }
+
   fn n(&self) -> usize {
     self.n
   }

src/stochastic/diffusion/gompertz.rs

@@ -67,6 +67,27 @@ impl SamplingExt<f32> for Gompertz<f32> {
     x
   }
 
+  #[cfg(feature = "simd")]
+  fn sample_simd(&self) -> Array1<f32> {
+    use crate::stats::distr::normal::SimdNormal;
+
+    let dt = self.t.unwrap_or(1.0) / (self.n - 1) as f32;
+    let gn = Array1::random(self.n - 1, SimdNormal::new(0.0, dt.sqrt()));
+
+    let mut x = Array1::<f32>::zeros(self.n);
+    x[0] = self.x0.unwrap_or(0.0).max(1e-8);
+
+    for i in 1..self.n {
+      let xi = x[i - 1].max(1e-8);
+      let drift = (self.a - self.b * xi.ln()) * xi * dt;
+      let diff = self.sigma * xi * gn[i - 1];
+      let next = xi + drift + diff;
+      x[i] = next.max(1e-8);
+    }
+
+    x
+  }
+
   fn n(&self) -> usize {
     self.n
   }

src/stochastic/diffusion/jacobi.rs

@@ -46,36 +46,6 @@ impl SamplingExt<f64> for Jacobi<f64> {
     jacobi
   }
 
-  #[cfg(feature = "simd")]
-  fn sample_simd(&self) -> Array1<f64> {
-    use crate::stats::distr::normal::SimdNormal;
-
-    assert!(self.alpha > 0.0, "alpha must be positive");
-    assert!(self.beta > 0.0, "beta must be positive");
-    assert!(self.sigma > 0.0, "sigma must be positive");
-    assert!(self.alpha < self.beta, "alpha must be less than beta");
-
-    let dt = self.t.unwrap_or(1.0) / (self.n - 1) as f64;
-    let gn = Array1::random(self.n - 1, SimdNormal::new(0.0, dt.sqrt() as f32));
-
-    let mut jacobi = Array1::<f64>::zeros(self.n);
-    jacobi[0] = self.x0.unwrap_or(0.0);
-
-    for i in 1..self.n {
-      jacobi[i] = match jacobi[i - 1] {
-        _ if jacobi[i - 1] <= 0.0 && i > 0 => 0.0,
-        _ if jacobi[i - 1] >= 1.0 && i > 0 => 1.0,
-        _ => {
-          jacobi[i - 1]
-            + (self.alpha - self.beta * jacobi[i - 1]) * dt
-            + self.sigma * (jacobi[i - 1] * (1.0 - jacobi[i - 1])).sqrt() * gn[i - 1] as f64
-        }
-      }
-    }
-
-    jacobi
-  }
-
   /// Number of time steps
   fn n(&self) -> usize {
     self.n

src/stochastic/diffusion/kimura.rs

@@ -71,6 +71,29 @@ impl SamplingExt<f32> for Kimura<f32> {
     x
   }
 
+  #[cfg(feature = "simd")]
+  fn sample_simd(&self) -> Array1<f32> {
+    use crate::stats::distr::normal::SimdNormal;
+
+    let dt = self.t.unwrap_or(1.0) / (self.n - 1) as f32;
+    let gn = Array1::random(self.n - 1, SimdNormal::new(0.0, dt.sqrt()));
+
+    let mut x = Array1::<f32>::zeros(self.n);
+    x[0] = self.x0.unwrap_or(0.0);
+
+    for i in 1..self.n {
+      let xi = x[i - 1].clamp(0.0, 1.0);
+      let sqrt_term = (xi * (1.0 - xi)).sqrt();
+      let drift = self.a * xi * (1.0 - xi) * dt;
+      let diff = self.sigma * sqrt_term * gn[i - 1];
+      let mut next = xi + drift + diff;
+      next = next.clamp(0.0, 1.0);
+      x[i] = next;
+    }
+
+    x
+  }
+
   fn n(&self) -> usize {
     self.n
   }

src/stochastic/diffusion/ou.rs

@@ -33,25 +33,6 @@ impl SamplingExt<f64> for OU<f64> {
     ou
   }
 
-  #[cfg(feature = "simd")]
-  // TODO: experimental
-  fn sample_simd(&self) -> Array1<f64> {
-    use crate::stats::distr::normal::SimdNormal;
-
-    let dt = self.t.unwrap_or(1.0) as f32 / (self.n - 1) as f32;
-    let gn = Array1::random(self.n, SimdNormal::new(0.0, dt.sqrt()));
-
-    let mut ou = Array1::<f64>::zeros(self.n);
-    ou[0] = self.x0.unwrap_or(0.0);
-
-    for i in 1..self.n {
-      ou[i] =
-        ou[i - 1] + self.theta * (self.mu - ou[i - 1]) * dt as f64 + self.sigma * gn[i - 1] as f64
-    }
-
-    ou
-  }
-
   /// Number of time steps
   fn n(&self) -> usize {
     self.n