Merge branch 'main' into plot-recipe

src/general/semidiscretization.jl

@@ -36,36 +36,44 @@ struct Semidiscretization{S, RU, RV, NS}
     ranges_v              :: RV
     neighborhood_searches :: NS
 
-    function Semidiscretization(systems...; neighborhood_search=GridNeighborhoodSearch,
-                                periodic_box_min_corner=nothing,
-                                periodic_box_max_corner=nothing)
-        # Check e.g. that the boundary systems are using a state equation if EDAC is not used.
-        # Other checks might be added here later.
-        check_configuration(systems)
-
-        sizes_u = [u_nvariables(system) * n_moving_particles(system)
-                   for system in systems]
-        ranges_u = Tuple((sum(sizes_u[1:(i - 1)]) + 1):sum(sizes_u[1:i])
-                         for i in eachindex(sizes_u))
-        sizes_v = [v_nvariables(system) * n_moving_particles(system)
-                   for system in systems]
-        ranges_v = Tuple((sum(sizes_v[1:(i - 1)]) + 1):sum(sizes_v[1:i])
-                         for i in eachindex(sizes_v))
-
-        # Create (and initialize) a tuple of n neighborhood searches for each of the n systems
-        # We will need one neighborhood search for each pair of systems.
-        searches = Tuple(Tuple(create_neighborhood_search(system, neighbor,
-                                                          Val(neighborhood_search),
-                                                          periodic_box_min_corner,
-                                                          periodic_box_max_corner)
-                               for neighbor in systems)
-                         for system in systems)
-
+    # Dispatch at `systems` to distinguish this constructor from the one below when
+    # 4 systems are passed.
+    # This is an internal constructor only used in `test/count_allocations.jl`.
+    function Semidiscretization(systems::Tuple, ranges_u, ranges_v, neighborhood_searches)
         new{typeof(systems), typeof(ranges_u),
-            typeof(ranges_v), typeof(searches)}(systems, ranges_u, ranges_v, searches)
+            typeof(ranges_v), typeof(neighborhood_searches)}(systems, ranges_u, ranges_v,
+                                                             neighborhood_searches)
     end
 end
 
+function Semidiscretization(systems...; neighborhood_search=GridNeighborhoodSearch,
+                            periodic_box_min_corner=nothing,
+                            periodic_box_max_corner=nothing)
+    # Check e.g. that the boundary systems are using a state equation if EDAC is not used.
+    # Other checks might be added here later.
+    check_configuration(systems)
+
+    sizes_u = [u_nvariables(system) * n_moving_particles(system)
+               for system in systems]
+    ranges_u = Tuple((sum(sizes_u[1:(i - 1)]) + 1):sum(sizes_u[1:i])
+                     for i in eachindex(sizes_u))
+    sizes_v = [v_nvariables(system) * n_moving_particles(system)
+               for system in systems]
+    ranges_v = Tuple((sum(sizes_v[1:(i - 1)]) + 1):sum(sizes_v[1:i])
+                     for i in eachindex(sizes_v))
+
+    # Create (and initialize) a tuple of n neighborhood searches for each of the n systems
+    # We will need one neighborhood search for each pair of systems.
+    searches = Tuple(Tuple(create_neighborhood_search(system, neighbor,
+                                                      Val(neighborhood_search),
+                                                      periodic_box_min_corner,
+                                                      periodic_box_max_corner)
+                           for neighbor in systems)
+                     for system in systems)
+
+    return Semidiscretization(systems, ranges_u, ranges_v, searches)
+end
+
 # Inline show function e.g. Semidiscretization(neighborhood_search=...)
 function Base.show(io::IO, semi::Semidiscretization)
     @nospecialize semi # reduce precompilation time

test/count_allocations.jl

