Geometry-Aware Continuous Flow Operator for 2D Shallow Water Equations

Abstract

Continuous Flow Operators (CFO) learn a continuous-time vector field from discrete physical trajectories, rather than a fixed temporal transition as neural operators do. For two dimensional shallow-water flow over variable bathymetry, a surrogate has to respect conservation laws, terrain-dependent forcing and equilibrium states, and it has to stay stable over long rollouts. We build a geometry-aware, well-balanced, physics-informed CFO for the 2D shallow-water equations. The conservative state is q= [h,hu,hv], and the bed elevation b(x,y) is kept as a separate conditioning field. Beyond the CFO flow-matching loss, we added a shallow-water residual, counterfactual bed-response loss, a lake-at-rest well-balanced loss and a rollout loss. Bathymetry enters through the geometric features [b,bx,by], which a learned geometry encoder processes before they are coupled to a U-FNO backbone. The reason for introduction of rollout loss was due to the accumulation of small local errors pile up during free rollout even though the recursive prediction in the one-step vector-field predictions can look accurate. To counter this, we fine-tune the final 64 ×64 model with a short differentiable horizon intervals RK4 rollout loss with four substeps. This helps in dropping the final Geometry-U-FNO model’s ID relative L2 error, mean Gaussian terrain-effect error, the mean OOD terrain-effect error and final ID mass drift at 64 ×64. We then evaluate the same 64 ×64 checkpoint zero-shot on a 128 ×128 grid. Applying the local U-Net branch directly at the finer resolution caused strong, resolution-dependent oscillations. We fixed this with a canonical U-Net evaluation. The 128 ×128 latent field is area-averaged down to the 64 ×64 training grid, the trained U-Net runs there, and its response is reconstructed at 128 ×128 by cubic interpolation. The Fourier and geometry branches stay at 128 ×128, and no neural training is done at that resolution. Which helped improve the errors in both Gaussian case and the unseen cases. Together, these results show stable zero-shot resolution transfer that keeps the continuous-time rollout formulation intact.