Add SeawaterDensity to Oceananigans.Models (#3329)

* First steps * Simpler because `S` and `T` tracers are already required What to do with `LinearEquationOfState` which does not have a reference density? * For now restrict to `BoussinesqEquationOfState` * Update density_model.jl * Working, need tests * Add tests * Include the test script Need to run the tests and look at some of my examples with a clean julia session before opening the PR * Tests pass locally from cold start * Only need export `SeawaterDensity` * Extend to `SewaterBuoyancy` with constant `S` or `T` Still needs to be a `BoussinesqEquationOfState`. * Better docstrings * Doctest and GPU fixes Still not sure if `GPU()` tests will be fixed but I think array type should be correct. I ran `doctest` locally on the `Oceananigans.Models` module and there was no error * Update src/Models/Models.jl Co-authored-by: Gregory L. Wagner <[email protected]> * File rename to `seawater_density.jl` * Simplify testing functions * Changes after @glwagner comments * Export in main + `KernelFunctionOperation` check to `AbstractOperation` --------- Co-authored-by: Gregory L. Wagner <[email protected]> Co-authored-by: Gregory L. Wagner <[email protected]>
CliMA · Oct 16, 2023 · 2354238 · 2354238
1 parent 10aff1c
commit 2354238
Show file tree

Hide file tree

Showing 4 changed files with 320 additions and 8 deletions.
diff --git a/src/Models/Models.jl b/src/Models/Models.jl
@@ -6,7 +6,8 @@ export
     HydrostaticFreeSurfaceModel,
     ExplicitFreeSurface, ImplicitFreeSurface, SplitExplicitFreeSurface,
     PrescribedVelocityFields, PressureField,
-    LagrangianParticles
+    LagrangianParticles,
+    seawater_density
 
 using Oceananigans: AbstractModel, fields, prognostic_fields
 using Oceananigans.Grids: AbstractGrid, halo_size, inflate_halo_size
@@ -23,15 +24,13 @@ import Oceananigans.TimeSteppers: reset!
 using Oceananigans.OutputReaders: update_field_time_series!, extract_field_timeseries
 
 # A prototype interface for AbstractModel.
-# 
+#
 # TODO: decide if we like this.
 #
 # We assume that model has some properties, eg:
 #   - model.clock::Clock
 #   - model.architecture.
 #   - model.timestepper with timestepper.G⁻ and timestepper.Gⁿ :spiral_eyes:
-#
-# Perhaps this is a little unclean.
 
 iteration(model::AbstractModel) = model.clock.iteration
 Base.time(model::AbstractModel) = model.clock.time
@@ -169,5 +168,7 @@ function default_nan_checker(model::OceananigansModels)
     return nan_checker
 end
 
-end # module
+# This is here so that `NonhydrostaticModel` and  `HydrsostaticFreeSurfaceModel`
+include("seawater_density.jl")
 
+end # module
diff --git a/src/Models/seawater_density.jl b/src/Models/seawater_density.jl
@@ -0,0 +1,131 @@
+using Oceananigans.AbstractOperations: AbstractOperation, KernelFunctionOperation
+using Oceananigans.BuoyancyModels: SeawaterBuoyancy, Zᶜᶜᶜ
+using Oceananigans.Fields: field
+using Oceananigans.Grids: Center
+using SeawaterPolynomials: BoussinesqEquationOfState
+import SeawaterPolynomials.ρ
+
+"Extend `SeawaterPolynomials.ρ` to compute density for a `KernelFunctionOperation` -
+**note** `eos` must be `BoussinesqEquationOfState` because a reference density is needed for the computation."
+@inline ρ(i, j, k, grid, eos, T, S, Z) = @inbounds ρ(T[i, j, k], S[i, j, k], Z[i, j, k], eos)
+
+"Return a `KernelFunctionOperation` to compute the in-situ `seawater_density`."
+seawater_density(grid, eos, temperature, salinity, geopotential_height) =
+    KernelFunctionOperation{Center, Center, Center}(ρ, grid, eos, temperature, salinity, geopotential_height)
+
+const ModelsWithBuoyancy = Union{NonhydrostaticModel, HydrostaticFreeSurfaceModel}
+
+validate_model_eos(eos:: BoussinesqEquationOfState) = nothing
+validate_model_eos(eos) = throw(ArgumentError("seawater_density is not defined for $eos."))
+
+# some nice fallbacks
+model_temperature(bf, model)     = model.tracers.T
+model_salinity(bf, model)        = model.tracers.S
+model_geopotential_height(model) = KernelFunctionOperation{Center, Center, Center}(Zᶜᶜᶜ, model.grid)
+
+const ConstantTemperatureSB = SeawaterBuoyancy{FT, EOS, <:Number, <:Nothing} where {FT, EOS}
+const ConstantSalinitySB    = SeawaterBuoyancy{FT, EOS, <:Nothing, <:Number} where {FT, EOS}
+
+model_temperature(b::ConstantTemperatureSB, model) = b.constant_temperature
+model_salinity(b::ConstantSalinitySB, model)       = b.constant_salinity
+
+"""
+    seawater_density(model; temperature, salinity, geopotential_height)
+
+Return a `KernelFunctionOperation` that computes the in-situ density of seawater
+with (gridded) `temperature`, `salinity`, and at `geopotential_height`. To compute the
+in-situ density, the 55 term polynomial approximation to the equation of state from
+[Roquet et al. (2015)](https://www.sciencedirect.com/science/article/pii/S1463500315000566?ref=pdf_download&fr=RR-2&rr=813416acba58557b) is used.
+By default the `seawater_density` extracts the geopotential height from the model to compute
+the in-situ density. To compute a potential density at some user chosen reference geopotential height,
+set `geopotential_height` to a constant for the density computation,
+
+```julia
+geopotential_height = 0 # sea-surface height
+σ₀ = seawater_density(model; geopotential_height)
+```
+
+**Note:** `seawater_density` must be passed a `BoussinesqEquationOfState` to compute the
+density. See the [relevant documentation](https://clima.github.io/OceananigansDocumentation/dev/model_setup/buoyancy_and_equation_of_state/#Idealized-nonlinear-equations-of-state)
+for how to set `SeawaterBuoyancy` using a `BoussinesqEquationOfState`.
+
+Example
+=======
+
+Compute a density `Field` using the `KernelFunctionOperation` returned from `seawater_density`
+
+```jldoctest density
+julia> using Oceananigans, Oceananigans.Models, SeawaterPolynomials.TEOS10
+
+julia> grid = RectilinearGrid(CPU(), size=(32, 32), x=(0, 2π), y=(0, 2π), topology=(Periodic, Periodic, Flat))
+32×32×1 RectilinearGrid{Float64, Periodic, Periodic, Flat} on CPU with 3×3×0 halo
+├── Periodic x ∈ [3.60072e-17, 6.28319) regularly spaced with Δx=0.19635
+├── Periodic y ∈ [3.60072e-17, 6.28319) regularly spaced with Δy=0.19635
+└── Flat z
+
+julia> tracers = (:S, :T)
+(:S, :T)
+
+julia> eos = TEOS10EquationOfState()
+BoussinesqEquationOfState{Float64}:
+    ├── seawater_polynomial: TEOS10SeawaterPolynomial{Float64}
+    └── reference_density: 1020.0
+
+julia> buoyancy = SeawaterBuoyancy(equation_of_state=eos)
+SeawaterBuoyancy{Float64}:
+├── gravitational_acceleration: 9.80665
+└── equation of state: BoussinesqEquationOfState{Float64}
+
+julia> model = NonhydrostaticModel(; grid, buoyancy, tracers)
+NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0)
+├── grid: 32×32×1 RectilinearGrid{Float64, Periodic, Periodic, Flat} on CPU with 3×3×0 halo
+├── timestepper: QuasiAdamsBashforth2TimeStepper
+├── tracers: (S, T)
+├── closure: Nothing
+├── buoyancy: SeawaterBuoyancy with g=9.80665 and BoussinesqEquationOfState{Float64} with ĝ = NegativeZDirection()
+└── coriolis: Nothing
+
+julia> set!(model, S = 34.7, T = 0.5)
+
+julia> density_field = Field(seawater_density(model))
+32×32×1 Field{Center, Center, Center} on RectilinearGrid on CPU
+├── grid: 32×32×1 RectilinearGrid{Float64, Periodic, Periodic, Flat} on CPU with 3×3×0 halo
+├── boundary conditions: FieldBoundaryConditions
+│   └── west: Periodic, east: Periodic, south: Periodic, north: Periodic, bottom: Nothing, top: Nothing, immersed: ZeroFlux
+├── operand: KernelFunctionOperation at (Center, Center, Center)
+├── status: time=0.0
+└── data: 38×38×1 OffsetArray(::Array{Float64, 3}, -2:35, -2:35, 1:1) with eltype Float64 with indices -2:35×-2:35×1:1
+    └── max=0.0, min=0.0, mean=0.0
+
+julia> compute!(density_field)
+32×32×1 Field{Center, Center, Center} on RectilinearGrid on CPU
+├── grid: 32×32×1 RectilinearGrid{Float64, Periodic, Periodic, Flat} on CPU with 3×3×0 halo
+├── boundary conditions: FieldBoundaryConditions
+│   └── west: Periodic, east: Periodic, south: Periodic, north: Periodic, bottom: Nothing, top: Nothing, immersed: ZeroFlux
+├── operand: KernelFunctionOperation at (Center, Center, Center)
+├── status: time=0.0
+└── data: 38×38×1 OffsetArray(::Array{Float64, 3}, -2:35, -2:35, 1:1) with eltype Float64 with indices -2:35×-2:35×1:1
+    └── max=1027.7, min=1027.7, mean=1027.7
+```
+
+Values for `temperature`, `salinity` and `geopotential_height` can be passed to
+`seawater_density` to override the defaults that are obtained from the `model`.
+"""
+function seawater_density(model::ModelsWithBuoyancy;
+                         temperature = model_temperature(model.buoyancy.model, model),
+                         salinity = model_salinity(model.buoyancy.model, model),
+                         geopotential_height = model_geopotential_height(model))
+
+    eos = model.buoyancy.model.equation_of_state
+    validate_model_eos(eos)
+    # Convert function or constant user input to AbstractField
+    grid = model.grid
+    loc = (Center, Center, Center)
+    temperature = field(loc, temperature, grid)
+    salinity = field(loc, salinity, grid)
+
+    geopotential_height = geopotential_height isa AbstractOperation ? geopotential_height :
+                                                                      field(loc, geopotential_height, grid)
+
+    return seawater_density(grid, eos, temperature, salinity, geopotential_height)
+end
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -8,8 +8,8 @@ CUDA.allowscalar() do
             include(String(test_file))
         end
     end
-    
-    # Core Oceananigans 
+
+    # Core Oceananigans
     if group == :unit || group == :all
         @testset "Unit tests" begin
             include("test_grids.jl")
@@ -99,6 +99,7 @@ CUDA.allowscalar() do
         @testset "Model and time stepping tests (part 3)" begin
             include("test_dynamics.jl")
             include("test_biogeochemistry.jl")
+            include("test_sewater_density.jl")
         end
     end
 
@@ -124,7 +125,7 @@ CUDA.allowscalar() do
             include("test_hydrostatic_free_surface_immersed_boundaries_implicit_solve.jl")
         end
     end
-    
+
     # Model enhancements: cubed sphere, distributed, etc
     if group == :multi_region || group == :all
         @testset "Multi Region tests" begin

diff --git a/test/test_sewater_density.jl b/test/test_sewater_density.jl
@@ -0,0 +1,179 @@
+include("dependencies_for_runtests.jl")
+
+using Oceananigans.Models
+using Oceananigans.AbstractOperations: AbstractOperation
+using Oceananigans.Models: model_temperature, model_salinity, model_geopotential_height
+using Oceananigans.Models: ConstantTemperatureSB, ConstantSalinitySB
+using SeawaterPolynomials: ρ, BoussinesqEquationOfState
+using SeawaterPolynomials: SecondOrderSeawaterPolynomial, RoquetEquationOfState
+using SeawaterPolynomials: TEOS10EquationOfState, TEOS10SeawaterPolynomial
+using Oceananigans.BuoyancyModels: Zᶜᶜᶜ
+
+tracers = (:S, :T)
+ST_testvals = (S = 34.7, T = 0.5)
+Roquet_eos = (RoquetEquationOfState(:Linear), RoquetEquationOfState(:Cabbeling),
+              RoquetEquationOfState(:CabbelingThermobaricity), RoquetEquationOfState(:Freezing),
+              RoquetEquationOfState(:SecondOrder), RoquetEquationOfState(:SimplestRealistic))
+
+TEOS10_eos = TEOS10EquationOfState()
+
+"Return and `Array` on `arch` that is `size(grid)` flled with `value`."
+function grid_size_value(arch, grid, value)
+
+    value_array = fill(value, size(grid))
+
+    return arch_array(arch, value_array)
+
+end
+
+"Check the error thrown for non-`BoussinesqEquationOfState`."
+function error_non_Boussinesq(arch, FT)
+
+    grid = RectilinearGrid(arch, FT, size=(3, 3, 3), extent=(1, 1, 1))
+    buoyancy = SeawaterBuoyancy()
+    model = NonhydrostaticModel(; grid, buoyancy, tracers)
+    seawater_density(model) # throws error
+
+    return nothing
+end
+
+"""
+    function eos_works(arch, FT, eos::BoussinesqEquationOfState)
+Check if using a `BoussinesqEquationOfState` returns a `KernelFunctionOperation`.
+"""
+function eos_works(arch, FT, eos::BoussinesqEquationOfState;
+                   constant_temperature = nothing, constant_salinity = nothing)
+
+    grid = RectilinearGrid(arch, FT, size=(3, 3, 3), extent=(1, 1, 1))
+    buoyancy = SeawaterBuoyancy(equation_of_state = eos; constant_temperature, constant_salinity)
+    model = NonhydrostaticModel(; grid, buoyancy, tracers)
+
+    return seawater_density(model) isa AbstractOperation
+end
+
+"""
+    function insitu_density(arch, FT, eos)
+Use the `KernelFunctionOperation` returned from `seawater_density` to compute a density `Field`
+and compare the computed values to density values explicitly calculate using
+`SeawaterPolynomials.ρ`. Similar function is used to test the potential density computation.
+"""
+function insitu_density(arch, FT, eos::BoussinesqEquationOfState;
+                        constant_temperature = nothing, constant_salinity = nothing)
+
+    grid = RectilinearGrid(arch, FT, size=(3, 3, 3), extent=(1, 1, 1))
+    buoyancy = SeawaterBuoyancy(equation_of_state = eos; constant_temperature, constant_salinity)
+    model = NonhydrostaticModel(; grid, buoyancy, tracers)
+
+    if !isnothing(constant_temperature)
+        set!(model, S = ST_testvals.S)
+    elseif !isnothing(constant_salinity)
+        set!(model, T = ST_testvals.T)
+    else
+        set!(model, S = ST_testvals.S, T = ST_testvals.T)
+    end
+
+    d_field = compute!(Field(seawater_density(model)))
+
+    # Computation from SeawaterPolynomials to check against
+    geopotential_height = model_geopotential_height(model)
+    T_vec = grid_size_value(arch, grid, ST_testvals.T)
+    S_vec = grid_size_value(arch, grid, ST_testvals.S)
+    eos_vec = grid_size_value(arch, grid, model.buoyancy.model.equation_of_state)
+    SWP_ρ = similar(interior(d_field))
+    SWP_ρ .= SeawaterPolynomials.ρ.(T_vec, S_vec, geopotential_height, eos_vec)
+
+    return all(CUDA.@allowscalar interior(d_field) .== SWP_ρ)
+end
+
+function potential_density(arch, FT, eos::BoussinesqEquationOfState;
+                           constant_temperature = nothing, constant_salinity = nothing)
+
+    grid = RectilinearGrid(arch, FT, size=(3, 3, 3), extent=(1, 1, 1))
+    buoyancy = SeawaterBuoyancy(equation_of_state = eos; constant_temperature, constant_salinity)
+    model = NonhydrostaticModel(; grid, buoyancy, tracers)
+
+    if !isnothing(constant_temperature)
+        set!(model, S = ST_testvals.S)
+    elseif !isnothing(constant_salinity)
+        set!(model, T = ST_testvals.T)
+    else
+        set!(model, S = ST_testvals.S, T = ST_testvals.T)
+    end
+
+    d_field = compute!(Field(seawater_density(model, geopotential_height = 0)))
+
+    # Computation from SeawaterPolynomials to check against
+    geopotential_height = grid_size_value(arch, grid, 0)
+    T_vec = grid_size_value(arch, grid, ST_testvals.T)
+    S_vec = grid_size_value(arch, grid, ST_testvals.S)
+    eos_vec = grid_size_value(arch, grid, model.buoyancy.model.equation_of_state)
+    SWP_ρ = similar(interior(d_field))
+    SWP_ρ .= SeawaterPolynomials.ρ.(T_vec, S_vec, geopotential_height, eos_vec)
+
+    return all(CUDA.@allowscalar interior(d_field) .== SWP_ρ)
+end
+
+@testset "Density models" begin
+    @info "Testing `seawater_density`..."
+
+    @testset "Error for non-`BoussinesqEquationOfState`" begin
+        @info "Testing error is thrown... "
+        for FT in float_types
+            for arch in archs
+                @test_throws ArgumentError error_non_Boussinesq(arch, FT)
+            end
+        end
+    end
+
+    @testset "seawater_density `KernelFunctionOperation` instantiation" begin
+        @info "Testing `AbstractOperation` is returned..."
+
+        for FT in float_types
+            for arch in archs
+                for eos in Roquet_eos
+                    @test eos_works(arch, FT, eos)
+                    @test eos_works(arch, FT, eos, constant_temperature = ST_testvals.T)
+                    @test eos_works(arch, FT, eos, constant_salinity = ST_testvals.S)
+                end
+                @test eos_works(arch, FT, TEOS10_eos)
+                @test eos_works(arch, FT, TEOS10_eos, constant_temperature = ST_testvals.T)
+                @test eos_works(arch, FT, TEOS10_eos, constant_salinity = ST_testvals.S)
+            end
+        end
+    end
+
+    @testset "In-situ density computation tests" begin
+        @info "Testing in-situ density compuation..."
+
+        for FT in float_types
+            for arch in archs
+                for eos in Roquet_eos
+                    @test insitu_density(arch, FT, eos)
+                    @test insitu_density(arch, FT, eos, constant_temperature = ST_testvals.T)
+                    @test insitu_density(arch, FT, eos, constant_salinity = ST_testvals.S)
+                end
+                @test insitu_density(arch, FT, TEOS10_eos)
+                @test insitu_density(arch, FT, TEOS10_eos, constant_temperature = ST_testvals.T)
+                @test insitu_density(arch, FT, TEOS10_eos, constant_salinity = ST_testvals.S)
+            end
+        end
+    end
+
+    @testset "Potential density computation tests" begin
+        @info "Testing a potential density comnputation..."
+
+        for FT in float_types
+            for arch in archs
+                for eos in Roquet_eos
+                    @test potential_density(arch, FT, eos)
+                    @test potential_density(arch, FT, eos, constant_temperature = ST_testvals.T)
+                    @test potential_density(arch, FT, eos, constant_salinity = ST_testvals.S)
+                end
+                @test potential_density(arch, FT, TEOS10_eos)
+                @test potential_density(arch, FT, TEOS10_eos, constant_temperature = ST_testvals.T)
+                @test potential_density(arch, FT, TEOS10_eos, constant_salinity = ST_testvals.S)
+            end
+        end
+    end
+
+end