Merge pull request #210 from OceanBioME/jsw/fix-fields-in-sinking-setup

Fixes `setup_sinking_velocities` and adds tests

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)

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