Skip to content

Diffusion Models API

deepuq.models.diffusion

Diffusion-model building blocks for scientific field reconstruction.

The classes here support conditional denoising notebooks where uncertainty is estimated from sample spread rather than Bayesian posterior moments.

ConditionalUNet2D

Bases: Module

A compact conditional U-Net denoiser for 2D diffusion notebooks.

Parameters:

Name Type Description Default
x_channels int

Number of noisy input channels to denoise.

 1
cond_channels int

Number of conditioning channels supplied alongside the noisy input.

 2
base_channels int

Base feature width used by the U-Net encoder/decoder.

 32
time_dim int

Width of the timestep embedding processed by the residual blocks.

 128
dropout_p float

Spatial dropout probability inside the residual blocks.

 0.0
use_coordinate_features bool

Whether to append normalized (x, y) coordinate channels.

 True
Shape contract
  • x_t: [batch, x_channels, height, width]
  • timesteps: [batch]
  • condition: [batch, cond_channels, height, width]
  • output: denoised tensor with shape [batch, x_channels, height, width]
Example 
model = ConditionalUNet2D(x_channels=1, cond_channels=2, base_channels=32)
eps_hat = model(x_t, timesteps, condition)

forward 

forward(
    x_t: Tensor, timesteps: Tensor, condition: Tensor
) -> torch.Tensor

Predict the noise or residual field for a conditioned diffusion step.

SinusoidalTimeEmbedding

Bases: Module

Sinusoidal timestep embedding used by diffusion denoisers.

Parameters:

Name Type Description Default
embedding_dim int

Width of the returned timestep embedding.

 required
Shape contract
  • input: timesteps with shape [batch]
  • output: embedding tensor with shape [batch, embedding_dim]