Allow forcing to be an array.

[ali-ramadhan]

Following up from #73 it would be nice if the forcing can be also be expressed as an array. This might be nice for a couple of applications:

1. A forcing that is constant and very expensive to compute (e.g. lots of calculations, branching statements, or intermediate calculations).
2. Forcing array can be passed around and filled elsewhere, e.g. by a biogeochemical agent-based model.

Pinging: @zhenwu0728

[glwagner]

This can be implemented by implementing a new method for update_source_terms! that dispatches on forcing::Forcing{Tu} where Tu <: AbstractArray. No changes are needed to the code except for this addition. When the user passes arrays to the Forcing constructor, the code will do the right thing as long as this method is implemented.

I would note unless we write a lot of new code, the user will be essentially limited to either providing an array for all fields, or forcing for all fields. To allow the user to implement either forcing functions or forcing arrays requires writing 5! = 120 new methods (right?) which is not desirable. A convenience constructor for Forcing that detects when the user passes arrays for any of the fields will help users avoid accidentally specifying forcing arrays for some fields and functions for others.

However, once we resolve the isbitstype problem (perhaps also using FieldVectors or LabeledArrays for velocities and tracers to retain high performance) we can introduce abstractions for the idea of an equation function for each field individually (which would also be stored in a LabeledArray), and permit more flexible equation specification by the user without inducing a combinatorial explosion of code length.

[ali-ramadhan]

That's a good point you bring up, sounds like this will be easier once #59 is resolved.

Instead of dispatching, I was actually thinking of using a macro here, e.g. @insert_forcing that inserts a zero if the forcing is nothing, a function call if the forcing is a function, and grabs a numbers from an array if it's an array.

Some pseudocode:

# Insert forcing for u-momentum equation.
macro insert_forcing_u(Fu)
    if Fu == nothing
        return 0
    else if typeof(Fu) == Function
        return Meta.parse("Fu(grid, velocities, tracers, i, j, k)")
    else if typeof(Fu) <: AbstractArray
        return Meta.parse("Fu[i, j, k]")
    end
end

Then the time-stepping might look like

Gu[i, j, k] = (-u∇u(...) + ... + @insert_forcing_u)

