diffusionx 0.12.0

A multi-threaded crate for random number generation and stochastic process simulation, with optional GPU acceleration.
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

use crate::{FloatExt, XResult, random::normal};
use rand_distr::{Distribution, StandardNormal};
use rayon::prelude::*;
use realfft::{FftNum, RealFftPlanner, num_complex::Complex};

/// Circulant embedding method for generating stationary Gaussian random fields with given correlation functions
pub struct CirculantEmbedding<F: Fn(usize) -> T, T: FloatExt + FftNum = f64> {
    /// Number of grid points
    size: usize,
    /// Correlation function, takes distance as input and returns correlation
    correlation_fn: F,
    /// embedding size
    embedding_size: usize,
    /// sqrt eigenvalues
    sqrt_eigenvalues: Option<Vec<T>>,
}

impl<F: Fn(usize) -> T + Send + Sync, T: FloatExt + FftNum> CirculantEmbedding<F, T> {
    /// Create a new `CirculantEmbedding`
    ///
    /// # Arguments
    ///
    /// * `size` - Number of grid points per dimension
    /// * `correlation_fn` - Correlation function, takes distance as input and returns correlation
    ///
    /// # Example
    ///
    /// ```rust
    /// use diffusionx::utils::CirculantEmbedding;
    ///
    /// let embedding = CirculantEmbedding::new(100, |r| (-r / 10.0).exp());
    /// ```
    pub fn new(size: usize, correlation_fn: F) -> Self {
        CirculantEmbedding {
            size,
            correlation_fn,
            embedding_size: next_power_of_2(2 * size),
            sqrt_eigenvalues: None,
        }
    }

    fn embed(&mut self) -> XResult<()> {
        let mut first_row = (0..self.embedding_size)
            .into_par_iter()
            .map(|i| {
                let dist = if i < self.size {
                    i
                } else {
                    self.embedding_size - i
                };
                (self.correlation_fn)(dist)
            })
            .collect::<Vec<_>>();

        let mut planner = RealFftPlanner::new();
        let fft = planner.plan_fft_forward(self.embedding_size);
        let mut eigenvalues = fft.make_output_vec();

        fft.process(&mut first_row, &mut eigenvalues)?;

        if eigenvalues.iter().any(|val| val.re < T::zero()) {
            self.embedding_size *= 2;
            self.embed()
        } else {
            self.sqrt_eigenvalues =
                Some(eigenvalues.into_iter().map(|val| val.re.sqrt()).collect());
            Ok(())
        }
    }

    /// Generate a one-dimensional stationary Gaussian random field
    pub fn generate(&mut self) -> XResult<Vec<T>>
    where
        StandardNormal: Distribution<T>,
    {
        if self.sqrt_eigenvalues.is_none() {
            self.embed()?;
        }

        let mut planner = RealFftPlanner::new();
        let ifft = planner.plan_fft_inverse(self.embedding_size);

        let mut modified_z = self
            .sqrt_eigenvalues
            .as_ref()
            .unwrap()
            .par_iter()
            .map(|&sqrt_lambda| {
                let re = sqrt_lambda * normal::standard_rand();
                let im = sqrt_lambda * normal::standard_rand();
                Complex::new(re, im)
            })
            .collect::<Vec<_>>();

        let mut result = ifft.make_output_vec();
        ifft.process(&mut modified_z, &mut result)?;
        let scale = T::one() / T::from(self.embedding_size).unwrap();
        result.iter_mut().for_each(|x| *x *= scale);
        result.truncate(self.size);

        Ok(result)
    }
}

fn next_power_of_2(mut n: usize) -> usize {
    if n == 0 {
        return 1;
    }
    n -= 1;
    n |= n >> 1;
    n |= n >> 2;
    n |= n >> 4;
    n |= n >> 8;
    n |= n >> 16;
    if size_of::<usize>() == 8 {
        n |= n >> 32;
    }
    n + 1
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_next_power_of_2() {
        assert_eq!(next_power_of_2(8), 8);
        assert_eq!(next_power_of_2(11), 16);
    }
}