stochastic-rs-stochastic 2.0.0

Stochastic process simulation.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136

//! # Poisson
//!
//! $$
//! \mathbb{P}(N=k)=e^{-\lambda}\frac{\lambda^k}{k!},\ k\in\mathbb N_0
//! $$
//!
use ndarray::Array1;
use ndarray::s;
use rand_distr::Distribution;
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 crate::traits::FloatExt;
use crate::traits::ProcessExt;

#[derive(Clone, Copy)]
pub struct Poisson<T: FloatExt, S: SeedExt = Unseeded> {
  /// Jump intensity (expected arrivals per unit time).
  pub lambda: T,
  /// Optional fixed number of sampled events.
  pub n: Option<usize>,
  /// Optional terminal time for horizon-based sampling.
  pub t_max: Option<T>,
  /// Seed strategy (compile-time: [`Unseeded`] or [`Deterministic`]).
  pub seed: S,
}

#[inline]
fn validate_n_or_tmax<T: FloatExt>(n: Option<usize>, t_max: Option<T>, type_name: &'static str) {
  if n.is_none() && t_max.is_none() {
    panic!("{type_name}: n or t_max must be provided");
  }
}

impl<T: FloatExt> Poisson<T> {
  pub fn new(lambda: T, n: Option<usize>, t_max: Option<T>) -> Self {
    validate_n_or_tmax(n, t_max, "Poisson");
    Poisson {
      lambda,
      n,
      t_max,
      seed: Unseeded,
    }
  }
}

impl<T: FloatExt> Poisson<T, Deterministic> {
  pub fn seeded(lambda: T, n: Option<usize>, t_max: Option<T>, seed: u64) -> Self {
    validate_n_or_tmax(n, t_max, "Poisson");
    Poisson {
      lambda,
      n,
      t_max,
      seed: Deterministic::new(seed),
    }
  }
}

impl<T: FloatExt, S: SeedExt> Poisson<T, S> {
  /// Sample with an explicit seed, used by callers like CompoundPoisson.
  pub fn sample_with_seed(&self, seed: u64) -> Array1<T> {
    self.sample_impl(&Deterministic::new(seed))
  }

  /// Core sampling — monomorphised per seed strategy, zero runtime branching.
  pub(crate) fn sample_impl<S2: SeedExt>(&self, seed: &S2) -> Array1<T> {
    let distr = SimdExp::from_seed_source(self.lambda, seed);

    if let Some(n) = self.n {
      let mut poisson = Array1::<T>::zeros(n);
      if n <= 1 {
        return poisson;
      }

      let mut tail_view = poisson.slice_mut(s![1..]);
      let tail = tail_view
        .as_slice_mut()
        .expect("Poisson output tail must be contiguous");
      let mut rng = seed.rng();
      distr.fill_slice(&mut rng, tail);

      let mut acc = T::zero();
      for x in tail.iter_mut() {
        acc += *x;
        *x = acc;
      }
      poisson
    } else if let Some(t_max) = self.t_max {
      let expected = if t_max > T::zero() {
        (self.lambda * t_max).to_f64().unwrap_or(0.0)
      } else {
        0.0
      };
      let cap = if expected.is_finite() && expected > 0.0 {
        (expected.ceil() as usize).saturating_add(1)
      } else {
        1
      };
      let mut poisson = Vec::with_capacity(cap);
      poisson.push(T::zero());
      if t_max <= T::zero() {
        return Array1::from(poisson);
      }

      let mut t = T::zero();
      let mut rng = seed.rng();

      while t < t_max {
        t += distr.sample(&mut rng);

        if t < t_max {
          poisson.push(t);
        }
      }

      Array1::from(poisson)
    } else {
      unreachable!("validate_n_or_tmax ensures at least one of n, t_max is set")
    }
  }
}

impl<T: FloatExt, S: SeedExt> ProcessExt<T> for Poisson<T, S> {
  type Output = Array1<T>;

  fn sample(&self) -> Self::Output {
    self.sample_impl(&self.seed)
  }
}

py_process_1d!(PyPoisson, Poisson,
  sig: (lambda_, n=None, t_max=None, seed=None, dtype=None),
  params: (lambda_: f64, n: Option<usize>, t_max: Option<f64>)
);