A Physics-informed Diffusion Model for High-fidelity Flow Field Reconstruction

This card is based on arXiv v2 (LaTeX source, 2023-02-10). The official code README cites the journal version (Journal of Computational Physics, 2023). Section, equation, table, figure and algorithm numbers below follow the arXiv v2 source.

In plain words

Precise fluid simulations are expensive, and cheap ones are blurry or missing detail. Earlier learned fixes usually need matched pairs of cheap and precise runs, and they lose accuracy when the cheap input looks different from what they were trained on.

This paper treats the job as removing noise. A diffusion model is trained only on precise flow fields. At use time, the cheap field is mixed with random noise, and the model's reverse process turns that noisy version back into a sharp field. Because the model never saw the cheap input directly, the same trained model can handle inputs from different sources without retraining. Optionally, the fluid equation itself is used as a hint: how far a candidate field is from obeying the equation steers the sampling. The tests are two- dimensional turbulent flows.

TL;DR

• A DDPM is trained only on high-fidelity 2D Kolmogorov-flow vorticity, with no paired low-fidelity data §2.2, §3.1.
• Reconstruction starts the reverse chain from a noised low-fidelity or padded sparse guide at an intermediate step t, not from pure noise §2.3.
• The PDE residual gradient can be fed to the noise network, or used in the sampler. Author-reported gains are largest in equation loss and energy spectrum; L2 is close to the direct-mapping model §4.
• Shipped code differs from the paper in several places (see Paper vs code).

Problem

Learned super-resolution of CFD fields usually fits one particular type of under-resolved input. Accuracy drops when the test input has a different resolution, blur or sampling pattern, so the model has to be retrained for each new input type §1. The paper targets two settings: uniform down-sampling (4x and 8x upsampling) and sparse, randomly placed observations (5% and 1.5625% of grid points) §4.

Method

Problem setup. A direct-mapping network learns f: X -> Y from low-fidelity to high-fidelity samples, and its accuracy depends on the test inputs matching the training distribution §2.1. The paper instead adds Gaussian noise to the low-fidelity input, which makes the training and test input distributions closer, then denoises with a pretrained DDPM §2.1, fig. 1.

Step 1. Diffusion model. The reverse chain starts at a standard Gaussian and has transitions parameterised as Gaussians eq. 1. The forward process adds noise with a per-step scale beta_t §2.2. Training uses the simplified noise-prediction objective: the network predicts the noise in a noised sample x_t at a random step t, and the loss is the squared error eq. 2. Here alpha_t = 1 - beta_t and alpha_bar_t is the running product of alpha §2.2.

Step 2. Guided sampling from an intermediate step. The standard sampler starts from x_T drawn from noise, which gives no control over the output §2.3. The paper instead noises the guide x^(g) to step t: x_t = sqrt(alpha_bar_t) x^(g) + sqrt(1 - alpha_bar_t) eps eq. 3, eq. 4. The reverse chain then runs from t to 0. The paper's sampler is the DDIM deterministic update, with sigma = 0 §2.3, alg. 2. The step t is chosen in the interval [0, T/2] for accurate, less noisy output §2.3.

Step 3. Iterative refinement for sparse inputs §2.3.1. For sparse observations, unmeasured points are first filled by nearest-neighbour padding. The output of one round then becomes the guide of the next round, for K rounds. The number of rounds needed grows with the gap between the guide and the true field §2.3.1. For 4x upsampling and 5% sparsity the paper uses K = 3 §3.1. For 8x and 1.5625% it uses K = 1 with larger t and more steps §3.1.

Step 4. Physics-informed guidance §2.3.1. The PDE is written as G(u) = 0 eq. 5. The residual r_t = G evaluated at x_t measures how far a field is from satisfying the equation. Its gradient c = dr_t/dx_t is the conditioning signal eq. 6. Two ways of using c are given:

1. Learned guidance (Algorithm 1 for training, Algorithm 2 for sampling). The noise network takes c as an extra input. During training, c is dropped with probability p_u, so the same network also works unconditionally alg. 1. At sampling time, the prediction mixes the conditional and unconditional outputs with a guidance strength w alg. 2, eq. 8. Equation 7 reads these network outputs as scaled score estimates, so the mix combines the data score with the residual-gradient term eq. 7, eq. 8.
2. Direct gradient eq. 9. The unmodified network is used, and the DDIM update subtracts lambda times the residual gradient. The paper gives this as an option and sets the step size lambda as a symbol only.

The paper's main results use the learned variant with "learned residual guidance" as input features §3.1, §4.

# Physics-informed DDPM training  [alg. 1]
# Input: x_0 (B, 3, H, W), three consecutive vorticity frames  [§3.1]
# Input: p_u, probability of dropping the condition (not stated in paper)

