Turbulence Closure Modeling with Differentiable Physics

SEPT 8, 2022

Vedant Puri

Advisor: Venkat Viswanathan

Mechanical Engineering PhD Qualifying Examination

Carnegie Mellon University

Agenda

  • Introduction
    • Turbulence Modeling
    • Differentiable Physics
  • Method and Results 
    • LinearSolve.jl
    • Benchmarks
  • Future Work

Introduction

Learning residual physics in multiphysics simulations

\frac{D}{D t} \begin{bmatrix} m \\ \vec{p} \\ E \\ \cdot \end{bmatrix} = \overbrace{\begin{bmatrix} \cdot \\ \cdot \\ \cdot \\ \cdot \end{bmatrix}}^{\text{source}} - \overbrace{\begin{bmatrix} \cdot \\ \cdot \\ \cdot \\ \cdot \end{bmatrix}}^{\text{sink}}

Need analytical forms of all sources and sinks!

Residual physics

=

Ground truth

-

Analytical form

Goal: Embed symmetry-representing functions (neural networks) inside conservation laws to learn residual physics.

Conservation Laws

Symmetries (Noether's theorem)

Conservation Law Symmetry
Linear momentum Translation
Angular momentum Rotation
Time Energy
\iff

Turbulence Closure Modeling

Navier-Stokes system

\partial_t \vec{v} + (\vec{v}\cdot\nabla)\vec{v} = -\nabla p + \frac{1}{\mathit{Re}}\Delta \vec{v} + f\\ \nabla\cdot\vec{v} = 0

(Pandya, 2013)

Mode coupling

u = \bar{u} + u'
\partial_t \bar{v}_i + \bar{v}_j \bar{v}_{i,j} = -\bar{p}_{,i} + \frac{1}{\mathit{Re}}\bar{v}_{i,jj} - \partial_j \overline{v'_i v'_j} \\ \bar{v}_{j,j} = 0

Closure model

Filtering

Large Eddy Simulation equations

Aim: Develop a framework to embed symmetry-respecting functions inside governing equations to improve closure models.

Differentiable Physics

Orthogonal Functions Deep Neural Networks






 

 

 

ML beats the Curse of Dimensionality - ish

f = \tilde{f} + \mathcal{O}(h)
\tilde{f}(x) = \Sigma_{i=1}^N f_i \phi_i(x)
\tilde{f}(x) = Z_L \circ (\dotsc (\sigma Z_0(x)))

\( N \) parameters, \(M\) points

\( h \sim N^{-c/d} \)

\( h \sim 1 / N \) (for 2-layer networks)

\( N \) points

\( \dfrac{d}{dx} \tilde{f}\sim \mathcal{O}(N^2) \) (exact)

\( \dfrac{d}{dx} \tilde{f} \sim \mathcal{O}(N) \) (exact, AD)

\( \int_\Omega \tilde{f} dx \sim \mathcal{O}(N) \) (quadrature, exact)

(Weinan, 2020)

\( \int_\Omega \tilde{f} dx \sim \mathcal{O}(M) \) (Monte-Carlo, approx)

Adding ML to the simulation pipeline

Mesh

Challenge: Embed ML inside PDE problems

\(A\underline{u} = M\underline{b} \)

Domain

Governing Equation

Boundary Constraint

\( NN_\theta \)

Discretization

\( \dfrac{d}{dt} \underline{u} = f(\underline{u}) \)

Solving

Discrete Problem

\(u(\underline{x},t) \)

\(u(\underline{x}) \)

Time-Stepper

Linear Solver

Solution

Loss

Backpropogation

Ground truth

\( NN_\theta \)

\( NN_\theta \)

Backpropogation via the adjoint method

\frac{d}{dt} u = f(u, \theta)
L(u) = \int_0^T l(u, t) dt

Interpolating Adjoint Method for ODEs

Forward pass

Backward pass

(Daulbaev, 2020)

\lambda(t) = \int_T^t -\lambda^T f + l_u dt
\lambda(T) = 0
u(0) = u_0
\frac{d}{d\theta} L = \int_0^T -\lambda^T f_\theta dt

(Johnson, 2006)

A u = b

Adjoint Method for linear systems

\frac{d}{d\theta} L = -\lambda^T(A_\theta u - b_\theta)
L(u)
A^T\lambda = L_u^T

AD tools are incomplete

  1. cannot differentiate many types of code (external libraries, mutating code, etc)
  2. generate suboptimal code for iterative/approximate algorithms

Method and Results

Numerical PDEs pipeline

Meshing

Linear System

-\Delta u = f
A \underline{u} = M \underline{f}
u(x) = \sum_{i} u_i \phi_i(x)
A\(Mf)

Sparse

Structured

+

+

u|_{\partial\Omega} = g

PDE System

Boundary Data

\{ \phi_i(x) \}_{i=1}^N

Function Space

Solve Scheme

Iterative Solver

Enforce Locality

\frac{\mathrm{d}}{\mathrm{d}x}(u) = D \underline{u}
\int_\Omega v u \,\text{d} x = \underline{v}^T M\underline{u}

Cheap

Domain

We are building an differentiable PDEs ecosystem

Strategy: Build out a unified PDEs ecosystem for Julia SciML!

Abstractions

Composable

AD support

Fast solvers

DL ecosystem

Large (NNs + solvers)

Multiple discretizations

Wants

Fast adjoints

Method

Interoperability

Optimized methods

High performance

High level

General unified interface

Wrap SOTA solvers

ML based discrs

\(\texttt{LinearSolve.jl}\) - high performance linear solvers

using LinearSolve, LinearAlgebra
A = rand(1000, 1000)
b = rand(1000)
prob = LinearProblem(A, b; u0=zero(1000))
sol = solve(prob)
sol = solve(prob, LUFactorization())
sol = solve(prob, IterativeSolversJL_GMRES())
Pl = Diagonal(A)
sol = solve(prob, KrylovJL_CG(), Pl=Pl,
            abstol=1e-8, maxiters=20)

cache = init(prob)
sol = solve(cache)

cache = LinearSolve.set_b(rand(1000))
sol2 = solve(cache)

using OrdinaryDiffEq, LinearSolve
prob = ODEProblem(A, rand(1000), (0, 1))
sol = solve(prob, SBDF2(
                        linsolve=KrylovJL_GMRES(),
                        precs=incompletelu,
                        tol=1e-8,
                        maxiters=20,
                       )

Unified, general, interface over linear solvers in Julia

\(\texttt{LinearSolve.jl}\) sped up large stiff PDEs solves in \(\texttt{SciML}\) ecosystem!

Adjoint support over interface

Cache reuse

Wraps it all!

Expose solver choice to user in downstream packages

Boundary value problem example

using AbstractPDEInterfaces, NodalPolynomialSpaces

N = 64
ν = 0.01

xdom = IntervalDomain(-1, 1, boundary_tags=(:left, :right))
ydom = IntervalDomain(-1, 1, boundary_tags=(:lower, :upper))
dom = xdom ⊗ ydom

space = GaussLobattoLegendreSpace(N, N; dom=dom)
(x, y,) = points(space)

bcs = Dict(
           :left => NeumannBC(),
           :right => DirichletBC(),

           :lower => DirichletBC(),
           :upper => NeumannBC(),
          )

discr = Galerkin()
lhsOp = diffusionOp(ν, space, discr)

f = @. 0*x + 1

prob = BoundaryValueProblem(lhsOp, f, bcs, space, discr,
                            abstol=1e-8, reltol=1e-8)

alg  = LinearBoundaryValueAlg(linalg=KrylovJL_CG())

@time sol = solve(prob, alg; verbose=false)
# 0.201703 seconds (39.71 k allocations: 49.417 MiB)
@test sol.resid < 1e-8
# Test Passed

Plug & play

modular

Domains

Function Spaces

\( \texttt{FourierSpaces.jl}  \)

Solve schemes

Boundary constraint

Galerkin

Collocation

Rectangular

Mapped

Dirichlet

Neumann

Unified API over PDE solvers

\( \texttt{CalculustCore.jl} \)

*wip

Meshed*

Implicit*

\( \texttt{SpectralElementSpaces.jl} \)*

Periodic*

Robin

\( A \underline{u} = M \underline{b} \)

\( \dfrac{d}{dt}\underline{u} = f(\underline{u}) \)

\( \texttt{NodalPolynomialSpaces.jl}  \)

\( \texttt{SciMLOperators.jl} \)

\( \texttt{LinearSolve.jl}  \)

\( \texttt{OrdinaryDiffEq.jl}  \)

CalculustJL
  • Plug-and-play ecosystem for PDE solvers
    • Make it easy to experiment with discretizations
    • You don't need to redefine Laplacian every time you write a new derivative method
    • Just need gradient operator, mass matrix to have full suite of vector-calculus operations
    • Lazily implemented operator algebra via SciMLOperators.jl
  • Modular
    • Switching from 'Fourier spectral collocation' to 'Chebyshev Galerkin projection' is a 1 line change
    • Domain management tools
    • Dictionary based boundary handling interface

Enabling ML-based discretizations

Learning the closure to Burgers turbulence

\partial_t u + u \partial_x u = \nu \partial_x ^2 u, x\in\mathbb{R}
\partial_t \bar{u} + \bar{u} \partial_x \bar{u} = \nu \partial_x ^2 \bar{u} - \frac{1}{2}\overline{u'^2}
u = \bar{u} + u'
,\,\nu = 0.01

Closure Model

Deterministic

Shocks

Energy Cascade

\eta := \text{NN}_\theta(u)

Eddy Viscosity Model

Vanilla Neural Network

\partial_t\nu_T + \alpha_\theta(\bar{u}\partial_x \nu_T) = \beta_(\theta\nu\partial_{xx}\nu_T)
\eta := \partial_x(\nu_T\partial_x\bar{u})

Adding physical priors improves performance

\eta := \partial_x(\nu_T\partial_x\bar{u})
\eta := \text{NN}_\theta(u)

Eddy Viscosity Model

\partial_t\nu_T + \alpha_\theta(\bar{u}\partial_x \nu_T) = \beta_(\theta\nu\partial_{xx}\nu_T)
\partial_t\eta_\theta + \bar{u}\partial_x \eta_\theta = \eta\partial_{xx}\eta_\theta
\eta := \text{NN}_\theta(u)

Transport Equation Model

\eta := \partial_x(\nu_T\partial_x\bar{u})

Eddy Viscosity with Transport

\eta(x,0) = \gamma_\theta(u_0)
\nu_T(x,0) = \gamma_\theta(u_0)

Vanilla Neural Network

\partial_t \bar{u} + \bar{u} \partial_x \bar{u} = \nu \partial_x ^2 \bar{u} + \eta

Conclusions

  • Neural networks are able to increase the accuracy of numerical solutions to PDEs on coarse grid
  • using neural networks in the form of neural ODEs yields better results than using neural networks to advance the solution by a specific amount in time;
  • using neural networks in the form of a closure term, i.e. correcting for the ‘standard’ discretised ODE, yields better results than predicting the entire time series using only the neural network.

 

SciMLOperators.jl
  • Common operator interface for SciML
    • Lazily implemented operator algebra
    • Fully cached so allocations are done upfront and reused
    • Matrix-free operators
    • Fast tensor products (for spectral elements)
    • Can wrap user-provided functions
    • GPU compatibility
  • Easy switch between fast in-place, AD-compatible out-of-place operations

Future Work

  • Continue evolving PDEInterfaces
    • wrap PDE solvers - Trixi.jl, Gridap.jl, ClimaCore.jl
    • meshed domains,
  • OrdinaryDiffEq automatic jacobians - even faster ODE integration, AD
  • Symbolic frontend via ModelingToolkit.jl
  • Create preconditioner library

I see \(\texttt{CalculustJL}\) as the overarching outcome of my PhD. Plan to have future work compatible with and extend the interface

Appendix

Components of PDE Problems*

\mathcal{L}(u) = 0,\, (\vec{x},t) \in \Omega,\times (0,T]
(u + \nabla_{\hat{n}}u)|_{x=\partial\Omega} = g(x, t)
u|_{t=0} = u_0(x)

Governing System

Boundary Data

Time Domain

** stationary

(\nabla,\, \int\cdot\,\mathrm{d}x)
\Omega \subset \mathbb{R}^d

* continuum scale

Space Domain

  • Discretize geometry (meshing)
  • Fill elements with function space
  • Solving
    • Collocation: strong form satisfied at some points
    • Projection (Galerkin): weak form enforced over domain. need notion of inner product
  • discretize governing equation
    • fast gradient computation
    • fast integration
  • project BC data to function space and evolve (solve for) coeffs
  • impose BC
    • either apply restriction matrices (need nodal basis indices)
    • or select function space that satisfies BC (eg Fourier)

Numerically solving PDE Problems

Setting up, and solving PDE problems requires specialized handling.

- Setting up is hard because you have to read geometry from CAD/ scans.

- Mesh has to accurately represent geom

- Obtaining physically correct initial and boundary conditions from data

- Solving is hard because we end up with large linear systems

- Critically, quality of mesh dictates time to solution. Mesh iteration is time consuming. Geom respecting meshes are v large

- Asserting BC is nontrivial

- Massive systems of equations. Large compute requirements

  • Deep Neural Networks
    • Composition of linear + nonlinear operations
    • Representation power of neural networks is in their structure
    • Good for high dimensional problems (empirical evidence)
    • Automatic Differentiation - Modify parameters past nonlinearities
  • Finite Elements
    • Special purpose architectures for solving PDEs in complex geometries
    • designed for low dimension problems; enforced sparsity

Autodiff: compiler magic that gives bwd pass based on chain rule. Lets you modify parameters past nonlinearities

Differentiable Physics

Tackle deeply nested, high dim problems in computational workflows with DiffPhys