Merge pull request #5 from jagoosw/jsw/precomputed-roughness-length

added precomputed roughness length and friction velocity + fix bug

Project.toml

@@ -1,7 +1,7 @@
 name = "Walrus"
 uuid = "bec5164f-1c7a-4c32-8768-be9354254d6a"
 authors = ["Jago Stong-Wright <[email protected]>"]
-version = "0.4.2"
+version = "0.4.3"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/Walrus.jl

@@ -1,5 +1,7 @@
 module Walrus
 
+export SimpleInterpolation
+
 export WallStress, WallStressBoundaryConditions
 
 export HomogeneousBodyHeating
@@ -26,12 +28,14 @@ adapt_structure(to, rv::ReturnValue) = ReturnValue(adapt(to, rv.value))
 display_input(return_value::ReturnValue) = return_value.value
 display_input(::Function) = "Function"
 
+include("interpolations.jl")
 include("wall_model.jl")
 include("radiative_transfer/radiative_transfer.jl")
 include("tidal_forcings.jl")
 include("wind_stress.jl")
 include("surface_heating.jl")
 
+using .Interpolations
 using .WallStressModel
 using .RadiativeTransfer
 using .TidalForcings

src/interpolations.jl

@@ -0,0 +1,42 @@
+module Interpolations
+
+export SimpleInterpolation
+
+using Adapt: adapt
+
+using Oceananigans.Architectures: arch_array, CPU
+
+import Adapt: adapt_structure
+
+struct SimpleInterpolation{R, V}
+     range :: R
+    values :: V
+end
+
+adapt_structure(to, itp::SimpleInterpolation) = SimpleInterpolation(adapt(to, itp.range),
+                                                                    adapt(to, itp.values))
+
+function SimpleInterpolation(range::Array, values; arch = CPU())
+    x₀ = minimum(range)
+
+    (range[2] - range[1] ≈ range[end] - range[end - 1]) || throw(ArgumentError("Interpolation range must be regularly spaced"))
+
+    Δx = range[2] - range[1]
+
+    return SimpleInterpolation((; x₀, Δx), arch_array(arch, values))
+end
+
+
+function (itp::SimpleInterpolation)(x)
+    n₁ = floor(Int, (x - itp.range.x₀) / itp.range.Δx)
+
+    x₁ = itp.range.x₀ + itp.range.Δx * n₁
+    x₂ = x₁ + itp.range.Δx
+
+    y₁ = @inbounds itp.values[n₁ + 1]
+    y₂ = @inbounds itp.values[n₁ + 2]
+
+    return y₁ + (x - x₁) * (y₂ - y₁) / (x₂ - x₁)
+end
+
+end # module

src/radiative_transfer/homogeneous_body_heating.jl

@@ -92,10 +92,9 @@ end
     zᶠ⁺ = znode(i, j, k + 1, grid, Center(), Center(), Face())
 
     t = clock.time
-    A = Azᶠᶠᶜ(i, j, k, grid)
 
