diff --git a/src/Fields/interpolate.jl b/src/Fields/interpolate.jl index 05f62a1406..86ae5a3043 100644 --- a/src/Fields/interpolate.jl +++ b/src/Fields/interpolate.jl @@ -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 @@ -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 @@ -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) @@ -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...) + @@ -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) @@ -338,4 +355,3 @@ function interpolate!(to_field::Field, from_field::AbstractField) return nothing end - diff --git a/src/Models/LagrangianParticleTracking/LagrangianParticleTracking.jl b/src/Models/LagrangianParticleTracking/LagrangianParticleTracking.jl index db81da1680..43c39ee5e1 100644 --- a/src/Models/LagrangianParticleTracking/LagrangianParticleTracking.jl +++ b/src/Models/LagrangianParticleTracking/LagrangianParticleTracking.jl @@ -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 @@ -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) diff --git a/src/Models/LagrangianParticleTracking/lagrangian_particle_advection.jl b/src/Models/LagrangianParticleTracking/lagrangian_particle_advection.jl index 5695ae05f7..443d144c4e 100644 --- a/src/Models/LagrangianParticleTracking/lagrangian_particle_advection.jl +++ b/src/Models/LagrangianParticleTracking/lagrangian_particle_advection.jl @@ -1,4 +1,4 @@ -using Oceananigans.Utils: instantiate +using Oceananigans.Utils: instantiate using Oceananigans.Models: total_velocities ##### @@ -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() @@ -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 @@ -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) @@ -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 @@ -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 @@ -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 @@ -166,4 +169,3 @@ function advect_lagrangian_particles!(particles, model, Δt) return nothing end - diff --git a/test/test_lagrangian_particle_tracking.jl b/test/test_lagrangian_particle_tracking.jl index 1e6c9b6ac4..7eaa8410f7 100644 --- a/test/test_lagrangian_particle_tracking.jl +++ b/test/test_lagrangian_particle_tracking.jl @@ -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) @@ -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 @@ -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) @@ -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)) @@ -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 @@ -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) @@ -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) @@ -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