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
-