Initial commit

2026-03-04 00:08:42 +00:00
commit 4d88dc0584
4449 changed files with 1556714 additions and 0 deletions
@@ -0,0 +1,647 @@
+//! Learnable tissue conductivity models for head modeling.
+//!
+//! Provides models for tissue conductivity that can be learned
+//! jointly with source localization or used with fixed values.
+
+use crate::error::{PinnError, PinnResult};
+use ndarray::{Array1, Array2};
+use serde::{Deserialize, Serialize};
+
+/// Standard tissue types in the head
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
+pub enum TissueType {
+    /// Brain gray matter
+    GrayMatter,
+    /// Brain white matter
+    WhiteMatter,
+    /// Cerebrospinal fluid
+    Csf,
+    /// Skull bone
+    Skull,
+    /// Scalp/skin
+    Scalp,
+    /// Air (e.g., sinuses)
+    Air,
+    /// Eye tissue
+    Eye,
+    /// Muscle tissue
+    Muscle,
+}
+
+impl TissueType {
+    /// Default isotropic conductivity (S/m)
+    pub fn default_conductivity(&self) -> f64 {
+        match self {
+            TissueType::GrayMatter => 0.33,
+            TissueType::WhiteMatter => 0.14,
+            TissueType::Csf => 1.79,
+            TissueType::Skull => 0.0042,
+            TissueType::Scalp => 0.43,
+            TissueType::Air => 1e-12,
+            TissueType::Eye => 1.5,
+            TissueType::Muscle => 0.2,
+        }
+    }
+
+    /// Conductivity range (min, max) for learning
+    pub fn conductivity_range(&self) -> (f64, f64) {
+        match self {
+            TissueType::GrayMatter => (0.1, 0.6),
+            TissueType::WhiteMatter => (0.05, 0.3),
+            TissueType::Csf => (1.0, 2.5),
+            TissueType::Skull => (0.001, 0.02),
+            TissueType::Scalp => (0.2, 0.7),
+            TissueType::Air => (1e-15, 1e-10),
+            TissueType::Eye => (0.5, 2.5),
+            TissueType::Muscle => (0.1, 0.4),
+        }
+    }
+}
+
+/// A tissue layer with conductivity
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct TissueLayer {
+    /// Tissue type
+    pub tissue_type: TissueType,
+    /// Name/label
+    pub name: String,
+    /// Isotropic conductivity (S/m)
+    pub conductivity: f64,
+    /// Inner radius (for spherical models)
+    pub inner_radius: Option<f64>,
+    /// Outer radius (for spherical models)
+    pub outer_radius: Option<f64>,
+    /// Whether conductivity is learnable
+    pub learnable: bool,
+}
+
+impl TissueLayer {
+    /// Create a new tissue layer
+    pub fn new(tissue_type: TissueType, name: impl Into<String>) -> Self {
+        Self {
+            tissue_type,
+            name: name.into(),
+            conductivity: tissue_type.default_conductivity(),
+            inner_radius: None,
+            outer_radius: None,
+            learnable: false,
+        }
+    }
+
+    /// Set conductivity
+    pub fn with_conductivity(mut self, sigma: f64) -> Self {
+        self.conductivity = sigma;
+        self
+    }
+
+    /// Set radii for spherical model
+    pub fn with_radii(mut self, inner: f64, outer: f64) -> Self {
+        self.inner_radius = Some(inner);
+        self.outer_radius = Some(outer);
+        self
+    }
+
+    /// Make learnable
+    pub fn learnable(mut self) -> Self {
+        self.learnable = true;
+        self
+    }
+
+    /// Check if point is inside this layer (spherical model)
+    pub fn contains(&self, r: f64) -> bool {
+        match (self.inner_radius, self.outer_radius) {
+            (Some(inner), Some(outer)) => r >= inner && r < outer,
+            (Some(inner), None) => r >= inner,
+            (None, Some(outer)) => r < outer,
+            (None, None) => true,
+        }
+    }
+}
+
+/// Conductivity model trait
+pub trait ConductivityModel {
+    /// Get conductivity at a point
+    fn conductivity_at(&self, point: &[f64; 3]) -> f64;
+
+    /// Get conductivity field at multiple points
+    fn conductivity_field(&self, points: &Array2<f64>) -> Array1<f64> {
+        let n = points.nrows();
+        let mut field = Array1::zeros(n);
+        for i in 0..n {
+            field[i] = self.conductivity_at(&[points[[i, 0]], points[[i, 1]], points[[i, 2]]]);
+        }
+        field
+    }
+
+    /// Whether this model has learnable parameters
+    fn is_learnable(&self) -> bool {
+        false
+    }
+
+    /// Get learnable parameters
+    fn get_parameters(&self) -> Vec<f64> {
+        Vec::new()
+    }
+
+    /// Set learnable parameters
+    fn set_parameters(&mut self, _params: &[f64]) -> PinnResult<()> {
+        Ok(())
+    }
+}
+
+/// Homogeneous (constant) conductivity
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct HomogeneousConductivity {
+    /// Conductivity value (S/m)
+    sigma: f64,
+    /// Whether learnable
+    learnable: bool,
+}
+
+impl HomogeneousConductivity {
+    /// Create with fixed conductivity
+    pub fn new(sigma: f64) -> Self {
+        Self {
+            sigma,
+            learnable: false,
+        }
+    }
+
+    /// Create with learnable conductivity
+    pub fn learnable(initial: f64) -> Self {
+        Self {
+            sigma: initial,
+            learnable: true,
+        }
+    }
+
+    /// Get conductivity value
+    pub fn sigma(&self) -> f64 {
+        self.sigma
+    }
+}
+
+impl ConductivityModel for HomogeneousConductivity {
+    fn conductivity_at(&self, _point: &[f64; 3]) -> f64 {
+        self.sigma
+    }
+
+    fn is_learnable(&self) -> bool {
+        self.learnable
+    }
+
+    fn get_parameters(&self) -> Vec<f64> {
+        if self.learnable {
+            vec![self.sigma]
+        } else {
+            Vec::new()
+        }
+    }
+
+    fn set_parameters(&mut self, params: &[f64]) -> PinnResult<()> {
+        if self.learnable {
+            if params.is_empty() {
+                return Err(PinnError::DimensionMismatch(
+                    "Expected 1 parameter for learnable conductivity".into(),
+                ));
+            }
+            self.sigma = params[0].max(1e-6); // Ensure positive
+        }
+        Ok(())
+    }
+}
+
+/// Layered spherical conductivity model
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct LayeredConductivity {
+    /// Tissue layers (ordered from inside to outside)
+    layers: Vec<TissueLayer>,
+    /// Center of the spherical model
+    center: [f64; 3],
+}
+
+impl LayeredConductivity {
+    /// Create a new layered model
+    pub fn new(center: [f64; 3]) -> Self {
+        Self {
+            layers: Vec::new(),
+            center,
+        }
+    }
+
+    /// Add a layer
+    pub fn add_layer(&mut self, layer: TissueLayer) {
+        self.layers.push(layer);
+        // Sort by outer radius
+        self.layers.sort_by(|a, b| {
+            let ra = a.outer_radius.unwrap_or(f64::MAX);
+            let rb = b.outer_radius.unwrap_or(f64::MAX);
+            ra.partial_cmp(&rb).unwrap()
+        });
+    }
+
+    /// Create standard 3-layer head model
+    pub fn three_layer(
+        center: [f64; 3],
+        brain_radius: f64,
+        skull_thickness: f64,
+        scalp_thickness: f64,
+    ) -> Self {
+        let skull_outer = brain_radius + skull_thickness;
+        let scalp_outer = skull_outer + scalp_thickness;
+
+        let mut model = Self::new(center);
+
+        model.add_layer(
+            TissueLayer::new(TissueType::GrayMatter, "brain").with_radii(0.0, brain_radius),
+        );
+
+        model.add_layer(
+            TissueLayer::new(TissueType::Skull, "skull").with_radii(brain_radius, skull_outer),
+        );
+
+        model.add_layer(
+            TissueLayer::new(TissueType::Scalp, "scalp").with_radii(skull_outer, scalp_outer),
+        );
+
+        model
+    }
+
+    /// Create standard 4-layer head model (with CSF)
+    pub fn four_layer(
+        center: [f64; 3],
+        brain_radius: f64,
+        csf_thickness: f64,
+        skull_thickness: f64,
+        scalp_thickness: f64,
+    ) -> Self {
+        let csf_outer = brain_radius + csf_thickness;
+        let skull_outer = csf_outer + skull_thickness;
+        let scalp_outer = skull_outer + scalp_thickness;
+
+        let mut model = Self::new(center);
+
+        model.add_layer(
+            TissueLayer::new(TissueType::GrayMatter, "brain").with_radii(0.0, brain_radius),
+        );
+
+        model.add_layer(
+            TissueLayer::new(TissueType::Csf, "csf").with_radii(brain_radius, csf_outer),
+        );
+
+        model.add_layer(
+            TissueLayer::new(TissueType::Skull, "skull").with_radii(csf_outer, skull_outer),
+        );
+
+        model.add_layer(
+            TissueLayer::new(TissueType::Scalp, "scalp").with_radii(skull_outer, scalp_outer),
+        );
+
+        model
+    }
+
+    /// Get layer containing a point
+    pub fn layer_at(&self, point: &[f64; 3]) -> Option<&TissueLayer> {
+        let dx = point[0] - self.center[0];
+        let dy = point[1] - self.center[1];
+        let dz = point[2] - self.center[2];
+        let r = (dx * dx + dy * dy + dz * dz).sqrt();
+
+        self.layers.iter().find(|&layer| layer.contains(r)).map(|v| v as _)
+    }
+
+    /// Get all layers
+    pub fn layers(&self) -> &[TissueLayer] {
+        &self.layers
+    }
+
+    /// Get mutable layers
+    pub fn layers_mut(&mut self) -> &mut [TissueLayer] {
+        &mut self.layers
+    }
+
+    /// Number of layers
+    pub fn n_layers(&self) -> usize {
+        self.layers.len()
+    }
+}
+
+impl ConductivityModel for LayeredConductivity {
+    fn conductivity_at(&self, point: &[f64; 3]) -> f64 {
+        self.layer_at(point).map_or(0.0, |l| l.conductivity)
+    }
+
+    fn is_learnable(&self) -> bool {
+        self.layers.iter().any(|l| l.learnable)
+    }
+
+    fn get_parameters(&self) -> Vec<f64> {
+        self.layers
+            .iter()
+            .filter(|l| l.learnable)
+            .map(|l| l.conductivity)
+            .collect()
+    }
+
+    fn set_parameters(&mut self, params: &[f64]) -> PinnResult<()> {
+        let learnable_count = self.layers.iter().filter(|l| l.learnable).count();
+        if params.len() != learnable_count {
+            return Err(PinnError::DimensionMismatch(format!(
+                "Expected {} parameters, got {}",
+                learnable_count,
+                params.len()
+            )));
+        }
+
+        let mut idx = 0;
+        for layer in &mut self.layers {
+            if layer.learnable {
+                let (min, max) = layer.tissue_type.conductivity_range();
+                layer.conductivity = params[idx].clamp(min, max);
+                idx += 1;
+            }
+        }
+
+        Ok(())
+    }
+}
+
+/// Neural network-based learnable conductivity
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct LearnableConductivity {
+    /// Base layered model
+    base: LayeredConductivity,
+    /// Perturbation weights [n_basis, n_layers]
+    perturbation_weights: Array2<f64>,
+    /// Radial basis function centers
+    rbf_centers: Array1<f64>,
+    /// RBF width
+    rbf_width: f64,
+    /// Maximum perturbation magnitude
+    max_perturbation: f64,
+}
+
+impl LearnableConductivity {
+    /// Create learnable conductivity from base model
+    pub fn new(base: LayeredConductivity, n_basis: usize) -> Self {
+        let n_layers = base.n_layers();
+
+        // Distribute RBF centers radially
+        let max_radius = base
+            .layers()
+            .iter()
+            .filter_map(|l| l.outer_radius)
+            .fold(0.0f64, f64::max);
+
+        let mut rbf_centers = Array1::zeros(n_basis);
+        for i in 0..n_basis {
+            rbf_centers[i] = max_radius * (i as f64 + 0.5) / n_basis as f64;
+        }
+
+        Self {
+            base,
+            perturbation_weights: Array2::zeros((n_basis, n_layers)),
+            rbf_centers,
+            rbf_width: max_radius / (n_basis as f64 * 2.0),
+            max_perturbation: 0.1, // 10% max perturbation
+        }
+    }
+
+    /// Compute RBF values at radius
+    fn rbf_values(&self, r: f64) -> Array1<f64> {
+        self.rbf_centers.mapv(|c| {
+            let d = (r - c) / self.rbf_width;
+            (-0.5 * d * d).exp()
+        })
+    }
+
+    /// Get perturbation at radius for each layer
+    fn perturbation(&self, r: f64) -> Array1<f64> {
+        let rbf = self.rbf_values(r);
+        let mut perturb: Array1<f64> = Array1::zeros(self.base.n_layers());
+
+        for j in 0..self.base.n_layers() {
+            for i in 0..self.rbf_centers.len() {
+                perturb[j] += self.perturbation_weights[[i, j]] * rbf[i];
+            }
+            // Clamp perturbation
+            perturb[j] = perturb[j].tanh() * self.max_perturbation;
+        }
+
+        perturb
+    }
+
+    /// Get base model
+    pub fn base(&self) -> &LayeredConductivity {
+        &self.base
+    }
+}
+
+impl ConductivityModel for LearnableConductivity {
+    fn conductivity_at(&self, point: &[f64; 3]) -> f64 {
+        let dx = point[0] - self.base.center[0];
+        let dy = point[1] - self.base.center[1];
+        let dz = point[2] - self.base.center[2];
+        let r = (dx * dx + dy * dy + dz * dz).sqrt();
+
+        // Find which layer and apply perturbation
+        let perturb = self.perturbation(r);
+
+        for (i, layer) in self.base.layers().iter().enumerate() {
+            if layer.contains(r) {
+                let (min, max) = layer.tissue_type.conductivity_range();
+                let sigma = layer.conductivity * (1.0 + perturb[i]);
+                return sigma.clamp(min, max);
+            }
+        }
+
+        0.0
+    }
+
+    fn is_learnable(&self) -> bool {
+        true
+    }
+
+    fn get_parameters(&self) -> Vec<f64> {
+        self.perturbation_weights.iter().copied().collect()
+    }
+
+    fn set_parameters(&mut self, params: &[f64]) -> PinnResult<()> {
+        if params.len() != self.perturbation_weights.len() {
+            return Err(PinnError::DimensionMismatch(format!(
+                "Expected {} parameters, got {}",
+                self.perturbation_weights.len(),
+                params.len()
+            )));
+        }
+
+        for (i, &p) in params.iter().enumerate() {
+            let (row, col) = (i / self.base.n_layers(), i % self.base.n_layers());
+            self.perturbation_weights[[row, col]] = p;
+        }
+
+        Ok(())
+    }
+}
+
+/// Anisotropic conductivity tensor
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct AnisotropicConductivity {
+    /// Conductivity tensor [3, 3] at each layer
+    tensors: Vec<Array2<f64>>,
+    /// Base layered model for geometry
+    geometry: LayeredConductivity,
+}
+
+impl AnisotropicConductivity {
+    /// Create from isotropic layered model
+    pub fn from_isotropic(geometry: LayeredConductivity) -> Self {
+        let tensors = geometry
+            .layers()
+            .iter()
+            .map(|l| {
+                let sigma = l.conductivity;
+                Array2::from_diag(&ndarray::arr1(&[sigma, sigma, sigma]))
+            })
+            .collect();
+
+        Self { tensors, geometry }
+    }
+
+    /// Set anisotropic tensor for a layer
+    pub fn set_tensor(&mut self, layer_idx: usize, tensor: Array2<f64>) -> PinnResult<()> {
+        if layer_idx >= self.tensors.len() {
+            return Err(PinnError::InvalidConfig(format!(
+                "Layer index {} out of range",
+                layer_idx
+            )));
+        }
+        if tensor.shape() != [3, 3] {
+            return Err(PinnError::DimensionMismatch("Tensor must be 3x3".into()));
+        }
+        self.tensors[layer_idx] = tensor;
+        Ok(())
+    }
+
+    /// Get conductivity tensor at a point
+    pub fn tensor_at(&self, point: &[f64; 3]) -> Array2<f64> {
+        let dx = point[0] - self.geometry.center[0];
+        let dy = point[1] - self.geometry.center[1];
+        let dz = point[2] - self.geometry.center[2];
+        let r = (dx * dx + dy * dy + dz * dz).sqrt();
+
+        for (i, layer) in self.geometry.layers().iter().enumerate() {
+            if layer.contains(r) {
+                return self.tensors[i].clone();
+            }
+        }
+
+        Array2::zeros((3, 3))
+    }
+
+    /// Compute σ∇Φ where σ is tensor
+    pub fn apply_to_gradient(&self, point: &[f64; 3], gradient: &[f64; 3]) -> [f64; 3] {
+        let tensor = self.tensor_at(point);
+        let grad = ndarray::arr1(gradient);
+        let result = tensor.dot(&grad);
+        [result[0], result[1], result[2]]
+    }
+}
+
+impl ConductivityModel for AnisotropicConductivity {
+    fn conductivity_at(&self, point: &[f64; 3]) -> f64 {
+        // Return trace / 3 as effective isotropic conductivity
+        let tensor = self.tensor_at(point);
+        (tensor[[0, 0]] + tensor[[1, 1]] + tensor[[2, 2]]) / 3.0
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_tissue_types() {
+        assert!((TissueType::GrayMatter.default_conductivity() - 0.33).abs() < 1e-10);
+        assert!((TissueType::Skull.default_conductivity() - 0.0042).abs() < 1e-10);
+    }
+
+    #[test]
+    fn test_homogeneous() {
+        let model = HomogeneousConductivity::new(0.33);
+        assert!((model.conductivity_at(&[0.0, 0.0, 0.0]) - 0.33).abs() < 1e-10);
+        assert!((model.conductivity_at(&[1.0, 2.0, 3.0]) - 0.33).abs() < 1e-10);
+    }
+
+    #[test]
+    fn test_layered_conductivity() {
+        let model = LayeredConductivity::three_layer(
+            [0.0, 0.0, 0.0],
+            0.08,  // brain radius
+            0.007, // skull thickness
+            0.006, // scalp thickness
+        );
+
+        assert_eq!(model.n_layers(), 3);
+
+        // Inside brain
+        let sigma = model.conductivity_at(&[0.0, 0.0, 0.05]);
+        assert!((sigma - 0.33).abs() < 1e-10);
+
+        // Inside skull
+        let sigma = model.conductivity_at(&[0.0, 0.0, 0.082]);
+        assert!((sigma - 0.0042).abs() < 1e-10);
+
+        // Inside scalp
+        let sigma = model.conductivity_at(&[0.0, 0.0, 0.09]);
+        assert!((sigma - 0.43).abs() < 1e-10);
+    }
+
+    #[test]
+    fn test_four_layer_model() {
+        let model = LayeredConductivity::four_layer(
+            [0.0, 0.0, 0.0],
+            0.078, // brain
+            0.002, // CSF
+            0.007, // skull
+            0.006, // scalp
+        );
+
+        assert_eq!(model.n_layers(), 4);
+
+        // CSF layer
+        let sigma = model.conductivity_at(&[0.0, 0.0, 0.079]);
+        assert!((sigma - 1.79).abs() < 1e-10);
+    }
+
+    #[test]
+    fn test_learnable_conductivity() {
+        let base = LayeredConductivity::three_layer([0.0, 0.0, 0.0], 0.08, 0.007, 0.006);
+
+        let model = LearnableConductivity::new(base, 5);
+
+        // Should return values close to base model
+        let sigma = model.conductivity_at(&[0.0, 0.0, 0.05]);
+        assert!((sigma - 0.33).abs() < 0.05);
+
+        // Has learnable parameters
+        assert!(model.is_learnable());
+        assert!(model.get_parameters().len() > 0);
+    }
+
+    #[test]
+    fn test_anisotropic_conductivity() {
+        let geometry = LayeredConductivity::three_layer([0.0, 0.0, 0.0], 0.08, 0.007, 0.006);
+
+        let model = AnisotropicConductivity::from_isotropic(geometry);
+
+        // Effective conductivity should equal isotropic value
+        let sigma = model.conductivity_at(&[0.0, 0.0, 0.05]);
+        assert!((sigma - 0.33).abs() < 1e-10);
+
+        // Tensor should be diagonal
+        let tensor = model.tensor_at(&[0.0, 0.0, 0.05]);
+        assert!((tensor[[0, 0]] - 0.33).abs() < 1e-10);
+        assert!(tensor[[0, 1]].abs() < 1e-10);
+    }
+}