Neural closure models

NeuralClorusure

These features are experimental, and require cloning IncompressibleNavierStokes from GitHub:

git clone https://github.com/agdestein/IncompressibleNavierStokes.jl
cd IncompressibleNavierStokes/lib/NeuralClosure

Large eddy simulation, a closure model is required. With IncompressibleNavierStokes, a neural closure model can be trained on filtered DNS data. The discrete DNS equations are given by

$$\begin{array}{rl}Mu& =0,\\ \frac{\mathrm{d}u}{\mathrm{d}t}& =F(u)-Gp.\end{array}$$

Applying a spatial filter $\mathrm{\Phi }$, the extracted large scale components $\overline{u}=\mathrm{\Phi }u$ are governed by the equation

$$\begin{array}{rl}M\overline{u}& =0,\\ \frac{\mathrm{d}\overline{u}}{\mathrm{d}t}& =F(\overline{u})+c-G\overline{p},\end{array}$$

where the discretizations $M$, $F$, and $G$ are adapted to the size of their inputs and $c=\overline{F(u)}-F(\overline{u})$ is a commutator error. We here assumed that $M$ and $\mathrm{\Phi }$ commute, which is the case for face averaging filters. Replacing $c$ with a parameterized closure model $m(\overline{u},\theta )\approx c$ gives the LES equations for the approximate large scale velocity $\overline{v}\approx \overline{u}$

$$\begin{array}{rl}M\overline{v}& =0,\\ \frac{\mathrm{d}\overline{v}}{\mathrm{d}t}& =F(\overline{v})+m(\overline{v},\theta )-G\overline{q}.\end{array}$$

NeuralClosure module

IncompressibleNavierStokes provides a NeuralClosure module.

# NeuralClosure.NeuralClosure — Module.

Neural closure modelling tools.

source

# NeuralClosure.collocate — Method.

collocate(u) -> Any

Interpolate velocity components to volume centers.

source

# NeuralClosure.create_closure — Method.

create_closure(layers...; rng)

Create neural closure model from layers.

source

# NeuralClosure.create_tensorclosure — Method.

create_tensorclosure(layers...; setup, rng)

Create tensor basis closure.

source

# NeuralClosure.decollocate — Method.

decollocate(u) -> Any

Interpolate closure force from volume centers to volume faces.

source

# NeuralClosure.wrappedclosure — Method.

wrappedclosure(
    m,
    setup
) -> NeuralClosure.var"#neuralclosure#1"

Wrap closure model and parameters so that it can be used in the solver.

source

Filters

The following filters are available:

# NeuralClosure.AbstractFilter — Type.

abstract type AbstractFilter

Discrete DNS filter.

Subtypes ConcreteFilter should implement the in-place method:

(::ConcreteFilter)(v, u, setup_les, compression)

which filters the DNS field u and put result in LES field v. Then the out-of place method:

(::ConcreteFilter)(u, setup_les, compression)

automatically becomes available.

Fields

source

# NeuralClosure.FaceAverage — Type.

struct FaceAverage <: NeuralClosure.AbstractFilter

Average fine grid velocity field over coarse volume face.

Fields

source

# NeuralClosure.VolumeAverage — Type.

struct VolumeAverage <: NeuralClosure.AbstractFilter

Average fine grid velocity field over coarse volume.

Fields

source

# NeuralClosure.reconstruct! — Method.

reconstruct!(u, v, setup_dns, setup_les, comp) -> Any

Reconstruct DNS velocity u from LES velocity v.

source

# NeuralClosure.reconstruct — Method.

reconstruct(v, setup_dns, setup_les, comp) -> Any

Reconstruct DNS velocity field. See also reconstruct!.

source

Training

To improve the model parameters, we exploit exact filtered DNS data $\overline{u}$ and exact commutator errors $c$ obtained through DNS. The model is trained by minimizing the a priori loss function

$$L^{\text{prior}}(\theta )=\parallel m(\overline{u},\theta )-c{\parallel }^2,$$

or the a posteriori loss function

$$L^{\text{post}}(\theta )=\parallel {\overline{v}}_{\theta }-\overline{u}{\parallel }^2,$$

where ${\overline{v}}_{\theta }$ is the solution to the LES equation for the given parameters $\theta$. The prior loss is easy to evaluate and easy to differentiate, as it does not involve solving the ODE. However, minimizing $L^{\text{prior}}$ does not take into account the effect of the prediction error on the LES solution error. The posterior loss does, but has a longer computational chain involving solving the LES ODE.

# NeuralClosure.create_callback — Method.

create_callback(
    err;
    θ,
    callbackstate,
    displayref,
    display_each_iteration,
    displayfig,
    filename
)

Create convergence plot for relative error between f(x, θ) and y. At each callback, plot is updated and current error is printed.

If state is nonempty, it also plots previous convergence.

If not using interactive GLMakie window, set display_each_iteration to true.

source

# NeuralClosure.create_dataloader_post — Method.

create_dataloader_post(trajectories; nunroll, device, rng)

Create trajectory dataloader.

source

# NeuralClosure.create_dataloader_prior — Method.

create_dataloader_prior(data; batchsize, device, rng)

Create dataloader that uses a batch of batchsize random samples from data at each evaluation. The batch is moved to device.

source

# NeuralClosure.create_loss_post — Method.

create_loss_post(
;
    setup,
    method,
    psolver,
    closure,
    nupdate,
    projectorder
)

