Add custom forcing velocity to particles

src/Models/LagrangianParticleTracking/LagrangianParticleTracking.jl

@@ -35,15 +35,17 @@ Base.show(io::IO, p::Particle) = print(io, "Particle at (",
                                        @sprintf("%-8s", prettysummary(p.y, true) * ", "),
                                        @sprintf("%-8s", prettysummary(p.z, true) * ")"))
 
-struct LagrangianParticles{P, R, T, D, Π} <: AbstractLagrangianParticles
-        properties :: P
-       restitution :: R
-    tracked_fields :: T
-          dynamics :: D
-        parameters :: Π
+struct LagrangianParticles{P, R, T, D, Π, F} <: AbstractLagrangianParticles
+              properties :: P
+             restitution :: R
+          tracked_fields :: T
+                dynamics :: D
+              parameters :: Π
+       advective_forcing :: F
 end
 
 @inline no_dynamics(args...) = nothing
+@inline no_advective_forcing(args...) = nothing
 
 """
     LagrangianParticles(; x, y, z, restitution=1.0, dynamics=no_dynamics, parameters=nothing)
@@ -54,7 +56,7 @@ Construct some `LagrangianParticles` that can be passed to a model. The particle
 `dynamics` is a function of `(lagrangian_particles, model, Δt)` that is called prior to advecting particles.
 `parameters` can be accessed inside the `dynamics` function.
 """
-function LagrangianParticles(; x, y, z, restitution=1.0, dynamics=no_dynamics, parameters=nothing)
+function LagrangianParticles(; x, y, z, restitution=1.0, dynamics=no_dynamics, parameters=nothing, advective_forcing=no_advective_forcing)
     size(x) == size(y) == size(z) ||
         throw(ArgumentError("x, y, z must all have the same size!"))
 
@@ -63,7 +65,7 @@ function LagrangianParticles(; x, y, z, restitution=1.0, dynamics=no_dynamics, p
 
     particles = StructArray{Particle}((x, y, z))
 
-    return LagrangianParticles(particles; restitution, dynamics, parameters)
+    return LagrangianParticles(particles; restitution, dynamics, parameters, advective_forcing)
 end
 
 """
@@ -83,15 +85,16 @@ function LagrangianParticles(particles::StructArray;
                              restitution = 1.0,
                              tracked_fields::NamedTuple=NamedTuple(),
                              dynamics = no_dynamics,
-                             parameters = nothing)
+                             parameters = nothing,
+                             advective_forcing=no_advective_forcing)
 
     for (field_name, tracked_field) in pairs(tracked_fields)
         field_name in propertynames(particles) ||
             throw(ArgumentError("$field_name is a tracked field but $(eltype(particles)) has no $field_name field! " *
                                 "You might have to define your own particle type."))
     end
 
-    return LagrangianParticles(particles, restitution, tracked_fields, dynamics, parameters)
+    return LagrangianParticles(particles, restitution, tracked_fields, dynamics, parameters, advective_forcing)
 end
 
 size(lagrangian_particles::LagrangianParticles) = size(lagrangian_particles.properties)
@@ -115,6 +118,16 @@ function Base.show(io::IO, lagrangian_particles::LagrangianParticles)
         "└── dynamics: ", prettysummary(lagrangian_particles.dynamics, false))
 end
 
+struct ParticleAdvectiveForcing{U, V, W}
+    u :: U
+    v :: V
+    w :: W
+end
+
+function ParticleAdvectiveForcing(; u=no_advective_forcing, v=no_advective_forcing, w=no_advective_forcing)
+    return ParticleAdvectiveForcing(u, v, w)
+end
+
 include("update_lagrangian_particle_properties.jl")
 include("lagrangian_particle_advection.jl")
 

src/Models/LagrangianParticleTracking/lagrangian_particle_advection.jl

@@ -37,21 +37,23 @@ const c = Center()
 Return a new particle position if the position `(x, y, z)` lies in an immersed cell by
 bouncing the particle off the immersed boundary with a coefficient or `restitution`.
 """
-@inline function bounce_immersed_particle((x, y, z), grid, restitution, previous_particle_indices)
+@inline function bounce_immersed_particle((x, y, z), ibg, restitution, previous_particle_indices)
+    X = flattened_node((x, y, z), ibg)
+
     # Determine current particle cell
-    i, j, k = fractional_indices(x, y, z, (c, c, c), grid.underlying_grid)
+    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)
 
