Fix bug when using Lagrangian particles on Flat topologies

src/Fields/interpolate.jl

@@ -15,7 +15,7 @@ using Oceananigans.Architectures: child_architecture
 """
     index_binary_search(val, vec, N)
 
-Return indices `low, high` of `vec`tor for which 
+Return indices `low, high` of `vec`tor for which
 
 ```julia
 vec[low] ≤ val && vec[high] ≥ val
@@ -29,9 +29,9 @@ Code credit: https://gist.github.com/cuongld2/8e4fed9ba44ea2b4598f90e7d5b6c612/1
     low = 0
     high = N - 1
 
-    while low + 1 < high 
+    while low + 1 < high
         mid = middle_point(low, high)
-        if @inbounds vec[mid + 1] == val 
+        if @inbounds vec[mid + 1] == val
             return (mid + 1, mid + 1)
         elseif @inbounds vec[mid + 1] < val
             low = mid
@@ -211,6 +211,23 @@ end
     return (ii, jj, kk)
 end
 
+"""
+    truncate_fractional_indices(fi, fj, fk)
+
+Truncate _fractional_ indices output from fractional indices `fi, fj, fk` to integer indices, dealing
+with `nothing` indices for `Flat` domains.
+"""
+@inline function truncate_fractional_indices(fi, fj, fk)
+    i = truncate_fractional_index(fi)
+    j = truncate_fractional_index(fj)
+    k = truncate_fractional_index(fk)
+    return (i, j, k)
+end
+
+@inline truncate_fractional_index(::Nothing) = 1
+@inline truncate_fractional_index(fi) = Base.unsafe_trunc(Int, fi)
+
+
 """
     interpolate(at_node, from_field, from_loc, from_grid)
 
@@ -277,7 +294,7 @@ end
     k⁻, k⁺, ζ = iz
 
     return @inbounds ϕ₁(ξ, η, ζ) * getindex(data, i⁻, j⁻, k⁻, in...) +
-                     ϕ₂(ξ, η, ζ) * getindex(data, i⁻, j⁻, k⁺, in...) +  
+                     ϕ₂(ξ, η, ζ) * getindex(data, i⁻, j⁻, k⁺, in...) +
                      ϕ₃(ξ, η, ζ) * getindex(data, i⁻, j⁺, k⁻, in...) +
                      ϕ₄(ξ, η, ζ) * getindex(data, i⁻, j⁺, k⁺, in...) +
                      ϕ₅(ξ, η, ζ) * getindex(data, i⁺, j⁻, k⁻, in...) +
@@ -316,7 +333,7 @@ function interpolate!(to_field::Field, from_field::AbstractField)
     to_arch   = architecture(to_field)
     from_arch = architecture(from_field)
 
-    # In case architectures are `Distributed` we 
+    # In case architectures are `Distributed` we
     # verify that the fields are on the same child architecture
     to_arch   = child_architecture(to_arch)
     from_arch = child_architecture(from_arch)
@@ -338,4 +355,3 @@ function interpolate!(to_field::Field, from_field::AbstractField)
 
     return nothing
 end
-

src/Models/LagrangianParticleTracking/LagrangianParticleTracking.jl

@@ -17,7 +17,7 @@ using Oceananigans.Grids: XYFlatGrid, YZFlatGrid, XZFlatGrid
 using Oceananigans.ImmersedBoundaries: immersed_cell
 using Oceananigans.Architectures: device, architecture
 using Oceananigans.Fields: interpolate, datatuple, compute!, location
-using Oceananigans.Fields: fractional_indices
+using Oceananigans.Fields: fractional_indices, truncate_fractional_indices
 using Oceananigans.TimeSteppers: AbstractLagrangianParticles
 using Oceananigans.Utils: prettysummary, launch!, SumOfArrays
 
@@ -135,7 +135,7 @@ step_lagrangian_particles!(::Nothing, model, Δt) = nothing
 function step_lagrangian_particles!(particles::LagrangianParticles, model, Δt)
     # Update the properties of the Lagrangian particles
     update_lagrangian_particle_properties!(particles, model, Δt)
-    
+
     # Compute dynamics
     particles.dynamics(particles, model, Δt)
 

src/Models/LagrangianParticleTracking/lagrangian_particle_advection.jl

@@ -1,4 +1,4 @@
-using Oceananigans.Utils: instantiate 
+using Oceananigans.Utils: instantiate
 using Oceananigans.Models: total_velocities
 
 #####
