implement an IMEXFluxTimeDiscretization for FluxBoundaryCondition

[simone-silvestri]

this PR introduces a way to implicitly handle (linear) boundary conditions

"""
    struct ImplicitExplicitFlux{E, C}

Condition for a vertical `Flux` boundary condition whose value is affine in the
boundary-cell field value `φ_b`,

J(φ_b) = Fₑ + λ φ_b ,

where `Fₑ` is the `explicit_flux` and `λ` is the `coefficient`. The explicit part `Fₑ` is integrated through the tendency
like an ordinary flux boundary condition, while the linear part `λ φ_b` is integrated implicitly by the vertical tridiagonal
solver. This removes the `Δz`-dependent CFL limit that an explicit flux imposes, and is unconditionally stable for dissipative
fluxes (`λ φ_b` a sink), such as drag and linear restoring.
"""
struct ImplicitExplicitFlux{E, C}
    explicit_flux :: E
    coefficient   :: C
end

At the beginning I thought we could implement also non-linear boundary conditions but I do think that at the moment we can do a lot of things with a linear implicit BC (a non linear would require probably AD to work properly). But happy to see what people think about it.

A step forward to solve #5630

[glwagner]

Can you show the user interface that you envision for this?

I propose something like

flux_bc = FluxBoundaryCondition(flux, time_discretization=ImplicitExplicitTimeDiscretization())

[simone-silvestri]

at the moment I had

flux_bc = ImplicitExplicitFluxBoundaryCondition(explicit_flux, implicit_coefficient; kwargs...)

We could go in that direction though. I would keep this kwarg only for the FluxBoundaryCondition though so we do not complicate excessively the PR.

[glwagner]

at the moment I had

flux_bc = ImplicitExplicitFluxBoundaryCondition(explicit_flux, implicit_coefficient; kwargs...)

We could go in that direction though. I would keep this kwarg only for the FluxBoundaryCondition though so we do not complicate excessively the PR.

You can add time_discretization to the Flux classification.

[simone-silvestri]

Added. To avoid passing different, conflicting kwargs to the FluxBoundaryCondition what do you think about

  FluxBoundaryCondition(Q)                           # explicit flux, constructs Flux{ExplicitTimeDiscretization}
  FluxBoundaryCondition(Fₑ; implicit_coefficient=λ)  # IMEX flux, constructs Flux{ImplicitExplicitTimeDiscretization}

[glwagner]

Added. To avoid passing different, conflicting kwargs to the FluxBoundaryCondition what do you think about

  FluxBoundaryCondition(Q)                           # explicit flux, constructs Flux{ExplicitTimeDiscretization}
  FluxBoundaryCondition(Fₑ; implicit_coefficient=λ)  # IMEX flux, constructs Flux{ImplicitExplicitTimeDiscretization}

What conflicting kwargs?

[simone-silvestri]

The IMEX flux requires two arguments, an "explicit" flux and an "implicit" coefficient. We need to pass both information to the FluxBoundaryCondition, however, if we pass both we already know that the time discretization is an IMEX so we can avoid that kwarg. We can rely on having just

FluxBoundaryCondition(Fₑ; implicit_coefficient=λ)

to build an IMEX bc. Other options are

FluxBoundaryCondition(Fₑ; implicit_coefficient=λ, time_discretization = ImplicitExplicitTimeDiscretization())

which is redundant or

FluxBoundaryCondition(Fₑ; time_discretization = ImplicitExplicitTimeDiscretization(implicit_coefficient=λ))

which I do not particularly like.
Otherwise we can think of two positional arguments

FluxBoundaryCondition(Fₑ, λ; time_discretization = ImplicitExplicitTimeDiscretization())

which might also be a bit redundant

[glwagner]

I propose this syntax:

FluxBoundaryCondition(Fₑ; time_discretization = ImplicitExplicitTimeDiscretization(λ))

ie omit the kwarg on ImplicitExplicitTimeDiscretization. I'm also ok with IMEXTimeDiscretization.

[simone-silvestri]

I am not so sure, it looks like the λ is a numerical parameter that has to do with the discretization rather than the actual physical boundary condition coefficient (or function).

[glwagner]

I am not so sure, it looks like the λ is a numerical parameter that has to do with the discretization rather than the actual physical boundary condition coefficient (or function).

There is only one argument/property to ImplicitExplicitTimeDiscretization right?

[simone-silvestri]

yes, what I mean is that the complete flux is Fₑ + λ ϕ, so λ is a physical quantity needed to compute the flux and, to me, having it separated from Fₑ and included in the time discretization makes it look like a numerical parameter that has to do with time stepping (like a cfl for example) rather than a physical quantity.

[glwagner]

yes, what I mean is that the complete flux is Fₑ + λ ϕ, so λ is a physical quantity needed to compute the flux and, to me, having it separated from Fₑ and included in the time discretization makes it look like a numerical parameter that has to do with time stepping (like a cfl for example) rather than a physical quantity.

what is ϕ ?

[glwagner]

I guess I was confused. I thought lambda would have been the forward weighting in the IMEX discretization. Since it's IMEX is there a second parameter involved?

[glwagner]

I don't know what implicit_coefficient means. It isn't any more clear that ImplicitExplicitTimeDiscretization(coefficient).

[simone-silvestri]

what is ϕ ?

it's the prognostic variable the boundary condition is applied to.

I don't know what implicit_coefficient means. It isn't any more clear that ImplicitExplicitTimeDiscretization(coefficient).

Ok, what about just passing a tuple of values to a flux as positional arguments? The problem to me is that ImplicitExplicitTimeDiscretization(coefficient) is just very misleading, not only in terms of name but in terms of functionality, for example in

