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

trx_rs/ops/
subset.rs
