Make continuity_equation! and density_diffusion! work on Ref values (#1130)

* Make `continuity_equation!` and `density_diffusion!` work on Ref values

* Fix writing drho to dv for summation density

* Implement copilot suggestions

* Fix density diffusion smoothing length

* Fix missing suggestion from #1043

* Fix

* Fix `continuity_equation!` functions

Author: efaulhaber

src/schemes/boundary/open_boundary/rhs.jl

@@ -77,14 +77,17 @@ function interact!(dv, v_particle_system, u_particle_system,
             @inbounds dv[i, particle] += dv_particle[i]
         end
 
-        v_diff = current_velocity(v_particle_system, particle_system, particle) -
-                 current_velocity(v_neighbor_system, neighbor_system, neighbor)
+        v_a = current_velocity(v_particle_system, particle_system, particle)
+        v_b = current_velocity(v_neighbor_system, neighbor_system, neighbor)
+        v_diff = v_a - v_b
 
         # Propagate `@inbounds` to the continuity equation, which accesses particle data
-        @inbounds continuity_equation!(dv, particle_system, neighbor_system,
-                                       v_particle_system, v_neighbor_system, particle,
-                                       neighbor, pos_diff, distance, m_b, rho_a, rho_b,
-                                       grad_kernel)
+        drho_particle = Ref(zero(rho_a))
+        @inbounds continuity_equation!(drho_particle,
+                                       particle_system, neighbor_system,
+                                       particle, neighbor, pos_diff, distance,
+                                       m_b, rho_a, rho_b, v_a, v_b, grad_kernel)
+        dv[end, particle] += drho_particle[]
 
         # Open boundary pressure evolution matches the corresponding fluid system:
         # - EDAC: Compute pressure evolution like the fluid system

src/schemes/boundary/wall_boundary/rhs.jl

@@ -45,11 +45,14 @@ function interact!(dv, v_particle_system, u_particle_system,
         rho_a = current_density(v_particle_system, particle_system, particle)
         rho_b = current_density(v_neighbor_system, neighbor_system, neighbor)
 
-        v_diff = current_velocity(v_particle_system, particle_system, particle) -
-                 current_velocity(v_neighbor_system, neighbor_system, neighbor)
+        v_a = current_velocity(v_particle_system, particle_system, particle)
+        v_b = current_velocity(v_neighbor_system, neighbor_system, neighbor)
 
-        continuity_equation!(dv, density_calculator(neighbor_system),
-                             m_b, rho_a, rho_b, v_diff, grad_kernel, particle)
+        drho_particle = Ref(zero(rho_a))
+        continuity_equation!(drho_particle, density_calculator(neighbor_system),
+                             m_b, rho_a, rho_b, v_a, v_b, grad_kernel, particle)
+
+        dv[end, particle] += drho_particle[]
     end
 
     return dv
@@ -58,13 +61,19 @@ end
 # This is the derivative of the density summation, which is compatible with the
 # `SummationDensity` pressure acceleration.
 # Energy preservation tests will fail with the other formulation.
-function continuity_equation!(dv, fluid_density_calculator::SummationDensity,
-                              m_b, rho_a, rho_b, v_diff, grad_kernel, particle)
-    dv[end, particle] += m_b * dot(v_diff, grad_kernel)
+@propagate_inbounds function continuity_equation!(drho_particle, ::SummationDensity,
+                                                  m_b, rho_a, rho_b, v_a, v_b,
+                                                  grad_kernel, particle)
+    drho_particle[] += m_b * dot(v_a - v_b, grad_kernel)
+
+    return drho_particle
 end
 
 # This is identical to the continuity equation of the fluid
-function continuity_equation!(dv, fluid_density_calculator::ContinuityDensity,
-                              m_b, rho_a, rho_b, v_diff, grad_kernel, particle)
-    dv[end, particle] += rho_a / rho_b * m_b * dot(v_diff, grad_kernel)
+@propagate_inbounds function continuity_equation!(drho_particle, ::ContinuityDensity,
+                                                  m_b, rho_a, rho_b, v_a, v_b,
+                                                  grad_kernel, particle)
+    drho_particle[] += rho_a / rho_b * m_b * dot(v_a - v_b, grad_kernel)
+
+    return drho_particle
 end

src/schemes/fluid/entropically_damped_sph/rhs.jl

@@ -93,19 +93,24 @@ function interact!(dv, v_particle_system, u_particle_system,
             @inbounds dv[i, particle] += dv_particle[][i]
         end
 
-        v_diff = current_velocity(v_particle_system, particle_system, particle) -
-                 current_velocity(v_neighbor_system, neighbor_system, neighbor)
+        v_a = current_velocity(v_particle_system, particle_system, particle)
+        v_b = current_velocity(v_neighbor_system, neighbor_system, neighbor)
+        v_diff = v_a - v_b
 
         pressure_evolution!(dv, particle_system, neighbor_system, v_diff, grad_kernel,
                             particle, neighbor, pos_diff, distance,
                             sound_speed, m_a, m_b, p_a, p_b, rho_a, rho_b, nu_edac)
 
+        drho_particle = Ref(zero(rho_a))
+
         # TODO If variable smoothing_length is used, this should use the neighbor smoothing length
         # Propagate `@inbounds` to the continuity equation, which accesses particle data
-        @inbounds continuity_equation!(dv, density_calculator, particle_system,
-                                       neighbor_system, v_particle_system,
-                                       v_neighbor_system, particle, neighbor,
-                                       pos_diff, distance, m_b, rho_a, rho_b, grad_kernel)
+        @inbounds continuity_equation!(drho_particle, density_calculator,
+                                       particle_system, neighbor_system,
+                                       particle, neighbor, pos_diff, distance,
+                                       m_b, rho_a, rho_b, v_a, v_b, grad_kernel)
+
+        @inbounds write_drho_particle!(dv, density_calculator, drho_particle, particle)
     end
 
     return dv

src/schemes/fluid/fluid.jl

@@ -128,51 +128,62 @@ function compute_density!(system, u, u_ode, semi, ::SummationDensity)
 end
 
 # With 'SummationDensity', density is calculated in wcsph/system.jl:compute_density!
-@inline function continuity_equation!(dv, density_calculator::SummationDensity,
+@inline function continuity_equation!(drho_particle, ::SummationDensity,
                                       particle_system, neighbor_system,
-                                      v_particle_system, v_neighbor_system,
                                       particle, neighbor, pos_diff, distance,
-                                      m_b, rho_a, rho_b, grad_kernel)
-    return dv
+                                      m_b, rho_a, rho_b, v_a, v_b, grad_kernel)
+    return drho_particle
 end
 
 # This formulation was chosen to be consistent with the used pressure_acceleration formulations
-@propagate_inbounds function continuity_equation!(dv, density_calculator::ContinuityDensity,
+@propagate_inbounds function continuity_equation!(drho_particle,
+                                                  ::ContinuityDensity,
                                                   particle_system::AbstractFluidSystem,
                                                   neighbor_system,
-                                                  v_particle_system, v_neighbor_system,
                                                   particle, neighbor, pos_diff, distance,
-                                                  m_b, rho_a, rho_b, grad_kernel)
-    return continuity_equation!(dv, particle_system, neighbor_system, v_particle_system,
-                                v_neighbor_system, particle, neighbor, pos_diff, distance,
-                                m_b, rho_a, rho_b, grad_kernel)
+                                                  m_b, rho_a, rho_b, v_a, v_b, grad_kernel)
+    continuity_equation!(drho_particle, particle_system, neighbor_system,
+                         particle, neighbor, pos_diff, distance,
+                         m_b, rho_a, rho_b, v_a, v_b, grad_kernel)
 end
 
-@propagate_inbounds function continuity_equation!(dv, particle_system, neighbor_system,
-                                                  v_particle_system, v_neighbor_system,
+@propagate_inbounds function continuity_equation!(drho_particle,
+                                                  particle_system, neighbor_system,
                                                   particle, neighbor, pos_diff, distance,
-                                                  m_b, rho_a, rho_b, grad_kernel)
-    vdiff = current_velocity(v_particle_system, particle_system, particle) -
-            current_velocity(v_neighbor_system, neighbor_system, neighbor)
+                                                  m_b, rho_a, rho_b, v_a, v_b, grad_kernel)
+    v_diff = v_a - v_b
 
-    vdiff += continuity_equation_shifting_term(shifting_technique(particle_system),
+    v_diff = continuity_equation_shifting_term(v_diff,
+                                               shifting_technique(particle_system),
                                                particle_system, neighbor_system,
                                                particle, neighbor, rho_a, rho_b)
 
     # Since this is one of the most performance critical functions, using fast divisions
     # here gives a significant speedup on GPUs.
     # See the docs page "Development" for more details on `div_fast`.
-    dv[end, particle] += div_fast(rho_a, rho_b) * m_b * dot(vdiff, grad_kernel)
+    drho_particle[] += div_fast(rho_a, rho_b) * m_b * dot(v_diff, grad_kernel)
 
     # Artificial density diffusion should only be applied to systems representing a fluid
     # with the same physical properties i.e. density and viscosity.
     # TODO: shouldn't be applied to particles on the interface (depends on PR #539)
     if particle_system === neighbor_system
-        density_diffusion!(dv, density_diffusion(particle_system),
-                           v_particle_system, particle, neighbor,
-                           pos_diff, distance, m_b, rho_a, rho_b, particle_system,
-                           grad_kernel)
+        density_diffusion!(drho_particle, density_diffusion(particle_system),
+                           particle_system, particle, neighbor,
+                           pos_diff, distance, m_b, rho_a, rho_b, grad_kernel)
     end
+
+    return drho_particle
+end
+
+@inline function write_drho_particle!(dv, density_calculator, drho_particle, particle)
+    return dv
+end
+
+@propagate_inbounds function write_drho_particle!(dv, ::ContinuityDensity,
+                                                  drho_particle, particle)
+    dv[end, particle] += drho_particle[]
+
+    return dv
 end
 
 function calculate_dt(v_ode, u_ode, cfl_number, system::AbstractFluidSystem, semi)

src/schemes/fluid/shifting_techniques.jl

@@ -48,11 +48,12 @@ end
     return dv_particle
 end
 
-# Additional term(s) in the continuity equation due to the shifting technique
-@inline function continuity_equation_shifting_term(shifting, particle_system,
+# Add additional term(s) in the continuity equation due to the shifting technique
+# and return the modified term.
+@inline function continuity_equation_shifting_term(v_diff, shifting, particle_system,
                                                    neighbor_system,
                                                    particle, neighbor, rho_a, rho_b)
-    return zero(SVector{ndims(particle_system), eltype(particle_system)})
+    return v_diff
 end
 
 @doc raw"""
@@ -356,7 +357,8 @@ end
 end
 
 # `ParticleShiftingTechnique{<:Any, <:Any, true}` means `modify_continuity_equation=true`
-@propagate_inbounds function continuity_equation_shifting_term(shifting::ParticleShiftingTechnique{<:Any,
+@propagate_inbounds function continuity_equation_shifting_term(v_diff,
+                                                               shifting::ParticleShiftingTechnique{<:Any,
                                                                                                    <:Any,
                                                                                                    true},
                                                                system, neighbor_system,
@@ -366,9 +368,10 @@ end
     delta_v_b = delta_v(neighbor_system, neighbor)
     delta_v_diff = delta_v_a - delta_v_b
 
-    second_term = second_continuity_equation_term(shifting.second_continuity_equation_term,
-                                                  delta_v_a, delta_v_b, rho_a, rho_b)
-    return delta_v_diff + second_term
+    shifting_term = second_continuity_equation_term(delta_v_diff,
+                                                    shifting.second_continuity_equation_term,
+                                                    delta_v_a, delta_v_b, rho_a, rho_b)
+    return v_diff + shifting_term
 end
 
 """
@@ -381,15 +384,16 @@ See [`ParticleShiftingTechnique`](@ref).
 """
 struct ContinuityEquationTermSun2019 end
 
-@propagate_inbounds function second_continuity_equation_term(::ContinuityEquationTermSun2019,
+@propagate_inbounds function second_continuity_equation_term(v_diff,
+                                                             ::ContinuityEquationTermSun2019,
                                                              delta_v_a, delta_v_b,
                                                              rho_a, rho_b)
-    return delta_v_a + rho_b / rho_a * delta_v_b
+    return v_diff + delta_v_a + rho_b / rho_a * delta_v_b
 end
 
-@inline function second_continuity_equation_term(second_continuity_equation_term,
+@inline function second_continuity_equation_term(v_diff, second_continuity_equation_term,
                                                  delta_v_a, delta_v_b, rho_a, rho_b)
-    return zero(delta_v_a)
+    return v_diff
 end
 
 # `ParticleShiftingTechnique{<:Any, true}` means `update_everystage=true`
@@ -606,15 +610,16 @@ end
     return pressure_acceleration_continuity_density(m_a, m_b, rho_a, rho_b, p_a, p_b, W_a)
 end
 
-@propagate_inbounds function continuity_equation_shifting_term(::TransportVelocityAdami{true},
+@propagate_inbounds function continuity_equation_shifting_term(v_diff,
+                                                               ::TransportVelocityAdami{true},
                                                                particle_system,
                                                                neighbor_system,
                                                                particle, neighbor,
                                                                rho_a, rho_b)
     delta_v_diff = delta_v(particle_system, particle) -
                    delta_v(neighbor_system, neighbor)
 
-    return delta_v_diff
+    return v_diff + delta_v_diff
 end
 
 function update_shifting!(system, shifting::TransportVelocityAdami, v, u, v_ode,

src/schemes/fluid/weakly_compressible_sph/density_diffusion.jl

@@ -213,13 +213,12 @@ function update!(density_diffusion::DensityDiffusionAntuono, v, u, system, semi)
     return density_diffusion
 end
 
-@propagate_inbounds function density_diffusion!(dv,
+@propagate_inbounds function density_diffusion!(drho_particle,
                                                 density_diffusion::AbstractDensityDiffusion,
-                                                v_particle_system, particle, neighbor,
-                                                pos_diff, distance, m_b, rho_a, rho_b,
                                                 particle_system::Union{AbstractFluidSystem,
                                                                        OpenBoundarySystem{<:BoundaryModelDynamicalPressureZhang}},
-                                                grad_kernel)
+                                                particle, neighbor, pos_diff, distance,
+                                                m_b, rho_a, rho_b, grad_kernel)
     # Density diffusion terms are all zero for distance zero.
     # If `skip_zero_distance` is `true`, we can assume that this function isn't called
     # for distance zero because these neighbors have already been skipped.
@@ -240,12 +239,12 @@ end
 
     (; delta) = density_diffusion
     sound_speed = system_sound_speed(particle_system)
-    dv[end, particle] += delta * smoothing_length_avg * sound_speed * density_diffusion_term
+    drho_particle[] += delta * smoothing_length_avg * sound_speed * density_diffusion_term
 end
 
 # Density diffusion `nothing` or interaction other than fluid-fluid
-@inline function density_diffusion!(dv, density_diffusion, v_particle_system, particle,
-                                    neighbor, pos_diff, distance, m_b, rho_a, rho_b,
-                                    particle_system, grad_kernel)
-    return dv
+@inline function density_diffusion!(drho_particle, density_diffusion, particle_system,
+                                    particle, neighbor, pos_diff, distance,
+                                    m_b, rho_a, rho_b, grad_kernel)
+    return drho_particle
 end

src/schemes/fluid/weakly_compressible_sph/rhs.jl

@@ -51,6 +51,9 @@ function interact!(dv, v_particle_system, u_particle_system,
         rho_b = @inbounds current_density(v_neighbor_system, neighbor_system, neighbor)
         rho_mean = (rho_a + rho_b) / 2
 
+        v_a = @inbounds current_velocity(v_particle_system, particle_system, particle)
+        v_b = @inbounds current_velocity(v_neighbor_system, neighbor_system, neighbor)
+
         # Determine correction factors.
         # This can be ignored, as these are all 1 when no correction is used.
         (viscosity_correction, pressure_correction,
@@ -114,12 +117,16 @@ function interact!(dv, v_particle_system, u_particle_system,
             # debug_array[i, particle] += dv_pressure[i]
         end
 
+        drho_particle = Ref(zero(rho_a))
+
         # TODO If variable smoothing_length is used, this should use the neighbor smoothing length
         # Propagate `@inbounds` to the continuity equation, which accesses particle data
-        @inbounds continuity_equation!(dv, density_calculator, particle_system,
-                                       neighbor_system, v_particle_system,
-                                       v_neighbor_system, particle, neighbor,
-                                       pos_diff, distance, m_b, rho_a, rho_b, grad_kernel)
+        @inbounds continuity_equation!(drho_particle, density_calculator,
+                                       particle_system, neighbor_system,
+                                       particle, neighbor, pos_diff, distance,
+                                       m_b, rho_a, rho_b, v_a, v_b, grad_kernel)
+
+        @inbounds write_drho_particle!(dv, density_calculator, drho_particle, particle)
     end
     # Debug example
     # periodic_box = neighborhood_search.periodic_box