def train_step(x_0, model, p_u):
    # Draw a random diffusion step per example  [alg. 1]
    t = uniform_int(1, T, size=B)
    # Noise the clean sample to step t  [alg. 1] [eq. 3]
    eps = normal_like(x_0)
    x_t = sqrt(alpha_bar[t]) * x_0 + sqrt(1 - alpha_bar[t]) * eps
    # Residual of the PDE at the noisy sample  [alg. 1] [eq. 5]
    r_t = G(x_t)
    # Drop the physics condition with probability p_u  [alg. 1]
    if random() < p_u:
        c = None
    else:
        # Physics-informed guidance c = dr_t/dx_t  [alg. 1] [eq. 6]
        c = gradient_of(r_t, wrt=x_t)
    # Squared error between the true and predicted noise  [alg. 1] [eq. 2]
    loss = squared_error(eps - model(x_t, t, c))
    return loss

# Physics-informed conditional sampling  [alg. 2]
# Input: x_g, the guide (low-fidelity or padded sparse field)
# Input: t, start step (t < T); tau, increasing sub-sequence from 0 to t
# Input: K rounds, guidance strength w, PDE operator G

def sample(x_g, t, tau, K, w, model):
    # Repeat the partial reverse chain K times  [alg. 2] [§2.3.1]
    for k in range(K):
        # Noise the current guide to step t  [alg. 2] [eq. 3]
        eps = normal_like(x_g)
        x_t = sqrt(alpha_bar[t]) * x_g + sqrt(1 - alpha_bar[t]) * eps
        # Walk the reduced step sequence from S down to 1  [alg. 2]
        for i in range(S, 0, -1):
            # Residual gradient at the current state  [alg. 2] [eq. 6]
            c = gradient_of(G(x_t[i]), wrt=x_t[i])
            eps_c = model(x_t[i], tau[i], c)
            eps_u = model(x_t[i], tau[i], None)
            # Mix conditional and unconditional predictions  [alg. 2] [eq. 8]
            eps_tilde = eps_c + w * (eps_c - eps_u)
            # Deterministic DDIM step, sigma = 0  [alg. 2] [§2.3]
            x_t[i - 1] = (
                sqrt(alpha_bar[tau[i-1]]) / sqrt(alpha_bar[tau[i]])
                * (x_t[i] - sqrt(1 - alpha_bar[tau[i]]) * eps_tilde)
                + sqrt(1 - alpha_bar[tau[i-1]]) * eps_tilde
            )
        # The output becomes the guide for the next round  [alg. 2]
        x_g = x_t[0]
    return x_g

Hyperparameters

Values below marked "not stated" are absent from the paper. Code values are listed where they exist, and the paper's wording is kept where the paper gives a value.

Results

Author-reported. Setup: 2D Kolmogorov flow at Re = 1000 §3.1. The same trained diffusion model is used for all tasks, without retraining §4. The direct-mapping baseline is a UNet trained on 64×64 → 256×256 pairs, so the other tasks are out-of-distribution for it tab. 1. Hardware: not stated.

Table 1 tab. 1: pointwise L2 norm and equation loss, lower is better.

Table 2 tab. 2: effect of Gaussian pre-filtering on the 1.5625% task.

Table 3 tab. 3: guidance-combination ablation. "Learned" is the network-input method, "Linear" is the direct gradient method eq. 9, and "Baseline" is no physics term. The upper block is the sparse task and the lower block is 8x. Column headers are as printed in the source. Printed iteration labels are reproduced as they appear.

Qualitative results: the diffusion model recovers more detail than bicubic interpolation on the harder 8x and sparse tasks [§4, fig. 3, fig. 4, fig. 5, fig. 6]. The direct-mapping output is blurrier and less physically coherent on out-of-distribution input §4, fig. 8. The diffusion output's vorticity distribution is sharper than the ground truth and drifts slightly, more so at higher sparsity §4, fig. 7. The energy spectrum matches better at high wavenumbers than bicubic §4, fig. 7, fig. 8. Gaussian filtering improves L2 for the direct-mapping model but removes high-frequency energy §4, fig. 8.

Paper vs code

Code: https://github.com/BaratiLab/Diffusion-based-Fluid-Super-resolution at commit a802a4b11e3b577b8d8c3f415cb8f9e75b373400, MIT (BaratiLab org; the README cites this paper). The paper and abstract do not link the repository. It was found by a GitHub search for the title words. The code sampling path uses the learned guided sampler (the 'conditional' config type).

