Stabilizing forcing function arguments and boundary condition function arguments

[glwagner]

We currently allow users to implement custom forcing functions with the signature:

F(i, j, k, grid, time, U, C, parameters)

where U is a named tuple of velocity fields, C is a named tuple of tracer fields, and parameters is the object passed to IncompressibleModel via the parameters keyword argument.

Boundary condition functions have the signature:

condition(i, j, grid, time, iteration, U, C, parameters)

To stabilize the API, we may want to get rid of parameters, include clock rather than time or iteration, and add a named tuple container called state that holds U, C, and any other state variables that we want to add (now or in the future) as accessible to forcing functions or boundary condition functions.

I'd propose that state be defined something like

state = (velocities=U, tracers=C, diffusivities=K, pressures=pressures, tendencies=G)

If we form state within the time-stepping loop, it could also simplify the function signatures for time-stepping kernels.

Related is #565.

[ali-ramadhan]

A state variable would be cool, especially if it shortens function signatures!

And yeah if state is not a model property but just a common function signature (presumably for internal/backend use) then there won't be much confusion.

I guess function signatures will become

F(i, j, k, grid, time, state, parameters)

which is quite nice.

This change will break a lot of scripts that define custom forcing functions and boundary conditions so should probably mention this in release notes.

[glwagner]

I was thinking even simpler:

F(i, j, k, grid, clock, state)

Any 'parameters' that are needed for forcing functions can be introduced by defining a callable object that carries around its parameters. We can introduce something simple like ParameterizedForcing for users, too, perhaps.

(Now that I think of it, we should probably separate SimpleForcing (forcing functions of x,y,z,t with no parameters) from SimpleParameterizedForcing (simple forcing functions with parameters) too...)

[ali-ramadhan]

Hmmm, I'm not 100% sure how the callable object would work with ParameterizedForcing but as long as it's easy/intuitive to use in scripts then that'll be great.

Do you have some pseudocode about what the setup would look like for users?

[glwagner]

Yes:

cool_forcing_function(i, j, k, grid, clock, state, parameters) = # something that uses the parameters argument
cool_forcing = ParameterizedForcing(cool_forcing_function, parameters=(a=1, b=2))

Pretty much the same way that it's used for SimpleForcing, except here we can make parameters a required keyword argument.

[glwagner]

The definition of ParameterizedForcing is

struct ParameterizedForcing{F, P}
    func :: F
    parameters :: P
end

(F::ParameterizedForcing)(i, j, k, grid, clock, state) = F.func(i, j, k, grid, clock, state, F.parameters)

easy peasy.

[ali-ramadhan]

Ah looks awesome then!

[glwagner]

Simple forcing is defined similarly here:

https://github.com/climate-machine/Oceananigans.jl/blob/master/src/Forcing/simple_forcing.jl

here's the docstring:

"""
    SimpleForcing([location=(Cell, Cell, Cell),] forcing; parameters=nothing)
Construct forcing for a field at `location` using `forcing::Function`, and optionally
with `parameters`. If `parameters=nothing`, `forcing` must have the signature
    `forcing(x, y, z, t)`;
otherwise it must have the signature
    `forcing(x, y, z, t, parameters)`.
Examples
========
```julia
julia> const a = 2.1
julia> fun_forcing(x, y, z, t) = a * exp(z) * cos(t)
julia> u_forcing = SimpleForcing(fun_forcing)
julia> parameterized_forcing(x, y, z, t, p) = p.μ * exp(z/p.λ) * cos(p.ω*t)
julia> v_forcing = SimpleForcing(parameterized_forcing, parameters=(μ=42, λ=0.1, ω=π))

"""

[glwagner]

Suggestions are definitely welcome for streamlining everything!

Note: in dedalus, this problem is handled by allowing users to accumulate parameters into something like model.parameters (in dedalus, this is essentially a dict whose keys are accessible by their names to any function/equation defined with a string expression).

Having a "global-like" variable (like model.parameters) is pretty convenient, but also leads to more complicated code. The more democratic callable object strategy is a bit simpler and more modular, but may involve a slightly more complicated API (though I'm not 100% sure about that --- users must decide!)

[ali-ramadhan]

I am strongly in favor of more localized parameters as it makes the code more modular and we can stop juggling around model.parameters.

If users do want a global-like parameters object, they can define one big named tuple and pass it to all the ParameterizedForcingFunctions and ParameterizedBoundaryFunctions.

[glwagner]

I've realized that it doesn't make sense to include either the tendencies or pressures in state, as neither of these is actually known at the time that boundary condition functions or forcing functions are called.

so we just have state = (velocities=U, tracers=C, diffusivities=K).

This also means that we need a special solution for implementing predictor velocity boundary conditions.