FluxBoundaryCondition(Fₑ; discrete_form = true, parameters = p, time_discretization = ImplicitExplicitTimeDiscretization(λ))

both the parameters and discrete_form would need to apply also to λ which is a little bit counterintuitive since it enters ImplicitExplicitTimeDiscretization.
Another way to do it would be

FluxBoundaryCondition(Fₑ; discrete_form = true, parameters = p, time_discretization = ImplicitExplicitTimeDiscretization(λ; discrete_form = true, parameters = p))

which I am not super sure about. Looks a little repetitive.
Basically there is symmetry between Fₑ and λ, they are connected, they are two parts of the same object, so one entering the flux constructor via a different path (ImplicitExplicitTimeDiscretization) sounds weird to me.

[simone-silvestri]

We could have a compromise:

flux_bc = ImplicitExplicitFluxBoundaryCondition(explicit_flux, implicit_coefficient; kwargs...)

with

ImplicitExplicitFluxBoundaryCondition(explicit_flux, implicit_coefficient; kwargs...) = FluxBoundaryCondition(explicit_flux; time_discretization = ImplicitExplicitTimeDiscretization(implicit_coefficient), kwargs...)

[glwagner]

what is ϕ ?

it's the prognostic variable the boundary condition is applied to.

ok, then I suggest:

FluxBoundaryCondition(Fₑ; time_discretization = IMEXFluxTimeDiscretization(λ))

[simone-silvestri]

Ok done!

[simone-silvestri]

Another example of the capabilities of this PR. This comes from a request from @yuchenma23 that needs to perform DNS simulations of double diffusion and he cannot achieve stability because, needing a very small grid size near the boundary to capture boundary layer behavior, the time step becomes miniscule even when using an implicit solver as the cfl limitation comes from the boundary fluxes that are treated always explicitly. This snippet shows how it is possible to treat even value boundary conditions implicitly with a bit of a reformulation and recast in flux boundary conditions.

const Lz = 0.1meters
const Nz = 512
const Δt = 100seconds
const stop_time = 30minutes
const κᵀ = 1e-6
const Sf = 3.5

stretched_faces(k) = Lz - Lz * tanh(Sf * (Nz+1-k) / (Nz)) / tanh(Sf)
Tᵢ = 0.001 * rand(Nz)

grid = RectilinearGrid(CPU();
                       size = Nz,
                       z = stretched_faces,
                       topology = (Flat,Flat, Bounded))

function diffusion_dns(Tbcs)
    molecular_closure = ScalarDiffusivity(VerticallyImplicitTimeDiscretization(), κ=κᵀ)
    model = NonhydrostaticModel(grid; timestepper=:RungeKutta3, tracers = (:T), closure = molecular_closure, boundary_conditions = (; T = Tbcs))
    set!(model, T=Tᵢ)
    simulation = Simulation(model; Δt, stop_time)

    wall_clock = Ref(time_ns())
    function print_progress(sim)
        progress = 100 * time(sim) / sim.stop_time
        elapsed  = (time_ns() - wall_clock[]) / 1e9
        @printf("[%05.2f%%], i=%d, t=%s, dt =%s, walltime=%s max|u,v,w,T|=(%.2e,%.2e,%.2e,%.2e)\n",
        progress, iteration(sim), prettytime(sim), prettytime(sim.Δt),prettytime(elapsed),
        maximum(abs, model.velocities.u),
        maximum(abs, model.velocities.v),
        maximum(abs, model.velocities.w),
        maximum(abs, model.tracers.T))
        wall_clock[] = time_ns()
    end

    add_callback!(simulation, print_progress, IterationInterval(10))
    Tframes = Vector{Vector{Float64}}()
    save_profile(sim) = push!(Tframes, Array(interior(sim.model.tracers.T, 1, 1, :)))
    add_callback!(simulation, save_profile, TimeInterval(stop_time / 200))
    @info "Running the simulation ..."
    run!(simulation)

    return Tframes
end

explicit_bcs = FieldBoundaryConditions(top = ValueBoundaryCondition(0.0))

@inline function implicit_coefficient(i, j, grid, clock, fields, κᵀ)
    Nz = size(grid, 3)
    Δz = Oceananigans.Operators.Δzᶜᶜᶠ(i, j, Nz+1, grid)
    return 2 * κᵀ / Δz
end

implicit_flux = FluxBoundaryCondition(0; time_discretization=IMEXFluxTimeDiscretization(implicit_coefficient), discrete_form=true, parameters=κᵀ)
implicit_bcs  = FieldBoundaryConditions(top = implicit_flux)

Texplicit = diffusion_dns(explicit_bcs)
Timplicit = diffusion_dns(implicit_bcs)

z = znodes(grid, Center(), Center(), Center())
n = Observable(1)

fig = Figure(size=(480, 560))
ax = Axis(fig[1, 1]; xlabel="T", ylabel="z (m)", title=@lift("frame $($n)"))
lines!(ax, @lift(Texplicit[$n]), z; color=:firebrick,  linewidth=4, label="explicit (Value)")
lines!(ax, @lift(Timplicit[$n]), z; color=:dodgerblue, linewidth=2, linestyle=:dash, label="implicit (IMEX)")
xlims!(ax, 0, 0.005)
axislegend(ax; position=:rb)

record(fig, "implicit_vs_explicit.mp4", 1:min(length(Texplicit), length(Timplicit)); framerate=20) do i
    n[] = i
end

implicit_vs_explicit.mp4

Of course it is not super optimal to have to recast a value boundary condition into a flux boundary condition, so I propose that after we merge this PR that implements implicit flux bcs and the plumbing of bcs in the implicit solver we can make the small changes required to natively support also value boundary conditions without needing to recast them.