(0.90.3) Fix bug in adapting ScalarDiffusivity and `ScalarBiharmoni…

…cDiffusivity` (#3413)

* Fix bug in adapting scalar diffusivity

* Use correct constructor for adapting ScalarDiffusivity and ScalarBiharmonicDiffusivity

* DiscreteBoundaryFunction -> DiscreteDiffusionFunction

* bump patch release

---------

Co-authored-by: Navid C. Constantinou <[email protected]>

Project.toml

@@ -1,7 +1,7 @@
 name = "Oceananigans"
 uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 authors = ["Climate Modeling Alliance and contributors"]
-version = "0.90.2"
+version = "0.90.3"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/TurbulenceClosures/discrete_diffusion_function.jl

@@ -81,13 +81,13 @@ const UnlocalizedDDF                 = DiscreteDiffusionFunction{<:Nothing, <:No
 const UnlocalizedUnparametrizedDDF   = DiscreteDiffusionFunction{<:Nothing, <:Nothing, <:Nothing, <:Nothing}
 
 @inline function getdiffusivity(dd::DiscreteDiffusionFunction{LX, LY, LZ},
-                              i, j, k, grid, location, clock, fields) where {LX, LY, LZ} 
+                                i, j, k, grid, location, clock, fields) where {LX, LY, LZ} 
     from = (LX(), LY(), LZ())
     return ℑxyz(i, j, k, grid, from, location, dd.func, clock, fields, dd.parameters)
 end
 
 @inline function getdiffusivity(dd::UnparameterizedDDF{LX, LY, LZ},
-                              i, j, k, grid, location, clock, fields) where {LX, LY, LZ}
+                                i, j, k, grid, location, clock, fields) where {LX, LY, LZ}
     from = (LX(), LY(), LZ())
     return ℑxyz(i, j, k, grid, from, location, dd.func, clock, fields)
 end
@@ -99,5 +99,5 @@ end
         dd.func(i, j, k, grid, location..., clock, fields)
 
 Adapt.adapt_structure(to, dd::DiscreteDiffusionFunction{LX, LY, LZ}) where {LX, LY, LZ} =
-     DiscreteBoundaryFunction{LX, LY, LZ}(Adapt.adapt(to, dd.func),
+    DiscreteDiffusionFunction{LX, LY, LZ}(Adapt.adapt(to, dd.func),
                                           Adapt.adapt(to, dd.parameters))

src/TurbulenceClosures/turbulence_closure_implementations/scalar_biharmonic_diffusivity.jl

@@ -103,5 +103,9 @@ end
 
 Base.show(io::IO, closure::ScalarBiharmonicDiffusivity) = print(io, summary(closure))
 
-adapt_structure(to, closure::ScalarBiharmonicDiffusivity) = ScalarBiharmonicDiffusivity(adapt(closure.ν),
-                                                                                        adapt(closure.κ))
+function Adapt.adapt_structure(to, closure::ScalarBiharmonicDiffusivity{F, <:Any, <:Any, N}) where {F, N}
+    ν = Adapt.adapt(to, closure.ν)
+    κ = Adapt.adapt(to, closure.κ)
+    return ScalarBiharmonicDiffusivity{F, N}(ν, κ)
+end
+

src/TurbulenceClosures/turbulence_closure_implementations/scalar_diffusivity.jl

@@ -1,11 +1,10 @@
-import Adapt: adapt_structure
+import Adapt
 import Oceananigans.Grids: required_halo_size
 using Oceananigans.Utils: prettysummary
 
 struct ScalarDiffusivity{TD, F, V, K, N} <: AbstractScalarDiffusivity{TD, F, N}
     ν :: V
     κ :: K
-
     ScalarDiffusivity{TD, F, N}(ν::V, κ::K) where {TD, F, V, K, N} = new{TD, F, V, K, N}(ν, κ)
 end
 
@@ -25,37 +24,42 @@ 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`)
+    * `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`, three-dimensional `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`.
+    * `ν`: Viscosity. `Number`, three-dimensional `AbstractArray`, `Field`, or `Function`.
+    
+    * `κ`: Diffusivity. `Number`, `AbstractArray`, `Field`, `Function`, or
+           `NamedTuple` of diffusivities with entries for each tracer.
+    
+    * `discrete_form`: `Boolean`; default: `False`.
 
-  - `discrete_form = true`: 
-    - with `loc = (nothing, nothing, nothing)` (default): functions of `(i, j, k, grid, ℓx, ℓy, ℓz)` with `ℓx`, `ℓy` and `ℓz` either `Face()` or `Center()`.
-    - with `loc = (ℓx, ℓy, ℓz)` with `ℓx`, `ℓy` and `ℓz` either `Face()` or `Center()`: functions of `(i, j, k, grid)`.
+When prescribing the viscosities or diffusivities as functions, depending on the
+value of keyword argument `discrete_form`, the constructor expects:
 
-  - `parameters`: `NamedTuple` with parameters used by the functions that compute viscosity and/or diffusivity; default: `nothing`.
+    * `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)` (default):
+          functions of `(i, j, k, grid, ℓx, ℓy, ℓz)` with `ℓx`, `ℓy`
+          and `ℓz` either `Face()` or `Center()`.
+        - with `loc = (ℓx, ℓy, ℓz)` with `ℓx`, `ℓy`
+          and `ℓz` either `Face()` or `Center()`: functions of `(i, j, k, grid)`.
 
+    * `parameters`: `NamedTuple` with parameters used by the functions
+      that compute viscosity and/or diffusivity; default: `nothing`.
 
 Examples
 ========
@@ -110,7 +114,8 @@ function ScalarDiffusivity(time_discretization=ExplicitTimeDiscretization(),
                            required_halo_size = 1)
 
     if formulation == HorizontalFormulation() && time_discretization == VerticallyImplicitTimeDiscretization()
-        throw(ArgumentError("VerticallyImplicitTimeDiscretization is only supported for `VerticalFormulation` or `ThreeDimensionalFormulation`"))
+        throw(ArgumentError("VerticallyImplicitTimeDiscretization is only supported for \
+                            `VerticalFormulation` or `ThreeDimensionalFormulation`"))
     end
 
     κ = convert_diffusivity(FT, κ; discrete_form, loc, parameters)
@@ -198,5 +203,9 @@ end
 
 Base.show(io::IO, closure::ScalarDiffusivity) = print(io, summary(closure))
 
-adapt_structure(to, closure::ScalarDiffusivity) = ScalarDiffusivity(adapt(closure.ν),
-                                                                    adapt(closure.κ))
+function Adapt.adapt_structure(to, closure::ScalarDiffusivity{TD, F, <:Any, <:Any, N}) where {TD, F, N}
+    ν = Adapt.adapt(to, closure.ν)
+    κ = Adapt.adapt(to, closure.κ)
+    return ScalarDiffusivity{TD, F, N}(ν, κ)
+end
+