diff --git a/src/Utils/sinking_velocity_fields.jl b/src/Utils/sinking_velocity_fields.jl index 26266c37d..257ef1b13 100644 --- a/src/Utils/sinking_velocity_fields.jl +++ b/src/Utils/sinking_velocity_fields.jl @@ -1,23 +1,32 @@ -using Oceananigans.Fields: ZFaceField, AbstractField +using Oceananigans.Fields: ZFaceField, AbstractField, location, Center, Face using Oceananigans.Forcings: maybe_constant_field using Oceananigans.Grids: AbstractGrid import Adapt: adapt_structure, adapt +const valid_sinking_velocity_locations = ((Center, Center, Face), (Nothing, Nothing, Face), (Nothing, Nothing, Nothing)) # nothings for constant fields + function setup_velocity_fields(drift_speeds, grid::AbstractGrid, open_bottom; smoothing_distance = 2) drift_velocities = [] - for w in values(drift_speeds) - if isa(values(w), Number) + for (name, w) in pairs(drift_speeds) + if isa(w, Number) w_field = ZFaceField(grid) for k=1:grid.Nz @inbounds w_field[:, :, k] .= - w * ifelse(open_bottom, 1.0, (1 - exp((1-k) / smoothing_distance))) end - w = w_field - elseif !isa(values(w), Tuple{AbstractField, AbstractField, AbstractField}) - error("Provided sinking speeds are an unsuitable value, must be a `NamedTuple` of scalar values (positive = sinking) with keys of tracer names, or `NamedTuple` of `VelocityFields`") + elseif isa(w, AbstractField) + location(w) in valid_sinking_velocity_locations || + @warn "The location of the sinking velocity field provided for $name is incorrect, it should be (Center, Center, Face)" + + open_bottom || @warn "The sinking velocity provided for $name is a field and therefore `open_bottom=false` can't be enforced automatically" + + w_field = w + else + @warn "Sinking speed provided for $name was not a number or field so may be unsiutable" + w_field = w end - push!(drift_velocities, w) + push!(drift_velocities, w_field) end return NamedTuple{keys(drift_speeds)}(drift_velocities) diff --git a/test/test_utils.jl b/test/test_utils.jl index b93cec4a9..b7c1c32b1 100644 --- a/test/test_utils.jl +++ b/test/test_utils.jl @@ -1,5 +1,9 @@ include("dependencies_for_runtests.jl") +using OceanBioME: setup_velocity_fields, valid_sinking_velocity_locations + +using Oceananigans.Fields: AbstractField, CenterField, ConstantField, FunctionField, ZFaceField + function test_negative_scaling(arch) grid = RectilinearGrid(arch, size = (1, 1, 1), extent = (1, 1, 1)) @@ -35,7 +39,32 @@ function test_negative_zeroing(arch) return (N ≈ 2) && (P ≈ 0.0) && (Z ≈ -1) end -@testset "Test Utils" begin +@testset "Test negative tracer handeling" begin @test test_negative_scaling(architecture) @test test_negative_zeroing(architecture) end + +scalar_sinking_speeds = (A = 1, B = 1.0) + +grid = RectilinearGrid(architecture, size = (1, 1, 10), extent = (1, 1, 10)) + +field_sinking_speeds = (C = ConstantField(1), D = FunctionField{Center, Center, Face}((x, y, z) -> z, grid), E = ZFaceField(grid)) + +sinking_speeds = merge(scalar_sinking_speeds, field_sinking_speeds) + +@testset "Test sinking velocity setup" begin + sinking_velocities = @test_nowarn setup_velocity_fields(sinking_speeds, grid, true) + + @test all(map(w -> isa(w, AbstractField) & (location(w) in valid_sinking_velocity_locations), values(sinking_velocities))) + + sinking_velocities = + @test_warn ("The sinking velocity provided for C is a field and therefore `open_bottom=false` can't be enforced automatically", + "The sinking velocity provided for D is a field and therefore `open_bottom=false` can't be enforced automatically", + "The sinking velocity provided for E is a field and therefore `open_bottom=false` can't be enforced automatically") setup_velocity_fields(sinking_speeds, grid, false) + + @test all([isa(w, AbstractField) & (location(w) in valid_sinking_velocity_locations) for w in values(sinking_velocities)]) + + @test all(map(w -> Array(interior(w, 1, 1, 1)) .== 0, sinking_velocities[(:A, :B)])) + + @test_warn "The location of the sinking velocity field provided for X is incorrect, it should be (Center, Center, Face)" setup_velocity_fields((X = CenterField(grid), ), grid, true) +end