osobhandClaude Opus 4.6
02d382d5f6
Consistent formatting pass: line wrapping, import sorting, trailing whitespace removal, let-chain indentation, merged derive attributes, and unsafe block reformatting. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
384 lines
11 KiB
Rust
384 lines
11 KiB
Rust
//! Physics-Informed Neural Operator (PINO) for cardiac electrophysiology.
|
|
//!
|
|
//! Implements a neural operator that learns the solution operator for the
|
|
//! monodomain equation with physics-based constraints.
|
|
|
|
use cardiosim_shared::{
|
|
HeartMesh, NeuralOperatorConfig, SimulationConfig, StimulationProtocol, VoltageField,
|
|
};
|
|
|
|
use crate::CardioSimError;
|
|
|
|
/// Physics-Informed Neural Operator for cardiac simulation.
|
|
#[derive(Debug)]
|
|
pub struct CardiacPINO {
|
|
config: NeuralOperatorConfig,
|
|
fno_layers: Vec<FourierLayer>,
|
|
lifting: MLPLayer,
|
|
projection: MLPLayer,
|
|
}
|
|
|
|
impl CardiacPINO {
|
|
/// Create a new cardiac PINO.
|
|
#[must_use]
|
|
pub fn new(config: NeuralOperatorConfig) -> Self {
|
|
let lifting = MLPLayer::new(config.input_channels, config.hidden_dim);
|
|
|
|
let fno_layers: Vec<FourierLayer> = (0..config.num_layers)
|
|
.map(|_| FourierLayer::new(config.hidden_dim, 16))
|
|
.collect();
|
|
|
|
let projection = MLPLayer::new(config.hidden_dim, config.output_channels);
|
|
|
|
Self {
|
|
config,
|
|
fno_layers,
|
|
lifting,
|
|
projection,
|
|
}
|
|
}
|
|
|
|
/// Solve using neural operator.
|
|
pub fn solve(
|
|
&self,
|
|
mesh: &HeartMesh,
|
|
protocol: &StimulationProtocol,
|
|
config: &SimulationConfig,
|
|
) -> Result<Vec<VoltageField>, CardioSimError> {
|
|
let n_vertices = mesh.vertices.len();
|
|
let n_outputs = (config.total_time / config.output_interval) as usize;
|
|
|
|
// Initialize input (voltage, recovery, coordinates)
|
|
let mut voltage = vec![-85.0_f32; n_vertices];
|
|
let recovery = vec![0.0_f32; n_vertices];
|
|
|
|
// Encode mesh coordinates
|
|
let coords: Vec<[f32; 3]> = mesh
|
|
.vertices
|
|
.iter()
|
|
.map(|v| [v.x / 100.0, v.y / 100.0, v.z / 100.0]) // Normalize
|
|
.collect();
|
|
|
|
let mut output = Vec::new();
|
|
|
|
// Simulate in chunks using neural operator
|
|
let _time_steps_per_inference = self.config.time_steps;
|
|
let mut t = 0.0;
|
|
|
|
for _ in 0..n_outputs {
|
|
// Prepare input tensor
|
|
let input = self.prepare_input(&voltage, &recovery, &coords, mesh, protocol, t);
|
|
|
|
// Forward through neural operator
|
|
let prediction = self.forward(&input);
|
|
|
|
// Extract voltage and recovery
|
|
let new_voltage: Vec<f32> = prediction.iter().take(n_vertices).copied().collect();
|
|
let _new_recovery: Vec<f32> = prediction
|
|
.iter()
|
|
.skip(n_vertices)
|
|
.take(n_vertices)
|
|
.copied()
|
|
.collect();
|
|
|
|
voltage = new_voltage;
|
|
|
|
output.push(VoltageField {
|
|
time: t,
|
|
voltages: voltage.clone(),
|
|
});
|
|
|
|
t += config.output_interval;
|
|
}
|
|
|
|
Ok(output)
|
|
}
|
|
|
|
fn prepare_input(
|
|
&self,
|
|
voltage: &[f32],
|
|
recovery: &[f32],
|
|
coords: &[[f32; 3]],
|
|
mesh: &HeartMesh,
|
|
protocol: &StimulationProtocol,
|
|
t: f32,
|
|
) -> Vec<f32> {
|
|
let n = voltage.len();
|
|
let mut input = Vec::with_capacity(n * self.config.input_channels);
|
|
|
|
for i in 0..n {
|
|
// Voltage (normalized)
|
|
input.push((voltage[i] + 85.0) / 120.0);
|
|
// Recovery
|
|
input.push(recovery[i]);
|
|
// Spatial coordinates
|
|
input.push(coords[i][0]);
|
|
input.push(coords[i][1]);
|
|
input.push(coords[i][2]);
|
|
// Fiber direction
|
|
input.push(mesh.fibers[i].x);
|
|
input.push(mesh.fibers[i].y);
|
|
input.push(mesh.fibers[i].z);
|
|
}
|
|
|
|
// Add stimulus encoding
|
|
for site in &protocol.sites {
|
|
for &stim_time in &site.times {
|
|
if t >= stim_time && t < stim_time + site.duration {
|
|
// Mark stimulated vertices
|
|
for (i, vertex) in mesh.vertices.iter().enumerate() {
|
|
let dist = vertex.distance_to(&site.center);
|
|
if dist < site.radius {
|
|
input[i * self.config.input_channels] += 0.5;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
input
|
|
}
|
|
|
|
fn forward(&self, input: &[f32]) -> Vec<f32> {
|
|
// Lifting layer
|
|
let mut hidden = self.lifting.forward(input);
|
|
|
|
// Fourier layers
|
|
for layer in &self.fno_layers {
|
|
hidden = layer.forward(&hidden);
|
|
}
|
|
|
|
// Projection layer
|
|
self.projection.forward(&hidden)
|
|
}
|
|
}
|
|
|
|
/// Fourier layer for FNO.
|
|
#[derive(Debug)]
|
|
struct FourierLayer {
|
|
hidden_dim: usize,
|
|
modes: usize,
|
|
weights_real: Vec<Vec<f32>>,
|
|
weights_imag: Vec<Vec<f32>>,
|
|
linear_weights: Vec<Vec<f32>>,
|
|
}
|
|
|
|
impl FourierLayer {
|
|
fn new(hidden_dim: usize, modes: usize) -> Self {
|
|
use rand::SeedableRng;
|
|
use rand_distr::{Distribution, Normal};
|
|
|
|
let mut rng = rand::rngs::StdRng::seed_from_u64(42);
|
|
let std = (2.0 / hidden_dim as f32).sqrt();
|
|
let normal = Normal::new(0.0_f32, std).unwrap();
|
|
|
|
let weights_real: Vec<Vec<f32>> = (0..modes)
|
|
.map(|_| (0..hidden_dim).map(|_| normal.sample(&mut rng)).collect())
|
|
.collect();
|
|
|
|
let weights_imag: Vec<Vec<f32>> = (0..modes)
|
|
.map(|_| (0..hidden_dim).map(|_| normal.sample(&mut rng)).collect())
|
|
.collect();
|
|
|
|
let linear_weights: Vec<Vec<f32>> = (0..hidden_dim)
|
|
.map(|_| (0..hidden_dim).map(|_| normal.sample(&mut rng)).collect())
|
|
.collect();
|
|
|
|
Self {
|
|
hidden_dim,
|
|
modes,
|
|
weights_real,
|
|
weights_imag,
|
|
linear_weights,
|
|
}
|
|
}
|
|
|
|
fn forward(&self, input: &[f32]) -> Vec<f32> {
|
|
let n = input.len() / self.hidden_dim;
|
|
if n == 0 {
|
|
return input.to_vec();
|
|
}
|
|
|
|
// Simplified spectral convolution
|
|
let mut output = vec![0.0_f32; input.len()];
|
|
|
|
// Linear path (skip connection)
|
|
for i in 0..n {
|
|
for j in 0..self.hidden_dim {
|
|
let input_val = input.get(i * self.hidden_dim + j).copied().unwrap_or(0.0);
|
|
|
|
for k in 0..self.hidden_dim {
|
|
let weight = self
|
|
.linear_weights
|
|
.get(j)
|
|
.and_then(|w| w.get(k))
|
|
.copied()
|
|
.unwrap_or(0.0);
|
|
output[i * self.hidden_dim + k] += input_val * weight;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Spectral path (simplified DFT-like operation)
|
|
for mode in 0..self.modes.min(n) {
|
|
let freq = 2.0 * std::f32::consts::PI * mode as f32 / n as f32;
|
|
|
|
for i in 0..n {
|
|
let phase = freq * i as f32;
|
|
let cos_p = phase.cos();
|
|
let sin_p = phase.sin();
|
|
|
|
for j in 0..self.hidden_dim {
|
|
let input_val = input.get(i * self.hidden_dim + j).copied().unwrap_or(0.0);
|
|
let weight_r = self
|
|
.weights_real
|
|
.get(mode)
|
|
.and_then(|w| w.get(j))
|
|
.copied()
|
|
.unwrap_or(0.0);
|
|
let weight_i = self
|
|
.weights_imag
|
|
.get(mode)
|
|
.and_then(|w| w.get(j))
|
|
.copied()
|
|
.unwrap_or(0.0);
|
|
|
|
output[i * self.hidden_dim + j] +=
|
|
input_val * (weight_r * cos_p - weight_i * sin_p);
|
|
}
|
|
}
|
|
}
|
|
|
|
// GELU activation
|
|
for v in &mut output {
|
|
*v = gelu(*v);
|
|
}
|
|
|
|
output
|
|
}
|
|
}
|
|
|
|
/// MLP layer.
|
|
#[derive(Debug)]
|
|
struct MLPLayer {
|
|
in_dim: usize,
|
|
out_dim: usize,
|
|
weights: Vec<Vec<f32>>,
|
|
bias: Vec<f32>,
|
|
}
|
|
|
|
impl MLPLayer {
|
|
fn new(in_dim: usize, out_dim: usize) -> Self {
|
|
use rand::SeedableRng;
|
|
use rand_distr::{Distribution, Normal};
|
|
|
|
let mut rng = rand::rngs::StdRng::seed_from_u64(123);
|
|
let std = (2.0 / in_dim as f32).sqrt();
|
|
let normal = Normal::new(0.0_f32, std).unwrap();
|
|
|
|
let weights: Vec<Vec<f32>> = (0..in_dim)
|
|
.map(|_| (0..out_dim).map(|_| normal.sample(&mut rng)).collect())
|
|
.collect();
|
|
|
|
let bias: Vec<f32> = (0..out_dim).map(|_| 0.0).collect();
|
|
|
|
Self {
|
|
in_dim,
|
|
out_dim,
|
|
weights,
|
|
bias,
|
|
}
|
|
}
|
|
|
|
fn forward(&self, input: &[f32]) -> Vec<f32> {
|
|
let batch_size = input.len() / self.in_dim;
|
|
if batch_size == 0 {
|
|
return vec![0.0; self.out_dim];
|
|
}
|
|
|
|
let mut output = vec![0.0_f32; batch_size * self.out_dim];
|
|
|
|
for b in 0..batch_size {
|
|
for j in 0..self.out_dim {
|
|
let mut sum = self.bias[j];
|
|
|
|
for i in 0..self.in_dim {
|
|
let input_idx = b * self.in_dim + i;
|
|
let input_val = input.get(input_idx).copied().unwrap_or(0.0);
|
|
let weight = self
|
|
.weights
|
|
.get(i)
|
|
.and_then(|w| w.get(j))
|
|
.copied()
|
|
.unwrap_or(0.0);
|
|
sum += input_val * weight;
|
|
}
|
|
|
|
output[b * self.out_dim + j] = sum;
|
|
}
|
|
}
|
|
|
|
output
|
|
}
|
|
}
|
|
|
|
/// GELU activation function.
|
|
fn gelu(x: f32) -> f32 {
|
|
0.5 * x * (1.0 + ((2.0_f32 / std::f32::consts::PI).sqrt() * (x + 0.044715 * x.powi(3))).tanh())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_pino_creation() {
|
|
let config = NeuralOperatorConfig::default();
|
|
let pino = CardiacPINO::new(config);
|
|
assert_eq!(pino.fno_layers.len(), 4);
|
|
}
|
|
|
|
#[test]
|
|
fn test_pino_solve() {
|
|
let config = NeuralOperatorConfig::default();
|
|
let pino = CardiacPINO::new(config);
|
|
|
|
let mesh = cardiosim_shared::get_sample_heart_mesh();
|
|
let protocol = cardiosim_shared::get_sample_protocol();
|
|
let sim_config = SimulationConfig {
|
|
total_time: 10.0,
|
|
output_interval: 5.0,
|
|
..Default::default()
|
|
};
|
|
|
|
let result = pino.solve(&mesh, &protocol, &sim_config);
|
|
assert!(result.is_ok());
|
|
|
|
let fields = result.unwrap();
|
|
assert!(!fields.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_fourier_layer() {
|
|
let layer = FourierLayer::new(32, 8);
|
|
let input = vec![0.1_f32; 32 * 10];
|
|
let output = layer.forward(&input);
|
|
assert_eq!(output.len(), input.len());
|
|
}
|
|
|
|
#[test]
|
|
fn test_mlp_layer() {
|
|
let layer = MLPLayer::new(16, 32);
|
|
let input = vec![0.1_f32; 16 * 5];
|
|
let output = layer.forward(&input);
|
|
assert_eq!(output.len(), 32 * 5);
|
|
}
|
|
|
|
#[test]
|
|
fn test_gelu() {
|
|
assert!((gelu(0.0) - 0.0).abs() < 0.01);
|
|
assert!(gelu(2.0) > 1.9);
|
|
assert!(gelu(-2.0) < 0.1);
|
|
}
|
|
}
|