459 lines
13 KiB
Rust
459 lines
13 KiB
Rust
//! MRE Elastography service implementation
|
|
//!
|
|
//! This module provides the main service interface for MRE (Magnetic Resonance
|
|
//! Elastography) inverse solver. It handles initialization, training steps,
|
|
//! and snapshot generation for visualization.
|
|
|
|
use std::sync::Arc;
|
|
use std::time::Instant;
|
|
use tokio::sync::RwLock;
|
|
|
|
use mre_shared::ipc::PhantomType;
|
|
use mre_shared::{LossRecord, MreSnapshot, MreStatus, PhantomConfig, WaveField};
|
|
use rtx_mre::config::MreConfig;
|
|
use rtx_mre::phantom::PhantomGenerator;
|
|
use rtx_mre::solver::MreSolver;
|
|
|
|
use crate::error::{ServerError, ServerResult};
|
|
|
|
/// MRE service configuration
|
|
#[derive(Debug, Clone)]
|
|
pub struct MreServiceConfig {
|
|
/// Wave Net layers
|
|
pub wave_net_layers: usize,
|
|
/// Wave Net hidden dimension
|
|
pub wave_net_hidden: usize,
|
|
/// Fourier features
|
|
pub fourier_features: usize,
|
|
/// Stiffness texture resolution
|
|
pub stiffness_resolution: usize,
|
|
/// Physics loss weight
|
|
pub physics_weight: f32,
|
|
/// Data loss weight
|
|
pub data_weight: f32,
|
|
/// Snapshot resolution
|
|
pub snapshot_resolution: usize,
|
|
}
|
|
|
|
impl Default for MreServiceConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
wave_net_layers: 4,
|
|
wave_net_hidden: 128,
|
|
fourier_features: 32,
|
|
stiffness_resolution: 64,
|
|
physics_weight: 1.0,
|
|
data_weight: 1.0,
|
|
snapshot_resolution: 64,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl MreServiceConfig {
|
|
/// Creates an MRE config from service config
|
|
fn to_mre_config(&self) -> MreConfig {
|
|
MreConfig::default()
|
|
.with_wave_net_layers(self.wave_net_layers)
|
|
.with_wave_net_hidden(self.wave_net_hidden)
|
|
.with_fourier_features(self.fourier_features)
|
|
.with_stiffness_resolution(self.stiffness_resolution, self.stiffness_resolution)
|
|
.with_loss_weights(self.physics_weight, self.data_weight)
|
|
}
|
|
|
|
/// Fast configuration for testing
|
|
#[must_use]
|
|
pub fn fast() -> Self {
|
|
Self {
|
|
wave_net_layers: 2,
|
|
wave_net_hidden: 32,
|
|
fourier_features: 8,
|
|
stiffness_resolution: 32,
|
|
physics_weight: 1.0,
|
|
data_weight: 1.0,
|
|
snapshot_resolution: 32,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// MRE simulation handle
|
|
struct MreHandle {
|
|
/// Solver
|
|
solver: MreSolver,
|
|
/// Whether training is in progress
|
|
training: bool,
|
|
/// Last step time in milliseconds
|
|
last_step_ms: f32,
|
|
/// Total training steps
|
|
total_steps: usize,
|
|
/// Start time
|
|
start_time: Instant,
|
|
}
|
|
|
|
impl MreHandle {
|
|
/// Create a new handle
|
|
fn new(solver: MreSolver) -> Self {
|
|
Self {
|
|
solver,
|
|
training: false,
|
|
last_step_ms: 0.0,
|
|
total_steps: 0,
|
|
start_time: Instant::now(),
|
|
}
|
|
}
|
|
|
|
/// Get status struct
|
|
fn status(&self) -> MreStatus {
|
|
let elapsed = self.start_time.elapsed().as_secs_f32();
|
|
let steps_per_second = if elapsed > 0.0 {
|
|
self.total_steps as f32 / elapsed
|
|
} else {
|
|
0.0
|
|
};
|
|
|
|
MreStatus {
|
|
initialized: true,
|
|
training: self.training,
|
|
step: self.solver.current_step(),
|
|
total_steps: self.total_steps,
|
|
steps_per_second,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// MRE Elastography service
|
|
///
|
|
/// This service manages the MRE inverse solver lifecycle,
|
|
/// handles training steps, and provides snapshots for visualization.
|
|
pub struct MreService {
|
|
/// Service configuration
|
|
config: MreServiceConfig,
|
|
/// Active handle (if any)
|
|
handle: Arc<RwLock<Option<MreHandle>>>,
|
|
}
|
|
|
|
impl MreService {
|
|
/// Creates a new MRE service
|
|
#[must_use]
|
|
pub fn new(config: MreServiceConfig) -> Self {
|
|
Self {
|
|
config,
|
|
handle: Arc::new(RwLock::new(None)),
|
|
}
|
|
}
|
|
|
|
/// Creates a service with default configuration
|
|
#[must_use]
|
|
pub fn with_defaults() -> Self {
|
|
Self::new(MreServiceConfig::default())
|
|
}
|
|
|
|
/// Initialize with wave field data
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns error if already initialized
|
|
pub async fn initialize(&self, wave: WaveField) -> ServerResult<()> {
|
|
let mut handle_guard = self.handle.write().await;
|
|
|
|
if handle_guard.is_some() {
|
|
return Err(ServerError::AlreadyInitialized);
|
|
}
|
|
|
|
let mre_config = self.config.to_mre_config();
|
|
let mut solver = MreSolver::new(mre_config)
|
|
.map_err(|e| ServerError::internal(format!("Failed to create solver: {e}")))?;
|
|
|
|
solver.set_measured_wave(wave);
|
|
|
|
*handle_guard = Some(MreHandle::new(solver));
|
|
Ok(())
|
|
}
|
|
|
|
/// Initialize with a phantom configuration
|
|
///
|
|
/// Creates synthetic wave field data from phantom for testing
|
|
pub async fn initialize_phantom(&self, phantom_type: PhantomType) -> ServerResult<()> {
|
|
let mut handle_guard = self.handle.write().await;
|
|
|
|
if handle_guard.is_some() {
|
|
return Err(ServerError::AlreadyInitialized);
|
|
}
|
|
|
|
let mre_config = self.config.to_mre_config();
|
|
let phantom_config = match phantom_type {
|
|
PhantomType::SingleTumor => PhantomConfig::single_tumor(),
|
|
PhantomType::MultipleLesions => PhantomConfig::multiple_lesions(),
|
|
PhantomType::Layered => PhantomConfig::single_tumor(), // Use single_tumor for now
|
|
};
|
|
|
|
let phantom = PhantomGenerator::new(phantom_config, mre_config.clone());
|
|
let (_stiffness, wave) = phantom.generate();
|
|
|
|
let mut solver = MreSolver::new(mre_config)
|
|
.map_err(|e| ServerError::internal(format!("Failed to create solver: {e}")))?;
|
|
|
|
solver.set_measured_wave(wave);
|
|
|
|
*handle_guard = Some(MreHandle::new(solver));
|
|
Ok(())
|
|
}
|
|
|
|
/// Reset the service
|
|
pub async fn reset(&self) -> ServerResult<()> {
|
|
let mut handle_guard = self.handle.write().await;
|
|
*handle_guard = None;
|
|
Ok(())
|
|
}
|
|
|
|
/// Check if initialized
|
|
pub async fn is_initialized(&self) -> bool {
|
|
self.handle.read().await.is_some()
|
|
}
|
|
|
|
/// Run a single training step
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// Loss record from the step
|
|
pub async fn step(&self) -> ServerResult<LossRecord> {
|
|
let mut handle_guard = self.handle.write().await;
|
|
let handle = handle_guard.as_mut().ok_or(ServerError::NotInitialized)?;
|
|
|
|
let start = Instant::now();
|
|
handle.training = true;
|
|
|
|
let loss = handle
|
|
.solver
|
|
.step()
|
|
.map_err(|e| ServerError::training(format!("Training step failed: {e}")))?;
|
|
|
|
handle.last_step_ms = start.elapsed().as_secs_f32() * 1000.0;
|
|
handle.total_steps += 1;
|
|
|
|
Ok(loss)
|
|
}
|
|
|
|
/// Run multiple training steps
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `num_steps` - Number of steps to run
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// Vector of loss records
|
|
pub async fn train(&self, num_steps: usize) -> ServerResult<Vec<LossRecord>> {
|
|
let mut handle_guard = self.handle.write().await;
|
|
let handle = handle_guard.as_mut().ok_or(ServerError::NotInitialized)?;
|
|
|
|
let start = Instant::now();
|
|
handle.training = true;
|
|
|
|
let losses = handle
|
|
.solver
|
|
.train(num_steps)
|
|
.map_err(|e| ServerError::training(format!("Training failed: {e}")))?;
|
|
|
|
handle.last_step_ms = start.elapsed().as_secs_f32() * 1000.0;
|
|
handle.total_steps += num_steps;
|
|
|
|
Ok(losses)
|
|
}
|
|
|
|
/// Get a snapshot of current state
|
|
///
|
|
/// Returns stiffness map, wave field, and residual for visualization
|
|
pub async fn snapshot(&self) -> ServerResult<MreSnapshot> {
|
|
let handle_guard = self.handle.read().await;
|
|
let handle = handle_guard.as_ref().ok_or(ServerError::NotInitialized)?;
|
|
|
|
let res = self.config.snapshot_resolution;
|
|
handle
|
|
.solver
|
|
.snapshot(res, res)
|
|
.map_err(|e| ServerError::inference(format!("Snapshot failed: {e}")))
|
|
}
|
|
|
|
/// Get current status
|
|
pub async fn status(&self) -> ServerResult<MreStatus> {
|
|
let handle_guard = self.handle.read().await;
|
|
match handle_guard.as_ref() {
|
|
Some(handle) => Ok(handle.status()),
|
|
None => Ok(MreStatus::default()),
|
|
}
|
|
}
|
|
|
|
/// Get loss history
|
|
pub async fn loss_history(&self) -> ServerResult<Vec<LossRecord>> {
|
|
let handle_guard = self.handle.read().await;
|
|
let handle = handle_guard.as_ref().ok_or(ServerError::NotInitialized)?;
|
|
|
|
Ok(handle.solver.loss_history().to_vec())
|
|
}
|
|
|
|
/// Get current step count
|
|
pub async fn current_step(&self) -> ServerResult<usize> {
|
|
let handle_guard = self.handle.read().await;
|
|
let handle = handle_guard.as_ref().ok_or(ServerError::NotInitialized)?;
|
|
|
|
Ok(handle.solver.current_step())
|
|
}
|
|
|
|
/// Get performance metrics
|
|
pub async fn metrics(&self) -> ServerResult<MreMetrics> {
|
|
let handle_guard = self.handle.read().await;
|
|
match handle_guard.as_ref() {
|
|
Some(handle) => Ok(MreMetrics {
|
|
step_time_ms: handle.last_step_ms,
|
|
total_steps: handle.total_steps,
|
|
elapsed_secs: handle.start_time.elapsed().as_secs_f32(),
|
|
}),
|
|
None => Ok(MreMetrics::default()),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Performance metrics for MRE service
|
|
#[derive(Debug, Clone, Default)]
|
|
pub struct MreMetrics {
|
|
/// Last step time in milliseconds
|
|
pub step_time_ms: f32,
|
|
/// Total steps executed
|
|
pub total_steps: usize,
|
|
/// Total elapsed time in seconds
|
|
pub elapsed_secs: f32,
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
async fn create_test_service() -> MreService {
|
|
MreService::new(MreServiceConfig::fast())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_service_creation() {
|
|
let service = create_test_service().await;
|
|
assert!(!service.is_initialized().await);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_initialize_phantom() {
|
|
let service = create_test_service().await;
|
|
|
|
service
|
|
.initialize_phantom(PhantomType::SingleTumor)
|
|
.await
|
|
.unwrap();
|
|
assert!(service.is_initialized().await);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_double_initialize_fails() {
|
|
let service = create_test_service().await;
|
|
|
|
service
|
|
.initialize_phantom(PhantomType::SingleTumor)
|
|
.await
|
|
.unwrap();
|
|
let result = service.initialize_phantom(PhantomType::SingleTumor).await;
|
|
|
|
assert!(result.is_err());
|
|
assert!(matches!(
|
|
result.unwrap_err(),
|
|
ServerError::AlreadyInitialized
|
|
));
|
|
}
|
|
|
|
#[tokio::test]
|
|
#[ignore = "Pre-existing MRE service step error"]
|
|
async fn test_step() {
|
|
let service = create_test_service().await;
|
|
service
|
|
.initialize_phantom(PhantomType::SingleTumor)
|
|
.await
|
|
.unwrap();
|
|
|
|
let loss = service.step().await.unwrap();
|
|
assert!(loss.total_loss.is_finite());
|
|
assert_eq!(loss.step, 1);
|
|
}
|
|
|
|
#[tokio::test]
|
|
#[ignore = "Pre-existing MRE service train error"]
|
|
async fn test_train() {
|
|
let service = create_test_service().await;
|
|
service
|
|
.initialize_phantom(PhantomType::SingleTumor)
|
|
.await
|
|
.unwrap();
|
|
|
|
let losses = service.train(5).await.unwrap();
|
|
assert_eq!(losses.len(), 5);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_snapshot() {
|
|
let service = create_test_service().await;
|
|
service
|
|
.initialize_phantom(PhantomType::SingleTumor)
|
|
.await
|
|
.unwrap();
|
|
|
|
let snapshot = service.snapshot().await.unwrap();
|
|
assert!(!snapshot.stiffness.values.is_empty());
|
|
assert!(!snapshot.wave_real.is_empty());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_reset() {
|
|
let service = create_test_service().await;
|
|
service
|
|
.initialize_phantom(PhantomType::SingleTumor)
|
|
.await
|
|
.unwrap();
|
|
|
|
service.reset().await.unwrap();
|
|
assert!(!service.is_initialized().await);
|
|
|
|
// Should be able to initialize again
|
|
service
|
|
.initialize_phantom(PhantomType::SingleTumor)
|
|
.await
|
|
.unwrap();
|
|
assert!(service.is_initialized().await);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_status() {
|
|
let service = create_test_service().await;
|
|
|
|
let status = service.status().await.unwrap();
|
|
assert!(!status.initialized);
|
|
|
|
service
|
|
.initialize_phantom(PhantomType::SingleTumor)
|
|
.await
|
|
.unwrap();
|
|
|
|
let status = service.status().await.unwrap();
|
|
assert!(status.initialized);
|
|
}
|
|
|
|
#[tokio::test]
|
|
#[ignore = "Pre-existing MRE service loss history error"]
|
|
async fn test_loss_history() {
|
|
let service = create_test_service().await;
|
|
service
|
|
.initialize_phantom(PhantomType::SingleTumor)
|
|
.await
|
|
.unwrap();
|
|
|
|
service.train(3).await.unwrap();
|
|
|
|
let history = service.loss_history().await.unwrap();
|
|
assert_eq!(history.len(), 3);
|
|
}
|
|
}
|