@@ -28,6 +28,13 @@ is outside the Periodic interval `(xᴸ, xᴿ)`.
 @inline enforce_boundary_conditions(::Periodic, x, xᴸ, xᴿ, Cʳ) = ifelse(x > xᴿ, xᴸ + (x - xᴿ),
                                                                  ifelse(x < xᴸ, xᴿ - (xᴸ - x), x))
 
+"""
+    enforce_boundary_conditions(::Flat, x, xᴸ, xᴿ, Cʳ)
+
+Do nothing on Flat dimensions.
+"""
+@inline enforce_boundary_conditions(::Flat, x, xᴸ, xᴿ, Cʳ) = x
+
 const f = Face()
 const c = Center()
 
@@ -41,23 +48,21 @@ bouncing the particle off the immersed boundary with a coefficient or `restituti
     X = flattened_node((x, y, z), ibg)
 
     # Determine current particle cell
-    i, j, k = fractional_indices(X, ibg.underlying_grid, (c, c, c))
-    i = Base.unsafe_trunc(Int, i)
-    j = Base.unsafe_trunc(Int, j)
-    k = Base.unsafe_trunc(Int, k)
-
+    fi, fj, fk = fractional_indices(X, ibg.underlying_grid, (c, c, c))
+    i, j, k = truncate_fractional_indices(fi, fj, fk)
+
     if immersed_cell(i, j, k, ibg)
         # Determine whether particle was _previously_ in a non-immersed cell
         i⁻, j⁻, k⁻ = previous_particle_indices
-       
+
         if !immersed_cell(i⁻, j⁻, k⁻, ibg)
             # Left-right bounds of the previous, non-immersed cell
             xᴿ, yᴿ, zᴿ = node(i⁻+1, j⁻+1, k⁻+1, ibg, f, f, f)
             xᴸ, yᴸ, zᴸ = node(i⁻,   j⁻,   k⁻,   ibg, f, f, f)
 
             Cʳ = restitution
-            x⁺ = enforce_boundary_conditions(Bounded(), x, xᴸ, xᴿ, Cʳ)    
-            y⁺ = enforce_boundary_conditions(Bounded(), y, yᴸ, yᴿ, Cʳ)    
+            x⁺ = enforce_boundary_conditions(Bounded(), x, xᴸ, xᴿ, Cʳ)
+            y⁺ = enforce_boundary_conditions(Bounded(), y, yᴸ, yᴿ, Cʳ)
             z⁺ = enforce_boundary_conditions(Bounded(), z, zᴸ, zᴿ, Cʳ)
 
         end
@@ -76,10 +81,8 @@ given `velocities`, time-step `Δt, and coefficient of `restitution`.
     X = flattened_node((x, y, z), grid)
 
     # Obtain current particle indices
-    i, j, k = fractional_indices(X, grid, c, c, c)
-    i = Base.unsafe_trunc(Int, i)
-    j = Base.unsafe_trunc(Int, j)
-    k = Base.unsafe_trunc(Int, k)
+    fi, fj, fk = fractional_indices(X, grid, c, c, c)
+    i, j, k = truncate_fractional_indices(fi, fj, fk)
 
     current_particle_indices = (i, j, k)
 
