PerFlow: Physics-Embedded Rectified Flow for Efficient Reconstruction and Uncertainty Quantification of Spatiotemporal Dynamics

1 Introduction

Partial differential equations (PDEs) govern various dynamical systems in science and engineering [Holton and Hakim, 2013; Tu et al., 2023; Strogatz, 2024]. While classical numerical solvers offer high accuracy and convergence guarantees, they require fine-grained spatiotemporal discretization, resulting in prohibitively high computational costs [Lapidus and Pinder, 1999; Tadmor, 2012]. To address this issue, neural surrogate models learn the solution operator between function spaces, aiming for rapid inference after training [Li et al., 2020; Lu et al., 2021; Zhang et al., 2025]. Most existing surrogate models are deterministic and trained on gridded inputs (e.g., complete initial conditions), yet real-world measurement data is often sparse and irregular, rendering field reconstruction an ill-posed and highly challenging problem [Zhang et al., 2015]. Furthermore, deterministic methods only provide point estimates and lack uncertainty estimation capability.

∗Corresponding authors: Rui Zhang and Hao Sun.

Recent advances in generative models have enabled probabilistic surrogates for PDE systems. Diffusion [Ho et al., 2020; Song et al., 2020b] and flow [Lipman et al., 2022] models learn a probability distribution over spatiotemporal fields, where stochastic sampling naturally supports uncertainty quantification. For forward modeling, generative PDE surrogates can capture multi-scale spectral features and quantify uncertainty, improving robustness to noise and distribution shift [Lippe et al., 2023; Shysheya et al., 2024; Li et al., 2025]. More importantly, for sparse reconstruction, they can support conditionally guided generation, offering multiple high-fidelity plausible fields consistent with limited measurements [Huang et al., 2024; Li et al., 2024a].

Many generative methods achieve controllable generation by enforcing data consistency and physics constraints during sampling via gradient-based guidance [Li et al., 2024a; Wang et al., 2025]. However, this introduces two major issues. Firstly, it requires calculating observation-consistency gradients and physics gradients at each step of the sampling process, resulting in the slow inference speeds [Li et al., 2024a]. Moreover, guidance tends to increase the required number of sampling steps and forces smaller step sizes for stable reconstructions. Secondly, it suffers from sensitivity to hyperparameters, one must carefully balance dataconsistency and physics-informed terms to conform the prior distribution and guidance information, where reconstruction quality is highly sensitive to these weights in practice [Huang et al., 2024]. Fig. 1 illustrates these issues on two representative guided diffusion baselines (DiffPDE [Huang et al., 2024] and S3GM [Li et al., 2024a]).

To this end, we propose the Physics-embedded rectified Flow (PerFlow, shown in Fig. 2), an efficient conditional generative framework for sparse reconstruction and uncertainty quantification under partial observations. Unlike sampling-time guidance, PerFlow decouples observation conditioning from physics enforcement, avoiding gradient-based guidance at inference. For observed information, PerFlow incorporates sparse measurements as conditions instead of conducting iterative data-fidelity gradients during sampling. For physics information, PerFlow encodes conservation laws through a constraint-preserving projection, which can ensure the entire generation remains on the physics-consistent manifold. Finally, by adopting rectified flow, PerFlow supports fast few-step sampling, enabling efficient conditional sampling. In summary, we make the following contributions:

1. We propose PerFlow, a physics-embedded rectified-flow framework for sparse-field reconstruction and uncertainty quantification, enabling guidance-free conditional sampling.
2. We enforce hard physical constraints via constraintpreserving projection, and provide theoretical guarantees that sampling trajectories remain on the physics-consistent manifold throughout sampling.
3. We evaluate the performance of PerFlow across various PDE benchmarks. PerFlow achieves competitive reconstruction accuracy with remarkably lower physics error, while providing ∼320× speedup over guided diffusion baselines.

2 RelatedWork

2.1 Neural Surrogates for PDEs