@@ -0,0 +1,75 @@
+# Wrapper for any neighborhood search that forwards `for_particle_neighbor` to the wrapped
+# neighborhood search, but doesn't do anything in the update step.
+# This is used in the example tests to test for zero allocations in the `kick!` function.
+struct NoUpdateNeighborhoodSearch{NHS}
+    nhs::NHS
+end
+
+# Copy a `Semidiscretization`, but wrap the neighborhood searches with
+# `NoUpdateNeighborhoodSearch`.
+function copy_semi_with_no_update_nhs(semi)
+    neighborhood_searches = Tuple(Tuple(NoUpdateNeighborhoodSearch(nhs)
+                                        for nhs in searches)
+                                  for searches in semi.neighborhood_searches)
+
+    return Semidiscretization(semi.systems, semi.ranges_u, semi.ranges_v,
+                              neighborhood_searches)
+end
+
+# Forward `for_particle_neighbor` to wrapped neighborhood search
+@inline function TrixiParticles.for_particle_neighbor(f, system_coords, neighbor_coords,
+                                                      neighborhood_search::NoUpdateNeighborhoodSearch;
+                                                      particles=axes(system_coords, 2),
+                                                      parallel=true)
+    TrixiParticles.for_particle_neighbor(f, system_coords, neighbor_coords,
+                                         neighborhood_search.nhs,
+                                         particles=particles, parallel=parallel)
+end
+
+# No update
+@inline TrixiParticles.update!(search::NoUpdateNeighborhoodSearch, coords_fun) = search
+
+# Count allocations of one call to the right-hand side (`kick!` + `drift!`)
+function count_rhs_allocations(sol, semi)
+    t = sol.t[end]
+    v_ode, u_ode = sol.u[end].x
+    dv_ode = similar(v_ode)
+    du_ode = similar(u_ode)
+
+    # Wrap neighborhood searches to avoid counting alloctations in the NHS update
+    semi_no_nhs_update = copy_semi_with_no_update_nhs(semi)
+
+    try
+        # Disable timers, which cause extra allocations
+        TrixiParticles.TimerOutputs.disable_debug_timings(TrixiParticles)
+
+        # Disable multithreading, which causes extra allocations
+        return disable_polyester_threads() do
+            # We need `@invokelatest` here to ensure that the most recent method of
+            # `TrixiParticles.timeit_debug_enabled()` is called, which is redefined in
+            # `disable_debug_timings` above.
+            return @invokelatest count_rhs_allocations_inner(dv_ode, du_ode, v_ode, u_ode,
+                                                             semi_no_nhs_update, t)
+        end
+    finally
+        # Enable timers again
+        @invokelatest TrixiParticles.TimerOutputs.enable_debug_timings(TrixiParticles)
+    end
+end
+
+# Function barrier to avoid type instabilites with `semi_no_nhs_update`, which will
+# cause extra allocations.
+@inline function count_rhs_allocations_inner(dv_ode, du_ode, v_ode, u_ode,
+                                             semi_no_nhs_update, t)
+    # Run RHS once to avoid counting allocations from compilation
+    TrixiParticles.kick!(dv_ode, v_ode, u_ode, semi_no_nhs_update, t)
+    TrixiParticles.drift!(du_ode, v_ode, u_ode, semi_no_nhs_update, t)
+
+    # Count allocations
+    allocations_kick = @allocated TrixiParticles.kick!(dv_ode, v_ode, u_ode,
+                                                       semi_no_nhs_update, t)
+    allocations_drift = @allocated TrixiParticles.drift!(du_ode, v_ode, u_ode,
+                                                         semi_no_nhs_update, t)
+
+    return allocations_kick + allocations_drift
+end

test/examples/dam_break_2d_corrections.jl

@@ -70,6 +70,7 @@
                       fluid_density=fluid_density, density_diffusion=nothing)
 
         @test sol.retcode == ReturnCode.Success
+        @test count_rhs_allocations(sol, semi) == 0
     end
 
     @testset verbose=true "$correction_name" for correction_name in keys(correction_dict)
@@ -94,5 +95,6 @@
                       boundary_layers=5)
 
         @test sol.retcode == ReturnCode.Success
+        @test count_rhs_allocations(sol, semi) == 0
     end
 end

test/examples/examples.jl

@@ -9,13 +9,15 @@
             @test sol.retcode == ReturnCode.Success
             # This error varies between serial and multithreaded runs
             @test isapprox(error_A, 0.0001717690010767381, atol=5e-7)
+            @test count_rhs_allocations(sol, semi) == 0
         end
 
         @trixi_testset "fluid/hydrostatic_water_column_2d.jl" begin
             @test_nowarn_mod trixi_include(@__MODULE__,
                                            joinpath(examples_dir(), "fluid",
                                                     "hydrostatic_water_column_2d.jl"))
             @test sol.retcode == ReturnCode.Success
+            @test count_rhs_allocations(sol, semi) == 0
         end
 
         @trixi_testset "fluid/hydrostatic_water_column_2d.jl with SummationDensity" begin
@@ -25,6 +27,7 @@
                                            fluid_density_calculator=SummationDensity(),
                                            clip_negative_pressure=true)
             @test sol.retcode == ReturnCode.Success
+            @test count_rhs_allocations(sol, semi) == 0
         end
 
         @trixi_testset "fluid/hydrostatic_water_column_3d.jl" begin
@@ -33,6 +36,7 @@
                                                     "hydrostatic_water_column_3d.jl"),
                                            tspan=(0.0, 0.1))
             @test sol.retcode == ReturnCode.Success
+            @test count_rhs_allocations(sol, semi) == 0
         end
 
         @trixi_testset "fluid/hydrostatic_water_column_3d.jl with SummationDensity" begin
@@ -43,13 +47,15 @@
                                            fluid_density_calculator=SummationDensity(),
                                            clip_negative_pressure=true)
             @test sol.retcode == ReturnCode.Success
+            @test count_rhs_allocations(sol, semi) == 0
         end
 
         @trixi_testset "fluid/hydrostatic_water_column_edac_2d.jl" begin
             @test_nowarn_mod trixi_include(@__MODULE__,
                                            joinpath(examples_dir(), "fluid",
                                                     "hydrostatic_water_column_edac_2d.jl"))
             @test sol.retcode == ReturnCode.Success
+            @test count_rhs_allocations(sol, semi) == 0
         end
 
         @trixi_testset "fluid/dam_break_2d.jl" begin
