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

trx_rs/
tractogram.rs
