Using LagrangianParticles with Flat topology hits wrong function dispatch when computing fractional indices

[xkykai]

When using LagrangianParticles in a grid with one direction having Flat topology, the code fails because the fractional index function hits the wrong line of code. Here is a MWE:

using Oceananigans
using StructArrays
using Printf

grid = RectilinearGrid(CPU(), Float64,
                       size = (2, 2),
                       halo = (5, 5),
                       x = (0, 1),
                       y = (0, 1),
                       topology = (Periodic, Bounded, Flat))

#%%
struct SimpleParticle{X}
    x :: X
    y :: X
    z :: X
end

x_particle = [0.5]
y_particle = [0.5]
z_particle = [0.5]

particles = StructArray{SimpleParticle}((x_particle, y_particle, z_particle))

lagrangian_particles = LagrangianParticles(particles)

#%%
model = NonhydrostaticModel(; 
            grid = grid,
            timestepper = :RungeKutta3,
            advection = WENO(order=9),
            particles = lagrangian_particles
            )

u, v, w = model.velocities

simulation = Simulation(model, Δt=0.1, stop_iteration=2)

wall_clock = [time_ns()]

function print_progress(sim)
    @printf("i: %d, t: %s, wall time: %s, max(u): (%6.3e, %6.3e, %6.3e) m/s, next Δt: %s\n",
            sim.model.clock.iteration,
            prettytime(sim.model.clock.time),
            prettytime(1e-9 * (time_ns() - wall_clock[1])),
            maximum(abs, sim.model.velocities.u),
            maximum(abs, sim.model.velocities.v),
            maximum(abs, sim.model.velocities.w),
            prettytime(sim.Δt))
    @info "x(particle): $(round.(lagrangian_particles.properties.x, digits=2)), y(particle): $(round.(lagrangian_particles.properties.y, digits=2)), z(particle): $(round.(lagrangian_particles.properties.z, digits=2))\n"

    wall_clock[1] = time_ns()

    return nothing
end

simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(1))

run!(simulation)

which gives an error output of

