PC-RF: Physics-Constrained
Rectified Flow

Technical overview

PC-RF studies how a generative downscaler can produce conditional fields while keeping its output constraints explicit. The model learns a rectified-flow velocity field. The sampler integrates that field from Gaussian noise, then applies a precipitation operator that clamps negative values and matches a coarse mean.

The project also exports the velocity network into AOTInductor and ONNX artifacts. Python and Rust samplers keep the integration loop and output checks outside the network graph. The measured engineering results cover backend timing, numerical parity, and a corrected distributed-training sweep.

The A100 benchmark JSON records the timing and parameter count. The DDP-correct CSV records the training rate. Both records are explained below.

Problem formulation

A coarse field supplies the large-scale context, while a conditional generator samples possible fine fields. The coarse precipitation mean also supplies the value that the output contract must preserve.

The recorded architecture operates on a three-channel \(128\times128\) field and seven conditioning channels. The condition contains upsampled precipitation, upsampled wind components, and four zero-filled slots. The conference experiments used no Sentinel-2 information.

The architecture record names src/pcrf/model.py, src/pcrf/data.py, and the conference numbers ledger as its sources.

Conditional rectified flow

The network learns the direction from noise to data.

The training procedure samples a reference field \(x_1\), Gaussian noise \(x_0\), and a time \(t\). It places the training state on the straight segment between the noise and reference.

\[x_t=(1-t)x_0+t x_1,\qquad t\sim\mathcal U(0,1),\qquad x_0\sim\mathcal N(0,I).\]

The displacement \(x_1-x_0\) supplies the supervised target. The conditional network predicts that displacement from the intermediate state, time, and coarse condition \(c\).

\[\mathcal L_{\mathrm{RF}}=\mathbb E\left[\left\|v_\theta(x_t,t,c)-(x_1-x_0)\right\|^2\right].\]

The quantity \(v_\theta\) is a rate of change in the generated tensor. It is not atmospheric wind. The generated wind channels are values inside that tensor.

The equations follow src/pcrf/flow.py and the rectified-flow construction of Liu, Gong, and Liu.

The U-Net uses coarse context and time conditioning.

The velocity network concatenates \(x_t\) and the coarse condition before its input convolution. Residual blocks and skip connections carry features across the U-Net. A sinusoidal time embedding enters the residual blocks, and time-only FiLM modulates the first feature map.

The recorded model uses resolution levels \(128,64,32\) above a \(16\times16\) bottleneck. Self-attention runs at the \(32\times32\) level and in the bottleneck. The benchmark architecture contains 7,565,251 parameters.

The model record and benchmark configuration source these dimensions. The conditioning path uses concatenation, and the attention is self-attention.

The sampler integrates the learned field.

The sampler starts from new Gaussian noise and advances the state through forward Euler steps. The conference configuration uses 20 steps.

\[x_{k+1}=x_k+\Delta t\,v_\theta(x_k,t_k,c),\qquad \Delta t=\frac1K,\qquad K=20.\]

The final precipitation projection acts after the integration loop. The sampler can draw several members from the same coarse condition by changing the initial noise.

The benchmark configuration records the step count. The conference implementation records the loop in src/pcrf/flow.py.

Soft penalties and their limits

The physics terms act on a one-step estimate of the clean field. The estimate extends the current velocity to the end of the interpolation interval.

\[\widetilde x_1=x_t+(1-t)v_\theta(x_t,t,c).\]

The original training objective adds negative-precipitation, domain-mean, and wind-derivative penalties to the flow loss.

\[\mathcal L=\mathcal L_{\mathrm{RF}}+\lambda_{\mathrm{nn}}\mathcal L_{\mathrm{nn}}+\lambda_{\mathrm{mass}}\mathcal L_{\mathrm{mass}}+\lambda_{\mathrm{wind}}\mathcal L_{\mathrm{wind}}.\]

\[\mathcal L_{\mathrm{nn}}=\mathbb E\left[\|\min(\widetilde p,0)\|^2\right],\qquad\mathcal L_{\mathrm{mass}}=\mathbb E\left[(\operatorname{mean}(\widetilde p)-m_c)^2\right].\]

The intended wind term penalizes squared horizontal divergence. The historical ERA5 array order and channel scaling prevented that calculation from representing physical divergence. The later audit supersedes the old derivative results, and the new study drops that penalty.

The released paper documents the original loss. The scope record follows the later research audit and distinguishes the original formulation from the proposed evaluation.

The precipitation projection

The operator clamps and rescales.

The operator first replaces negative precipitation with zero. It then multiplies the entire precipitation channel by a scalar that matches the coarse mean.

\[q_i=\max(\widehat p_i,0),\qquad \mu=\frac1n\sum_{i=1}^{n}q_i.\]

\[p_i^*=q_i\frac{m_c}{\mu}\qquad\text{when }\mu>0\text{ and }m_c\ge0.\]