With the development of deep learning, neural surrogate models have emerged as a data-driven approach to accelerating PDE simulation [Azizzadenesheli et al., 2024; Zhang et al., 2024; Bhaganagar and Chambers, 2025]. The main idea of neural surrogates is using a neural network to learn the mapping between function space, such as from PDE initial conditions to the corresponding solution. Compared with classical numerical methods, neural surrogate models do not need fine-grained spatiotemporal discretization and thus achieve remarkable speedup. DeepONet and its variants [Lu et al., 2021; He et al., 2024] use branch and trunk networks to encode the physical fields and spatiotemporal coordinates, respectively. Fourier neural operator and its variants [Li et al., 2020; Rahman et al., 2022; Xiong et al., 2024; Li et al., 2024b; Wan et al., 2025] conduct function transformation in the frequency domain to achieve discrete invariance. Moreover, graph-based models [Eliasof et al., 2021; Zeng et al., 2024] and transformer-based models [Li et al., 2023; Wu et al., 2024] also extend neural surrogates to the irregular regime with their flexibility. However, most existing models use deterministic frameworks and conduct the training process across the complete fields, thus struggling with sparsity and uncertainty.

2.2 Generative Modeling of PDE Systems

Generative models offer a complementary approach to learning spatiotemporal systems by modeling distributions over trajectories rather than producing a single point estimate, such as variational autoencoders [Kingma and Welling, 2013], normalizing flows [Rezende and Mohamed, 2015], GANs [Goodfellow et al., 2020], and diffusion models [Ho et al., 2020; Song et al., 2020a; Lu et al., 2025] that sample via iterative denoising. Recently, continuous-time flowbased formulations, such as flow matching [Lipman et al., 2022] and rectified flow [Liu et al., 2022], learn a timedependent velocity field and generate samples by integrating an ODE, reducing the number of function evaluations required for sampling. For physical systems, these models have been used for probabilistic forward simulation and for various downstream tasks [Xu et al., 2022; Du et al., 2024; Wang et al., 2025; Hu et al., 2025b], including parameter inference, system control, and sparse reconstruction [Wei et al., 2024; Li et al., 2024a; Huang et al., 2024; Hu et al., 2025a]. A common paradigm is to model the joint distribution of the full spatiotemporal fields and incorporate measurements and physics information through sampling-time guidance [Huang et al., 2024]. However, this framework requires iterative denoising with per-step backpropagation, leading to slow inference and requiring more sampling steps or smaller step sizes for stable reconstructions. Moreover, these methods suffer from hyperparameter sensitivity, since reconstruction quality depends strongly on the relative weights used to balance observation and physics terms.

3 Method

3.1 Problem Setup

We study sparse reconstruction and uncertainty quantification for PDE-governed physical fields. Let Ω ⊂ Rd be a spatial domain and [0, T] a time interval. A broad class of systems can be written as

where x denotes the spatiotemporal field, θ denotes PDE parameters, and b represents boundary conditions. In practice, only sparse and irregular measurements are available. We model sparse observations with a binary mask M, where Mi = 1 indicates an observed entry and Mi = 0 otherwise:

Starting from sparse measurements y, PerFlow infers the complete field x and represents uncertainty through samples from the conditional distribution

3.2 Overview of PerFlow

Most posterior-guided diffusion/flow approaches enforce data consistency and physics constraints during sampling via gradient-based guidance, resulting in slow inference and sensitivity to the guidance weights. To address these issues, PerFlow adopts rectified flow for efficient ODE-based sampling and decouples observation conditioning from physical guidance. PerFlow has the following key designs:

1. Rectified flow for few-step inference. Based on rectified flow, PerFlow samples by integrating a conditional ODE with a small number of function evaluations, enabling efficient reconstruction and uncertainty quantification.
2. Guidance-free conditioning. We regard sparse reconstruction as conditional generation and explicitly provide the sparse measurements as conditions, avoiding sampling-time data-fidelity gradients.
3. Constraint-preserving projection. We use constraintpreserving projection to ensure that trajectories remain on the physics-consistent manifold throughout sampling.

