Make DEM run on GPUs and with Float32 (#979)

* Make DEM run on GPUs

* Add test

* Fix Metal tests

* Fix DEM coefficients

---------

Co-authored-by: Niklas Neher <73897120+LasNikas@users.noreply.github.com>

Author: efaulhaber

examples/dem/rectangular_tank_2d.jl

@@ -35,8 +35,10 @@ tank = RectangularTank(particle_spacing, (rock_width, rock_height),
 
 # Move the rock particles up to let them fall
 tank.fluid.coordinates[2, :] .+= 0.5
-# small perturbation
-tank.fluid.coordinates .+= 0.01 .* (2 .* rand(size(tank.fluid.coordinates)) .- 1)
+# Small perturbation
+for i in 1:size(tank.fluid.coordinates, 2)
+    tank.fluid.coordinates[:, i] .+= 0.01 .* (2 .* rand(2) .- 1)
+end
 
 # Create a contact model.
 # Option 1: Hertzian contact model (uses elastic modulus and Poisson's ratio)
@@ -54,8 +56,9 @@ boundary_system = BoundaryDEMSystem(tank.boundary, 10e7)
 # ==========================================================================================
 # ==== Simulation
 # ==========================================================================================
-
-semi = Semidiscretization(rock_system, boundary_system)
+semi = Semidiscretization(rock_system, boundary_system;
+                          neighborhood_search=GridNeighborhoodSearch{2}(),
+                          parallelization_backend=PolyesterBackend())
 
 tspan = (0.0, 4.0)
 ode = semidiscretize(semi, tspan)

src/general/gpu.jl

@@ -20,6 +20,8 @@ Adapt.@adapt_structure TotalLagrangianSPHSystem
 Adapt.@adapt_structure BoundaryZone
 Adapt.@adapt_structure SystemBuffer
 Adapt.@adapt_structure OpenBoundarySystem
+Adapt.@adapt_structure DEMSystem
+Adapt.@adapt_structure BoundaryDEMSystem
 Adapt.@adapt_structure RCRWindkesselModel
 
 KernelAbstractions.get_backend(::PtrArray) = KernelAbstractions.CPU()

src/schemes/boundary/dem_boundary/system.jl

@@ -16,18 +16,29 @@ struct BoundaryDEMSystem{NDIMS, ELTYPE <: Real, IC,
     normal_stiffness  :: ELTYPE
     buffer            :: Nothing
 
-    function BoundaryDEMSystem(initial_condition, normal_stiffness)
-        coordinates = initial_condition.coordinates
-        radius = 0.5 * initial_condition.particle_spacing *
-                 ones(length(initial_condition.mass))
-        NDIMS = size(coordinates, 1)
+    function BoundaryDEMSystem(initial_condition, coordinates, radius,
+                               normal_stiffness, buffer)
+        NDIMS = ndims(initial_condition)
 
-        return new{NDIMS, eltype(coordinates), typeof(initial_condition), typeof(radius),
+        return new{NDIMS, eltype(initial_condition), typeof(initial_condition),
+                   typeof(radius),
                    typeof(coordinates)}(initial_condition, coordinates, radius,
-                                        normal_stiffness, nothing)
+                                        normal_stiffness, buffer)
     end
 end
 
+# The default constructor needs to be accessible for Adapt.jl to work with this struct.
+# See the comments in general/gpu.jl for more details.
+function BoundaryDEMSystem(initial_condition, normal_stiffness)
+    ELTYPE = eltype(initial_condition)
+    coordinates = initial_condition.coordinates
+    radius = initial_condition.particle_spacing *
+             ones(ELTYPE, length(initial_condition.mass)) / 2
+
+    return BoundaryDEMSystem(initial_condition, coordinates, radius,
+                             normal_stiffness, nothing)
+end
+
 @inline function Base.eltype(system::BoundaryDEMSystem)
     eltype(system.coordinates)
 end

src/schemes/structure/discrete_element_method/contact_models.jl

@@ -88,7 +88,7 @@ end
     r_star = (r_a * r_b) / (r_a + r_b)
 
     # Non-linear stiffness for Hertzian contact
-    elastic_force_per_overlap = (4 / 3) * E_star * sqrt(r_star * overlap)
+    elastic_force_per_overlap = 4 * E_star * sqrt(r_star * overlap) / 3
 
     # Compute effective mass for damping
     if neighbor_system isa DEMSystem

src/schemes/structure/discrete_element_method/rhs.jl

@@ -12,7 +12,7 @@ function interact!(dv, v_particle_system, u_particle_system, v_neighbor_system,
                                                                                 neighbor,
                                                                                 pos_diff,
                                                                                 distance
-        # See `src/general/smoothing_kernels.jl` for more details.
+        # See `src/general/smoothing_kernels.jl` for more details
         distance^2 < eps(first(particle_system.radius)^2) && return
 
         # Retrieve particle properties
@@ -26,7 +26,7 @@ function interact!(dv, v_particle_system, u_particle_system, v_neighbor_system,
             # Compute the unit normal vector (from neighbor to particle)
             normal = pos_diff / distance
 
-            # Compute Normal Force by Dispatching on the Contact Model.
+            # Compute Normal Force by Dispatching on the Contact Model
             F_normal = collision_force_normal(particle_system.contact_model,
                                               particle_system, neighbor_system,
                                               overlap, normal,
@@ -46,7 +46,7 @@ function interact!(dv, v_particle_system, u_particle_system, v_neighbor_system,
             end
 
             # Apply a simple position correction to mitigate overlap.
-            # TODO: use update callback
+            # TODO: use update callback, changing `u` is not allowed here.
             position_correction!(neighbor_system, u_particle_system, overlap, normal,
                                  particle)
         end
@@ -63,10 +63,14 @@ end
                                             overlap, normal,
                                             v_particle_system, v_neighbor_system,
                                             particle, neighbor, F_normal)
-    # Tangential force parameters
-    friction_coefficient = 0.5       # Coulomb friction coefficient [Cundall and Strack, 1979]
-    tangential_stiffness = 1e3       # Tangential spring constant
-    tangential_damping = 0.001       # Damping coefficient for tangential force
+    # Tangential force parameters. Avoid hardcoding double values.
+    ELTYPE = eltype(particle_system)
+    # Coulomb friction coefficient [Cundall and Strack, 1979]
+    friction_coefficient = convert(ELTYPE, 0.5)
+    # Tangential spring constant
+    tangential_stiffness = 1000
+    # Damping coefficient for tangential force
+    tangential_damping = convert(ELTYPE, 0.001)
 
     # Compute relative velocity and extract the tangential component.
     v_a = current_velocity(v_particle_system, particle_system, particle)
@@ -79,7 +83,7 @@ end
 
     # Coulomb friction: limit the tangential force to μ * |F_normal|
     max_tangent = friction_coefficient * norm(F_normal)
-    if norm(F_t) > max_tangent && norm(F_t) > 0.0
+    if norm(F_t) > max_tangent && norm(F_t) > 0
         F_t = F_t * (max_tangent / norm(F_t))
     end
 
@@ -94,6 +98,6 @@ end
 @inline function position_correction!(neighbor_system::BoundaryDEMSystem,
                                       u_particle_system, overlap, normal, particle)
     for i in 1:ndims(neighbor_system)
-        u_particle_system[i, particle] -= 0.5 * overlap * normal[i]
+        u_particle_system[i, particle] -= overlap * normal[i] / 2
     end
 end

src/schemes/structure/discrete_element_method/system.jl

@@ -38,35 +38,35 @@ struct DEMSystem{NDIMS, ELTYPE <: Real, IC, ARRAY1D, ST,
     acceleration        :: SVector{NDIMS, ELTYPE}
     source_terms        :: ST
     contact_model       :: CM
+end
 
-    function DEMSystem(initial_condition, contact_model; damping_coefficient=0.0001,
-                       acceleration=ntuple(_ -> 0.0,
-                                           ndims(initial_condition)), source_terms=nothing,
-                       radius=nothing)
-        NDIMS = ndims(initial_condition)
-        ELTYPE = eltype(initial_condition)
-
-        mass = copy(initial_condition.mass)
+# The default constructor needs to be accessible for Adapt.jl to work with this struct.
+# See the comments in general/gpu.jl for more details.
+function DEMSystem(initial_condition, contact_model; damping_coefficient=0.0001,
+                   acceleration=ntuple(_ -> 0.0,
+                                       ndims(initial_condition)), source_terms=nothing,
+                   radius=nothing)
+    NDIMS = ndims(initial_condition)
+    ELTYPE = eltype(initial_condition)
 
-        if isnothing(radius)
-            radius = 0.5 * initial_condition.particle_spacing * ones(length(mass))
-        else
-            mass = (radius / (0.5 * initial_condition.particle_spacing))^3 * mass
-            radius = radius * ones(length(mass))
-        end
+    mass = copy(initial_condition.mass)
 
-        # Make acceleration an SVector
-        acceleration_ = SVector(acceleration...)
-        if length(acceleration_) != NDIMS
-            throw(ArgumentError("`acceleration` must be of length $NDIMS for a $(NDIMS)D problem"))
-        end
+    if isnothing(radius)
+        radius = initial_condition.particle_spacing * ones(ELTYPE, length(mass)) / 2
+    else
+        mass = (radius / (initial_condition.particle_spacing / 2))^3 * mass
+        radius = radius * ones(ELTYPE, length(mass))
+    end
 
-        return new{NDIMS, ELTYPE, typeof(initial_condition),
-                   typeof(mass), typeof(source_terms),
-                   typeof(contact_model)}(initial_condition, mass, radius,
-                                          damping_coefficient, acceleration_, source_terms,
-                                          contact_model)
+    # Make acceleration an SVector
+    acceleration_ = SVector(acceleration...)
+    if length(acceleration_) != NDIMS
+        throw(ArgumentError("`acceleration` must be of length $NDIMS for a $(NDIMS)D problem"))
     end
+
+    return DEMSystem(initial_condition, mass, radius,
+                     damping_coefficient, acceleration_, source_terms,
+                     contact_model)
 end
 
 function Base.show(io::IO, system::DEMSystem)
@@ -118,12 +118,12 @@ end
 timer_name(::DEMSystem) = "solid"
 
 function TrixiParticles.write_u0!(u0, system::DEMSystem)
-    u0 .= system.initial_condition.coordinates
+    copyto!(u0, system.initial_condition.coordinates)
     return u0
 end
 
 function TrixiParticles.write_v0!(v0, system::DEMSystem)
-    v0 .= system.initial_condition.velocity
+    copyto!(v0, system.initial_condition.velocity)
     return v0
 end
 

test/examples/gpu.jl

@@ -482,6 +482,31 @@ end
         end
     end
 
+    @testset verbose=true "DEM" begin
+        @trixi_testset "dem/rectangular_tank_2d.jl" begin
+            @trixi_test_nowarn trixi_include_changeprecision(Float32, @__MODULE__,
+                                                             joinpath(examples_dir(), "dem",
+                                                                      "rectangular_tank_2d.jl"),
+                                                             ode=nothing, sol=nothing)
+            # Neighborhood search with `FullGridCellList` for GPU compatibility
+            min_corner = minimum(tank.boundary.coordinates, dims=2)
+            max_corner = maximum(tank.boundary.coordinates, dims=2)
+            cell_list = FullGridCellList(; min_corner, max_corner)
+            neighborhood_search = GridNeighborhoodSearch{2}(; cell_list,
+                                                            update_strategy=ParallelUpdate())
+
+            @trixi_test_nowarn trixi_include_changeprecision(Float32, @__MODULE__,
+                                                             joinpath(examples_dir(), "dem",
+                                                                      "rectangular_tank_2d.jl"),
+                                                             tspan=(0.0f0, 0.05f0),
+                                                             neighborhood_search=neighborhood_search,
+                                                             parallelization_backend=Main.parallelization_backend)
+            @test sol.retcode == ReturnCode.Success
+            backend = TrixiParticles.KernelAbstractions.get_backend(sol.u[end].x[1])
+            @test backend == Main.parallelization_backend
+        end
+    end
+
     @testset verbose=true "Postprocessing $TRIXIPARTICLES_TEST_" begin
         @testset verbose=true "Interpolation" begin
             # Run the dam break example to get a solution