• Equation metric. The paper normalises the equation loss by the mean squared vorticity norm eq. 12. The code's residual metric is the mean squared residual, with no normalisation rs256_guided_diffusion.py. Reported equation-loss numbers may not be comparable to values computed by the code.
• Guidance quantity. The paper's conditioning is the derivative of the residual with respect to the input eq. 6. The code uses the gradient of the mean squared residual, divided by the scaler's scale rs256_guided_diffusion.py rs256_guided_diffusion.py. The same gradient is used in training losses.py.
• Update rule in the learned path. Algorithm 2 has no gradient term in the update alg. 2. The code's guided DDIM step subtracts the residual gradient from the update, with a fixed weight of 1 and no lambda denoising_step.py. The shipped conditional configs set the guidance weight to 0 kmflow_re1000_rs256_conditional.yml, so the update term is active while the mixing term is off.
• Guidance mixing. The formula matches alg. 2: (w+1) * model(x, t, c) - w * model(x, t) denoising_step.py. The paper gives no value for w. The code default is 0.
• Rounds and steps. The paper uses K = 3 rounds for 4x and 5% with S = 160, 114, 80 §3.1. The code runs sample_step rounds, default 1 main.py. In each round it shrinks both the noise level and the step count by 0.7^k rs256_guided_diffusion.py rs256_guided_diffusion.py. Applied to 160 this gives 112 and 78, not the paper's 114 and 80. The paper does not say t decays.
• Shipped run command. The README example runs the unguided config (kmflow_re1000_rs256.yml) with sample_step 1, t 240, r 30, which is not the paper's K = 3 setting.
• Smoothing. The paper says σ = 5 §4. The shipped sparse conditional config sets scale 7, used as both kernel size and σ kmflow_re1000_rs256_sparse_recons_conditional.yml rs256_guided_diffusion.py.
• Clamping. The DDPM step paths clamp the predicted x_0 to [-1, 1] denoising_step.py denoising_step.py. The DDIM path used by the learned sampler does not denoising_step.py.
• Training details. The conditioning drop probability is 0.1 by default, and the paper gives no number losses.py. Training samples use antithetic time sampling diffusion_tub.py.
• Input preparation. The code reads precomputed sparse arrays from an external archive (keys u3232 and u3232_nearest). The nearest-neighbour padding step is not implemented in the repository kmflow_re1000_rs256_sparse_recons_conditional.yml.
• Test windows. The code takes all 318 three-frame windows of each 320-frame test sequence rs256_guided_diffusion.py.
• Stream function sign. The paper writes ψ = ∇^{-2} ω §3.1. The code uses ψ̂ = ω̂ / |k|², which is (−∇²)^{-1} ω, with u = ∂ψ/∂y and v = −∂ψ/∂x rs256_guided_diffusion.py. This is consistent with ω = ∂v/∂x − ∂u/∂y. Use the code's sign convention.
• Forcing and residual. The code's forcing and residual match the paper's equation 10 rs256_guided_diffusion.py. Time derivatives use a central difference over three frames rs256_guided_diffusion.py. This matches §3.1.

Pitfalls

• Timestep convention. The paper indexes steps from 1 eq. 3, alg. 1. The code is 0-based: training draws t from [0, T) and indexes the cumulative product directly losses.py. Noising in sampling uses index total − 1 rs256_guided_diffusion.py, and compute_alpha looks up t + 1 denoising_step.py. Pick one convention and check it against the code before porting.
• Step size of the learned residual term. The paper gives no lambda for the learned path. The code subtracts the raw gradient, scaled by the standardisation factor, at weight 1 denoising_step.py. Its magnitude depends on the residual's scale and on the normalisation the paper leaves unstated.
• Normalisation is unstated. The paper does not describe the scaling. The code standardises with training mean and std, and the guidance gradient is divided by that std rs256_guided_diffusion.py.
• Noise level and guidance strength. Both affect output sharpness. The paper reports a sharper, drifting vorticity distribution at higher sparsity §4, fig. 7. No setting for w is given in the paper, and the shipped conditional configs use 0.
• Decreasing schedule. The code shrinks t and S per round by 0.7^k. The paper's S values do not follow that rule exactly (see Paper vs code).
• Label mismatch. Table 1 writes 1.6%; the text and Table 2 use 1.5625%.
• Reporting. Each result is a single number, with no seeds, runs or spread stated. The code can repeat runs and log per-iteration statistics.
• Small test set. Four test sequences are used for all tasks §3.1.
• L2 versus equation loss. The paper argues that L2 is insensitive to blurring and to physical consistency §4. The learned method's L2 gain over the direct-mapping baseline is small, so the claim rests mostly on the equation and spectrum metrics.
• Cost. Sampling is a multi-step chain with K rounds, and each step calls the network twice (conditional and unconditional) plus a residual gradient. No timing or memory figures are given in the paper.

Lineage

• Builds on: DDPM (Ho et al., 2020) for the training objective and noise model §2.2; DDIM (Song et al., 2020) for the deterministic sampler §2.3; SDEdit (Meng et al., 2021) for starting the reverse chain from a noised guide §2.3; classifier-free guidance (Ho and Salimans, 2022) for the conditional/unconditional mix §2.3.1.
• Compared against: a direct-mapping UNet trained on paired data, and bicubic interpolation §4, tab. 1.
• Superseded by: not known from the wiki.