trx_rs/
stream.rs

1use std::path::Path;
2
3use crate::dtype::TrxScalar;
4use crate::error::Result;
5use crate::header::Header;
6use crate::mmap_backing::MmapBacking;
7use crate::trx_file::{TrxFile, TrxParts};
8
9/// Incremental builder for constructing a TRX file streamline-by-streamline.
10///
11/// Buffers positions and offsets in memory, then finalizes into a `TrxFile`
12/// or writes directly to disk.
13pub struct TrxStream<P: TrxScalar> {
14    positions: Vec<[P; 3]>,
15    offsets: Vec<u32>,
16    header: Header,
17}
18
19impl<P: TrxScalar> TrxStream<P> {
20    /// Create a new stream builder with the given affine and dimensions.
21    pub fn new(voxel_to_rasmm: [[f64; 4]; 4], dimensions: [u64; 3]) -> Self {
22        Self {
23            positions: Vec::new(),
24            offsets: vec![0],
25            header: Header {
26                voxel_to_rasmm,
27                dimensions,
28                nb_streamlines: 0,
29                nb_vertices: 0,
30                extra: Default::default(),
31            },
32        }
33    }
34
35    /// Push a single streamline (slice of 3D points).
36    pub fn push_streamline(&mut self, points: &[[P; 3]]) {
37        self.positions.extend_from_slice(points);
38        self.offsets.push(self.positions.len() as u32);
39        self.header.nb_streamlines += 1;
40        self.header.nb_vertices += points.len() as u64;
41    }
42
43    /// Number of streamlines added so far.
44    pub fn nb_streamlines(&self) -> usize {
45        self.header.nb_streamlines as usize
46    }
47
48    /// Number of vertices added so far.
49    pub fn nb_vertices(&self) -> usize {
50        self.header.nb_vertices as usize
51    }
52
53    /// Finalize into an in-memory `TrxFile`.
54    pub fn finalize(self) -> TrxFile<P> {
55        let pos_bytes = crate::mmap_backing::vec_to_bytes(self.positions);
56        let off_bytes = crate::mmap_backing::vec_to_bytes(self.offsets);
57
58        TrxFile::from_parts(TrxParts {
59            header: self.header,
60            positions_backing: MmapBacking::Owned(pos_bytes),
61            offsets_backing: MmapBacking::Owned(off_bytes),
62            dps: Default::default(),
63            dpv: Default::default(),
64            groups: Default::default(),
65            dpg: Default::default(),
66            tempdir: None,
67        })
68    }
69
70    /// Finalize and save to the given path.
71    pub fn finalize_to(self, path: &Path) -> Result<()> {
72        let trx = self.finalize();
73        trx.save(path)
74    }
75}
76
77#[cfg(test)]
78mod tests {
79    use super::*;
80
81    #[test]
82    fn stream_build_and_finalize() {
83        let mut stream = TrxStream::<f32>::new(Header::identity_affine(), [100, 100, 100]);
84
85        stream.push_streamline(&[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]);
86        stream.push_streamline(&[[7.0, 8.0, 9.0]]);
87
88        assert_eq!(stream.nb_streamlines(), 2);
89        assert_eq!(stream.nb_vertices(), 3);
90
91        let trx = stream.finalize();
92        assert_eq!(trx.nb_streamlines(), 2);
93        assert_eq!(trx.nb_vertices(), 3);
94        assert_eq!(trx.streamline(0), &[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]);
95        assert_eq!(trx.streamline(1), &[[7.0, 8.0, 9.0]]);
96    }
97}
trx_rs/stream.rs

trx_rs/
stream.rs
