Skip to main content

trx_rs/
transform.rs

1//! Apply spatial transforms to streamline coordinates.
2//!
3//! Streamlines are stored as a single contiguous `Vec<[f32; 3]>` of
4//! vertices in RAS+ mm. Spatially transforming them is a pointwise
5//! operation: for each vertex `p`, compute `chain.map_point(p)` and
6//! overwrite. There is no notion of "pull"-style resampling — streamlines
7//! aren't a sampled image, they're a list of mm coordinates — so the
8//! transform's chain must map *source-coords → target-coords* (i.e., the
9//! direction the points are moving). For ANTs paired h5s in BIDS naming,
10//! that means the *inverse-named* file: e.g. to warp streamlines
11//! ACPC → MNI, pass `from-MNI_to-ACPC.h5`.
12//!
13//! Use [`apply_transform_in_place`] to mutate a [`Tractogram`] in situ,
14//! or [`apply_transform`] for the consuming variant. With the `parallel`
15//! crate feature enabled (default), the per-point loop is rayon-parallel
16//! across streamline vertices.
17
18use itk_transforms_rs::TransformChain;
19
20use crate::tractogram::Tractogram;
21
22/// Apply `chain` to every vertex of `tractogram` in place.
23///
24/// Streamline boundaries, DPS, DPV, DPG, and groups are preserved
25/// unchanged — only `positions` is mutated.
26pub fn apply_transform_in_place(tractogram: &mut Tractogram, chain: &TransformChain) {
27    let positions = tractogram.positions_mut();
28    map_positions(positions, chain);
29}
30
31/// Consuming variant of [`apply_transform_in_place`]. Returns the
32/// transformed [`Tractogram`].
33pub fn apply_transform(mut tractogram: Tractogram, chain: &TransformChain) -> Tractogram {
34    apply_transform_in_place(&mut tractogram, chain);
35    tractogram
36}
37
38#[inline]
39fn map_one(p: &mut [f32; 3], chain: &TransformChain) {
40    let q = chain.map_point([p[0] as f64, p[1] as f64, p[2] as f64]);
41    *p = [q[0] as f32, q[1] as f32, q[2] as f32];
42}
43
44#[cfg(feature = "parallel")]
45fn map_positions(positions: &mut [[f32; 3]], chain: &TransformChain) {
46    use rayon::prelude::*;
47    positions.par_iter_mut().for_each(|p| map_one(p, chain));
48}
49
50#[cfg(not(feature = "parallel"))]
51fn map_positions(positions: &mut [[f32; 3]], chain: &TransformChain) {
52    for p in positions.iter_mut() {
53        map_one(p, chain);
54    }
55}
56
57#[cfg(test)]
58mod tests {
59    use super::*;
60    use itk_transforms_rs::{Affine3, TransformChain};
61    use nalgebra::Matrix4;
62
63    use crate::header::Header;
64    use crate::tractogram::Tractogram;
65
66    fn one_streamline(points: Vec<[f32; 3]>) -> Tractogram {
67        let mut t = Tractogram::with_header(Header {
68            voxel_to_rasmm: Header::identity_affine(),
69            dimensions: [10, 10, 10],
70            nb_streamlines: 0,
71            nb_vertices: 0,
72            extra: Default::default(),
73        });
74        t.push_streamline(&points).unwrap();
75        t
76    }
77
78    #[test]
79    fn identity_chain_is_noop() {
80        let pts = vec![[1.0_f32, 2.0, 3.0], [4.0, 5.0, 6.0]];
81        let mut t = one_streamline(pts.clone());
82        let mut chain = TransformChain::new();
83        chain.push_affine(Affine3::identity());
84        apply_transform_in_place(&mut t, &chain);
85        for (a, b) in t.positions().iter().zip(pts.iter()) {
86            assert_eq!(a, b);
87        }
88    }
89
90    #[test]
91    fn translation_shifts_every_point() {
92        let pts = vec![[0.0_f32, 0.0, 0.0], [1.0, 2.0, 3.0]];
93        let mut t = one_streamline(pts);
94
95        let mut m = Matrix4::identity();
96        m[(0, 3)] = 10.0;
97        m[(1, 3)] = -5.0;
98        m[(2, 3)] = 100.0;
99        let mut chain = TransformChain::new();
100        chain.push_affine(Affine3::from_matrix(m));
101
102        apply_transform_in_place(&mut t, &chain);
103        let p0 = t.positions()[0];
104        let p1 = t.positions()[1];
105        assert!((p0[0] - 10.0).abs() < 1e-6);
106        assert!((p0[1] + 5.0).abs() < 1e-6);
107        assert!((p0[2] - 100.0).abs() < 1e-6);
108        assert!((p1[0] - 11.0).abs() < 1e-6);
109        assert!((p1[1] + 3.0).abs() < 1e-6);
110        assert!((p1[2] - 103.0).abs() < 1e-6);
111    }
112
113    #[test]
114    fn rotation_z_90_maps_x_to_y() {
115        // Rz(90°) sends +x → +y, +y → -x. Matrix in RAS+:
116        let theta = std::f64::consts::FRAC_PI_2;
117        let (s, c) = (theta.sin(), theta.cos());
118        let mut m = Matrix4::identity();
119        m[(0, 0)] = c;
120        m[(0, 1)] = -s;
121        m[(1, 0)] = s;
122        m[(1, 1)] = c;
123        let mut chain = TransformChain::new();
124        chain.push_affine(Affine3::from_matrix(m));
125
126        let pts = vec![[1.0_f32, 0.0, 0.0], [0.0, 1.0, 0.0]];
127        let mut t = one_streamline(pts);
128        apply_transform_in_place(&mut t, &chain);
129
130        let p0 = t.positions()[0];
131        let p1 = t.positions()[1];
132        assert!(p0[0].abs() < 1e-6 && (p0[1] - 1.0).abs() < 1e-6 && p0[2].abs() < 1e-6);
133        assert!((p1[0] + 1.0).abs() < 1e-6 && p1[1].abs() < 1e-6 && p1[2].abs() < 1e-6);
134    }
135
136    #[test]
137    fn streamline_count_and_offsets_unchanged() {
138        let pts = vec![[1.0_f32, 0.0, 0.0]; 17];
139        let mut t = one_streamline(pts);
140        let n_before = t.positions().len();
141        let offsets_before = t.offsets().to_vec();
142        let mut chain = TransformChain::new();
143        chain.push_affine(Affine3::identity());
144        apply_transform_in_place(&mut t, &chain);
145        assert_eq!(t.positions().len(), n_before);
146        assert_eq!(t.offsets(), offsets_before.as_slice());
147    }
148}