From 09863806df4563d5e55bb2fe3374e611cbe0051f Mon Sep 17 00:00:00 2001
From: YoungHui Park <ypark234@wisc.edu>
Date: Wed, 29 Jul 2020 16:29:45 -0500
Subject: [PATCH] Add comments and minor changes

---
 include/openmc/simulation.h |  8 ++---
 src/simulation.cpp          | 72 ++++++++++++++++++++-----------------
 src/state_point.cpp         |  4 +++
 3 files changed, 47 insertions(+), 37 deletions(-)

diff --git a/include/openmc/simulation.h b/include/openmc/simulation.h
index 9422590cf..bd10c96c7 100644
--- a/include/openmc/simulation.h
+++ b/include/openmc/simulation.h
@@ -46,7 +46,7 @@ extern std::vector<double> k_generation;
 extern std::vector<int64_t> work_index;
 
 extern int64_t total_surf_banks;  //!< Total number of surface source banks
-extern int64_t max_bank_size;     //!< Maximum bank size for MPI
+extern int64_t max_bank_size;     //!< Maximum bank size from a process
 extern std::vector<int64_t> surf_src_index;
 
 } // namespace simulation
@@ -61,12 +61,12 @@ void allocate_banks();
 //! Determine number of particles to transport per process
 void calculate_work();
 
+//! Determine number of surface source banks per process and their sum
+void query_surf_src_size();
+
 //! Initialize nuclear data before a simulation
 void initialize_data();
 
-//! Get the total number of surface source banks and populate surf_src_index
-void query_surf_src_size();
-
 //! Initialize a batch
 void initialize_batch();
 
diff --git a/src/simulation.cpp b/src/simulation.cpp
index f56a543bc..b71de773b 100644
--- a/src/simulation.cpp
+++ b/src/simulation.cpp
@@ -570,6 +570,45 @@ void calculate_work()
   }
 }
 
+void query_surf_src_size()
+{
+  int64_t total;
+  if (mpi::master) {
+    simulation::surf_src_index.resize(mpi::n_procs + 1);
+    simulation::surf_src_index[0] = 0;
+  }
+
+#ifdef OPENMC_MPI
+  std::vector<int64_t> bank_size;
+  bank_size.resize(mpi::n_procs);
+
+  // Collect the number of surface source banks from all processes
+  int64_t size = simulation::surf_src_bank.size();
+  MPI_Gather(&size, 1, MPI_INT64_T,
+             bank_size.data(), 1, MPI_INT64_T,
+             0, mpi::intracomm);
+
+  if (mpi::master) {
+    // Populate the surf_src_index with cumulative sum of the number of
+    // surface source banks per process
+    for (int i = 1; i < mpi::n_procs + 1; ++i) {
+      simulation::surf_src_index[i] = \
+        simulation::surf_src_index[i - 1] + bank_size[i - 1];
+    }
+    // Set maximum bank size
+    simulation::max_bank_size = *std::max_element(bank_size.begin(),
+                                                  bank_size.end());
+    total = simulation::surf_src_index[mpi::n_procs];
+  }
+#else
+  total = simulation::surf_src_bank.size();
+  simulation::surf_src_index[mpi::n_procs] = total;
+#endif
+  // Set total number of surface source banks
+  simulation::total_surf_banks = total;
+
+}
+
 void initialize_data()
 {
   // Determine minimum/maximum energy for incident neutron/photon data
@@ -639,39 +678,6 @@ void initialize_data()
     data::energy_min[neutron]) / settings::n_log_bins;
 }
 
-void query_surf_src_size()
-{
-  int64_t total;
-  if (mpi::master) {
-    simulation::surf_src_index.resize(mpi::n_procs + 1);
-    simulation::surf_src_index[0] = 0;
-  }
-
-#ifdef OPENMC_MPI
-  std::vector<int64_t> bank_size;
-  bank_size.resize(mpi::n_procs);
-
-  int64_t size = simulation::surf_src_bank.size();
-  MPI_Gather(&size, 1, MPI_INT64_T,
-             bank_size.data(), 1, MPI_INT64_T,
-             0, mpi::intracomm);
-
-  if (mpi::master) {
-    for (int i = 1; i < mpi::n_procs + 1; ++i) {
-      simulation::surf_src_index[i] = simulation::surf_src_index[i - 1] + bank_size[i - 1];
-    }
-    simulation::max_bank_size = *std::max_element(bank_size.begin(), bank_size.end());
-    total = simulation::surf_src_index[mpi::n_procs];
-  }
-#else
-  total = simulation::surf_src_bank.size();
-  simulation::surf_src_index[mpi::n_procs] = total;
-#endif
-
-  simulation::total_surf_banks = total;
-
-}
-
 #ifdef OPENMC_MPI
 void broadcast_results() {
   // Broadcast tally results so that each process has access to results
diff --git a/src/state_point.cpp b/src/state_point.cpp
index 9740220a8..735601b55 100644
--- a/src/state_point.cpp
+++ b/src/state_point.cpp
@@ -566,11 +566,15 @@ write_source_bank(hid_t group_id, bool surf_src_bank)
 {
   hid_t banktype = h5banktype();
 
+  // Set total and individual process dataspace sizes for source bank
   int64_t dims_size = settings::n_particles;
   int64_t count_size = simulation::work_per_rank;
+
+  // Set vectors for source bank and starting bank index of each process
   std::vector<int64_t> wi = simulation::work_index;
   std::vector<Particle::Bank> src_bank = simulation::source_bank;
 
+  // Reset dataspace sizes and vectors for surface source bank
   if (surf_src_bank) {
     dims_size = simulation::total_surf_banks;
     count_size = simulation::surf_src_bank.size();