fix: possion processes

src/jumps/bates.rs

@@ -24,7 +24,7 @@ pub fn bates_1996(
   let mut s = Array1::<f64>::zeros(n);
   let mut v = Array1::<f64>::zeros(n);
 
-  let [jump_times, jumps] = compound_poisson(n, lambda, None, Some(t.unwrap_or(1.0)));
+  let [jump_times, jumps] = compound_poisson(n, lambda, None, Some(t.unwrap_or(1.0)), None, None);
 
   s[0] = s0.unwrap_or(0.0);
   v[0] = v0.unwrap_or(0.0);

src/processes/poisson.rs

@@ -1,57 +1,51 @@
-use std::ops::Div;
-
 use ndarray::Array1;
-use ndarray_rand::RandomExt;
 use rand::{thread_rng, Rng};
+use rand_distr::{Distribution, Exp, Normal};
 
-pub fn poisson(n: usize, lambda: usize, t: Option<f64>) -> Vec<f64> {
+pub fn poisson(n: usize, lambda: usize, t_max: Option<f64>) -> Vec<f64> {
   if n == 0 || lambda == 0 {
     panic!("lambda, t and n must be positive integers");
   }
 
-  if t.is_none() {
-    let exponentials =
-      Array1::<f64>::random(n - 1, rand_distr::Exp::new(1.0.div(lambda as f64)).unwrap());
-    let mut p = Array1::<f64>::zeros(n);
-    p[0] = 0.0;
-
-    for i in 1..n {
-      p[i] = p[i - 1] + exponentials[i - 1];
-    }
-
-    p.to_vec()
-  } else {
-    let mut p = vec![0.0];
-    let mut _t = 0.0;
+  let t_max = t_max.unwrap_or(1.0);
+  let mut times = vec![0.0];
+  let exp = Exp::new(lambda as f64).unwrap();
 
-    while _t < t.unwrap() {
-      let exponential = thread_rng().sample(rand_distr::Exp::new(1.0.div(lambda as f64)).unwrap());
-      _t += exponential;
-      p.push(_t);
+  while times.last().unwrap() < &t_max {
+    let inter_arrival = exp.sample(&mut thread_rng());
+    let next_time = times.last().unwrap() + inter_arrival;
+    if next_time > t_max {
+      break;
     }
-
-    p
+    times.push(next_time);
   }
+
+  times
 }
 
 pub fn compound_poisson(
   n: usize,
   lambda: usize,
   jumps: Option<Vec<f64>>,
-  t: Option<f64>,
+  t_max: Option<f64>,
+  jump_mean: Option<f64>,
+  jump_std: Option<f64>,
 ) -> [Vec<f64>; 2] {
   if n == 0 {
     panic!("n must be a positive integer");
   }
 
-  let times = poisson(n, lambda, t);
+  let _t_max = t_max.unwrap_or(1.0);
+  let times = poisson(n, lambda, Some(_t_max));
   let mut cp = Array1::<f64>::zeros(times.len());
   let jumps = match jumps {
     Some(jumps) => jumps,
-    None => thread_rng()
-      .sample_iter(rand_distr::StandardNormal)
-      .take(n)
-      .collect(),
+    None => {
+      let _jump_mean = jump_mean.unwrap_or(1.0);
+      let _jump_std = jump_std.unwrap_or(0.0);
+      let norm = Normal::new(_jump_mean, _jump_std).unwrap();
+      thread_rng().sample_iter(norm).take(n).collect()
+    }
   };
 
   for i in 1..times.len() {
@@ -60,3 +54,27 @@ pub fn compound_poisson(
 
   [times, cp.to_vec()]
 }
+
+#[cfg(test)]
+mod tests {
+  use super::*;
+
+  #[test]
+  fn test_poisson() {
+    let n = 1000;
+    let lambda = 10;
+    let t = 10.0;
+    let p = poisson(n, lambda, Some(t));
+    println!("{:?}", p.len());
+  }
+
+  #[test]
+  fn test_compound_poisson() {
+    let n = 1000;
+    let lambda = 10;
+    let t = 10.0;
+    let cp = compound_poisson(n, lambda, None, Some(t), Some(1.0), Some(0.0));
+    println!("{:?}", cp[0].len());
+    println!("{:?}", cp[1].len());
+  }
+}