An End-to-End PyTorch Interface for Differentiable PDE Solvers: A RANS Model-Correction Study

The gist The work presents an end-to-end strategy for solving inverse problems constrained by Partial Differential Equations within a fully differentiable Machine Learning framework.

How it works

1. The approach combines a baseline differentiable PDE solver, which predicts the state w from the nonlinear system R(w) = 0, with a generic additive, parametrized and differentiable correction f(w), with trainable parameters p.

2. The PDE is reformulated as an implicit layer, enabling its integration into arbitrary objective functions by treating it as an implicit layer.

3. Two algorithmic problems are addressed: first, for given parameters p, seeking the state w that solves R(w) + fp(w) = 0 by minimizing the norm of the residual J(w) = R(w) + fp(w)2. Second, for a given cost functional J (w, p), seeking parameters p such that R(w) + fp(w) = 0 by minimizing J (2.4).

4. For solving for the state with fixed parameters, residual minimization is used by minimizing the norm of the residual Q = [R(w) + fp(w)]2. This involves computing the differential of J (w) using an explicit operation and back-propagation through the PDE solver and closure term.

5. For solving for parameters, a Lagrangian L (w, n, p) is introduced, leading to Karush-Kuhn-Tucker optimality conditions that allow for gradient descent strategy where the gradients are computed as D M-1 (∂Jp)T - (∂fp(w))T (∂wR + ∂wfp(w)) - T (∂wJ)T.

6. The end-to-end differentiable learning pipeline is organized around a modular structure composed of three fundamental components: enforcing boundary conditions, evaluating the discrete residual of the governing equations, and applying external or model-derived forcing terms.

Data assimilation and closure with compressible RANS equations

1. The steady compressible RANS equations are written as a system (3.1). To close the system, models relating PR and QT to the state variables are required, most commonly relying on the Boussinesq eddy viscosity assumption (Boussinesq 1877).

2. One approach defines the baseline model with state vector w = [ρ, ρU, ρV, ρE] and the discrete version of (∗) in (3.3) for R(w). The correction term f (w, α) in (2.11) then represents the left-over part of the equation, the discretization of term (∗∗), with α = μt.

3. In the data assimilation setting, μt is treated as a spatially varying tunable parameter, μt(x). In contrast, the closure problem (2.12) consists of identifying a functional relationship between the eddy viscosity and the state, μt,p(w).

4. An alternative approach incorporates a turbulence model directly into R, considering the SA model where the eddy viscosity μt depends analytically on an additional variable ν˜ (i.e., μt(ν˜)) governed by (Spalart and Allmaras 1992). The state vector is augmented as w = [ρ, ρU, ρV, ρE, ρν˜], and the baseline operator R now includes the full (∗) + (∗∗) equation (3.3), and the turbulence variable equation (3.5).

The end-to-end differentiable learning pipeline applied to BROADCAST

1. The BROADCAST solver was chosen as core CFD solver because of the possibility of automatically differentiating its Fortran routines thanks to the software Tapenade (Hascoët 2004). Every routine can be linearized to obtain the Jacobian-vector product (JVP) and the vector-Jacobian product (VJP).

2. The discrete equations are written as M w n+1 − w n Δt = R(w n+1) + P fp(w n+1), where M is a diagonal matrix representing the volume of each cell.

3. The computational flow follows the sketch introduced in Section 2.3, structured around a main module delivering R(w), RANS, composed of three internal layers: BC, NS and FORCE. Each layer serves as an interface to underlying Fortran routines from the BROADCAST solver (Poulain et al 2023).

4. The custom layers (BC, NS, and FORCE) can be composed into a PyTorch nn.Module: RANS, which automatically handles the chaining of gradients and the definition of the backward pass.

Results

1. The routines are first validated on NASA’s 2D Wall-Mounted Hump for data assimilation of the β parameter (using both explicit and implicit strategies).

2. In § 4.3, we demonstrate that, given the SA solution for the VKI LS-59 case, wSA = [ρ, ρU, ρE], we are able to satisfactorily reconstruct the SA eddy viscosity field μSA t using the explicit residual minimization strategy by tuning α = μt.

