(0.91.14) Extend Lagrangian advection to immersed grids and add tests

Project.toml

@@ -1,7 +1,7 @@
 name = "Oceananigans"
 uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 authors = ["Climate Modeling Alliance and contributors"]
-version = "0.91.13"
+version = "0.91.14"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/Models/LagrangianParticleTracking/lagrangian_particle_advection.jl

@@ -38,6 +38,13 @@ Do nothing on Flat dimensions.
 const f = Face()
 const c = Center()
 
+"""
+    immersed_boundary_topology(grid_topology)
+
+Unless `Flat`, immersed boundaries are treated as `Bounded` regardless of underlying grid topology.
+"""
+immersed_boundary_topology(grid_topology) = ifelse(grid_topology == Flat, Flat(), Bounded())
+
 """
     bounce_immersed_particle((x, y, z), grid, restitution, previous_particle_indices)
 
@@ -48,27 +55,34 @@ bouncing the particle off the immersed boundary with a coefficient or `restituti
     X = flattened_node((x, y, z), ibg)
 
     # Determine current particle cell
-    fi, fj, fk = fractional_indices(X, ibg.underlying_grid, (c, c, c))
+    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
+    # 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)
+    # Left bounds of the previous cell
+    xᴿ = xnode(i⁻ + 1, j, k, ibg, f, f, f)
+    yᴿ = ynode(i, j⁻ + 1, k, ibg, f, f, f)
+    zᴿ = znode(i, j, k⁻ + 1, ibg, f, f, f)
 
-            Cʳ = restitution
-            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ʳ)
+    # Right bounds of the previous cell
+    xᴸ = xnode(i⁻, j,  k,  ibg, f, f, f)
+    yᴸ = ynode(i,  j⁻, k,  ibg, f, f, f)
+    zᴸ = znode(i,  j,  k⁻, ibg, f, f, f)
 
-        end
-    end
+    Cʳ = restitution
+    tx, ty, tz = map(immersed_boundary_topology, topology(ibg))
+    xb⁺ = enforce_boundary_conditions(tx, x, xᴸ, xᴿ, Cʳ)
+    yb⁺ = enforce_boundary_conditions(ty, y, yᴸ, yᴿ, Cʳ)
+    zb⁺ = enforce_boundary_conditions(tz, z, zᴸ, zᴿ, Cʳ)
 
-    return x⁺, y⁺, z⁺
+    immersed = immersed_cell(i, j, k, ibg)
+    x⁺ = ifelse(immersed, xb⁺, x)
+    y⁺ = ifelse(immersed, yb⁺, y)
+    z⁺ = ifelse(immersed, zb⁺, z)
+
+    return (x⁺, y⁺, z⁺)
 end
 
 """
@@ -134,7 +148,7 @@ given `velocities`, time-step `Δt, and coefficient of `restitution`.
     z⁺ = enforce_boundary_conditions(tz, z⁺, zᴸ, zᴿ, Cʳ)
     if grid isa ImmersedBoundaryGrid
         previous_particle_indices = current_particle_indices # particle has been advected
-        x⁺, y⁺, z⁺ = bounce_immersed_particle((x⁺, y⁺, z⁺), grid, Cʳ, previous_particle_indices)
+        (x⁺, y⁺, z⁺) = bounce_immersed_particle((x⁺, y⁺, z⁺), grid, Cʳ, previous_particle_indices)
     end
 
     return (x⁺, y⁺, z⁺)
@@ -145,11 +159,13 @@ end
 #     * Sphere metric for `LatitudeLongitudeGrid` and geographic coordinates
 @inline x_metric(i, j, grid::RectilinearGrid) = 1
 @inline x_metric(i, j, grid::LatitudeLongitudeGrid{FT}) where FT = @inbounds 1 / (grid.radius * hack_cosd(grid.φᵃᶜᵃ[j])) * FT(360 / 2π)
+@inline x_metric(i, j, grid::ImmersedBoundaryGrid) = x_metric(i, j, grid.underlying_grid)
 
 @inline y_metric(i, j, grid::RectilinearGrid) = 1
 @inline y_metric(i, j, grid::LatitudeLongitudeGrid{FT}) where FT = 1 / grid.radius * FT(360 / 2π)
+@inline y_metric(i, j, grid::ImmersedBoundaryGrid) = y_metric(i, j, grid.underlying_grid)
 
-@kernel function _advect_particles!(particles, restitution, grid::AbstractUnderlyingGrid, Δt, velocities)
+@kernel function _advect_particles!(particles, restitution, grid::AbstractGrid, Δt, velocities)
     p = @index(Global)
 
     @inbounds begin

test/test_lagrangian_particle_tracking.jl

@@ -269,6 +269,19 @@ lagrangian_particle_test_grid(arch, ::Periodic, z) =
 lagrangian_particle_test_grid(arch, ::Flat, z) =
     RectilinearGrid(arch; topology=(Periodic, Flat, Bounded), size=(5, 5), x=(-1, 1), z)
 
+lagrangian_particle_test_grid_expanded(arch, ::Periodic, z) =
+    RectilinearGrid(arch; topology=(Periodic, Periodic, Bounded), size=(5, 5, 5), x=(-1, 1), y=(-1, 1), z = 2 .*z)
+lagrangian_particle_test_grid_expanded(arch, ::Flat, z) =
+    RectilinearGrid(arch; topology=(Periodic, Flat, Bounded), size=(5, 5), x=(-1, 1), z = 2 .*z)
+
+function lagrangian_particle_test_immersed_grid(arch, y_topo, z)
+    underlying_grid = lagrangian_particle_test_grid_expanded(arch, y_topo, z)
+    z_immersed_boundary(x, z) = ifelse(z < -1, true, ifelse(z > 1, true, false))
+    z_immersed_boundary(x, y, z) = z_immersed_boundary(x, z)
+    GFB = GridFittedBoundary(z_immersed_boundary)
+    return ImmersedBoundaryGrid(underlying_grid, GFB)
+end
+
 @testset "Lagrangian particle tracking" begin
     timesteppers = (:QuasiAdamsBashforth2, :RungeKutta3)
     y_topologies = (Periodic(), Flat())
@@ -278,5 +291,11 @@ lagrangian_particle_test_grid(arch, ::Flat, z) =
         @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)
+
+        if z isa NTuple{2} # Test immersed regular grids
+            @info "  Testing Lagrangian particle tracking [$(typeof(arch)), $timestepper] with y $(typeof(y_topo)) on vertically $z_grid_type immersed grid ..."
+            grid = lagrangian_particle_test_immersed_grid(arch, y_topo, z)
+            run_simple_particle_tracking_tests(grid, timestepper)
+        end
     end
 end