Initial commit

crates/specialized/rtx-neuro-fem/src/leadfield.rs

@@ -0,0 +1,693 @@
+//! Lead field (gain matrix) computation for MEG/EEG forward modeling.
+//!
+//! The lead field matrix relates neural source currents to measured
+//! sensor signals: M = L * J, where L is the lead field.
+
+use crate::assembly::{FemAssembler, TransferMatrix};
+use crate::error::{FemError, FemResult};
+use crate::mesh::HeadMesh;
+use crate::solver::{FemSolver, SolverConfig};
+use nalgebra::Vector3;
+use ndarray::{Array1, Array2, Axis};
+use rayon::prelude::*;
+use serde::{Deserialize, Serialize};
+use std::time::Instant;
+
+/// Source type for lead field computation
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+pub enum SourceType {
+    /// Current dipole
+    Dipole,
+    /// Monopole (point source)
+    Monopole,
+}
+
+/// Lead field configuration
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct LeadFieldConfig {
+    /// Source type
+    pub source_type: SourceType,
+    /// Number of orientations per source (1 for fixed, 3 for free)
+    pub n_orientations: usize,
+    /// Whether to compute MEG lead field
+    pub compute_meg: bool,
+    /// Whether to compute EEG lead field
+    pub compute_eeg: bool,
+    /// Apply average reference to EEG
+    pub eeg_average_reference: bool,
+    /// Solver configuration
+    pub solver_config: SolverConfig,
+    /// Parallel computation
+    pub parallel: bool,
+    /// Verbose output
+    pub verbose: bool,
+}
+
+impl Default for LeadFieldConfig {
+    fn default() -> Self {
+        Self {
+            source_type: SourceType::Dipole,
+            n_orientations: 3,
+            compute_meg: false,
+            compute_eeg: true,
+            eeg_average_reference: true,
+            solver_config: SolverConfig::default(),
+            parallel: true,
+            verbose: false,
+        }
+    }
+}
+
+/// Source space definition
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct SourceSpace {
+    /// Source positions (n_sources x 3)
+    pub positions: Array2<f64>,
+    /// Source normals/orientations (n_sources x 3, optional)
+    pub normals: Option<Array2<f64>>,
+    /// Element indices containing each source
+    pub element_indices: Vec<Option<usize>>,
+}
+
+impl SourceSpace {
+    /// Create source space from positions
+    pub fn from_positions(positions: Array2<f64>) -> Self {
+        let n_sources = positions.nrows();
+        Self {
+            positions,
+            normals: None,
+            element_indices: vec![None; n_sources],
+        }
+    }
+
+    /// Create source space with normals
+    pub fn with_normals(positions: Array2<f64>, normals: Array2<f64>) -> FemResult<Self> {
+        if positions.nrows() != normals.nrows() {
+            return Err(FemError::DimensionMismatch(
+                "Positions and normals must have same number of rows".into(),
+            ));
+        }
+
+        let n_sources = positions.nrows();
+        Ok(Self {
+            positions,
+            normals: Some(normals),
+            element_indices: vec![None; n_sources],
+        })
+    }
+
+    /// Create regular grid source space
+    pub fn regular_grid(center: &Vector3<f64>, extent: &Vector3<f64>, spacing: f64) -> Self {
+        let nx = (extent.x / spacing).ceil() as usize;
+        let ny = (extent.y / spacing).ceil() as usize;
+        let nz = (extent.z / spacing).ceil() as usize;
+
+        let n_sources = nx * ny * nz;
+        let mut positions = Array2::zeros((n_sources, 3));
+
+        let mut idx = 0;
+        for i in 0..nx {
+            for j in 0..ny {
+                for k in 0..nz {
+                    let x = center.x - extent.x / 2.0 + i as f64 * spacing;
+                    let y = center.y - extent.y / 2.0 + j as f64 * spacing;
+                    let z = center.z - extent.z / 2.0 + k as f64 * spacing;
+
+                    positions[[idx, 0]] = x;
+                    positions[[idx, 1]] = y;
+                    positions[[idx, 2]] = z;
+                    idx += 1;
+                }
+            }
+        }
+
+        Self::from_positions(positions)
+    }
+
+    /// Locate sources in mesh elements
+    pub fn locate_in_mesh(&mut self, mesh: &HeadMesh) {
+        for (i, row) in self.positions.axis_iter(Axis(0)).enumerate() {
+            let pos = Vector3::new(row[0], row[1], row[2]);
+            self.element_indices[i] = mesh.find_element(&pos);
+        }
+    }
+
+    /// Number of sources
+    pub fn n_sources(&self) -> usize {
+        self.positions.nrows()
+    }
+
+    /// Get source position
+    pub fn position(&self, idx: usize) -> Vector3<f64> {
+        Vector3::new(
+            self.positions[[idx, 0]],
+            self.positions[[idx, 1]],
+            self.positions[[idx, 2]],
+        )
+    }
+
+    /// Get source normal (or default)
+    pub fn normal(&self, idx: usize) -> Vector3<f64> {
+        if let Some(ref normals) = self.normals {
+            Vector3::new(normals[[idx, 0]], normals[[idx, 1]], normals[[idx, 2]])
+        } else {
+            Vector3::new(0.0, 0.0, 1.0)
+        }
+    }
+}
+
+/// EEG/MEG sensor positions
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct SensorArray {
+    /// Sensor positions (n_sensors x 3)
+    pub positions: Array2<f64>,
+    /// Sensor orientations for MEG (n_sensors x 3)
+    pub orientations: Option<Array2<f64>>,
+    /// Sensor labels
+    pub labels: Vec<String>,
+    /// Sensor type
+    pub sensor_type: SensorType,
+}
+
+/// Sensor type
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+pub enum SensorType {
+    /// EEG electrodes
+    EEG,
+    /// MEG magnetometers
+    Magnetometer,
+    /// MEG gradiometers
+    Gradiometer,
+}
+
+impl SensorArray {
+    /// Create EEG sensor array
+    pub fn eeg(positions: Array2<f64>, labels: Vec<String>) -> Self {
+        Self {
+            positions,
+            orientations: None,
+            labels,
+            sensor_type: SensorType::EEG,
+        }
+    }
+
+    /// Create MEG sensor array
+    pub fn meg(
+        positions: Array2<f64>,
+        orientations: Array2<f64>,
+        labels: Vec<String>,
+        sensor_type: SensorType,
+    ) -> Self {
+        Self {
+            positions,
+            orientations: Some(orientations),
+            labels,
+            sensor_type,
+        }
+    }
+
+    /// Number of sensors
+    pub fn n_sensors(&self) -> usize {
+        self.positions.nrows()
+    }
+
+    /// Get sensor positions as vectors
+    pub fn position_vectors(&self) -> Vec<Vector3<f64>> {
+        (0..self.n_sensors())
+            .map(|i| {
+                Vector3::new(
+                    self.positions[[i, 0]],
+                    self.positions[[i, 1]],
+                    self.positions[[i, 2]],
+                )
+            })
+            .collect()
+    }
+}
+
+/// Computed lead field matrix
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct LeadField {
+    /// Lead field matrix (n_sensors x n_sources * n_orientations)
+    pub matrix: Array2<f64>,
+    /// Number of sensors
+    pub n_sensors: usize,
+    /// Number of sources
+    pub n_sources: usize,
+    /// Number of orientations per source
+    pub n_orientations: usize,
+    /// Source positions
+    pub source_positions: Array2<f64>,
+    /// Sensor positions
+    pub sensor_positions: Array2<f64>,
+    /// Computation time (seconds)
+    pub computation_time: f64,
+    /// Sources inside mesh
+    pub sources_in_mesh: usize,
+}
+
+impl LeadField {
+    /// Get lead field for a single source and orientation
+    pub fn get_column(&self, source_idx: usize, orientation: usize) -> Array1<f64> {
+        let col_idx = source_idx * self.n_orientations + orientation;
+        self.matrix.column(col_idx).to_owned()
+    }
+
+    /// Get lead field for all orientations of a source
+    pub fn get_source(&self, source_idx: usize) -> Array2<f64> {
+        let start = source_idx * self.n_orientations;
+        let end = start + self.n_orientations;
+        self.matrix.slice(ndarray::s![.., start..end]).to_owned()
+    }
+
+    /// Apply lead field: M = L * J
+    pub fn apply(&self, source_currents: &Array1<f64>) -> FemResult<Array1<f64>> {
+        let expected_len = self.n_sources * self.n_orientations;
+        if source_currents.len() != expected_len {
+            return Err(FemError::DimensionMismatch(format!(
+                "Source currents length {} doesn't match expected {}",
+                source_currents.len(),
+                expected_len
+            )));
+        }
+
+        Ok(self.matrix.dot(source_currents))
+    }
+}
+
+/// Lead field computer
+pub struct LeadFieldComputer<'a> {
+    /// Configuration
+    config: LeadFieldConfig,
+    /// FEM assembler
+    assembler: &'a FemAssembler,
+    /// FEM solver
+    solver: FemSolver,
+    /// Transfer matrix for EEG
+    eeg_transfer: Option<TransferMatrix>,
+}
+
+impl<'a> LeadFieldComputer<'a> {
+    /// Create lead field computer
+    pub fn new(assembler: &'a FemAssembler, config: LeadFieldConfig) -> FemResult<Self> {
+        let solver = FemSolver::from_assembler(assembler, config.solver_config.clone())?;
+
+        Ok(Self {
+            config,
+            assembler,
+            solver,
+            eeg_transfer: None,
+        })
+    }
+
+    /// Set EEG electrode positions
+    pub fn set_eeg_electrodes(&mut self, electrodes: &[Vector3<f64>]) -> FemResult<()> {
+        self.eeg_transfer = Some(TransferMatrix::compute(self.assembler.mesh(), electrodes)?);
+        Ok(())
+    }
+
+    /// Compute EEG lead field
+    pub fn compute_eeg(
+        &self,
+        sources: &SourceSpace,
+        sensors: &SensorArray,
+    ) -> FemResult<LeadField> {
+        let start_time = Instant::now();
+
+        let n_sensors = sensors.n_sensors();
+        let n_sources = sources.n_sources();
+        let n_orient = self.config.n_orientations;
+        let n_cols = n_sources * n_orient;
+
+        if self.config.verbose {
+            eprintln!(
+                "Computing EEG lead field: {} sensors, {} sources, {} orientations",
+                n_sensors, n_sources, n_orient
+            );
+        }
+
+        // Compute transfer matrix if not already set
+        let transfer = if let Some(ref t) = self.eeg_transfer {
+            t.clone()
+        } else {
+            TransferMatrix::compute(self.assembler.mesh(), &sensors.position_vectors())?
+        };
+
+        // Compute lead field columns
+        let mut leadfield = Array2::zeros((n_sensors, n_cols));
+        let mut sources_in_mesh = 0;
+
+        // Define orientation vectors
+        let orientations: Vec<Vector3<f64>> = if n_orient == 1 {
+            vec![Vector3::new(0.0, 0.0, 1.0)]
+        } else {
+            vec![
+                Vector3::new(1.0, 0.0, 0.0),
+                Vector3::new(0.0, 1.0, 0.0),
+                Vector3::new(0.0, 0.0, 1.0),
+            ]
+        };
+
+        if self.config.parallel {
+            // Parallel computation over sources
+            let results: Vec<(usize, Vec<Array1<f64>>)> = (0..n_sources)
+                .into_par_iter()
+                .filter_map(|src_idx| {
+                    let pos = sources.position(src_idx);
+
+                    // Find element containing source
+                    let elem_idx = match sources.element_indices.get(src_idx) {
+                        Some(Some(idx)) => *idx,
+                        _ => self.assembler.mesh().find_element(&pos)?,
+                    };
+
+                    // Compute for each orientation
+                    let mut cols = Vec::with_capacity(n_orient);
+
+                    for orient in &orientations {
+                        // Create dipole RHS
+                        let rhs = self.create_dipole_rhs(src_idx, elem_idx, orient);
+
+                        // Solve
+                        if let Ok(result) = self.solver.solve(&rhs) {
+                            let nodal_potential = Array1::from_vec(result.solution);
+                            let sensor_potential = transfer.apply(&nodal_potential);
+                            cols.push(sensor_potential);
+                        } else {
+                            cols.push(Array1::zeros(n_sensors));
+                        }
+                    }
+
+                    Some((src_idx, cols))
+                })
+                .collect();
+
+            // Assemble results
+            for (src_idx, cols) in results {
+                sources_in_mesh += 1;
+                for (orient_idx, col) in cols.iter().enumerate() {
+                    let col_idx = src_idx * n_orient + orient_idx;
+                    for (row_idx, &val) in col.iter().enumerate() {
+                        leadfield[[row_idx, col_idx]] = val;
+                    }
+                }
+            }
+        } else {
+            // Sequential computation
+            for src_idx in 0..n_sources {
+                let pos = sources.position(src_idx);
+
+                // Find element
+                let elem_idx = match sources.element_indices.get(src_idx) {
+                    Some(Some(idx)) => *idx,
+                    _ => {
+                        if let Some(idx) = self.assembler.mesh().find_element(&pos) {
+                            idx
+                        } else {
+                            continue;
+                        }
+                    }
+                };
+
+                sources_in_mesh += 1;
+
+                for (orient_idx, orient) in orientations.iter().enumerate() {
+                    let rhs = self.create_dipole_rhs(src_idx, elem_idx, orient);
+
+                    if let Ok(result) = self.solver.solve(&rhs) {
+                        let nodal_potential = Array1::from_vec(result.solution);
+                        let sensor_potential = transfer.apply(&nodal_potential);
+
+                        let col_idx = src_idx * n_orient + orient_idx;
+                        for (row_idx, &val) in sensor_potential.iter().enumerate() {
+                            leadfield[[row_idx, col_idx]] = val;
+                        }
+                    }
+                }
+
+                if self.config.verbose && src_idx % 100 == 0 {
+                    eprintln!("Computed source {}/{}", src_idx + 1, n_sources);
+                }
+            }
+        }
+
+        // Apply average reference
+        if self.config.eeg_average_reference {
+            leadfield = self.apply_average_reference(leadfield);
+        }
+
+        let computation_time = start_time.elapsed().as_secs_f64();
+
+        if self.config.verbose {
+            eprintln!(
+                "Lead field computed in {:.2}s ({} sources in mesh)",
+                computation_time, sources_in_mesh
+            );
+        }
+
+        Ok(LeadField {
+            matrix: leadfield,
+            n_sensors,
+            n_sources,
+            n_orientations: n_orient,
+            source_positions: sources.positions.clone(),
+            sensor_positions: sensors.positions.clone(),
+            computation_time,
+            sources_in_mesh,
+        })
+    }
+
+    /// Create RHS vector for a dipole at given position with given orientation
+    fn create_dipole_rhs(
+        &self,
+        _source_idx: usize,
+        element_idx: usize,
+        orientation: &Vector3<f64>,
+    ) -> Array1<f64> {
+        let n_nodes = self.assembler.n_nodes();
+        let mut rhs = Array1::zeros(n_nodes);
+
+        let mesh = self.assembler.mesh();
+        let elem = &mesh.elements[element_idx];
+        let grads = elem.shape_gradients(&mesh.nodes);
+
+        // Dipole source: ∫ J_p · ∇φ dV = J_p · ∇φ * V (constant for linear tet)
+        // RHS contribution to node i: q * orientation · grad_i * volume
+        let q = 1.0; // Unit dipole moment
+
+        for (local_idx, &global_idx) in elem.nodes.iter().enumerate() {
+            let grad = &grads[local_idx];
+            rhs[global_idx] = q * orientation.dot(grad) * elem.volume;
+        }
+
+        rhs
+    }
+
+    /// Apply average reference to lead field
+    fn apply_average_reference(&self, mut leadfield: Array2<f64>) -> Array2<f64> {
+        let n_sensors = leadfield.nrows();
+
+        if n_sensors == 0 {
+            return leadfield;
+        }
+
+        // Subtract mean across sensors for each column
+        for mut col in leadfield.axis_iter_mut(Axis(1)) {
+            let mean = col.sum() / n_sensors as f64;
+            col -= mean;
+        }
+
+        leadfield
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::assembly::FemAssembler;
+    use crate::conductivity::TissueConductivity;
+    use crate::mesh::HeadMesh;
+
+    fn create_test_assembler() -> FemAssembler {
+        let mesh = HeadMesh::three_layer_sphere(0.08, 0.007, 0.006, 2, 1).unwrap();
+        let conductivity = TissueConductivity::default_isotropic();
+        let mut assembler = FemAssembler::with_defaults(mesh, conductivity);
+        assembler.assemble_global().unwrap();
+        assembler
+    }
+
+    #[test]
+    fn test_source_space_creation() {
+        let mut positions = Array2::zeros((10, 3));
+        for i in 0..10 {
+            positions[[i, 0]] = i as f64 * 0.01;
+            positions[[i, 1]] = 0.0;
+            positions[[i, 2]] = 0.0;
+        }
+
+        let sources = SourceSpace::from_positions(positions);
+        assert_eq!(sources.n_sources(), 10);
+    }
+
+    #[test]
+    fn test_source_space_grid() {
+        let center = Vector3::new(0.0, 0.0, 0.0);
+        let extent = Vector3::new(0.04, 0.04, 0.04);
+        let spacing = 0.01;
+
+        let sources = SourceSpace::regular_grid(&center, &extent, spacing);
+        assert!(sources.n_sources() > 0);
+    }
+
+    #[test]
+    fn test_sensor_array() {
+        let mut positions = Array2::zeros((19, 3));
+        let labels: Vec<String> = (0..19).map(|i| format!("E{}", i)).collect();
+
+        let sensors = SensorArray::eeg(positions, labels);
+        assert_eq!(sensors.n_sensors(), 19);
+        assert_eq!(sensors.sensor_type, SensorType::EEG);
+    }
+
+    #[test]
+    fn test_leadfield_computer_creation() {
+        let assembler = create_test_assembler();
+        let config = LeadFieldConfig::default();
+
+        let computer = LeadFieldComputer::new(&assembler, config);
+        assert!(computer.is_ok());
+    }
+
+    #[test]
+    fn test_create_dipole_rhs() {
+        let assembler = create_test_assembler();
+        let config = LeadFieldConfig::default();
+        let computer = LeadFieldComputer::new(&assembler, config).unwrap();
+
+        let orientation = Vector3::new(0.0, 0.0, 1.0);
+        let rhs = computer.create_dipole_rhs(0, 0, &orientation);
+
+        assert_eq!(rhs.len(), assembler.n_nodes());
+
+        // RHS should be sparse (only 4 nodes per element)
+        let nonzeros = rhs.iter().filter(|&&v| v.abs() > 1e-15).count();
+        assert!(nonzeros <= 4);
+    }
+
+    #[test]
+    fn test_source_locate_in_mesh() {
+        let mesh = HeadMesh::three_layer_sphere(0.08, 0.007, 0.006, 2, 1).unwrap();
+
+        let mut positions = Array2::zeros((5, 3));
+        // Points inside mesh
+        positions[[0, 2]] = 0.02; // Center-ish
+        positions[[1, 0]] = 0.04;
+        positions[[2, 1]] = 0.03;
+        // Point outside
+        positions[[3, 2]] = 0.5;
+        // Another inside
+        positions[[4, 2]] = 0.05;
+
+        let mut sources = SourceSpace::from_positions(positions);
+        sources.locate_in_mesh(&mesh);
+
+        // Some should be found
+        let found = sources
+            .element_indices
+            .iter()
+            .filter(|x| x.is_some())
+            .count();
+        assert!(found >= 2);
+    }
+
+    #[test]
+    fn test_average_reference() {
+        let assembler = create_test_assembler();
+        let config = LeadFieldConfig {
+            eeg_average_reference: true,
+            ..Default::default()
+        };
+        let computer = LeadFieldComputer::new(&assembler, config).unwrap();
+
+        // Create test matrix
+        let mut matrix = Array2::zeros((10, 5));
+        for i in 0..10 {
+            for j in 0..5 {
+                matrix[[i, j]] = (i + j) as f64;
+            }
+        }
+
+        let ref_matrix = computer.apply_average_reference(matrix);
+
+        // Each column should have zero mean
+        for j in 0..5 {
+            let col_sum: f64 = ref_matrix.column(j).sum();
+            assert!(col_sum.abs() < 1e-10, "Column {} sum = {}", j, col_sum);
+        }
+    }
+
+    #[test]
+    fn test_leadfield_apply() {
+        // Create a simple lead field
+        let matrix = Array2::from_shape_vec(
+            (3, 6),
+            vec![
+                1.0, 0.0, 0.0, 2.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 2.0, 0.0, 0.0, 0.0, 1.0, 0.0,
+                0.0, 2.0,
+            ],
+        )
+        .unwrap();
+
+        let leadfield = LeadField {
+            matrix,
+            n_sensors: 3,
+            n_sources: 2,
+            n_orientations: 3,
+            source_positions: Array2::zeros((2, 3)),
+            sensor_positions: Array2::zeros((3, 3)),
+            computation_time: 0.0,
+            sources_in_mesh: 2,
+        };
+
+        // Apply to source currents
+        let currents = Array1::from_vec(vec![1.0, 0.0, 0.0, 0.0, 1.0, 0.0]);
+        let sensors = leadfield.apply(&currents).unwrap();
+
+        assert_eq!(sensors.len(), 3);
+        assert!((sensors[0] - 1.0).abs() < 1e-10);
+        assert!((sensors[1] - 2.0).abs() < 1e-10);
+        assert!((sensors[2] - 0.0).abs() < 1e-10);
+    }
+
+    #[test]
+    fn test_leadfield_get_source() {
+        let matrix = Array2::from_shape_vec(
+            (2, 6),
+            vec![
+                1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0,
+            ],
+        )
+        .unwrap();
+
+        let leadfield = LeadField {
+            matrix,
+            n_sensors: 2,
+            n_sources: 2,
+            n_orientations: 3,
+            source_positions: Array2::zeros((2, 3)),
+            sensor_positions: Array2::zeros((2, 3)),
+            computation_time: 0.0,
+            sources_in_mesh: 2,
+        };
+
+        // Get first source
+        let src0 = leadfield.get_source(0);
+        assert_eq!(src0.dim(), (2, 3));
+        assert!((src0[[0, 0]] - 1.0).abs() < 1e-10);
+        assert!((src0[[0, 2]] - 3.0).abs() < 1e-10);
+
+        // Get second source
+        let src1 = leadfield.get_source(1);
+        assert!((src1[[0, 0]] - 4.0).abs() < 1e-10);
+    }
+}