Source term of an average field

[andrewooish]

I need to implement a source term on the u velocity momentum equation that looks something like -u/Ub, where u(x,y,z,t) is the instantaneous x component of velocity and Ub is the instantaneous volume average of u(x,y,z,t), i.e. Ub(t). I have tried various expressions for the forcing function but none of the two below seem to work. For some reason, Oceananigans does not like (or rather I cannot get it to like) the Average() function in the forcing function. Any help would be greatly appreciated. Thank you.

u_forcing_func(x, z, t, u) = -u/(Average(u, dims=(1,2, 3)))
u_forcing_func(x, z, t, u) = -u/(Field(Average(u, dims=(1,2, 3))))


u_forcing = Forcing(u_forcing_func, field_dependencies=(:u))
model = NonhydrostaticModel(; grid,
                            advection = WENO(),
                            timestepper = :RungeKutta3,
                            tracers = :b,
                            forcing=(u=u_forcing,),
                            closure = ScalarDiffusivity(; ν, κ),
                            boundary_conditions = (; b=b_bcs))

[glwagner]

This requires defining the average in advance, and then updating it in a callback. Like this:

Ub_forcing = Ref(zero(grid)) # can also use Field{Nothing, Nothing, Nothing}(grid)
u_forcing_func(x, z, t, u, Ub) = -u / Ub[]
u_forcing = Forcing(u_forcing_func, field_dependencies=:u, parameters=Ub_forcing)

# define model and simulation

Ub = Field(Average(model.velocities.u))

using GPUArrays: @allowscalar # req for GPU

function update_forcing!(sim)
    compute!(Ub)
    # update Ub used in the forcing function
    Ub_forcing[] = @allowscalar first(Ub)
    return nothing
end

add_callback!(simulation, update_forcing!)

[andrewooish]

Thank you very much for your prompt response. Yes, this works perfectly.