Physics‑Informed Neural Networks for Navier‑Stokes Flow Parameter Identification

Abstract

The article presents a complete implementation of a PINN‑based inverse solver for the incompressible Navier‑Stokes equations governing cylinder wake flow. By letting the network output a stream function, the continuity constraint is satisfied analytically. High‑order spatial derivatives required by the PDE residual are obtained efficiently with nested forward‑mode automatic differentiation (JVP) in JAX, avoiding the compilation bottleneck of reverse‑mode. A two‑stage optimization strategy (Adam for global exploration followed by L‑BFGS‑B for local refinement) yields rapid convergence. Experiments on synthetic data from a spectral element solver show parameter identification errors below 0.2 % for the convection coefficient and around 5 % for the viscosity coefficient, with negligible degradation under 1 % Gaussian noise, and a total training time of roughly 6 minutes on an NVIDIA RTX 4090.

1. Introduction

The Navier‑Stokes equations describe viscous incompressible fluid motion and appear in aerospace, marine, and biomedical applications. In many practical situations the physical parameters (e.g., viscosity) are unknown and must be inferred from limited observations, forming a classic PDE parameter‑identification inverse problem. Traditional approaches rely on adjoint equations or optimal control, which are complex and sensitive to initial guesses. Physics‑informed neural networks (PINNs) embed the governing equations as soft constraints in the loss function, allowing simultaneous training of the solution field and unknown parameters.

2. Methodology

2.1 Stream‑function parameterization

Instead of directly predicting the velocity components (u, v), the network outputs a scalar stream function ψ. The velocity field is recovered as u = ∂ψ/∂y and v = –∂ψ/∂x, which automatically satisfies the incompressibility condition ∇· v = 0. This reduces the number of constraint equations but requires third‑order derivatives of ψ.

2.2 Nested JVP for high‑order derivatives

JAX provides two automatic‑differentiation modes: reverse‑mode (grad/VJP) and forward‑mode (jvp). For the Navier‑Stokes problem the input dimension (x, y, t) is low while the output dimension (velocity components and pressure) is high, making forward‑mode more efficient for high‑order derivatives. By nesting jax.jvp calls, the required first‑, second‑, and third‑order spatial derivatives are computed with a single compiled graph. The implementation requires 15 JVP calls (four of them triple‑nested) to assemble all terms of the PDE residual.

2.3 Two‑stage optimization (Adam + L‑BFGS‑B)

The loss consists of four equally weighted terms: data mismatches for u and v, and PDE residuals for the momentum equations in x and y. Phase 1 uses Adam for 200 k iterations to bring the parameters into a reasonable basin. Phase 2 flattens the JAX pytree with jax.flatten_util.ravel_pytree and hands the resulting vector to scipy.optimize.minimize with the L‑BFGS‑B algorithm, which converges in a handful of iterations thanks to the good initialization from Adam.

3. Implementation Details

3.1 Network architecture

The network is a fully‑connected feed‑forward model with eight hidden layers of equal width. Input coordinates are min‑max normalized to the interval [‑1, 1] before passing through tanh activations. The final linear layer outputs two scalars (ψ and an auxiliary variable for pressure gradient construction). Xavier initialization is used for all weight matrices.

def net_batch(params, X, lb, ub):
    H = 2.0 * (X - lb) / (ub - lb) - 1.0  # [N, 3] → [-1, 1]^3
    for W, b in params[:-1]:
        H = jnp.tanh(H @ W + b)          # hidden layers
    W, b = params[-1]
    return H @ W + b                     # output layer, [N, 2]

3.2 Nested JVP implementation