ERROR: BoundsError: attempt to access Tuple{Float64, Float64} at index [3]
Stacktrace:
  [1] getindex(t::Tuple, i::Int64)
    @ Base .\tuple.jl:31
  [2] fractional_z_index
    @ c:\Users\xinle\MIT\lagrangian_particles_rightmost_cell\Oceananigans.jl\src\Fields\interpolate.jl:110 [inlined]
  [3] fractional_indices
    @ c:\Users\xinle\MIT\lagrangian_particles_rightmost_cell\Oceananigans.jl\src\Fields\interpolate.jl:133 [inlined]
  [4] advect_particle
    @ c:\Users\xinle\MIT\lagrangian_particles_rightmost_cell\Oceananigans.jl\src\Models\LagrangianParticleTracking\lagrangian_particle_advection.jl:75 [inlined]
  [5] macro expansion
    @ c:\Users\xinle\MIT\lagrangian_particles_rightmost_cell\Oceananigans.jl\src\Models\LagrangianParticleTracking\lagrangian_particle_advection.jl:145 [inlined]
  [6] cpu__advect_particles!
    @ C:\Users\xinle\.julia\packages\KernelAbstractions\WoCk1\src\macros.jl:276 [inlined]
  [7] cpu__advect_particles!(__ctx__::KernelAbstractions.CompilerMetadata{…}, particles::StructVector{…}, restitution::Float64, grid::RectilinearGrid{…}, Δt::Float64, velocities::@NamedTuple{…})
    @ Oceananigans.Models.LagrangianParticleTracking .\none:0
  [8] __thread_run(tid::Int64, len::Int64, rem::Int64, obj::KernelAbstractions.Kernel{…}, ndrange::Nothing, iterspace::KernelAbstractions.NDIteration.NDRange{…}, args::Tuple{…}, dynamic::KernelAbstractions.NDIteration.NoDynamicCheck)
    @ KernelAbstractions C:\Users\xinle\.julia\packages\KernelAbstractions\WoCk1\src\cpu.jl:115
  [9] __run(obj::KernelAbstractions.Kernel{…}, ndrange::Nothing, iterspace::KernelAbstractions.NDIteration.NDRange{…}, args::Tuple{…}, dynamic::KernelAbstractions.NDIteration.NoDynamicCheck, static_threads::Bool)
    @ KernelAbstractions C:\Users\xinle\.julia\packages\KernelAbstractions\WoCk1\src\cpu.jl:82
 [10] (::KernelAbstractions.Kernel{…})(::StructVector{…}, ::Vararg{…}; ndrange::Nothing, workgroupsize::Nothing)
    @ KernelAbstractions C:\Users\xinle\.julia\packages\KernelAbstractions\WoCk1\src\cpu.jl:44
 [11] (::KernelAbstractions.Kernel{…})(::StructVector{…}, ::Vararg{…})
    @ KernelAbstractions C:\Users\xinle\.julia\packages\KernelAbstractions\WoCk1\src\cpu.jl:37
 [12] advect_lagrangian_particles!(particles::LagrangianParticles{…}, model::NonhydrostaticModel{…}, Δt::Float64)
    @ Oceananigans.Models.LagrangianParticleTracking c:\Users\xinle\MIT\lagrangian_particles_rightmost_cell\Oceananigans.jl\src\Models\LagrangianParticleTracking\lagrangian_particle_advection.jl:161
 [13] step_lagrangian_particles!
    @ c:\Users\xinle\MIT\lagrangian_particles_rightmost_cell\Oceananigans.jl\src\Models\LagrangianParticleTracking\LagrangianParticleTracking.jl:131 [inlined]
 [14] step_lagrangian_particles!
    @ c:\Users\xinle\MIT\lagrangian_particles_rightmost_cell\Oceananigans.jl\src\Models\NonhydrostaticModels\NonhydrostaticModels.jl:76 [inlined]
 [15] time_step!(model::NonhydrostaticModel{…}, Δt::Float64; callbacks::Tuple{}, compute_tendencies::Bool)
    @ Oceananigans.TimeSteppers c:\Users\xinle\MIT\lagrangian_particles_rightmost_cell\Oceananigans.jl\src\TimeSteppers\runge_kutta_3.jl:110
 [16] time_step!
    @ c:\Users\xinle\MIT\lagrangian_particles_rightmost_cell\Oceananigans.jl\src\TimeSteppers\runge_kutta_3.jl:81 [inlined]
 [17] time_step!(sim::Simulation{…})
    @ Oceananigans.Simulations c:\Users\xinle\MIT\lagrangian_particles_rightmost_cell\Oceananigans.jl\src\Simulations\run.jl:122
 [18] run!(sim::Simulation{…}; pickup::Bool)
    @ Oceananigans.Simulations c:\Users\xinle\MIT\lagrangian_particles_rightmost_cell\Oceananigans.jl\src\Simulations\run.jl:97
 [19] run!(sim::Simulation{…})
    @ Oceananigans.Simulations c:\Users\xinle\MIT\lagrangian_particles_rightmost_cell\Oceananigans.jl\src\Simulations\run.jl:85
 [20] top-level scope
    @ c:\Users\xinle\MIT\lagrangian_particles_rightmost_cell\Oceananigans.jl\lagrangian_particle_MWE.jl:59

The error indicates that fractional_z_index function hits

Oceananigans.jl/src/Fields/interpolate.jl

Line 133 in d4bcc09

@inline function fractional_z_index(z::FT, locs, grid::ZRegGrid) where FT

instead of

Oceananigans.jl/src/Fields/interpolate.jl

Line 129 in d4bcc09

@inline fractional_z_index(z, locs, grid::ZFlatGrid) = zero(grid)

, which is the intended function dispatch.

Note: doing something like

struct SimpleParticle{X}
    x :: X
    y :: X
end

is not supported, but perhaps this is a separate discussion.

[xkykai]

Duplicate of #3550