modewise_dataset/
lib.rs

1//! Finds recordings on disk by file patterns and loads them as pipeline inputs.
2//!
3//! Patterns use `{entity}` placeholders, for example
4//! `sub-{subject}/func/sub-{subject}_task-{task}_run-{run}_bold.h5`, so BIDS and
5//! BIDS-like layouts of any modality are described the same way.
6
7mod channels;
8mod events_tsv;
9mod participants;
10mod pattern;
11mod spec;
12
13use std::collections::BTreeMap;
14use std::path::{Path, PathBuf};
15
16use modewise_brain::ChannelTable;
17use modewise_hdf5::read_series;
18use modewise_pipeline::{Artifact, InputSource};
19use modewise_signals::{MultichannelSeries, SamplingRate};
20use thiserror::Error;
21
22pub use channels::{TableFormat, read_table, schaefer_channel, tian_channel};
23pub use events_tsv::read as read_events_tsv;
24pub use participants::{ParticipantTable, ParticipantsSpec, read_participants};
25pub use pattern::Pattern;
26pub use spec::{ChannelsSpec, DatasetSpec, Format, InputSpec, LayoutName, TableSpec};
27
28use spec::SamplingRateSource;
29
30#[derive(Debug, Error)]
31pub enum DatasetError {
32    #[error("cannot read dataset spec: {0}")]
33    Spec(String),
34    #[error("invalid pattern {pattern:?}: {message}")]
35    Pattern { pattern: String, message: String },
36    #[error("cannot list {path}: {message}")]
37    Io { path: String, message: String },
38    #[error("filter on {0:?}, which the first input's pattern does not contain")]
39    UnknownFilterEntity(String),
40}
41
42/// One matched set of entity values, with the file for each input.
43#[derive(Debug, Clone, PartialEq)]
44pub struct Recording {
45    /// Entity values in the order they appear in the first input's pattern.
46    pub entities: Vec<(String, String)>,
47    pub paths: BTreeMap<String, PathBuf>,
48}
49
50impl Recording {
51    /// Stable name such as `subject-01_task-rest_run-1`.
52    pub fn id(&self) -> String {
53        let parts: Vec<String> = self
54            .entities
55            .iter()
56            .map(|(k, v)| format!("{k}-{v}"))
57            .collect();
58        parts.join("_")
59    }
60
61    fn entity_map(&self) -> BTreeMap<String, String> {
62        self.entities.iter().cloned().collect()
63    }
64}
65
66/// A dataset spec with its recordings discovered.
67pub struct Dataset {
68    spec: DatasetSpec,
69    /// The spec's folder, which atlas and channel-table paths are relative to.
70    base: PathBuf,
71    root: PathBuf,
72    recordings: Vec<Recording>,
73    channels: Option<ChannelTable>,
74    participants: Option<modewise_study::Participants>,
75}
76
77impl Dataset {
78    /// Reads a dataset TOML file. A relative `root` is resolved against the file's folder.
79    pub fn from_file(path: &Path) -> Result<Self, DatasetError> {
80        let text = std::fs::read_to_string(path).map_err(|e| DatasetError::Io {
81            path: path.display().to_string(),
82            message: e.to_string(),
83        })?;
84        let base = path.parent().unwrap_or(Path::new("."));
85        Self::discover(DatasetSpec::from_toml(&text)?, base)
86    }
87
88    pub fn discover(spec: DatasetSpec, base: &Path) -> Result<Self, DatasetError> {
89        let root = base.join(&spec.root);
90        let Some(first) = spec.inputs.first() else {
91            return Err(DatasetError::Spec("declare at least one [[input]]".into()));
92        };
93
94        let primary = Pattern::parse(&first.pattern)?;
95        if let Some(unknown) = spec.filter.keys().find(|k| !primary.entities().contains(k)) {
96            return Err(DatasetError::UnknownFilterEntity(unknown.clone()));
97        }
98        let others = spec.inputs[1..]
99            .iter()
100            .map(|i| Ok((i.name.clone(), Pattern::parse(&i.pattern)?)))
101            .collect::<Result<Vec<_>, DatasetError>>()?;
102
103        let mut recordings = Vec::new();
104        for (path, entities) in primary.find(&root)? {
105            let keep = spec.filter.iter().all(|(entity, allowed)| {
106                entities
107                    .iter()
108                    .any(|(k, v)| k == entity && allowed.contains(v))
109            });
110            if !keep {
111                continue;
112            }
113            let values: BTreeMap<String, String> = entities.iter().cloned().collect();
114            let mut paths = BTreeMap::from([(first.name.clone(), path)]);
115            for (name, pattern) in &others {
116                paths.insert(name.clone(), root.join(pattern.fill(&values)?));
117            }
118            recordings.push(Recording { entities, paths });
119        }
120        let channels = match &spec.channels {
121            Some(c) => {
122                let tables = c
123                    .tables
124                    .iter()
125                    .map(|t| read_table(&base.join(&t.path), t.format))
126                    .collect::<Result<Vec<_>, _>>()
127                    .map_err(DatasetError::Spec)?;
128                Some(ChannelTable::concat(tables).map_err(|e| DatasetError::Spec(e.to_string()))?)
129            }
130            None => None,
131        };
132        let participants = match &spec.participants {
133            Some(p) => Some(read_participants(&root, p).map_err(DatasetError::Spec)?),
134            None => None,
135        };
136        Ok(Self {
137            spec,
138            base: base.to_path_buf(),
139            root,
140            recordings,
141            channels,
142            participants,
143        })
144    }
145
146    /// Participant information from `[participants]`, if the spec has one.
147    pub fn participants(&self) -> Option<&modewise_study::Participants> {
148        self.participants.as_ref()
149    }
150
151    /// The channel table from `[channels]`, if the spec has one.
152    pub fn channels(&self) -> Option<&ChannelTable> {
153        self.channels.as_ref()
154    }
155
156    pub fn recordings(&self) -> &[Recording] {
157        &self.recordings
158    }
159
160    pub fn recording(&self, id: &str) -> Option<&Recording> {
161        self.recordings.iter().find(|r| r.id() == id)
162    }
163
164    fn load_input(&self, input: &InputSpec, recording: &Recording) -> Result<Artifact, String> {
165        let path = &recording.paths[&input.name];
166        match input.format {
167            Format::Hdf5 => self.load_series(input, recording).map(Artifact::Series),
168            Format::Nifti => self.load_nifti(input, recording).map(Artifact::Series),
169            Format::EventsTsv => {
170                let column = input.label_column.as_deref().unwrap_or("trial_type");
171                events_tsv::read(path, column).map(Artifact::Events)
172            }
173        }
174    }
175
176    /// Channel names: a labels file, or the channel table when `names_for` lists this input.
177    fn channel_names(&self, input: &InputSpec) -> Result<Option<Vec<String>>, String> {
178        let named = self
179            .spec
180            .channels
181            .as_ref()
182            .is_some_and(|c| c.names_for.contains(&input.name));
183        Ok(match (&input.labels_file, &self.channels) {
184            (Some(file), _) => Some(read_labels(&self.root.join(file))?),
185            (None, Some(table)) if named => Some(table.names()),
186            _ => None,
187        })
188    }
189
190    fn sampling_rate(
191        &self,
192        input: &InputSpec,
193        recording: &Recording,
194    ) -> Result<Option<f64>, String> {
195        Ok(match input.sampling_rate_source() {
196            SamplingRateSource::Hz(hz) if hz.is_nan() => None,
197            SamplingRateSource::Hz(hz) => Some(hz),
198            SamplingRateSource::Sidecar(sidecar) => {
199                let relative = Pattern::parse(sidecar)
200                    .and_then(|p| p.fill(&recording.entity_map()))
201                    .map_err(|e| e.to_string())?;
202                Some(sidecar_rate(&self.root.join(relative))?)
203            }
204        })
205    }
206
207    fn load_nifti(
208        &self,
209        input: &InputSpec,
210        recording: &Recording,
211    ) -> Result<MultichannelSeries, String> {
212        let bold =
213            modewise_nifti::read_bold(&recording.paths[&input.name]).map_err(|e| e.to_string())?;
214        let hz =
215            match self.sampling_rate(input, recording)? {
216                Some(hz) => hz,
217                None => 1.0 / bold.repetition_time_s.ok_or(
218                    "the NIfTI header has no repetition time; give `sampling_rate_hz` or `sidecar`",
219                )?,
220            };
221        let standardize = match input.standardize {
222            spec::StandardizeName::None => modewise_parcellate::Standardize::None,
223            spec::StandardizeName::ZscoreSample => modewise_parcellate::Standardize::ZscoreSample,
224        };
225        let mut rows = Vec::new();
226        for atlas in &input.atlases {
227            let path = self.base.join(atlas);
228            let labels = modewise_nifti::read_labels(&path).map_err(|e| e.to_string())?;
229            let (signals, _) = modewise_parcellate::parcellate(
230                bold.data.view(),
231                &bold.affine,
232                labels.data.view(),
233                &labels.affine,
234                standardize,
235            )
236            .map_err(|e| format!("{}: {e}", path.display()))?;
237            rows.push(signals);
238        }
239        let views: Vec<_> = rows.iter().map(|r| r.view()).collect();
240        let data = ndarray::concatenate(ndarray::Axis(0), &views).map_err(|e| e.to_string())?;
241        let labels = self
242            .channel_names(input)?
243            .unwrap_or_else(|| (0..data.nrows()).map(|c| format!("ch{c}")).collect());
244        let fs = SamplingRate::new(hz).map_err(|e| e.to_string())?;
245        MultichannelSeries::new(data, fs, labels).map_err(|e| e.to_string())
246    }
247
248    fn load_series(
249        &self,
250        input: &InputSpec,
251        recording: &Recording,
252    ) -> Result<MultichannelSeries, String> {
253        let hz = self.sampling_rate(input, recording)?.unwrap_or(f64::NAN);
254        let fs = SamplingRate::new(hz).map_err(|e| e.to_string())?;
255        let labels = self.channel_names(input)?;
256        let dataset = input.dataset.as_deref().unwrap_or_default();
257        read_series(
258            &recording.paths[&input.name],
259            dataset,
260            input.layout.into(),
261            fs,
262            labels,
263        )
264        .map_err(|e| e.to_string())
265    }
266}
267
268impl InputSource for Dataset {
269    fn input_names(&self) -> Vec<String> {
270        let mut names: Vec<String> = self.spec.inputs.iter().map(|i| i.name.clone()).collect();
271        names.extend(self.spec.channels.as_ref().map(|c| c.name.clone()));
272        names
273    }
274
275    fn recording_ids(&self) -> Vec<String> {
276        self.recordings.iter().map(Recording::id).collect()
277    }
278
279    fn load(&self, id: &str) -> Result<BTreeMap<String, Artifact>, String> {
280        let recording = self
281            .recording(id)
282            .ok_or_else(|| format!("no recording {id:?}"))?;
283        let mut loaded = self
284            .spec
285            .inputs
286            .iter()
287            .map(|input| {
288                let series = self
289                    .load_input(input, recording)
290                    .map_err(|e| format!("input {:?}: {e}", input.name))?;
291                Ok((input.name.clone(), series))
292            })
293            .collect::<Result<BTreeMap<_, _>, String>>()?;
294        if let (Some(spec), Some(table)) = (&self.spec.channels, &self.channels) {
295            loaded.insert(spec.name.clone(), Artifact::Channels(table.clone()));
296        }
297        Ok(loaded)
298    }
299}
300
301/// Reads `SamplingFrequency` (Hz) or `RepetitionTime` (s) from a BIDS JSON sidecar.
302fn sidecar_rate(path: &Path) -> Result<f64, String> {
303    let text =
304        std::fs::read_to_string(path).map_err(|e| format!("sidecar {}: {e}", path.display()))?;
305    let json: serde_json::Value =
306        serde_json::from_str(&text).map_err(|e| format!("sidecar {}: {e}", path.display()))?;
307    if let Some(hz) = json.get("SamplingFrequency").and_then(|v| v.as_f64()) {
308        return Ok(hz);
309    }
310    if let Some(tr) = json.get("RepetitionTime").and_then(|v| v.as_f64()) {
311        return Ok(1.0 / tr);
312    }
313    Err(format!(
314        "sidecar {} has neither SamplingFrequency nor RepetitionTime",
315        path.display()
316    ))
317}
318
319/// One channel label per non-empty line.
320fn read_labels(path: &Path) -> Result<Vec<String>, String> {
321    let text = std::fs::read_to_string(path)
322        .map_err(|e| format!("labels file {}: {e}", path.display()))?;
323    Ok(text
324        .lines()
325        .map(str::trim)
326        .filter(|l| !l.is_empty())
327        .map(String::from)
328        .collect())
329}