Error using `ScalarDiffusivity` where the viscosity or diffusivity are functions with parameters

[ali-ramadhan]

I believe the below MWE should work according to the ScalarDiffusivity docstring, but from the error it seems to be expecting funky_diffusion(x, y, z, t) instead of funky_diffusion(x, y, z, t, p).

https://github.com/CliMA/Oceananigans.jl/blob/fe056fb44ce7173ce9e7eaa4f5c349d6ee2b61a4/src/TurbulenceClosures/turbulence_closure_implementations/scalar_diffusivity.jl#L12-L67

MWE:

using Oceananigans

grid = RectilinearGrid(CPU(), size=(3, 4, 5), extent=(1, 1, 1))

@inline funky_diffusion(x, y, z, t, p) = p.A + p.M * (x+y+z+t)

params = (; A=1.2, M=0.7)

closure = ScalarDiffusivity(;
    ν = funky_diffusion,
    κ = funky_diffusion,
    parameters = params
)

model = NonhydrostaticModel(; grid, closure)

time_step!(model, 0.1)

Error:

ERROR: MethodError: no method matching funky_diffusion(::Float64, ::Float64, ::Float64, ::Float64)

Closest candidates are:
  funky_diffusion(::Any, ::Any, ::Any, ::Any, ::Any)
   @ Main REPL[2]:1

Stacktrace:
  [1] νᶜᶜᶜ
    @ ~/atdepth/Oceananigans.jl/src/TurbulenceClosures/abstract_scalar_diffusivity_closure.jl:309 [inlined]
  [2] νᶜᶜᶜ
    @ ~/atdepth/Oceananigans.jl/src/TurbulenceClosures/abstract_scalar_diffusivity_closure.jl:84 [inlined]
  [3] ν_σᶜᶜᶜ
    @ ~/atdepth/Oceananigans.jl/src/TurbulenceClosures/abstract_scalar_diffusivity_closure.jl:154 [inlined]
  [4] viscous_flux_ux
    @ ~/atdepth/Oceananigans.jl/src/TurbulenceClosures/abstract_scalar_diffusivity_closure.jl:159 [inlined]
  [5] viscous_flux_ux
    @ ~/atdepth/Oceananigans.jl/src/TurbulenceClosures/implicit_explicit_time_discretization.jl:43 [inlined]
  [6] _viscous_flux_ux
    @ ~/atdepth/Oceananigans.jl/src/TurbulenceClosures/closure_kernel_operators.jl:4 [inlined]
  [7] Ax_qᶜᶜᶜ
    @ ~/atdepth/Oceananigans.jl/src/Operators/products_between_fields_and_grid_metrics.jl:12 [inlined]
  [8] δxᶠᵃᵃ
    @ ~/atdepth/Oceananigans.jl/src/Operators/difference_operators.jl:21 [inlined]
  [9] ∂ⱼ_τ₁ⱼ
    @ ~/atdepth/Oceananigans.jl/src/TurbulenceClosures/closure_kernel_operators.jl:24 [inlined]
 [10] u_velocity_tendency
    @ ~/atdepth/Oceananigans.jl/src/Models/NonhydrostaticModels/nonhydrostatic_tendency_kernel_functions.jl:71 [inlined]
 [11] cpu_compute_Gu!
    @ ~/.julia/packages/KernelAbstractions/491pi/src/macros.jl:291 [inlined]
 [12] cpu_compute_Gu!(__ctx__::KernelAbstractions.CompilerMetadata{…}, Gu::Field{…}, grid::RectilinearGrid{…}, ::Nothing, args::Tuple{…})
    @ Oceananigans.Models.NonhydrostaticModels ./none:0
 [13] __thread_run(tid::Int64, len::Int64, rem::Int64, obj::KernelAbstractions.Kernel{…}, ndrange::Nothing, iterspace::KernelAbstractions.NDIteration.NDRange{…}, args::Tuple{…}, dynamic::KernelAbstractions.NDIteration.DynamicCheck)
    @ KernelAbstractions ~/.julia/packages/KernelAbstractions/491pi/src/cpu.jl:144
 [14] __run(obj::KernelAbstractions.Kernel{…}, ndrange::Nothing, iterspace::KernelAbstractions.NDIteration.NDRange{…}, args::Tuple{…}, dynamic::KernelAbstractions.NDIteration.DynamicCheck, static_threads::Bool)
    @ KernelAbstractions ~/.julia/packages/KernelAbstractions/491pi/src/cpu.jl:111
 [15] (::KernelAbstractions.Kernel{…})(::Field{…}, ::Vararg{…}; ndrange::Nothing, workgroupsize::Nothing)
    @ KernelAbstractions ~/.julia/packages/KernelAbstractions/491pi/src/cpu.jl:46
 [16] (::KernelAbstractions.Kernel{…})(::Field{…}, ::Vararg{…})
    @ KernelAbstractions ~/.julia/packages/KernelAbstractions/491pi/src/cpu.jl:39
 [17] launch!(::CPU, ::RectilinearGrid{…}, ::Symbol, ::typeof(Oceananigans.Models.NonhydrostaticModels.compute_Gu!), ::Field{…}, ::Vararg{…}; include_right_boundaries::Bool, reduced_dimensions::Tuple{}, location::Nothing, active_cells_map::Nothing, kwargs::@Kwargs{})
    @ Oceananigans.Utils ~/atdepth/Oceananigans.jl/src/Utils/kernel_launching.jl:168
 [18] launch!
    @ ~/atdepth/Oceananigans.jl/src/Utils/kernel_launching.jl:154 [inlined]
 [19] #compute_interior_tendency_contributions!#17
    @ ~/atdepth/Oceananigans.jl/src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl:102 [inlined]
 [20] compute_interior_tendency_contributions!
    @ ~/atdepth/Oceananigans.jl/src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl:54 [inlined]
 [21] compute_tendencies!(model::NonhydrostaticModel{…}, callbacks::Vector{…})
    @ Oceananigans.Models.NonhydrostaticModels ~/atdepth/Oceananigans.jl/src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl:32
 [22] #apply_regionally!#56
    @ ~/atdepth/Oceananigans.jl/src/Utils/multi_region_transformation.jl:121 [inlined]
 [23] apply_regionally!
    @ ~/atdepth/Oceananigans.jl/src/Utils/multi_region_transformation.jl:118 [inlined]
 [24] macro expansion
    @ ~/atdepth/Oceananigans.jl/src/Utils/multi_region_transformation.jl:206 [inlined]
 [25] update_state!(model::NonhydrostaticModel{…}, callbacks::Vector{…}; compute_tendencies::Bool)
    @ Oceananigans.Models.NonhydrostaticModels ~/atdepth/Oceananigans.jl/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl:52
 [26] update_state!
    @ ~/atdepth/Oceananigans.jl/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl:20 [inlined]
 [27] time_step!(model::NonhydrostaticModel{…}, Δt::Float64; callbacks::Vector{…})
    @ Oceananigans.TimeSteppers ~/atdepth/Oceananigans.jl/src/TimeSteppers/runge_kutta_3.jl:85
 [28] time_step!(model::NonhydrostaticModel{…}, Δt::Float64)
    @ Oceananigans.TimeSteppers ~/atdepth/Oceananigans.jl/src/TimeSteppers/runge_kutta_3.jl:81
 [29] top-level scope
    @ REPL[6]:1
