Files
rustytorch/crates/specialized/rtx-neuro/src/protocol.rs
T
2026-03-04 00:08:42 +00:00

389 lines
10 KiB
Rust

//! Protocol and study organization system.
use crate::{NeuroResult, Recording};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
/// A protocol (study/experiment) containing multiple subjects
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Protocol {
/// Protocol name
pub name: String,
/// Root directory for this protocol
pub path: PathBuf,
/// Protocol description
pub description: String,
/// Creation date
pub created: DateTime<Utc>,
/// Last modified date
pub modified: DateTime<Utc>,
/// Subjects in this protocol
subjects: HashMap<String, Subject>,
/// Group-level analyses
group_analyses: Vec<GroupAnalysis>,
/// Protocol settings
pub settings: ProtocolSettings,
}
impl Protocol {
/// Create a new protocol
#[must_use]
pub fn new(name: impl Into<String>, path: impl Into<PathBuf>) -> Self {
let now = Utc::now();
Self {
name: name.into(),
path: path.into(),
description: String::new(),
created: now,
modified: now,
subjects: HashMap::new(),
group_analyses: Vec::new(),
settings: ProtocolSettings::default(),
}
}
/// Add a subject to the protocol
pub fn add_subject(&mut self, subject: Subject) {
self.subjects.insert(subject.id.clone(), subject);
self.modified = Utc::now();
}
/// Get a subject by ID
#[must_use]
pub fn get_subject(&self, id: &str) -> Option<&Subject> {
self.subjects.get(id)
}
/// Get mutable reference to a subject
pub fn get_subject_mut(&mut self, id: &str) -> Option<&mut Subject> {
self.modified = Utc::now();
self.subjects.get_mut(id)
}
/// Remove a subject
pub fn remove_subject(&mut self, id: &str) -> Option<Subject> {
self.modified = Utc::now();
self.subjects.remove(id)
}
/// List all subject IDs
#[must_use]
pub fn subject_ids(&self) -> Vec<&str> {
self.subjects.keys().map(String::as_str).collect()
}
/// Number of subjects
#[must_use]
pub fn n_subjects(&self) -> usize {
self.subjects.len()
}
/// Iterate over subjects
pub fn subjects(&self) -> impl Iterator<Item = &Subject> {
self.subjects.values()
}
/// Add a group analysis
pub fn add_group_analysis(&mut self, analysis: GroupAnalysis) {
self.group_analyses.push(analysis);
self.modified = Utc::now();
}
/// Get group analyses
#[must_use]
pub fn group_analyses(&self) -> &[GroupAnalysis] {
&self.group_analyses
}
/// Save protocol to disk
pub fn save(&self) -> NeuroResult<()> {
let protocol_file = self.path.join("protocol.json");
let json = serde_json::to_string_pretty(self).map_err(|e| {
crate::NeuroError::Io(std::io::Error::new(
std::io::ErrorKind::InvalidData,
e.to_string(),
))
})?;
std::fs::write(&protocol_file, json)?;
Ok(())
}
/// Load protocol from disk
pub fn load(path: impl AsRef<Path>) -> NeuroResult<Self> {
let path = path.as_ref();
let protocol_file = path.join("protocol.json");
let json = std::fs::read_to_string(&protocol_file)?;
let protocol: Self = serde_json::from_str(&json)
.map_err(|e| crate::NeuroError::InvalidFormat(e.to_string()))?;
Ok(protocol)
}
}
/// Protocol-wide settings
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProtocolSettings {
/// Default sampling frequency for new recordings
pub default_sfreq: f64,
/// Default line frequency (50 or 60 Hz)
pub line_freq: f64,
/// Use GPU acceleration if available
pub use_gpu: bool,
/// Number of parallel workers
pub n_jobs: usize,
}
impl Default for ProtocolSettings {
fn default() -> Self {
Self {
default_sfreq: 1000.0,
line_freq: 60.0,
use_gpu: true,
n_jobs: 4,
}
}
}
/// A subject (participant) in a study
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Subject {
/// Subject ID
pub id: String,
/// Subject name (optional, for display)
pub name: Option<String>,
/// Subject comments/notes
pub comments: String,
/// Anatomy data
pub anatomy: Option<Anatomy>,
/// Raw recordings
recordings: Vec<Recording>,
/// Conditions/sessions
conditions: HashMap<String, Condition>,
}
impl Subject {
/// Create a new subject
#[must_use]
pub fn new(id: impl Into<String>) -> Self {
Self {
id: id.into(),
name: None,
comments: String::new(),
anatomy: None,
recordings: Vec::new(),
conditions: HashMap::new(),
}
}
/// Add a recording
pub fn add_recording(&mut self, recording: Recording) {
self.recordings.push(recording);
}
/// Get recordings
#[must_use]
pub fn recordings(&self) -> &[Recording] {
&self.recordings
}
/// Add a condition
pub fn add_condition(&mut self, condition: Condition) {
self.conditions.insert(condition.name.clone(), condition);
}
/// Get a condition by name
#[must_use]
pub fn get_condition(&self, name: &str) -> Option<&Condition> {
self.conditions.get(name)
}
/// List condition names
#[must_use]
pub fn condition_names(&self) -> Vec<&str> {
self.conditions.keys().map(String::as_str).collect()
}
/// Set anatomy
pub fn set_anatomy(&mut self, anatomy: Anatomy) {
self.anatomy = Some(anatomy);
}
}
/// Anatomical data for a subject
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Anatomy {
/// Path to MRI volume
pub mri_path: Option<PathBuf>,
/// Path to cortical surface (left hemisphere)
pub surface_lh: Option<PathBuf>,
/// Path to cortical surface (right hemisphere)
pub surface_rh: Option<PathBuf>,
/// Path to head surface (for BEM/visualization)
pub surface_head: Option<PathBuf>,
/// Fiducial points (NAS, LPA, RPA) in MRI coordinates
pub fiducials: HashMap<String, [f64; 3]>,
/// Coordinate transform from MRI to head coordinates
pub mri_to_head: Option<[[f64; 4]; 4]>,
/// Whether anatomy is from template (vs individual MRI)
pub is_template: bool,
/// Template name if using template
pub template_name: Option<String>,
}
impl Anatomy {
/// Create a new anatomy
#[must_use]
pub fn new() -> Self {
Self {
mri_path: None,
surface_lh: None,
surface_rh: None,
surface_head: None,
fiducials: HashMap::new(),
mri_to_head: None,
is_template: false,
template_name: None,
}
}
/// Create anatomy from template
#[must_use]
pub fn from_template(name: impl Into<String>) -> Self {
Self {
is_template: true,
template_name: Some(name.into()),
..Self::new()
}
}
/// Set fiducial point
pub fn set_fiducial(&mut self, name: impl Into<String>, pos: [f64; 3]) {
self.fiducials.insert(name.into(), pos);
}
}
impl Default for Anatomy {
fn default() -> Self {
Self::new()
}
}
/// A condition/task within a subject's data
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Condition {
/// Condition name
pub name: String,
/// Description
pub description: String,
/// Associated data files
pub data_files: Vec<PathBuf>,
/// Epochs file (if computed)
pub epochs_file: Option<PathBuf>,
/// Evoked/average file (if computed)
pub evoked_file: Option<PathBuf>,
/// Source estimate files
pub source_files: Vec<PathBuf>,
}
impl Condition {
/// Create a new condition
#[must_use]
pub fn new(name: impl Into<String>) -> Self {
Self {
name: name.into(),
description: String::new(),
data_files: Vec::new(),
epochs_file: None,
evoked_file: None,
source_files: Vec::new(),
}
}
/// Add a data file
pub fn add_data_file(&mut self, path: impl Into<PathBuf>) {
self.data_files.push(path.into());
}
}
/// Group-level analysis
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GroupAnalysis {
/// Analysis name
pub name: String,
/// Description
pub description: String,
/// Subjects included
pub subjects: Vec<String>,
/// Conditions compared
pub conditions: Vec<String>,
/// Output files
pub outputs: Vec<PathBuf>,
/// Creation date
pub created: DateTime<Utc>,
}
impl GroupAnalysis {
/// Create a new group analysis
#[must_use]
pub fn new(name: impl Into<String>) -> Self {
Self {
name: name.into(),
description: String::new(),
subjects: Vec::new(),
conditions: Vec::new(),
outputs: Vec::new(),
created: Utc::now(),
}
}
/// Add subjects
pub fn add_subjects<S: Into<String>>(&mut self, subjects: impl IntoIterator<Item = S>) {
self.subjects.extend(subjects.into_iter().map(Into::into));
}
/// Add conditions
pub fn add_conditions<S: Into<String>>(&mut self, conditions: impl IntoIterator<Item = S>) {
self.conditions
.extend(conditions.into_iter().map(Into::into));
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
#[test]
fn test_protocol_creation() {
let protocol = Protocol::new("test_study", "/tmp/test");
assert_eq!(protocol.name, "test_study");
assert_eq!(protocol.n_subjects(), 0);
}
#[test]
fn test_add_subject() {
let mut protocol = Protocol::new("test", "/tmp/test");
protocol.add_subject(Subject::new("sub-01"));
protocol.add_subject(Subject::new("sub-02"));
assert_eq!(protocol.n_subjects(), 2);
assert!(protocol.get_subject("sub-01").is_some());
assert!(protocol.get_subject("sub-03").is_none());
}
#[test]
fn test_protocol_save_load() {
let temp_dir = TempDir::new().unwrap();
let path = temp_dir.path();
let mut protocol = Protocol::new("test", path);
protocol.add_subject(Subject::new("sub-01"));
protocol.save().unwrap();
let loaded = Protocol::load(path).unwrap();
assert_eq!(loaded.name, "test");
assert_eq!(loaded.n_subjects(), 1);
}
}