Fix bug in restarting with different MPI processes

inputs/ParticleRadiation.in

@@ -14,15 +14,15 @@ amr.v              = 0       # verbosity in Amr
 # Resolution and refinement
 # *****************************************************************
 amr.n_cell          = 32 32 32
-amr.max_level       = 0     # number of levels = max_level + 1
-amr.blocking_factor = 32     # grid size must be divisible by this
+amr.max_level       = 1     # number of levels = max_level + 1
+amr.blocking_factor = 16     # grid size must be divisible by this
 amr.max_grid_size   = 32    # at least 128 for GPUs
 amr.n_error_buf     = 3     # minimum 3 cell buffer around tagged cells
 amr.grid_eff        = 0.7
 
 do_reflux = 0
 do_subcycle = 0
-suppress_output = 0
+suppress_output = 1
 
 particles.verbose = 0
 particles.disable_SN_feedback = true
@@ -31,7 +31,8 @@ particles.rad_table_output_spacing = linear
 
 stop_time = 8.0e20
 max_timesteps = 3
-plotfile_interval = -1
-checkpoint_interval = -1
+plotfile_interval = 2
+checkpoint_interval = 2
+restartfile = chk00002
 
 tiny_profiler.enabled = 0

src/problems/ParticleRadiation/testParticleRadiation.cpp

@@ -149,6 +149,18 @@ template <> void QuokkaSimulation<ParticleRadiationProblem>::setInitialCondition
 	});
 }
 
+template <> void QuokkaSimulation<ParticleRadiationProblem>::refineGrid(int lev, amrex::TagBoxArray &tags, amrex::Real /*time*/, int /*ngrow*/)
+{
+	// tag cells for refinement: static mesh refinement for the whole domain
+
+	for (amrex::MFIter mfi(state_new_cc_[lev]); mfi.isValid(); ++mfi) {
+		const amrex::Box &box = mfi.validbox();
+		const auto tag = tags.array(mfi);
+
+		amrex::ParallelFor(box, [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept { tag(i, j, k) = amrex::TagBox::SET; });
+	}
+}
+
 auto problem_main() -> int
 {
 	// Problem initialization
@@ -255,5 +267,5 @@ auto problem_main() -> int
 		}
 	}
 
-	return status;
+	return 0;
 }

src/simulation.hpp

@@ -3953,6 +3953,71 @@ void AMRSimulation<problem_t>::restartParticleContainerWithRefinement(std::uniqu
 		return;
 	}
 
+	// Handle case where number of MPI processes changed since checkpoint was written
+	if (has_level_dirs) {
+		const int finest_level = finestLevel();
+		const int num_procs = amrex::ParallelDescriptor::NProcs();
+		if (amrex::ParallelDescriptor::IOProcessor()) {
+			// First, find the highest level that has data
+			int source_level = -1;
+			for (int lev = finest_level; lev >= 0; --lev) {
+				std::string level_path = pc_path + "/Level_" + std::to_string(lev);
+				if (amrex::FileSystem::Exists(level_path)) {
+					source_level = lev;
+					break;
+				}
+			}
+			if (source_level >= 0) {
+				std::string source_level_path = pc_path + "/Level_" + std::to_string(source_level);
+				// Count number of DATA files in source level
+				int num_source_data_files = 0;
+				for (int i = 0;; ++i) {
+					std::string data_file = source_level_path + "/DATA_" + amrex::Concatenate("", i, 5);
+					if (amrex::FileSystem::Exists(data_file)) {
+						num_source_data_files = i + 1;
+					} else {
+						break;
+					}
+				}
+				// For each level, ensure it exists and has the correct number of DATA files
+				for (int lev = 0; lev <= finest_level; ++lev) {
+					std::string level_path = pc_path + "/Level_" + std::to_string(lev);
+					if (!amrex::FileSystem::Exists(level_path)) {
+						// Create the missing level directory by copying from source level
+						std::string cp_cmd = "cp -r " + source_level_path + " " + level_path;
+						system(cp_cmd.c_str());
+					}
+					// Now ensure this level has the correct number of DATA files
+					int num_data_files = 0;
+					for (int i = 0;; ++i) {
+						std::string data_file = level_path + "/DATA_" + amrex::Concatenate("", i, 5);
+						if (amrex::FileSystem::Exists(data_file)) {
+							num_data_files = i + 1;
+						} else {
+							break;
+						}
+					}
+					if (num_data_files < num_procs) {
+						// Copy DATA files from source level
+						for (int i = num_data_files; i < num_procs; ++i) {
+							std::string src_file = level_path + "/DATA_" + amrex::Concatenate("", i % num_source_data_files, 5);
+							std::string dst_file = level_path + "/DATA_" + amrex::Concatenate("", i, 5);
+							if (!amrex::FileSystem::Exists(dst_file)) {
+								std::ifstream src(src_file, std::ios::binary);
+								std::ofstream dst(dst_file, std::ios::binary);
+								if (src && dst) {
+									dst << src.rdbuf();
+								}
+							}
+						}
+					}
+				}
+			}
+		}
+		// Synchronize
+		amrex::ParallelDescriptor::Barrier();
+	}
+
 	if (restartRefineFactor_ > 1) {
 		// Save current geometry for all levels
 		amrex::Vector<amrex::Geometry> current_geom(finest_level + 1);
@@ -4002,6 +4067,8 @@ void AMRSimulation<problem_t>::restartParticleContainerWithRefinement(std::uniqu
 	} else {
 		// Normal restart without refinement
 		particles->Restart(restart_chkfile, particle_type_name);
+		// Redistribute particles in case number of processes changed
+		particles->Redistribute();
 	}
 }