Some type information was truncated. Use `show(err)` to see complete types.

[ali-ramadhan]

Hmmm, I think the issue is that nothing happens with the parameters in convert_diffusivity if discrete_form = false.

https://github.com/CliMA/Oceananigans.jl/blob/fe056fb44ce7173ce9e7eaa4f5c349d6ee2b61a4/src/TurbulenceClosures/turbulence_closure_utils.jl#L19-L22

[glwagner]

Looks like it isn't actually supported. We'd need a ContinuousDiffusionFunction or something like that.

[ali-ramadhan]

Ah less a bug and more a feature request then. I can clarify what's supported in the docstring. Can always just use the discrete form.

[glwagner]

What's needed is something like

struct ContinuousDiffusionFunction{F, P}
    func :: F
    parameters :: P
end

@inline (K::ContinuousDiffusionFunction)(x, y, z, t) = K.func(x, y, z, t, K.parameters)

@inline function convert_diffusivity(FT, κ; discrete_form=false, loc=(nothing, nothing, nothing), parameters=nothing) 
    if discrete_form
        return DiscreteDiffusionFunction(κ; loc, parameters)
    elseif isnothing(parameters)
        return ContinuousDiffusionFunction(κ, parameters)
    else
        return κ 
    end 
 end 

I think anyways. reduced dimensionality grids may be different

[simone-silvestri]

I guess to have a continuous diffusion function that has the same features of the discrete version (with field dependency and parameters), we could implement something very similar to the ContinuousForcing. That would require a regularization.

[glwagner]

Pretty sure what I wrote will work. May need adapt as well.