Files
rustytorch/crates/specialized/rtx-digital-twin/src/geometry.rs
T
2026-03-04 00:08:42 +00:00

461 lines
14 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//! Organ geometry representation for digital twins.
//!
//! This module provides structures for representing patient-specific organ
//! geometry from segmented medical imaging data.
use serde::{Deserialize, Serialize};
use crate::error::{DigitalTwinError, Result};
use crate::tissue::{TissueDatabase, TissueProperties, TissueType};
/// Tissue label at a voxel position.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct TissueLabel(pub u8);
impl TissueLabel {
/// Create a new tissue label.
pub fn new(label: u8) -> Self {
Self(label)
}
/// Get the tissue type for this label.
pub fn tissue_type(&self) -> TissueType {
TissueType::from_label(self.0)
}
/// Get the raw label value.
pub fn value(&self) -> u8 {
self.0
}
}
impl From<TissueType> for TissueLabel {
fn from(tt: TissueType) -> Self {
Self(tt.label())
}
}
/// Per-voxel data including tissue label and physical state.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VoxelData {
/// Tissue label (from segmentation)
pub label: TissueLabel,
/// Temperature [°C] - for bioheat simulations
pub temperature: f32,
/// Damage parameter (Arrhenius) [-]
pub damage: f32,
/// Custom scalar field (user-defined)
pub scalar_field: f32,
}
impl Default for VoxelData {
fn default() -> Self {
Self {
label: TissueLabel(0), // Air
temperature: 37.0, // Body temperature
damage: 0.0,
scalar_field: 0.0,
}
}
}
impl VoxelData {
/// Create voxel data with specific tissue label.
pub fn with_label(label: TissueLabel) -> Self {
Self {
label,
..Default::default()
}
}
/// Check if voxel is inside the body (not air).
pub fn is_tissue(&self) -> bool {
self.label.0 != 0
}
}
/// 3D organ geometry representation.
///
/// Stores the patient-specific geometry as a voxelized representation
/// with tissue labels and physical state at each voxel.
#[derive(Debug, Clone)]
pub struct OrganGeometry {
/// Voxel data array [x * y * z]
data: Vec<VoxelData>,
/// Shape [x, y, z]
shape: [usize; 3],
/// Voxel spacing [dx, dy, dz] in mm
spacing: [f32; 3],
/// Origin position in world coordinates [x, y, z] in mm
origin: [f32; 3],
/// Tissue property database
tissue_db: TissueDatabase,
}
impl OrganGeometry {
/// Create new organ geometry filled with air.
pub fn new(shape: [usize; 3], spacing: [f32; 3]) -> Self {
let size = shape[0] * shape[1] * shape[2];
Self {
data: vec![VoxelData::default(); size],
shape,
spacing,
origin: [0.0, 0.0, 0.0],
tissue_db: TissueDatabase::standard(),
}
}
/// Create geometry from segmentation labels.
///
/// # Arguments
/// * `labels` - Flattened array of tissue labels [x * y * z]
/// * `shape` - Volume shape [x, y, z]
/// * `spacing` - Voxel spacing in mm
pub fn from_labels(labels: &[u8], shape: [usize; 3], spacing: [f32; 3]) -> Result<Self> {
let expected_size = shape[0] * shape[1] * shape[2];
if labels.len() != expected_size {
return Err(DigitalTwinError::InvalidGeometry(format!(
"Label array size {} doesn't match shape {:?} (expected {})",
labels.len(),
shape,
expected_size
)));
}
let data: Vec<VoxelData> = labels
.iter()
.map(|&label| VoxelData::with_label(TissueLabel(label)))
.collect();
Ok(Self {
data,
shape,
spacing,
origin: [0.0, 0.0, 0.0],
tissue_db: TissueDatabase::standard(),
})
}
/// Set the origin position.
pub fn set_origin(&mut self, origin: [f32; 3]) {
self.origin = origin;
}
/// Set custom tissue database.
pub fn set_tissue_database(&mut self, db: TissueDatabase) {
self.tissue_db = db;
}
/// Get the shape of the geometry.
pub fn shape(&self) -> [usize; 3] {
self.shape
}
/// Get voxel spacing in mm.
pub fn spacing(&self) -> [f32; 3] {
self.spacing
}
/// Get origin position.
pub fn origin(&self) -> [f32; 3] {
self.origin
}
/// Get total number of voxels.
pub fn num_voxels(&self) -> usize {
self.data.len()
}
/// Get physical dimensions in mm.
pub fn dimensions(&self) -> [f32; 3] {
[
self.shape[0] as f32 * self.spacing[0],
self.shape[1] as f32 * self.spacing[1],
self.shape[2] as f32 * self.spacing[2],
]
}
/// Convert index to 3D coordinates.
pub fn index_to_coords(&self, index: usize) -> [usize; 3] {
let z = index / (self.shape[0] * self.shape[1]);
let remainder = index % (self.shape[0] * self.shape[1]);
let y = remainder / self.shape[0];
let x = remainder % self.shape[0];
[x, y, z]
}
/// Convert 3D coordinates to index.
pub fn coords_to_index(&self, x: usize, y: usize, z: usize) -> usize {
z * self.shape[0] * self.shape[1] + y * self.shape[0] + x
}
/// Check if coordinates are within bounds.
pub fn in_bounds(&self, x: usize, y: usize, z: usize) -> bool {
x < self.shape[0] && y < self.shape[1] && z < self.shape[2]
}
/// Get voxel data at coordinates.
pub fn get(&self, x: usize, y: usize, z: usize) -> Option<&VoxelData> {
if self.in_bounds(x, y, z) {
Some(&self.data[self.coords_to_index(x, y, z)])
} else {
None
}
}
/// Get mutable voxel data at coordinates.
pub fn get_mut(&mut self, x: usize, y: usize, z: usize) -> Option<&mut VoxelData> {
if self.in_bounds(x, y, z) {
let idx = self.coords_to_index(x, y, z);
Some(&mut self.data[idx])
} else {
None
}
}
/// Get voxel data by flat index.
pub fn get_by_index(&self, index: usize) -> Option<&VoxelData> {
self.data.get(index)
}
/// Get mutable voxel data by flat index.
pub fn get_by_index_mut(&mut self, index: usize) -> Option<&mut VoxelData> {
self.data.get_mut(index)
}
/// Get tissue properties at coordinates.
pub fn tissue_properties(&self, x: usize, y: usize, z: usize) -> Option<TissueProperties> {
self.get(x, y, z)
.map(|voxel| self.tissue_db.get_or_default(voxel.label.tissue_type()))
}
/// Get all voxel data as a slice.
pub fn data(&self) -> &[VoxelData] {
&self.data
}
/// Get mutable access to all voxel data.
pub fn data_mut(&mut self) -> &mut [VoxelData] {
&mut self.data
}
/// Get reference to tissue database.
pub fn tissue_db(&self) -> &TissueDatabase {
&self.tissue_db
}
/// Set tissue label at coordinates.
pub fn set_label(&mut self, x: usize, y: usize, z: usize, label: TissueLabel) -> bool {
if let Some(voxel) = self.get_mut(x, y, z) {
voxel.label = label;
true
} else {
false
}
}
/// Set temperature at coordinates.
pub fn set_temperature(&mut self, x: usize, y: usize, z: usize, temp: f32) -> bool {
if let Some(voxel) = self.get_mut(x, y, z) {
voxel.temperature = temp;
true
} else {
false
}
}
/// Get temperature field as flat array.
pub fn temperature_field(&self) -> Vec<f32> {
self.data.iter().map(|v| v.temperature).collect()
}
/// Set temperature field from flat array.
pub fn set_temperature_field(&mut self, field: &[f32]) -> Result<()> {
if field.len() != self.data.len() {
return Err(DigitalTwinError::ShapeMismatch {
expected: self.shape,
got: [field.len(), 1, 1],
});
}
for (voxel, &temp) in self.data.iter_mut().zip(field.iter()) {
voxel.temperature = temp;
}
Ok(())
}
/// Get thermal conductivity field (k) from tissue properties.
pub fn thermal_conductivity_field(&self) -> Vec<f32> {
self.data
.iter()
.map(|v| {
self.tissue_db
.get_or_default(v.label.tissue_type())
.thermal_conductivity
})
.collect()
}
/// Get density field (ρ) from tissue properties.
pub fn density_field(&self) -> Vec<f32> {
self.data
.iter()
.map(|v| self.tissue_db.get_or_default(v.label.tissue_type()).density)
.collect()
}
/// Get specific heat field (c) from tissue properties.
pub fn specific_heat_field(&self) -> Vec<f32> {
self.data
.iter()
.map(|v| {
self.tissue_db
.get_or_default(v.label.tissue_type())
.specific_heat
})
.collect()
}
/// Get perfusion rate field (ω) from tissue properties.
pub fn perfusion_field(&self) -> Vec<f32> {
self.data
.iter()
.map(|v| {
self.tissue_db
.get_or_default(v.label.tissue_type())
.perfusion_rate
})
.collect()
}
/// Count voxels of each tissue type.
pub fn tissue_histogram(&self) -> std::collections::HashMap<TissueType, usize> {
let mut counts = std::collections::HashMap::new();
for voxel in &self.data {
*counts.entry(voxel.label.tissue_type()).or_insert(0) += 1;
}
counts
}
/// Create a spherical region centered at world coordinates.
///
/// # Arguments
/// * `center` - Center position in mm [x, y, z]
/// * `radius` - Radius in mm
/// * `label` - Tissue label to assign
pub fn create_sphere(&mut self, center: [f32; 3], radius: f32, label: TissueLabel) {
let radius_sq = radius * radius;
for z in 0..self.shape[2] {
for y in 0..self.shape[1] {
for x in 0..self.shape[0] {
// World position of this voxel
let pos = [
self.origin[0] + x as f32 * self.spacing[0],
self.origin[1] + y as f32 * self.spacing[1],
self.origin[2] + z as f32 * self.spacing[2],
];
// Distance squared from center
let dist_sq = (pos[0] - center[0]).powi(2)
+ (pos[1] - center[1]).powi(2)
+ (pos[2] - center[2]).powi(2);
if dist_sq <= radius_sq {
self.set_label(x, y, z, label);
}
}
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_geometry_creation() {
let geom = OrganGeometry::new([10, 10, 10], [1.0, 1.0, 1.0]);
assert_eq!(geom.shape(), [10, 10, 10]);
assert_eq!(geom.num_voxels(), 1000);
assert_eq!(geom.dimensions(), [10.0, 10.0, 10.0]);
}
#[test]
fn test_from_labels() {
let labels: Vec<u8> = (0..27).map(|i| if i == 13 { 6 } else { 0 }).collect();
let geom = OrganGeometry::from_labels(&labels, [3, 3, 3], [1.0, 1.0, 1.0]).unwrap();
// Center voxel should be liver (label 6)
let center = geom.get(1, 1, 1).unwrap();
assert_eq!(center.label.value(), 6);
assert_eq!(center.label.tissue_type(), TissueType::Liver);
}
#[test]
fn test_coordinate_conversion() {
let geom = OrganGeometry::new([5, 6, 7], [1.0, 1.0, 1.0]);
let index = geom.coords_to_index(2, 3, 4);
let coords = geom.index_to_coords(index);
assert_eq!(coords, [2, 3, 4]);
}
#[test]
fn test_tissue_properties_lookup() {
let labels = vec![6u8; 8]; // All liver
let geom = OrganGeometry::from_labels(&labels, [2, 2, 2], [1.0, 1.0, 1.0]).unwrap();
let props = geom.tissue_properties(0, 0, 0).unwrap();
assert_eq!(props.tissue_type, TissueType::Liver);
assert!(props.thermal_conductivity > 0.5);
}
#[test]
fn test_temperature_field() {
let mut geom = OrganGeometry::new([3, 3, 3], [1.0, 1.0, 1.0]);
// Set center to 42°C
geom.set_temperature(1, 1, 1, 42.0);
let temp_field = geom.temperature_field();
let center_idx = geom.coords_to_index(1, 1, 1);
assert_eq!(temp_field[center_idx], 42.0);
}
#[test]
fn test_create_sphere() {
let mut geom = OrganGeometry::new([20, 20, 20], [1.0, 1.0, 1.0]);
// Create tumor sphere at center
let center = [10.0, 10.0, 10.0];
geom.create_sphere(center, 5.0, TissueLabel::from(TissueType::Tumor));
// Center should be tumor
let center_voxel = geom.get(10, 10, 10).unwrap();
assert_eq!(center_voxel.label.tissue_type(), TissueType::Tumor);
// Corner should be air
let corner_voxel = geom.get(0, 0, 0).unwrap();
assert_eq!(corner_voxel.label.tissue_type(), TissueType::Air);
}
#[test]
fn test_tissue_histogram() {
let mut labels = vec![0u8; 100]; // Mostly air
labels[50] = 6; // One liver voxel
labels[51] = 6;
labels[52] = 3; // One muscle voxel
let geom = OrganGeometry::from_labels(&labels, [10, 10, 1], [1.0, 1.0, 1.0]).unwrap();
let hist = geom.tissue_histogram();
assert_eq!(hist.get(&TissueType::Air), Some(&97));
assert_eq!(hist.get(&TissueType::Liver), Some(&2));
assert_eq!(hist.get(&TissueType::Muscle), Some(&1));
}
}