314 lines
10 KiB
Rust
314 lines
10 KiB
Rust
//! MRE inverse solver combining Wave Net and Stiffness Texture
|
|
//!
|
|
//! The solver trains both components jointly to minimize:
|
|
//! - Data loss: Match the measured wave field
|
|
//! - Physics loss: Satisfy the Helmholtz equation
|
|
|
|
use crate::config::MreConfig;
|
|
use crate::helmholtz::{HelmholtzResidual, compute_data_loss};
|
|
use crate::phantom::{generate_grid_coords, generate_random_coords};
|
|
use crate::stiffness_texture::StiffnessTexture;
|
|
use crate::wave_net::WaveNet;
|
|
use anyhow::Result;
|
|
use mre_shared::{LossRecord, MreSnapshot, StiffnessField, WaveField};
|
|
use rtx_tensor::{Device, Tensor};
|
|
|
|
/// MRE inverse solver
|
|
pub struct MreSolver {
|
|
/// Configuration
|
|
config: MreConfig,
|
|
/// Wave field neural network
|
|
wave_net: WaveNet,
|
|
/// Learnable stiffness texture
|
|
stiffness: StiffnessTexture,
|
|
/// Helmholtz residual computer
|
|
helmholtz: HelmholtzResidual,
|
|
/// Measured wave field (input data)
|
|
measured_wave: Option<WaveField>,
|
|
/// Current training step
|
|
step: usize,
|
|
/// Loss history
|
|
loss_history: Vec<LossRecord>,
|
|
/// Device
|
|
device: Device,
|
|
}
|
|
|
|
impl MreSolver {
|
|
/// Create a new MRE solver
|
|
pub fn new(config: MreConfig) -> Result<Self> {
|
|
let device = Device::try_default()?; // Use default device (CUDA if available)
|
|
|
|
let wave_net = WaveNet::new(&config, &device)?;
|
|
let stiffness = StiffnessTexture::new(&config, &device)?;
|
|
let helmholtz = HelmholtzResidual::new(&config, &device);
|
|
|
|
Ok(Self {
|
|
config,
|
|
wave_net,
|
|
stiffness,
|
|
helmholtz,
|
|
measured_wave: None,
|
|
step: 0,
|
|
loss_history: Vec::new(),
|
|
device,
|
|
})
|
|
}
|
|
|
|
/// Initialize with measured wave field data
|
|
pub fn set_measured_wave(&mut self, wave: WaveField) {
|
|
self.measured_wave = Some(wave);
|
|
}
|
|
|
|
/// Initialize stiffness from a known field (for comparison)
|
|
pub fn set_initial_stiffness(&mut self, field: &StiffnessField) -> Result<()> {
|
|
self.stiffness = StiffnessTexture::from_field(field, &self.config, &self.device)?;
|
|
Ok(())
|
|
}
|
|
|
|
/// Run a single training step
|
|
///
|
|
/// # Returns
|
|
/// LossRecord containing total, data, and physics loss
|
|
pub fn step(&mut self) -> Result<LossRecord> {
|
|
let wave = self
|
|
.measured_wave
|
|
.as_ref()
|
|
.ok_or_else(|| anyhow::anyhow!("No measured wave field set"))?;
|
|
|
|
// 1. Sample collocation points for physics loss
|
|
let (x_coll, y_coll) = self.sample_collocation_points()?;
|
|
|
|
// 2. Forward pass through Wave Net with derivatives
|
|
let wave_derivs = self.wave_net.forward_with_derivs(&x_coll, &y_coll)?;
|
|
|
|
// 3. Sample stiffness and gradients at collocation points
|
|
let (mu, mu_x, mu_y) = self.stiffness.sample_with_gradients(&x_coll, &y_coll)?;
|
|
|
|
// 4. Compute physics residual
|
|
let (r_real, r_imag) = self
|
|
.helmholtz
|
|
.compute_residual(&wave_derivs, &mu, &mu_x, &mu_y)?;
|
|
let physics_loss = self.helmholtz.physics_loss(&r_real, &r_imag)?;
|
|
|
|
// 5. Sample data points and compute data loss
|
|
let (x_data, y_data, meas_real, meas_imag) = self.sample_data_points(wave)?;
|
|
let (pred_real, pred_imag) = self.wave_net.forward(&x_data, &y_data)?;
|
|
let data_loss = compute_data_loss(&pred_real, &pred_imag, &meas_real, &meas_imag)?;
|
|
|
|
// 6. Total loss (weighted sum)
|
|
let total_loss = self.config.physics_weight * physics_loss
|
|
+ self.config.data_weight * data_loss
|
|
+ self.config.tv_weight * self.stiffness.tv_loss()?;
|
|
|
|
// 7. Update stiffness texture using gradient descent
|
|
// Compute gradient of physics loss w.r.t stiffness
|
|
let stiffness_grad =
|
|
self.helmholtz
|
|
.compute_stiffness_gradient(&wave_derivs, &r_real, &r_imag)?;
|
|
|
|
// Scatter gradients back to texture
|
|
let texture_grad =
|
|
self.stiffness
|
|
.compute_texture_gradient(&x_coll, &y_coll, &stiffness_grad)?;
|
|
|
|
// Apply gradient update
|
|
self.stiffness
|
|
.apply_gradient(&texture_grad, self.config.learning_rate_stiffness)?;
|
|
|
|
// 8. Record loss
|
|
self.step += 1;
|
|
let record = LossRecord {
|
|
step: self.step,
|
|
total_loss,
|
|
data_loss,
|
|
physics_loss,
|
|
};
|
|
self.loss_history.push(record.clone());
|
|
|
|
Ok(record)
|
|
}
|
|
|
|
/// Run multiple training steps
|
|
pub fn train(&mut self, num_steps: usize) -> Result<Vec<LossRecord>> {
|
|
let mut losses = Vec::with_capacity(num_steps);
|
|
for _ in 0..num_steps {
|
|
losses.push(self.step()?);
|
|
}
|
|
Ok(losses)
|
|
}
|
|
|
|
/// Get current snapshot for visualization
|
|
pub fn snapshot(&self, grid_nx: usize, grid_ny: usize) -> Result<MreSnapshot> {
|
|
// Generate regular grid in non-dimensional coordinates
|
|
let (x_grid, y_grid) =
|
|
generate_grid_coords(grid_nx, grid_ny, 0.0, 1.0, 0.0, 1.0, &self.device)?;
|
|
|
|
// Evaluate wave field
|
|
let (wave_real, wave_imag) = self.wave_net.forward(&x_grid, &y_grid)?;
|
|
let wave_real_data = wave_real.to_cpu()?;
|
|
let wave_imag_data = wave_imag.to_cpu()?;
|
|
|
|
// Evaluate residual
|
|
let wave_derivs = self.wave_net.forward_with_derivs(&x_grid, &y_grid)?;
|
|
let (mu, mu_x, mu_y) = self.stiffness.sample_with_gradients(&x_grid, &y_grid)?;
|
|
let (r_real, r_imag) = self
|
|
.helmholtz
|
|
.compute_residual(&wave_derivs, &mu, &mu_x, &mu_y)?;
|
|
|
|
// Residual magnitude
|
|
let r_real_data = r_real.to_cpu()?;
|
|
let r_imag_data = r_imag.to_cpu()?;
|
|
let residual: Vec<f32> = r_real_data
|
|
.iter()
|
|
.zip(&r_imag_data)
|
|
.map(|(&r, &i)| (r * r + i * i).sqrt())
|
|
.collect();
|
|
|
|
// Get stiffness field
|
|
let stiffness = self.stiffness.to_field(&self.config)?;
|
|
|
|
// Get latest loss
|
|
let loss = self.loss_history.last().cloned().unwrap_or_default();
|
|
|
|
Ok(MreSnapshot {
|
|
stiffness,
|
|
wave_real: wave_real_data,
|
|
wave_imag: wave_imag_data,
|
|
residual,
|
|
grid_resolution: (grid_nx, grid_ny),
|
|
step: self.step,
|
|
loss,
|
|
})
|
|
}
|
|
|
|
/// Sample random collocation points for physics loss
|
|
fn sample_collocation_points(&self) -> Result<(Tensor, Tensor)> {
|
|
generate_random_coords(
|
|
self.config.num_collocation,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
1.0,
|
|
&self.device,
|
|
)
|
|
}
|
|
|
|
/// Sample data points from measured wave field
|
|
fn sample_data_points(&self, wave: &WaveField) -> Result<(Tensor, Tensor, Tensor, Tensor)> {
|
|
let n = wave.len().min(self.config.num_data_points);
|
|
|
|
// Random sample indices using rand crate
|
|
use rand::seq::SliceRandom;
|
|
let mut rng = rand::thread_rng();
|
|
let mut indices: Vec<usize> = (0..wave.len()).collect();
|
|
indices.shuffle(&mut rng);
|
|
indices.truncate(n);
|
|
|
|
// Extract points and values
|
|
let mut x_data = Vec::with_capacity(n);
|
|
let mut y_data = Vec::with_capacity(n);
|
|
let mut real_data = Vec::with_capacity(n);
|
|
let mut imag_data = Vec::with_capacity(n);
|
|
|
|
for &idx in &indices {
|
|
let pt = &wave.points[idx];
|
|
// Convert to non-dimensional
|
|
let x_nd = self.config.nondim.nondim_length(pt.x);
|
|
let y_nd = self.config.nondim.nondim_length(pt.y);
|
|
|
|
x_data.push(x_nd);
|
|
y_data.push(y_nd);
|
|
real_data.push(wave.real[idx]);
|
|
imag_data.push(wave.imag[idx]);
|
|
}
|
|
|
|
let x = Tensor::from_data(x_data, vec![n], &self.device)?;
|
|
let y = Tensor::from_data(y_data, vec![n], &self.device)?;
|
|
let real = Tensor::from_data(real_data, vec![n], &self.device)?;
|
|
let imag = Tensor::from_data(imag_data, vec![n], &self.device)?;
|
|
|
|
Ok((x, y, real, imag))
|
|
}
|
|
|
|
/// Get current training step
|
|
pub fn current_step(&self) -> usize {
|
|
self.step
|
|
}
|
|
|
|
/// Get loss history
|
|
pub fn loss_history(&self) -> &[LossRecord] {
|
|
&self.loss_history
|
|
}
|
|
|
|
/// Reset solver to initial state
|
|
pub fn reset(&mut self) -> Result<()> {
|
|
self.wave_net = WaveNet::new(&self.config, &self.device)?;
|
|
self.stiffness = StiffnessTexture::new(&self.config, &self.device)?;
|
|
self.step = 0;
|
|
self.loss_history.clear();
|
|
Ok(())
|
|
}
|
|
|
|
/// Get configuration
|
|
pub fn config(&self) -> &MreConfig {
|
|
&self.config
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::phantom::PhantomGenerator;
|
|
|
|
fn get_test_config() -> MreConfig {
|
|
MreConfig::fast()
|
|
.with_stiffness_resolution(16, 16)
|
|
.with_wave_net_layers(2)
|
|
.with_wave_net_hidden(32)
|
|
.with_fourier_features(8)
|
|
}
|
|
|
|
#[test]
|
|
fn test_solver_creation() {
|
|
let config = get_test_config();
|
|
let solver = MreSolver::new(config).unwrap();
|
|
assert_eq!(solver.current_step(), 0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_solver_with_phantom() {
|
|
let config = get_test_config();
|
|
let mut solver = MreSolver::new(config.clone()).unwrap();
|
|
|
|
// Generate phantom
|
|
let phantom = PhantomGenerator::tumor_phantom(config);
|
|
let (_stiffness, wave) = phantom.generate();
|
|
|
|
solver.set_measured_wave(wave);
|
|
|
|
// Run a few steps
|
|
for _ in 0..3 {
|
|
let loss = solver.step().unwrap();
|
|
assert!(loss.total_loss.is_finite());
|
|
}
|
|
|
|
assert_eq!(solver.current_step(), 3);
|
|
}
|
|
|
|
#[test]
|
|
fn test_snapshot() {
|
|
let config = get_test_config();
|
|
let mut solver = MreSolver::new(config.clone()).unwrap();
|
|
|
|
let phantom = PhantomGenerator::tumor_phantom(config);
|
|
let (_stiffness, wave) = phantom.generate();
|
|
solver.set_measured_wave(wave);
|
|
|
|
let snapshot = solver.snapshot(16, 16).unwrap();
|
|
|
|
assert_eq!(snapshot.grid_resolution, (16, 16));
|
|
assert_eq!(snapshot.wave_real.len(), 256);
|
|
assert_eq!(snapshot.stiffness.resolution, (16, 16));
|
|
}
|
|
}
|