@@ -58,6 +64,7 @@
                                                     "dam_break_2d.jl"),
                                            tspan=(0.0, 0.1))
             @test sol.retcode == ReturnCode.Success
+            @test count_rhs_allocations(sol, semi) == 0
         end
 
         @trixi_testset "fluid/dam_break_3d.jl" begin
@@ -66,6 +73,7 @@
                                                     "dam_break_3d.jl"),
                                            tspan=(0.0, 0.1), fluid_particle_spacing=0.1)
             @test sol.retcode == ReturnCode.Success
+            @test count_rhs_allocations(sol, semi) == 0
         end
 
         @trixi_testset "fluid/falling_water_column_2d.jl" begin
@@ -74,6 +82,7 @@
                                                     "falling_water_column_2d.jl"),
                                            tspan=(0.0, 0.4))
             @test sol.retcode == ReturnCode.Success
+            @test count_rhs_allocations(sol, semi) == 0
         end
 
         @trixi_testset "fluid/periodic_channel_2d.jl" begin
@@ -82,6 +91,7 @@
                                                     "periodic_channel_2d.jl"),
                                            tspan=(0.0, 0.4))
             @test sol.retcode == ReturnCode.Success
+            @test count_rhs_allocations(sol, semi) == 0
         end
 
         include("dam_break_2d_corrections.jl")
@@ -94,6 +104,7 @@
                                                     "oscillating_beam_2d.jl"),
                                            tspan=(0.0, 0.1))
             @test sol.retcode == ReturnCode.Success
+            @test count_rhs_allocations(sol, semi) == 0
         end
     end
 
@@ -104,6 +115,7 @@
                                                     "falling_water_column_2d.jl"),
                                            tspan=(0.0, 0.4))
             @test sol.retcode == ReturnCode.Success
+            @test count_rhs_allocations(sol, semi) == 0
         end
 
         @trixi_testset "fsi/dam_break_2d.jl" begin
@@ -115,6 +127,7 @@
                                            tspan=(0.0, 0.4),
                                            dtmax=1e-3)
             @test sol.retcode == ReturnCode.Success
+            @test count_rhs_allocations(sol, semi) == 0
         end
 
         @trixi_testset "fsi/dam_break_gate_2d.jl" begin
@@ -124,6 +137,7 @@
                                            tspan=(0.0, 0.4),
                                            dtmax=1e-3)
             @test sol.retcode == ReturnCode.Success
+            @test count_rhs_allocations(sol, semi) == 0
         end
 
         @trixi_testset "fsi/falling_spheres_2d.jl" begin
@@ -132,6 +146,7 @@
                                                     "falling_spheres_2d.jl"),
                                            tspan=(0.0, 1.0))
             @test sol.retcode == ReturnCode.Success
+            @test count_rhs_allocations(sol, semi) == 0
         end
     end
 
@@ -141,6 +156,7 @@
                                            joinpath(examples_dir(), "n_body",
                                                     "n_body_solar_system.jl"))
             @test sol.retcode == ReturnCode.Success
+            @test count_rhs_allocations(sol, semi) == 0
         end
 
         @trixi_testset "n_body/n_body_benchmark_trixi.jl" begin

test/rectangular_patch.jl

@@ -0,0 +1,47 @@
+# Rectangular patch of particles, optionally with a perturbation in position and/or quantities
+function rectangular_patch(particle_spacing, size; density=1000.0, pressure=0.0, seed=1,
+                           perturbation_factor=1.0, perturbation_factor_position=1.0,
+                           set_function=nothing, offset=ntuple(_ -> 0.0, length(size)))
+    # Fixed seed to ensure reproducibility
+    Random.seed!(seed)
+
+    # Center particle at the origin (assuming odd size)
+    min_corner = -particle_spacing / 2 .* size
+    ic = RectangularShape(particle_spacing, size, min_corner,
+                          density=density, pressure=pressure)
+
+    perturb!(ic.coordinates, perturbation_factor_position * 0.5 * particle_spacing)
+
+    # Don't perturb center particle position
+    center_particle = ceil(Int, prod(size) / 2)
+    ic.coordinates[:, center_particle] .= 0.0
+
+    for i in 1:Base.size(ic.coordinates, 2)
+        ic.coordinates[:, i] .+= offset
+    end
+
+    if set_function !== nothing
+        for i in 1:Base.size(ic.coordinates, 2)
+            coord = ic.coordinates[:, i]
+            ic.mass[i], ic.density[i], ic.pressure[i], ic.velocity[:, i] = set_function(coord)
+        end
+    end
+
+    if perturbation_factor > eps()
+        perturb!(ic.mass, perturbation_factor * 0.1 * ic.mass[1])
+        perturb!(ic.density, perturbation_factor * 0.1 * density)
+        perturb!(ic.pressure, perturbation_factor * 2000)
+        perturb!(ic.velocity, perturbation_factor * 0.5 * particle_spacing)
+    end
+
+    return ic
+end
+
+function perturb!(data, amplitude)
+    for i in eachindex(data)
+        # Perturbation in the interval (-amplitude, amplitude)
+        data[i] += 2 * amplitude * rand() - amplitude
+    end
+
+    return data
+end