3.3 Rectified Flow Preliminaries

Rectified flow (RF) [Liu et al., 2022] learns a continuous-time transport that maps a simple base distribution to a target conditional data distribution. Let p0 be a base distribution and p1(· | c) the target conditional distribution given c. RF learns a time-dependent velocity field predictor vθ(t, x, c) and generates samples by integrating

Training is based on flow matching along an interpolation path. Given x1 ∼ p1(· | c) and x0 ∼ p0, we construct

and use the rectified-flow target velocity

Compared to diffusion-based sampling, RF admits efficient ODE solvers and typically requires fewer function evaluations, which is crucial for fast PDE reconstruction.

3.4 Conditioning on Sparse Observations

Unlike posterior-guided diffusion/flow methods that enforce observation consistency via sampling-time gradients, PerFlow casts sparse reconstruction as conditional generation. The observation information is incorporated through the condition input, enabling guidance-free inference without perstep backpropagation.

Given sparse measurements y and mask M, we use a mask-aware condition

where M is broadcast to match channels if needed. We then learn a conditional velocity field vθ(t, x, c) parameterized by a U-Net backbone. Training and inference procedures are detailed in Sec. 3.6.

3.5 Constraint-preserving Projection

A central design of PerFlow is to restrict the rectified-flow dynamics to a physics-consistent manifold that encodes hard PDE priors (Fig. 2a). Instead of injecting physics through gradient guidance, we propose the constraint-preserving projection to enforce essential constraints into the model architecture and sampling pipeline so that generated fields satisfy them by construction.

Where physics is embedded. PerFlow enforces hard constraints at two points: (i) the physics-embedded prior, which is obtained by applying a prior projection to the sampled noise, thereby yielding a physics-feasible initial state x(0); and (ii) the constraint-preserving velocity projection, where the learned velocity field is mapped to satisfy the corresponding constraints required by each ODE updates. We implement these designs using a set of physics-embedded operators, whose underlying physical constraints are summarized in Fig. 2c for clarity.

1. Divergence-free. For 2D incompressible flow, we enforce incompressibility with a stream-function construction,

where we represent the velocity field u = (u, v) with a scalar stream function ψ:

which satisfies ∇·u = 0 by construction. For the physicsembedded prior, we sample ψ from a Gaussian random field, and project it to a divergence-free initial state u(0) = ∇⊥ψ. For the constraint-preserving velocity projection, the network predicts ψ, and we map it to (u, v) to obtain divergence-free velocity updates.

During ODE integration, we also keep boundary values fixed by applying the same boundary mask to the velocity, i.e., we set boundary updates to zero using (1 - B) ⊙ (·).

(3) Global conservation. For scalar conserved quantities, we enforce a fixed total mass using the mean-correction operator Πmass, which shifts the field to match the target mass:

where m0 represents the total mass. For velocity updates, we apply the corresponding zero-mean correction so that the ODE updates do not change the total mass. In practice, we apply Πmass when sampling the physics-embedded prior, and apply the mean-correction to the predicted velocity field whenever this constraint is required.

By encoding these operators into both the physicsembedded prior and the constraint-preserving velocity projection, PerFlow starts from a physics-feasible state and performs rectified-flow sampling directly on the physicsconsistent manifold without sampling-time physics guidance. Furthermore, we provide theoretical guarantees for the continuous-time dynamics in Sec. 4, where we prove that the physical constraint satisfaction can be maintained throughout sampling.

3.6 Training Objective and Conditional Sampling

We introduce PerFlow's training and inference procedures with sparse observations (Fig. 2b). The complete algorithmic procedure is provided in Appendix 2.