All the different macros can be defined in an @eval loop or maybe with a single @insert_forcing definition is it knows which forcing is being inserted so it can inject the correct `Fu", "Fv", etc. in the expression.

[glwagner]

What's the motivation for using a macro rather than multiple dispatch? To my eye it looks like your macro is basically performing multiple dispatch with an if statement. But Julia has multiple dispatch built in; we don't have to implement it ourselves.

An argument against macros is that they make the code more obscure. It's harder to figure out what is happening because you have to find the definition of the macro.

Come to think of it, the user can also just define a forcing function that indexes into some constant array. Why is this not a good solution?

[ali-ramadhan]

What's the motivation for using a macro rather than multiple dispatch?

The main motivation being that we don't have to write extra functions that dispatch on the forcing, thus simplifying the time stepping code. As you point out we don't want to write 5! = 120 new functions. The update_source_terms! function is already 52 lines long so I'd rather avoid having to dispatch on this function.

An argument against macros is that they make the code more obscure. It's harder to figure out what is happening because you have to find the definition of the macro.

I think this is context-dependent. The purpose of a macro with a name like @insert_forcing_u or @insert_forcing_term is pretty clear. If we used dispatch then you'd still have to scroll through multiple function definitions.

Come to think of it, the user can also just define a forcing function that indexes into some constant array. Why is this not a good solution?

I think this would work pretty well. I couldn't figure out how to pass in the array to be indexed so that it can fit in the Fu(grid, velocities, tracers, i, j, k) signature and be available for the user to fill. Where in the model should we store the forcing array in this case?

Hmmm, actually we could make the function accept the forcing struct, e.g.

Fu(grid, velocities, tracers, forcing, i, j, k)

but then we'd have to have arrays in the forcing struct, e.g.

struct Forcing{Tu,Tv,Tw,TT,TS,TA<:AbstractArray}
    u::Tu
    v::Tv
    w::Tw
    T::TT
    S::TS
    u_arr::TA
    v_arr::TA
    w_arr::TA
    T_arr::TA
    S_arr::TA
end

and then the forcing function is just

Fu(grid, velocities, tracers, forcing, i, j, k) = forcing.u_arr[i, j, k]

Either we have 5 array types so fields with an actual forcing function get nothing for the array or we have 1 array type TA and set the arrays for forcings with a function to something like Array{Float64}(undef, 0).

Might be a little too ugly but I think yeah we should be able to accommodate array forcings with the current framework somehow (and avoid complicating the time stepping code).

[glwagner]

Where in the model should we store the forcing array in this case?

The user would define an array in a script, declare it const to the compiler, and then write a function that indexes into it as a global:

# define a
forcing(..., i, j, k) = a[i, j, k]

We can also include a constructor for Forcing that allows the user to pass some function that defines a constant array, and set up the same functionality internally.

The main motivation being that we don't have to write extra functions that dispatch on the forcing, thus simplifying the time stepping code. As you point out we don't want to write 5! = 120 new functions. The update_source_terms! function is already 52 lines long so I'd rather avoid having to dispatch on this function.

I think the problem is that our functions are trying to do too much at once. We need smaller functions that perform more atomic operations so we can dispatch on atomic operations. I don't think we need to re-invent multiple dispatch with macros. We just need to refactor the code so we can use multiple dispatch effectively.

[glwagner]

Let me make things concrete. Right now we have:

"Store previous value of the source term and calculate current source term."
function update_source_terms!(::Val{Dev}, fCor, χ, ρ₀, κh, κv, 𝜈h, 𝜈v, Nx, Ny, Nz, Δx, Δy, Δz,
                              u, v, w, T, S, pHY′, Gu, Gv, Gw, GT, GS, Gpu, Gpv, Gpw, GpT, GpS, F) where Dev
 
# ...
                # u-momentum equation
                @inbounds Gu[i, j, k] = (-u∇u(u, v, w, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k)
                                            + fCor*avg_xy(v, Nx, Ny, i, j, k)
                                            - δx_c2f(pHY′, Nx, i, j, k) / (Δx * ρ₀)
                                            + 𝜈∇²u(u, 𝜈h, 𝜈v, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k)
                                            + F.u(u, v, w, T, S, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k))
# ...

A simple change we could make would be instead write

"Calculate the right-hand-side of the u-momentum equation at I, j, k."
u_eqn(args..., F::Function i, j, k) = stuff +  F(u, v, w, T, S, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k)
u_eqn(args..., F::AbstractArray i, j, k) = stuff +  F[i, j, k]
u_eqn(args..., F::Nothing, i, j, k) = stuff

"Store previous value of the source term and calculate current source term."
function update_source_terms!(::Val{Dev}, fCor, χ, ρ₀, κh, κv, 𝜈h, 𝜈v, Nx, Ny, Nz, Δx, Δy, Δz,
                              u, v, w, T, S, pHY′, Gu, Gv, Gw, GT, GS, Gpu, Gpv, Gpw, GpT, GpS, F) where Dev
 
# ...
                # u-momentum equation
                @inbounds Gu[i, j, k] = u_eqn(args..., F.u, i, j, k)
# ...

We could write even less code if we created an abstraction for the right hand side, something like

struct Equation{TF}
   G::Function
   F::TF
end

(eq::Equation{TF})(args..., i, j, k) where TF <: Function = eq.G(args..., i, j, k) + eq.F(args..., i, j, k)
(eq::Equation{TF})(args..., i, j, k) where TF <: AbstractArray = eq.G(args..., i, j, k) + eq.F(args..., i, j, k)
(eq::Equation{TF})(args..., i, j, k) where TF <: Nothing = eq.G(args..., i, j, k) 

u_eqn = Equation(Gu, Fu)

...
@inbounds Gu[i, j, k] = u_eqn(args..., i, j, k)

We can then load all the equations we have into a FieldVector or LabeledArray to make things even better and do something like

@loop for k in (1:Nz; blockIdx().z)
        @loop for j in (1:Ny; (blockIdx().y - 1) * blockDim().y + threadIdx().y)
            @loop for i in (1:Nx; (blockIdx().x - 1) * blockDim().x + threadIdx().x)
                for (Gφ, i) in enumerate(G)
                      φ_eqn = equation[i]
                      Gφ[i, j, k] = φ_eqn(args... i, j, k)
                end
          end
    end
end

With a time-stepping kernel of that form we can easily add and subtract tracers, equations, sub grid closure variables, etc. I think the inner loop gets unrolled when the array is static, so the compiled code is no different from what we currently have.

This is somewhere down the line hopefully. Maybe v0.6... or 1.0. Heh.

[glwagner]

Let's wait until isbitstype issue is solved before doing this stuff so we can do it without painfully writing out 10,000 arguments.

[ali-ramadhan]

I really like what you've sketched out. It definitely looks like a more robust solution than peppering macros around. I think some will argue that we shouldn't completely abstract away the equations as they might want to see something that looks like a momentum equation, but we should just write good code. I agree that we should revisit once #59 is resolved.

Hopefully once I finish working on vchuravy/GPUifyLoops.jl#18 we can even have something more like

@loop for I in (eachindex(grid); gpuIndex3D())
    for (Gφ, i) in enumerate(G)
        φ_eqn = equation[i]
        Gφ[i, j, k] = φ_eqn(args... i, j, k)
    end
end

Note: I think @zhenwu0728 needs this feature a bit sooner than v1.0 but we can work on hacking in a fix on his fork.

[glwagner]

I think some will argue that we shouldn't completely abstract away the equations as they might want to see something that looks like a momentum equation

I'm definitely willing to debate pros and cons when the code is ready for it!

What about my solution to define an array in the script and reference it the user-specified function? As in

# define a (maybe const a = ...)
u_forcing(args..., i, j, k) = a[i, j, k]

I'm surprised there are arrays you can't define with a function though.

[glwagner]

By the way, the abstraction may actually increase clarity and the 'appearance of a momentum equation', because the definition of the momentum equation can be done somewhere easy to find, rather than buried in a time-stepping loop. With the struct I proposed you would still have to write

Gu(args..., i, j, k) = # lots of momentum-equation looking things
u_eqn = Equation(Gu, Fu)

The abstraction lets you put this code in a file called, for example, equations.jl (or momentum_equations.jl!!), rather than in time_stepping.jl.

[ali-ramadhan]

That's a good point about just designing the abstractions to make the equations transparent. Having a momentum_equations.jl file instead of burying them in the time stepping would be the ultimate in clarity.

What about my solution to define an array in the script and reference it the user-specified function?

That's definitely the easiest solution and should work with the current code. @zhenwu0728 we should try it.

[zhenwu0728]

What about my solution to define an array in the script and reference it the user-specified function?

That's definitely the easiest solution and should work with the current code. @zhenwu0728 we should try it.

Definitely.
By the way, are all these arguments necessary when defining a forcing function?
I mean can I do

external_N_input(i, j, k) = a[i, j, k]

instead of

external_N_input(u, v, w, T, S, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k) = a[i, j, k]

Also I want to do something like this

function irradiance(u, v, w, T, S, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k, surface_irra)
    surface_irra*exp(-kw*Δz*k)
end

Is that possible?

[glwagner]

The arguments are necessary --- right now.

If we change the call signature to have i, j, k in front, you could legitimately write

external_N_input(i, j, k, args...) = a[i, j, k]

writing args... to stand for arbitrary arguments is valid julia syntax.

You may also be able to use Varargs as the code is written now. I'm not sure.

What is surface_irra? You can do, for example

surface_irradiance = 400
kw = 20
function irradiance(u, v, w, T, S, Nx, Ny, Nz, Δx, Δy, Δz, i, j, k)
   surface_irradiance*exp(-kw*Δz*(k-Nz))
end

We should make z an argument to the forcing function...

[glwagner]

@ali-ramadhan we should set up some kind of Stack Overflow - like system (discourse?) for answering user questions.

[glwagner]

@zhenwu0728 see:

https://docs.julialang.org/en/v1/manual/functions/index.html#Varargs-Functions-1

for a bit of info on how to define julia functions.

[glwagner]

See also

https://docs.julialang.org/en/v1/manual/variables-and-scoping/

to understand how scoping works, and how using the const keyword can help produce performant code when working with global variables.

The last example at that link is:

julia> const x = 1
1

julia> f() = x
f (generic function with 1 method)

julia> f()
1

julia> x = 2
WARNING: redefining constant x
2

julia> f()
1

which is illustrative.

[ali-ramadhan]

I think this is resolved. You just define a const array and the forcing function accesses it. See #110 (comment)

Although might make sense to benchmark this.

X-Ref: #370