Files
rustytorch/demos/rtx-cardiosim-demo/src/lib.rs
T

433 lines
12 KiB
Rust

//! `CardioSim` cardiac electrophysiology simulation demo.
//!
//! This crate implements cardiac electrophysiology simulation using
//! physics-informed neural operators (PINO).
//!
//! # Architecture
//!
//! The simulation consists of:
//! - **Monodomain Solver**: Traditional PDE solver for comparison
//! - **Neural Operator**: PINO for fast inference
//! - **Ionic Models**: Mitchell-Schaeffer, FitzHugh-Nagumo, etc.
//! - **ECG Computation**: Virtual ECG from voltage fields
//!
//! # Example
//!
//! ```ignore
//! use rtx_cardiosim_demo::{run_simulation, SimulationRequest};
//!
//! let result = run_simulation(request).await?;
//! println!("Computed {} time steps in {:.2}s",
//! result.stats.time_steps,
//! result.stats.compute_time);
//! ```
pub mod ecg;
pub mod ionic_models;
pub mod monodomain;
pub mod neural_operator;
pub mod sample_data;
use cardiosim_shared::{
APDMap, ActivationMap, ArrhythmiaEvent, ArrhythmiaType, ECGRequest, ECGResult, HeartMesh,
NeuralOperatorConfig, Point3D, SimulationConfig, SimulationRequest, SimulationResult,
SimulationStats, VoltageField,
};
use thiserror::Error;
/// Errors that can occur during cardiac simulation.
#[derive(Debug, Error)]
pub enum CardioSimError {
/// Invalid mesh
#[error("Invalid mesh: {0}")]
InvalidMesh(String),
/// Simulation error
#[error("Simulation error: {0}")]
SimulationError(String),
/// Neural operator error
#[error("Neural operator error: {0}")]
NeuralOperatorError(String),
/// Configuration error
#[error("Configuration error: {0}")]
ConfigError(String),
}
/// Run cardiac electrophysiology simulation.
pub async fn run_simulation(
request: SimulationRequest,
) -> Result<SimulationResult, CardioSimError> {
use std::time::Instant;
let start = Instant::now();
// Validate mesh
validate_mesh(&request.mesh)?;
tracing::info!(
"Running cardiac simulation on mesh with {} vertices",
request.mesh.vertices.len()
);
// Choose solver based on configuration
let (voltage_fields, speedup_factor) = if request.config.use_neural_operator {
let operator = neural_operator::CardiacPINO::new(NeuralOperatorConfig::default());
let fields = operator.solve(&request.mesh, &request.protocol, &request.config)?;
(fields, 1000.0) // Neural operators are ~1000x faster
} else {
let solver = monodomain::MonodomainSolver::new(&request.config);
let fields = solver.solve(&request.mesh, &request.protocol)?;
(fields, 1.0)
};
// Compute activation and APD maps
let activation_map = compute_activation_map(&voltage_fields);
let apd_map = compute_apd_map(&voltage_fields);
// Detect arrhythmias
let arrhythmias = detect_arrhythmias(&voltage_fields, &activation_map);
// Extract time points
let times: Vec<f32> = voltage_fields.iter().map(|f| f.time).collect();
let elapsed = start.elapsed();
let stats = SimulationStats {
compute_time: elapsed.as_secs_f32(),
time_steps: voltage_fields.len(),
speedup_factor,
inference_time: elapsed.as_millis() as f32 / voltage_fields.len() as f32,
};
Ok(SimulationResult {
times,
voltage_fields,
activation_map,
apd_map,
arrhythmias,
stats,
})
}
/// Compute ECG from simulation result.
pub async fn compute_ecg(request: ECGRequest) -> Result<ECGResult, CardioSimError> {
let ecg_computer = ecg::ECGComputer::new(&request.leads);
ecg_computer.compute(&request.simulation)
}
/// Validate mesh.
fn validate_mesh(mesh: &HeartMesh) -> Result<(), CardioSimError> {
if mesh.vertices.is_empty() {
return Err(CardioSimError::InvalidMesh(
"Mesh has no vertices".to_string(),
));
}
if mesh.triangles.is_empty() {
return Err(CardioSimError::InvalidMesh(
"Mesh has no triangles".to_string(),
));
}
if mesh.fibers.len() != mesh.vertices.len() {
return Err(CardioSimError::InvalidMesh(
"Fiber directions don't match vertex count".to_string(),
));
}
Ok(())
}
/// Compute activation time map.
fn compute_activation_map(voltage_fields: &[VoltageField]) -> ActivationMap {
if voltage_fields.is_empty() {
return ActivationMap {
activation_times: vec![],
conduction_velocity: vec![],
};
}
let n_vertices = voltage_fields[0].voltages.len();
let mut activation_times = vec![f32::MAX; n_vertices];
// Find first time each vertex crosses -30 mV (activation threshold)
let threshold = -30.0;
for field in voltage_fields {
for (i, &voltage) in field.voltages.iter().enumerate() {
if voltage > threshold && activation_times[i] == f32::MAX {
activation_times[i] = field.time;
}
}
}
// Replace unactivated vertices with NaN
for t in &mut activation_times {
if *t == f32::MAX {
*t = f32::NAN;
}
}
// Compute conduction velocity (simplified)
let conduction_velocity = vec![0.8; n_vertices]; // ~0.8 m/s typical
ActivationMap {
activation_times,
conduction_velocity,
}
}
/// Compute APD map.
fn compute_apd_map(voltage_fields: &[VoltageField]) -> APDMap {
if voltage_fields.is_empty() {
return APDMap {
apd50: vec![],
apd90: vec![],
dispersion: 0.0,
};
}
let n_vertices = voltage_fields[0].voltages.len();
let mut apd50 = vec![0.0; n_vertices];
let mut apd90 = vec![0.0; n_vertices];
// Track activation and repolarization times
let mut activation_times = vec![f32::MAX; n_vertices];
let mut repol50_times = vec![f32::MAX; n_vertices];
let mut repol90_times = vec![f32::MAX; n_vertices];
let mut max_voltages = vec![f32::MIN; n_vertices];
for field in voltage_fields {
for (i, &voltage) in field.voltages.iter().enumerate() {
// Track max voltage
if voltage > max_voltages[i] {
max_voltages[i] = voltage;
}
// Activation (upstroke)
if voltage > -30.0 && activation_times[i] == f32::MAX {
activation_times[i] = field.time;
}
// Repolarization at 50%
if activation_times[i] < f32::MAX && repol50_times[i] == f32::MAX {
let v_max = max_voltages[i];
let v_rest = -85.0;
let v_50 = v_rest + 0.5 * (v_max - v_rest);
if voltage < v_50 {
repol50_times[i] = field.time;
}
}
// Repolarization at 90%
if activation_times[i] < f32::MAX && repol90_times[i] == f32::MAX {
let v_max = max_voltages[i];
let v_rest = -85.0;
let v_90 = v_rest + 0.1 * (v_max - v_rest);
if voltage < v_90 {
repol90_times[i] = field.time;
}
}
}
}
// Calculate APDs
for i in 0..n_vertices {
if activation_times[i] < f32::MAX {
if repol50_times[i] < f32::MAX {
apd50[i] = repol50_times[i] - activation_times[i];
}
if repol90_times[i] < f32::MAX {
apd90[i] = repol90_times[i] - activation_times[i];
}
}
}
// Calculate dispersion
let valid_apd90: Vec<f32> = apd90.iter().copied().filter(|&x| x > 0.0).collect();
let dispersion = if valid_apd90.len() > 1 {
let mean: f32 = valid_apd90.iter().sum::<f32>() / valid_apd90.len() as f32;
let variance: f32 =
valid_apd90.iter().map(|x| (x - mean).powi(2)).sum::<f32>() / valid_apd90.len() as f32;
variance.sqrt()
} else {
0.0
};
APDMap {
apd50,
apd90,
dispersion,
}
}
/// Detect arrhythmias in simulation.
fn detect_arrhythmias(
voltage_fields: &[VoltageField],
activation_map: &ActivationMap,
) -> Vec<ArrhythmiaEvent> {
let mut events = Vec::new();
if voltage_fields.is_empty() {
return events;
}
// Check for normal rhythm
let max_activation = activation_map
.activation_times
.iter()
.filter(|&&t| t.is_finite())
.fold(0.0_f32, |a, &b| a.max(b));
let min_activation = activation_map
.activation_times
.iter()
.filter(|&&t| t.is_finite())
.fold(f32::MAX, |a, &b| a.min(b));
let activation_spread = max_activation - min_activation;
// Normal activation should complete within ~150ms
if activation_spread < 150.0 {
events.push(ArrhythmiaEvent {
arrhythmia_type: ArrhythmiaType::Normal,
start_time: 0.0,
end_time: Some(voltage_fields.last().map_or(0.0, |f| f.time)),
location: None,
severity: 0.0,
});
} else if activation_spread > 500.0 {
// Likely reentry
events.push(ArrhythmiaEvent {
arrhythmia_type: ArrhythmiaType::Reentry,
start_time: min_activation,
end_time: None,
location: Some(Point3D::new(0.0, 0.0, 0.0)),
severity: 0.7,
});
} else if activation_spread > 300.0 {
// Possible conduction abnormality
events.push(ArrhythmiaEvent {
arrhythmia_type: ArrhythmiaType::Block,
start_time: min_activation,
end_time: Some(max_activation),
location: None,
severity: 0.4,
});
}
events
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_run_simulation() {
let mesh = cardiosim_shared::get_sample_heart_mesh();
let protocol = cardiosim_shared::get_sample_protocol();
let config = SimulationConfig::default();
let request = SimulationRequest {
mesh,
config,
protocol,
};
let result = run_simulation(request).await;
assert!(result.is_ok());
let sim = result.unwrap();
assert!(!sim.voltage_fields.is_empty());
assert!(!sim.activation_map.activation_times.is_empty());
}
#[tokio::test]
async fn test_neural_operator_simulation() {
let mesh = cardiosim_shared::get_sample_heart_mesh();
let protocol = cardiosim_shared::get_sample_protocol();
let config = SimulationConfig {
use_neural_operator: true,
total_time: 100.0,
..Default::default()
};
let request = SimulationRequest {
mesh,
config,
protocol,
};
let result = run_simulation(request).await;
assert!(result.is_ok());
let sim = result.unwrap();
assert!(sim.stats.speedup_factor > 1.0);
}
#[test]
fn test_validate_mesh() {
let mesh = cardiosim_shared::get_sample_heart_mesh();
assert!(validate_mesh(&mesh).is_ok());
}
#[test]
fn test_validate_empty_mesh() {
let mesh = HeartMesh {
name: "Empty".to_string(),
vertices: vec![],
triangles: vec![],
tetrahedra: None,
fibers: vec![],
sheets: vec![],
regions: vec![],
vertex_regions: vec![],
};
assert!(validate_mesh(&mesh).is_err());
}
#[test]
fn test_compute_activation_map() {
let fields = vec![
VoltageField {
time: 0.0,
voltages: vec![-85.0, -85.0, -85.0],
},
VoltageField {
time: 10.0,
voltages: vec![20.0, -85.0, -85.0],
},
VoltageField {
time: 20.0,
voltages: vec![-20.0, 20.0, -85.0],
},
];
let map = compute_activation_map(&fields);
assert!((map.activation_times[0] - 10.0).abs() < 0.001);
assert!((map.activation_times[1] - 20.0).abs() < 0.001);
}
#[test]
fn test_detect_arrhythmias() {
let fields = vec![
VoltageField {
time: 0.0,
voltages: vec![-85.0; 100],
},
VoltageField {
time: 50.0,
voltages: vec![20.0; 100],
},
];
let activation_map = compute_activation_map(&fields);
let events = detect_arrhythmias(&fields, &activation_map);
assert!(!events.is_empty());
}
}