What is the best way to implement the forcing term with derivatives?

[Leezhard]

Hi all, I had a problem about how to implement the simulation of 2D Oldroyd-B equations (x periodic and z bounded), where several derivatives are needed in the forcing terms. For example, I need to add:

$$ {u_{forcing}} = \frac{{\left( {1 - \beta } \right)}}{{{\mathop{\rm Re}\nolimits} }}\left( {\frac{{\partial {\tau _{xx}}}}{{\partial x}} + \frac{{\partial {\tau _{xz}}}}{{\partial z}}} \right) $$

$$ {\tau _{xx,forcing}} = \frac{2}{W}\frac{{\partial u}}{{\partial x}} - \frac{1}{W}{\tau _{xx}} + 2\left( {{\tau _{xx}}\frac{{\partial u}}{{\partial x}} + {\tau _{xz}}\frac{{\partial u}}{{\partial z}}} \right) $$

I have tried the discrete form in the following way, but this way seems not such "automatic":

function u_forcing_func(i, j, k, grid, clock, F)
    @inbounds begin
        dτxxdx = (F.τxx[i+1, j, k] - F.τxx[i-1, j, k]) / (2*dx)
        dτxzdz = (F.τxz[i, j, k+1] - F.τxz[i, j, k-1]) / (2*dz)
        return (1-β)/Re * (dτxxdx + dτxzdz)
    end
end
u_forcing = Forcing(u_forcing_func, discrete_form=true)

function τxx_forcing_func(i, j, k, grid, clock, F)
    @inbounds begin
        dudx = (F.u[i+1, j, k] - F.u[i-1, j, k]) / (2*dx)
        dudz = (F.u[i, j, k+1] - F.u[i, j, k-1]) / (2*dz)
        return 2/W * dudx - F.τxx[i, j, k]/W + 2*(F.τxx[i, j, k]*dudx + F.τxz[i, j, k]*dudz)
    end
end
τxx_forcing = Forcing(τxx_forcing_func, discrete_form=true)

Is it possible to use the derivative operator and how to implement it? I tried to use the operator but failed:

u_forcing = Forcing((x, z, t, F) ->
  (1 - β)/Re * (∂x(F.τxx) + ∂z(F.τxz)))

τxx_forcing = Forcing((x, z, t, F) ->
  2/W * ∂x(F.u) - F.τxx/W + 2*(F.τxx*∂x(F.u) + F.τxz*∂z(F.u)))

Any guidance would be appreciated. Thanks a lot!

[jagoosw]

You can use the derivative operators in the discrete form like:

using Oceananigans.Operators: ∂xᶠᶜᶜ, ∂zᶠᶜᶜ

function τxx_forcing_func(i, j, k, grid, clock, F)
    @inbounds begin
        dudx = ∂zᶠᶜᶜ(i, j, k, grid, F.τxx)
        dudz = ∂xᶠᶜᶜ(i, j, k, grid, F.τxz)
        return 2/W * dudx - F.τxx[i, j, k]/W + 2*(F.τxx[i, j, k]*dudx + F.τxz[i, j, k]*dudz)
    end
end
τxx_forcing = Forcing(τxx_forcing_func, discrete_form=true)

If you wanted to wrap this up to make the functions a bit more compact, the only other way that comes to mind is to make fields of the derivatives and then put them into the model as auxiliary fields, something like:

u, v, w = VelocityFields(grid)

τxx = ...

∂x_τxx = ∂x(τxx)
∂z_τxz = ∂z(τxz)

τxx_forcing_function(x, z, t, ∂x_τxx, ∂z_τxz, p) = (1-p.β)/p.Re * (x_τxx + z_τxz))

τxx_forcing = Forcing(τxx_forcing_function, field_dependancies = (:∂x_τxx, :∂z_τxz), parameters = ...)

model = NonhydristaticModel(; grid, velocities = (; u, v, w), auxiliary_fields = (; ∂x_τxx, ∂z_τxz), forcing = ...)

[Leezhard]

Thank you very much, that indeed helps a lot!

[glwagner]

Just wanted to comment that I think @jagoosw's first suggestion via the discrete form is probably the best way. We don't have an interface for including the derivatives of fields in forcing functions! It's possible to write such an interface, but given some of the difficulties we've encountered with the continuous form re: parameter space, GPU compilation, etc, I worry that a successful implementation may require a bit of work or even a redesign of Forcing itself. So the best option for now is the manual method (which requires taking care with staggered locations, etc).

I'll also note that it could be possible to implement this model as a new closure (eg extending the methods defined in Oceananigans.TurbulenceClosures), but that will probably be more effort at least at first.