def compute_ns_derivs(net_params, X, lb, ub):
    f_psi = lambda X_: net_batch(net_params, X_, lb, ub)[:, 0]
    f_p   = lambda X_: net_batch(net_params, X_, lb, ub)[:, 1]
    # basis vectors for forward‑mode differentiation
    tx = jnp.zeros_like(X).at[:, 0].set(1.0)
    ty = jnp.zeros_like(X).at[:, 1].set(1.0)
    tt = jnp.zeros_like(X).at[:, 2].set(1.0)
    # first‑order derivatives
    _, psi_y = jax.jvp(f_psi, (X,), (ty,))
    _, psi_x = jax.jvp(f_psi, (X,), (tx,))
    u, v = psi_y, -psi_x
    # second‑order (nested JVP)
    dpsi_dy_fn = lambda X_: jax.jvp(f_psi, (X_,), (ty,))[1]
    _, u_t = jax.jvp(dpsi_dy_fn, (X,), (tt,))
    _, u_x = jax.jvp(dpsi_dy_fn, (X,), (tx,))
    # third‑order (nested twice)
    d2psi_dydx_fn = lambda X_: jax.jvp(dpsi_dy_fn, (X_,), (tx,))[1]
    _, u_xx = jax.jvp(d2psi_dydx_fn, (X,), (tx,))
    # ... analogous calls for the remaining derivatives
    return (u, v, ...)  # full derivative set

3.3 Loss function with has_aux

def ns_loss_full(all_params, X, u_data, v_data, lb, ub):
    net_params, lam1, lam2 = all_params
    (u, v, p, u_t, u_x, u_y, u_xx, u_yy,
     v_t, v_x, v_y, v_xx, v_yy, p_x, p_y) = compute_ns_derivs(...)
    f_u = u_t + lam1*(u*u_x + v*u_y) + p_x - lam2*(u_xx + u_yy)
    f_v = v_t + lam1*(u*v_x + v*v_y) + p_y - lam2*(v_xx + v_yy)
    ld_u = jnp.mean((u_data - u)**2)
    ld_v = jnp.mean((v_data - v)**2)
    lp_fu = jnp.mean(f_u**2)
    lp_fv = jnp.mean(f_v**2)
    total = ld_u + ld_v + lp_fu + lp_fv
    return total, (ld_u, ld_v, lp_fu, lp_fv)

loss_val_grad = jax.value_and_grad(ns_loss_full, has_aux=True)

4. Experiments and Results

Configuration : Reference solution generated with the Nektar spectral element solver (cylinder wake, Re = 100). Spatial domain discretized with 5 000 unstructured points; 200 time snapshots; 5 000 points randomly sampled (0.5 % of the full grid) for training. Network: 8 hidden layers, ≈3 062 trainable weights plus two scalar parameters (convection and viscosity). Computations performed in float32 on the GPU, with float64 conversion for the L‑BFGS‑B stage.

Training dynamics : The total loss converges similarly for clean and 1 % noisy data. Adam reduces the loss to a plateau after ~200 k steps; L‑BFGS‑B then reaches the final tolerance in 6–7 iterations. Training time: 367 s (Adam) + 26 s (L‑BFGS‑B) ≈ 6.5 min for the clean case; the noisy case is slightly faster due to JIT cache reuse.

Parameter identification : The convection coefficient (true value = 1.0) is recovered as 0.998355 (error ≈ 0.16 %). The viscosity coefficient (true value = 0.01) is recovered as 0.010552 (error ≈ 5.5 %). Adding 1 % Gaussian noise changes the errors only marginally (0.11 % and 5.1 % respectively), demonstrating the regularizing effect of the PDE residual.

Field reconstruction : Relative L2 errors for the velocity components fall below 0.5 % for the clean data and remain comparable for noisy data. Visual comparison shows excellent agreement in the far‑field and minor discrepancies confined to the vortex‑shedding region downstream of the cylinder.

5. Conclusions and Outlook

The tutorial confirms that (1) stream‑function parameterization is an effective way to enforce incompressibility, (2) nested forward‑mode JVP dramatically reduces the cost of computing third‑order derivatives needed for Navier‑Stokes PINNs, and (3) a two‑stage Adam + L‑BFGS‑B optimizer yields fast convergence with total wall‑clock time under seven minutes on a modern GPU. Parameter identification accuracy correlates with the sensitivity of the loss to each coefficient: the convection term (order 1) is identified with sub‑0.2 % error, while the viscosity term (order 0.01) exhibits larger relative error (~5 %). The PDE residual provides inherent denoising, as 1 % observation noise has little impact on the final estimates. Future work includes adaptive loss weighting, extension to three‑dimensional vector potentials, causal PINN training for long‑time predictions, multi‑fidelity data fusion, and domain‑decomposition strategies for industrial‑scale fluid‑mechanics problems.