Implement AMR with MPI

examples/2d/elixir_advection_amr_coarsen_once.jl

@@ -4,8 +4,10 @@ using OrdinaryDiffEq
 using Trixi
 
 
-# define new structs inside a module to allow re-evaluating the file
-module TrixiExtension
+# Define new structs inside a module to allow re-evaluating the file.
+# This module name needs to be unique among all examples, otherwise Julia will throw warnings 
+# if multiple test cases using the same module name are run in the same session.
+module TrixiExtensionCoarsen
 
 using Trixi
 
@@ -34,9 +36,9 @@ function (indicator::IndicatorAlwaysCoarsen)(u::AbstractArray{<:Any,4},
   return alpha
 end
 
-end # module TrixiExtension
+end # module TrixiExtensionCoarsen
 
-import .TrixiExtension
+import .TrixiExtensionCoarsen
 
 ###############################################################################
 # semidiscretization of the linear advection equation
@@ -83,7 +85,7 @@ save_solution = SaveSolutionCallback(interval=100,
                                      solution_variables=cons2prim)
 
 
-amr_controller = ControllerThreeLevel(semi, TrixiExtension.IndicatorAlwaysCoarsen(semi),
+amr_controller = ControllerThreeLevel(semi, TrixiExtensionCoarsen.IndicatorAlwaysCoarsen(semi),
                                       base_level=2, max_level=2,
                                       med_threshold=0.1, max_threshold=0.6)
 amr_callback = AMRCallback(semi, amr_controller,

examples/2d/elixir_advection_amr_refine_twice.jl

@@ -4,8 +4,10 @@ using OrdinaryDiffEq
 using Trixi
 
 
-# define new structs inside a module to allow re-evaluating the file
-module TrixiExtension
+# Define new structs inside a module to allow re-evaluating the file.
+# This module name needs to be unique among all examples, otherwise Julia will throw warnings 
+# if multiple test cases using the same module name are run in the same session.
+module TrixiExtensionRefine
 
 using Trixi
 
@@ -34,9 +36,9 @@ function (indicator::IndicatorAlwaysRefine)(u::AbstractArray{<:Any,4},
   return alpha
 end
 
-end # module TrixiExtension
+end # module TrixiExtensionRefine
 
-import .TrixiExtension
+import .TrixiExtensionRefine
 
 ###############################################################################
 # semidiscretization of the linear advection equation
@@ -83,7 +85,7 @@ save_solution = SaveSolutionCallback(interval=100,
                                      solution_variables=cons2prim)
 
 
-amr_controller = ControllerThreeLevel(semi, TrixiExtension.IndicatorAlwaysRefine(semi),
+amr_controller = ControllerThreeLevel(semi, TrixiExtensionRefine.IndicatorAlwaysRefine(semi),
                                       base_level=4, max_level=4,
                                       med_threshold=0.1, max_threshold=0.6)
 amr_callback = AMRCallback(semi, amr_controller,

src/callbacks_step/amr.jl

@@ -173,7 +173,7 @@ function (amr_callback::AMRCallback)(u_ode::AbstractVector, mesh::TreeMesh,
 
   if mpi_isparallel()
     # Collect lambda for all elements
-    lambda_global = Vector{eltype(lambda)}(undef, cache.mpi_cache.n_elements_global)
+    lambda_global = Vector{eltype(lambda)}(undef, nelementsglobal(dg, cache))
     # Use parent because n_elements_by_rank is an OffsetArray
     recvbuf = MPI.VBuffer(lambda_global, parent(cache.mpi_cache.n_elements_by_rank))
     MPI.Allgatherv!(lambda, recvbuf, mpi_comm())

src/callbacks_step/save_restart_dg.jl

@@ -148,7 +148,7 @@ function load_restart_file(mesh::ParallelTreeMesh, equations, dg::DG, cache, res
   if !mpi_isroot()
     # Receive nodal data from root
     for v in eachvariable(equations)
-      u[v, .., :] = MPI.Scatterv(eltype(u)[], node_counts, mpi_root(), mpi_comm())
+      MPI.Scatterv!(nothing, @view(u[v, .., :]), mpi_root(), mpi_comm())
     end
 
     return u_ode
@@ -180,7 +180,8 @@ function load_restart_file(mesh::ParallelTreeMesh, equations, dg::DG, cache, res
 
       # Read variable
       println("Reading variables_$v ($name)...")
-      u[v, .., :] = MPI.Scatterv(read(file["variables_$v"]), node_counts, mpi_root(), mpi_comm())
+      sendbuf = MPI.VBuffer(read(file["variables_$v"]), node_counts)
+      MPI.Scatterv!(sendbuf, @view(u[v, .., :]), mpi_root(), mpi_comm())
     end
   end
 

src/callbacks_step/save_solution_dg.jl

@@ -122,7 +122,7 @@ function save_solution_file(u, time, dt, timestep,
     attributes(file)["equations"] = get_name(equations)
     attributes(file)["polydeg"] = polydeg(dg)
     attributes(file)["n_vars"] = n_vars
-    attributes(file)["n_elements"] = cache.mpi_cache.n_elements_global
+    attributes(file)["n_elements"] = nelementsglobal(dg, cache)
     attributes(file)["mesh_file"] = splitdir(mesh.current_filename)[2]
     attributes(file)["time"] = convert(Float64, time) # Ensure that `time` is written as a double precision scalar
     attributes(file)["dt"] = convert(Float64, dt) # Ensure that `dt` is written as a double precision scalar

src/mesh/parallel.jl

@@ -6,10 +6,10 @@ based on leaf cell count and tree structure.
 If `allow_coarsening` is `true`, the algorithm will keep leaf cells together
 on one rank when needed for local coarsening (i.e. when all children of a cell are leaves).
 """
-function partition!(mesh::ParallelTreeMesh, allow_coarsening=true)
+function partition!(mesh::ParallelTreeMesh; allow_coarsening=true)
   # Determine number of leaf cells per rank
   leaves = leaf_cells(mesh.tree)
-  @assert length(leaves) > mpi_nranks()
+  @assert length(leaves) > mpi_nranks() "Too many ranks to properly partition the mesh!"
   n_leaves_per_rank = OffsetArray(fill(div(length(leaves), mpi_nranks()), mpi_nranks()),
                                   0:(mpi_nranks() - 1))
   for d in 0:(rem(length(leaves), mpi_nranks()) - 1)
@@ -49,6 +49,8 @@ function partition!(mesh::ParallelTreeMesh, allow_coarsening=true)
       end
     end
 
+    @assert all(n_leaves_per_rank .> 0) "Too many ranks to properly partition the mesh!"
+
     mesh.n_cells_by_rank[d] = last_id - mesh.first_cell_by_rank[d] + 1
     mesh.tree.mpi_ranks[mesh.first_cell_by_rank[d]:last_id] .= d
 

src/solvers/dg/dg.jl

@@ -161,6 +161,7 @@ end
 
 @inline nnodes(dg::DG)             = nnodes(dg.basis)
 @inline nelements(dg::DG, cache)   = nelements(cache.elements)
+@inline nelementsglobal(dg::DG, cache) = mpi_isparallel() ? cache.mpi_cache.n_elements_global : nelements(dg, cache)
 @inline ninterfaces(dg::DG, cache) = ninterfaces(cache.interfaces)
 @inline nboundaries(dg::DG, cache) = nboundaries(cache.boundaries)
 @inline nmortars(dg::DG, cache)    = nmortars(cache.mortars)

test/Project.toml

@@ -5,6 +5,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+SimpleMock = "a896ed2c-15a5-4479-b61d-a0e88e2a1d25"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]

test/test_manual.jl

@@ -1,6 +1,7 @@
 module TestManual
 
 using Test
+using SimpleMock
 using Documenter
 using Trixi
 
@@ -59,6 +60,116 @@ isdir(outdir) && rm(outdir, recursive=true)
     end
   end
 
+  @testset "ParallelTreeMesh" begin
+    @testset "partition!" begin
+      @testset "mpi_nranks() = 2" begin
+        mock((Trixi.mpi_nranks) => () -> 2) do _
+          mesh = TreeMesh{2, Trixi.ParallelTree{2}}(30, (0.0, 0.0), 1)
+          # Refine twice
+          Trixi.refine!(mesh.tree)
+          Trixi.refine!(mesh.tree)
+
+          # allow_coarsening = true
+          Trixi.partition!(mesh)
+          # Use parent for OffsetArray
+          @test parent(mesh.n_cells_by_rank) == [11, 10]
+          @test mesh.tree.mpi_ranks[1:21] == 
+              [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
+          @test parent(mesh.first_cell_by_rank) == [1, 12]
+
+          # allow_coarsening = false
+          Trixi.partition!(mesh; allow_coarsening=false)
+          @test parent(mesh.n_cells_by_rank) == [11, 10]
+          @test mesh.tree.mpi_ranks[1:21] == 
+              [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
+          @test parent(mesh.first_cell_by_rank) == [1, 12]
+        end
+      end
+
+      @testset "mpi_nranks() = 3" begin
+        mock((Trixi.mpi_nranks) => () -> 3) do _
+          mesh = TreeMesh{2, Trixi.ParallelTree{2}}(100, (0.0, 0.0), 1)
+          # Refine twice
+          Trixi.refine!(mesh.tree)
+          Trixi.refine!(mesh.tree)
+
+          # allow_coarsening = true
+          Trixi.partition!(mesh)
+          # Use parent for OffsetArray
+          @test parent(mesh.n_cells_by_rank) == [11, 5, 5]
+          @test mesh.tree.mpi_ranks[1:21] == 
+              [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2]
+          @test parent(mesh.first_cell_by_rank) == [1, 12, 17]
+
+          # allow_coarsening = false
+          Trixi.partition!(mesh; allow_coarsening=false)
+          @test parent(mesh.n_cells_by_rank) == [9, 6, 6]
+          @test mesh.tree.mpi_ranks[1:21] == 
+              [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2]
+          @test parent(mesh.first_cell_by_rank) == [1, 10, 16]
+        end
+      end
+
+      @testset "mpi_nranks() = 9" begin
+        mock((Trixi.mpi_nranks) => () -> 9) do _
+          mesh = TreeMesh{2, Trixi.ParallelTree{2}}(1000, (0.0, 0.0), 1)
+          # Refine twice
+          Trixi.refine!(mesh.tree)
+          Trixi.refine!(mesh.tree)
+          Trixi.refine!(mesh.tree)
+          Trixi.refine!(mesh.tree)
+
+          # allow_coarsening = true
+          Trixi.partition!(mesh)
+          # Use parent for OffsetArray
+          @test parent(mesh.n_cells_by_rank) == [44, 37, 38, 37, 37, 37, 38, 37, 36]
+          @test parent(mesh.first_cell_by_rank) == [1, 45, 82, 120, 157, 194, 231, 269, 306]
+        end
+      end
+
+      @testset "mpi_nranks() = 3 non-uniform" begin
+        mock((Trixi.mpi_nranks) => () -> 3) do _
+          mesh = TreeMesh{2, Trixi.ParallelTree{2}}(100, (0.0, 0.0), 1)
+          # Refine whole tree
+          Trixi.refine!(mesh.tree)
+          # Refine left leaf
+          Trixi.refine!(mesh.tree, [2])
+
+          # allow_coarsening = true
+          Trixi.partition!(mesh)
+          # Use parent for OffsetArray
+          @test parent(mesh.n_cells_by_rank) == [6, 1, 2]
+          @test mesh.tree.mpi_ranks[1:9] == [0, 0, 0, 0, 0, 0, 1, 2, 2]
+          @test parent(mesh.first_cell_by_rank) == [1, 7, 8]
+
+          # allow_coarsening = false
+          Trixi.partition!(mesh; allow_coarsening=false)
+          @test parent(mesh.n_cells_by_rank) == [5, 2, 2]
+          @test mesh.tree.mpi_ranks[1:9] == [0, 0, 0, 0, 0, 1, 1, 2, 2]
+          @test parent(mesh.first_cell_by_rank) == [1, 6, 8]
+        end
+      end
+
+      @testset "not enough ranks" begin
+        mock((Trixi.mpi_nranks) => () -> 3) do _
+          mesh = TreeMesh{2, Trixi.ParallelTree{2}}(100, (0.0, 0.0), 1)
+
+          # Only one leaf
+          @test_throws AssertionError(
+            "Too many ranks to properly partition the mesh!") Trixi.partition!(mesh)
+
+          # Refine to 4 leaves
+          Trixi.refine!(mesh.tree)
+
+          # All four leaves will need to be on one rank to allow coarsening
+          @test_throws AssertionError(
+            "Too many ranks to properly partition the mesh!") Trixi.partition!(mesh)
+          @test_nowarn Trixi.partition!(mesh; allow_coarsening=false)
+        end
+      end
+    end
+  end
+
   @testset "interpolation" begin
     @testset "nodes and weights" begin
       @test Trixi.gauss_nodes_weights(1) == ([0.0], [2.0])