Create a-posteriori loss function.

source

# NeuralClosure.create_loss_prior — Method.

create_loss_prior(loss, f) -> NeuralClosure.var"#54#55"

Wrap loss function loss(batch, θ).

The function loss should take inputs like loss(f, x, y, θ).

source

# NeuralClosure.create_relerr_post — Method.

create_relerr_post(
;
    data,
    setup,
    method,
    psolver,
    closure_model,
    nupdate,
    projectorder
)

Create a-posteriori relative error.

source

# NeuralClosure.create_relerr_prior — Method.

create_relerr_prior(f, x, y) -> NeuralClosure.var"#56#57"

Create a-priori error.

source

# NeuralClosure.create_relerr_symmetry_post — Method.

create_relerr_symmetry_post(
;
    u,
    setup,
    method,
    psolver,
    Δt,
    nstep,
    g
)

Create a-posteriori symmetry error.

source

# NeuralClosure.create_relerr_symmetry_prior — Method.

create_relerr_symmetry_prior(; u, setup, g)

Create a-priori equivariance error.

source

# NeuralClosure.mean_squared_error — Method.

mean_squared_error(f, x, y, θ; normalize, λ) -> Any

Compute MSE between f(x, θ) and y.

The MSE is further divided by normalize(y).

source

# NeuralClosure.train — Method.

train(
    dataloader,
    loss,
    optstate,
    θ;
    niter,
    ncallback,
    callback,
    callbackstate
) -> NamedTuple{(:optstate, :θ, :callbackstate), <:Tuple{Any, Any, Nothing}}

Update parameters θ to minimize loss(dataloader(), θ) using the optimiser opt for niter iterations.

Return the a new named tuple (; opt, θ, callbackstate) with updated state and parameters.

source

Neural architectures

We provide neural architectures: A convolutional neural network (CNN), group convolutional neural networks (G-CNN) and a Fourier neural operator (FNO).

# NeuralClosure.cnn — Method.

cnn(
;
    setup,
    radii,
    channels,
    activations,
    use_bias,
    channel_augmenter,
    rng
)

Create CNN closure model. Return a tuple (closure, θ) where θ are the initial parameters and closure(u, θ) predicts the commutator error.

source

# NeuralClosure.GroupConv2D — Type.

struct GroupConv2D{C} <: LuxCore.AbstractLuxLayer

Group-equivariant convolutional layer – with respect to the p4 group. The layer is equivariant to rotations and translations of the input vector field.

The kwargs are passed to the Conv layer.

The layer has three variants:

• If islifting then it lifts a vector input (u1, u2) into a rotation-state vector (v1, v2, v3, v4).

• If isprojecting, it projects a rotation-state vector (u1, u2, u3, v4) into a vector (v1, v2).

• Otherwise, it cyclically transforms the rotation-state vector (u1, u2, u3, u4) into a new rotation-state vector (v1, v2, v3, v4).

Fields

• islifting

• isprojecting

• cin

• cout

• conv

source

# NeuralClosure.gcnn — Method.

gcnn(; setup, radii, channels, activations, use_bias, rng)

Create CNN closure model. Return a tuple (closure, θ) where θ are the initial parameters and closure(u, θ) predicts the commutator error.

source

# NeuralClosure.rot2 — Function.

rot2(u, r)

Rotate the field u by 90 degrees counter-clockwise r - 1 times.

source

# NeuralClosure.rot2stag — Method.

rot2stag(u, g) -> Tuple{Any, Any}

Rotate staggered grid velocity field. See also rot2.

source

# NeuralClosure.FourierLayer — Type.

struct FourierLayer{D, A, F} <: LuxCore.AbstractLuxLayer

Fourier layer operating on uniformly discretized functions.

Some important sizes:

• dimension: Spatial dimension, e.g. Dimension(2) or Dimension(3).

• (nx..., cin, nsample): Input size

• (nx..., cout, nsample): Output size

• nx = fill(n, dimension()): Number of points in each spatial dimension

• n ≥ kmax: Same number of points in each spatial dimension, must be larger than cut-off wavenumber

• kmax: Cut-off wavenumber

• nsample: Number of input samples (treated independently)

Fields

• dimension

• kmax

• cin

• cout

• σ

• init_weight

source

# NeuralClosure.fno — Method.

fno(; setup, kmax, c, σ, ψ, rng, kwargs...)

Create FNO closure model. Return a tuple (closure, θ) where θ are the initial parameters and closure(V, θ) predicts the commutator error.

source

Data generation

# NeuralClosure.create_io_arrays — Method.

create_io_arrays(data, setups) -> Any

Create $(\overline{u},c)$ pairs for a-priori training.

source

# NeuralClosure.create_les_data — Method.

create_les_data(
;
    D,
    Re,
    lims,
    nles,
    ndns,
    filters,
    tburn,
    tsim,
    savefreq,
    Δt,
    create_psolver,
    ArrayType,
    icfunc,
    rng,
    kwargs...
)

Create filtered DNS data.

source

# NeuralClosure.filtersaver — Method.

filtersaver(
    dns,
    les,
    filters,
    compression,
    psolver_dns,
    psolver_les;
    nupdate
) -> NamedTuple{(:initialize, :finalize), <:Tuple{NeuralClosure.var"#100#104"{Int64}, NeuralClosure.var"#103#107"}}

Save filtered DNS data.

source