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

534 lines
16 KiB
Rust

//! MRE-to-mesh mapping utilities.
//!
//! This module provides functionality to map MRE (Magnetic Resonance Elastography)
//! property fields onto finite element mesh nodes.
use crate::error::{Mri2FeError, Result};
use nalgebra::{DVector, Point3};
use rayon::prelude::*;
use rtx_medical_io::Volume;
use rtx_mesh_gen::TetrahedralMesh;
use rtx_segmentation::prelude::*;
use serde::{Deserialize, Serialize};
/// Configuration for MRE mapping.
#[derive(Debug, Clone)]
pub struct MreMappingConfig {
/// Number of K-means segments for spatial grouping.
pub n_segments: usize,
/// Interpolation method.
pub interpolation: InterpolationMethod,
/// Search radius for local averaging (mm).
pub search_radius: Option<f64>,
/// Mask threshold for valid MRE values.
pub mask_threshold: f64,
}
impl Default for MreMappingConfig {
fn default() -> Self {
Self {
n_segments: 5,
interpolation: InterpolationMethod::Trilinear,
search_radius: None,
mask_threshold: 0.0,
}
}
}
/// Interpolation method for MRE values.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum InterpolationMethod {
/// Nearest neighbor interpolation.
NearestNeighbor,
/// Trilinear interpolation.
Trilinear,
/// Inverse distance weighting.
InverseDistance,
}
/// MRE property type.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MrePropertyType {
/// Complex shear modulus (G' and G'').
ComplexShear,
/// Wave amplitude.
Amplitude,
/// Phase velocity.
PhaseVelocity,
/// Stiffness (μ) and damping (ξ).
StiffnessDamping,
}
/// Result of MRE mapping operation.
#[derive(Debug, Clone)]
pub struct MreMappingResult {
/// Mapped storage modulus values at mesh nodes.
pub storage_modulus: Vec<f64>,
/// Mapped loss modulus values at mesh nodes.
pub loss_modulus: Vec<f64>,
/// Mask indicating valid values (1.0) vs extrapolated (0.0).
pub valid_mask: Vec<f64>,
/// Segment assignments for each node.
pub segments: Option<Vec<usize>>,
/// Statistics about the mapping.
pub stats: MappingStats,
}
/// Statistics about the MRE mapping.
#[derive(Debug, Clone, Default)]
pub struct MappingStats {
/// Number of nodes mapped.
pub nodes_mapped: usize,
/// Number of nodes with valid MRE data.
pub nodes_valid: usize,
/// Number of nodes extrapolated.
pub nodes_extrapolated: usize,
/// Mean storage modulus.
pub mean_storage: f64,
/// Mean loss modulus.
pub mean_loss: f64,
/// Standard deviation of storage modulus.
pub std_storage: f64,
/// Standard deviation of loss modulus.
pub std_loss: f64,
}
/// Map MRE properties onto mesh nodes.
pub fn map_mre_to_mesh(
mesh: &TetrahedralMesh,
storage_volume: &Volume,
loss_volume: &Volume,
mask: Option<&Volume>,
config: &MreMappingConfig,
) -> Result<MreMappingResult> {
let n_nodes = mesh.vertices.len();
if n_nodes == 0 {
return Err(Mri2FeError::MappingError("Empty mesh".to_string()));
}
// Get volume properties
let shape = storage_volume.shape();
let spacing = storage_volume.spacing();
let origin = storage_volume.origin();
// Map each node to volume coordinates and sample
let node_values: Vec<(f64, f64, bool)> = mesh
.vertices
.par_iter()
.map(|vertex| {
let node = vertex.position;
// Convert world coordinates to voxel coordinates
let vx = ((node.x - origin[0]) / spacing[0]).floor() as i64;
let vy = ((node.y - origin[1]) / spacing[1]).floor() as i64;
let vz = ((node.z - origin[2]) / spacing[2]).floor() as i64;
// Check if within bounds
if vx < 0
|| vy < 0
|| vz < 0
|| vx >= shape[0] as i64
|| vy >= shape[1] as i64
|| vz >= shape[2] as i64
{
return (0.0, 0.0, false);
}
let (vx, vy, vz) = (vx as usize, vy as usize, vz as usize);
// Check mask if provided
if let Some(m) = mask {
if m.get(vx, vy, vz).unwrap_or(0.0) <= config.mask_threshold {
return (0.0, 0.0, false);
}
}
match config.interpolation {
InterpolationMethod::NearestNeighbor => {
let g_prime = storage_volume.get(vx, vy, vz).unwrap_or(0.0);
let g_double_prime = loss_volume.get(vx, vy, vz).unwrap_or(0.0);
(g_prime, g_double_prime, true)
}
InterpolationMethod::Trilinear => {
let g_prime = trilinear_interp(storage_volume, &node, &origin, &spacing);
let g_double_prime = trilinear_interp(loss_volume, &node, &origin, &spacing);
(g_prime, g_double_prime, true)
}
InterpolationMethod::InverseDistance => {
let radius = config.search_radius.unwrap_or(spacing[0] * 2.0);
let g_prime = idw_interp(storage_volume, &node, &origin, &spacing, radius);
let g_double_prime = idw_interp(loss_volume, &node, &origin, &spacing, radius);
(g_prime, g_double_prime, true)
}
}
})
.collect();
// Separate results
let mut storage_modulus = Vec::with_capacity(n_nodes);
let mut loss_modulus = Vec::with_capacity(n_nodes);
let mut valid_mask = Vec::with_capacity(n_nodes);
let mut nodes_valid = 0;
for (g_prime, g_double_prime, valid) in &node_values {
storage_modulus.push(*g_prime);
loss_modulus.push(*g_double_prime);
valid_mask.push(if *valid { 1.0 } else { 0.0 });
if *valid {
nodes_valid += 1;
}
}
// Optionally perform K-means segmentation
let segments = if config.n_segments > 1 && nodes_valid > config.n_segments {
Some(segment_nodes(
&storage_modulus,
&loss_modulus,
&valid_mask,
config.n_segments,
)?)
} else {
None
};
// Calculate statistics
let stats = calculate_stats(&storage_modulus, &loss_modulus, &valid_mask, nodes_valid);
Ok(MreMappingResult {
storage_modulus,
loss_modulus,
valid_mask,
segments,
stats,
})
}
/// Trilinear interpolation in a volume.
fn trilinear_interp(
volume: &Volume,
point: &Point3<f64>,
origin: &[f64; 3],
spacing: &[f64; 3],
) -> f64 {
let shape = volume.shape();
// Get fractional voxel coordinates
let fx = (point.x - origin[0]) / spacing[0];
let fy = (point.y - origin[1]) / spacing[1];
let fz = (point.z - origin[2]) / spacing[2];
let x0 = fx.floor() as i64;
let y0 = fy.floor() as i64;
let z0 = fz.floor() as i64;
let x1 = x0 + 1;
let y1 = y0 + 1;
let z1 = z0 + 1;
// Check bounds
if x0 < 0
|| y0 < 0
|| z0 < 0
|| x1 >= shape[0] as i64
|| y1 >= shape[1] as i64
|| z1 >= shape[2] as i64
{
// Fall back to nearest neighbor
let vx = fx.round().clamp(0.0, (shape[0] - 1) as f64) as usize;
let vy = fy.round().clamp(0.0, (shape[1] - 1) as f64) as usize;
let vz = fz.round().clamp(0.0, (shape[2] - 1) as f64) as usize;
return volume.get(vx, vy, vz).unwrap_or(0.0);
}
let (x0, y0, z0) = (x0 as usize, y0 as usize, z0 as usize);
let (x1, y1, z1) = (x1 as usize, y1 as usize, z1 as usize);
// Interpolation weights
let xd = fx - fx.floor();
let yd = fy - fy.floor();
let zd = fz - fz.floor();
// Sample 8 corners
let c000 = volume.get(x0, y0, z0).unwrap_or(0.0);
let c001 = volume.get(x0, y0, z1).unwrap_or(0.0);
let c010 = volume.get(x0, y1, z0).unwrap_or(0.0);
let c011 = volume.get(x0, y1, z1).unwrap_or(0.0);
let c100 = volume.get(x1, y0, z0).unwrap_or(0.0);
let c101 = volume.get(x1, y0, z1).unwrap_or(0.0);
let c110 = volume.get(x1, y1, z0).unwrap_or(0.0);
let c111 = volume.get(x1, y1, z1).unwrap_or(0.0);
// Trilinear interpolation
let c00 = c000 * (1.0 - xd) + c100 * xd;
let c01 = c001 * (1.0 - xd) + c101 * xd;
let c10 = c010 * (1.0 - xd) + c110 * xd;
let c11 = c011 * (1.0 - xd) + c111 * xd;
let c0 = c00 * (1.0 - yd) + c10 * yd;
let c1 = c01 * (1.0 - yd) + c11 * yd;
c0 * (1.0 - zd) + c1 * zd
}
/// Inverse distance weighted interpolation.
fn idw_interp(
volume: &Volume,
point: &Point3<f64>,
origin: &[f64; 3],
spacing: &[f64; 3],
radius: f64,
) -> f64 {
let shape = volume.shape();
// Center voxel
let cx = ((point.x - origin[0]) / spacing[0]).round() as i64;
let cy = ((point.y - origin[1]) / spacing[1]).round() as i64;
let cz = ((point.z - origin[2]) / spacing[2]).round() as i64;
// Search radius in voxels
let rx = (radius / spacing[0]).ceil() as i64;
let ry = (radius / spacing[1]).ceil() as i64;
let rz = (radius / spacing[2]).ceil() as i64;
let mut sum = 0.0;
let mut weight_sum = 0.0;
for dz in -rz..=rz {
for dy in -ry..=ry {
for dx in -rx..=rx {
let vx = cx + dx;
let vy = cy + dy;
let vz = cz + dz;
if vx < 0
|| vy < 0
|| vz < 0
|| vx >= shape[0] as i64
|| vy >= shape[1] as i64
|| vz >= shape[2] as i64
{
continue;
}
let (vx, vy, vz) = (vx as usize, vy as usize, vz as usize);
// World position of this voxel
let vox_x = origin[0] + vx as f64 * spacing[0];
let vox_y = origin[1] + vy as f64 * spacing[1];
let vox_z = origin[2] + vz as f64 * spacing[2];
let dist = ((point.x - vox_x).powi(2)
+ (point.y - vox_y).powi(2)
+ (point.z - vox_z).powi(2))
.sqrt();
if dist <= radius {
let weight = if dist < 1e-10 {
1e10
} else {
1.0 / (dist * dist)
};
sum += weight * volume.get(vx, vy, vz).unwrap_or(0.0);
weight_sum += weight;
}
}
}
}
if weight_sum > 0.0 {
sum / weight_sum
} else {
0.0
}
}
/// Segment nodes using K-means on MRE values.
fn segment_nodes(
storage: &[f64],
loss: &[f64],
valid_mask: &[f64],
n_segments: usize,
) -> Result<Vec<usize>> {
// Collect valid data points
let mut data: Vec<DVector<f64>> = Vec::new();
let mut valid_indices: Vec<usize> = Vec::new();
for i in 0..storage.len() {
if valid_mask[i] > 0.5 {
data.push(DVector::from_vec(vec![storage[i], loss[i]]));
valid_indices.push(i);
}
}
if data.len() < n_segments {
return Err(Mri2FeError::MappingError(format!(
"Not enough valid points ({}) for {} segments",
data.len(),
n_segments
)));
}
// Run K-means
let kmeans = KMeans::with_k(n_segments);
let result = kmeans.fit(&data)?;
// Map back to all nodes
let mut segments = vec![0; storage.len()];
for (i, &orig_idx) in valid_indices.iter().enumerate() {
segments[orig_idx] = result.labels[i];
}
Ok(segments)
}
/// Calculate mapping statistics.
fn calculate_stats(
storage: &[f64],
loss: &[f64],
valid_mask: &[f64],
nodes_valid: usize,
) -> MappingStats {
let n_nodes = storage.len();
// Calculate means
let (sum_storage, sum_loss): (f64, f64) = storage
.iter()
.zip(loss.iter())
.zip(valid_mask.iter())
.filter(|(_, m)| **m > 0.5)
.map(|((s, l), _)| (*s, *l))
.fold((0.0, 0.0), |(ss, sl), (s, l)| (ss + s, sl + l));
let mean_storage = if nodes_valid > 0 {
sum_storage / nodes_valid as f64
} else {
0.0
};
let mean_loss = if nodes_valid > 0 {
sum_loss / nodes_valid as f64
} else {
0.0
};
// Calculate standard deviations
let (var_storage, var_loss): (f64, f64) = storage
.iter()
.zip(loss.iter())
.zip(valid_mask.iter())
.filter(|(_, m)| **m > 0.5)
.map(|((s, l), _)| ((*s - mean_storage).powi(2), (*l - mean_loss).powi(2)))
.fold((0.0, 0.0), |(vs, vl), (s, l)| (vs + s, vl + l));
let std_storage = if nodes_valid > 1 {
(var_storage / (nodes_valid - 1) as f64).sqrt()
} else {
0.0
};
let std_loss = if nodes_valid > 1 {
(var_loss / (nodes_valid - 1) as f64).sqrt()
} else {
0.0
};
MappingStats {
nodes_mapped: n_nodes,
nodes_valid,
nodes_extrapolated: n_nodes - nodes_valid,
mean_storage,
mean_loss,
std_storage,
std_loss,
}
}
/// Convert stiffness (μ) and damping (ξ) to complex shear modulus.
pub fn stiffness_damping_to_complex_shear(mu: f64, xi: f64) -> (f64, f64) {
// G' = μ(1 - ξ²) / (1 + ξ²)
// G'' = 2μξ / (1 + ξ²)
let xi_sq = xi * xi;
let denom = 1.0 + xi_sq;
let g_prime = mu * (1.0 - xi_sq) / denom;
let g_double_prime = 2.0 * mu * xi / denom;
(g_prime, g_double_prime)
}
/// Convert complex shear modulus to stiffness and damping.
pub fn complex_shear_to_stiffness_damping(g_prime: f64, g_double_prime: f64) -> (f64, f64) {
// μ = sqrt(G'² + G''²)
// ξ = G'' / (G' + μ)
let mu = (g_prime * g_prime + g_double_prime * g_double_prime).sqrt();
let xi = if (g_prime + mu).abs() > 1e-10 {
g_double_prime / (g_prime + mu)
} else {
0.0
};
(mu, xi)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_stiffness_damping_conversion() {
let mu = 1000.0;
let xi = 0.1;
let (g_prime, g_double_prime) = stiffness_damping_to_complex_shear(mu, xi);
let (mu_back, xi_back) = complex_shear_to_stiffness_damping(g_prime, g_double_prime);
assert!((mu - mu_back).abs() < 1e-6);
assert!((xi - xi_back).abs() < 1e-6);
}
#[test]
fn test_trilinear_interp() {
// Create a simple volume with identity affine
let identity_affine = [
[1.0, 0.0, 0.0, 0.0],
[0.0, 1.0, 0.0, 0.0],
[0.0, 0.0, 1.0, 0.0],
[0.0, 0.0, 0.0, 1.0],
];
let volume = Volume::new(
vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0],
[2, 2, 2],
[1.0, 1.0, 1.0],
[0.0, 0.0, 0.0],
identity_affine,
);
let origin = [0.0, 0.0, 0.0];
let spacing = [1.0, 1.0, 1.0];
// Center of volume
let point = Point3::new(0.5, 0.5, 0.5);
let val = trilinear_interp(&volume, &point, &origin, &spacing);
// Should be average of all 8 corners = 4.5
assert!((val - 4.5).abs() < 1e-6);
}
#[test]
fn test_mapping_config_default() {
let config = MreMappingConfig::default();
assert_eq!(config.n_segments, 5);
assert_eq!(config.interpolation, InterpolationMethod::Trilinear);
assert!(config.search_radius.is_none());
}
#[test]
fn test_mapping_stats() {
let storage = vec![100.0, 200.0, 300.0, 0.0];
let loss = vec![10.0, 20.0, 30.0, 0.0];
let valid_mask = vec![1.0, 1.0, 1.0, 0.0];
let stats = calculate_stats(&storage, &loss, &valid_mask, 3);
assert_eq!(stats.nodes_mapped, 4);
assert_eq!(stats.nodes_valid, 3);
assert_eq!(stats.nodes_extrapolated, 1);
assert!((stats.mean_storage - 200.0).abs() < 1e-6);
assert!((stats.mean_loss - 20.0).abs() < 1e-6);
}
}