Training (Conditional RF with physics embedding). During training, we take a physical field x1 from the training set and sample a masking M to simulate sparse measurements y = M⊙ x1, thus forming the condition c = concat(y,M). After that, we sample a random noise ˜x0 from an initial distribution p0 and project it to a physicsconsistent initial state x0 = Π(˜x0) using the projection operators in Sec. 3.5. We then construct xt and v⋆ as in Eqs. 5-6,

In our implementation, vθ is made physics-consistent by constraint-preserving projection (Sec. 3.5), thus the learned dynamics is restricted to the physics-consistent manifold automatically.

Inference (Sparse reconstruction and uncertainty quantification). Given sparse observations (y,M), we set c = concat(y,M) to form a condition. We sample K i.i.d. physics-consistent initial states and integrate the conditional ODE in Eq. 4 from t = 0 to 1 using a few-step solver to obtain samples ˆx(k) = x(k) (1). After that, we can estimate the uncertainty from the ensemble {ˆx(k) }Kk=1 via posterior mean and variance.

4 Theory

This section states theoretical guarantees for PerFlow's hardconstraint embedding. After discretization, many physics priors (e.g., conservation and boundary conditions) can be written as affine linear constraints

where A stacks the corresponding constraint operators and p encodes prescribed values. PerFlow generates samples by integrating the conditional ODE

where c is the observation condition. By design, PerFlow obtains a feasible initial state x(0) ∈ S via prior projection and uses a constraint-preserving velocity field (Sec. 3.5). The following theorem formalizes the resulting trajectory-level constraint preservation.

Theorem 1 (Constraint invariance of continuous dynamics). Assume the feasible set is S = {x : Ax = p}. Assume vθ(t, x, c) is continuous in t and locally Lipschitz in x so that the ODE (13) admits a unique solution on [0, 1]. If (i) x(0) ∈ S and (ii) the velocity field is tangent to S, i.e.,

then the solution satisfies x(t) ∈ S for all t ∈ [0, 1].

In PerFlow, we parameterize the velocity field in a constrained form (e.g., stream-function for incompressibility) or use an exact projection to the network output to enforce the corresponding constraints (mean correction for mass conservation). As a result, for any feasible state x ∈ S, the effective velocity satisfies Avθ(t, x, c) = 0, which is exactly the tangency condition in Eq. 14. Therefore, with prior-projected initialization and constraint-preserving velocity projection, physical constraint satisfaction can be maintained throughout the entire sampling (detailed in Appendix 1). This trajectory-level guarantee differs from sampling-time guidance, which encourages physical consistency through additional correction terms during sampling but does not enforce feasibility for every ODE update. For nonlinear manifolds M = {x : g(x) = 0}, trajectory invariance would require Jg(x)vθ(t, x, c) = 0, where Jg(x) is the Jacobian of g at x; we leave this extension for future work.

5 Results

5.1 Benchmarks, Baselines, andMetrics

Benchmarks. We evaluate PerFlow on four kinds of PDE systems, including: (i) 1D Burgers with periodic boundary conditions, (ii) 2D Poisson with Dirichlet boundary conditions, (iii) 2D Darcy with Dirichlet boundary conditions, and (iv) 2D incompressible Navier-Stokes (NS) on a periodic domain. The default discretizations and measurements are as follows: Burgers uses 128 time steps on a 128-point spatial grid with k = 5 observed locations; Poisson and Darcy use 128 × 128 grids with k = 500 observed points; NS uses 2D velocity fields (u, v) on a 64 × 64 grid over T = 10 frames, with k = 200 observed locations per frame. Observation masks are randomly sampled. Additional dataset details are provided in Appendix 5 and Appendix 6.

Baselines. We compare PerFlow with (i) deterministic surrogates and (ii) guided generative baselines. For deterministic models, we train FNO [Li et al., 2020], U-Net [Ronneberger et al., 2015], and FFNO [Tran et al., 2023] as surrogate solvers, learning the forward mapping between full fields governed by the PDE. For sparse reconstruction, we randomly initialize the full fields as input variables and iteratively optimize them by minimizing the mismatch between the model's output and the sparse observations. The reported inference time includes the entire optimization process. Furthermore, we also compare PerFlow with conditional FNO, which learns a deterministic mapping from sparse observations to the full field directly, and the corresponding results are reported in Appendix 8.

