fix: distributions beta,gamma,geometric,nig

Author: dancixx

src/stats/distr/beta.rs

@@ -16,7 +16,7 @@ pub struct SimdBeta {
 
 impl SimdBeta {
   pub fn new(alpha: f32, beta: f32) -> Self {
-    assert!(alpha >= 1.0 && beta >= 1.0);
+    assert!(alpha > 0.0 && beta > 0.0);
     Self {
       alpha,
       beta,

src/stats/distr/gamma.rs

@@ -15,7 +15,7 @@ pub struct SimdGamma {
 
 impl SimdGamma {
   pub fn new(alpha: f32, scale: f32) -> Self {
-    assert!(alpha >= 1.0 && scale > 0.0);
+    assert!(alpha > 0.0 && scale > 0.0);
     Self {
       alpha,
       scale,
@@ -25,26 +25,40 @@ impl SimdGamma {
     }
   }
 
-  /// Bulk fill using Marsaglia–Tsang; uses scalar acceptance per sample but reduces per-call overhead.
+  /// Bulk fill using Marsaglia–Tsang (for alpha >= 1) or Ahrens-Dieter (for alpha < 1)
   pub fn fill_slice<R: Rng + ?Sized>(&self, rng: &mut R, out: &mut [f32]) {
-    let d = self.alpha - 1.0 / 3.0;
-    let c = 1.0 / (3.0 * d).sqrt();
-    for x in out.iter_mut() {
-      let val = loop {
-        let z = self.normal.sample(rng);
-        let v = (1.0 + c * z).powi(3);
-        if v <= 0.0 {
-          continue;
-        }
+    if self.alpha < 1.0 {
+      // For alpha < 1, use the transformation: if X ~ Gamma(alpha+1, scale), then X*U^(1/alpha) ~ Gamma(alpha, scale)
+      let gamma_plus_one = SimdGamma::new(self.alpha + 1.0, self.scale);
+      for x in out.iter_mut() {
+        let g = gamma_plus_one.sample(rng);
         let u: f32 = rng.gen_range(0.0..1.0);
-        if u < 1.0 - 0.0331 * z.powi(4) {
-          break self.scale * d * v;
-        }
-        if u.ln() < 0.5 * z * z + d * (1.0 - v + v.ln()) {
-          break self.scale * d * v;
-        }
-      };
-      *x = val;
+        *x = g * u.powf(1.0 / self.alpha);
+      }
+    } else {
+      // Marsaglia-Tsang for alpha >= 1
+      let d = self.alpha - 1.0 / 3.0;
+      let c = 1.0 / (9.0 * d).sqrt();
+      for x in out.iter_mut() {
+        let val = loop {
+          let z = self.normal.sample(rng);
+          let v = (1.0 + c * z).powi(3);
+          if v <= 0.0 {
+            continue;
+          }
+          let u: f32 = rng.gen_range(0.0..1.0);
+          let z2 = z * z;
+          // Quick acceptance
+          if u < 1.0 - 0.0331 * z2 * z2 {
+            break d * v;
+          }
+          // Log acceptance
+          if u.ln() < 0.5 * z2 + d * (1.0 - v + v.ln()) {
+            break d * v;
+          }
+        };
+        *x = self.scale * val;
+      }
     }
   }
 

src/stats/distr/geometric.rs

@@ -22,6 +22,8 @@ impl SimdGeometric {
     use crate::stats::distr::fill_f32_zero_one;
     use wide::f32x8;
 
+    // rand_distr Geometric returns number of failures before first success (starts at 0)
+    // Formula: floor(ln(U) / ln(1-p)) where U ~ Uniform(0,1)
     let ln1p = (1.0 - self.p).ln();
     let inv_ln1p = f32x8::splat(1.0 / ln1p);
     let mut u = [0.0f32; 8];
@@ -30,10 +32,11 @@ impl SimdGeometric {
     for chunk in &mut chunks {
       fill_f32_zero_one(rng, &mut u);
       let v = f32x8::from(u);
-      let g = (v.ln() * inv_ln1p).floor() + f32x8::splat(1.0);
+      // Number of failures before success (can be 0)
+      let g = (v.ln() * inv_ln1p).floor();
       let mut tmp = g.to_array();
       for t in &mut tmp {
-        *t = (*t).max(1.0);
+        *t = (*t).max(0.0);
       }
       for (o, t) in chunk.iter_mut().zip(tmp.iter()) {
         *o = *t as u32;
@@ -43,10 +46,10 @@ impl SimdGeometric {
     if !rem.is_empty() {
       fill_f32_zero_one(rng, &mut u);
       let v = f32x8::from(u);
-      let g = (v.ln() * inv_ln1p).floor() + f32x8::splat(1.0);
+      let g = (v.ln() * inv_ln1p).floor();
       let tmp = g.to_array();
       for i in 0..rem.len() {
-        let val = tmp[i].max(1.0);
+        let val = tmp[i].max(0.0);
         rem[i] = val as u32;
       }
     }

src/stats/distr/normal_inverse_gauss.rs

@@ -18,9 +18,19 @@ pub struct SimdNormalInverseGauss {
 }
 
 impl SimdNormalInverseGauss {
+  /// Create a NIG with 4 parameters (alpha, beta, delta, mu)
+  /// This is the full parameterization.
+  /// NIG(α, β, δ, μ) where α > |β|, δ > 0
   pub fn new(alpha: f32, beta: f32, delta: f32, mu: f32) -> Self {
-    // Typically alpha> |beta|, delta>0, etc.
-    let ig = SimdInverseGauss::new(delta * (alpha * alpha - beta * beta).sqrt(), delta);
+    assert!(
+      alpha > 0.0 && alpha > beta.abs(),
+      "NIG: alpha must be > |beta|"
+    );
+    assert!(delta > 0.0, "NIG: delta must be positive");
+    let gamma = (alpha * alpha - beta * beta).sqrt();
+    let ig_mean = delta / gamma;
+    let ig_shape = delta * delta;
+    let ig = SimdInverseGauss::new(ig_mean, ig_shape);
     let normal = SimdNormal::new(0.0, 1.0);
     Self {
       alpha,