Merge pull request #2098 from myerspat/otimize-photon-calcxs

Optimize microscopic photoelectric cross section calculation

include/openmc/photon.h

@@ -36,7 +36,6 @@ public:
   int threshold;
   double n_electrons;
   double binding_energy;
-  xt::xtensor<double, 1> cross_section;
   vector<Transition> transitions;
 };
 
@@ -80,8 +79,11 @@ public:
   Tabulated1D coherent_anomalous_real_;
   Tabulated1D coherent_anomalous_imag_;
 
-  // Photoionization and atomic relaxation data
+  // Photoionization and atomic relaxation data. Subshell cross sections are
+  // stored separately to improve memory access pattern when calculating the
+  // total cross section
   vector<ElectronSubshell> shells_;
+  xt::xtensor<double, 2> cross_sections_;
 
   // Compton profile data
   xt::xtensor<double, 2> profile_pdf_;

src/photon.cpp

@@ -14,7 +14,9 @@
 #include "openmc/settings.h"
 
 #include "xtensor/xbuilder.hpp"
+#include "xtensor/xmath.hpp"
 #include "xtensor/xoperation.hpp"
+#include "xtensor/xslice.hpp"
 #include "xtensor/xview.hpp"
 
 #include <cmath>
@@ -44,6 +46,8 @@ vector<unique_ptr<PhotonInteraction>> elements;
 
 PhotonInteraction::PhotonInteraction(hid_t group)
 {
+  using namespace xt::placeholders;
+
   // Set index of element in global vector
   index_ = data::elements.size();
 
@@ -125,6 +129,7 @@ PhotonInteraction::PhotonInteraction(hid_t group)
   }
 
   shells_.resize(n_shell);
+  cross_sections_ = xt::zeros<double>({energy_.size(), n_shell});
 
   // Create mapping from designator to index
   std::unordered_map<int, int> shell_map;
@@ -155,13 +160,15 @@ PhotonInteraction::PhotonInteraction(hid_t group)
     read_attribute(tgroup, "num_electrons", shell.n_electrons);
 
     // Read subshell cross section
+    xt::xtensor<double, 1> xs;
     dset = open_dataset(tgroup, "xs");
     read_attribute(dset, "threshold_idx", shell.threshold);
     close_dataset(dset);
-    read_dataset(tgroup, "xs", shell.cross_section);
+    read_dataset(tgroup, "xs", xs);
 
-    auto& xs = shell.cross_section;
-    xs = xt::where(xs > 0.0, xt::log(xs), -500.0);
+    auto cross_section =
+      xt::view(cross_sections_, xt::range(shell.threshold, _), i);
+    cross_section = xt::where(xs > 0, xt::log(xs), 0);
 
     if (object_exists(tgroup, "transitions")) {
       // Determine dimensions of transitions
@@ -566,18 +573,13 @@ void PhotonInteraction::calculate_xs(Particle& p) const
 
   // Calculate microscopic photoelectric cross section
   xs.photoelectric = 0.0;
-  for (const auto& shell : shells_) {
-    // Check threshold of reaction
-    int i_start = shell.threshold;
-    if (i_grid < i_start)
-      continue;
+  const auto& xs_lower = xt::row(cross_sections_, i_grid);
+  const auto& xs_upper = xt::row(cross_sections_, i_grid + 1);
 
-    // Evaluation subshell photoionization cross section
-    xs.photoelectric +=
-      std::exp(shell.cross_section(i_grid - i_start) +
-               f * (shell.cross_section(i_grid + 1 - i_start) -
-                     shell.cross_section(i_grid - i_start)));
-  }
+  for (int i = 0; i < xs_upper.size(); ++i)
+    if (xs_lower(i) != 0)
+      xs.photoelectric +=
+        std::exp(xs_lower(i) + f * (xs_upper(i) - xs_lower(i)));
 
   // Calculate microscopic pair production cross section
   xs.pair_production = std::exp(

src/physics.cpp

@@ -29,6 +29,7 @@
 
 #include <algorithm> // for max, min, max_element
 #include <cmath>     // for sqrt, exp, log, abs, copysign
+#include <xtensor/xview.hpp>
 
 namespace openmc {
 
@@ -344,25 +345,26 @@ void sample_photon_reaction(Particle& p)
 
   // Photoelectric effect
   double prob_after = prob + micro.photoelectric;
+
   if (prob_after > cutoff) {
+    // Get grid index, interpolation factor, and bounding subshell
+    // cross sections
+    int i_grid = micro.index_grid;
+    double f = micro.interp_factor;
+    const auto& xs_lower = xt::row(element.cross_sections_, i_grid);
+    const auto& xs_upper = xt::row(element.cross_sections_, i_grid + 1);
+
     for (int i_shell = 0; i_shell < element.shells_.size(); ++i_shell) {
       const auto& shell {element.shells_[i_shell]};
 
-      // Get grid index and interpolation factor
-      int i_grid = micro.index_grid;
-      double f = micro.interp_factor;
-
       // Check threshold of reaction
-      int i_start = shell.threshold;
-      if (i_grid < i_start)
+      if (xs_lower(i_shell) == 0)
         continue;
 
-      // Evaluation subshell photoionization cross section
-      double xs = std::exp(shell.cross_section(i_grid - i_start) +
-                           f * (shell.cross_section(i_grid + 1 - i_start) -
-                                 shell.cross_section(i_grid - i_start)));
+      //  Evaluation subshell photoionization cross section
+      prob += std::exp(
+        xs_lower(i_shell) + f * (xs_upper(i_shell) - xs_lower(i_shell)));
 
-      prob += xs;
       if (prob > cutoff) {
         double E_electron = p.E() - shell.binding_energy;