-    return heating.surface_flux(x, y, t) * (exp(- α * abs(zᶠ⁺)) - exp(- α * abs(zᶠ))) * A / (ρ * cₚ)
-end
+    return heating.surface_flux(x, y, t) * (exp(- α * abs(zᶠ⁺)) - exp(- α * abs(zᶠ))) / ((zᶠ⁺ - zᶠ) * (ρ * cₚ))
+end 
 
 summary(::HomogeneousBodyHeating) = string("Single band light attenuation and body heating model")
 show(io::IO, body_heating::HomogeneousBodyHeating) = println(io, string(summary(body_heating), " with: \n",

src/radiative_transfer/radiative_transfer.jl

@@ -17,7 +17,6 @@ using Oceananigans.Fields: CenterField
 using Oceananigans.Forcings: Forcing
 using Oceananigans.Grids: node, znode, Center, Face
 using Oceananigans.Utils: launch!
-using Oceananigans.Operators: Azᶠᶠᶜ
 
 using KernelAbstractions.Extras: @unroll
 

src/surface_heating.jl

@@ -119,10 +119,10 @@ julia> using Walrus
 julia> using Oceananigans
 
 julia> wind_stress = WindStress(; reference_wind_speed = 0., reference_wind_direction = 90.)
-(::WindStress{Walrus.ReturnValue{Float64}, Walrus.ReturnValue{Float64}, LogarithmicNeutralWind{Float64}, Float64}) (generic function with 2 methods)
+(::WindStress{Walrus.ReturnValue{Float64}, Walrus.ReturnValue{Float64}, LogarithmicNeutralWind{Float64, Nothing}, Float64}) (generic function with 2 methods)
 
 julia> surface_heat_exchange = SurfaceHeatExchange(; wind_stress)
-(::SurfaceHeatExchange{WindStress{Walrus.ReturnValue{Float64}, Walrus.ReturnValue{Float64}, LogarithmicNeutralWind{Float64}, Float64}, Walrus.ReturnValue{Int64}, Walrus.SurfaceHeatingModel.EmpiricalLatentHeatVaporisation{Float64}, Walrus.SurfaceHeatingModel.AugustRocheMagnusVapourPressure{Float64}, Float64}) (generic function with 1 method)
+(::SurfaceHeatExchange{WindStress{Walrus.ReturnValue{Float64}, Walrus.ReturnValue{Float64}, LogarithmicNeutralWind{Float64, Nothing}, Float64}, Walrus.ReturnValue{Int64}, Walrus.SurfaceHeatingModel.EmpiricalLatentHeatVaporisation{Float64}, Walrus.SurfaceHeatingModel.AugustRocheMagnusVapourPressure{Float64}, Float64}) (generic function with 1 method)
 
 julia> boundary_conditions = (; T = FieldBoundaryConditions(top = FluxBoundaryCondition(surface_heat_exchange, field_dependencies = (:T, :u, :v))))
 (T = Oceananigans.FieldBoundaryConditions, with boundary conditions
@@ -131,7 +131,7 @@ julia> boundary_conditions = (; T = FieldBoundaryConditions(top = FluxBoundaryCo
 ├── south: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing)
 ├── north: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing)
 ├── bottom: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing)
-├── top: FluxBoundaryCondition: ContinuousBoundaryFunction (::SurfaceHeatExchange{WindStress{Walrus.ReturnValue{Float64}, Walrus.ReturnValue{Float64}, LogarithmicNeutralWind{Float64}, Float64}, Walrus.ReturnValue{Int64}, Walrus.SurfaceHeatingModel.EmpiricalLatentHeatVaporisation{Float64}, Walrus.SurfaceHeatingModel.AugustRocheMagnusVapourPressure{Float64}, Float64}) at (Nothing, Nothing, Nothing)
+├── top: FluxBoundaryCondition: ContinuousBoundaryFunction (::SurfaceHeatExchange{WindStress{Walrus.ReturnValue{Float64}, Walrus.ReturnValue{Float64}, LogarithmicNeutralWind{Float64, Nothing}, Float64}, Walrus.ReturnValue{Int64}, Walrus.SurfaceHeatingModel.EmpiricalLatentHeatVaporisation{Float64}, Walrus.SurfaceHeatingModel.AugustRocheMagnusVapourPressure{Float64}, Float64}) at (Nothing, Nothing, Nothing)
 └── immersed: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing),)
 
 ```
@@ -190,10 +190,10 @@ Example
 julia> using Walrus
 
 julia> wind_stress = WindStress(; reference_wind_speed = 0., reference_wind_direction = 90.)
-(::WindStress{Walrus.ReturnValue{Float64}, Walrus.ReturnValue{Float64}, LogarithmicNeutralWind{Float64}, Float64}) (generic function with 2 methods)
+(::WindStress{Walrus.ReturnValue{Float64}, Walrus.ReturnValue{Float64}, LogarithmicNeutralWind{Float64, Nothing}, Float64}) (generic function with 2 methods)
 
 julia> surface_heat_exchange = SurfaceHeatExchangeBoundaryCondition(; wind_stress)
-FluxBoundaryCondition: ContinuousBoundaryFunction (::SurfaceHeatExchange{WindStress{Walrus.ReturnValue{Float64}, Walrus.ReturnValue{Float64}, LogarithmicNeutralWind{Float64}, Float64}, Walrus.ReturnValue{Int64}, Walrus.SurfaceHeatingModel.EmpiricalLatentHeatVaporisation{Float64}, Walrus.SurfaceHeatingModel.AugustRocheMagnusVapourPressure{Float64}, Float64}) at (Nothing, Nothing, Nothing)
+FluxBoundaryCondition: ContinuousBoundaryFunction (::SurfaceHeatExchange{WindStress{Walrus.ReturnValue{Float64}, Walrus.ReturnValue{Float64}, LogarithmicNeutralWind{Float64, Nothing}, Float64}, Walrus.ReturnValue{Int64}, Walrus.SurfaceHeatingModel.EmpiricalLatentHeatVaporisation{Float64}, Walrus.SurfaceHeatingModel.AugustRocheMagnusVapourPressure{Float64}, Float64}) at (Nothing, Nothing, Nothing)
 
 ```
 """
@@ -230,7 +230,7 @@ end
 
     params = (; κ, ν, aᶜ, b, g)
 
-    z₀ = find_velocity_roughness_length(wind_speed, 10, params)
+    z₀ = find_velocity_roughness_length(drag_coefficient, wind_speed, 10, params)
 
     ū = κ * wind_speed / log(10 / z₀)
 

src/wall_model.jl

@@ -11,29 +11,57 @@ export WallStress, WallStressBoundaryConditions
 
 using Roots
 
+using Adapt: adapt
+
 using Oceananigans.BoundaryConditions: FluxBoundaryCondition
 using Oceananigans.Fields: Center, Face
 using Oceananigans.Grids: znode
 
+using Oceananigans.Architectures: arch_array, CPU, architecture
+
+using Walrus.Interpolations: SimpleInterpolation
+
 import Adapt: adapt_structure
 import Base: summary, show
 
+struct WallStress{FT, U} <: Function
+    von_Karman_constant :: FT
+    kinematic_viscosity :: FT
+                      B :: FT
+
+      friction_velocities :: U
+end
+
+adapt_structure(to, ws::WallStress) = WallStress(ws.von_Karman_constant, ws.kinematic_viscosity,
+                                                 ws.B, adapt(to, ws.friction_velocities))
+
 """
     WallStress(; von_Karman_constant = 0.4,
                  kinematic_viscosity = 1e-6,
-                 B = 5.2)
+                 B = 5.2,
+                 precomputed_friction_velocities = false,
+                 precompute_speeds = [0:25/100000:25;],
+                 grid = nothing,
+                 arch = isnothing(grid) ? CPU() : architecture(grid))
 
 Returns a wall stress model for LES simulation with default parameters similar
 to that proposed in [SCHUMANN1975376](@citet), [HARTEL1996283](@citet),
 [Piomelli1989](@citet), and [taylor2007](@citet).
 
