`LagrangianParticles` get moved at the right of `Periodic` topology when it shouldn't be

[xkykai]

I noticed that in the Periodic directions, particles that have coordinates within the last cell at the end of the domain get moved into the first cell when I think it should not be. Here's a MWE:

using Oceananigans
using Oceananigans.Units
using StructArrays
using Printf
using Random
using Statistics

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

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

n_particles = 5
x_particle = collect(range(0, stop=1, length=n_particles))
y_particle = collect(range(0, stop=1, length=n_particles))
z_particle = collect(fill(-0.5, n_particles))

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.1seconds, 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)

Here's the output log of the script:

[ Info: Initializing simulation...
i: 0, t: 0 seconds, wall time: 74.284 ms, max(u): (0.000e+00, 0.000e+00, 0.000e+00) m/s, next Δt: 100 ms
[ Info: x(particle): [0.0, 0.25, 0.5, 0.75, 1.0], y(particle): [0.0, 0.25, 0.5, 0.75, 1.0], z(particle): [-0.5, -0.5, -0.5, -0.5, -0.5]
[ Info:     ... simulation initialization complete (72.159 ms)
[ Info: Executing initial time step...
[ Info:     ... initial time step complete (92.717 ms).
i: 1, t: 100 ms, wall time: 96.815 ms, max(u): (0.000e+00, 0.000e+00, 0.000e+00) m/s, next Δt: 100 ms
[ Info: x(particle): [0.0, 0.25, 0.5, 0.25, 0.5], y(particle): [0.0, 0.25, 0.5, 0.25, 0.5], z(particle): [-0.5, -0.5, -0.5, -0.5, -0.5]
[ Info: Simulation is stopping after running for 171.290 ms.
[ Info: Model iteration 2 equals or exceeds stop iteration 2.
i: 2, t: 200 ms, wall time: 2.341 ms, max(u): (0.000e+00, 0.000e+00, 0.000e+00) m/s, next Δt: 100 ms
[ Info: x(particle): [0.0, 0.25, 0.5, 0.25, 0.5], y(particle): [0.0, 0.25, 0.5, 0.25, 0.5], z(particle): [-0.5, -0.5, -0.5, -0.5, -0.5]

We see that the particles with x- and y- coordinates larger than 0.5 get moved after the first time step when I think it shouldn't be.

This is because of the following lines: https://github.com/CliMA/Oceananigans.jl/blob/3e2650373e9c73231681535aadb5b720a830dc97/src/Models/LagrangianParticleTracking/lagrangian_particle_advection.jl#L102C4-L104C27

Instead of

    iᴿ = length(f, tx, Nx)
    jᴿ = length(f, ty, Ny)
    kᴿ = length(f, tz, Nz)

I think it should be

    iᴿ = length(f, tx, Nx) + ifelse(tx == Periodic(), 1, 0)
    jᴿ = length(f, ty, Ny) + ifelse(ty == Periodic(), 1, 0)
    kᴿ = length(f, tz, Nz) + ifelse(tz == Periodic(), 1, 0)

The change gives an output log of

[ Info: Initializing simulation...
i: 0, t: 0 seconds, wall time: 59.256 ms, max(u): (0.000e+00, 0.000e+00, 0.000e+00) m/s, next Δt: 100 ms
[ Info: x(particle): [0.0, 0.25, 0.5, 0.75, 1.0], y(particle): [0.0, 0.25, 0.5, 0.75, 1.0], z(particle): [-0.5, -0.5, -0.5, -0.5, -0.5]
[ Info:     ... simulation initialization complete (58.949 ms)
[ Info: Executing initial time step...
[ Info:     ... initial time step complete (86.845 ms).
i: 1, t: 100 ms, wall time: 91.015 ms, max(u): (0.000e+00, 0.000e+00, 0.000e+00) m/s, next Δt: 100 ms
[ Info: x(particle): [0.0, 0.25, 0.5, 0.75, 1.0], y(particle): [0.0, 0.25, 0.5, 0.75, 1.0], z(particle): [-0.5, -0.5, -0.5, -0.5, -0.5]
[ Info: Simulation is stopping after running for 152.524 ms.
[ Info: Model iteration 2 equals or exceeds stop iteration 2.
i: 2, t: 200 ms, wall time: 3.068 ms, max(u): (0.000e+00, 0.000e+00, 0.000e+00) m/s, next Δt: 100 ms
[ Info: x(particle): [0.0, 0.25, 0.5, 0.75, 1.0], y(particle): [0.0, 0.25, 0.5, 0.75, 1.0], z(particle): [-0.5, -0.5, -0.5, -0.5, -0.5]

This then ensures that we are only moving particles in the periodic direction only if they are truly outside of the rightmost cell.

[jagoosw]

This seems correct for Periodic topologies but wouldn't it be the case for Bounded too. Do particles get bounced if you put them just to the right of the N'th face in a Bounded direction?

[xkykai]

This seems correct for Periodic topologies but wouldn't it be the case for Bounded too. Do particles get bounced if you put them just to the right of the N'th face in a Bounded direction?

I've put up a similar test script in the PR: https://github.com/CliMA/Oceananigans.jl/pull/3416#issuecomment-1875732332

Let me know if you think that is the expected behaviour, and what else should I test.

[jagoosw]

Oops sorry meant to leave my comment in the PR, will change to there