diff --git a/src/jumps/bates.rs b/src/jumps/bates.rs index 52c96d3..5115f58 100644 --- a/src/jumps/bates.rs +++ b/src/jumps/bates.rs @@ -24,7 +24,7 @@ pub fn bates_1996( let mut s = Array1::::zeros(n); let mut v = Array1::::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); diff --git a/src/processes/poisson.rs b/src/processes/poisson.rs index a46d8b6..f000009 100644 --- a/src/processes/poisson.rs +++ b/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) -> Vec { +pub fn poisson(n: usize, lambda: usize, t_max: Option) -> Vec { if n == 0 || lambda == 0 { panic!("lambda, t and n must be positive integers"); } - if t.is_none() { - let exponentials = - Array1::::random(n - 1, rand_distr::Exp::new(1.0.div(lambda as f64)).unwrap()); - let mut p = Array1::::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>, - t: Option, + t_max: Option, + jump_mean: Option, + jump_std: Option, ) -> [Vec; 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::::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()); + } +}