From b271399f638eef3919af1bef5024c9be6d016add Mon Sep 17 00:00:00 2001
From: Paul Romano <paul.k.romano@gmail.com>
Date: Mon, 26 Oct 2020 16:54:40 -0500
Subject: [PATCH] Allow read_source_bank to be distributed or not

---
 include/openmc/state_point.h |  2 +-
 src/source.cpp               |  4 ++--
 src/state_point.cpp          | 37 ++++++++++++++++++++----------------
 3 files changed, 24 insertions(+), 19 deletions(-)

diff --git a/include/openmc/state_point.h b/include/openmc/state_point.h
index 31dffc9ee..c50ca355b 100644
--- a/include/openmc/state_point.h
+++ b/include/openmc/state_point.h
@@ -14,7 +14,7 @@ namespace openmc {
 void load_state_point();
 void write_source_point(const char* filename);
 void write_source_bank(hid_t group_id);
-void read_source_bank(hid_t group_id, std::vector<Particle::Bank>& sites);
+void read_source_bank(hid_t group_id, std::vector<Particle::Bank>& sites, bool distribute);
 void write_tally_results_nr(hid_t file_id);
 void restart_set_keff();
 
diff --git a/src/source.cpp b/src/source.cpp
index fd2d6cd8f..a7bb80747 100644
--- a/src/source.cpp
+++ b/src/source.cpp
@@ -254,8 +254,8 @@ FileSource::FileSource(std::string path)
     fatal_error("Specified starting source file not a source file type.");
   }
 
-  // Read in the source bank
-  read_source_bank(file_id, sites_);
+  // Read in the source particles
+  read_source_bank(file_id, sites_, false);
 
   // Close file
   file_close(file_id);
diff --git a/src/state_point.cpp b/src/state_point.cpp
index 6fcc900a6..66706a676 100644
--- a/src/state_point.cpp
+++ b/src/state_point.cpp
@@ -489,7 +489,7 @@ void load_state_point()
     }
 
     // Read source
-    read_source_bank(file_id, simulation::source_bank);
+    read_source_bank(file_id, simulation::source_bank, true);
 
   }
 
@@ -653,7 +653,7 @@ write_source_bank(hid_t group_id)
 }
 
 
-void read_source_bank(hid_t group_id, std::vector<Particle::Bank>& sites)
+void read_source_bank(hid_t group_id, std::vector<Particle::Bank>& sites, bool distribute)
 {
   hid_t banktype = h5banktype();
 
@@ -663,22 +663,27 @@ void read_source_bank(hid_t group_id, std::vector<Particle::Bank>& sites)
   hsize_t n_sites;
   H5Sget_simple_extent_dims(dspace, &n_sites, nullptr);
 
-  // Make sure source bank is big enough
-  sites.resize(n_sites);
+  // Make sure vector is big enough in case where we're reading entire source on
+  // each process
+  if (!distribute) sites.resize(n_sites);
 
-  // Determine number of source particles to read from each rank and offset
-  hsize_t min_sites = n_sites / mpi::n_procs;
-  hsize_t remainder = n_sites % mpi::n_procs;
-  hsize_t n_sites_local = (mpi::rank < remainder) ? min_sites + 1 : min_sites;
-  hsize_t offset = (mpi::rank <= remainder) ?
-    mpi::rank*(min_sites + 1) :
-    remainder*(min_sites + 1) + (mpi::rank - remainder)*min_sites;
+  hid_t memspace;
+  if (distribute) {
+    if (simulation::work_index[mpi::n_procs] > n_sites) {
+      fatal_error("Number of source sites in source file is less "
+                  "than number of source particles per generation.");
+    }
 
-  // Create another data space but for each proc individually
-  hid_t memspace = H5Screate_simple(1, &n_sites_local, nullptr);
+    // Create another data space but for each proc individually
+    hsize_t n_sites_local = simulation::work_per_rank;
+    memspace = H5Screate_simple(1, &n_sites_local, nullptr);
 
-  // Select hyperslab for each process
-  H5Sselect_hyperslab(dspace, H5S_SELECT_SET, &offset, nullptr, &n_sites_local, nullptr);
+    // Select hyperslab for each process
+    hsize_t offset = simulation::work_index[mpi::rank];
+    H5Sselect_hyperslab(dspace, H5S_SELECT_SET, &offset, nullptr, &n_sites_local, nullptr);
+  } else {
+    memspace = H5S_ALL;
+  }
 
 #ifdef PHDF5
     // Read data in parallel
@@ -692,7 +697,7 @@ void read_source_bank(hid_t group_id, std::vector<Particle::Bank>& sites)
 
   // Close all ids
   H5Sclose(dspace);
-  H5Sclose(memspace);
+  if (distribute) H5Sclose(memspace);
   H5Dclose(dset);
   H5Tclose(banktype);
 }