The proof is direct. Each \(q_i\) is nonnegative, and the rescale factor is nonnegative. The new mean equals \(\mu(m_c/\mu)=m_c\). The equality holds up to floating-point rounding in an implementation.

The rescale preserves ratios among positive cells. It cannot place a rain band, restore a missing storm, or repair an incorrect reference. The clamp changes the field where values were negative.

The operator belongs to established work on constrained generation. The page credits Hess et al. and Harder et al. and makes no novelty claim for the global rescale.

The all-clamped case needs an explicit rule.

An all-clamped field has \(\mu=0\). A positive target mean then makes the division undefined. The conference Python serving copy floors the denominator, while the Rust copy skips rescaling for a sufficiently small clamped mean. Both leave an all-clamped field at zero and fail its positive target.

The conference serving copies also skip rescaling when the target mean is zero. A clamped field that remains wet then fails the dry-target contract. An output gate must check the result rather than assume the runtime enforced that contract.

The legacy serving MCE helper returns zero for sufficiently small targets. That branch can hide the dry-target failure, so an absolute mean-error check is needed.

The newer local research implementation adds a uniform fallback that fills the selected cells with \(m_c\). It also supports masks, area weights, and block constraints. Those changes do not establish that the conference Python and Rust serving runtimes received the same update.

The implementation record separates src/pcrf/physics.py from src/pcrf/serving/math.py and services/inference-rs/src/physics.rs.

The certificate checks a contract.

The non-negativity violation rate (NVR) counts negative precipitation cells. The mass-conservation error (MCE) compares the generated mean with the coarse mean.

\[\mathrm{NVR}=\frac1n\sum_i\mathbf1[p_i<0],\qquad \mathrm{MCE}=\frac{|\operatorname{mean}(p)-m_c|}{|m_c|}\quad(m_c>0).\]

A dry target needs an absolute-error check because the relative denominator vanishes. A useful release gate must also check finite values and verify the mean after projection. A constant field can pass the precipitation contract, so the certificate is necessary but insufficient for scientific use.

Export and serving architecture

The export boundary is the velocity network.

The exporter captures \(v_\theta(x,t,c)\), not a complete request workflow. AOTInductor packages compiled PyTorch execution into a .pt2 artifact. ONNX exports the same network into a graph that ONNX Runtime can execute.

The exported spatial shape stays fixed. The batch dimension is dynamic within its recorded export bound. The integration loop, random-state setup, terminal precipitation projection, and output checks remain runtime responsibilities.

The benchmark compares PyTorch eager, torch.compile, AOTInductor, and ONNX Runtime's CUDA execution provider. The native Rust service uses ONNX Runtime and implements its own sampler and precipitation operator.

The implementation record names export_aoti.py, export_onnx.py, bench_serving.py, and the Rust service. The PyTorch AOTInductor documentation describes the packaging mechanism.

The model and control plane have different jobs.

The runtime scope record identifies the source modules and recorded smoke evidence for each path.

A100 backend benchmark

The benchmark uses the recorded ERA5 architecture with random weights because the trained checkpoints were no longer available. It measures execution cost and numerical consistency, and it measures no climate skill.

The recorded device is an NVIDIA A100-SXM4-40GB. The run uses PyTorch 2.14.0+cu130 and ONNX Runtime 1.30.0. The velocity measurement uses 10 warmups and 100 timed calls. The full-sample measurement uses 2 warmups and 30 timed samples. A full sample contains 20 Euler steps and the terminal projection.

The original JSON records the environment, random-weight flag, iteration counts, and all measurements.

The table rounds the recorded measurements using the conference ledger's display precision. Batch-32 throughput uses the mean full-sample time.

The AOTInductor median is about 2.1 times faster than eager for one full sample in this run. The ratio is \(91.7/43.3\). The timing excludes HTTP transport, queueing, request parsing, cold startup, and the Go or Temporal workflow.

Numerical parity depends on the precision policy.

The default PyTorch configuration enables TF32 for cuDNN convolutions and disables it for matrix multiplication. The recorded maximum absolute velocity error is about \(10^{-3}\) for the compiled PyTorch backends and about \(1.5\times10^{-3}\) for ONNX Runtime.

A separate parity-only run disables both PyTorch TF32 switches. The compiled PyTorch errors then fall below \(2\times10^{-6}\). ONNX Runtime remains near \(10^{-3}\), and the PyTorch switch does not control that backend.

The parity comparison covers velocity outputs at batch sizes 1, 4, and 8. It does not establish matched final ensemble distributions or unchanged climate skill. Every timed backend reports zero post-projection NVR on the benchmark inputs. The largest recorded MCE is about \(7.45\times10^{-8}\).

The default benchmark and TF32-disabled parity record support these results. The page makes no claim about the cause of ONNX Runtime's remaining gap.

Distributed training and the DDP lesson

The earlier training loop wrapped the model in DistributedDataParallel (DDP), then called a custom loss through rf.module. The call bypassed the wrapper's forward path, so the ranks trained independent replicas without the intended gradient synchronization.

