trx_rs/
tractogram.rs

1use half::f16;
2use std::collections::HashMap;
3
4use crate::any_trx_file::{AnyTrxFile, PositionsRef};
5use crate::dtype::{DType, TrxScalar};
6use crate::error::{Result, TrxError};
7use crate::header::Header;
8use crate::mmap_backing::{vec_to_bytes, MmapBacking};
9use crate::trx_file::{DataArray, DataPerGroup, TrxFile, TrxParts};
10
11/// Neutral in-memory streamline representation used for cross-format conversion.
12///
13/// Positions are stored in world/RASMM-style coordinates as `f32` triplets.
14/// `dps` (data-per-streamline) and `dpv` (data-per-vertex) ride along
15/// unchanged through any operation that doesn't reorder streamlines or
16/// vertices — including [`crate::transform::apply_transform_in_place`],
17/// which only mutates the `positions` buffer.
18#[derive(Debug)]
19pub struct Tractogram {
20    header: Header,
21    positions: Vec<[f32; 3]>,
22    offsets: Vec<u32>,
23    dps: HashMap<String, DataArray>,
24    dpv: HashMap<String, DataArray>,
25    groups: HashMap<String, Vec<u32>>,
26    dpg: DataPerGroup,
27}
28
29impl Clone for Tractogram {
30    fn clone(&self) -> Self {
31        Self {
32            header: self.header.clone(),
33            positions: self.positions.clone(),
34            offsets: self.offsets.clone(),
35            dps: clone_arrays(&self.dps),
36            dpv: clone_arrays(&self.dpv),
37            groups: self.groups.clone(),
38            dpg: clone_dpg(&self.dpg),
39        }
40    }
41}
42
43impl Tractogram {
44    /// Create an empty tractogram with identity metadata.
45    pub fn new() -> Self {
46        Self::with_header(Header {
47            voxel_to_rasmm: Header::identity_affine(),
48            dimensions: [1, 1, 1],
49            nb_streamlines: 0,
50            nb_vertices: 0,
51            extra: Default::default(),
52        })
53    }
54
55    /// Create an empty tractogram with a caller-provided header.
56    pub fn with_header(mut header: Header) -> Self {
57        header.nb_streamlines = 0;
58        header.nb_vertices = 0;
59        Self {
60            header,
61            positions: Vec::new(),
62            offsets: vec![0],
63            dps: HashMap::new(),
64            dpv: HashMap::new(),
65            groups: HashMap::new(),
66            dpg: HashMap::new(),
67        }
68    }
69
70    /// Build a tractogram by copying streamline geometry from a typed TRX file.
71    pub fn from_trx<P: TrxScalar>(trx: &TrxFile<P>) -> Self {
72        let positions = trx
73            .positions()
74            .iter()
75            .map(|point| [point[0].to_f32(), point[1].to_f32(), point[2].to_f32()])
76            .collect();
77        Self {
78            header: trx.header().clone(),
79            positions,
80            offsets: trx.offsets().to_vec(),
81            dps: clone_arrays(trx.dps_arrays()),
82            dpv: clone_arrays(trx.dpv_arrays()),
83            groups: clone_groups(trx.group_arrays()),
84            dpg: clone_dpg(trx.dpg_arrays()),
85        }
86    }
87
88    /// Build a tractogram by copying streamline geometry from a dtype-erased TRX file.
89    pub fn from_any_trx(trx: &AnyTrxFile) -> Self {
90        let positions = match trx.positions_ref() {
91            PositionsRef::F16(data) => data
92                .iter()
93                .map(|point| [point[0].to_f32(), point[1].to_f32(), point[2].to_f32()])
94                .collect(),
95            PositionsRef::F32(data) => data.to_vec(),
96            PositionsRef::F64(data) => data
97                .iter()
98                .map(|point| [point[0] as f32, point[1] as f32, point[2] as f32])
99                .collect(),
100        };
101
102        Self {
103            header: trx.header().clone(),
104            positions,
105            offsets: trx.offsets_vec(),
106            dps: trx.with_typed(
107                |inner| clone_arrays(inner.dps_arrays()),
108                |inner| clone_arrays(inner.dps_arrays()),
109                |inner| clone_arrays(inner.dps_arrays()),
110            ),
111            dpv: trx.with_typed(
112                |inner| clone_arrays(inner.dpv_arrays()),
113                |inner| clone_arrays(inner.dpv_arrays()),
114                |inner| clone_arrays(inner.dpv_arrays()),
115            ),
116            groups: trx.groups_owned().into_iter().collect(),
117            dpg: trx.with_typed(
118                |inner| clone_dpg(inner.dpg_arrays()),
119                |inner| clone_dpg(inner.dpg_arrays()),
120                |inner| clone_dpg(inner.dpg_arrays()),
121            ),
122        }
123    }
124
125    pub fn header(&self) -> &Header {
126        &self.header
127    }
128
129    pub fn set_spatial_metadata(&mut self, voxel_to_rasmm: [[f64; 4]; 4], dimensions: [u64; 3]) {
130        self.header.voxel_to_rasmm = voxel_to_rasmm;
131        self.header.dimensions = dimensions;
132    }
133
134    pub fn set_header(&mut self, header: Header) {
135        self.header = header;
136    }
137
138    pub fn extra(&self) -> &HashMap<String, serde_json::Value> {
139        &self.header.extra
140    }
141
142    pub fn extra_mut(&mut self) -> &mut HashMap<String, serde_json::Value> {
143        &mut self.header.extra
144    }
145
146    pub fn positions(&self) -> &[[f32; 3]] {
147        &self.positions
148    }
149
150    /// Mutable view of the contiguous `positions` buffer.
151    ///
152    /// Useful for in-place spatial transforms (see [`crate::transform`]).
153    /// Streamline boundaries (`offsets`) are unaffected by point edits, so
154    /// callers may mutate the slice freely without invalidating any
155    /// metadata.
156    pub fn positions_mut(&mut self) -> &mut [[f32; 3]] {
157        &mut self.positions
158    }
159
160    pub fn offsets(&self) -> &[u32] {
161        &self.offsets
162    }
163
164    pub fn group_names(&self) -> impl Iterator<Item = &str> {
165        self.groups.keys().map(String::as_str)
166    }
167
168    pub fn group(&self, name: &str) -> Option<&[u32]> {
169        self.groups.get(name).map(Vec::as_slice)
170    }
171
172    pub fn groups(&self) -> &HashMap<String, Vec<u32>> {
173        &self.groups
174    }
175
176    pub fn insert_group(&mut self, name: impl Into<String>, members: Vec<u32>) {
177        self.groups.insert(name.into(), members);
178    }
179
180    pub fn insert_dpg(
181        &mut self,
182        group: impl Into<String>,
183        name: impl Into<String>,
184        data: DataArray,
185    ) {
186        self.dpg
187            .entry(group.into())
188            .or_default()
189            .insert(name.into(), data);
190    }
191
192    pub fn dpg(&self) -> &DataPerGroup {
193        &self.dpg
194    }
195
196    pub fn dps_arrays(&self) -> &HashMap<String, DataArray> {
197        &self.dps
198    }
199
200    pub fn dpv_arrays(&self) -> &HashMap<String, DataArray> {
201        &self.dpv
202    }
203
204    pub fn dps_names(&self) -> impl Iterator<Item = &str> {
205        self.dps.keys().map(String::as_str)
206    }
207
208    pub fn dpv_names(&self) -> impl Iterator<Item = &str> {
209        self.dpv.keys().map(String::as_str)
210    }
211
212    /// Insert (or replace) a DPS field. The data array's row count must
213    /// equal the current streamline count.
214    pub fn insert_dps(&mut self, name: impl Into<String>, data: DataArray) {
215        self.dps.insert(name.into(), data);
216    }
217
218    /// Insert (or replace) a DPV field. The data array's row count must
219    /// equal the current vertex count.
220    pub fn insert_dpv(&mut self, name: impl Into<String>, data: DataArray) {
221        self.dpv.insert(name.into(), data);
222    }
223
224    pub fn subset_streamlines(&self, indices: &[usize]) -> Result<Self> {
225        let mut tractogram = Tractogram::with_header(self.header.clone());
226        let mut remap = HashMap::with_capacity(indices.len());
227
228        for (new_index, &old_index) in indices.iter().enumerate() {
229            let window = self.offsets.get(old_index..=old_index + 1).ok_or_else(|| {
230                TrxError::Argument(format!("streamline index {old_index} out of bounds"))
231            })?;
232            tractogram.push_streamline(&self.positions[window[0] as usize..window[1] as usize])?;
233            remap.insert(old_index as u32, new_index as u32);
234        }
235
236        for (name, members) in &self.groups {
237            let remapped = members
238                .iter()
239                .filter_map(|member| remap.get(member).copied())
240                .collect::<Vec<_>>();
241            if remapped.is_empty() {
242                continue;
243            }
244            tractogram.insert_group(name.clone(), remapped);
245            if let Some(entries) = self.dpg.get(name) {
246                for (entry_name, array) in entries {
247                    tractogram.insert_dpg(name.clone(), entry_name.clone(), array.clone_owned());
248                }
249            }
250        }
251
252        Ok(tractogram)
253    }
254
255    /// Add a streamline to the tractogram.
256    pub fn push_streamline(&mut self, points: &[[f32; 3]]) -> Result<()> {
257        self.positions.extend_from_slice(points);
258        let next_offset = u32::try_from(self.positions.len())
259            .map_err(|_| TrxError::Argument("tractogram has more than u32::MAX vertices".into()))?;
260        self.offsets.push(next_offset);
261        self.header.nb_streamlines += 1;
262        self.header.nb_vertices = self.positions.len() as u64;
263        Ok(())
264    }
265
266    /// Number of streamlines.
267    pub fn nb_streamlines(&self) -> usize {
268        self.offsets.len().saturating_sub(1)
269    }
270
271    /// Number of vertices.
272    pub fn nb_vertices(&self) -> usize {
273        self.positions.len()
274    }
275
276    /// Borrow a single streamline as a position slice.
277    pub fn streamline(&self, index: usize) -> &[[f32; 3]] {
278        let start = self.offsets[index] as usize;
279        let end = self.offsets[index + 1] as usize;
280        &self.positions[start..end]
281    }
282
283    /// Iterate over streamlines.
284    pub fn streamlines(&self) -> impl Iterator<Item = &[[f32; 3]]> {
285        self.offsets.windows(2).map(|window| {
286            let start = window[0] as usize;
287            let end = window[1] as usize;
288            &self.positions[start..end]
289        })
290    }
291
292    /// Materialize a TRX file using the requested positions dtype.
293    pub fn to_trx(&self, dtype: DType) -> Result<AnyTrxFile> {
294        match dtype {
295            DType::Float16 => Ok(AnyTrxFile::F16(self.to_trx_typed::<f16>()?)),
296            DType::Float32 => Ok(AnyTrxFile::F32(self.to_trx_typed::<f32>()?)),
297            DType::Float64 => Ok(AnyTrxFile::F64(self.to_trx_typed::<f64>()?)),
298            other => Err(TrxError::DType(format!(
299                "TRX positions must be float16, float32, or float64, got {other}"
300            ))),
301        }
302    }
303
304    fn to_trx_typed<P>(&self) -> Result<TrxFile<P>>
305    where
306        P: TrxScalar + FromF32,
307    {
308        let positions: Vec<[P; 3]> = self
309            .positions
310            .iter()
311            .map(|point| {
312                [
313                    P::from_f32(point[0]),
314                    P::from_f32(point[1]),
315                    P::from_f32(point[2]),
316                ]
317            })
318            .collect();
319
320        let mut header = self.header.clone();
321        header.nb_streamlines = self.nb_streamlines() as u64;
322        header.nb_vertices = self.nb_vertices() as u64;
323
324        Ok(TrxFile::from_parts(TrxParts {
325            header,
326            positions_backing: MmapBacking::Owned(vec_to_bytes(positions)),
327            offsets_backing: MmapBacking::Owned(vec_to_bytes(self.offsets.clone())),
328            dps: clone_arrays(&self.dps),
329            dpv: clone_arrays(&self.dpv),
330            groups: groups_to_data_arrays(&self.groups),
331            dpg: clone_dpg(&self.dpg),
332            tempdir: None,
333        }))
334    }
335}
336
337impl<P: TrxScalar> From<&TrxFile<P>> for Tractogram {
338    fn from(value: &TrxFile<P>) -> Self {
339        Self::from_trx(value)
340    }
341}
342
343impl From<&AnyTrxFile> for Tractogram {
344    fn from(value: &AnyTrxFile) -> Self {
345        Self::from_any_trx(value)
346    }
347}
348
349impl Default for Tractogram {
350    fn default() -> Self {
351        Self::new()
352    }
353}
354
355fn clone_groups(groups: &HashMap<String, DataArray>) -> HashMap<String, Vec<u32>> {
356    groups
357        .iter()
358        .map(|(name, arr)| (name.clone(), arr.to_u32_vec()))
359        .collect()
360}
361
362fn clone_arrays(arrays: &HashMap<String, DataArray>) -> HashMap<String, DataArray> {
363    arrays
364        .iter()
365        .map(|(name, arr)| (name.clone(), arr.clone_owned()))
366        .collect()
367}
368
369fn groups_to_data_arrays(groups: &HashMap<String, Vec<u32>>) -> HashMap<String, DataArray> {
370    groups
371        .iter()
372        .map(|(name, members)| {
373            (
374                name.clone(),
375                DataArray::owned_bytes(vec_to_bytes(members.clone()), 1, DType::UInt32),
376            )
377        })
378        .collect()
379}
380
381fn clone_dpg(dpg: &DataPerGroup) -> DataPerGroup {
382    dpg.iter()
383        .map(|(group, entries)| {
384            (
385                group.clone(),
386                entries
387                    .iter()
388                    .map(|(name, arr)| (name.clone(), arr.clone_owned()))
389                    .collect(),
390            )
391        })
392        .collect()
393}
394
395trait FromF32 {
396    fn from_f32(value: f32) -> Self;
397}
398
399impl FromF32 for f16 {
400    fn from_f32(value: f32) -> Self {
401        f16::from_f32(value)
402    }
403}
404
405impl FromF32 for f32 {
406    fn from_f32(value: f32) -> Self {
407        value
408    }
409}
410
411impl FromF32 for f64 {
412    fn from_f32(value: f32) -> Self {
413        value as f64
414    }
415}
416
417#[cfg(test)]
418mod tests {
419    use super::*;
420
421    #[test]
422    fn streamline_push_updates_counts() {
423        let mut tractogram = Tractogram::new();
424        tractogram
425            .push_streamline(&[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
426            .unwrap();
427        tractogram.push_streamline(&[[7.0, 8.0, 9.0]]).unwrap();
428
429        assert_eq!(tractogram.nb_streamlines(), 2);
430        assert_eq!(tractogram.nb_vertices(), 3);
431        assert_eq!(tractogram.offsets, vec![0, 2, 3]);
432        assert_eq!(tractogram.streamline(1), &[[7.0, 8.0, 9.0]]);
433    }
434}
trx_rs/tractogram.rs

trx_rs/
tractogram.rs
