trx_rs/ops/
subset.rs

1use std::collections::HashMap;
2
3use crate::dtype::TrxScalar;
4use crate::error::{Result, TrxError};
5use crate::header::Header;
6use crate::mmap_backing::MmapBacking;
7use crate::trx_file::{DataArray, DataPerGroup, TrxFile, TrxParts};
8
9/// Extract a subset of streamlines by index, producing a new `TrxFile`.
10///
11/// All DPS, DPV, and group arrays are remapped accordingly.
12pub fn subset_streamlines<P: TrxScalar>(trx: &TrxFile<P>, indices: &[usize]) -> Result<TrxFile<P>> {
13    let (new_positions, new_offsets) =
14        collect_positions_and_offsets(trx.positions(), trx.offsets(), indices)?;
15
16    let nb_streamlines = indices.len() as u64;
17    let nb_vertices = new_positions.len() as u64;
18
19    // Remap DPS (data per streamline)
20    let new_dps = remap_dps(trx.dps_arrays(), indices);
21
22    // Remap DPV (data per vertex)
23    let new_dpv = remap_dpv(trx.dpv_arrays(), trx.offsets(), indices);
24
25    // Remap groups
26    let new_groups = remap_groups(trx.group_arrays(), indices);
27    let new_dpg = remap_dpg(trx.dpg_arrays(), &new_groups);
28
29    let header = Header {
30        voxel_to_rasmm: trx.header().voxel_to_rasmm,
31        dimensions: trx.header().dimensions,
32        nb_streamlines,
33        nb_vertices,
34        extra: trx.header().extra.clone(),
35    };
36
37    let pos_bytes = crate::mmap_backing::vec_to_bytes(new_positions);
38    let off_bytes = crate::mmap_backing::vec_to_bytes(new_offsets);
39
40    Ok(TrxFile::from_parts(TrxParts {
41        header,
42        positions_backing: MmapBacking::Owned(pos_bytes),
43        offsets_backing: MmapBacking::Owned(off_bytes),
44        dps: new_dps,
45        dpv: new_dpv,
46        groups: new_groups,
47        dpg: new_dpg,
48    }))
49}
50
51/// Remap DPS arrays: select rows by streamline index.
52fn remap_dps(dps: &HashMap<String, DataArray>, indices: &[usize]) -> HashMap<String, DataArray> {
53    dps.iter()
54        .map(|(name, arr)| {
55            let row_bytes = arr.ncols() * arr.dtype().size_of();
56            let dst = copy_row_bytes(arr.as_bytes(), row_bytes, indices);
57            (
58                name.clone(),
59                DataArray::from_backing(MmapBacking::Owned(dst), arr.ncols(), arr.dtype()),
60            )
61        })
62        .collect()
63}
64
65/// Remap DPV arrays: select vertex ranges corresponding to selected streamlines.
66fn remap_dpv(
67    dpv: &HashMap<String, DataArray>,
68    offsets: &[u32],
69    indices: &[usize],
70) -> HashMap<String, DataArray> {
71    dpv.iter()
72        .map(|(name, arr)| {
73            let row_bytes = arr.ncols() * arr.dtype().size_of();
74            let dst = copy_vertex_ranges(arr.as_bytes(), row_bytes, offsets, indices);
75            (
76                name.clone(),
77                DataArray::from_backing(MmapBacking::Owned(dst), arr.ncols(), arr.dtype()),
78            )
79        })
80        .collect()
81}
82
83/// Remap groups: update streamline indices to reflect the new ordering.
84fn remap_groups(
85    groups: &HashMap<String, DataArray>,
86    indices: &[usize],
87) -> HashMap<String, DataArray> {
88    // Build a reverse map: old_index → new_index
89    let mut old_to_new: HashMap<usize, u32> = HashMap::new();
90    for (new_idx, &old_idx) in indices.iter().enumerate() {
91        old_to_new.insert(old_idx, new_idx as u32);
92    }
93
94    let mut out = HashMap::new();
95    for (name, arr) in groups {
96        let old_members: &[u32] = arr.cast_slice();
97        let new_members: Vec<u32> = old_members
98            .iter()
99            .filter_map(|&m| old_to_new.get(&(m as usize)).copied())
100            .collect();
101        let bytes = crate::mmap_backing::vec_to_bytes(new_members);
102        out.insert(
103            name.clone(),
104            DataArray::from_backing(MmapBacking::Owned(bytes), 1, crate::dtype::DType::UInt32),
105        );
106    }
107    out
108}
109
110fn remap_dpg(dpg: &DataPerGroup, groups: &HashMap<String, DataArray>) -> DataPerGroup {
111    let mut out = HashMap::new();
112    for (group_name, entries) in dpg {
113        if let Some(group_members) = groups.get(group_name) {
114            if group_members.as_bytes().is_empty() {
115                continue;
116            }
117            out.insert(
118                group_name.clone(),
119                entries
120                    .iter()
121                    .map(|(name, arr)| (name.clone(), arr.clone_owned()))
122                    .collect(),
123            );
124        }
125    }
126    out
127}
128
129/// Axis-aligned bounding box for a single streamline.
130#[derive(Debug, Clone, Copy, PartialEq)]
131pub struct StreamlineAabb {
132    min: [f32; 3],
133    max: [f32; 3],
134}
135
136impl StreamlineAabb {
137    /// Minimum corner of the bounding box (x, y, z).
138    pub fn min(&self) -> [f32; 3] {
139        self.min
140    }
141
142    /// Maximum corner of the bounding box (x, y, z).
143    pub fn max(&self) -> [f32; 3] {
144        self.max
145    }
146
147    /// Returns true if this AABB overlaps the given query box.
148    pub fn overlaps_box(&self, min: [f32; 3], max: [f32; 3]) -> bool {
149        self.min[0] <= max[0]
150            && self.max[0] >= min[0]
151            && self.min[1] <= max[1]
152            && self.max[1] >= min[1]
153            && self.min[2] <= max[2]
154            && self.max[2] >= min[2]
155    }
156}
157
158/// Compute per-streamline AABBs for all streamlines in a `TrxFile`.
159///
160/// Returns a `Vec` of length `nb_streamlines`, where each entry is
161/// `[min_x, min_y, min_z, max_x, max_y, max_z]` in f32.
162///
163/// This is intended to be computed once and reused across multiple queries,
164/// matching trx-cpp's `build_streamline_aabbs()` / AABB cache pattern.
165pub fn build_streamline_aabbs<P: TrxScalar>(trx: &TrxFile<P>) -> Vec<StreamlineAabb> {
166    build_streamline_aabbs_from_iter(
167        trx.offsets(),
168        trx.streamlines().map(|streamline| {
169            streamline
170                .iter()
171                .map(|point| [point[0].to_f32(), point[1].to_f32(), point[2].to_f32()])
172        }),
173    )
174}
175
176/// Compute per-streamline AABBs from flat position and offset slices.
177///
178/// Useful when positions are already in a flat `[[f32; 3]]` array rather than
179/// inside a `TrxFile`.
180pub fn build_streamline_aabbs_from_slices(
181    positions: &[[f32; 3]],
182    offsets: &[u32],
183) -> Vec<StreamlineAabb> {
184    build_streamline_aabbs_from_iter(
185        offsets,
186        offsets.windows(2).map(|window| {
187            positions[window[0] as usize..window[1] as usize]
188                .iter()
189                .copied()
190        }),
191    )
192}
193
194/// Query pre-computed AABBs against a query box, returning matching streamline indices.
195///
196/// This is the fast path: O(N) with 6 float comparisons per streamline,
197/// matching trx-cpp's `query_aabb()` implementation.
198pub fn query_aabb_cached(aabbs: &[StreamlineAabb], min: [f64; 3], max: [f64; 3]) -> Vec<usize> {
199    let min = [min[0] as f32, min[1] as f32, min[2] as f32];
200    let max = [max[0] as f32, max[1] as f32, max[2] as f32];
201
202    aabbs
203        .iter()
204        .enumerate()
205        .filter_map(|(index, aabb)| aabb.overlaps_box(min, max).then_some(index))
206        .collect()
207}
208
209/// Find streamline indices whose AABB intersects the given query box.
210///
211/// Convenience wrapper that builds AABBs on the fly. For repeated queries on the
212/// same data, use [`build_streamline_aabbs`] + [`query_aabb_cached`] instead.
213///
214/// `min` and `max` define the query AABB corners. Comparisons are done in f32,
215/// matching trx-cpp's behavior.
216pub fn query_aabb<P: TrxScalar>(trx: &TrxFile<P>, min: [f64; 3], max: [f64; 3]) -> Vec<usize> {
217    let aabbs = build_streamline_aabbs(trx);
218    query_aabb_cached(&aabbs, min, max)
219}
220
221fn collect_positions_and_offsets<P: TrxScalar>(
222    positions: &[[P; 3]],
223    offsets: &[u32],
224    indices: &[usize],
225) -> Result<(Vec<[P; 3]>, Vec<u32>)> {
226    let mut new_positions = Vec::new();
227    let mut new_offsets = vec![0];
228
229    for &idx in indices {
230        let window = offsets
231            .get(idx..=idx + 1)
232            .ok_or_else(|| TrxError::Argument(format!("streamline index {idx} out of bounds")))?;
233        new_positions.extend_from_slice(&positions[window[0] as usize..window[1] as usize]);
234        new_offsets.push(
235            u32::try_from(new_positions.len())
236                .map_err(|_| TrxError::Argument("subset would exceed u32::MAX vertices".into()))?,
237        );
238    }
239
240    Ok((new_positions, new_offsets))
241}
242
243fn copy_row_bytes(src: &[u8], row_bytes: usize, indices: &[usize]) -> Vec<u8> {
244    let mut dst = Vec::with_capacity(indices.len() * row_bytes);
245    for &idx in indices {
246        let start = idx * row_bytes;
247        dst.extend_from_slice(&src[start..start + row_bytes]);
248    }
249    dst
250}
251
252fn copy_vertex_ranges(src: &[u8], row_bytes: usize, offsets: &[u32], indices: &[usize]) -> Vec<u8> {
253    let mut dst = Vec::new();
254    for &idx in indices {
255        let start = offsets[idx] as usize * row_bytes;
256        let end = offsets[idx + 1] as usize * row_bytes;
257        dst.extend_from_slice(&src[start..end]);
258    }
259    dst
260}
261
262fn build_streamline_aabbs_from_iter<I, J>(offsets: &[u32], streamlines: I) -> Vec<StreamlineAabb>
263where
264    I: Iterator<Item = J>,
265    J: Iterator<Item = [f32; 3]>,
266{
267    let mut aabbs = Vec::with_capacity(offsets.len().saturating_sub(1));
268
269    for streamline in streamlines {
270        let mut min = [f32::INFINITY; 3];
271        let mut max = [f32::NEG_INFINITY; 3];
272        for point in streamline {
273            for axis in 0..3 {
274                min[axis] = min[axis].min(point[axis]);
275                max[axis] = max[axis].max(point[axis]);
276            }
277        }
278        aabbs.push(StreamlineAabb { min, max });
279    }
280
281    aabbs
282}
283
284#[cfg(test)]
285mod tests {
286    use super::*;
287    use crate::stream::TrxStream;
288
289    fn make_test_trx() -> TrxFile<f32> {
290        let mut stream = TrxStream::<f32>::new(Header::identity_affine(), [100, 100, 100]);
291        stream.push_streamline(&[[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]]);
292        stream.push_streamline(&[[10.0, 10.0, 10.0], [11.0, 11.0, 11.0], [12.0, 12.0, 12.0]]);
293        stream.push_streamline(&[[20.0, 20.0, 20.0]]);
294        stream.finalize()
295    }
296
297    #[test]
298    fn subset_basic() {
299        let trx = make_test_trx();
300        let sub = subset_streamlines(&trx, &[0, 2]).unwrap();
301        assert_eq!(sub.nb_streamlines(), 2);
302        assert_eq!(sub.nb_vertices(), 3);
303        assert_eq!(sub.streamline(0), &[[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]]);
304        assert_eq!(sub.streamline(1), &[[20.0, 20.0, 20.0]]);
305    }
306
307    #[test]
308    fn query_aabb_basic() {
309        let trx = make_test_trx();
310        let hits = query_aabb(&trx, [9.0, 9.0, 9.0], [15.0, 15.0, 15.0]);
311        assert_eq!(hits, vec![1]);
312    }
313
314    #[test]
315    fn query_aabb_cached_basic() {
316        let trx = make_test_trx();
317        let aabbs = build_streamline_aabbs(&trx);
318        assert_eq!(aabbs.len(), 3);
319
320        // Streamline 0: [0,0,0] to [1,1,1]
321        assert_eq!(aabbs[0].min(), [0.0, 0.0, 0.0]);
322        assert_eq!(aabbs[0].max(), [1.0, 1.0, 1.0]);
323
324        // Streamline 1: [10,10,10] to [12,12,12]
325        assert_eq!(aabbs[1].min(), [10.0, 10.0, 10.0]);
326        assert_eq!(aabbs[1].max(), [12.0, 12.0, 12.0]);
327
328        // Query should match same results as non-cached version
329        let hits = query_aabb_cached(&aabbs, [9.0, 9.0, 9.0], [15.0, 15.0, 15.0]);
330        assert_eq!(hits, vec![1]);
331
332        // Query that hits all
333        let hits = query_aabb_cached(&aabbs, [-1.0, -1.0, -1.0], [25.0, 25.0, 25.0]);
334        assert_eq!(hits, vec![0, 1, 2]);
335    }
336}
trx_rs/ops/subset.rs

trx_rs/ops/
subset.rs
