From 687d43b0ab6e3e3ddfe867767c210f0a1cfb38a7 Mon Sep 17 00:00:00 2001
From: Patrick Myers <pmyers@anl.gov>
Date: Wed, 25 May 2022 09:06:53 -0500
Subject: [PATCH] added external xtensor within PhotonInteraction for all shell
 photoelectric xs to utilize contiguous memory in calculate_xs

---
 include/openmc/photon.h |  2 +-
 src/photon.cpp          | 26 +++++++++++---------------
 src/physics.cpp         |  6 +++---
 3 files changed, 15 insertions(+), 19 deletions(-)

diff --git a/include/openmc/photon.h b/include/openmc/photon.h
index e8ec82d107..50ee228c44 100644
--- a/include/openmc/photon.h
+++ b/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;
 };
 
@@ -82,6 +81,7 @@ public:
 
   // Photoionization and atomic relaxation data
   vector<ElectronSubshell> shells_;
+  xt::xtensor<double, 2> cross_sections_;
 
   // Compton profile data
   xt::xtensor<double, 2> profile_pdf_;
diff --git a/src/photon.cpp b/src/photon.cpp
index b0cd8aa87d..3f9ff6dbf3 100644
--- a/src/photon.cpp
+++ b/src/photon.cpp
@@ -16,6 +16,9 @@
 #include "xtensor/xbuilder.hpp"
 #include "xtensor/xoperation.hpp"
 #include "xtensor/xview.hpp"
+#include "xtensor/xmath.hpp"
+#include "xtensor/xslice.hpp"
+#include "xtensor/xtensor_forward.hpp"
 
 #include <cmath>
 #include <fmt/core.h>
@@ -125,6 +128,7 @@ PhotonInteraction::PhotonInteraction(hid_t group)
   }
 
   shells_.resize(n_shell);
+  cross_sections_.resize({energy_.size(), n_shell});
 
   // Create mapping from designator to index
   std::unordered_map<int, int> shell_map;
@@ -155,13 +159,14 @@ 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, shell.threshold + xs.size()), i);
+    cross_section = xt::where(xs > 0.0, xt::log(xs), -500.0);
 
     if (object_exists(tgroup, "transitions")) {
       // Determine dimensions of transitions
@@ -565,19 +570,10 @@ void PhotonInteraction::calculate_xs(Particle& p) const
     incoherent_(i_grid) + f * (incoherent_(i_grid + 1) - incoherent_(i_grid)));
 
   // 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_upper = xt::row(cross_sections_, i_grid);
+  const auto& xs_lower = 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)));
-  }
+  xs.photoelectric = xt::sum(xt::exp(xs_upper + f * (xs_upper - xs_lower)))[0];
 
   // Calculate microscopic pair production cross section
   xs.pair_production = std::exp(
diff --git a/src/physics.cpp b/src/physics.cpp
index 13dd063de5..38d74dbede 100644
--- a/src/physics.cpp
+++ b/src/physics.cpp
@@ -358,9 +358,9 @@ void sample_photon_reaction(Particle& p)
         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)));
+      double xs = std::exp(element.cross_sections_(i_grid - i_start) +
+                           f * (element.cross_sections_(i_grid + 1 - i_start) -
+                                 element.cross_sections_(i_grid - i_start)));
 
       prob += xs;
       if (prob > cutoff) {