rustytorch/demos/mre-shared/src/fields.rs

//! Field types for wave and stiffness data

use crate::geometry::Point2D;
use serde::{Deserialize, Serialize};

/// Complex wave field u(x,y) measured from MRI
/// Represents tissue displacement under mechanical excitation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WaveField {
    /// Sample points in physical coordinates
    pub points: Vec<Point2D>,
    /// Real component of complex wave displacement
    pub real: Vec<f32>,
    /// Imaginary component of complex wave displacement
    pub imag: Vec<f32>,
}

impl WaveField {
    /// Create a new wave field with the given data
    pub fn new(points: Vec<Point2D>, real: Vec<f32>, imag: Vec<f32>) -> Self {
        debug_assert_eq!(points.len(), real.len());
        debug_assert_eq!(points.len(), imag.len());
        Self { points, real, imag }
    }

    /// Number of sample points
    #[must_use]
    pub fn len(&self) -> usize {
        self.points.len()
    }

    /// Check if empty
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.points.is_empty()
    }

    /// Get wave magnitude |u| = sqrt(real^2 + imag^2)
    #[must_use]
    pub fn magnitude(&self) -> Vec<f32> {
        self.real
            .iter()
            .zip(&self.imag)
            .map(|(&r, &i)| (r * r + i * i).sqrt())
            .collect()
    }

    /// Get wave phase angle
    #[must_use]
    pub fn phase(&self) -> Vec<f32> {
        self.real
            .iter()
            .zip(&self.imag)
            .map(|(&r, &i)| i.atan2(r))
            .collect()
    }

    /// Create a regular grid wave field for testing
    #[must_use]
    pub fn regular_grid(nx: usize, ny: usize, domain: (f32, f32)) -> Self {
        let (w, h) = domain;
        let dx = w / (nx - 1) as f32;
        let dy = h / (ny - 1) as f32;

        let mut points = Vec::with_capacity(nx * ny);
        for j in 0..ny {
            for i in 0..nx {
                points.push(Point2D::new(i as f32 * dx, j as f32 * dy));
            }
        }

        let n = nx * ny;
        Self {
            points,
            real: vec![0.0; n],
            imag: vec![0.0; n],
        }
    }
}

/// Stiffness field mu(x,y) - the unknown we're solving for
/// Represents tissue shear modulus in kPa
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StiffnessField {
    /// Grid resolution (nx, ny)
    pub resolution: (usize, usize),
    /// Stiffness values in row-major order [ny][nx], in kPa
    pub values: Vec<f32>,
    /// Physical bounds (x_min, x_max, y_min, y_max)
    pub bounds: (f32, f32, f32, f32),
}

impl StiffnessField {
    /// Create a uniform stiffness field
    #[must_use]
    pub fn uniform(resolution: (usize, usize), bounds: (f32, f32, f32, f32), value: f32) -> Self {
        let n = resolution.0 * resolution.1;
        Self {
            resolution,
            values: vec![value; n],
            bounds,
        }
    }

    /// Get stiffness at grid index
    #[must_use]
    pub fn at(&self, i: usize, j: usize) -> f32 {
        let (nx, _ny) = self.resolution;
        self.values[j * nx + i]
    }

    /// Set stiffness at grid index
    pub fn set(&mut self, i: usize, j: usize, value: f32) {
        let (nx, _ny) = self.resolution;
        self.values[j * nx + i] = value;
    }

    /// Get grid spacing
    #[must_use]
    pub fn spacing(&self) -> (f32, f32) {
        let (nx, ny) = self.resolution;
        let (x_min, x_max, y_min, y_max) = self.bounds;
        let dx = (x_max - x_min) / (nx - 1).max(1) as f32;
        let dy = (y_max - y_min) / (ny - 1).max(1) as f32;
        (dx, dy)
    }

    /// Get physical coordinates for grid index
    #[must_use]
    pub fn coords_at(&self, i: usize, j: usize) -> Point2D {
        let (dx, dy) = self.spacing();
        let (x_min, _, y_min, _) = self.bounds;
        Point2D::new(x_min + i as f32 * dx, y_min + j as f32 * dy)
    }

    /// Get min and max stiffness values
    #[must_use]
    pub fn min_max(&self) -> (f32, f32) {
        let min = self.values.iter().copied().fold(f32::INFINITY, f32::min);
        let max = self
            .values
            .iter()
            .copied()
            .fold(f32::NEG_INFINITY, f32::max);
        (min, max)
    }

    /// Total number of grid points
    #[must_use]
    pub fn len(&self) -> usize {
        self.resolution.0 * self.resolution.1
    }

    /// Check if empty
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.values.is_empty()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_wave_field_magnitude() {
        let field = WaveField::new(vec![Point2D::new(0.0, 0.0)], vec![3.0], vec![4.0]);
        let mag = field.magnitude();
        assert!((mag[0] - 5.0).abs() < 1e-6);
    }

    #[test]
    fn test_stiffness_field() {
        let mut field = StiffnessField::uniform((10, 10), (0.0, 1.0, 0.0, 1.0), 3.0);
        assert!((field.at(5, 5) - 3.0).abs() < 1e-6);

        field.set(5, 5, 10.0);
        assert!((field.at(5, 5) - 10.0).abs() < 1e-6);
    }

    #[test]
    fn test_spacing() {
        let field = StiffnessField::uniform((11, 11), (0.0, 1.0, 0.0, 1.0), 3.0);
        let (dx, dy) = field.spacing();
        assert!((dx - 0.1).abs() < 1e-6);
        assert!((dy - 0.1).abs() < 1e-6);
    }
}