I’m computing PDE residuals for the The Well datasets (for example turbulent_radiative_layer_2D and shear_flow) using finite-difference methods, but the results are much larger than I expected. Because the data were generated by simulation, I assumed the residuals should be nearly zero.
What I’m doing
I use a central-difference scheme (second-order in the interior, one-sided at the boundaries) to compute derivatives. Here is my derivative function in PyTorch:
import torch
def compute_derivative(u, delta, order=1, axis=1):
"""
Compute numerical derivatives along a specified axis.
Args:
u: Input tensor with shape [B, T, X, Y] (or [B, T, X, Y, C] for vector fields).
delta: grid spacing along the chosen axis.
order: 1 for first derivative, 2 for second derivative.
axis: 1=time, 2=x, 3=y.
"""
derivative = torch.zeros_like(u)
if order == 1:
# second-order central differences in interior, one-sided at boundaries
if axis == 1: # time
derivative[:, 1:-1, :, :] = (u[:, 2:, :, :] - u[:, :-2, :, :]) / (2.0 * delta)
derivative[:, 0:1, :, :] = (-3.0*u[:, 0:1, :, :] + 4.0*u[:, 1:2, :, :] - u[:, 2:3, :, :]) / (2.0 * delta)
derivative[:, -1:, :, :] = (3.0*u[:, -1:, :, :] - 4.0*u[:, -2:-1, :, :] + u[:, -3:-2, :, :]) / (2.0 * delta)
elif axis == 2: # x
derivative[:, :, 1:-1, :] = (u[:, :, 2:, :] - u[:, :, :-2, :]) / (2.0 * delta)
derivative[:, :, 0:1, :] = (-3.0*u[:, :, 0:1, :] + 4.0*u[:, :, 1:2, :] - u[:, :, 2:3, :]) / (2.0 * delta)
derivative[:, :, -1:, :] = (3.0*u[:, :, -1:, :] - 4.0*u[:, :, -2:-1, :] + u[:, :, -3:-2, :]) / (2.0 * delta)
elif axis == 3: # y
derivative[:, :, :, 1:-1] = (u[:, :, :, 2:] - u[:, :, :, :-2]) / (2.0 * delta)
derivative[:, :, :, 0:1] = (-3.0*u[:, :, :, 0:1] + 4.0*u[:, :, :, 1:2] - u[:, :, :, 2:3]) / (2.0 * delta)
derivative[:, :, :, -1:] = (3.0*u[:, :, :, -1:] - 4.0*u[:, :, :, -2:-1] + u[:, :, :, -3:-2]) / (2.0 * delta)
elif order == 2:
# second-order second derivative
if axis == 1:
derivative[:, 1:-1, :, :] = (u[:, 2:, :, :] - 2*u[:, 1:-1, :, :] + u[:, :-2, :, :]) / (delta**2)
derivative[:, 0:1, :, :] = (u[:, 2:3, :, :] - 2*u[:, 1:2, :, :] + u[:, 0:1, :, :]) / (delta**2)
derivative[:, -1:, :, :] = (u[:, -1:, :, :] - 2*u[:, -2:-1, :, :] + u[:, -3:-2, :, :]) / (delta**2)
elif axis == 2:
derivative[:, :, 1:-1, :] = (u[:, :, 2:, :] - 2*u[:, :, 1:-1, :] + u[:, :, :-2, :]) / (delta**2)
derivative[:, :, 0:1, :] = (u[:, :, 2:3, :] - 2*u[:, :, 1:2, :] + u[:, :, 0:1, :]) / (delta**2)
derivative[:, :, -1:, :] = (u[:, :, -1:, :] - 2*u[:, :, -2:-1, :] + u[:, :, -3:-2, :]) / (delta**2)
elif axis == 3:
derivative[:, :, :, 1:-1] = (u[:, :, :, 2:] - 2*u[:, :, :, 1:-1] + u[:, :, :, :-2]) / (delta**2)
derivative[:, :, :, 0:1] = (u[:, :, :, 2:3] - 2*u[:, :, :, 1:2] + u[:, :, :, 0:1]) / (delta**2)
derivative[:, :, :, -1:] = (u[:, :, :, -1:] - 2*u[:, :, :, -2:-1] + u[:, :, :, -3:-2]) / (delta**2)
return derivative
From the HDF5 file I load fields with shapes:
pressure: (4, 200, 256, 512)tracer: (4, 200, 256, 512)velocity: (4, 200, 256, 512, 2)
I pass velocity[...,0] and velocity[...,1] to my function as velocity_x and velocity_y.
The PDEs I use
- Momentum (component-form): $[ \partial_t \mathbf{u} + \nabla p ;-; \nu \Delta \mathbf{u} = -,\mathbf{u}!\cdot!\nabla\mathbf{u} ]$
- Tracer: $[ \partial_t s ;-; D,\Delta s = -,\mathbf{u}!\cdot!\nabla s ]$ where $\nu = 1/\text{Reynolds}$ and $D = \nu/\text{Schmidt}$.
The time, spatial derivatives, and advection terms are computed via:
# time derivatives (first order)
du_x_dt = compute_derivative(velocity_x, dt, order=1, axis=1)
du_y_dt = compute_derivative(velocity_y, dt, order=1, axis=1)
ds_dt = compute_derivative(tracer, dt, order=1, axis=1)
# spatial derivatives
dp_dx = compute_derivative(pressure, dx, order=1, axis=2)
dp_dy = compute_derivative(pressure, dy, order=1, axis=3)
# Laplacians (second order)
laplacian_u_x = compute_derivative(velocity_x, dx, order=2, axis=2) \
+ compute_derivative(velocity_x, dy, order=2, axis=3)
laplacian_u_y = compute_derivative(velocity_y, dx, order=2, axis=2) \
+ compute_derivative(velocity_y, dy, order=2, axis=3)
laplacian_s = compute_derivative(tracer, dx, order=2, axis=2) \
+ compute_derivative(tracer, dy, order=2, axis=3)
# advection terms
advection_u_x = velocity_x * compute_derivative(velocity_x, dx, order=1, axis=2) \
+ velocity_y * compute_derivative(velocity_x, dy, order=1, axis=3)
advection_u_y = velocity_x * compute_derivative(velocity_y, dx, order=1, axis=2) \
+ velocity_y * compute_derivative(velocity_y, dy, order=1, axis=3)
advection_s = velocity_x * compute_derivative(tracer, dx, order=1, axis=2) \
+ velocity_y * compute_derivative(tracer, dy, order=1, axis=3)
Residual tensors (shape [B, T, X, Y]) are computed as:
momentum_x_residual = du_x_dt + dp_dx - nu * laplacian_u_x + advection_u_x
momentum_y_residual = du_y_dt + dp_dy - nu * laplacian_u_y + advection_u_y
tracer_residual = ds_dt - D * laplacian_s + advection_s
And then I compute the mean-squared residuals:
res_mom_x = ((momentum_x_residual)**2).mean().item()
res_mom_y = ((momentum_y_residual)**2).mean().item()
res_tracer = ((tracer_residual)**2).mean().item()
# yes — this is the mean squared residual (MSE)
Observed output:
The momentum equation residual norms: 0.18399687111377716, 0.1050143912434578 The tracer equation residual norm: 2.9967291355133057
I expected values near zero but I'm getting ~0.1 to ~3.0.
What I’ve already checked
- Data shapes and indexing appear correct.
- The derivative implementation uses standard central-difference in interior + one-sided at boundaries.
- Conversion to
torch.float32on device is done.
Questions / things I would appreciate input on
- Am I overlooking any common pitfalls that could cause such large residuals?
- Could the dataset use different conventions (e.g., a staggered grid, pressure/velocity staggering, nondimensional units or scaling, periodic boundary conditions, ghost cells, coordinate transforms) that I should account for?
- Could there be errors in my finite-difference discretization or in the sign conventions of the PDE terms?
- Are there recommended sanity checks I should run to localize the issue?
Additional info (available on request):
- A small sample slice of
pressure,velocity_x,velocity_y,tracer. - Numeric values of
Reynolds,Schmidt,dt,dx,dy. - Statistic summaries (min/mean/max/std) of the fields.
- Residual field histograms / visualization.
Thanks in advance for any pointers—much appreciated!
The total code is :
import torch
import torch.nn.functional as F
from utils.criterion import SimpleLpLoss, masked_mse
from turbulent_radiative_layer_2D import visualize_residuals, visualize_residuals_abs
import numpy as np
import torch
from typing import Union, Dict
ArrayLike = Union[np.ndarray, torch.Tensor]
def compute_physics_loss(dt, dx, dy, Reynolds, Schmidt,
pressure, tracer, velocity_x, velocity_y, device):
nu = 1.0 / Reynolds
D = nu / Schmidt
# 转换为torch tensor
# device = pressure.device
if not isinstance(pressure, torch.Tensor):
pressure = torch.tensor(pressure, dtype=torch.float32, device=device)
tracer = torch.tensor(tracer, dtype=torch.float32, device=device)
velocity_x = torch.tensor(velocity_x, dtype=torch.float32, device=device)
velocity_y = torch.tensor(velocity_y, dtype=torch.float32, device=device)
# time derivatives (one-order)
du_x_dt = compute_derivative(velocity_x, dt, order=1, axis=1)
du_y_dt = compute_derivative(velocity_y, dt, order=1, axis=1)
ds_dt = compute_derivative(tracer, dt, order=1, axis=1)
# spatial derivatives
# pressure gradients (one-order)
dp_dx = compute_derivative(pressure, dx, order=1, axis=2)
dp_dy = compute_derivative(pressure, dy, order=1, axis=3)
# velocity Laplacian (second-order)
du_x_dx2 = compute_derivative(velocity_x, dx, order=2, axis=2)
du_x_dy2 = compute_derivative(velocity_x, dy, order=2, axis=3)
laplacian_u_x = du_x_dx2 + du_x_dy2
du_y_dx2 = compute_derivative(velocity_y, dx, order=2, axis=2)
du_y_dy2 = compute_derivative(velocity_y, dy, order=2, axis=3)
laplacian_u_y = du_y_dx2 + du_y_dy2
# tracer Laplacian (second-order)
ds_dx2 = compute_derivative(tracer, dx, order=2, axis=2)
ds_dy2 = compute_derivative(tracer, dy, order=2, axis=3)
laplacian_s = ds_dx2 + ds_dy2
# convection term (u·∇u)
du_x_dx = compute_derivative(velocity_x, dx, order=1, axis=2)
du_x_dy = compute_derivative(velocity_x, dy, order=1, axis=3)
advection_u_x = velocity_x * du_x_dx + velocity_y * du_x_dy
du_y_dx = compute_derivative(velocity_y, dx, order=1, axis=2)
du_y_dy = compute_derivative(velocity_y, dy, order=1, axis=3)
advection_u_y = velocity_x * du_y_dx + velocity_y * du_y_dy
# convection term (u·∇s)
ds_dx = compute_derivative(tracer, dx, order=1, axis=2)
ds_dy = compute_derivative(tracer, dy, order=1, axis=3)
advection_s = velocity_x * ds_dx + velocity_y * ds_dy
# Momentum Equation Residuals
momentum_x_residual = du_x_dt + dp_dx - nu * laplacian_u_x + advection_u_x
momentum_y_residual = du_y_dt + dp_dy - nu * laplacian_u_y + advection_u_y
# Tracer Equation Residual
tracer_residual = ds_dt - D * laplacian_s + advection_s
residuals = {
'momentum_x': momentum_x_residual,
'momentum_y': momentum_y_residual,
'tracer': tracer_residual,
}
return residuals
def compute_derivative(u, delta, order=1, axis=1):
derivative = torch.zeros_like(u)
if order == 1:
if axis == 1:
derivative[:, 1:-1, :, :] = (u[:, 2:, :, :] - u[:, :-2, :, :]) / (2.0 * delta)
derivative[:, 0:1, :, :] = (-3.0 * u[:, 0:1, :, :] + 4.0 * u[:, 1:2, :, :] - u[:, 2:3, :, :]) / (2.0 * delta)
derivative[:, -1:, :, :] = (3.0 * u[:, -1:, :, :] - 4.0 * u[:, -2:-1, :, :] + u[:, -3:-2, :, :]) / (2.0 * delta)
elif axis == 2:
derivative[:, :, 1:-1, :] = (u[:, :, 2:, :] - u[:, :, :-2, :]) / (2.0 * delta)
derivative[:, :, 0:1, :] = (-3.0 * u[:, :, 0:1, :] + 4.0 * u[:, :, 1:2, :] - u[:, :, 2:3, :]) / (2.0 * delta)
derivative[:, :, -1:, :] = (3.0 * u[:, :, -1:, :] - 4.0 * u[:, :, -2:-1, :] + u[:, :, -3:-2, :]) / (2.0 * delta)
elif axis == 3:
derivative[:, :, :, 1:-1] = (u[:, :, :, 2:] - u[:, :, :, :-2]) / (2.0 * delta)
derivative[:, :, :, 0:1] = (-3.0 * u[:, :, :, 0:1] + 4.0 * u[:, :, :, 1:2] - u[:, :, :, 2:3]) / (2.0 * delta)
derivative[:, :, :, -1:] = (3.0 * u[:, :, :, -1:] - 4.0 * u[:, :, :, -2:-1] + u[:, :, :, -3:-2]) / (2.0 * delta)
elif order == 2:
if axis == 1:
derivative[:, 1:-1, :, :] = (u[:, 2:, :, :] - 2 * u[:, 1:-1, :, :] + u[:, :-2, :, :]) / (delta ** 2)
derivative[:, 0:1, :, :] = (u[:, 2:3, :, :] - 2 * u[:, 1:2, :, :] + u[:, 0:1, :, :]) / (delta ** 2)
derivative[:, -1:, :, :] = (u[:, -1:, :, :] - 2 * u[:, -2:-1, :, :] + u[:, -3:-2, :, :]) / (delta ** 2)
elif axis == 2:
derivative[:, :, 1:-1, :] = (u[:, :, 2:, :] - 2 * u[:, :, 1:-1, :] + u[:, :, :-2, :]) / (delta ** 2)
derivative[:, :, 0:1, :] = (u[:, :, 2:3, :] - 2 * u[:, :, 1:2, :] + u[:, :, 0:1, :]) / (delta ** 2)
derivative[:, :, -1:, :] = (u[:, :, -1:, :] - 2 * u[:, :, -2:-1, :] + u[:, :, -3:-2, :]) / (delta ** 2)
elif axis == 3:
derivative[:, :, :, 1:-1] = (u[:, :, :, 2:] - 2 * u[:, :, :, 1:-1] + u[:, :, :, :-2]) / (delta ** 2)
derivative[:, :, :, 0:1] = (u[:, :, :, 2:3] - 2 * u[:, :, :, 1:2] + u[:, :, :, 0:1]) / (delta ** 2)
derivative[:, :, :, -1:] = (u[:, :, :, -1:] - 2 * u[:, :, :, -2:-1] + u[:, :, :, -3:-2]) / (delta ** 2)
return derivative
def compute_residual_norms(residuals):
norms = {}
for key, residual in residuals.items():
# L2 Norm (Root Mean Square)
l2_norm = torch.sqrt(torch.mean(residual**2))
# L∞ Norm (Max Absolute Value)
linf_norm = torch.max(torch.abs(residual))
# L1 Norm (Mean Absolute Value)
l1_norm = torch.mean(torch.abs(residual))
norms[key] = {
'L2': l2_norm.item(),
'Linf': linf_norm.item(),
'L1': l1_norm.item()
}
return norms
if __name__ == "__main__":
import h5py
filename = 'dataset/shear_flow_Reynolds_5e5_Schmidt_2e-1.hdf5'
device = 'cuda' if torch.cuda.is_available() else 'cpu'
with h5py.File(filename, 'r') as f:
pressure = f['t0_fields/pressure'][:]
tracer = f['t0_fields/tracer'][:]
velocity = f['t1_fields/velocity'][:]
Reynolds = float(f['scalars/Reynolds'][()])
Schmidt = float(f['scalars/Schmidt'][()])
time = f['dimensions/time'][:]
x = f['dimensions/x'][:]
y = f['dimensions/y'][:]
dt = time[1] - time[0]
dx = x[1] - x[0]
dy = y[1] - y[0]
print(f'dt={dt}, dx={dx}, dy={dy}, Reynolds={Reynolds}, Schmidt={Schmidt}')
print(f'pressure shape: {pressure.shape}, tracer shape: {tracer.shape}, velocity shape: {velocity.shape}')
# pressure shape: (4, 200, 256, 512), tracer shape: (4, 200, 256, 512), velocity shape: (4, 200, 256, 512, 2)
residuals = compute_physics_loss(dt, dx, dy, Reynolds, Schmidt,
pressure, tracer, velocity[..., 0], velocity[..., 1],
device)
res_mom_x = ((residuals['momentum_x'])**2).mean().item()
res_mom_y = ((residuals['momentum_y'])**2).mean().item()
res_tracer = ((residuals['tracer'])**2).mean().item()
print(f'\nThe momentum equation residual norms: {res_mom_x}, {res_mom_y}')
print(f'The tracer equation residual norm: {res_tracer}')
norms = compute_residual_norms(residuals)
# Residual Statistics
print("\nResidual Statistics:")
for key, residual in residuals.items():
print(f"\n{key}:")
print(f" Shape: {residual.shape}")
print(f" Mean: {residual.mean().item():.6e}")
print(f" Std: {residual.std().item():.6e}")
print(f" Min: {residual.min().item():.6e}")
print(f" Max: {residual.max().item():.6e}")
```
