From 07a0ff5072863c3129fbe6e6a17f4f1acf2a3fa0 Mon Sep 17 00:00:00 2001
From: amandalund <alund@anl.gov>
Date: Sun, 3 Oct 2021 19:25:56 -0500
Subject: [PATCH] Add a helper function to calculate the maximum size of the
 stack used in atomic relaxation

---
 include/openmc/photon.h |  3 +++
 src/photon.cpp          | 57 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+)

diff --git a/include/openmc/photon.h b/include/openmc/photon.h
index f5a3d8e81b..989d6dfa9c 100644
--- a/include/openmc/photon.h
+++ b/include/openmc/photon.h
@@ -104,6 +104,9 @@ public:
 private:
   void compton_doppler(
     double alpha, double mu, double* E_out, int* i_shell, uint64_t* seed) const;
+
+  int calc_max_stack_size() const;
+  int calc_helper(std::unordered_map<int, int>& visited, int i_shell) const;
 };
 
 //==============================================================================
diff --git a/src/photon.cpp b/src/photon.cpp
index b1720019f3..385583ee33 100644
--- a/src/photon.cpp
+++ b/src/photon.cpp
@@ -5,6 +5,7 @@
 #include "openmc/constants.h"
 #include "openmc/distribution_multi.h"
 #include "openmc/hdf5_interface.h"
+#include "openmc/message_passing.h"
 #include "openmc/nuclide.h"
 #include "openmc/particle.h"
 #include "openmc/random_dist.h"
@@ -187,6 +188,14 @@ PhotonInteraction::PhotonInteraction(hid_t group)
   }
   close_group(rgroup);
 
+  // Check the maximum size of the atomic relaxation stack
+  auto max_size = this->calc_max_stack_size();
+  if (max_size > MAX_STACK_SIZE && mpi::master) {
+    warning("The subshell vacancy stack in atomic relaxation can grow up to " +
+            std::to_string(max_size) + ", but the stack size limit is set to " +
+            std::to_string(MAX_STACK_SIZE) + ".");
+  }
+
   // Determine number of electron shells
   rgroup = open_group(group, "compton_profiles");
 
@@ -314,6 +323,54 @@ PhotonInteraction::~PhotonInteraction()
   data::element_map.erase(name_);
 }
 
+int PhotonInteraction::calc_max_stack_size() const
+{
+  // Table to store solutions to sub-problems
+  std::unordered_map<int, int> visited;
+
+  // Find the maximum possible size of the stack used to store holes created
+  // during atomic relaxation, checking over every subshell the initial hole
+  // could be in
+  int max_size = 0;
+  for (int i_shell = 0; i_shell < shells_.size(); ++i_shell) {
+    max_size = std::max(max_size, this->calc_helper(visited, i_shell));
+  }
+  return max_size;
+}
+
+int PhotonInteraction::calc_helper(
+  std::unordered_map<int, int>& visited, int i_shell) const
+{
+  // No transitions for this subshell, so this is the only shell in the stack
+  const auto& shell {shells_[i_shell]};
+  if (shell.transitions.empty()) {
+    return 1;
+  }
+
+  // Check the table to see if the maximum stack size has already been
+  // calculated for this shell
+  auto it = visited.find(i_shell);
+  if (it != visited.end()) {
+    return it->second;
+  }
+
+  int max_size = 0;
+  for (const auto& transition : shell.transitions) {
+    // If this is a non-radiative transition two vacancies are created and
+    // the stack grows by one; if this is a radiative transition only one
+    // vacancy is created and the stack size stays the same
+    int size = 0;
+    if (transition.secondary_subshell != -1) {
+      size = this->calc_helper(visited, transition.secondary_subshell) + 1;
+    }
+    size =
+      std::max(size, this->calc_helper(visited, transition.primary_subshell));
+    max_size = std::max(max_size, size);
+  }
+  visited[i_shell] = max_size;
+  return max_size;
+}
+
 void PhotonInteraction::compton_scatter(double alpha, bool doppler,
   double* alpha_out, double* mu, int* i_shell, uint64_t* seed) const
 {