The correction adds a RectifiedFlow.forward() that returns the training loss. The loop calls rf(...) through the DDP wrapper. The revised throughput calculation uses the all-reduced sample count and the slowest rank's epoch duration.

# The DDP wrapper receives the forward call.
loss, logs = rf(reference_field, condition, coarse_mean)
loss.backward()

The verified replacement sweep runs on one A100 node with one, two, and four workers. It uses synthetic fields, a \(128\times128\) tile, base width 64, batch 32 per worker, and 40 epochs.

\[\mathrm{efficiency}(W)=\frac{\mathrm{throughput}(W)}{W\,\mathrm{throughput}(1)}.\]

The corrected CSV and sweep configuration record support the table. The values are rounded for display, and the efficiency uses the original CSV values.

A real two-rank CPU regression check compared gradients on different input data before the replacement GPU run. The corrected measurements support a four-worker synthetic throughput claim. The earlier eight-worker and inter-node curves remain invalid as data-parallel results.

What the climate evaluation can establish

The conference evaluation paired ERA5 coarse precipitation with ERA5-Land precipitation. The audit found problems with the reference, time windows, wind metric, and baseline configuration. The historical quality tables are archival, and the page does not reuse them to rank models.

The audit scope record captures the later findings. The ECMWF ERA5-Land documentation describes precipitation forcing and accumulation conventions.

A mechanism check has a narrower purpose.

The project also includes a fixed-seed synthetic-field check without a trained network. It tests whether the chosen operators change their intended residuals on a constructed field. It does not establish generalization, and it cannot separate a fidelity change caused by penalties from one caused by projection when both act together.

A useful evaluation scores each ensemble member before and after projection from the same sampled noise. It reports precipitation-only quality metrics with fixed scaling, fair CRPS, negative-cell and negative-mass rates, and masked or block-level mean errors.

The physical contract and the reference-based scores answer different questions. A conservation check asks whether the generated field agrees with its specified budget. A skill score asks whether the generated field agrees with the independent fine reference.

Proposed evaluation improvements

A stronger evaluation would use an independent precipitation reference and a self-coarsened control. It would separate output validity from predictive skill and isolate the effect of each constraint. The following work is proposed, and the page claims no completed model result from it.

The data pipeline must align the physical quantities.

The proposed ERA5-to-IMERG pair uses matching accumulation windows, explicit units, preserved masks, and conservative remapping. The self-coarsened ERA5 control isolates the constraint mechanism when the fine field and coarse budget agree by construction.

The local preparation code supports six-hour precipitation windows. It can de-accumulate ERA5-Land, convert IMERG rates to depths, and reject incomplete windows. The existence of the preparation code does not establish that the independent-reference study has been run.

The comparison must isolate the constraint choice.

The proposed design crosses penalty on or off with projection on or off. Each arm receives matched training and sampling budgets. Each before-and-after projection score comes from the same members, and every baseline can receive the same projection.

A seeded pilot will estimate the variation before the study sets its final repeat count. Paired day-block resampling will preserve daily dependence when estimating score differences. The study will report uncertainty and projection statistics beside the average scores.

The projection must reveal what it changed.

The audit should record the fraction of clamped cells, the removed negative mass, the rescale factor, and every fallback event. The masked, area-weighted operator should check means on the stated support. A block operator should check each coarse block when the grids nest.

The proposed-study record distinguishes the local preparation and evaluation foundations from the unrun model experiments. The plan follows the later research handoff.

Evidence and reproducibility

The downloadable records preserve the measurements used by this page. The evidence index records source filenames, file hashes, configurations, and claim limits. The records support inspection of the reported engineering runs, and they do not supply the missing climate checkpoints.

A reproduction needs more than a timing number.

A serving reproduction must freeze the graph shape, parameter count, weights or weight seed, backend versions, GPU, TF32 policy, warmup, and timed boundary. It must run output checks after the same sampler and projection. It must also state whether the runtime transfers tensors across the host-device boundary.

A scientific reproduction also needs immutable input and reference versions, matched time windows, masks, split dates, training seeds, saved checkpoints, and per-sample outputs. The current page does not claim those assets are available for the superseded climate comparison.

The page uses the following evidence labels.

The A100 benchmark is recorded GPU runtime evidence with random weights. The DDP sweep is recorded GPU training evidence with synthetic data. The browser demonstration is an illustrative calculation. The older climate evaluation is superseded archival evidence. The independent-reference model study is proposed work.

Related methods and primary sources

Liu, Gong, and Liu introduce rectified flow. Their flow construction supplies the noise-to-data learning objective used here.

Hess et al. constrain generated precipitation budgets. Their work establishes prior art for global precipitation preservation.

Harder et al. study hard constraints for climate downscaling. Their constraint layers establish prior art for multiplicative mean constraints across downscalers.

ECMWF documents ERA5-Land forcing and accumulation. The reference provenance and time conventions determine what an evaluation can measure.

PyTorch documents AOTInductor export and deployment. The artifact boundary helps distinguish network execution from an entire served request.