trx_rs/fit/
tractogram.rs

1//! Apply Catmull-Rom fitting at the [`Tractogram`] level.
2//!
3//! Produces a new tractogram with simplified positions and a
4//! `header.extra["catmull_rom_fitted"]` marker so downstream tools can
5//! recognise the file as ready-to-edit (no second pass of fitting needed).
6
7use std::collections::HashMap;
8
9use serde::{Deserialize, Serialize};
10
11use crate::dtype::DType;
12use crate::error::{Result, TrxError};
13use crate::mmap_backing::vec_to_bytes;
14use crate::tractogram::Tractogram;
15use crate::trx_file::DataArray;
16
17use super::catmull_rom::simplify_streamline;
18
19/// JSON key written under [`Tractogram::extra_mut`] to mark a tractogram as
20/// already Catmull-Rom-fitted.
21pub const FITTED_MARKER_KEY: &str = "catmull_rom_fitted";
22
23/// Schema for the [`FITTED_MARKER_KEY`] payload. Stable across releases of
24/// this crate; consumers should treat unknown fields as ignorable.
25#[derive(Clone, Debug, Serialize, Deserialize)]
26pub struct FittedMarker {
27    pub version: u32,
28    pub epsilon_mm: f32,
29    pub fitter: String,
30}
31
32impl FittedMarker {
33    pub fn new(epsilon_mm: f32) -> Self {
34        Self {
35            version: 1,
36            epsilon_mm,
37            fitter: "trx-rs".to_string(),
38        }
39    }
40}
41
42#[derive(Clone, Debug)]
43pub struct SimplifyOptions {
44    /// Hausdorff tolerance in the same units as the tractogram positions
45    /// (RAS+ mm for TRX). Smaller values keep more control points.
46    pub epsilon_mm: f32,
47    /// If `Some`, only this group's streamlines are written to the output.
48    /// `None` keeps all streamlines and preserves every group.
49    pub group: Option<String>,
50    /// Default `width` value (mm) written as a per-vertex DPV column. Lets
51    /// downstream editors honour future width-driven dispersal without
52    /// requiring the user to enter values up-front.
53    pub default_width_mm: f32,
54    /// Default `tension` value written as a per-vertex DPV column.
55    pub default_tension: f32,
56    /// Number of streamlines to process in parallel (per Rayon default if
57    /// `None`). Currently the implementation is single-threaded; this knob
58    /// is reserved for a future `parallel` feature gate.
59    pub _parallel_chunks: Option<usize>,
60}
61
62impl Default for SimplifyOptions {
63    fn default() -> Self {
64        Self {
65            epsilon_mm: 1.0,
66            group: None,
67            default_width_mm: 1.0,
68            default_tension: 0.5,
69            _parallel_chunks: None,
70        }
71    }
72}
73
74#[derive(Clone, Copy, Debug, Default)]
75pub struct SimplifyStats {
76    pub input_streamlines: usize,
77    pub output_streamlines: usize,
78    pub input_vertices: usize,
79    pub output_vertices: usize,
80    pub groups_preserved: usize,
81}
82
83impl SimplifyStats {
84    pub fn vertex_compression_ratio(&self) -> f32 {
85        if self.output_vertices == 0 {
86            0.0
87        } else {
88            self.input_vertices as f32 / self.output_vertices as f32
89        }
90    }
91}
92
93/// Build a new [`Tractogram`] from `input` with each streamline simplified
94/// to a Catmull-Rom-fittable control polygon. Preserves the source header
95/// (so spatial metadata round-trips) and writes the
96/// [`FITTED_MARKER_KEY`] marker.
97///
98/// When `opts.group` is set, only streamlines belonging to that group are
99/// kept; output groups are remapped to the new (compacted) streamline
100/// indices.
101pub fn simplify_tractogram(
102    input: &Tractogram,
103    opts: &SimplifyOptions,
104) -> Result<(Tractogram, SimplifyStats)> {
105    if !opts.epsilon_mm.is_finite() || opts.epsilon_mm <= 0.0 {
106        return Err(TrxError::Argument(format!(
107            "epsilon_mm must be a positive finite value, got {}",
108            opts.epsilon_mm
109        )));
110    }
111
112    // Decide which streamlines we're keeping. `keep_input_idx[i] = Some(out_idx)`
113    // when input streamline `i` should be written; `None` means dropped.
114    let nb_in = input.nb_streamlines();
115    let keep_set: Option<std::collections::HashSet<u32>> = match &opts.group {
116        Some(name) => {
117            let members = input
118                .group(name)
119                .ok_or_else(|| TrxError::Argument(format!("group `{name}` not found")))?;
120            Some(members.iter().copied().collect())
121        }
122        None => None,
123    };
124
125    let mut output = Tractogram::with_header(input.header().clone());
126    let mut stats = SimplifyStats::default();
127    let mut new_widths: Vec<f32> = Vec::new();
128    let mut new_tensions: Vec<f32> = Vec::new();
129    // Map old streamline index → new streamline index (only set for kept).
130    let mut remap: Vec<Option<u32>> = vec![None; nb_in];
131
132    for (input_idx, remap_slot) in remap.iter_mut().enumerate() {
133        if let Some(keep) = &keep_set {
134            if !keep.contains(&(input_idx as u32)) {
135                continue;
136            }
137        }
138        let dense = input.streamline(input_idx);
139        stats.input_streamlines += 1;
140        stats.input_vertices += dense.len();
141
142        let simplified = simplify_streamline(dense, opts.epsilon_mm);
143        let out_idx = output.nb_streamlines() as u32;
144        output.push_streamline(&simplified)?;
145        let cp_count = simplified.len();
146        stats.output_vertices += cp_count;
147        new_widths.extend(std::iter::repeat_n(opts.default_width_mm, cp_count));
148        new_tensions.extend(std::iter::repeat_n(opts.default_tension, cp_count));
149        *remap_slot = Some(out_idx);
150    }
151    stats.output_streamlines = output.nb_streamlines();
152
153    if !new_widths.is_empty() {
154        output.insert_dpv("width", scalar_dpv(new_widths));
155        output.insert_dpv("tension", scalar_dpv(new_tensions));
156    }
157
158    let mut group_remap: HashMap<String, Vec<u32>> = HashMap::new();
159    for (name, members) in input.groups() {
160        if let Some(filter) = &opts.group {
161            if name != filter {
162                continue;
163            }
164        }
165        let mut remapped: Vec<u32> = members
166            .iter()
167            .filter_map(|&idx| remap.get(idx as usize).copied().flatten())
168            .collect();
169        if remapped.is_empty() {
170            continue;
171        }
172        remapped.sort_unstable();
173        group_remap.insert(name.clone(), remapped);
174    }
175    stats.groups_preserved = group_remap.len();
176    for (name, members) in group_remap {
177        output.insert_group(name, members);
178    }
179
180    let marker = FittedMarker::new(opts.epsilon_mm);
181    let json = serde_json::to_value(&marker)
182        .map_err(|e| TrxError::Argument(format!("serialise fitted marker: {e}")))?;
183    output
184        .extra_mut()
185        .insert(FITTED_MARKER_KEY.to_string(), json);
186
187    Ok((output, stats))
188}
189
190fn scalar_dpv(values: Vec<f32>) -> DataArray {
191    DataArray::owned_bytes(vec_to_bytes(values), 1, DType::Float32)
192}
193
194#[cfg(test)]
195mod tests {
196    use super::*;
197    use crate::tractogram::Tractogram;
198
199    fn dense_streamline(seed: f32, n: usize) -> Vec<[f32; 3]> {
200        (0..=n)
201            .map(|i| {
202                let t = i as f32 / n as f32;
203                let x = seed + 50.0 * t;
204                let y = 10.0 * (t * std::f32::consts::PI).sin();
205                let z = 5.0 * t;
206                [x, y, z]
207            })
208            .collect()
209    }
210
211    fn build_input(streamline_count: usize, vertices: usize) -> Tractogram {
212        let mut t = Tractogram::new();
213        for s in 0..streamline_count {
214            let pts = dense_streamline(s as f32 * 1.0, vertices);
215            t.push_streamline(&pts).unwrap();
216        }
217        t.insert_group("even", (0..streamline_count as u32).step_by(2).collect());
218        t.insert_group("odd", (1..streamline_count as u32).step_by(2).collect());
219        t
220    }
221
222    #[test]
223    fn simplify_compresses_and_marks() {
224        let input = build_input(4, 100);
225        let (output, stats) = simplify_tractogram(&input, &SimplifyOptions::default()).unwrap();
226
227        assert_eq!(stats.input_streamlines, 4);
228        assert_eq!(stats.output_streamlines, 4);
229        assert!(stats.output_vertices < stats.input_vertices);
230        assert!(stats.vertex_compression_ratio() > 4.0);
231
232        // Marker present.
233        let marker = output.extra().get(FITTED_MARKER_KEY).unwrap();
234        let parsed: FittedMarker = serde_json::from_value(marker.clone()).unwrap();
235        assert_eq!(parsed.version, 1);
236        assert!((parsed.epsilon_mm - 1.0).abs() < 1e-6);
237
238        // DPVs match new vertex count.
239        let widths = output.dpv_arrays().get("width").unwrap();
240        assert_eq!(widths.dtype(), DType::Float32);
241        assert_eq!(widths.cast_slice::<f32>().len(), output.nb_vertices());
242
243        // Both groups remap.
244        assert_eq!(stats.groups_preserved, 2);
245        assert_eq!(output.group("even").unwrap().len(), 2);
246        assert_eq!(output.group("odd").unwrap().len(), 2);
247    }
248
249    #[test]
250    fn group_filter_drops_others() {
251        let input = build_input(4, 100);
252        let opts = SimplifyOptions {
253            group: Some("even".to_string()),
254            ..SimplifyOptions::default()
255        };
256        let (output, stats) = simplify_tractogram(&input, &opts).unwrap();
257        assert_eq!(stats.input_streamlines, 2);
258        assert_eq!(stats.output_streamlines, 2);
259        assert_eq!(output.groups().len(), 1);
260        assert_eq!(output.group("even").unwrap(), &[0, 1]);
261    }
262
263    #[test]
264    fn unknown_group_is_an_error() {
265        let input = build_input(2, 50);
266        let opts = SimplifyOptions {
267            group: Some("nonexistent".to_string()),
268            ..SimplifyOptions::default()
269        };
270        let err = simplify_tractogram(&input, &opts).unwrap_err();
271        assert!(err.to_string().contains("nonexistent"));
272    }
273
274    #[test]
275    fn invalid_epsilon_rejected() {
276        let input = build_input(1, 50);
277        for bad in [0.0, -1.0, f32::NAN, f32::INFINITY] {
278            let opts = SimplifyOptions {
279                epsilon_mm: bad,
280                ..SimplifyOptions::default()
281            };
282            assert!(simplify_tractogram(&input, &opts).is_err());
283        }
284    }
285}
trx_rs/fit/tractogram.rs

trx_rs/fit/
tractogram.rs
