Skip to content

Commit

Permalink
Fix bug when using Lagrangian particles on Flat topologies (#3550)
Browse files Browse the repository at this point in the history 
* rename output of fractional_indices to fi, fj, fk

* created truncate_fractional_indices

* use truncate_fractional_indices in advect_particles and bounce_immersed_particle

* enforce_boundary_conditions for Flat dimensions does nothing

* Improve docstring for truncate_fractional_indices

Co-authored-by: Navid C. Constantinou <[email protected]>

* rename _truncate_fractional_index to truncate_fractional_index

* Added tests for Lagrangian particles on 2D grids.
Modified run_simple_particle_tracking_tests to accept grid as input

* code alignment

* infer architecture from grid

---------

Co-authored-by: Navid C. Constantinou <[email protected]>
  • Loading branch information
@Jamie-Hilditch @navidcy
Jamie-Hilditch and navidcy committed Apr 22, 2024
1 parent d621542 commit 024b83b
Show file tree
Hide file tree
Showing 4 changed files with 88 additions and 73 deletions.
28 changes: 22 additions & 6 deletions src/Fields/interpolate.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ using Oceananigans.Architectures: child_architecture
"""
index_binary_search(val, vec, N)
Return indices `low, high` of `vec`tor for which
Return indices `low, high` of `vec`tor for which
```julia
vec[low] ≤ val && vec[high] ≥ val
Expand All @@ -29,9 +29,9 @@ Code credit: https://gist.github.com/cuongld2/8e4fed9ba44ea2b4598f90e7d5b6c612/1
low = 0
high = N - 1

while low + 1 < high
while low + 1 < high
mid = middle_point(low, high)
if @inbounds vec[mid + 1] == val
if @inbounds vec[mid + 1] == val
return (mid + 1, mid + 1)
elseif @inbounds vec[mid + 1] < val
low = mid
Expand Down Expand Up @@ -211,6 +211,23 @@ 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)


"""
interpolate(at_node, from_field, from_loc, from_grid)
Expand Down Expand Up @@ -277,7 +294,7 @@ end
k⁻, k⁺, ζ = iz

return @inbounds ϕ₁(ξ, η, ζ) * getindex(data, i⁻, j⁻, k⁻, in...) +
ϕ₂(ξ, η, ζ) * getindex(data, i⁻, j⁻, k⁺, in...) +
ϕ₂(ξ, η, ζ) * getindex(data, i⁻, j⁻, k⁺, in...) +
ϕ₃(ξ, η, ζ) * getindex(data, i⁻, j⁺, k⁻, in...) +
ϕ₄(ξ, η, ζ) * getindex(data, i⁻, j⁺, k⁺, in...) +
ϕ₅(ξ, η, ζ) * getindex(data, i⁺, j⁻, k⁻, in...) +
Expand Down Expand Up @@ -316,7 +333,7 @@ function interpolate!(to_field::Field, from_field::AbstractField)
to_arch = architecture(to_field)
from_arch = architecture(from_field)

# In case architectures are `Distributed` we
# In case architectures are `Distributed` we
# verify that the fields are on the same child architecture
to_arch = child_architecture(to_arch)
from_arch = child_architecture(from_arch)
Expand All @@ -338,4 +355,3 @@ function interpolate!(to_field::Field, from_field::AbstractField)

return nothing
end

4 changes: 2 additions & 2 deletions src/Models/LagrangianParticleTracking/LagrangianParticleTracking.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
using Oceananigans.Fields: fractional_indices, truncate_fractional_indices
using Oceananigans.TimeSteppers: AbstractLagrangianParticles
using Oceananigans.Utils: prettysummary, launch!, SumOfArrays

Expand Down Expand Up @@ -135,7 +135,7 @@ step_lagrangian_particles!(::Nothing, model, Δt) = nothing
function step_lagrangian_particles!(particles::LagrangianParticles, model, Δt)
# Update the properties of the Lagrangian particles
update_lagrangian_particle_properties!(particles, model, Δt)

# Compute dynamics
particles.dynamics(particles, model, Δt)

Expand Down
42 changes: 22 additions & 20 deletions src/Models/LagrangianParticleTracking/lagrangian_particle_advection.jl
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
using Oceananigans.Utils: instantiate
using Oceananigans.Utils: instantiate
using Oceananigans.Models: total_velocities

#####
Expand Down Expand Up @@ -28,6 +28,13 @@ is outside the Periodic interval `(xᴸ, xᴿ)`.
@inline enforce_boundary_conditions(::Periodic, x, xᴸ, xᴿ, Cʳ) = ifelse(x > xᴿ, xᴸ + (x - xᴿ),
ifelse(x < xᴸ, xᴿ - (xᴸ - x), x))

"""
enforce_boundary_conditions(::Flat, x, xᴸ, xᴿ, Cʳ)
Do nothing on Flat dimensions.
"""
@inline enforce_boundary_conditions(::Flat, x, xᴸ, xᴿ, Cʳ) = x

const f = Face()
const c = Center()

Expand All @@ -41,23 +48,21 @@ bouncing the particle off the immersed boundary with a coefficient or `restituti
X = flattened_node((x, y, z), ibg)

# Determine current particle cell
i, j, k = fractional_indices(X, ibg.underlying_grid, (c, c, c))
i = Base.unsafe_trunc(Int, i)
j = Base.unsafe_trunc(Int, j)
k = Base.unsafe_trunc(Int, k)

fi, fj, fk = fractional_indices(X, ibg.underlying_grid, (c, c, c))
i, j, k = truncate_fractional_indices(fi, fj, fk)

if immersed_cell(i, j, k, ibg)
# Determine whether particle was _previously_ in a non-immersed cell
i⁻, j⁻, k⁻ = previous_particle_indices

if !immersed_cell(i⁻, j⁻, k⁻, ibg)
# Left-right bounds of the previous, non-immersed cell
xᴿ, yᴿ, zᴿ = node(i⁻+1, j⁻+1, k⁻+1, ibg, f, f, f)
xᴸ, yᴸ, zᴸ = node(i⁻, j⁻, k⁻, ibg, f, f, f)

Cʳ = restitution
x⁺ = enforce_boundary_conditions(Bounded(), x, xᴸ, xᴿ, Cʳ)
y⁺ = enforce_boundary_conditions(Bounded(), y, yᴸ, yᴿ, Cʳ)
x⁺ = enforce_boundary_conditions(Bounded(), x, xᴸ, xᴿ, Cʳ)
y⁺ = enforce_boundary_conditions(Bounded(), y, yᴸ, yᴿ, Cʳ)
z⁺ = enforce_boundary_conditions(Bounded(), z, zᴸ, zᴿ, Cʳ)

end
Expand All @@ -76,10 +81,8 @@ given `velocities`, time-step `Δt, and coefficient of `restitution`.
X = flattened_node((x, y, z), grid)

# Obtain current particle indices
i, j, k = fractional_indices(X, grid, c, c, c)
i = Base.unsafe_trunc(Int, i)
j = Base.unsafe_trunc(Int, j)
k = Base.unsafe_trunc(Int, k)
fi, fj, fk = fractional_indices(X, grid, c, c, c)
i, j, k = truncate_fractional_indices(fi, fj, fk)

current_particle_indices = (i, j, k)

Expand All @@ -91,7 +94,7 @@ given `velocities`, time-step `Δt, and coefficient of `restitution`.
# Advect particles, calculating the advection metric for a curvilinear grid.
# Note that all supported grids use length coordinates in the vertical, so we do not
# transform the vertical velocity nor invoke the k-index.
ξ = x_metric(i, j, grid)
ξ = x_metric(i, j, grid)
η = y_metric(i, j, grid)

x⁺ = x + ξ * u * Δt
Expand Down Expand Up @@ -137,7 +140,7 @@ end
@inline y_metric(i, j, grid::RectilinearGrid) = 1
@inline y_metric(i, j, grid::LatitudeLongitudeGrid{FT}) where FT = 1 / grid.radius * FT(360 / 2π)

@kernel function _advect_particles!(particles, restitution, grid::AbstractUnderlyingGrid, Δt, velocities)
@kernel function _advect_particles!(particles, restitution, grid::AbstractUnderlyingGrid, Δt, velocities)
p = @index(Global)

@inbounds begin
Expand All @@ -146,12 +149,12 @@ end
z = particles.z[p]
end

x⁺, y⁺, z⁺ = advect_particle((x, y, z), p, restitution, grid, Δt, velocities)
x⁺, y⁺, z⁺ = advect_particle((x, y, z), p, restitution, grid, Δt, velocities)

@inbounds begin
particles.x[p] = x⁺
particles.y[p] = y⁺
particles.z[p] = z⁺
particles.x[p] = x⁺
particles.y[p] = y⁺
particles.z[p] = z⁺
end
end

Expand All @@ -166,4 +169,3 @@ function advect_lagrangian_particles!(particles, model, Δt)

return nothing
end

87 changes: 42 additions & 45 deletions test/test_lagrangian_particle_tracking.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,16 +2,16 @@ include("dependencies_for_runtests.jl")

using NCDatasets
using StructArrays
using Oceananigans.Architectures: on_architecture
using Oceananigans.Architectures: architecture, on_architecture

struct TestParticle{T}
x :: T
y :: T
z :: T
u :: T
v :: T
w :: T
s :: T
x::T
y::T
z::T
u::T
v::T
w::T
s::T
end

function particle_tracking_simulation(; grid, particles, timestepper=:RungeKutta3, velocities=nothing)
Expand All @@ -22,41 +22,31 @@ function particle_tracking_simulation(; grid, particles, timestepper=:RungeKutta
jld2_filepath = "test_particles.jld2"
sim.output_writers[:particles_jld2] =
JLD2OutputWriter(model, (; particles=model.particles),
filename="test_particles", schedule=IterationInterval(1))
filename="test_particles", schedule=IterationInterval(1))

nc_filepath = "test_particles.nc"
sim.output_writers[:particles_nc] =
NetCDFOutputWriter(model, model.particles, filename=nc_filepath, schedule=IterationInterval(1))
NetCDFOutputWriter(model, model.particles, filename=nc_filepath, schedule=IterationInterval(1))

sim.output_writers[:checkpointer] = Checkpointer(model, schedule=IterationInterval(1),
dir = ".", prefix = "particles_checkpoint")
dir=".", prefix="particles_checkpoint")

return sim, jld2_filepath, nc_filepath
end

function run_simple_particle_tracking_tests(arch, timestepper; vertically_stretched=false)
topo = (Periodic, Periodic, Bounded)

Nx = Ny = Nz = 5
if vertically_stretched
# Slightly stretched at the top
z = [-1, -0.5, 0.0, 0.4, 0.7, 1]
else
z = (-1, 1)
end
function run_simple_particle_tracking_tests(grid, timestepper)

grid = RectilinearGrid(arch; topology=topo, size=(Nx, Ny, Nz),
x=(-1, 1), y=(-1, 1), z)
arch = architecture(grid)

P = 10

#####
##### Test default particle
#####
xs = on_architecture(arch, 0.6*ones(P))
ys = on_architecture(arch, 0.58*ones(P))
zs = on_architecture(arch, 0.8*ones(P))

xs = on_architecture(arch, 0.6 * ones(P))
ys = on_architecture(arch, 0.58 * ones(P))
zs = on_architecture(arch, 0.8 * ones(P))

particles = LagrangianParticles(x=xs, y=ys, z=zs)
@test particles isa LagrangianParticles
Expand All @@ -77,21 +67,21 @@ function run_simple_particle_tracking_tests(arch, timestepper; vertically_stretc
##### Test Boundary restitution
#####

initial_z = CUDA.@allowscalar grid.zᵃᵃᶜ[grid.Nz - 1]
top_boundary = CUDA.@allowscalar grid.zᵃᵃᶠ[grid.Nz + 1]
initial_z = CUDA.@allowscalar grid.zᵃᵃᶜ[grid.Nz-1]
top_boundary = CUDA.@allowscalar grid.zᵃᵃᶠ[grid.Nz+1]

x, y, z = on_architecture.(Ref(arch), ([0.0], [0.0], [initial_z]))

particles = LagrangianParticles(; x, y, z)
u, v, w = VelocityFields(grid)
u, v, w = VelocityFields(grid)

Δt = 0.01
interior(w, :, :, grid.Nz) .= (0.1 + top_boundary - initial_z) / Δt
interior(w, :, :, grid.Nz-1) .= (0.2 + top_boundary - initial_z) / Δt
interior(w, :, :, grid.Nz) .= (0.1 + top_boundary - initial_z) / Δt
interior(w, :, :, grid.Nz - 1) .= (0.2 + top_boundary - initial_z) / Δt

velocities = PrescribedVelocityFields(; u, v, w)

model = HydrostaticFreeSurfaceModel(; grid, particles, velocities, buoyancy=nothing, tracers = ())
model = HydrostaticFreeSurfaceModel(; grid, particles, velocities, buoyancy=nothing, tracers=())

time_step!(model, Δt)

Expand All @@ -101,7 +91,7 @@ function run_simple_particle_tracking_tests(arch, timestepper; vertically_stretc
#####
##### Test custom particle "SpeedTrackingParticle"
#####

xs = on_architecture(arch, zeros(P))
ys = on_architecture(arch, zeros(P))
zs = on_architecture(arch, 0.5 * ones(P))
Expand All @@ -115,7 +105,7 @@ function run_simple_particle_tracking_tests(arch, timestepper; vertically_stretc

velocities = VelocityFields(grid)
u, v, w = velocities
speed = Field(√(u*u + v*v + w*w))
speed = Field(√(u * u + v * v + w * w))
tracked_fields = merge(velocities, (; s=speed))

# applying v component of advection with background field to ensure it is included
Expand All @@ -128,7 +118,7 @@ function run_simple_particle_tracking_tests(arch, timestepper; vertically_stretc

model = NonhydrostaticModel(; grid, timestepper,
velocities, particles=lagrangian_particles,
background_fields = (v = background_v, ))
background_fields=(v=background_v,))

set!(model, u=1)

Expand All @@ -144,13 +134,13 @@ function run_simple_particle_tracking_tests(arch, timestepper; vertically_stretc
NetCDFOutputWriter(model, model.particles, filename=nc_filepath, schedule=IterationInterval(1))

sim.output_writers[:checkpointer] = Checkpointer(model, schedule=IterationInterval(1),
dir = ".", prefix = "particles_checkpoint")
dir=".", prefix="particles_checkpoint")

rm(jld2_filepath)
rm(nc_filepath)
rm("particles_checkpoint_iteration1.jld2")

sim, jld2_filepath, nc_filepath = particle_tracking_simulation(; grid, particles=lagrangian_particles, timestepper, velocities)
sim, jld2_filepath, nc_filepath = particle_tracking_simulation(; grid, particles=lagrangian_particles, timestepper, velocities)
model = sim.model
run!(sim)

Expand Down Expand Up @@ -274,12 +264,19 @@ function run_simple_particle_tracking_tests(arch, timestepper; vertically_stretc
return nothing
end

@testset "Lagrangian particle tracking" begin
for arch in archs, timestepper in (:QuasiAdamsBashforth2, :RungeKutta3)
@info " Testing uniform grid Lagrangian particle tracking [$(typeof(arch)), $timestepper]..."
run_simple_particle_tracking_tests(arch, timestepper; vertically_stretched=false)
lagrangian_particle_test_grid(arch, ::Periodic, z) =
RectilinearGrid(arch; topology=(Periodic, Periodic, Bounded), size=(5, 5, 5), x=(-1, 1), y=(-1, 1), z)
lagrangian_particle_test_grid(arch, ::Flat, z) =
RectilinearGrid(arch; topology=(Periodic, Flat, Bounded), size=(5, 5), x=(-1, 1), z)

@info " Testing stretched grid Lagrangian particle tracking [$(typeof(arch)), $timestepper]..."
run_simple_particle_tracking_tests(arch, timestepper; vertically_stretched=true)
@testset "Lagrangian particle tracking" begin
timesteppers = (:QuasiAdamsBashforth2, :RungeKutta3)
y_topologies = (Periodic(), Flat())
vertical_grids = (uniform=(-1, 1), stretched=[-1, -0.5, 0.0, 0.4, 0.7, 1])

for arch in archs, timestepper in timesteppers, y_topo in y_topologies, (z_grid_type, z) in pairs(vertical_grids)
@info " Testing Lagrangian particle tracking [$(typeof(arch)), $timestepper] with y $(typeof(y_topo)) on vertically $z_grid_type grid ..."
grid = lagrangian_particle_test_grid(arch, y_topo, z)
run_simple_particle_tracking_tests(grid, timestepper)
end
end

0 comments on commit 024b83b

Please sign in to comment.