Extract updates from PR #3850

src/TurbulenceClosures/abstract_scalar_biharmonic_diffusivity_closure.jl

@@ -5,12 +5,14 @@ using Oceananigans.Grids: peripheral_node
 
 Abstract type for closures with scalar biharmonic diffusivities.
 """
-abstract type AbstractScalarBiharmonicDiffusivity{F, N} <: AbstractTurbulenceClosure{ExplicitTimeDiscretization, N} end
+abstract type AbstractScalarBiharmonicDiffusivity{F, N, VI} <: AbstractTurbulenceClosure{ExplicitTimeDiscretization, N} end
 
 @inline formulation(::AbstractScalarBiharmonicDiffusivity{F}) where {F} = F()
 
 const ASBD = AbstractScalarBiharmonicDiffusivity
 
+const VectorInvariantASBD = AbstractScalarBiharmonicDiffusivity{<:HorizontalFormulation, <:Nothing, <:VectorInvariantForm}
+
 #####
 ##### Coefficient extractors
 #####
@@ -71,27 +73,34 @@ const AVBD = AbstractScalarBiharmonicDiffusivity{<:VerticalFormulation}
 ##### Biharmonic-specific viscous operators
 #####
 
-# See https://mitgcm.readthedocs.io/en/latest/algorithm/algorithm.html#horizontal-dissipation
-@inline function δ★ᶜᶜᶜ(i, j, k, grid, u, v)
+@inline ∇²u_vector_invariantᶠᶜᶜ(i, j, k, grid, u, v) = δxᶠᶜᶜ(i, j, k, grid, div_xyᶜᶜᶜ, u, v) - δyᶠᶜᶜ(i, j, k, grid, ζ₃ᶠᶠᶜ, u, v)
+@inline ∇²v_vector_invariantᶜᶠᶜ(i, j, k, grid, u, v) = δxᶜᶠᶜ(i, j, k, grid, ζ₃ᶠᶠᶜ, u, v)     + δyᶜᶠᶜ(i, j, k, grid, div_xyᶜᶜᶜ, u, v)
 
-    # These closures seem to be needed to help the compiler infer types
-    # (either of u and v or of the function arguments)
-    @inline Δy_∇²u(i, j, k, grid, u) = Δy_qᶠᶜᶜ(i, j, k, grid, biharmonic_mask_x, ∇²hᶠᶜᶜ, u)
-    @inline Δx_∇²v(i, j, k, grid, v) = Δx_qᶜᶠᶜ(i, j, k, grid, biharmonic_mask_y, ∇²hᶜᶠᶜ, v)
+# These closures seem to be needed to help the compiler infer types
+# (either of u and v or of the function arguments)
+@inline Δy_∇²u(i, j, k, grid, closure, u, v) = Δy_qᶠᶜᶜ(i, j, k, grid, biharmonic_mask_x, ∇²hᶠᶜᶜ, u)
+@inline Δx_∇²v(i, j, k, grid, closure, u, v) = Δx_qᶜᶠᶜ(i, j, k, grid, biharmonic_mask_y, ∇²hᶜᶠᶜ, v)
 
-    return 1 / Azᶜᶜᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, Δy_∇²u, u) +
-                                       δyᵃᶜᵃ(i, j, k, grid, Δx_∇²v, v))
-end
+# These closures seem to be needed to help the compiler infer types
+# (either of u and v or of the function arguments)
+@inline Δy_∇²u(i, j, k, grid, ::VectorInvariantASBD, u, v) = Δy_qᶠᶜᶜ(i, j, k, grid, biharmonic_mask_x, ∇²u_vector_invariantᶠᶜᶜ, u, v)
+@inline Δx_∇²v(i, j, k, grid, ::VectorInvariantASBD, u, v) = Δx_qᶜᶠᶜ(i, j, k, grid, biharmonic_mask_y, ∇²v_vector_invariantᶜᶠᶜ, u, v)
 
-@inline function ζ★ᶠᶠᶜ(i, j, k, grid, u, v)
+# See https://mitgcm.readthedocs.io/en/latest/algorithm/algorithm.html#horizontal-dissipation
+@inline function δ★ᶜᶜᶜ(i, j, k, grid, closure, u, v)
+    return 1 / Azᶜᶜᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, Δy_∇²u, closure, u, v) +
+                                       δyᵃᶜᵃ(i, j, k, grid, Δx_∇²v, closure, u, v))
+end
 
-    # These closures seem to be needed to help the compiler infer types
-    # (either of u and v or of the function arguments)
-    @inline Δy_∇²v(i, j, k, grid, v) = Δy_qᶜᶠᶜ(i, j, k, grid, biharmonic_mask_y, ∇²hᶜᶠᶜ, v)
-    @inline Δx_∇²u(i, j, k, grid, u) = Δx_qᶠᶜᶜ(i, j, k, grid, biharmonic_mask_x, ∇²hᶠᶜᶜ, u)
+@inline function ζ★ᶠᶠᶜ(i, j, k, grid, closure, u, v)
+    return 1 / Azᶠᶠᶜ(i, j, k, grid) * (δxᶠᵃᵃ(i, j, k, grid, Δy_∇²v, closure, u, v) -
+                                       δyᵃᶠᵃ(i, j, k, grid, Δx_∇²u, closure, u, v))
+end
 
-    return 1 / Azᶠᶠᶜ(i, j, k, grid) * (δxᶠᵃᵃ(i, j, k, grid, Δy_∇²v, v) -
-                                       δyᵃᶠᵃ(i, j, k, grid, Δx_∇²u, u))
+@inline function ζ★ᶠᶠᶜ(i, j, k, grid, ::VectorInvariantASBD, u, v)
+    ∇²u = ∇²u_vector_invariantᶠᶜᶜ(i, j, k, grid, u, v)
+    ∇²v = ∇²v_vector_invariantᶜᶠᶜ(i, j, k, grid, u, v)
+    return ζ₃ᶠᶠᶜ(i, j, k, grid, ∇²u, ∇²v)
 end
 
 #####

src/TurbulenceClosures/abstract_scalar_diffusivity_closure.jl

@@ -27,6 +27,13 @@ Specifies a horizontally-isotropic, `VectorInvariant`, `ScalarDiffusivity`.
 """
 struct HorizontalFormulation <: AbstractDiffusivityFormulation end
 