For generative baselines, we evaluate S3GM [Li et al., 2024a] and DiffPDE [Huang et al., 2024], each with a fewstep sampling (50 steps) and a long-run setting following the sampling steps used in the original papers (2000 steps), denoted as S3GM-50 (2000) and DiffPDE-50 (2000), respectively. These two settings highlight the accuracy-efficiency trade-off: 2000 steps reflect the high-budget performance of guided baselines, while 50 steps compare them with PerFlow under a similar inference budget. Implementation details and baseline descriptions are provided in Appendix 3 and Appendix 4. The code will be released at https://github.com/ intell-sci-comput/PerFlow.

Metrics. In this paper, we report reconstructed accuracy via relative errors, including Rel-ℓ2 = ∥ˆx - x∥2/∥x∥2 and Rel-ℓ1 = ∥ˆx - x∥1/∥x∥1. We also report the physics error (Phys-Err) for each PDE, such as divergence for NS and mass conservation for Burgers. Full definitions, including component averaging across variables and the PDE-specific PhysErr computation, are provided in Appendix 7.

5.2 Sparse Reconstruction of Physical Fields

Fig. 3 shows the sparse reconstruction results for each PDE. Among these methods, deterministic models often offer oversmoothed reconstructed results and introduce artifacts in unobserved regions. Meanwhile, guided diffusion baselines, like DiffPDE and S3GM, can recover several detail patterns and improve observation consistency when sufficient sampling steps are taken, but still struggle in the few-step regime.

Table 1 summarizes the quantitative results across 20 test cases for each PDE. PerFlow significantly outperforms baselines in reconstruction accuracy while strictly enforcing physical constraints. Compared to deterministic surrogates, PerFlow reduces Rel-ℓ2 errors by 29.8%-94.0% across the four benchmarks. This is because solving the inverse problem via iterative optimization with a deterministic model is ill-posed and sensitive to initialization, resulting in higher errors and prolonged inference times (29-55s). Moreover, compared with conditional FNO, PerFlow still achieves a remarkably high accuracy due to its ability to handle ill-posed scenarios (Appendix 8).

When compared with generative models under the same 50-step budget, PerFlow improves reconstruction accuracy by 1-2 orders of magnitude over DiffPDE-50 and S3GM-50. Simultaneously, PerFlow is 5.7-6.4× faster than DiffPDE-50 and 3.0-4.2× faster than S3GM-50. This efficiency stems from PerFlow's design. Guided baselines require per-step backpropagation for observation and physics guidance, while PerFlow performs guidance-free sampling with lightweight projections. For example, on Burgers, mass projection introduces negligible overhead, with 0.79±0.12s using projection versus 0.76 ± 0.06s without it. Even when the guided baselines are extended to 2000 sampling steps, PerFlow maintains superior accuracy on three of the four benchmarks while achieving significantly strong physics consistency. Notably, PerFlow provides up to a ∼320× speedup over DiffPDE2000. In summary, PerFlow can not only offer high fidelity comparable to converged diffusion models, but also enjoy fast inference speed due to few-step solvers. Additional robustness analyses on base distributions and noisy observations are provided in the Appendix 9. Moreover, long-run guided baselines and varying-observation results are reported in Appendix 10 and Appendix 11.

5.3 Uncertainty Quantification

