Lagrangian particles on Flat topology

[Jamie-Hilditch]

I was trying to add some Lagrangian particles to a 2D simulation I was running but Lagrangian particle advection currently assumes that the grid is bounded or periodic in all dimensions. An MWE (on Julia 1.10.1+0.x64.w64.mingw32 and Oceananigans version 0.90.12) that reproduces the issue is:

using Oceananigans
NX, NZ = 16, 16
grid = RectilinearGrid(size=(NX,NZ), x=(0,1), z=(0,1), topology=(Periodic, Flat, Bounded))
# initialise some particles at the cell centres
x₀ = reshape(xnodes(grid, Center()) * ones(NZ)', NX * NZ)
y₀ = zeros(NX * NZ)
z₀ = reshape(ones(NX) * znodes(grid, Center())', NX * NZ)
model = NonhydrostaticModel(; grid, particles=LagrangianParticles(x=x₀, y=y₀, z=z₀))
simulation = Simulation(model; Δt=1, stop_iteration=5)
run!(simulation)

with output

[ Info: Initializing simulation...
[ Info:     ... simulation initialization complete (115.810 ms)
[ Info: Executing initial time step...
ERROR: LoadError: MethodError: no method matching unsafe_trunc(::Type{Int64}, ::Nothing)

Closest candidates are:
  unsafe_trunc(::Type{Int64}, ::Union{Float16, Float32, Float64})
   @ Base float.jl:336
  unsafe_trunc(::Type{<:Integer}, ::BFloat16s.BFloat16)
   @ BFloat16s C:\Users\hildi\.julia\packages\BFloat16s\u3WQc\src\bfloat16.jl:288
  unsafe_trunc(::Type{T}, ::BigFloat) where T<:Integer
   @ Base mpfr.jl:358
  ...

Stacktrace:
  [1] advect_particle
    @ C:\Users\hildi\.julia\packages\Oceananigans\kBe5X\src\Models\LagrangianParticleTracking\lagrangian_particle_advection.jl:81 [inlined]

...

The offending line is

Oceananigans.jl/src/Models/LagrangianParticleTracking/lagrangian_particle_advection.jl

Line 81 in ce4faba

j = Base.unsafe_trunc(Int, j)

which fails because

Oceananigans.jl/src/Models/LagrangianParticleTracking/lagrangian_particle_advection.jl

Line 79 in ce4faba

i, j, k = fractional_indices(X, grid, c, c, c)

returns j = nothing when y is Flat.

However, I think that in general advect_particle is not written with Flat topologies in mind. For example, there are no boundary conditions implemented for the Flat topology.

The ability to do 2D (and 1D) Lagrangian particles would be quite useful. Although, it's not obvious what it means to advect a particle in the Flat direction. Skipping that direction entirely or simply accumulating the displacement in that direction (as I do below) could both be valid in some contexts. (For my original use case I don't care what happens in y.)

Finally, for this particular simple example, propagating j = nothing through unsafe_trunc and defining some boundary conditions was sufficient to make the example run.

using Oceananigans

Base.unsafe_trunc(::T, ::Nothing) where {T} = nothing
Oceananigans.Models.LagrangianParticleTracking.enforce_boundary_conditions(::Flat,  x, xᴸ, xᴿ, Cʳ) = x

NX, NZ = 16, 16
grid = RectilinearGrid(size=(NX,NZ), x=(0,1), z=(0,1), topology=(Periodic, Flat, Bounded))
# initialise some particles at the cell centres
x₀ = reshape(xnodes(grid, Center()) * ones(NZ)', NX * NZ)
y₀ = zeros(NX * NZ)
z₀ = reshape(ones(NX) * znodes(grid, Center())', NX * NZ)
model = NonhydrostaticModel(; grid, particles=LagrangianParticles(x=x₀, y=y₀, z=z₀))
simulation = Simulation(model; Δt=1, stop_iteration=5)
run!(simulation)

I'm not sure whether propagating j = nothing forward is a sensible or desirable thing to do, or if it even works in other cases, but I thought that I'd point out that it works here.

[glwagner]

Thank you for opening this @Jamie-Hilditch ! Of course we want to advect particles in Flat geometries :-D

I think this is a bug because apparently there is some attempt to support Flat, since

Oceananigans.jl/src/Models/LagrangianParticleTracking/lagrangian_particle_advection.jl

Line 76 in ce4faba

X = flattened_node((x, y, z), grid)

calls these functions

Oceananigans.jl/src/Models/LagrangianParticleTracking/LagrangianParticleTracking.jl

Lines 121 to 128 in ce4faba

	# To support interpolation on Flat grids
	@inline flattened_node((x, y, z), grid) = (x, y, z)
	@inline flattened_node((x, y, z), grid::XFlatGrid) = (y, z)
	@inline flattened_node((x, y, z), grid::YFlatGrid) = (x, z)
	@inline flattened_node((x, y, z), grid::ZFlatGrid) = (x, y)
	@inline flattened_node((x, y, z), grid::YZFlatGrid) = tuple(x)
	@inline flattened_node((x, y, z), grid::XZFlatGrid) = tuple(y)
	@inline flattened_node((x, y, z), grid::XYFlatGrid) = tuple(z)

Something must have broken when we changed fractional_indices...

We may want

# Deal with `nothing` indices, which are output by `fractional_indices` in `Flat` domains
@inline truncate_index(::Nothing) = 1
@inline truncate_index(i) = Base.unsafe_trunc(Int, i)

and then

fi, fj, fk = fractional_indices(X, grid, c, c, c)
i = truncate_index(fi)
j = truncate_index(fj)
k = truncate_index(fk)

Does this solve the problem? (I'm also suggesting a small change in notation to clarify the difference between a "fractional index" and a truncated index.)

@Jamie-Hilditch if you like, you can contribute this change --- or I can add it. Up to you, let us know.

[Jamie-Hilditch]

Yes, this change plus adding the code to "enforce" the boundary conditions appears to fix the bug. The fractional indices are truncated in two places - the advection code and in bounce_immersed_particle. I've fixed them both but not done any tests with an immersed boundary. I'll open a PR and you can take a look.

[glwagner]

@Jamie-Hilditch great!

We could also change the name of enforce_boundary_conditions or whatever it is now to make it more accurate 😄