+"""
+    struct VectorInvariantForm end
+
+Specifies a `VectorInvariant` `ScalarDiffusivity`.
+"""
+struct VectorInvariantForm <: AbstractDiffusivityFormulation end
+
 """
     struct HorizontalDivergenceFormulation end
 

src/TurbulenceClosures/turbulence_closure_implementations/scalar_biharmonic_diffusivity.jl

@@ -1,15 +1,15 @@
-import Oceananigans.Grids: required_halo_size
+import Oceananigans.Grids: required_halo_size_x, required_halo_size_y, required_halo_size_z
 using Oceananigans.Utils: prettysummary
 
 """
     struct ScalarBiharmonicDiffusivity{F, N, K} <: AbstractScalarBiharmonicDiffusivity{F}
 
 Holds viscosity and diffusivities for models with prescribed isotropic diffusivities.
 """
-struct ScalarBiharmonicDiffusivity{F, V, K, N} <: AbstractScalarBiharmonicDiffusivity{F, N}
+struct ScalarBiharmonicDiffusivity{F, N, VI, V, K} <: AbstractScalarBiharmonicDiffusivity{F, N, VI}
     ν :: V
     κ :: K
-    ScalarBiharmonicDiffusivity{F, N}(ν::V, κ::K) where {F, V, K, N} = new{F, V, K, N}(ν, κ)
+    ScalarBiharmonicDiffusivity{F, N, VI}(ν::V, κ::K) where {F, N, VI, V, K} = new{F, N, VI, V, K}(ν, κ)
 end
 
 # Aliases that allow specify the floating type, assuming that the discretization is Explicit in time
@@ -50,6 +50,9 @@ Keyword arguments
 
 * `discrete_form`: `Boolean`; default: `false`.
 
+* `required_halo_size = 2`: the required halo size for the closure. This value should be an integer.
+  change only if using a function for `ν` or `κ` that requires a halo size larger than 1 to compute.
+
 When prescribing the viscosities or diffusivities as functions, depending on the
 value of keyword argument `discrete_form`, the constructor expects:
 
@@ -73,18 +76,29 @@ function ScalarBiharmonicDiffusivity(formulation = ThreeDimensionalFormulation()
                                      ν = 0,
                                      κ = 0,
                                      discrete_form = false,
+                                     vector_invariant_form = true, 
                                      loc = (nothing, nothing, nothing),
                                      parameters = nothing,
-                                     required_halo_size = 2)
+                                     required_halo_size::Int = 2) 
 
     ν = convert_diffusivity(FT, ν; discrete_form, loc, parameters)
     κ = convert_diffusivity(FT, κ; discrete_form, loc, parameters)
-    return ScalarBiharmonicDiffusivity{typeof(formulation), required_halo_size}(ν, κ)
+
+    VI = vector_invariant_form ? VectorInvariantForm : nothing
+
+    # Force a type-stable constructor if ν and κ are numbers
+    # This particular short-circuiting of the required_halo_size kwargs is necessary to perform parameter
+    # estimation of the diffusivity coefficients using autodiff.
+    if ν isa Number && κ isa Number
+        return ScalarBiharmonicDiffusivity{typeof(formulation), 2, VI}(ν, κ)
+    end
+
+    return ScalarBiharmonicDiffusivity{typeof(formulation), required_halo_size, VI}(ν, κ)
 end
 
-function with_tracers(tracers, closure::ScalarBiharmonicDiffusivity{F, N}) where {F, N}
+function with_tracers(tracers, closure::ScalarBiharmonicDiffusivity{F, N, VI}) where {F, N, VI}
     κ = tracer_diffusivities(tracers, closure.κ)
-    return ScalarBiharmonicDiffusivity{F, N}(closure.ν, κ)
+    return ScalarBiharmonicDiffusivity{F, N, VI}(closure.ν, κ)
 end
 
 @inline viscosity(closure::ScalarBiharmonicDiffusivity, K) = closure.ν
@@ -116,4 +130,4 @@ function on_architecture(to, closure::ScalarBiharmonicDiffusivity{F, <:Any, <:An
     ν = on_architecture(to, closure.ν)
     κ = on_architecture(to, closure.κ)
     return ScalarBiharmonicDiffusivity{F, N}(ν, κ)
-end
+end