A main advantage of generative modeling is its ability to quantify posterior uncertainty under sparse observations. To verify this, we evaluate this capability on the Darcy and Burgers benchmarks, where the inverse problem is inherently illposed due to sparse and irregular measurements. Armed with the generative nature of PerFlow, we generate an ensemble of 20 independent predictions for each test case by sampling distinct initial noise states. We then calculate the predicted mean as the final reconstructed field and the entry-wise standard deviation across the ensemble to quantify the model's uncertainty. Fig. 4 shows the uncertainty quantification results. In all cases, the predictive mean accurately recovers the ground truth structures. Furthermore, we observe a strong spatial correlation between the estimated uncertainty and the actual reconstruction error. Regions exhibiting high variance typically coincide with larger deviations from the ground truth. For instance, uncertainty peaks at the sharp phase boundaries in the Darcy coefficient (Fig. 4a) and concentrates in the gaps between temporal snapshots for the Burgers equation (Fig. 4c). This alignment indicates that PerFlow effectively identifies challenging regions where the solution is less confident, serving as a reliable indicator of reconstruction errors.

5.4 Few-step Sampling

We study how reconstruction accuracy varies with the number of sampling steps for S3GM, DiffPDE, and PerFlow (Fig. 5). Across these PDEs, PerFlow achieves strong accuracy in the few-step sampling regime and remains stable as the number of steps increases. In contrast, other guided diffusion methods exhibit a much slower and more severe trade-off between accuracy and the number of steps, because of the need to balance data-fidelity guidance and physics guidance. PerFlow addresses this issue via avoiding sampling-time backpropagation and embedding physics as hard constraints, thereby achieving efficient and accurate reconstruction.

5.5 Ablation Studies

We conduct ablations on 1D Burgers equation to study (i) the effect of physics embedding and (ii) the choice of generative procedure. Table 2 compares PerFlow with several variants, including RF-50, which removes the hard constraint projector from PerFlow; FM-50, which replaces rectified flow with a standard flow-matching model under the same 50-step sampling [Lipman et al., 2022]; DDIM-50 [Song et al., 2020a] and DDPM-50 [Ho et al., 2020], two diffusion baselines sampled with 50 steps, and DDPM-1000 with 1000 steps.

Comparing PerFlow with RF-50 reveals the contribution of hard constraint embedding. While both methods achieve similar reconstruction accuracy, PerFlow reduces physics error by several orders of magnitude (Table 2), indicating that constraint-preserving projection is crucial for producing physically consistent reconstructions. Furthermore, RFbased generation in PerFlow is more effective than other diffusion models with the same number of sampling (Table 2). DDPM-50 fails catastrophically, whereas DDIM-50 improves substantially but remains inferior to PerFlow in both accuracy and physical consistency. Although DDPM-1000 can achieve competitive Rel-ℓ2/ℓ1, it still exhibits larger physics error and requires an order-of-magnitude more sampling steps.

6 Conclusion and FutureWork

We propose PerFlow, a physics-embedded rectified-flow framework for sparse reconstruction and uncertainty quantification in PDE-governed systems. Compared with posteriorguided generative methods, PerFlow decouples observation conditioning from physics enforcement by performing guidance-free observation conditioning in rectified-flow dynamics and encoding hard physics via constraint-preserving projection. Theoretically, sampling trajectories are guaranteed to remain on the physics-consistent manifold. Numerically, PerFlow achieves competitive accuracy and strong physics consistency across various PDE datasets, while enabling efficient conditional sampling (e.g., 50 steps) and up to ∼320× faster inference than 2000-step diffusion baselines.

There are two future works that we are interested in. Firstly, we can combine PerFlow with sensor placement approaches, which can maximize reconstruction quality and uncertainty reduction [Ma et al., 2025]. Secondly, we can extend PerFlow to other downstream tasks, because fast conditional sampling on a physics-consistent manifold can provide reliable control and planning [Hu et al., 2025a].

Acknowledgements

The work is supported by the Beijing Natural Science Foundation (No. F261002) and the National Natural Science Foundation of China (No. 62276269 and No. 62506367). R.Z. acknowledges support from the China Postdoctoral Science Foundation under Grant Number 2025M771582 and the Postdoctoral Fellowship Program of CPSF under Grant Number GZB20250408. The full version is available at http://arxiv.org/abs/2605.03548.