pathwise_geo/
simulate.rs

1use crate::sde::ManifoldSDE;
2use crate::scheme::euler::GeodesicEuler;
3use crate::scheme::milstein::GeodesicMilstein;
4use cartan_core::{Manifold, ParallelTransport};
5use pathwise_core::state::{Increment, NoiseIncrement};
6use rand::SeedableRng;
7
8// Same derivation as pathwise_core::rng::splitmix64 — must stay in sync.
9// pathwise-geo cannot take a dependency on pathwise-core::rng (private module),
10// so this is an intentional local copy.
11#[inline]
12fn splitmix64(mut x: u64) -> u64 {
13    x = x.wrapping_add(0x9e3779b97f4a7c15);
14    x = (x ^ (x >> 30)).wrapping_mul(0xbf58476d1ce4e5b9);
15    x = (x ^ (x >> 27)).wrapping_mul(0x94d049bb133111eb);
16    x ^ (x >> 31)
17}
18
19/// Simulate manifold SDE paths using GeodesicEuler.
20///
21/// Returns `Vec<Vec<M::Point>>` with outer index = path, inner index = time step.
22/// Each inner vec has `n_steps + 1` points (including the initial condition x0).
23#[allow(clippy::too_many_arguments)]
24pub fn manifold_simulate<M, D, G>(
25    sde: &ManifoldSDE<M, D, G>,
26    scheme: &GeodesicEuler,
27    x0: M::Point,
28    t0: f64,
29    t1: f64,
30    n_paths: usize,
31    n_steps: usize,
32    seed: u64,
33) -> Vec<Vec<M::Point>>
34where
35    M: Manifold + Clone + Send + Sync,
36    M::Point: Clone + Send + Sync,
37    M::Tangent: Clone + Send + Sync,
38    D: Fn(&M::Point, f64) -> M::Tangent + Send + Sync,
39    G: Fn(&M::Point, f64) -> M::Tangent + Send + Sync,
40{
41    let dt = (t1 - t0) / n_steps as f64;
42    let base_seed = splitmix64(seed);
43    (0..n_paths)
44        .map(|i| {
45            let path_seed = splitmix64(base_seed.wrapping_add(i as u64));
46            let mut rng = rand::rngs::SmallRng::seed_from_u64(path_seed);
47            let mut path = Vec::with_capacity(n_steps + 1);
48            let mut x = x0.clone();
49            path.push(x.clone());
50            for step in 0..n_steps {
51                let t = t0 + step as f64 * dt;
52                let inc = <f64 as NoiseIncrement>::sample(&mut rng, dt);
53                x = scheme.step(sde, &x, t, dt, &inc);
54                path.push(x.clone());
55            }
56            path
57        })
58        .collect()
59}
60
61/// Internal trait to unify geodesic scheme step dispatch.
62///
63/// Implementors provide a single `step_geo` method that advances x by one step.
64/// The M: ParallelTransport bound is required for the GeodesicMilstein impl;
65/// GeodesicEuler satisfies it but does not use the transport operation.
66pub trait GeoScheme<M, D, G>
67where
68    M: Manifold + ParallelTransport,
69    D: Fn(&M::Point, f64) -> M::Tangent,
70    G: Fn(&M::Point, f64) -> M::Tangent,
71{
72    /// Advance x by one scheme step.
73    fn step_geo(
74        &self,
75        sde: &ManifoldSDE<M, D, G>,
76        x: &M::Point,
77        t: f64,
78        dt: f64,
79        inc: &Increment<f64>,
80    ) -> M::Point;
81}
82
83impl<M, D, G> GeoScheme<M, D, G> for GeodesicEuler
84where
85    M: Manifold + ParallelTransport,
86    D: Fn(&M::Point, f64) -> M::Tangent + Send + Sync,
87    G: Fn(&M::Point, f64) -> M::Tangent + Send + Sync,
88{
89    fn step_geo(
90        &self,
91        sde: &ManifoldSDE<M, D, G>,
92        x: &M::Point,
93        t: f64,
94        dt: f64,
95        inc: &Increment<f64>,
96    ) -> M::Point {
97        self.step(sde, x, t, dt, inc)
98    }
99}
100
101impl<M, D, G> GeoScheme<M, D, G> for GeodesicMilstein
102where
103    M: Manifold + ParallelTransport,
104    D: Fn(&M::Point, f64) -> M::Tangent + Send + Sync,
105    G: Fn(&M::Point, f64) -> M::Tangent + Send + Sync,
106{
107    fn step_geo(
108        &self,
109        sde: &ManifoldSDE<M, D, G>,
110        x: &M::Point,
111        t: f64,
112        dt: f64,
113        inc: &Increment<f64>,
114    ) -> M::Point {
115        self.step(sde, x, t, dt, inc)
116    }
117}
118
119/// Simulate manifold SDE paths using any `GeoScheme` implementor.
120///
121/// Generic version of `manifold_simulate` that works with GeodesicEuler,
122/// GeodesicMilstein, or any future scheme implementing `GeoScheme`.
123/// Returns `Vec<Vec<M::Point>>` with outer index = path, inner index = time step.
124/// Each inner vec has `n_steps + 1` points (including the initial condition x0).
125#[allow(clippy::too_many_arguments)]
126pub fn manifold_simulate_with_scheme<M, D, G, SC>(
127    sde: &ManifoldSDE<M, D, G>,
128    scheme: &SC,
129    x0: M::Point,
130    t0: f64,
131    t1: f64,
132    n_paths: usize,
133    n_steps: usize,
134    seed: u64,
135) -> Vec<Vec<M::Point>>
136where
137    M: Manifold + ParallelTransport + Clone + Send + Sync,
138    M::Point: Clone + Send + Sync,
139    M::Tangent: Clone + Send + Sync,
140    D: Fn(&M::Point, f64) -> M::Tangent + Send + Sync,
141    G: Fn(&M::Point, f64) -> M::Tangent + Send + Sync,
142    SC: GeoScheme<M, D, G>,
143{
144    let dt = (t1 - t0) / n_steps as f64;
145    let base_seed = splitmix64(seed);
146    (0..n_paths)
147        .map(|i| {
148            let path_seed = splitmix64(base_seed.wrapping_add(i as u64));
149            let mut rng = rand::rngs::SmallRng::seed_from_u64(path_seed);
150            let mut path = Vec::with_capacity(n_steps + 1);
151            let mut x = x0.clone();
152            path.push(x.clone());
153            for step in 0..n_steps {
154                let t = t0 + step as f64 * dt;
155                let inc = <f64 as NoiseIncrement>::sample(&mut rng, dt);
156                x = scheme.step_geo(sde, &x, t, dt, &inc);
157                path.push(x.clone());
158            }
159            path
160        })
161        .collect()
162}
163
164/// Flatten manifold SDE paths into a 3-D array by applying `log` from a reference point.
165///
166/// Shape: `(n_paths, n_steps+1, tangent_dim)`.
167/// Requires `M::Tangent: AsRef<[f64]>`, which holds for `nalgebra::SVector<f64, N>`.
168/// If `log` fails for a point (cut locus), that entry is filled with NaN.
169pub fn paths_to_array<M>(
170    paths: &[Vec<M::Point>],
171    manifold: &M,
172    ref_point: &M::Point,
173) -> ndarray::Array3<f64>
174where
175    M: Manifold,
176    M::Point: Clone,
177    M::Tangent: AsRef<[f64]>,
178{
179    let n_paths = paths.len();
180    if n_paths == 0 {
181        return ndarray::Array3::zeros((0, 0, 0));
182    }
183    let n_steps_plus1 = paths[0].len();
184    // Determine tangent dimension from the first valid log.
185    let dim = manifold
186        .log(ref_point, &paths[0][0])
187        .map(|v| v.as_ref().len())
188        .unwrap_or(0);
189    if dim == 0 {
190        return ndarray::Array3::zeros((n_paths, n_steps_plus1, 0));
191    }
192    let mut out = ndarray::Array3::from_elem((n_paths, n_steps_plus1, dim), f64::NAN);
193    for (i, path) in paths.iter().enumerate() {
194        for (j, point) in path.iter().enumerate() {
195            if let Ok(tangent) = manifold.log(ref_point, point) {
196                let slice = tangent.as_ref();
197                for (k, &v) in slice.iter().enumerate().take(dim) {
198                    out[[i, j, k]] = v;
199                }
200            }
201        }
202    }
203    out
204}
pathwise_geo/simulate.rs

pathwise_geo/
simulate.rs
