Fix index calculation for Lagrangian particles in periodic directions

[xkykai]

Closes #3415

[xkykai]

Here's a test script:

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

Random.seed!(123)

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

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

x_particle = collect(0:0.25:1.5)
y_particle = collect(0:0.25:1.5)
z_particle = collect(fill(-0.5, length(x_particle)))

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)

In the test script, the domain is initialized to be Periodic, Bounded, Bounded, so particles should be shifted if $x \geq 1$, bounced if $y > 1$, $z<-1$.

Here's the output of the script:

[ Info: Initializing simulation...
i: 0, t: 0 seconds, wall time: 5.450 seconds, 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, 1.25, 1.5], y(particle): [0.0, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5], z(particle): [-0.5, -0.5, -0.5, -0.5, -0.5, -0.5, -0.5]
[ Info:     ... simulation initialization complete (3.384 seconds)
[ Info: Executing initial time step...
[ Info:     ... initial time step complete (7.395 seconds).
i: 1, t: 100 ms, wall time: 7.414 seconds, 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, 0.25, 0.5], y(particle): [0.0, 0.25, 0.5, 0.75, 1.0, 0.75, 0.5], z(particle): [-0.5, -0.5, -0.5, -0.5, -0.5, -0.5, -0.5]
[ Info: Simulation is stopping after running for 10.934 seconds.
[ Info: Model iteration 2 equals or exceeds stop iteration 2.
i: 2, t: 200 ms, wall time: 3.597 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, 0.25, 0.5], y(particle): [0.0, 0.25, 0.5, 0.75, 1.0, 0.75, 0.5], z(particle): [-0.5, -0.5, -0.5, -0.5, -0.5, -0.5, -0.5]

We see that the particles that are outside of the domain in $x$ is shifted and particles that are outside of the domain in $y$ are bounced in the expected manner.

[jagoosw]

Oh I see, that looks good. My confusion was because length(Face(), Periodic(), 2) = 2 whereas length(Face(), Bounded(), 2) = 3.

I guess instead of using tx/ty/etc. here it could be change to length(Face(), Bounded(), N) always.

[Yixiao-Zhang]

I suggest defining a new function for the rightmost cell index instead of using length. I do not think that length is relevant here.

interface_index_rightmost(::Bounded, N) = N + 1
interface_index_rightmost(::Periodic, N) = N + 1
interface_index_rightmost(::Flat, N) = N

Then

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

I am not sure whether interface_index_rightmost or something similar has been defined in grid_utils.jl or not. Probably it is a good thing to have in grid_utils.jl.

[glwagner]

I suggest defining a new function for the rightmost cell index instead of using length. I do not think that length is relevant here.

interface_index_rightmost(::Bounded, N) = N + 1
interface_index_rightmost(::Periodic, N) = N + 1
interface_index_rightmost(::Flat, N) = N

Then

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

I am not sure whether interface_index_rightmost or something similar has been defined in grid_utils.jl or not. Probably it is a good thing to have in grid_utils.jl.

That is indeed different from length. My only comment is to use English language conventions for names like everything else, so it should be called rightmost_interface_index (perhaps last_interface_index?).

[xkykai]

This seems to be ready to merge? @glwagner @simone-silvestri @jagoosw

[glwagner]

Does it resolve #3415?

[xkykai]

This resolves issue https://github.com/CliMA/Oceananigans.jl/issues/3415. See this script:

using Oceananigans
using StructArrays
using Printf

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

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

x_particle = collect(0:0.25:1.5)
y_particle = collect(0:0.25:1.5)
z_particle = collect(fill(-0.5, length(x_particle)))

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 outputs

[ Info: Initializing simulation...
i: 0, t: 0 seconds, wall time: 4.726 seconds, 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, 1.25, 1.5], y(particle): [0.0, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5], z(particle): [-0.5, -0.5, -0.5, -0.5, -0.5, -0.5, -0.5]
[ Info:     ... simulation initialization complete (3.406 seconds)
[ Info: Executing initial time step...
[ Info:     ... initial time step complete (6.270 seconds).
i: 1, t: 100 ms, wall time: 6.286 seconds, 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, 0.75, 0.5], y(particle): [0.0, 0.25, 0.5, 0.75, 1.0, 0.25, 0.5], z(particle): [-0.5, -0.5, -0.5, -0.5, -0.5, -0.5, -0.5]
[ Info: Simulation is stopping after running for 9.706 seconds.
[ Info: Model iteration 2 equals or exceeds stop iteration 2.
i: 2, t: 200 ms, wall time: 2.858 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, 0.75, 0.5], y(particle): [0.0, 0.25, 0.5, 0.75, 1.0, 0.25, 0.5], z(particle): [-0.5, -0.5, -0.5, -0.5, -0.5, -0.5, -0.5]

Particles are bounced and shifted appropriately, and it also works in Flat dimensions as well. This is now ready to be merged.

[glwagner]

Can you change the first post so merging this issue closes #3415 (and makes the connection clear)

[xkykai]

Is what I've done linking #3415 and #3416 correct?

[glwagner]

Is what I've done linking #3415 and #3416 correct?

Yes, I think "closes" or "resolves" is good. Maybe there are other words, there might be documentation. Thanks.

[navidcy]

I'm merging this