Fix bug in advect_particle (#3923)

* this should work

* remove all the @show

* restart tests

* will this work?

* test more than one position

* bugfix

* this should work

* remove unused truncate

* new interpolate

* I think this works

* try option 2

* no need for flip anymore

* fix gpu tests

* change comment and restore docstring

* fix the interpolate

* add (nothing nothing nothing) interpolate

* do not reuse names

* bugfix

* change ᵃ to ᶜ

* back to previous code

* add the :a in the metrics

Author: simone-silvestri

src/Fields/interpolate.jl

@@ -1,4 +1,4 @@
-using Oceananigans.Grids: topology, node, _node,
+using Oceananigans.Grids: topology, node, _node, φnode, λnode,
                           xspacings, yspacings, zspacings, λspacings, φspacings,
                           XFlatGrid, YFlatGrid, ZFlatGrid,
                           XYFlatGrid, YZFlatGrid, XZFlatGrid,
@@ -7,6 +7,8 @@ using Oceananigans.Grids: topology, node, _node,
                           ZRegOrthogonalSphericalShellGrid,
                           RectilinearGrid, LatitudeLongitudeGrid
 
+using Oceananigans.Operators: Δx, Δy, Δz
+
 using Oceananigans.Architectures: child_architecture
 
 # GPU-compatile middle point calculation
@@ -66,64 +68,64 @@ end
 
 @inline function fractional_x_index(x, locs, grid::XRegularRG)
     x₀ = xnode(1, 1, 1, grid, locs...)
-    Δx = @inbounds first(xspacings(grid, locs...))
+    dx = Δx(1, 1, 1, grid, locs...)
     FT = eltype(grid)
-    return convert(FT, (x - x₀) / Δx)
+    return convert(FT, (x - x₀) / dx) + 1 # 1 - based indexing 
 end
 
 @inline function fractional_x_index(λ, locs, grid::XRegularLLG)
     λ₀ = λnode(1, 1, 1, grid, locs...)
-    Δλ = @inbounds first(λspacings(grid, locs...))
+    λ₁ = λnode(2, 1, 1, grid, locs...)
     FT = eltype(grid)
-    return convert(FT, (λ - λ₀) / Δλ)
+    return convert(FT, (λ - λ₀) / (λ₁ - λ₀)) + 1 # 1 - based indexing 
 end
 
 @inline function fractional_x_index(x, locs, grid::RectilinearGrid)
     loc = @inbounds locs[1]
      Tx = topology(grid, 1)()
      Nx = length(loc, Tx, grid.Nx)
      xn = xnodes(grid, locs...)
-    return fractional_index(x, xn, Nx) - 1
+    return fractional_index(x, xn, Nx) 
 end
 
 @inline function fractional_x_index(x, locs, grid::LatitudeLongitudeGrid)
     loc = @inbounds locs[1]
      Tx = topology(grid, 1)()
      Nx = length(loc, Tx, grid.Nx)
      xn = λnodes(grid, locs...)
-    return fractional_index(x, xn, Nx) - 1
+    return fractional_index(x, xn, Nx) 
 end
 
 @inline fractional_y_index(y, locs, grid::YFlatGrid) = zero(grid)
 
 @inline function fractional_y_index(y, locs, grid::YRegularRG)
     y₀ = ynode(1, 1, 1, grid, locs...)
-    Δy = @inbounds first(yspacings(grid, locs...))
+    dy = Δy(1, 1, 1, grid, locs...)
     FT = eltype(grid)
-    return convert(FT, (y - y₀) / Δy)
+    return convert(FT, (y - y₀) / dy) + 1 # 1 - based indexing 
 end
 
 @inline function fractional_y_index(φ, locs, grid::YRegularLLG)
-    φ₀ = φnode(1, 1, 1, grid, locs...)
-    Δφ = @inbounds first(φspacings(grid, locs...))
+    φ₀ = φnode(1, 1, 1, grid, locs...)    
+    φ₁ = φnode(1, 2, 1, grid, locs...)
     FT = eltype(grid)
-    return convert(FT, (φ - φ₀) / Δφ)
+    return convert(FT, (φ - φ₀) / (φ₁ - φ₀)) + 1 # 1 - based indexing 
 end
 
 @inline function fractional_y_index(y, locs, grid::RectilinearGrid)
     loc = @inbounds locs[2]
      Ty = topology(grid, 2)()
      Ny = length(loc, Ty, grid.Ny)
      yn = ynodes(grid, locs...)
-    return fractional_index(y, yn, Ny) - 1
+    return fractional_index(y, yn, Ny) 
 end
 
 @inline function fractional_y_index(y, locs, grid::LatitudeLongitudeGrid)
     loc = @inbounds locs[2]
      Ty = topology(grid, 2)()
      Ny = length(loc, Ty, grid.Ny)
      yn = φnodes(grid, locs...)
-    return fractional_index(y, yn, Ny) - 1
+    return fractional_index(y, yn, Ny)
 end
 
 @inline fractional_z_index(z, locs, grid::ZFlatGrid) = zero(grid)
@@ -132,16 +134,16 @@ ZRegGrid = Union{ZRegularRG, ZRegularLLG, ZRegOrthogonalSphericalShellGrid}
 
 @inline function fractional_z_index(z::FT, locs, grid::ZRegGrid) where FT
     z₀ = znode(1, 1, 1, grid, locs...)
-    Δz = @inbounds first(zspacings(grid, locs...))
-    return convert(FT, (z - z₀) / Δz)
+    dz = Δz(1, 1, 1, grid, locs...)
+    return convert(FT, (z - z₀) / dz) + 1 # 1 - based indexing 
 end
 
 @inline function fractional_z_index(z, locs, grid)
     loc = @inbounds locs[3]
      Tz = topology(grid, 3)()
      Nz = length(loc, Tz, grid.Nz)
      zn = znodes(grid, loc)
-    return fractional_index(z, zn, Nz) - 1
+    return fractional_index(z, zn, Nz) 
 end
 
 """
@@ -210,22 +212,7 @@ 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)
-
+@inline _fractional_indices(at_node, grid, ::Nothing, ::Nothing, ::Nothing) = (nothing, nothing, nothing)
 
 """
     interpolate(at_node, from_field, from_loc, from_grid)
@@ -260,15 +247,12 @@ right of `i`, and `ξ` is the fractional distance between `i` and the
 left bound `i⁻`, such that `ξ ∈ [0, 1)`.
 """
 @inline function interpolator(fractional_idx)
-    # We use mod and trunc as CUDA.modf is not defined.
     # For why we use Base.unsafe_trunc instead of trunc see:
     # https://github.com/CliMA/Oceananigans.jl/issues/828
     # https://github.com/CliMA/Oceananigans.jl/pull/997
 
     i⁻ = Base.unsafe_trunc(Int, fractional_idx)
-    i⁻ = Int(i⁻ + 1) # convert to "proper" integer?
-    shift = Int(sign(fractional_idx))
-    i⁺ = i⁻ + shift
+    i⁺ = i⁻ + 1
     ξ = mod(fractional_idx, 1)
 
     return (i⁻, i⁺, ξ)

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, truncate_fractional_indices
+using Oceananigans.Fields: fractional_indices
 using Oceananigans.TimeSteppers: AbstractLagrangianParticles
 using Oceananigans.Utils: prettysummary, launch!, SumOfArrays
 

src/Models/LagrangianParticleTracking/lagrangian_particle_advection.jl

@@ -1,5 +1,6 @@
 using Oceananigans.Utils: instantiate, KernelParameters
 using Oceananigans.Models: total_velocities
+using Oceananigans.Fields: interpolator
 
 #####
 ##### Boundary conditions for Lagrangian particles
@@ -54,30 +55,35 @@ bouncing the particle off the immersed boundary with a coefficient or `restituti
 @inline function bounce_immersed_particle((x, y, z), ibg, restitution, previous_particle_indices)
     X = flattened_node((x, y, z), ibg)
 
-    # Determine current particle cell
-    fi, fj, fk = fractional_indices(X, ibg.underlying_grid, c, c, c)
-    i, j, k = truncate_fractional_indices(fi, fj, fk)
+    # Determine current particle cell from the interfaces
+    fi, fj, fk = fractional_indices(X, ibg.underlying_grid, f, f, f)
+    
+    i, i⁺, _ = interpolator(fi)
+    j, j⁺, _ = interpolator(fj)
+    k, k⁺, _ = interpolator(fk)
 
     # Determine whether particle was _previously_ in a non-immersed cell
     i⁻, j⁻, k⁻ = previous_particle_indices
 
-    # Left bounds of the previous cell
-    xᴿ = ξnode(i⁻ + 1, j⁻ + 1, k⁻ + 1, ibg, f, f, f)
-    yᴿ = ηnode(i⁻ + 1, j⁻ + 1, k⁻ + 1, ibg, f, f, f)
-    zᴿ = rnode(i⁻ + 1, j⁻ + 1, k⁻ + 1, ibg, f, f, f)
+    tx, ty, tz = map(immersed_boundary_topology, topology(ibg))
 
     # Right bounds of the previous cell
+    xᴿ = ξnode(i⁺, j,  k, ibg, f, f, f)
+    yᴿ = ηnode(i,  j⁺, k, ibg, f, f, f)
+    zᴿ = rnode(i,  j,  k⁺, ibg, f, f, f)
+
+    # Left bounds of the previous cell
     xᴸ = ξnode(i⁻, j⁻, k⁻, ibg, f, f, f)
     yᴸ = ηnode(i⁻, j⁻, k⁻, ibg, f, f, f)
     zᴸ = rnode(i⁻, j⁻, k⁻, ibg, f, f, f)
 
     Cʳ = restitution
-    tx, ty, tz = map(immersed_boundary_topology, topology(ibg))
+    
     xb⁺ = enforce_boundary_conditions(tx, x, xᴸ, xᴿ, Cʳ)
     yb⁺ = enforce_boundary_conditions(ty, y, yᴸ, yᴿ, Cʳ)
     zb⁺ = enforce_boundary_conditions(tz, z, zᴸ, zᴿ, Cʳ)
 
-    immersed = immersed_cell(i, j, k, ibg)
+    immersed = immersed_cell(i⁺, j⁺, k⁺, ibg)
     x⁺ = ifelse(immersed, xb⁺, x)
     y⁺ = ifelse(immersed, yb⁺, y)
     z⁺ = ifelse(immersed, zb⁺, z)
@@ -90,7 +96,7 @@ end
 
 Return the index of the rightmost cell interface for a grid with `topology` and `N` cells.
 """
-rightmost_interface_index(::Bounded, N) = N + 1
+rightmost_interface_index(::Bounded, N)  = N + 1
 rightmost_interface_index(::Periodic, N) = N + 1
 rightmost_interface_index(::Flat, N) = N
 
@@ -103,9 +109,12 @@ given `velocities`, time-step `Δt, and coefficient of `restitution`.
 @inline function advect_particle((x, y, z), p, restitution, grid, Δt, velocities)
     X = flattened_node((x, y, z), grid)
 
-    # Obtain current particle indices
-    fi, fj, fk = fractional_indices(X, grid, c, c, c)
-    i, j, k = truncate_fractional_indices(fi, fj, fk)
+    # Obtain current particle indices, looking at the interfaces
+    fi, fj, fk = fractional_indices(X, grid, f, f, f)
+    
+    i, i⁺, _ = interpolator(fi)
+    j, j⁺, _ = interpolator(fj)
+    k, k⁺, _ = interpolator(fk)
 
     current_particle_indices = (i, j, k)
 
@@ -137,15 +146,16 @@ given `velocities`, time-step `Δt, and coefficient of `restitution`.
     yᴸ = ηnode(i, 1, k, grid, f, f, f)
     zᴸ = rnode(i, j, 1, grid, f, f, f)
 
-    xᴿ = ξnode(iᴿ, j, k, grid, f, f, f)
-    yᴿ = ηnode(i, jᴿ, k, grid, f, f, f)
-    zᴿ = rnode(i, j, kᴿ, grid, f, f, f)
+    xᴿ = ξnode(iᴿ, j,  k,  grid, f, f, f)
+    yᴿ = ηnode(i,  jᴿ, k,  grid, f, f, f)
+    zᴿ = rnode(i,  j,  kᴿ, grid, f, f, f)
 
     # Enforce boundary conditions for particles.
     Cʳ = restitution
     x⁺ = enforce_boundary_conditions(tx, x⁺, xᴸ, xᴿ, Cʳ)
     y⁺ = enforce_boundary_conditions(ty, y⁺, yᴸ, yᴿ, Cʳ)
     z⁺ = enforce_boundary_conditions(tz, z⁺, zᴸ, zᴿ, Cʳ)
+
     if grid isa ImmersedBoundaryGrid
         previous_particle_indices = current_particle_indices # particle has been advected
         (x⁺, y⁺, z⁺) = bounce_immersed_particle((x⁺, y⁺, z⁺), grid, Cʳ, previous_particle_indices)

src/Operators/spacings_and_areas_and_volumes.jl

@@ -36,7 +36,7 @@ The operators in this file fall into three categories:
 #####
 
 # Convenience Functions for all grids
-for LX in (:ᶜ, :ᶠ), LY in (:ᶜ, :ᶠ)
+for LX in (:ᶜ, :ᶠ, :ᵃ), LY in (:ᶜ, :ᶠ, :ᵃ)
 
     x_spacing_1D = Symbol(:Δx, LX, :ᵃ, :ᵃ)
     x_spacing_2D = Symbol(:Δx, LX, LY, :ᵃ)
@@ -197,7 +197,7 @@ end
 #####
 #####
 
-for LX in (:ᶜ, :ᶠ), LY in (:ᶜ, :ᶠ)
+for LX in (:ᶜ, :ᶠ, :ᵃ), LY in (:ᶜ, :ᶠ, :ᵃ)
 
     x_spacing_2D = Symbol(:Δx, LX, LY, :ᵃ)
     y_spacing_2D = Symbol(:Δy, LX, LY, :ᵃ)

test/test_field.jl

@@ -6,6 +6,11 @@ using Oceananigans.Grids: total_length
 using Oceananigans.Fields: ReducedField, has_velocities
 using Oceananigans.Fields: VelocityFields, TracerFields, interpolate, interpolate!
 using Oceananigans.Fields: reduced_location
+using Oceananigans.Fields: fractional_indices, interpolator
+using Oceananigans.Grids: ξnode, ηnode, rnode
+
+using Random
+using CUDA: @allowscalar
 
 """
     correct_field_size(grid, FieldType, Tx, Ty, Tz)
@@ -404,6 +409,46 @@ end
         end
     end
 
+    @testset "Unit interpolation" begin
+        for arch in archs
+            hu = (-1, 1)
+            hs = range(-1, 1, length=21)
+            zu = (-100, 0)
+            zs = range(-100, 0, length=33)
+
+            for latitude in (hu, hs), longitude in (hu, hs), z in (zu, zs), loc in (Center(), Face())
+                @info "    Testing interpolation for $(latitude) latitude and longitude, $(z) z on $(typeof(loc))s..."
+                grid = LatitudeLongitudeGrid(arch; size = (20, 20, 32), longitude, latitude, z, halo = (5, 5, 5))
+            
+                # Test random positions, 
+                # set seed for reproducibility
+                Random.seed!(1234)
+                Xs = [(2rand()-1, 2rand()-1, -100rand()) for p in 1:20]
+
+                for X in Xs
+                    (x, y, z)  = X 
+                    fi, fj, fk = @allowscalar fractional_indices(X, grid, loc, loc, loc)
+
+                    i⁻, i⁺, _ = interpolator(fi)
+                    j⁻, j⁺, _ = interpolator(fj)
+                    k⁻, k⁺, _ = interpolator(fk)
+
+                    x⁻ = @allowscalar ξnode(i⁻, j⁻, k⁻, grid, loc, loc, loc)
+                    y⁻ = @allowscalar ηnode(i⁻, j⁻, k⁻, grid, loc, loc, loc)
+                    z⁻ = @allowscalar rnode(i⁻, j⁻, k⁻, grid, loc, loc, loc)
+
+                    x⁺ = @allowscalar ξnode(i⁺, j⁺, k⁺, grid, loc, loc, loc)
+                    y⁺ = @allowscalar ηnode(i⁺, j⁺, k⁺, grid, loc, loc, loc)
+                    z⁺ = @allowscalar rnode(i⁺, j⁺, k⁺, grid, loc, loc, loc)
+
+                    @test x⁻ ≤ x ≤ x⁺
+                    @test y⁻ ≤ y ≤ y⁺
+                    @test z⁻ ≤ z ≤ z⁺
+                end 
+            end
+        end
+    end
+
     @testset "Field interpolation" begin
         @info "  Testing field interpolation..."