Skip to main content

fdars_core/alignment/
nd.rs

1//! Multidimensional (R^d) SRSF transforms and elastic alignment.
2
3use super::srsf::reparameterize_curve;
4use super::{
5    dp_alignment_core, dp_edge_weight, dp_grid_solve, dp_lambda_penalty, dp_path_to_gamma,
6};
7use crate::helpers::{cumulative_trapz, l2_distance, simpsons_weights};
8use crate::matrix::{FdCurveSet, FdMatrix};
9
10/// Result of aligning multidimensional (R^d) curves.
11#[derive(Debug, Clone, PartialEq)]
12pub struct AlignmentResultNd {
13    /// Optimal warping function (length m), same for all dimensions.
14    pub gamma: Vec<f64>,
15    /// Aligned curve: d vectors, each length m.
16    pub f_aligned: Vec<Vec<f64>>,
17    /// Elastic distance after alignment.
18    pub distance: f64,
19}
20
21/// Scale derivative vector at one point by 1/√‖f'‖, writing into result_dims.
22#[inline]
23fn srsf_scale_point(derivs: &[FdMatrix], result_dims: &mut [FdMatrix], i: usize, j: usize) {
24    let d = derivs.len();
25    let norm_sq: f64 = derivs.iter().map(|dd| dd[(i, j)].powi(2)).sum();
26    let norm = norm_sq.sqrt();
27    if norm < 1e-15 {
28        for k in 0..d {
29            result_dims[k][(i, j)] = 0.0;
30        }
31    } else {
32        let scale = 1.0 / norm.sqrt();
33        for k in 0..d {
34            result_dims[k][(i, j)] = derivs[k][(i, j)] * scale;
35        }
36    }
37}
38
39/// Compute the SRSF transform for multidimensional (R^d) curves.
40///
41/// For f: \[0,1\] → R^d, the SRSF is q(t) = f'(t) / √‖f'(t)‖ where ‖·‖ is the
42/// Euclidean norm in R^d. For d=1 this reduces to `sign(f') · √|f'|`.
43///
44/// # Arguments
45/// * `data` — Set of n curves in R^d, each with m evaluation points
46/// * `argvals` — Evaluation points (length m)
47///
48/// # Returns
49/// `FdCurveSet` of SRSF values with the same shape as input.
50pub fn srsf_transform_nd(data: &FdCurveSet, argvals: &[f64]) -> FdCurveSet {
51    let d = data.ndim();
52    let n = data.ncurves();
53    let m = data.npoints();
54
55    if d == 0 || n == 0 || m == 0 || argvals.len() != m {
56        return FdCurveSet {
57            dims: (0..d).map(|_| FdMatrix::zeros(n, m)).collect(),
58        };
59    }
60
61    let derivs: Vec<FdMatrix> = data
62        .dims
63        .iter()
64        .map(|dim_mat| crate::fdata::deriv_1d(dim_mat, argvals, 1))
65        .collect();
66
67    let mut result_dims: Vec<FdMatrix> = (0..d).map(|_| FdMatrix::zeros(n, m)).collect();
68    for i in 0..n {
69        for j in 0..m {
70            srsf_scale_point(&derivs, &mut result_dims, i, j);
71        }
72    }
73
74    FdCurveSet { dims: result_dims }
75}
76
77/// Reconstruct an R^d curve from its SRSF.
78///
79/// Given d-dimensional SRSF vectors and initial point f0, reconstructs:
80/// `f_k(t) = f0_k + ∫₀ᵗ q_k(s) · ‖q(s)‖ ds` for each dimension k.
81///
82/// # Arguments
83/// * `q` — SRSF: d vectors, each length m
84/// * `argvals` — Evaluation points (length m)
85/// * `f0` — Initial values in R^d (length d)
86///
87/// # Returns
88/// Reconstructed curve: d vectors, each length m.
89pub fn srsf_inverse_nd(q: &[Vec<f64>], argvals: &[f64], f0: &[f64]) -> Vec<Vec<f64>> {
90    let d = q.len();
91    if d == 0 {
92        return Vec::new();
93    }
94    let m = q[0].len();
95    if m == 0 {
96        return vec![Vec::new(); d];
97    }
98
99    // Compute ||q(t)|| at each time point
100    let norms: Vec<f64> = (0..m)
101        .map(|j| {
102            let norm_sq: f64 = q.iter().map(|qk| qk[j].powi(2)).sum();
103            norm_sq.sqrt()
104        })
105        .collect();
106
107    // For each dimension, integrand = q_k(t) * ||q(t)||
108    let mut result = Vec::with_capacity(d);
109    for k in 0..d {
110        let integrand: Vec<f64> = (0..m).map(|j| q[k][j] * norms[j]).collect();
111        let integral = cumulative_trapz(&integrand, argvals);
112        let curve: Vec<f64> = integral.iter().map(|&v| f0[k] + v).collect();
113        result.push(curve);
114    }
115
116    result
117}
118
119/// Core DP alignment for R^d SRSFs.
120///
121/// Same DP grid and coprime neighborhood as `dp_alignment_core`, but edge weight
122/// is the sum of `dp_edge_weight` over d dimensions.
123fn dp_alignment_core_nd(
124    q1: &[Vec<f64>],
125    q2: &[Vec<f64>],
126    argvals: &[f64],
127    lambda: f64,
128) -> Vec<f64> {
129    let d = q1.len();
130    let m = argvals.len();
131    if m < 2 || d == 0 {
132        return argvals.to_vec();
133    }
134
135    // For d=1, delegate to existing implementation for exact backward compat
136    if d == 1 {
137        return dp_alignment_core(&q1[0], &q2[0], argvals, lambda);
138    }
139
140    // Normalize each dimension's SRSF to unit L2 norm
141    let q1n: Vec<Vec<f64>> = q1
142        .iter()
143        .map(|qk| {
144            let norm = qk.iter().map(|&v| v * v).sum::<f64>().sqrt().max(1e-10);
145            qk.iter().map(|&v| v / norm).collect()
146        })
147        .collect();
148    let q2n: Vec<Vec<f64>> = q2
149        .iter()
150        .map(|qk| {
151            let norm = qk.iter().map(|&v| v * v).sum::<f64>().sqrt().max(1e-10);
152            qk.iter().map(|&v| v / norm).collect()
153        })
154        .collect();
155
156    let path = dp_grid_solve(m, m, |sr, sc, tr, tc| {
157        let w: f64 = (0..d)
158            .map(|k| dp_edge_weight(&q1n[k], &q2n[k], argvals, sc, tc, sr, tr))
159            .sum();
160        w + dp_lambda_penalty(argvals, sc, tc, sr, tr, lambda)
161    });
162
163    dp_path_to_gamma(&path, argvals)
164}
165
166/// Align an R^d curve f2 to f1 using the elastic framework.
167///
168/// Finds the optimal warping γ (shared across all dimensions) such that
169/// f2∘γ is as close as possible to f1 in the elastic metric.
170///
171/// # Arguments
172/// * `f1` — Target curves (d dimensions)
173/// * `f2` — Curves to align (d dimensions)
174/// * `argvals` — Evaluation points (length m)
175/// * `lambda` — Penalty weight (0.0 = no penalty)
176pub fn elastic_align_pair_nd(
177    f1: &FdCurveSet,
178    f2: &FdCurveSet,
179    argvals: &[f64],
180    lambda: f64,
181) -> AlignmentResultNd {
182    let d = f1.ndim();
183    let m = f1.npoints();
184
185    // Compute SRSFs
186    let q1_set = srsf_transform_nd(f1, argvals);
187    let q2_set = srsf_transform_nd(f2, argvals);
188
189    // Extract first curve from each dimension
190    let q1: Vec<Vec<f64>> = q1_set.dims.iter().map(|dm| dm.row(0)).collect();
191    let q2: Vec<Vec<f64>> = q2_set.dims.iter().map(|dm| dm.row(0)).collect();
192
193    // DP alignment using summed cost over dimensions
194    let gamma = dp_alignment_core_nd(&q1, &q2, argvals, lambda);
195
196    // Apply warping to f2 in each dimension
197    let f_aligned: Vec<Vec<f64>> = f2
198        .dims
199        .iter()
200        .map(|dm| {
201            let row = dm.row(0);
202            reparameterize_curve(&row, argvals, &gamma)
203        })
204        .collect();
205
206    // Compute elastic distance: sum of squared L2 distances between aligned SRSFs
207    let f_aligned_set = {
208        let dims: Vec<FdMatrix> = f_aligned
209            .iter()
210            .map(|fa| {
211                FdMatrix::from_slice(fa, 1, m).expect("dimension invariant: data.len() == n * m")
212            })
213            .collect();
214        FdCurveSet { dims }
215    };
216    let q_aligned = srsf_transform_nd(&f_aligned_set, argvals);
217    let weights = simpsons_weights(argvals);
218
219    let mut dist_sq = 0.0;
220    for k in 0..d {
221        let q1k = q1_set.dims[k].row(0);
222        let qak = q_aligned.dims[k].row(0);
223        let d_k = l2_distance(&q1k, &qak, &weights);
224        dist_sq += d_k * d_k;
225    }
226
227    AlignmentResultNd {
228        gamma,
229        f_aligned,
230        distance: dist_sq.sqrt(),
231    }
232}
233
234/// Elastic distance between two R^d curves.
235///
236/// Aligns f2 to f1 and returns the post-alignment SRSF distance.
237pub fn elastic_distance_nd(f1: &FdCurveSet, f2: &FdCurveSet, argvals: &[f64], lambda: f64) -> f64 {
238    elastic_align_pair_nd(f1, f2, argvals, lambda).distance
239}