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).

Oceananigans.jl/src/TurbulenceClosures/turbulence_closure_implementations/scalar_diffusivity.jl

Lines 12 to 67 in fe056fb

	"""
	ScalarDiffusivity(time_discretization = ExplicitTimeDiscretization(),
	formulation = ThreeDimensionalFormulation(), FT = Float64;
	ν = 0,
	κ = 0,
	discrete_form = false,
	loc = (nothing, nothing, nothing),
	parameters = nothing)
	Return `ScalarDiffusivity` turbulence closure with viscosity `ν` and tracer diffusivities `κ`
	for each tracer field in `tracers`. If a single `κ` is provided, it is applied to all tracers.
	Otherwise `κ` must be a `NamedTuple` with values for every tracer individually.
	Arguments
	=========
	* `time_discretization`: either `ExplicitTimeDiscretization()` (default)
	or `VerticallyImplicitTimeDiscretization()`.
	* `formulation`:
	- `HorizontalFormulation()` for diffusivity applied in the horizontal direction(s)
	- `VerticalFormulation()` for diffusivity applied in the vertical direction,
	- `ThreeDimensionalFormulation()` (default) for diffusivity applied isotropically to all directions
	* `FT`: the float datatype (default: `Float64`)
	Keyword arguments
	=================
	* `ν`: Viscosity. `Number`, `AbstractArray`, `Field`, or `Function`.
	* `κ`: Diffusivity. `Number`, `AbstractArray`, `Field`, `Function`, or
	`NamedTuple` of diffusivities with entries for each tracer.
	* `discrete_form`: `Boolean`; default: `false`.
	When prescribing the viscosities or diffusivities as functions, depending on the
	value of keyword argument `discrete_form`, the constructor expects:
	* `discrete_form = false` (default): functions of the grid's native coordinates
	and time, e.g., `(x, y, z, t)` for a `RectilinearGrid` or `(λ, φ, z, t)` for
	a `LatitudeLongitudeGrid`.
	* `discrete_form = true`:
	- with `loc = (nothing, nothing, nothing)` and `parameters = nothing` (default):
	functions of `(i, j, k, grid, ℓx, ℓy, ℓz, clock, fields)` with `ℓx`, `ℓy`,
	and `ℓz` either `Face()` or `Center()`.
	- with `loc = (ℓx, ℓy, ℓz)` with `ℓx`, `ℓy`, and `ℓz` either
	`Face()` or `Center()` and `parameters = nothing`: functions of `(i, j, k, grid, clock, fields)`.
	- with `loc = (nothing, nothing, nothing)` and specified `parameters`:
	functions of `(i, j, k, grid, ℓx, ℓy, ℓz, clock, fields, parameters)`.
	- with `loc = (ℓx, ℓy, ℓz)` and specified `parameters`:
	functions of `(i, j, k, grid, clock, fields, parameters)`.
	* `parameters`: `NamedTuple` with parameters used by the functions
	that compute viscosity and/or diffusivity; default: `nothing`.

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.

Oceananigans.jl/src/TurbulenceClosures/turbulence_closure_utils.jl

Lines 19 to 22 in fe056fb

	@inline function convert_diffusivity(FT, κ; discrete_form=false, loc=(nothing, nothing, nothing), parameters=nothing)
	discrete_form && return DiscreteDiffusionFunction(κ; loc, parameters)
	return κ
	end

[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.