`closure` not working in `ShallowWaterModel`

[francispoulin]

This is a minimum working example that shows how we (@writingindy and I) have tried to use closure in the VectorInvariantFormulation of the ShallowWaterModel but get a TaskFailedException error, copied below.

@simone-silvestri , do you have any idea as to how this can be fixed?

using Oceananigans
using Oceananigans.Models.ShallowWaterModels: VectorInvariantFormulation
using Oceananigans.Advection: VelocityStencil, VorticityStencil
using Oceananigans.TurbulenceClosures

grid = RectilinearGrid(size = (16, 16), x = (0, 1), y = (0, 1), 
                   topology = (Periodic, Periodic, Flat))

model = ShallowWaterModel(grid = grid,
                          momentum_advection = WENO5(vector_invariant = VelocityStencil()),
                          tracers = (:A),
                          closure = HorizontalScalarBiharmonicDiffusivity(κ=1e-6),
                          gravitational_acceleration = 1.0,
                          formulation = VectorInvariantFormulation()
                          )

set!(model, h=1)
simulation = Simulation(model, Δt = 0.001, stop_time = 10.0)

run!(simulation)

Error:

ERROR: LoadError: TaskFailedException

    nested task error: MethodError: no method matching getindex(::Nothing, ::Int64)
    Stacktrace:
      [1] call
        @ ~/.julia/packages/Cassette/34vIw/src/context.jl:456 [inlined]
      [2] fallback
        @ ~/.julia/packages/Cassette/34vIw/src/context.jl:454 [inlined]
      [3] _overdub_fallback(::Any, ::Vararg{Any, N} where N)
        @ ~/.julia/packages/Cassette/34vIw/src/overdub.jl:586 [inlined]
      [4] overdub
        @ ~/.julia/packages/Cassette/34vIw/src/overdub.jl:586 [inlined]
      [5] diffusive_flux_x(::Int64, ::Int64, ::Int64, ::RectilinearGrid{Float64, Periodic, Periodic, Flat, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float

[simone-silvestri]

Is it there also in the ConservativeFormulation case?

[francispoulin]

Good question, and I tried it and, yes, in happens in both formulations.

[simone-silvestri]

The signature of the ∇_dot_qᶜ in solution_and_tracers.jl (line 107) seems to be wrong.

Try changing it with - ∇_dot_qᶜ(i, j, k, grid, closure, diffusivities, val_tracer_index, solution, tracers, clock, nothing)

[simone-silvestri]

If you want you can open a PR to fix it, maybe the one where we deal with closures + ShallowWaterModel

[francispoulin]

I am happy to open a PR but when I try the line you suggested , unfortunately, there is still a problem.

Some good news! This fixes it for the vector invariant formulation.

However, when I try the conervative form it complains about not knowing u. See the start of the output below.

ERROR: LoadError: TaskFailedException

    nested task error: type NamedTuple has no field u
    Stacktrace:
      [1] getproperty(x::NamedTuple{(:uh, :vh, :h), Tuple{Field{Face, Center, Center, Nothing, RectilinearGrid{Float64, Periodic, Periodic, Flat, Float64, Float64, Float64, OffsetArrays.

[simone-silvestri]

Ah, I see the problem... all the viscous_flux are formulated in terms of velocities U.u, U.v, U.w which the shallow water model does not necessarily have. This will be slightly more involved to fix than just changing the signature of the function.

It depends how you want to formulate the viscous operator, as a diffusion of u or a diffusion of uh. If the former is always the desired formulation, both in the vector invariant and the conservative equations, then it might be useful to add the velocities as a field of the shallow water model and pass them to the tendency kernel.

So indeed, we should open a PR to adapt the shallow water model to work with closures

[francispoulin]

I agree this is more complicated. I will create a PR to propose the change you suggested. Maybe next week we can discuss closure for the conservative model, and do that in a separate PR?