137 lines
3.6 KiB
Rust
137 lines
3.6 KiB
Rust
//! Configuration types for PIDDM demo.
|
|
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
/// PDE type for diffusion model training/inference.
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
|
|
#[serde(rename_all = "snake_case")]
|
|
pub enum PdeType {
|
|
/// Poisson equation: ∇²u = f
|
|
#[default]
|
|
Poisson,
|
|
/// Heat equation: ∂u/∂t = α∇²u
|
|
Heat,
|
|
/// Darcy flow: -∇·(k∇p) = f
|
|
Darcy,
|
|
/// Burgers equation: ∂u/∂t + u·∇u = ν∇²u
|
|
Burgers,
|
|
}
|
|
|
|
impl PdeType {
|
|
/// Get display name for PDE type.
|
|
#[must_use]
|
|
pub fn display_name(&self) -> &'static str {
|
|
match self {
|
|
Self::Poisson => "Poisson Equation",
|
|
Self::Heat => "Heat Equation",
|
|
Self::Darcy => "Darcy Flow",
|
|
Self::Burgers => "Burgers Equation",
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Noise scheduler type.
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
|
|
#[serde(rename_all = "snake_case")]
|
|
pub enum SchedulerType {
|
|
/// Denoising Diffusion Probabilistic Models
|
|
#[default]
|
|
DDPM,
|
|
/// Denoising Diffusion Implicit Models
|
|
DDIM,
|
|
}
|
|
|
|
/// Configuration for PIDDM training.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct PiddmTrainingConfig {
|
|
/// PDE type to solve
|
|
pub pde_type: PdeType,
|
|
/// Grid resolution
|
|
pub resolution: u32,
|
|
/// Number of training epochs
|
|
pub epochs: u32,
|
|
/// Learning rate
|
|
pub learning_rate: f64,
|
|
/// Batch size
|
|
pub batch_size: u32,
|
|
/// Number of diffusion timesteps
|
|
pub diffusion_steps: u32,
|
|
/// Physics loss weight
|
|
pub physics_weight: f32,
|
|
/// Scheduler type
|
|
pub scheduler: SchedulerType,
|
|
/// `UNet` hidden dimension
|
|
pub hidden_dim: usize,
|
|
/// Number of `UNet` layers
|
|
pub num_layers: usize,
|
|
}
|
|
|
|
impl Default for PiddmTrainingConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
pde_type: PdeType::Poisson,
|
|
resolution: 64,
|
|
epochs: 100,
|
|
learning_rate: 1e-4,
|
|
batch_size: 16,
|
|
diffusion_steps: 1000,
|
|
physics_weight: 0.1,
|
|
scheduler: SchedulerType::DDPM,
|
|
hidden_dim: 64,
|
|
num_layers: 4,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Configuration for PIDDM sampling/inference.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct PiddmSamplingConfig {
|
|
/// Number of samples to generate
|
|
pub num_samples: u32,
|
|
/// Grid resolution
|
|
pub resolution: u32,
|
|
/// Number of denoising steps (can be less than training steps for DDIM)
|
|
pub sampling_steps: u32,
|
|
/// Physics guidance strength (0 = no guidance, 1 = full guidance)
|
|
pub guidance_strength: f32,
|
|
/// Scheduler type
|
|
pub scheduler: SchedulerType,
|
|
}
|
|
|
|
impl Default for PiddmSamplingConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
num_samples: 4,
|
|
resolution: 64,
|
|
sampling_steps: 50,
|
|
guidance_strength: 0.5,
|
|
scheduler: SchedulerType::DDIM,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_pde_type_display() {
|
|
assert_eq!(PdeType::Poisson.display_name(), "Poisson Equation");
|
|
assert_eq!(PdeType::Heat.display_name(), "Heat Equation");
|
|
}
|
|
|
|
#[test]
|
|
fn test_training_config_default() {
|
|
let config = PiddmTrainingConfig::default();
|
|
assert_eq!(config.resolution, 64);
|
|
assert_eq!(config.diffusion_steps, 1000);
|
|
}
|
|
|
|
#[test]
|
|
fn test_sampling_config_default() {
|
|
let config = PiddmSamplingConfig::default();
|
|
assert_eq!(config.sampling_steps, 50);
|
|
assert_eq!(config.guidance_strength, 0.5);
|
|
}
|
|
}
|