-    if immersed_cell(i, j, k, grid)
+    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⁻, grid)
+        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, grid, f, f, f)
-            xᴸ, yᴸ, zᴸ = node(i⁻,   j⁻,   k⁻,   grid, f, f, f)
+            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ʳ)    
@@ -64,26 +66,49 @@ bouncing the particle off the immersed boundary with a coefficient or `restituti
     return x⁺, y⁺, z⁺
 end
 
+"""
+    particle_u_velocity(u_fluid, particles, p, Δt)
+
+a particle-specific advecting velocity such as 
+- sinking or rising for buoyant particles
+- electromagnetic forces for charged particles
+- swimming for biological particles
+- diffusing velocity for brownian motion
+Inputs are the fluid velocity, particle properties `particles`, and the particle index `p`
+
+Returns the fluid velocity by default (for non-buoyant, non-drifting particles). Has to be extended to obtain the desired effect
+"""
+@inline particle_u_velocity(u_fluid, u_advective_forcing, particles, model, p, Δt) = u_fluid + u_advective_forcing(particles, model, Δt)
+@inline particle_v_velocity(v_fluid, v_advective_forcing, particles, model, p, Δt) = v_fluid + v_advective_forcing(particles, model, Δt)
+@inline particle_w_velocity(w_fluid, w_advective_forcing, particles, model, p, Δt) = w_fluid + w_advective_forcing(particles, model, Δt)
+
 """
     advect_particle((x, y, z), p, restitution, grid, Δt, velocities)
 
 Return new position `(x⁺, y⁺, z⁺)` for a particle at current position (x, y, z),
 given `velocities`, time-step `Δt, and coefficient of `restitution`.
 """
-@inline function advect_particle((x, y, z), p, restitution, grid, Δt, velocities)
+@inline function advect_particle((x, y, z), p, particles, model, restitution, grid, Δt, velocities)
+    X = flattened_node((x, y, z), grid)
+
     # Obtain current particle indices
-    i, j, k = fractional_indices(x, y, z, (c, c, c), grid)
+    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)
 
     current_particle_indices = (i, j, k)
 
     # Interpolate velocity to particle position
-    u = interpolate(velocities.u, f, c, c, grid, x, y, z)
-    v = interpolate(velocities.v, c, f, c, grid, x, y, z)
-    w = interpolate(velocities.w, c, c, f, grid, x, y, z)
-
+    u_fluid = interpolate(X, velocities.u, (f, c, c), grid) 
+    v_fluid = interpolate(X, velocities.v, (c, f, c), grid) 
+    w_fluid = interpolate(X, velocities.w, (c, c, f), grid)
+
+    # Particle velocity
+    u = particle_u_velocity(u_fluid, advective_forcing.u, particles, model, p, Δt)
+    v = particle_v_velocity(v_fluid, advective_forcing.v, particles, model, p, Δt)
+    w = particle_w_velocity(w_fluid, advective_forcing.w, particles, model, p, Δt)
+
     # 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.
@@ -133,7 +158,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, model, restitution, grid::AbstractUnderlyingGrid, Δt, velocities) 
     p = @index(Global)
 
     @inbounds begin
@@ -142,7 +167,7 @@ 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, particles, model, restitution, grid, Δt, velocities) 
 
     @inbounds begin
         particles.x[p] = x⁺ 
@@ -158,8 +183,7 @@ function advect_lagrangian_particles!(particles, model, Δt)
     worksize = length(particles)
 
     advect_particles_kernel! = _advect_particles!(device(arch), workgroup, worksize)
-    advect_particles_kernel!(particles.properties, particles.restitution, model.grid, Δt, total_velocities(model))
+    advect_particles_kernel!(particles.properties, model, particles.restitution, model.grid, Δt, total_velocities(model))
 
     return nothing
 end
-