+Friction velocities will be precomputed at `precompute_speeds` if 
+`precomputed_friction_velocities` is true and `grid` is provided.
 
 Keyword Arguments
 =================
 
 - `von_Karman_constant`: the von Karman wall stress constant
 - `kinematic_viscosity`: kinematic viscosity of the water above the wall
 - `B`: wall stress constant 
+- `precomputed_friction_velocities`: precompute friction velocities?
+- `precompute_speeds`: bottom water speeds to precompute friction velocities for, 
+  this should encompas the range of speeds possible in your simulation
+- `grid`: the grid to precompute the friction velocities for
+- `arch`: architecture to adapt precomputed friction velocities for
 
 Example
 =======
@@ -44,37 +72,72 @@ julia> using Walrus: WallStress
 julia> using Oceananigans
 
 julia> wall_stress = WallStress()
-(::WallStress{Float64}) (generic function with 1 method)
+(::WallStress{Float64, Nothing}) (generic function with 1 method)
 julia> boundary_conditions = (u = FieldBoundaryConditions(bottom = FluxBoundaryCondition(wall_stress, discrete_form = true, parameters = Val(:x))),
                               v = FieldBoundaryConditions(bottom = FluxBoundaryCondition(wall_stress, discrete_form = true, parameters = Val(:y))))
 (u = Oceananigans.FieldBoundaryConditions, with boundary conditions
 ├── west: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing)
 ├── east: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing)
 ├── south: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing)
 ├── north: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing)
