From ca2169d6b1fbec8762736767f85e9c1215a60768 Mon Sep 17 00:00:00 2001
From: Benedict Geihe <bgeihe@uni-koeln.de>
Date: Tue, 12 Mar 2024 14:51:27 +0100
Subject: [PATCH] write mpi_rank per cell to hdf5 output

---
 src/meshes/mesh_io.jl     | 3 +++
 src/meshes/serial_tree.jl | 3 +++
 2 files changed, 6 insertions(+)

diff --git a/src/meshes/mesh_io.jl b/src/meshes/mesh_io.jl
index 28e6efa8c57..baff9b12a6c 100644
--- a/src/meshes/mesh_io.jl
+++ b/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
diff --git a/src/meshes/serial_tree.jl b/src/meshes/serial_tree.jl
index 143ac19f6ee..2568663ccaa 100644
--- a/src/meshes/serial_tree.jl
+++ b/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)