3. In § 4.4.2, a CNN-based algebraic turbulence model μt,p(w) is constructed that reproduces the SA eddy viscosity μSA t over a custom dataset. The optimization is carried out using the explicit strategy by minimizing the residual of a corrected model on the custom dataset.

4. The CNN-based algebraic eddy viscosity model can be regarded as strongly local, since its kernel width is limited to that of the spatial scheme, i.e., a few grid cells in each spatial direction.

Discussions and conclusions

1. The first two test cases show that both optimization strategies are fully operational, correctly propagating gradients through the entire computational chain in a truly end-to-end fashion.

2. The results further confirm the effectiveness of the proposed approach in predicting the eddy viscosity directly from the conservative flow variables.

3. The framework is general and extends naturally to PDE-constrained systems with embedded closure models, establishing a consistent and scalable paradigm for integrating trainable components within physics-based solvers via differentiable formulations and adjoint-based optimization.

A. PyTorch routines

All routines described in this work are implemented in Python, while the underlying CFD solver, including its numerical schemes, is implemented in Fortran and compiled (e.g., using Intel or gfortran compilers) as Python-callable libraries. These are interfaced with PyTorch via ctypes and wrapped using torch.autograd.Function.

B. Completeness and sensitivity of the dataset

1500 distinct parameter sets were generated, resulting in 840 successfully converged cases, having defined a tolerance for the value of the RANS residual.

A univariate statistical analysis was performed, constructing a histogram for each parameter to illustrate the distribution of the converged cases within the original design space.

A variance-based global sensitivity analysis was conducted using Sobol indices to assess the influence of each input parameter on the variance of the output.

C. Computing the gradients for the backward operation in PyTorch

The change of variable is performed by introducing a new variable θtilde = Np at the definition of the problem, and then directly computing partial derivative of the Lagrangian with respect to θtilde as D Np-1 (∂Jp)T + (∂fp(w))T M-1 ∂Jp = D Np-1 (∂J)T + (∂fp(w))T M-1 ∂J, since M is a symmetric matrix. This is precisely as computed in (2.25), but with respect to the new variable θtilde.

D. Architecture of the developed Convolutional Neural Network

The NN employed in this work is designed to approximate the relationship between the mean-flow conservative variables and the turbulent eddy viscosity field, μt (w). The architecture combines one-dimensional convolutional layers with fully connected (linear) layers and nonlinear activation functions. The choice of one-dimensional convolutions is guided by the structured nature of the mesh and ensures invariance with respect to rotations aligned with the grid directions. A key aspect of the design concerns the spatial support of the convolutional kernels, which is kept smaller than the stencil width of the underlying numerical scheme.

E. Layer Type Specification

The implementation uses a class BC(torch.autograd.Function) where @staticmethod def forward(u): return bc.f90(u). This shows how individual computational blocks are defined and wrapped as differentiable layers.

**F.

1. A fully differentiable framework for solving inverse problems constrained by PDEs allows for the optimization of parameters within a fully differentiable Python workflow by reformulating the PDE as an implicit layer, enabling end-to-end optimization that unifies simulation and learning within a single workflow.

2. The system can perform data assimilation by optimizing spatially varying parameters, such as eddy viscosity, where the framework is applied to reconstruct the full state of a nonlinear multiscale system from partial observations by solving an inverse problem.

3. The improved AI system can generate high-fidelity surrogate models for turbulence closures using Convolutional Neural Networks (CNNs) trained on datasets generated via Design of Experiments (DOE), such as the VKI LS-59 turbine blade, to predict eddy viscosity fields that accurately reconstruct the eddy viscosity in the wake of profiles.

4. The system can perform algebraic turbulence modeling by training a CNN to model closure relations, enabling the creation of models like a CNN-based algebraic eddy viscosity model that are capable of reproducing results from one-equation RANS models, even in regions where classical models perform poorly.

5. The architecture allows for the integration of traditional solvers with ML components by implementing custom layers that override PyTorch’s default backward behavior by solving an associated adjoint system to compute exact gradients with respect to the trainable parameters.