499 lines
16 KiB
Rust
499 lines
16 KiB
Rust
//! Physics-Informed Denoising Diffusion Model (PIDDM).
|
||
//!
|
||
//! This module implements the core PIDDM model that combines diffusion
|
||
//! models with physics-informed constraints for PDE solving.
|
||
|
||
use rand::Rng;
|
||
use rand::prelude::*;
|
||
use thiserror::Error;
|
||
|
||
use rtx_backend::Backend;
|
||
use rtx_tensor::GenericTensor;
|
||
|
||
use crate::scheduler::{DDPMScheduler, NoiseScheduler};
|
||
use crate::unet::DiffusionUNet;
|
||
|
||
/// Errors that can occur in PIDDM operations.
|
||
#[derive(Error, Debug)]
|
||
pub enum PIDDMError {
|
||
/// Invalid configuration
|
||
#[error("Invalid configuration: {0}")]
|
||
InvalidConfig(String),
|
||
|
||
/// Shape mismatch
|
||
#[error("Shape mismatch: expected {expected}, got {got}")]
|
||
ShapeMismatch {
|
||
/// The expected shape description
|
||
expected: String,
|
||
/// The actual shape received
|
||
got: String,
|
||
},
|
||
|
||
/// Training error
|
||
#[error("Training error: {0}")]
|
||
TrainingError(String),
|
||
|
||
/// Sampling error
|
||
#[error("Sampling error: {0}")]
|
||
SamplingError(String),
|
||
}
|
||
|
||
/// Configuration for PIDDM model.
|
||
#[derive(Debug, Clone)]
|
||
pub struct PIDDMConfig {
|
||
/// Weight for physics loss (λ in combined loss = L_diffusion + λ * L_physics)
|
||
pub physics_weight: f32,
|
||
/// Whether to use physics loss during sampling (physics guidance)
|
||
pub physics_guided_sampling: bool,
|
||
/// Guidance strength during sampling (if physics_guided_sampling is true)
|
||
pub guidance_strength: f32,
|
||
/// Whether to clamp predictions to [-1, 1]
|
||
pub clamp_predictions: bool,
|
||
}
|
||
|
||
impl Default for PIDDMConfig {
|
||
fn default() -> Self {
|
||
Self {
|
||
physics_weight: 0.1,
|
||
physics_guided_sampling: false,
|
||
guidance_strength: 0.0,
|
||
clamp_predictions: true,
|
||
}
|
||
}
|
||
}
|
||
|
||
impl PIDDMConfig {
|
||
/// Create config optimized for PDE solving.
|
||
pub fn for_pde() -> Self {
|
||
Self {
|
||
physics_weight: 0.5, // Higher physics weight for PDE applications
|
||
physics_guided_sampling: true,
|
||
guidance_strength: 0.1,
|
||
clamp_predictions: false, // PDEs may have unbounded solutions
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Physics-Informed Denoising Diffusion Model.
|
||
///
|
||
/// PIDDM combines diffusion models with physics constraints by:
|
||
/// 1. Adding physics residual loss during training
|
||
/// 2. Optionally using physics guidance during sampling
|
||
///
|
||
/// # Training Loss
|
||
///
|
||
/// The training loss combines diffusion loss and physics loss:
|
||
/// ```text
|
||
/// L = L_diffusion + λ * L_physics
|
||
/// = ||ε - ε_θ(x_t, t)||² + λ * ||L[x̂_0] - f||²
|
||
/// ```
|
||
///
|
||
/// where:
|
||
/// - ε is the noise added to the clean sample
|
||
/// - ε_θ is the predicted noise from the UNet
|
||
/// - x̂_0 is the predicted clean sample
|
||
/// - L[·] is the differential operator (e.g., Laplacian)
|
||
/// - f is the forcing term
|
||
pub struct PIDDM<B: Backend<FloatElem = f32>> {
|
||
/// Noise scheduler
|
||
scheduler: DDPMScheduler,
|
||
/// UNet noise predictor
|
||
unet: DiffusionUNet<B>,
|
||
/// Configuration
|
||
config: PIDDMConfig,
|
||
}
|
||
|
||
impl<B: Backend<FloatElem = f32>> PIDDM<B> {
|
||
/// Create a new PIDDM model.
|
||
pub fn new(scheduler: DDPMScheduler, unet: DiffusionUNet<B>, config: PIDDMConfig) -> Self {
|
||
Self {
|
||
scheduler,
|
||
unet,
|
||
config,
|
||
}
|
||
}
|
||
|
||
/// Get reference to the scheduler.
|
||
pub fn scheduler(&self) -> &DDPMScheduler {
|
||
&self.scheduler
|
||
}
|
||
|
||
/// Get reference to the UNet.
|
||
pub fn unet(&self) -> &DiffusionUNet<B> {
|
||
&self.unet
|
||
}
|
||
|
||
/// Get reference to the configuration.
|
||
pub fn config(&self) -> &PIDDMConfig {
|
||
&self.config
|
||
}
|
||
|
||
/// Perform a training step.
|
||
///
|
||
/// # Arguments
|
||
/// * `x0` - Clean samples [batch, channels, height, width]
|
||
/// * `physics_fn` - Function that computes physics residual from predicted x0
|
||
/// * `device` - Device
|
||
///
|
||
/// # Returns
|
||
/// Tuple of (total_loss, diffusion_loss, physics_loss)
|
||
pub fn training_step<F>(
|
||
&self,
|
||
x0: &GenericTensor<B, 4>,
|
||
physics_fn: F,
|
||
device: &B::Device,
|
||
) -> (f32, f32, f32)
|
||
where
|
||
F: Fn(&GenericTensor<B, 4>) -> GenericTensor<B, 4>,
|
||
{
|
||
let mut rng = thread_rng();
|
||
|
||
// Sample random timestep
|
||
let t = rng.gen_range(0..self.scheduler.num_timesteps());
|
||
|
||
// Sample noise
|
||
let noise = GenericTensor::<B, 4>::randn(x0.shape(), device);
|
||
|
||
// Create noisy sample
|
||
let xt = self.scheduler.add_noise(x0, &noise, t);
|
||
|
||
// Predict noise
|
||
let pred_noise = self.unet.forward(&xt, t, device);
|
||
|
||
// Compute diffusion loss: ||ε - ε_θ||²
|
||
let diff = noise.sub(&pred_noise);
|
||
let diff_sq = diff.mul(&diff);
|
||
let diffusion_loss = diff_sq.mean().to_vec()[0];
|
||
|
||
// Predict x0 from noisy sample and predicted noise
|
||
let pred_x0 = self.scheduler.predict_x0(&xt, &pred_noise, t);
|
||
|
||
// Optionally clamp predictions
|
||
let pred_x0 = if self.config.clamp_predictions {
|
||
clamp(&pred_x0, -1.0, 1.0)
|
||
} else {
|
||
pred_x0
|
||
};
|
||
|
||
// Compute physics loss: ||L[x̂_0] - f||²
|
||
let physics_residual = physics_fn(&pred_x0);
|
||
let physics_sq = physics_residual.mul(&physics_residual);
|
||
let physics_loss = physics_sq.mean().to_vec()[0];
|
||
|
||
// Combined loss
|
||
let total_loss = diffusion_loss + self.config.physics_weight * physics_loss;
|
||
|
||
(total_loss, diffusion_loss, physics_loss)
|
||
}
|
||
|
||
/// Sample from the model using DDPM sampling.
|
||
///
|
||
/// # Arguments
|
||
/// * `shape` - Shape of samples to generate [batch, channels, height, width]
|
||
/// * `device` - Device
|
||
///
|
||
/// # Returns
|
||
/// Generated samples
|
||
pub fn sample(&self, shape: [usize; 4], device: &B::Device) -> GenericTensor<B, 4> {
|
||
// Start from pure noise
|
||
let mut x = GenericTensor::<B, 4>::randn(shape, device);
|
||
|
||
// Reverse diffusion process
|
||
for t in (0..self.scheduler.num_timesteps()).rev() {
|
||
// Predict noise
|
||
let pred_noise = self.unet.forward(&x, t, device);
|
||
|
||
// Sample noise for stochastic step (not needed at t=0)
|
||
let noise = if t > 0 {
|
||
Some(GenericTensor::<B, 4>::randn(shape, device))
|
||
} else {
|
||
None
|
||
};
|
||
|
||
// Reverse step
|
||
x = self.scheduler.step(&x, &pred_noise, t, noise.as_ref());
|
||
|
||
// Optional: clamp intermediate results
|
||
if self.config.clamp_predictions {
|
||
x = clamp(&x, -1.0, 1.0);
|
||
}
|
||
}
|
||
|
||
x
|
||
}
|
||
|
||
/// Sample with physics guidance.
|
||
///
|
||
/// During sampling, we modify the predicted x0 to reduce physics residual:
|
||
/// ```text
|
||
/// x̂_0 = x̂_0 - γ * ∇_{x̂_0} ||L[x̂_0] - f||²
|
||
/// ```
|
||
///
|
||
/// # Arguments
|
||
/// * `shape` - Shape of samples to generate
|
||
/// * `physics_fn` - Function that computes physics residual
|
||
/// * `physics_grad_fn` - Function that computes gradient of physics loss w.r.t. input
|
||
/// * `device` - Device
|
||
///
|
||
/// # Returns
|
||
/// Generated samples satisfying physics constraints
|
||
pub fn sample_with_guidance<F, G>(
|
||
&self,
|
||
shape: [usize; 4],
|
||
physics_fn: F,
|
||
physics_grad_fn: G,
|
||
device: &B::Device,
|
||
) -> GenericTensor<B, 4>
|
||
where
|
||
F: Fn(&GenericTensor<B, 4>) -> GenericTensor<B, 4>,
|
||
G: Fn(&GenericTensor<B, 4>) -> GenericTensor<B, 4>,
|
||
{
|
||
let mut x = GenericTensor::<B, 4>::randn(shape, device);
|
||
|
||
for t in (0..self.scheduler.num_timesteps()).rev() {
|
||
// Predict noise
|
||
let pred_noise = self.unet.forward(&x, t, device);
|
||
|
||
// Predict x0
|
||
let mut pred_x0 = self.scheduler.predict_x0(&x, &pred_noise, t);
|
||
|
||
// Apply physics guidance (gradient descent on physics loss)
|
||
if self.config.physics_guided_sampling && self.config.guidance_strength > 0.0 {
|
||
// Compute gradient of physics loss
|
||
let grad = physics_grad_fn(&pred_x0);
|
||
|
||
// Update x0 to reduce physics residual
|
||
pred_x0 = pred_x0.sub(&grad.mul_scalar(self.config.guidance_strength));
|
||
}
|
||
|
||
// Clamp if configured
|
||
if self.config.clamp_predictions {
|
||
pred_x0 = clamp(&pred_x0, -1.0, 1.0);
|
||
}
|
||
|
||
// Reconstruct noise from corrected x0
|
||
let alpha_cumprod = self.scheduler.alpha_cumprod(t);
|
||
let corrected_noise = x
|
||
.sub(&pred_x0.mul_scalar(alpha_cumprod.sqrt()))
|
||
.mul_scalar(1.0 / (1.0 - alpha_cumprod).sqrt());
|
||
|
||
// Sample noise for stochastic step
|
||
let noise = if t > 0 {
|
||
Some(GenericTensor::<B, 4>::randn(shape, device))
|
||
} else {
|
||
None
|
||
};
|
||
|
||
// Reverse step with corrected noise prediction
|
||
x = self.scheduler.step(&x, &corrected_noise, t, noise.as_ref());
|
||
}
|
||
|
||
x
|
||
}
|
||
|
||
/// Compute the Laplacian of a 2D field using finite differences.
|
||
///
|
||
/// This is a helper for common physics computations.
|
||
/// ∇²u ≈ (u[i+1,j] + u[i-1,j] + u[i,j+1] + u[i,j-1] - 4*u[i,j]) / h²
|
||
pub fn laplacian_2d(field: &GenericTensor<B, 4>, h: f32) -> GenericTensor<B, 4> {
|
||
let shape = field.shape();
|
||
let [batch, channels, height, width] = [shape[0], shape[1], shape[2], shape[3]];
|
||
|
||
let field_vec = field.to_vec();
|
||
let mut laplacian = vec![0.0f32; batch * channels * height * width];
|
||
|
||
let h2 = h * h;
|
||
|
||
for b in 0..batch {
|
||
for c in 0..channels {
|
||
for i in 1..height - 1 {
|
||
for j in 1..width - 1 {
|
||
let idx = |ii: usize, jj: usize| -> usize {
|
||
b * channels * height * width + c * height * width + ii * width + jj
|
||
};
|
||
|
||
let center = field_vec[idx(i, j)];
|
||
let left = field_vec[idx(i - 1, j)];
|
||
let right = field_vec[idx(i + 1, j)];
|
||
let down = field_vec[idx(i, j - 1)];
|
||
let up = field_vec[idx(i, j + 1)];
|
||
|
||
laplacian[idx(i, j)] = (left + right + down + up - 4.0 * center) / h2;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
GenericTensor::<B, 4>::from_slice(&laplacian, shape, &field.device())
|
||
}
|
||
|
||
/// Create a Poisson residual function.
|
||
///
|
||
/// For -∇²u = f, the residual is: R = -∇²u - f
|
||
pub fn poisson_residual(
|
||
u: &GenericTensor<B, 4>,
|
||
f: &GenericTensor<B, 4>,
|
||
h: f32,
|
||
) -> GenericTensor<B, 4> {
|
||
let laplacian = Self::laplacian_2d(u, h);
|
||
// Residual = -∇²u - f = -laplacian - f
|
||
laplacian.mul_scalar(-1.0).sub(f)
|
||
}
|
||
}
|
||
|
||
/// Clamp tensor values to [min, max].
|
||
fn clamp<B: Backend<FloatElem = f32>, const D: usize>(
|
||
x: &GenericTensor<B, D>,
|
||
min: f32,
|
||
max: f32,
|
||
) -> GenericTensor<B, D> {
|
||
let x_vec = x.to_vec();
|
||
let clamped: Vec<f32> = x_vec.iter().map(|&v| v.clamp(min, max)).collect();
|
||
GenericTensor::<B, D>::from_slice(&clamped, x.shape(), &x.device())
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use crate::unet::UNetConfig;
|
||
use rtx_backend_cpu::{CpuBackend, CpuDevice};
|
||
|
||
fn create_test_model() -> PIDDM<CpuBackend> {
|
||
let device = CpuDevice::default();
|
||
|
||
let scheduler = DDPMScheduler::new(100, 1e-4, 0.02); // Small for testing
|
||
let unet_config = UNetConfig {
|
||
in_channels: 1,
|
||
out_channels: 1,
|
||
base_channels: 16,
|
||
hidden_mult: 2,
|
||
time_embed_dim: 32,
|
||
};
|
||
let spatial_size = 64; // 8x8
|
||
let unet = DiffusionUNet::new(unet_config, spatial_size, &device);
|
||
let config = PIDDMConfig::default();
|
||
|
||
PIDDM::new(scheduler, unet, config)
|
||
}
|
||
|
||
#[test]
|
||
fn test_piddm_creation() {
|
||
let model = create_test_model();
|
||
assert_eq!(model.scheduler().num_timesteps(), 100);
|
||
}
|
||
|
||
#[test]
|
||
fn test_laplacian_2d() {
|
||
let device = CpuDevice::default();
|
||
|
||
// Create a simple quadratic field: u(x,y) = x² + y²
|
||
// For this field, ∇²u = 2 + 2 = 4 (constant)
|
||
let n = 8;
|
||
let h = 1.0 / (n - 1) as f32;
|
||
|
||
let mut data = vec![0.0f32; n * n];
|
||
for i in 0..n {
|
||
for j in 0..n {
|
||
let x = i as f32 * h;
|
||
let y = j as f32 * h;
|
||
data[i * n + j] = x * x + y * y;
|
||
}
|
||
}
|
||
|
||
let field = GenericTensor::<CpuBackend, 4>::from_slice(&data, [1, 1, n, n], &device);
|
||
|
||
let laplacian = PIDDM::<CpuBackend>::laplacian_2d(&field, h);
|
||
let laplacian_vec = laplacian.to_vec();
|
||
|
||
// Interior points should have Laplacian ≈ 4
|
||
let center = n / 2;
|
||
let center_val = laplacian_vec[center * n + center];
|
||
assert!(
|
||
(center_val - 4.0).abs() < 0.5,
|
||
"Laplacian should be close to 4, got {}",
|
||
center_val
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_training_step() {
|
||
let model = create_test_model();
|
||
let device = CpuDevice::default();
|
||
|
||
// Create dummy clean sample
|
||
let x0 = GenericTensor::<CpuBackend, 4>::rand([1, 1, 8, 8], &device);
|
||
|
||
// Dummy physics function (just returns zeros - no physics constraint)
|
||
let physics_fn = |_x: &GenericTensor<CpuBackend, 4>| {
|
||
GenericTensor::<CpuBackend, 4>::zeros([1, 1, 8, 8], &device)
|
||
};
|
||
|
||
let (total, diffusion, physics) = model.training_step(&x0, physics_fn, &device);
|
||
|
||
// Losses should be non-negative
|
||
assert!(total >= 0.0, "Total loss should be non-negative");
|
||
assert!(diffusion >= 0.0, "Diffusion loss should be non-negative");
|
||
assert!(physics >= 0.0, "Physics loss should be non-negative");
|
||
|
||
// Physics loss should be 0 (since residual function returns zeros)
|
||
assert!(
|
||
physics < 1e-10,
|
||
"Physics loss should be ~0 for zero residual"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_sample_shape() {
|
||
let model = create_test_model();
|
||
let device = CpuDevice::default();
|
||
|
||
// Note: This is a simplified test - full sampling is slow
|
||
// We just test that the forward pass works with the expected shapes
|
||
let shape = [1, 1, 8, 8];
|
||
|
||
// Test single step
|
||
let x = GenericTensor::<CpuBackend, 4>::randn(shape, &device);
|
||
let pred_noise = model.unet().forward(&x, 50, &device);
|
||
|
||
let pred_shape = pred_noise.shape();
|
||
assert_eq!(pred_shape, shape);
|
||
}
|
||
|
||
#[test]
|
||
fn test_poisson_residual() {
|
||
let device = CpuDevice::default();
|
||
|
||
let n = 8;
|
||
let h = 1.0 / (n - 1) as f32;
|
||
|
||
// Create u = sin(πx)sin(πy), which satisfies -∇²u = 2π²sin(πx)sin(πy)
|
||
let pi = std::f32::consts::PI;
|
||
let mut u_data = vec![0.0f32; n * n];
|
||
let mut f_data = vec![0.0f32; n * n];
|
||
|
||
for i in 0..n {
|
||
for j in 0..n {
|
||
let x = i as f32 * h;
|
||
let y = j as f32 * h;
|
||
u_data[i * n + j] = (pi * x).sin() * (pi * y).sin();
|
||
f_data[i * n + j] = 2.0 * pi * pi * (pi * x).sin() * (pi * y).sin();
|
||
}
|
||
}
|
||
|
||
let u = GenericTensor::<CpuBackend, 4>::from_slice(&u_data, [1, 1, n, n], &device);
|
||
let f = GenericTensor::<CpuBackend, 4>::from_slice(&f_data, [1, 1, n, n], &device);
|
||
|
||
let residual = PIDDM::<CpuBackend>::poisson_residual(&u, &f, h);
|
||
let residual_vec = residual.to_vec();
|
||
|
||
// Interior residual should be small (discretization error)
|
||
let center = n / 2;
|
||
let center_residual = residual_vec[center * n + center];
|
||
assert!(
|
||
center_residual.abs() < 1.0,
|
||
"Poisson residual should be small at center, got {}",
|
||
center_residual
|
||
);
|
||
}
|
||
}
|