Oceananigans custom structures for functions

[m3azevedo]

Hello and thank you for maintaining this amazingly intuitive model,

I have started using Oceananigans this year and am currently developing some custom functions for building boundaries and forcings. Sorry if this is a bit of basic question, but I noticed most functions are wrapped in custom structures. Why is the code written this way and is there an advantage for building these structures instead of passing the function directly to where they are used?

Exemplifying, the relaxation scheme creates a mask structure, which is then passed as one of the fields for the relaxation struct, then the functions are inlined as in the shortened code below. Why not pass the masking function directly to the relaxation structure?

struct Relaxation{R, M, T}
      rate :: R
      mask :: M
    target :: T
end

@inline (f::Relaxation)(x, y, z, t, field) = f.rate * f.mask(x, y, z) * (f.target(x, y, z, t) - field)

struct GaussianMask{D, T}
    center :: T
     width :: T

    function GaussianMask{D}(; center, width) where D
        T = promote_type(typeof(center), typeof(width))
        return new{D, T}(center, width)
    end
end

@inline (g::GaussianMask{:x})(x, y, z) = exp(-(x - g.center)^2 / (2 * g.width^2))
@inline (g::GaussianMask{:y})(x, y, z) = exp(-(y - g.center)^2 / (2 * g.width^2))
@inline (g::GaussianMask{:z})(x, y, z) = exp(-(z - g.center)^2 / (2 * g.width^2))

[glwagner]

Thanks! <3

The purpose of this design is to support a user interface wherein masks can be plainly specified as functions of (x, y, z) (which may or may not close over parameters in a user script).

To support parameterized functions within this framework in the source code, those parameterized functions must be attached to callable objects that store the parameters. For example, the GaussianMask depends on center and width. So

@inline (g::GaussianMask{:x})(x, y, z) = exp(-(x - g.center)^2 / (2 * g.width^2))

makes g::GaussianMask into a callable object that is a function of (x, y, z), but also provides access to the two parameters g.center and g.width.

You can use this in your own code too, if you don't want to use masking functions that are closures (ie they capture parameters from the global scope). It's a bit neater to capture parameters via callable objects (your code is more portable).

Does that help / is that the answer you are looking for?

[m3azevedo]

Thanks, it was a clarifying explanation. It seems wrapping the functions in structures is the way to go when implementing custom functions on a GPU architecture.

[glwagner]

It's not GPU specific. What other design would you propose?