From 5d9655256b408855124c8ef95092298b74bea29b Mon Sep 17 00:00:00 2001
From: YoungHui Park <ypark234@wisc.edu>
Date: Wed, 22 Jul 2020 04:54:35 -0500
Subject: [PATCH] Add new variables and functions to query global surface
 source bank size and index

---
 include/openmc/simulation.h |  6 ++++++
 src/simulation.cpp          | 26 ++++++++++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/include/openmc/simulation.h b/include/openmc/simulation.h
index 32e2c6730..8c4f044ce 100644
--- a/include/openmc/simulation.h
+++ b/include/openmc/simulation.h
@@ -45,6 +45,9 @@ extern const RegularMesh* ufs_mesh;
 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 std::vector<int64_t> surf_src_index;
+
 } // namespace simulation
 
 //==============================================================================
@@ -60,6 +63,9 @@ void calculate_work();
 //! 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 e8dc30fcc..2ff136527 100644
--- a/src/simulation.cpp
+++ b/src/simulation.cpp
@@ -270,6 +270,8 @@ const RegularMesh* ufs_mesh {nullptr};
 std::vector<double> k_generation;
 std::vector<int64_t> work_index;
 
+int64_t total_surf_banks;
+std::vector<int64_t> surf_src_index;
 
 } // namespace simulation
 
@@ -385,6 +387,9 @@ void finalize_batch()
 
   // Write out surface source if requested.
   if (settings::surface_source && simulation::current_batch == settings::n_batches) {
+
+    query_surf_src_size();
+
     auto filename = settings::path_output + "surface_source.h5";
     write_source_point(filename.c_str(), true);  //!!
   }
@@ -630,6 +635,27 @@ void initialize_data()
     data::energy_min[neutron]) / settings::n_log_bins;
 }
 
+void query_surf_src_size()
+{
+  int64_t i_bank = 0;
+  simulation::surf_src_index.resize(mpi::n_procs + 1);
+  simulation::surf_src_index[0] = 0;
+  for (int i = 0; i < mpi::n_procs; ++i) {
+    int64_t work_i = 0;
+
+    // Set number of particles
+    if (mpi::rank == i) work_i = simulation::surf_src_bank.size();
+#ifdef OPENMC_MPI
+    MPI_Reduce(&work_i, &simulation::total_surf_banks, 1, MPI_INT,
+               MPI_SUM, 0, mpi::intracomm);
+#endif
+
+    // Set index into source bank for rank i
+    i_bank += work_i;
+    simulation::surf_src_index[i + 1] = i_bank;
+  }
+}
+
 #ifdef OPENMC_MPI
 void broadcast_results() {
   // Broadcast tally results so that each process has access to results