Time splitting interface ideas

[charleskawczynski]

Description

@thomasgibson and I were talking about #1634, and we started discussing time stepping. We prototyped some ideas and I just wanted to just record here. Not sure if this idea will work with asynchronous LHS/RHS evaluation, though.

module ClimateMachine

  module BalanceLaws
    export BalanceLaw, Model, terms
    abstract type BalanceLaw end
    struct Term1 end; export Term1
    struct Term2 end; export Term2
    struct Term3 end; export Term3 # stiff
    struct Model{SM}
      submodel::SM
    end
    # Specification proposed in #1634
    terms(::Model) = (Term1(), Term2(), Term3())
  end

  module TimeSteppers
    import ..BalanceLaws
    export SplitModel, set_explicit, set_implicit
    struct SplitModel{M,TE,TI}
      model::M
      terms_exp::TE
      terms_imp::TI
    end
    explicit_terms(model, terms_imp) =
        filter(x->all(x ≠ y for y in terms_imp), BalanceLaws.terms(model))
    function SplitModel(model::M; terms_imp=()) where {M}
      terms_exp = explicit_terms(model, terms_imp)
      args = (model,terms_exp,terms_imp)
      return SplitModel{typeof.(args)...}(args...)
    end
    function set_explicit(split_model)
      eval(:(BalanceLaws.terms(::typeof($(split_model.model))) = $(split_model.terms_exp)))
      return nothing
    end
    function set_implicit(split_model)
      eval(:(BalanceLaws.terms(::typeof($(split_model.model))) = $(split_model.terms_imp)))
      return nothing
    end
    BalanceLaws.terms(sm::SplitModel) = BalanceLaws.terms(sm.model)
  end
end

using .ClimateMachine.BalanceLaws
using .ClimateMachine.TimeSteppers
smodel = SplitModel(Model(1); terms_imp=(Term3(),));

set_implicit(smodel)
BalanceLaws.terms(smodel) # compute RHS with implicit terms
set_explicit(smodel)
BalanceLaws.terms(smodel) # compute RHS with explicit terms

[charleskawczynski]

A bit more brainstorming/scratchpad session with @thomasgibson:

solver = MultiRateSolver(
    iv = implicit(::Vertical) = (
        Rate{3}(Gravity),
        Rate{3}(Advection),
    ),
    ih = implicit(::Horizontal) = (
        Rate{3}(Gravity),
        Rate{3}(Advection),
    ),
    ev = explicit(::Vertical) = (
        Rate{3}(Gravity),
        Rate{3}(Advection),
    ),
    eh = explicit(::Horizontal) = (
        Rate{3}(Gravity),
        Rate{3}(Radiation),
    ),
    return MultiRateSolver{}(iv, ih, ev, eh)
)

terms = solver.iv(direction)

solver = MultiRateSolver(
  # Fully explicit solver
  Rate{1}(
    explicit(::EveryDirection) = (Radiation),
  ),
  # Fully explicit solver
  Rate{2}(
    explicit(::HorizontalDirection) = (Diffusion),
  ),
  # implicit-explicit method (ARK IMEX)
  Rate{3}(
    # HEVI
    explicit(::HorizontalDirection) = (Advection, Pressure),
    implicit(::VerticalDirection) = (Advection, Pressure, Diffusion),
  )
)

[charleskawczynski]

I was thinking more about this recently. Here's what I think is a full and clear picture of the most general explicit/explicit term selection I can think of:

G = ∇Y
∂_t Y_i + (c_h_F₁_implicit*∇_h•F₁(Y))_i + c_h_F₁_explicit*∇_h•F₁(Y))_i) +
          (c_v_F₁_implicit*∇_v•F₁(Y))_i + c_v_F₁_explicit*∇_v•F₁(Y))_i) +
          (c_h_F₂_implicit*∇_h•F₂(Y,G)))_i + c_h_F₂_explicit*∇_h•F₂(Y,G)))_i) +
          (c_v_F₂_implicit*∇_v•F₂(Y,G)))_i + c_v_F₂_explicit*∇_v•F₂(Y,G)))_i) =
          c_h_S_implicit*(S(Y,G))_i + c_h_S_explicit*(S(Y,G))_i
          c_v_S_implicit*(S(Y,G))_i + c_v_S_explicit*(S(Y,G))_i

Where we must define the coefficients:

c_h_F₁_implicit, c_h_F₁_explicit
c_v_F₁_implicit, c_v_F₁_explicit
c_h_F₂_implicit, c_h_F₂_explicit
c_v_F₂_implicit, c_v_F₂_explicit
c_h_S_implicit, c_h_S_explicit
c_v_S_implicit, c_v_S_explicit

We can think about grouping these 12 coefficiencts in different ways:

Group coefficients by tendency types

c_F₁ = (c_h_F₁_implicit, c_h_F₁_explicit, c_v_F₁_implicit, c_v_F₁_explicit) c_F₂ = (c_h_F₂_implicit, c_h_F₂_explicit, c_v_F₂_implicit, c_v_F₂_explicit) c_S = (c_h_S_implicit, c_h_S_explicit, c_v_S_implicit, c_v_S_explicit)

Group coefficients by implicit/explicit

c_impicit = (c_h_F₁_implicit, c_h_F₂_implicit, c_h_S_implicit, c_v_F₁_implicit, c_v_F₂_implicit, c_v_S_implicit) c_expicit = (c_h_F₁_explicit, c_h_F₂_explicit, c_h_S_explicit, c_v_F₁_explicit, c_v_F₂_explicit, c_v_S_explicit)

Group coefficients by horizontal/vertical

c_h = (c_h_F₁_implicit, c_h_F₂_implicit, c_h_S_implicit, c_h_F₁_explicit, c_h_F₂_explicit, c_h_S_explicit) c_v = (c_v_F₁_implicit, c_v_F₂_implicit, c_v_S_implicit, c_v_F₁_explicit, c_v_F₂_explicit, c_v_S_explicit)

Group coefficients by explicit/implicit and horizontal/vertical

c_h_implicit = (c_h_F₁_implicit, c_h_F₂_implicit, c_h_S_implicit) c_h_explicit = (c_h_F₁_explicit, c_h_F₂_explicit, c_h_S_explicit) c_v_implicit = (c_v_F₁_implicit, c_v_F₂_implicit, c_v_S_implicit) c_v_explicit = (c_v_F₁_explicit, c_v_F₂_explicit, c_v_S_explicit)

Defining coefficients

I think the first question we need to answer is how we can reduce the number of coefficients that one must prescribe to fully define the explicit/explicit term selection. Some rules and simplifications that come to mind include:

• General rule: c_impicit .+ c_expicit = 1
• General rule: c_h = (0 or 1)
• General rule: c_v = (0 or 1)
• using SingleStackConfiguration: c_h = 0
• using IMEXSolverType: ?
• using HEVI: c_h_explicit,c_h_implicit,c_v_explicit,c_v_implicit = (1,0,0,1) ?

More granularity

If we need more granularity (e.g., want certain terms in F₁ to be implicit), we can reinterpret each coefficient as a tuple that acts on each individual term inside F₁, F₂, and S.

Implementation notes

We can define a struct that contains all of this information and use this in the back-end to control decision making on whether certain kernels need to be launched or not.

Additional notes

Based on the docs for discrete_splitting, it seems that there may be an orthogonal question of PDE level or discretized level splitting, but I think that this could be handled separately.

cc @thomasgibson @simonbyrne