@@ -91,7 +94,7 @@ given `velocities`, time-step `Δt, and coefficient of `restitution`.
     # Advect particles, calculating the advection metric for a curvilinear grid.
     # Note that all supported grids use length coordinates in the vertical, so we do not
     # transform the vertical velocity nor invoke the k-index.
-    ξ = x_metric(i, j, grid) 
+    ξ = x_metric(i, j, grid)
     η = y_metric(i, j, grid)
 
     x⁺ = x + ξ * u * Δt
@@ -137,7 +140,7 @@ end
 @inline y_metric(i, j, grid::RectilinearGrid) = 1
 @inline y_metric(i, j, grid::LatitudeLongitudeGrid{FT}) where FT = 1 / grid.radius * FT(360 / 2π)
 
-@kernel function _advect_particles!(particles, restitution, grid::AbstractUnderlyingGrid, Δt, velocities) 
+@kernel function _advect_particles!(particles, restitution, grid::AbstractUnderlyingGrid, Δt, velocities)
     p = @index(Global)
 
     @inbounds begin
@@ -146,12 +149,12 @@ end
         z = particles.z[p]
     end
 
-    x⁺, y⁺, z⁺ = advect_particle((x, y, z), p, restitution, grid, Δt, velocities) 
+    x⁺, y⁺, z⁺ = advect_particle((x, y, z), p, restitution, grid, Δt, velocities)
 
     @inbounds begin
-        particles.x[p] = x⁺ 
-        particles.y[p] = y⁺ 
-        particles.z[p] = z⁺ 
+        particles.x[p] = x⁺
+        particles.y[p] = y⁺
+        particles.z[p] = z⁺
     end
 end
 
@@ -166,4 +169,3 @@ function advect_lagrangian_particles!(particles, model, Δt)
 
     return nothing
 end
-

test/test_lagrangian_particle_tracking.jl

@@ -2,16 +2,16 @@ include("dependencies_for_runtests.jl")
 
 using NCDatasets
 using StructArrays
-using Oceananigans.Architectures: on_architecture
+using Oceananigans.Architectures: architecture, on_architecture
 
 struct TestParticle{T}
-    x :: T
-    y :: T
-    z :: T
-    u :: T
-    v :: T
-    w :: T
-    s :: T
+    x::T
+    y::T
+    z::T
+    u::T
+    v::T
+    w::T
+    s::T
 end
 
 function particle_tracking_simulation(; grid, particles, timestepper=:RungeKutta3, velocities=nothing)
@@ -22,41 +22,31 @@ function particle_tracking_simulation(; grid, particles, timestepper=:RungeKutta
     jld2_filepath = "test_particles.jld2"
     sim.output_writers[:particles_jld2] =
         JLD2OutputWriter(model, (; particles=model.particles),
-                         filename="test_particles", schedule=IterationInterval(1))
+            filename="test_particles", schedule=IterationInterval(1))
 
     nc_filepath = "test_particles.nc"
     sim.output_writers[:particles_nc] =
-        NetCDFOutputWriter(model, model.particles, filename=nc_filepath, schedule=IterationInterval(1)) 
+        NetCDFOutputWriter(model, model.particles, filename=nc_filepath, schedule=IterationInterval(1))
 
     sim.output_writers[:checkpointer] = Checkpointer(model, schedule=IterationInterval(1),
-                                                     dir = ".", prefix = "particles_checkpoint")
+                                                     dir=".", prefix="particles_checkpoint")
 
     return sim, jld2_filepath, nc_filepath
 end
 
-function run_simple_particle_tracking_tests(arch, timestepper; vertically_stretched=false)
-    topo = (Periodic, Periodic, Bounded)
-
-    Nx = Ny = Nz = 5
-    if vertically_stretched
-        # Slightly stretched at the top
-        z = [-1, -0.5, 0.0, 0.4, 0.7, 1]
-    else
-        z = (-1, 1)
-    end
+function run_simple_particle_tracking_tests(grid, timestepper)
 
-    grid = RectilinearGrid(arch; topology=topo, size=(Nx, Ny, Nz),
-                           x=(-1, 1), y=(-1, 1), z)
+    arch = architecture(grid)
 
     P = 10
-    
+
     #####
     ##### Test default particle
     #####
-    
-    xs = on_architecture(arch, 0.6*ones(P))
-    ys = on_architecture(arch, 0.58*ones(P))
-    zs = on_architecture(arch, 0.8*ones(P))
+
+    xs = on_architecture(arch, 0.6 * ones(P))
+    ys = on_architecture(arch, 0.58 * ones(P))
+    zs = on_architecture(arch, 0.8 * ones(P))
 
     particles = LagrangianParticles(x=xs, y=ys, z=zs)
     @test particles isa LagrangianParticles
@@ -77,21 +67,21 @@ function run_simple_particle_tracking_tests(arch, timestepper; vertically_stretc
     ##### Test Boundary restitution
     #####
 
-    initial_z    = CUDA.@allowscalar grid.zᵃᵃᶜ[grid.Nz - 1]
-    top_boundary = CUDA.@allowscalar grid.zᵃᵃᶠ[grid.Nz + 1]
+    initial_z = CUDA.@allowscalar grid.zᵃᵃᶜ[grid.Nz-1]
+    top_boundary = CUDA.@allowscalar grid.zᵃᵃᶠ[grid.Nz+1]
 
     x, y, z = on_architecture.(Ref(arch), ([0.0], [0.0], [initial_z]))
 
     particles = LagrangianParticles(; x, y, z)
-    u, v, w   = VelocityFields(grid)
+    u, v, w = VelocityFields(grid)
 
     Δt = 0.01
-    interior(w, :, :, grid.Nz)   .= (0.1 + top_boundary - initial_z) / Δt
-    interior(w, :, :, grid.Nz-1) .= (0.2 + top_boundary - initial_z) / Δt
+    interior(w, :, :, grid.Nz) .= (0.1 + top_boundary - initial_z) / Δt
+    interior(w, :, :, grid.Nz - 1) .= (0.2 + top_boundary - initial_z) / Δt
 
     velocities = PrescribedVelocityFields(; u, v, w)
 
-    model = HydrostaticFreeSurfaceModel(; grid, particles, velocities, buoyancy=nothing, tracers = ())
+    model = HydrostaticFreeSurfaceModel(; grid, particles, velocities, buoyancy=nothing, tracers=())
 
     time_step!(model, Δt)
 
@@ -101,7 +91,7 @@ function run_simple_particle_tracking_tests(arch, timestepper; vertically_stretc
     #####
     ##### Test custom particle "SpeedTrackingParticle"
     #####
-    
+
     xs = on_architecture(arch, zeros(P))
     ys = on_architecture(arch, zeros(P))
     zs = on_architecture(arch, 0.5 * ones(P))
@@ -115,7 +105,7 @@ function run_simple_particle_tracking_tests(arch, timestepper; vertically_stretc
 
     velocities = VelocityFields(grid)
     u, v, w = velocities
-    speed = Field(√(u*u + v*v + w*w))
+    speed = Field(√(u * u + v * v + w * w))
     tracked_fields = merge(velocities, (; s=speed))
 
     # applying v component of advection with background field to ensure it is included
@@ -128,7 +118,7 @@ function run_simple_particle_tracking_tests(arch, timestepper; vertically_stretc
 
     model = NonhydrostaticModel(; grid, timestepper,
                                   velocities, particles=lagrangian_particles,
-                                  background_fields = (v = background_v, ))
+                                  background_fields=(v=background_v,))
 
     set!(model, u=1)
 
@@ -144,13 +134,13 @@ function run_simple_particle_tracking_tests(arch, timestepper; vertically_stretc
         NetCDFOutputWriter(model, model.particles, filename=nc_filepath, schedule=IterationInterval(1))
 
     sim.output_writers[:checkpointer] = Checkpointer(model, schedule=IterationInterval(1),
-                                                     dir = ".", prefix = "particles_checkpoint")
+                                                    dir=".", prefix="particles_checkpoint")
 
     rm(jld2_filepath)
     rm(nc_filepath)
     rm("particles_checkpoint_iteration1.jld2")
 
-    sim, jld2_filepath, nc_filepath = particle_tracking_simulation(; grid, particles=lagrangian_particles, timestepper, velocities)    
+    sim, jld2_filepath, nc_filepath = particle_tracking_simulation(; grid, particles=lagrangian_particles, timestepper, velocities)
     model = sim.model
     run!(sim)
 
@@ -274,12 +264,19 @@ function run_simple_particle_tracking_tests(arch, timestepper; vertically_stretc
     return nothing
 end
 
-@testset "Lagrangian particle tracking" begin
-    for arch in archs, timestepper in (:QuasiAdamsBashforth2, :RungeKutta3)
-        @info "  Testing uniform grid Lagrangian particle tracking [$(typeof(arch)), $timestepper]..."
-        run_simple_particle_tracking_tests(arch, timestepper; vertically_stretched=false)
+lagrangian_particle_test_grid(arch, ::Periodic, z) =
+    RectilinearGrid(arch; topology=(Periodic, Periodic, Bounded), size=(5, 5, 5), x=(-1, 1), y=(-1, 1), z)
+lagrangian_particle_test_grid(arch, ::Flat, z) =
+    RectilinearGrid(arch; topology=(Periodic, Flat, Bounded), size=(5, 5), x=(-1, 1), z)
 
-        @info "  Testing stretched grid Lagrangian particle tracking [$(typeof(arch)), $timestepper]..."
-        run_simple_particle_tracking_tests(arch, timestepper; vertically_stretched=true)
+@testset "Lagrangian particle tracking" begin
+    timesteppers = (:QuasiAdamsBashforth2, :RungeKutta3)
+    y_topologies = (Periodic(), Flat())
+    vertical_grids = (uniform=(-1, 1), stretched=[-1, -0.5, 0.0, 0.4, 0.7, 1])
+
+    for arch in archs, timestepper in timesteppers, y_topo in y_topologies, (z_grid_type, z) in pairs(vertical_grids)
+        @info "  Testing Lagrangian particle tracking [$(typeof(arch)), $timestepper] with y $(typeof(y_topo)) on vertically $z_grid_type grid ..."
+        grid = lagrangian_particle_test_grid(arch, y_topo, z)
+        run_simple_particle_tracking_tests(grid, timestepper)
     end
 end