trx_rs/
any_trx_file.rs

1use half::f16;
2use std::io::BufReader;
3use std::path::Path;
4
5use crate::dtype::DType;
6use crate::error::{Result, TrxError};
7use crate::header::Header;
8use crate::io::filename::TrxFilename;
9use crate::trx_file::{DataArrayInfo, TrxFile};
10
11/// References to positions data, dispatched by runtime dtype.
12pub enum PositionsRef<'a> {
13    F16(&'a [[f16; 3]]),
14    F32(&'a [[f32; 3]]),
15    F64(&'a [[f64; 3]]),
16}
17
18/// A type-erased TRX container that can hold any position dtype.
19///
20/// Use this when the position dtype is not known at compile time (e.g. CLI tools
21/// that accept arbitrary `.trx` files).
22pub enum AnyTrxFile {
23    F16(TrxFile<f16>),
24    F32(TrxFile<f32>),
25    F64(TrxFile<f64>),
26}
27
28impl AnyTrxFile {
29    /// Load a TRX file, detecting the positions dtype at runtime.
30    pub fn load(path: &Path) -> Result<Self> {
31        let dtype = detect_positions_dtype(path)?;
32        match dtype {
33            DType::Float16 => Ok(AnyTrxFile::F16(TrxFile::<f16>::load(path)?)),
34            DType::Float32 => Ok(AnyTrxFile::F32(TrxFile::<f32>::load(path)?)),
35            DType::Float64 => Ok(AnyTrxFile::F64(TrxFile::<f64>::load(path)?)),
36            other => Err(TrxError::DType(format!(
37                "positions dtype {other} is not a float type"
38            ))),
39        }
40    }
41
42    /// Get a reference to the positions, dispatched by dtype.
43    pub fn positions_ref(&self) -> PositionsRef<'_> {
44        match self {
45            AnyTrxFile::F16(f) => PositionsRef::F16(f.positions()),
46            AnyTrxFile::F32(f) => PositionsRef::F32(f.positions()),
47            AnyTrxFile::F64(f) => PositionsRef::F64(f.positions()),
48        }
49    }
50
51    /// The positions dtype.
52    pub fn dtype(&self) -> DType {
53        match self {
54            AnyTrxFile::F16(_) => DType::Float16,
55            AnyTrxFile::F32(_) => DType::Float32,
56            AnyTrxFile::F64(_) => DType::Float64,
57        }
58    }
59
60    /// The header.
61    pub fn header(&self) -> &Header {
62        match self {
63            AnyTrxFile::F16(f) => f.header(),
64            AnyTrxFile::F32(f) => f.header(),
65            AnyTrxFile::F64(f) => f.header(),
66        }
67    }
68
69    /// Number of streamlines.
70    pub fn nb_streamlines(&self) -> usize {
71        match self {
72            AnyTrxFile::F16(f) => f.nb_streamlines(),
73            AnyTrxFile::F32(f) => f.nb_streamlines(),
74            AnyTrxFile::F64(f) => f.nb_streamlines(),
75        }
76    }
77
78    /// Number of vertices.
79    pub fn nb_vertices(&self) -> usize {
80        match self {
81            AnyTrxFile::F16(f) => f.nb_vertices(),
82            AnyTrxFile::F32(f) => f.nb_vertices(),
83            AnyTrxFile::F64(f) => f.nb_vertices(),
84        }
85    }
86
87    /// Dispatch to a closure with a concrete `&TrxFile<P>`.
88    pub fn with_typed<R>(
89        &self,
90        on_f16: impl FnOnce(&TrxFile<f16>) -> R,
91        on_f32: impl FnOnce(&TrxFile<f32>) -> R,
92        on_f64: impl FnOnce(&TrxFile<f64>) -> R,
93    ) -> R {
94        match self {
95            AnyTrxFile::F16(f) => on_f16(f),
96            AnyTrxFile::F32(f) => on_f32(f),
97            AnyTrxFile::F64(f) => on_f64(f),
98        }
99    }
100
101    pub fn positions_f32(&self) -> Vec<[f32; 3]> {
102        match self.positions_ref() {
103            PositionsRef::F16(data) => data
104                .iter()
105                .map(|point| [point[0].to_f32(), point[1].to_f32(), point[2].to_f32()])
106                .collect(),
107            PositionsRef::F32(data) => data.to_vec(),
108            PositionsRef::F64(data) => data
109                .iter()
110                .map(|point| [point[0] as f32, point[1] as f32, point[2] as f32])
111                .collect(),
112        }
113    }
114
115    pub fn offsets_vec(&self) -> Vec<u32> {
116        self.with_typed(
117            TrxFile::<f16>::offsets_vec,
118            TrxFile::<f32>::offsets_vec,
119            TrxFile::<f64>::offsets_vec,
120        )
121    }
122
123    pub fn dpv_entries(&self) -> Vec<(String, DataArrayInfo)> {
124        self.with_typed(
125            |trx| {
126                trx.iter_dpv()
127                    .map(|(name, info)| (name.to_string(), info))
128                    .collect()
129            },
130            |trx| {
131                trx.iter_dpv()
132                    .map(|(name, info)| (name.to_string(), info))
133                    .collect()
134            },
135            |trx| {
136                trx.iter_dpv()
137                    .map(|(name, info)| (name.to_string(), info))
138                    .collect()
139            },
140        )
141    }
142
143    pub fn dps_entries(&self) -> Vec<(String, DataArrayInfo)> {
144        self.with_typed(
145            |trx| {
146                trx.iter_dps()
147                    .map(|(name, info)| (name.to_string(), info))
148                    .collect()
149            },
150            |trx| {
151                trx.iter_dps()
152                    .map(|(name, info)| (name.to_string(), info))
153                    .collect()
154            },
155            |trx| {
156                trx.iter_dps()
157                    .map(|(name, info)| (name.to_string(), info))
158                    .collect()
159            },
160        )
161    }
162
163    pub fn groups_owned(&self) -> Vec<(String, Vec<u32>)> {
164        self.with_typed(
165            TrxFile::<f16>::group_entries_owned,
166            TrxFile::<f32>::group_entries_owned,
167            TrxFile::<f64>::group_entries_owned,
168        )
169    }
170
171    pub fn dpg_group_entries(&self) -> Vec<(String, Vec<(String, DataArrayInfo)>)> {
172        self.with_typed(
173            collect_dpg_group_entries,
174            collect_dpg_group_entries,
175            collect_dpg_group_entries,
176        )
177    }
178
179    pub fn scalar_dpv_f32(&self, name: &str) -> Result<Vec<f32>> {
180        self.with_typed(
181            |trx| trx.scalar_dpv_f32(name),
182            |trx| trx.scalar_dpv_f32(name),
183            |trx| trx.scalar_dpv_f32(name),
184        )
185    }
186
187    pub fn scalar_dps_f32(&self, name: &str) -> Result<Vec<f32>> {
188        self.with_typed(
189            |trx| trx.scalar_dps_f32(name),
190            |trx| trx.scalar_dps_f32(name),
191            |trx| trx.scalar_dps_f32(name),
192        )
193    }
194
195    /// Return a new `AnyTrxFile` with the header replaced.
196    pub fn with_updated_header(self, header: Header) -> Self {
197        match self {
198            AnyTrxFile::F16(f) => AnyTrxFile::F16(f.with_updated_header(header)),
199            AnyTrxFile::F32(f) => AnyTrxFile::F32(f.with_updated_header(header)),
200            AnyTrxFile::F64(f) => AnyTrxFile::F64(f.with_updated_header(header)),
201        }
202    }
203
204    pub fn save(&self, path: &Path) -> Result<()> {
205        self.with_typed(
206            |trx| trx.save(path),
207            |trx| trx.save(path),
208            |trx| trx.save(path),
209        )
210    }
211
212    pub fn convert_positions_dtype(&self, dtype: DType) -> Result<Self> {
213        match dtype {
214            DType::Float16 => self.with_typed(
215                |trx| Ok(Self::F16(trx.clone_with_positions_dtype::<f16>())),
216                |trx| Ok(Self::F16(trx.clone_with_positions_dtype::<f16>())),
217                |trx| Ok(Self::F16(trx.clone_with_positions_dtype::<f16>())),
218            ),
219            DType::Float32 => self.with_typed(
220                |trx| Ok(Self::F32(trx.clone_with_positions_dtype::<f32>())),
221                |trx| Ok(Self::F32(trx.clone_with_positions_dtype::<f32>())),
222                |trx| Ok(Self::F32(trx.clone_with_positions_dtype::<f32>())),
223            ),
224            DType::Float64 => self.with_typed(
225                |trx| Ok(Self::F64(trx.clone_with_positions_dtype::<f64>())),
226                |trx| Ok(Self::F64(trx.clone_with_positions_dtype::<f64>())),
227                |trx| Ok(Self::F64(trx.clone_with_positions_dtype::<f64>())),
228            ),
229            other => Err(TrxError::DType(format!(
230                "TRX positions must be float16, float32, or float64, got {other}"
231            ))),
232        }
233    }
234}
235
236impl std::fmt::Debug for AnyTrxFile {
237    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
238        match self {
239            AnyTrxFile::F16(t) => t.fmt(f),
240            AnyTrxFile::F32(t) => t.fmt(f),
241            AnyTrxFile::F64(t) => t.fmt(f),
242        }
243    }
244}
245
246fn collect_dpg_group_entries<P: crate::dtype::TrxScalar>(
247    trx: &TrxFile<P>,
248) -> Vec<(String, Vec<(String, DataArrayInfo)>)> {
249    trx.dpg_group_names()
250        .into_iter()
251        .map(|group| {
252            let entries = trx
253                .dpg_entries(group)
254                .expect("group name came from dpg_group_names()");
255            (group.to_string(), entries)
256        })
257        .collect()
258}
259
260/// Detect the positions dtype from a TRX path (directory or zip).
261pub fn detect_positions_dtype(path: &Path) -> Result<DType> {
262    if path.is_dir() {
263        detect_positions_dtype_dir(path)
264    } else if path.is_file() {
265        detect_positions_dtype_zip(path)
266    } else {
267        Err(TrxError::FileNotFound(path.to_path_buf()))
268    }
269}
270
271fn detect_positions_dtype_dir(dir: &Path) -> Result<DType> {
272    for entry in std::fs::read_dir(dir)? {
273        let entry = entry?;
274        let name = entry.file_name();
275        let name_str = name.to_string_lossy();
276        if name_str.starts_with("positions.") {
277            let parsed = TrxFilename::parse(&name_str)?;
278            return Ok(parsed.dtype);
279        }
280    }
281
282    let header_path = dir.join("header.json");
283    if header_path.is_file() {
284        let file = std::fs::File::open(&header_path)?;
285        let reader = BufReader::new(file);
286        let header: Header = serde_json::from_reader(reader)?;
287        if header.nb_vertices == 0 {
288            return Ok(DType::Float16);
289        }
290    }
291    Err(TrxError::FileNotFound(dir.join("positions")))
292}
293
294fn detect_positions_dtype_zip(path: &Path) -> Result<DType> {
295    let file = std::fs::File::open(path)?;
296    let reader = BufReader::new(file);
297    let mut archive = zip::ZipArchive::new(reader)?;
298
299    for i in 0..archive.len() {
300        let name = archive.name_for_index(i).unwrap_or("");
301        let basename = name.rsplit('/').next().unwrap_or(name);
302        if basename.starts_with("positions.") {
303            let parsed = TrxFilename::parse(basename)?;
304            return Ok(parsed.dtype);
305        }
306    }
307
308    if let Ok(mut header_entry) = archive.by_name("header.json") {
309        let mut bytes = Vec::new();
310        if std::io::Read::read_to_end(&mut header_entry, &mut bytes).is_ok() {
311            if let Ok(header) = serde_json::from_slice::<Header>(&bytes) {
312                if header.nb_vertices == 0 {
313                    return Ok(DType::Float16);
314                }
315            }
316        }
317    }
318
319    Err(TrxError::FileNotFound(path.join("positions")))
320}
321
322#[cfg(test)]
323mod tests {
324    use super::*;
325    use tempfile::TempDir;
326
327    #[test]
328    fn detect_positions_dtype_missing_positions_non_empty_errors() {
329        let dir = TempDir::new().unwrap();
330        let header_path = dir.path().join("header.json");
331        let header = Header {
332            voxel_to_rasmm: Header::identity_affine(),
333            dimensions: [100, 100, 100],
334            nb_streamlines: 1,
335            nb_vertices: 10,
336            extra: Default::default(),
337        };
338        let json = serde_json::to_string(&header).unwrap();
339        std::fs::write(&header_path, json).unwrap();
340
341        let result = detect_positions_dtype(dir.path());
342        assert!(result.is_err());
343    }
344
345    #[test]
346    fn detect_positions_dtype_missing_positions_empty_succeeds() {
347        let dir = TempDir::new().unwrap();
348        let header_path = dir.path().join("header.json");
349        let header = Header {
350            voxel_to_rasmm: Header::identity_affine(),
351            dimensions: [100, 100, 100],
352            nb_streamlines: 0,
353            nb_vertices: 0,
354            extra: Default::default(),
355        };
356        let json = serde_json::to_string(&header).unwrap();
357        std::fs::write(&header_path, json).unwrap();
358
359        let result = detect_positions_dtype(dir.path()).unwrap();
360        assert_eq!(result, DType::Float16);
361    }
362}
trx_rs/any_trx_file.rs

trx_rs/
any_trx_file.rs
