write mpi_rank per cell to hdf5 output

src/meshes/mesh_io.jl

@@ -44,6 +44,7 @@ function save_mesh_file(mesh::TreeMesh, output_directory, timestep,
         file["neighbor_ids"] = @view mesh.tree.neighbor_ids[:, 1:n_cells]
         file["levels"] = @view mesh.tree.levels[1:n_cells]
         file["coordinates"] = @view mesh.tree.coordinates[:, 1:n_cells]
+        file["mpi_ranks"] = @view mesh.tree.mpi_ranks[1:n_cells]
     end
 
     return filename
@@ -88,6 +89,7 @@ function save_mesh_file(mesh::TreeMesh, output_directory, timestep,
         file["neighbor_ids"] = @view mesh.tree.neighbor_ids[:, 1:n_cells]
         file["levels"] = @view mesh.tree.levels[1:n_cells]
         file["coordinates"] = @view mesh.tree.coordinates[:, 1:n_cells]
+        file["mpi_ranks"] = @view mesh.tree.mpi_ranks[1:n_cells]
     end
 
     return filename
@@ -346,6 +348,7 @@ function load_mesh!(mesh::SerialTreeMesh, mesh_file::AbstractString)
         mesh.tree.neighbor_ids[:, 1:n_cells] = read(file["neighbor_ids"])
         mesh.tree.levels[1:n_cells] = read(file["levels"])
         mesh.tree.coordinates[:, 1:n_cells] = read(file["coordinates"])
+        mesh.tree.mpi_ranks[1:n_cells] = read(file["mpi_ranks"])
     end
 
     return mesh

src/meshes/serial_tree.jl

@@ -32,6 +32,7 @@ mutable struct SerialTree{NDIMS} <: AbstractTree{NDIMS}
     levels::Vector{Int}
     coordinates::Matrix{Float64}
     original_cell_ids::Vector{Int}
+    mpi_ranks::Vector{Int}
 
     capacity::Int
     length::Int
@@ -56,6 +57,7 @@ mutable struct SerialTree{NDIMS} <: AbstractTree{NDIMS}
         t.levels = fill(typemin(Int), capacity + 1)
         t.coordinates = fill(NaN, NDIMS, capacity + 1)
         t.original_cell_ids = fill(typemin(Int), capacity + 1)
+        t.mpi_ranks = fill(typemin(Int), capacity + 1)
 
         t.capacity = capacity
         t.length = 0
@@ -104,6 +106,7 @@ function init!(t::SerialTree, center::AbstractArray{Float64}, length::Real,
     t.levels[1] = 0
     set_cell_coordinates!(t, t.center_level_0, 1)
     t.original_cell_ids[1] = 0
+    t.mpi_ranks[1] = 0
 
     # Set neighbor ids: for each periodic direction, the level-0 cell is its own neighbor
     if all(periodicity)