-├── bottom: FluxBoundaryCondition: DiscreteBoundaryFunction (::WallStress{Float64}) with parameters Val{:x}
+├── bottom: FluxBoundaryCondition: DiscreteBoundaryFunction (::WallStress{Float64, Nothing}) with parameters Val{:x}
 ├── top: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing)
 └── immersed: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing), v = Oceananigans.FieldBoundaryConditions, with boundary conditions
 ├── west: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing)
 ├── east: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing)
 ├── south: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing)
 ├── north: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing)
-├── bottom: FluxBoundaryCondition: DiscreteBoundaryFunction (::WallStress{Float64}) with parameters Val{:y}
+├── bottom: FluxBoundaryCondition: DiscreteBoundaryFunction (::WallStress{Float64, Nothing}) with parameters Val{:y}
 ├── top: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing)
 └── immersed: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing))
 
 ```
 """
-@kwdef struct WallStress{FT} <: Function
-    von_Karman_constant :: FT = 0.4
-    kinematic_viscosity :: FT = 1.15e-6
-                      B :: FT = 5.2
-end
+function WallStress(; von_Karman_constant::FT = 0.4,
+                      kinematic_viscosity::FT = 1.15e-6,
+                      B::FT = 5.2,
+                      precomputed_friction_velocities = false,
+                      precompute_speeds = [0:25/100000:25;],
+                      grid = nothing,
+                      arch = isnothing(grid) ? CPU() : architecture(grid)) where FT
+
+    if precomputed_friction_velocities
+        tmp = WallStress(von_Karman_constant, kinematic_viscosity, B, nothing)
+
+        ν = kinematic_viscosity
+        κ = von_Karman_constant
+
+        z₁ = znode(1, grid, Center()) - znode(1, grid, Face())
+
+        n = 100000
+        velocities = zeros(length(precompute_speeds))
+
+        for (n, speed) in enumerate(precompute_speeds)
+            params = (; z₁, κ, ν, B)
+            velocities[n] = find_friction_velocity(tmp, speed, params)
+        end
+
+        friction_velocities = SimpleInterpolation(precompute_speeds, velocities; arch)
+    else
+        friction_velocities = nothing
+    end
 
-adapt_structure(to, ws::WallStress) = ws
+    return WallStress(von_Karman_constant, kinematic_viscosity, B, friction_velocities)
+end
 
 @inline stress_velocity(uₜ, params) = log(params.z₁ * uₜ / (params.ν + eps(0.0))) / (params.κ + eps(0.0)) + params.B - params.U₁ / (uₜ + eps(0.0))
 
+@inline function find_friction_velocity(::WallStress{<:Any, Nothing}, U₁, params)
+    uₜ = 0
+
+    U₁ == 0 || (uₜ = find_zero(stress_velocity, (0., Inf), Bisection(); p = merge(params, (; U₁)), maxiters = 10^5))
+
+    return uₜ
+end
+
+@inline find_friction_velocity(ws::WallStress, U₁, params) = 
+    ws.friction_velocities(U₁)
+
 @inline function (wall_stress::WallStress)(i, j, grid, clock, model_fields, ::Val{direction}) where direction
     ν = wall_stress.kinematic_viscosity
     κ = wall_stress.von_Karman_constant
@@ -89,9 +152,7 @@ adapt_structure(to, ws::WallStress) = ws
 
     U₁ = sqrt(u ^ 2 + v ^ 2)
 
-    uₜ = find_zero(stress_velocity, (0., Inf), Bisection(); p = (; U₁, z₁, κ, ν, B), maxiters = 10^5)
-
-    U₁ == 0 && (uₜ = 0)
+    uₜ = find_friction_velocity(wall_stress, U₁, (; z₁, κ, ν, B))
 
     return ifelse(direction == :x, -u / (U₁ + eps(0.0)) * uₜ ^ 2, -v / (U₁ + eps(0.0)) * uₜ ^ 2)
 end
@@ -100,7 +161,9 @@ end
 """
     WallStressBoundaryConditions(; von_Karman_constant = 0.4,
                                    kinematic_viscosity = 1e-6,
-                                   B = 5.2)
+                                   B = 5.2,
+                                   precompute_speeds = [0:25/100000:25;],
+                                   grid = nothing)
 
 Convenience constructor to setup `WallStress` boundary conditions.
 
@@ -110,6 +173,10 @@ Keyword Arguments
 - `von_Karman_constant`: the von Karman wall stress constant
 - `kinematic_viscosity`: kinematic viscosity of the water above the wall
 - `B`: wall stress constant 
+- `precomputed_friction_velocities`: precompute friction velocities?
+- `precompute_speeds`: bottom water speeds to precompute friction velocities for, 
+  this should encompas the range of speeds possible in your simulation
+- `grid`: the grid to precompute the friction velocities for
 
 Example
 =======
@@ -120,33 +187,39 @@ julia> using Walrus: WallStressBoundaryConditions
 julia> using Oceananigans
 
 julia> stress_boundary_conditions = WallStressBoundaryConditions()
-(u = FluxBoundaryCondition: DiscreteBoundaryFunction (::WallStress{Float64}) with parameters Val{:x}, v = FluxBoundaryCondition: DiscreteBoundaryFunction (::WallStress{Float64}) with parameters Val{:y})
+(u = FluxBoundaryCondition: DiscreteBoundaryFunction (::WallStress{Float64, Nothing}) with parameters Val{:x}, v = FluxBoundaryCondition: DiscreteBoundaryFunction (::WallStress{Float64, Nothing}) with parameters Val{:y})
 julia> boundary_conditions = (u = FieldBoundaryConditions(bottom = stress_boundary_conditions.u),
                               v = FieldBoundaryConditions(bottom = stress_boundary_conditions.v))
 (u = Oceananigans.FieldBoundaryConditions, with boundary conditions
 ├── west: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing)
 ├── east: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing)
 ├── south: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing)
 ├── north: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing)
-├── bottom: FluxBoundaryCondition: DiscreteBoundaryFunction (::WallStress{Float64}) with parameters Val{:x}
+├── bottom: FluxBoundaryCondition: DiscreteBoundaryFunction (::WallStress{Float64, Nothing}) with parameters Val{:x}
 ├── top: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing)
 └── immersed: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing), v = Oceananigans.FieldBoundaryConditions, with boundary conditions
 ├── west: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing)
 ├── east: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing)
 ├── south: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing)
 ├── north: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing)
-├── bottom: FluxBoundaryCondition: DiscreteBoundaryFunction (::WallStress{Float64}) with parameters Val{:y}
+├── bottom: FluxBoundaryCondition: DiscreteBoundaryFunction (::WallStress{Float64, Nothing}) with parameters Val{:y}
 ├── top: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing)
 └── immersed: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing))
 ```
 """
-function WallStressBoundaryConditions(; von_Karman_constant = 0.4,
-                                        kinematic_viscosity = 1e-6,
-                                        B = 5.2)
-
-    wall_stress_instance = WallStress(von_Karman_constant,
-                                      kinematic_viscosity,
-                                      B)
+function WallStressBoundaryConditions(; von_Karman_constant::FT = 0.4,
+                                        kinematic_viscosity::FT = 1.15e-6,
+                                        B::FT = 5.2,
+                                        precomputed_friction_velocities = false,
+                                        precompute_speeds = [0:25/100000:25;],
+                                        grid = nothing) where FT
+
+    wall_stress_instance = WallStress(; von_Karman_constant,
+                                        kinematic_viscosity,
+                                        B,
+                                        precomputed_friction_velocities,
+                                        precompute_speeds,
+                                        grid)
 
     u = FluxBoundaryCondition(wall_stress_instance, discrete_form = true, parameters = Val(:x))