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
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151

use crate::{
    FloatExt, XResult,
    gpu::{
        GPUMoment,
        metal::{MetalLibrary, OU_METALLIB, load_library},
    },
    simulation::continuous::OrnsteinUhlenbeck as OU,
};
use objc2::rc::Retained;
use std::sync::LazyLock;

static LIBRARY: LazyLock<XResult<Retained<MetalLibrary>>> =
    LazyLock::new(|| load_library(OU_METALLIB));

subscribe_metal_gpu_function!(LIBRARY, mean, "mean", (start_position: f32, theta: f32, sigma: f32, duration: f32, time_step: f32));

subscribe_metal_gpu_function!(LIBRARY, msd, "msd", (start_position: f32, theta: f32, sigma: f32, duration: f32, time_step: f32));

subscribe_metal_gpu_function!(LIBRARY, raw_moment, "raw_moment", (start_position: f32, theta: f32, sigma: f32, order: i32, duration: f32, time_step: f32));

subscribe_metal_gpu_function!(LIBRARY, frac_raw_moment, "frac_raw_moment", (start_position: f32, theta: f32, sigma: f32, order: f32, duration: f32, time_step: f32));

subscribe_metal_central_moment_gpu_function!(LIBRARY, central_moment, "central_moment", (start_position: f32, theta: f32, sigma: f32, duration: f32, time_step: f32), i32);

subscribe_metal_frac_central_moment_gpu_function!(
    LIBRARY,
    frac_central_moment,
    "frac_central_moment",
    (
        start_position: f32,
        theta: f32,
        sigma: f32,
        duration: f32,
        time_step: f32
    ),
    before_order = (),
    after_order = (start_position, theta, sigma, duration, time_step)
);

impl<T: FloatExt> GPUMoment for OU<T> {
    fn central_moment_gpu(
        &self,
        duration: f32,
        order: i32,
        particles: usize,
        time_step: f32,
    ) -> XResult<f32> {
        central_moment(
            self.get_start_position().to_f32().unwrap(),
            self.get_theta().to_f32().unwrap(),
            self.get_sigma().to_f32().unwrap(),
            duration,
            time_step,
            order,
            particles,
        )
    }

    fn raw_moment_gpu(
        &self,
        duration: f32,
        order: i32,
        particles: usize,
        time_step: f32,
    ) -> XResult<f32> {
        raw_moment(
            self.get_start_position().to_f32().unwrap(),
            self.get_theta().to_f32().unwrap(),
            self.get_sigma().to_f32().unwrap(),
            order,
            duration,
            time_step,
            particles,
        )
    }

    fn mean_gpu(&self, duration: f32, particles: usize, time_step: f32) -> XResult<f32> {
        mean(
            self.get_start_position().to_f32().unwrap(),
            self.get_theta().to_f32().unwrap(),
            self.get_sigma().to_f32().unwrap(),
            duration,
            time_step,
            particles,
        )
    }

    fn msd_gpu(&self, duration: f32, particles: usize, time_step: f32) -> XResult<f32> {
        msd(
            self.get_start_position().to_f32().unwrap(),
            self.get_theta().to_f32().unwrap(),
            self.get_sigma().to_f32().unwrap(),
            duration,
            time_step,
            particles,
        )
    }

    fn frac_central_moment_gpu(
        &self,
        duration: f32,
        order: f32,
        particles: usize,
        time_step: f32,
    ) -> XResult<f32> {
        frac_central_moment(
            self.get_start_position().to_f32().unwrap(),
            self.get_theta().to_f32().unwrap(),
            self.get_sigma().to_f32().unwrap(),
            duration,
            time_step,
            order,
            particles,
        )
    }

    fn frac_raw_moment_gpu(
        &self,
        duration: f32,
        order: f32,
        particles: usize,
        time_step: f32,
    ) -> XResult<f32> {
        frac_raw_moment(
            self.get_start_position().to_f32().unwrap(),
            self.get_theta().to_f32().unwrap(),
            self.get_sigma().to_f32().unwrap(),
            order,
            duration,
            time_step,
            particles,
        )
    }
}

#[cfg(test)]
mod tests {
    use crate::gpu::GPUMoment;
    use crate::simulation::continuous::OrnsteinUhlenbeck;

    #[test]
    fn test_gpu_moment() {
        let ou = OrnsteinUhlenbeck::<f32>::default();
        ou.mean_gpu(1.0, 100, 0.1).unwrap();
        ou.msd_gpu(1.0, 100, 0.1).unwrap();
        ou.raw_moment_gpu(1.0, 2, 100, 0.1).unwrap();
        ou.frac_raw_moment_gpu(1.0, 1.4, 100, 0.1).unwrap();
        ou.central_moment_gpu(1.0, 2, 100, 0.1).unwrap();
        ou.frac_central_moment_gpu(1.0, 1.5, 100, 0.1).unwrap();
    }
}