1use itk_transforms_rs::TransformChain;
19
20use crate::tractogram::Tractogram;
21
22pub fn apply_transform_in_place(tractogram: &mut Tractogram, chain: &TransformChain) {
27 let positions = tractogram.positions_mut();
28 map_positions(positions, chain);
29}
30
31pub 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 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}