Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,423 @@
|
||||
//! Learnable stiffness texture for MRE inverse problem
|
||||
//!
|
||||
//! The stiffness field mu(x,y) is represented as a 2D texture (grid) rather
|
||||
//! than a neural network. This allows for sharp discontinuities (tumor edges)
|
||||
//! that neural networks tend to over-smooth.
|
||||
//!
|
||||
//! Key features:
|
||||
//! - Bilinear interpolation for smooth sampling between grid points
|
||||
//! - Finite difference gradients (mu_x, mu_y) for physics loss
|
||||
//! - Direct gradient descent on texture pixels
|
||||
|
||||
use crate::config::MreConfig;
|
||||
use anyhow::Result;
|
||||
use mre_shared::StiffnessField;
|
||||
use rtx_tensor::{Device, Tensor};
|
||||
|
||||
/// Learnable stiffness texture
|
||||
pub struct StiffnessTexture {
|
||||
/// Stiffness values [ny, nx] in non-dimensional units
|
||||
values: Tensor,
|
||||
/// Grid dimensions
|
||||
nx: usize,
|
||||
ny: usize,
|
||||
/// Physical domain bounds (non-dimensional: [0, 1] x [0, 1])
|
||||
x_min: f32,
|
||||
x_max: f32,
|
||||
y_min: f32,
|
||||
y_max: f32,
|
||||
/// Minimum allowed stiffness (non-dimensional)
|
||||
min_stiffness: f32,
|
||||
/// Device
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl StiffnessTexture {
|
||||
/// Create a new stiffness texture with uniform initial value
|
||||
pub fn new(config: &MreConfig, device: &Device) -> Result<Self> {
|
||||
let nx = config.stiffness_nx;
|
||||
let ny = config.stiffness_ny;
|
||||
|
||||
// Initialize with uniform value = 1.0 (reference stiffness in non-dim)
|
||||
let init_value = 1.0;
|
||||
let n = nx * ny;
|
||||
let data = vec![init_value; n];
|
||||
let values = Tensor::from_data(data, vec![ny, nx], device)?;
|
||||
|
||||
Ok(Self {
|
||||
values,
|
||||
nx,
|
||||
ny,
|
||||
x_min: 0.0,
|
||||
x_max: 1.0, // Non-dimensional domain
|
||||
y_min: 0.0,
|
||||
y_max: 1.0,
|
||||
min_stiffness: config.min_stiffness,
|
||||
device: device.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Initialize from a ground truth stiffness field
|
||||
pub fn from_field(field: &StiffnessField, config: &MreConfig, device: &Device) -> Result<Self> {
|
||||
let (nx, ny) = field.resolution;
|
||||
|
||||
// Convert to non-dimensional values
|
||||
let values_nondim: Vec<f32> = field
|
||||
.values
|
||||
.iter()
|
||||
.map(|&v| config.nondim.nondim_stiffness_kpa(v))
|
||||
.collect();
|
||||
|
||||
let values = Tensor::from_data(values_nondim, vec![ny, nx], device)?;
|
||||
|
||||
Ok(Self {
|
||||
values,
|
||||
nx,
|
||||
ny,
|
||||
x_min: 0.0,
|
||||
x_max: 1.0,
|
||||
y_min: 0.0,
|
||||
y_max: 1.0,
|
||||
min_stiffness: config.min_stiffness,
|
||||
device: device.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Sample stiffness at arbitrary points using bilinear interpolation
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `x` - x-coordinates [batch] (non-dimensional, in [0, 1])
|
||||
/// * `y` - y-coordinates [batch]
|
||||
///
|
||||
/// # Returns
|
||||
/// Tensor [batch] with interpolated stiffness values
|
||||
pub fn sample(&self, x: &Tensor, y: &Tensor) -> Result<Tensor> {
|
||||
let batch_size = x.shape().dims()[0];
|
||||
let x_data = x.to_cpu()?;
|
||||
let y_data = y.to_cpu()?;
|
||||
let values_data = self.values.to_cpu()?;
|
||||
|
||||
let mut result = Vec::with_capacity(batch_size);
|
||||
|
||||
for i in 0..batch_size {
|
||||
let xi = x_data[i];
|
||||
let yi = y_data[i];
|
||||
|
||||
let mu = self.bilinear_interpolate(&values_data, xi, yi);
|
||||
result.push(mu.max(self.min_stiffness));
|
||||
}
|
||||
|
||||
Ok(Tensor::from_data(result, vec![batch_size], &self.device)?)
|
||||
}
|
||||
|
||||
/// Sample stiffness and its spatial gradients
|
||||
///
|
||||
/// Uses finite differences on the texture for gradient computation.
|
||||
///
|
||||
/// # Returns
|
||||
/// (mu, mu_x, mu_y) - stiffness and its gradients
|
||||
pub fn sample_with_gradients(
|
||||
&self,
|
||||
x: &Tensor,
|
||||
y: &Tensor,
|
||||
) -> Result<(Tensor, Tensor, Tensor)> {
|
||||
let batch_size = x.shape().dims()[0];
|
||||
let x_data = x.to_cpu()?;
|
||||
let y_data = y.to_cpu()?;
|
||||
let values_data = self.values.to_cpu()?;
|
||||
|
||||
let dx = (self.x_max - self.x_min) / (self.nx - 1).max(1) as f32;
|
||||
let dy = (self.y_max - self.y_min) / (self.ny - 1).max(1) as f32;
|
||||
|
||||
let mut mu_result = Vec::with_capacity(batch_size);
|
||||
let mut mu_x_result = Vec::with_capacity(batch_size);
|
||||
let mut mu_y_result = Vec::with_capacity(batch_size);
|
||||
|
||||
for i in 0..batch_size {
|
||||
let xi = x_data[i];
|
||||
let yi = y_data[i];
|
||||
|
||||
// Sample mu at current point
|
||||
let mu = self.bilinear_interpolate(&values_data, xi, yi);
|
||||
|
||||
// Central difference for gradients
|
||||
let mu_xp = self.bilinear_interpolate(&values_data, xi + dx, yi);
|
||||
let mu_xm = self.bilinear_interpolate(&values_data, xi - dx, yi);
|
||||
let mu_yp = self.bilinear_interpolate(&values_data, xi, yi + dy);
|
||||
let mu_ym = self.bilinear_interpolate(&values_data, xi, yi - dy);
|
||||
|
||||
let mu_x = (mu_xp - mu_xm) / (2.0 * dx);
|
||||
let mu_y = (mu_yp - mu_ym) / (2.0 * dy);
|
||||
|
||||
mu_result.push(mu.max(self.min_stiffness));
|
||||
mu_x_result.push(mu_x);
|
||||
mu_y_result.push(mu_y);
|
||||
}
|
||||
|
||||
let mu = Tensor::from_data(mu_result, vec![batch_size], &self.device)?;
|
||||
let mu_x = Tensor::from_data(mu_x_result, vec![batch_size], &self.device)?;
|
||||
let mu_y = Tensor::from_data(mu_y_result, vec![batch_size], &self.device)?;
|
||||
|
||||
Ok((mu, mu_x, mu_y))
|
||||
}
|
||||
|
||||
/// Bilinear interpolation helper
|
||||
fn bilinear_interpolate(&self, data: &[f32], x: f32, y: f32) -> f32 {
|
||||
// Convert to texture coordinates
|
||||
let tx = (x - self.x_min) / (self.x_max - self.x_min);
|
||||
let ty = (y - self.y_min) / (self.y_max - self.y_min);
|
||||
|
||||
// Scale to grid indices
|
||||
let fx = tx * (self.nx - 1) as f32;
|
||||
let fy = ty * (self.ny - 1) as f32;
|
||||
|
||||
// Floor indices
|
||||
let ix0 = (fx.floor() as usize).min(self.nx - 1);
|
||||
let iy0 = (fy.floor() as usize).min(self.ny - 1);
|
||||
let ix1 = (ix0 + 1).min(self.nx - 1);
|
||||
let iy1 = (iy0 + 1).min(self.ny - 1);
|
||||
|
||||
// Fractional parts (weights)
|
||||
let wx1 = fx - ix0 as f32;
|
||||
let wy1 = fy - iy0 as f32;
|
||||
let wx0 = 1.0 - wx1;
|
||||
let wy0 = 1.0 - wy1;
|
||||
|
||||
// Gather values at four corners
|
||||
let v00 = data[iy0 * self.nx + ix0];
|
||||
let v01 = data[iy0 * self.nx + ix1];
|
||||
let v10 = data[iy1 * self.nx + ix0];
|
||||
let v11 = data[iy1 * self.nx + ix1];
|
||||
|
||||
// Bilinear interpolation
|
||||
v00 * wx0 * wy0 + v01 * wx1 * wy0 + v10 * wx0 * wy1 + v11 * wx1 * wy1
|
||||
}
|
||||
|
||||
/// Apply gradient update to texture values
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `grad` - Gradient tensor [ny, nx]
|
||||
/// * `lr` - Learning rate
|
||||
pub fn apply_gradient(&mut self, grad: &Tensor, lr: f32) -> Result<()> {
|
||||
// values = values - lr * grad
|
||||
let update = grad.mul_scalar(lr)?;
|
||||
self.values = self.values.sub(&update)?;
|
||||
|
||||
// Clamp to minimum stiffness
|
||||
self.clamp_values()?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Compute gradient of stiffness texture from physics residual
|
||||
///
|
||||
/// Given the physics residual R and its gradient dR/dmu, compute
|
||||
/// the gradient dL/d(texture) by scattering residual contributions
|
||||
/// back to texture pixels.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `x` - sample x-coordinates [batch]
|
||||
/// * `y` - sample y-coordinates [batch]
|
||||
/// * `grad_mu` - gradient dL/dmu at each sample point [batch]
|
||||
pub fn compute_texture_gradient(
|
||||
&self,
|
||||
x: &Tensor,
|
||||
y: &Tensor,
|
||||
grad_mu: &Tensor,
|
||||
) -> Result<Tensor> {
|
||||
let batch_size = x.shape().dims()[0];
|
||||
let x_data = x.to_cpu()?;
|
||||
let y_data = y.to_cpu()?;
|
||||
let grad_data = grad_mu.to_cpu()?;
|
||||
|
||||
// Accumulate gradients using bilinear splatting (inverse of interpolation)
|
||||
let mut texture_grad = vec![0.0f32; self.ny * self.nx];
|
||||
|
||||
for i in 0..batch_size {
|
||||
let xi = x_data[i];
|
||||
let yi = y_data[i];
|
||||
let grad = grad_data[i];
|
||||
|
||||
// Convert to texture coordinates
|
||||
let tx = (xi - self.x_min) / (self.x_max - self.x_min);
|
||||
let ty = (yi - self.y_min) / (self.y_max - self.y_min);
|
||||
|
||||
let fx = tx * (self.nx - 1) as f32;
|
||||
let fy = ty * (self.ny - 1) as f32;
|
||||
|
||||
let ix0 = (fx.floor() as usize).min(self.nx - 1);
|
||||
let iy0 = (fy.floor() as usize).min(self.ny - 1);
|
||||
let ix1 = (ix0 + 1).min(self.nx - 1);
|
||||
let iy1 = (iy0 + 1).min(self.ny - 1);
|
||||
|
||||
let wx1 = fx - ix0 as f32;
|
||||
let wy1 = fy - iy0 as f32;
|
||||
let wx0 = 1.0 - wx1;
|
||||
let wy0 = 1.0 - wy1;
|
||||
|
||||
// Splat gradient to four corners (inverse bilinear)
|
||||
texture_grad[iy0 * self.nx + ix0] += grad * wx0 * wy0;
|
||||
texture_grad[iy0 * self.nx + ix1] += grad * wx1 * wy0;
|
||||
texture_grad[iy1 * self.nx + ix0] += grad * wx0 * wy1;
|
||||
texture_grad[iy1 * self.nx + ix1] += grad * wx1 * wy1;
|
||||
}
|
||||
|
||||
Ok(Tensor::from_data(
|
||||
texture_grad,
|
||||
vec![self.ny, self.nx],
|
||||
&self.device,
|
||||
)?)
|
||||
}
|
||||
|
||||
/// Clamp values to minimum stiffness
|
||||
fn clamp_values(&mut self) -> Result<()> {
|
||||
let data = self.values.to_cpu()?;
|
||||
let clamped: Vec<f32> = data.iter().map(|&v| v.max(self.min_stiffness)).collect();
|
||||
self.values = Tensor::from_data(clamped, vec![self.ny, self.nx], &self.device)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Convert to StiffnessField for visualization
|
||||
pub fn to_field(&self, config: &MreConfig) -> Result<StiffnessField> {
|
||||
let data = self.values.to_cpu()?;
|
||||
|
||||
// Convert from non-dim to kPa
|
||||
let values_kpa: Vec<f32> = data
|
||||
.iter()
|
||||
.map(|&v| config.nondim.dim_stiffness_kpa(v))
|
||||
.collect();
|
||||
|
||||
Ok(StiffnessField {
|
||||
resolution: (self.nx, self.ny),
|
||||
values: values_kpa,
|
||||
bounds: (
|
||||
config.nondim.dim_length(self.x_min),
|
||||
config.nondim.dim_length(self.x_max),
|
||||
config.nondim.dim_length(self.y_min),
|
||||
config.nondim.dim_length(self.y_max),
|
||||
),
|
||||
})
|
||||
}
|
||||
|
||||
/// Get grid dimensions
|
||||
pub fn resolution(&self) -> (usize, usize) {
|
||||
(self.nx, self.ny)
|
||||
}
|
||||
|
||||
/// Get raw values tensor
|
||||
pub fn values(&self) -> &Tensor {
|
||||
&self.values
|
||||
}
|
||||
|
||||
/// Compute Total Variation (TV) regularization loss
|
||||
///
|
||||
/// TV = sum(|mu(i+1,j) - mu(i,j)| + |mu(i,j+1) - mu(i,j)|)
|
||||
///
|
||||
/// Encourages piecewise-constant stiffness maps (sharp edges)
|
||||
pub fn tv_loss(&self) -> Result<f32> {
|
||||
let data = self.values.to_cpu()?;
|
||||
let mut tv = 0.0f32;
|
||||
|
||||
for j in 0..self.ny {
|
||||
for i in 0..self.nx {
|
||||
let idx = j * self.nx + i;
|
||||
let v = data[idx];
|
||||
|
||||
// Horizontal difference
|
||||
if i + 1 < self.nx {
|
||||
let v_right = data[idx + 1];
|
||||
tv += (v_right - v).abs();
|
||||
}
|
||||
|
||||
// Vertical difference
|
||||
if j + 1 < self.ny {
|
||||
let v_below = data[idx + self.nx];
|
||||
tv += (v_below - v).abs();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(tv)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn get_test_config() -> MreConfig {
|
||||
MreConfig::fast()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_texture_creation() {
|
||||
let config = get_test_config();
|
||||
let device = Device::try_default().unwrap();
|
||||
let texture = StiffnessTexture::new(&config, &device).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
texture.resolution(),
|
||||
(config.stiffness_nx, config.stiffness_ny)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_uniform_sampling() {
|
||||
let config = get_test_config();
|
||||
let device = Device::try_default().unwrap();
|
||||
let texture = StiffnessTexture::new(&config, &device).unwrap();
|
||||
|
||||
let x = Tensor::from_data(vec![0.25, 0.5, 0.75], vec![3], &device).unwrap();
|
||||
let y = Tensor::from_data(vec![0.25, 0.5, 0.75], vec![3], &device).unwrap();
|
||||
|
||||
let mu = texture.sample(&x, &y).unwrap();
|
||||
let mu_data = mu.to_cpu().unwrap();
|
||||
|
||||
// All should be 1.0 (initial uniform value)
|
||||
for val in &mu_data {
|
||||
assert!((*val - 1.0).abs() < 1e-5);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_bilinear_interpolation() {
|
||||
let config = MreConfig::fast().with_stiffness_resolution(3, 3);
|
||||
let device = Device::try_default().unwrap();
|
||||
|
||||
// Create texture with known pattern
|
||||
let mut texture = StiffnessTexture::new(&config, &device).unwrap();
|
||||
|
||||
// Set corners to different values
|
||||
// [0,0]=1, [0,2]=2, [2,0]=3, [2,2]=4
|
||||
let values = vec![1.0, 1.5, 2.0, 2.0, 2.5, 3.0, 3.0, 3.5, 4.0];
|
||||
texture.values = Tensor::from_data(values, vec![3, 3], &device).unwrap();
|
||||
|
||||
// Sample at center (0.5, 0.5) should interpolate
|
||||
let x = Tensor::from_data(vec![0.5], vec![1], &device).unwrap();
|
||||
let y = Tensor::from_data(vec![0.5], vec![1], &device).unwrap();
|
||||
let mu = texture.sample(&x, &y).unwrap();
|
||||
let mu_val = mu.to_cpu().unwrap()[0];
|
||||
|
||||
// Center value should be 2.5
|
||||
assert!((mu_val - 2.5).abs() < 1e-5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_gradient_computation() {
|
||||
let config = get_test_config();
|
||||
let device = Device::try_default().unwrap();
|
||||
let texture = StiffnessTexture::new(&config, &device).unwrap();
|
||||
|
||||
let x = Tensor::from_data(vec![0.5], vec![1], &device).unwrap();
|
||||
let y = Tensor::from_data(vec![0.5], vec![1], &device).unwrap();
|
||||
|
||||
let (mu, mu_x, mu_y) = texture.sample_with_gradients(&x, &y).unwrap();
|
||||
|
||||
// For uniform texture, gradients should be ~0
|
||||
let mu_x_val = mu_x.to_cpu().unwrap()[0];
|
||||
let mu_y_val = mu_y.to_cpu().unwrap()[0];
|
||||
|
||||
assert!(mu_x_val.abs() < 1e-5);
|
||||
assert!(mu_y_val.abs() < 1e-5);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user