paulromano · Jul 21, 2023 · Jun 26, 2023 · Jun 27, 2023 · Jun 27, 2023 · Jul 5, 2023
diff --git a/include/openmc/mgxs.h b/include/openmc/mgxs.h
@@ -16,20 +16,6 @@
 
 namespace openmc {
 
-//==============================================================================
-// Cache contains the cached data for an MGXS object
-//==============================================================================
-
-struct CacheData {
-  double sqrtkT; // last temperature corresponding to t
-  int t;         // temperature index
-  int a;         // angle index
-  // last angle that corresponds to a
-  double u;
-  double v;
-  double w;
-};
-
 //==============================================================================
 // MGXS contains the mgxs data for a nuclide/material
 //==============================================================================
@@ -96,10 +82,9 @@ class Mgxs {
   bool equiv(const Mgxs& that);
 
 public:
-  std::string name;        // name of dataset, e.g., UO2
-  double awr;              // atomic weight ratio
-  bool fissionable;        // Is this fissionable
-  vector<CacheData> cache; // index and data cache
+  std::string name; // name of dataset, e.g., UO2
+  double awr;       // atomic weight ratio
+  bool fissionable; // Is this fissionable
 
   Mgxs() = default;
 
@@ -135,13 +120,15 @@ class Mgxs {
   //! @param mu Cosine of the change-in-angle, for scattering quantities;
   //!   use nullptr if irrelevant.
   //! @param dg delayed group index; use nullptr if irrelevant.
+  //! @param t Temperature index.
+  //! @param a Angle index.
   //! @return Requested cross section value.
-  double get_xs(
-    MgxsType xstype, int gin, const int* gout, const double* mu, const int* dg);
+  double get_xs(MgxsType xstype, int gin, const int* gout, const double* mu,
+    const int* dg, int t, int a);
 
-  inline double get_xs(MgxsType xstype, int gin)
+  inline double get_xs(MgxsType xstype, int gin, int t, int a)
   {
-    return get_xs(xstype, gin, nullptr, nullptr, nullptr);
+    return get_xs(xstype, gin, nullptr, nullptr, nullptr, t, a);
   }
 
   //! \brief Samples the fission neutron energy and if prompt or delayed.
@@ -150,7 +137,10 @@ class Mgxs {
   //! @param dg Sampled delayed group index.
   //! @param gout Sampled outgoing energy group.
   //! @param seed Pseudorandom seed pointer
-  void sample_fission_energy(int gin, int& dg, int& gout, uint64_t* seed);
+  //! @param t Temperature index.
+  //! @param a Angle index.
+  void sample_fission_energy(
+    int gin, int& dg, int& gout, uint64_t* seed, int t, int a);
 
   //! \brief Samples the outgoing energy and angle from a scatter event.
   //!
@@ -159,23 +149,37 @@ class Mgxs {
   //! @param mu Sampled cosine of the change-in-angle.
   //! @param wgt Weight of the particle to be adjusted.
   //! @param seed Pseudorandom seed pointer.
+  //! @param t Temperature index.
+  //! @param a Angle index.
   void sample_scatter(
-    int gin, int& gout, double& mu, double& wgt, uint64_t* seed);
+    int gin, int& gout, double& mu, double& wgt, uint64_t* seed, int t, int a);
 
   //! \brief Calculates cross section quantities needed for tracking.
   //!
   //! @param p The particle whose attributes set which MGXS to get.
   void calculate_xs(Particle& p);
 
-  //! \brief Sets the temperature index in cache given a temperature
+  //! \brief Sets the temperature index in particle cache.
+  //!
+  //! @param p Particle.
+  void set_temperature_index(Particle& p);
+
+  //! \brief Gets the temperature index given a temperature.
   //!
   //! @param sqrtkT Temperature of the material.
-  void set_temperature_index(double sqrtkT);
+  //! @return The temperature index corresponding to sqrtkT.
+  int get_temperature_index(double sqrtkT);
+
+  //! \brief Sets the angle index in particle cache.
+  //!
+  //! @param p Particle.
+  void set_angle_index(Particle& p);
 
-  //! \brief Sets the angle index in cache given a direction
+  //! \brief Gets the angle index given a direction.
   //!
   //! @param u Incoming particle direction.
-  void set_angle_index(Direction u);
+  //! @return The angle index corresponding to u.
+  int get_angle_index(Direction& u);
 
   //! \brief Provide const access to list of XsData held by this
   const vector<XsData>& get_xsdata() const { return xs; }

diff --git a/include/openmc/particle_data.h b/include/openmc/particle_data.h
@@ -170,6 +170,18 @@ struct MacroXS {
   double pair_production; //!< macroscopic pair production xs
 };
 
+//==============================================================================
+// Cache contains the cached data for an MGXS object
+//==============================================================================
+
+struct CacheDataMG {
+  int material {-1}; //!< material index
+  double sqrtkT;     //!< last temperature corresponding to t
+  int t {0};         //!< temperature index
+  int a {0};         //!< angle index
+  Direction u;       //!< angle that corresponds to a
+};
+
 //==============================================================================
 // Information about nearest boundary crossing
 //==============================================================================
@@ -231,6 +243,7 @@ class ParticleData {
   vector<NuclideMicroXS> neutron_xs_; //!< Microscopic neutron cross sections
   vector<ElementMicroXS> photon_xs_;  //!< Microscopic photon cross sections
   MacroXS macro_xs_;                  //!< Macroscopic cross sections
+  CacheDataMG mg_xs_cache_;           //!< Multigroup XS cache
 
   int64_t id_;                                //!< Unique ID
   ParticleType type_ {ParticleType::neutron}; //!< Particle type (n, p, e, etc.)
@@ -347,6 +360,8 @@ class ParticleData {
   ElementMicroXS& photon_xs(int i) { return photon_xs_[i]; }
   MacroXS& macro_xs() { return macro_xs_; }
   const MacroXS& macro_xs() const { return macro_xs_; }
+  CacheDataMG& mg_xs_cache() { return mg_xs_cache_; }
+  const CacheDataMG& mg_xs_cache() const { return mg_xs_cache_; }
 
   int64_t& id() { return id_; }
   const int64_t& id() const { return id_; }

diff --git a/src/mgxs.cpp b/src/mgxs.cpp
@@ -5,10 +5,6 @@
 #include <cstdlib>
 #include <sstream>
 
-#ifdef _OPENMP
-#include <omp.h>
-#endif
-
 #include "xtensor/xadapt.hpp"
 #include "xtensor/xmath.hpp"
 #include "xtensor/xsort.hpp"
@@ -46,15 +42,6 @@ void Mgxs::init(const std::string& in_name, double in_awr,
   n_azi = in_azimuthal.size();
   polar = in_polar;
   azimuthal = in_azimuthal;
-
-  // Set the cross section index cache
-#ifdef _OPENMP
-  int n_threads = omp_get_max_threads();
-#else
-  int n_threads = 1;
-#endif
-  cache.resize(n_threads);
-  // vector.resize() will value-initialize the members of cache[:]
 }
 
 //==============================================================================
@@ -425,18 +412,10 @@ void Mgxs::combine(const vector<Mgxs*>& micros, const vector<double>& scalars,
 
 //==============================================================================
 
-double Mgxs::get_xs(
-  MgxsType xstype, int gin, const int* gout, const double* mu, const int* dg)
+double Mgxs::get_xs(MgxsType xstype, int gin, const int* gout, const double* mu,
+  const int* dg, int t, int a)
 {
-  // This method assumes that the temperature and angle indices are set
-#ifdef _OPENMP
-  int tid = omp_get_thread_num();
-  XsData* xs_t = &xs[cache[tid].t];
-  int a = cache[tid].a;
-#else
-  XsData* xs_t = &xs[cache[0].t];
-  int a = cache[0].a;
-#endif
+  XsData* xs_t = &xs[t];
   double val;
   switch (xstype) {
   case MgxsType::TOTAL:
@@ -538,19 +517,14 @@ double Mgxs::get_xs(
 
 //==============================================================================
 
-void Mgxs::sample_fission_energy(int gin, int& dg, int& gout, uint64_t* seed)
+void Mgxs::sample_fission_energy(
+  int gin, int& dg, int& gout, uint64_t* seed, int t, int a)
 {
-  // This method assumes that the temperature and angle indices are set
-#ifdef _OPENMP
-  int tid = omp_get_thread_num();
-#else
-  int tid = 0;
-#endif
-  XsData* xs_t = &xs[cache[tid].t];
-  double nu_fission = xs_t->nu_fission(cache[tid].a, gin);
+  XsData* xs_t = &xs[t];
+  double nu_fission = xs_t->nu_fission(a, gin);
 
   // Find the probability of having a prompt neutron
-  double prob_prompt = xs_t->prompt_nu_fission(cache[tid].a, gin);
+  double prob_prompt = xs_t->prompt_nu_fission(a, gin);
 
   // sample random numbers
   double xi_pd = prn(seed) * nu_fission;
@@ -566,7 +540,7 @@ void Mgxs::sample_fission_energy(int gin, int& dg, int& gout, uint64_t* seed)
     // sample the outgoing energy group
     double prob_gout = 0.;
     for (gout = 0; gout < num_groups; ++gout) {
-      prob_gout += xs_t->chi_prompt(cache[tid].a, gin, gout);
+      prob_gout += xs_t->chi_prompt(a, gin, gout);
       if (xi_gout < prob_gout)
         break;
     }
@@ -576,7 +550,7 @@ void Mgxs::sample_fission_energy(int gin, int& dg, int& gout, uint64_t* seed)
 
     // get the delayed group
     for (dg = 0; dg < num_delayed_groups; ++dg) {
-      prob_prompt += xs_t->delayed_nu_fission(cache[tid].a, dg, gin);
+      prob_prompt += xs_t->delayed_nu_fission(a, dg, gin);
       if (xi_pd < prob_prompt)
         break;
     }
@@ -587,7 +561,7 @@ void Mgxs::sample_fission_energy(int gin, int& dg, int& gout, uint64_t* seed)
     // sample the outgoing energy group
     double prob_gout = 0.;
     for (gout = 0; gout < num_groups; ++gout) {
-      prob_gout += xs_t->chi_delayed(cache[tid].a, dg, gin, gout);
+      prob_gout += xs_t->chi_delayed(a, dg, gin, gout);
       if (xi_gout < prob_gout)
         break;
     }
@@ -597,35 +571,39 @@ void Mgxs::sample_fission_energy(int gin, int& dg, int& gout, uint64_t* seed)
 //==============================================================================
 
 void Mgxs::sample_scatter(
-  int gin, int& gout, double& mu, double& wgt, uint64_t* seed)
+  int gin, int& gout, double& mu, double& wgt, uint64_t* seed, int t, int a)
 {
-  // This method assumes that the temperature and angle indices are set
   // Sample the data
-#ifdef _OPENMP
-  int tid = omp_get_thread_num();
-#else
-  int tid = 0;
-#endif
-  xs[cache[tid].t].scatter[cache[tid].a]->sample(gin, gout, mu, wgt, seed);
+  xs[t].scatter[a]->sample(gin, gout, mu, wgt, seed);
 }
 
 //==============================================================================
 
 void Mgxs::calculate_xs(Particle& p)
 {
-  // Set our indices
-#ifdef _OPENMP
-  int tid = omp_get_thread_num();
-#else
-  int tid = 0;
-#endif
-  set_temperature_index(p.sqrtkT());
-  set_angle_index(p.u_local());
-  XsData* xs_t = &xs[cache[tid].t];
-  p.macro_xs().total = xs_t->total(cache[tid].a, p.g());
-  p.macro_xs().absorption = xs_t->absorption(cache[tid].a, p.g());
+  // If the material is different, then we need to do a full lookup
+  if (p.material() != p.mg_xs_cache().material) {
+    set_temperature_index(p);
+    set_angle_index(p);
+    p.mg_xs_cache().material = p.material();
+  } else {
+    // If material is the same, but temperature is different, need to
+    // find the new temperature index
+    if (p.sqrtkT() != p.mg_xs_cache().sqrtkT) {
+      set_temperature_index(p);
+    }
+    // If the material is the same, but angle is different, need to
+    // find the new angle index
+    if (p.u_local() != p.mg_xs_cache().u) {
+      set_angle_index(p);
+    }
+  }
+  int temperature = p.mg_xs_cache().t;
+  int angle = p.mg_xs_cache().a;
+  p.macro_xs().total = xs[temperature].total(angle, p.g());
+  p.macro_xs().absorption = xs[temperature].absorption(angle, p.g());
   p.macro_xs().nu_fission =
-    fissionable ? xs_t->nu_fission(cache[tid].a, p.g()) : 0.;
+    fissionable ? xs[temperature].nu_fission(angle, p.g()) : 0.;
 }
 
 //==============================================================================
@@ -643,32 +621,26 @@ bool Mgxs::equiv(const Mgxs& that)
 
 //==============================================================================
 
-void Mgxs::set_temperature_index(double sqrtkT)
+int Mgxs::get_temperature_index(double sqrtkT)
 {
-  // See if we need to find the new index
-#ifdef _OPENMP
-  int tid = omp_get_thread_num();
-#else
-  int tid = 0;
-#endif
-  if (sqrtkT != cache[tid].sqrtkT) {
-    cache[tid].t = xt::argmin(xt::abs(kTs - sqrtkT * sqrtkT))[0];
-    cache[tid].sqrtkT = sqrtkT;
-  }
+  return xt::argmin(xt::abs(kTs - sqrtkT * sqrtkT))[0];
 }
 
 //==============================================================================
 
-void Mgxs::set_angle_index(Direction u)
+void Mgxs::set_temperature_index(Particle& p)
 {
-  // See if we need to find the new index
-#ifdef _OPENMP
-  int tid = omp_get_thread_num();
-#else
-  int tid = 0;
-#endif
-  if (!is_isotropic && ((u.x != cache[tid].u) || (u.y != cache[tid].v) ||
-                         (u.z != cache[tid].w))) {
+  p.mg_xs_cache().t = get_temperature_index(p.sqrtkT());
+  p.mg_xs_cache().sqrtkT = p.sqrtkT();
+}
+
+//==============================================================================
+
+int Mgxs::get_angle_index(Direction& u)
+{
+  if (is_isotropic) {
+    return 0;
+  } else {
     // convert direction to polar and azimuthal angles
     double my_pol = std::acos(u.z);
     double my_azi = std::atan2(u.y, u.x);
@@ -679,12 +651,18 @@ void Mgxs::set_angle_index(Direction u)
     delta_angle = 2. * PI / n_azi;
     int a = std::floor((my_azi + PI) / delta_angle);
 
-    cache[tid].a = n_azi * p + a;
+    return n_azi * p + a;
+  }
+}
+
+//==============================================================================
 
-    // store this direction as the last one used
-    cache[tid].u = u.x;
-    cache[tid].v = u.y;
-    cache[tid].w = u.z;
+void Mgxs::set_angle_index(Particle& p)
+{
+  // See if we need to find the new index
+  if (!is_isotropic) {
+    p.mg_xs_cache().a = get_angle_index(p.u_local());
+    p.mg_xs_cache().u = p.u_local();
   }
 }
 

diff --git a/src/physics_mg.cpp b/src/physics_mg.cpp
@@ -72,8 +72,8 @@ void sample_reaction(Particle& p)
 
 void scatter(Particle& p)
 {
-  data::mg.macro_xs_[p.material()].sample_scatter(
-    p.g_last(), p.g(), p.mu(), p.wgt(), p.current_seed());
+  data::mg.macro_xs_[p.material()].sample_scatter(p.g_last(), p.g(), p.mu(),
+    p.wgt(), p.current_seed(), p.mg_xs_cache().t, p.mg_xs_cache().a);
 
   // Rotate the angle
   p.u() = rotate_angle(p.u(), p.mu(), nullptr, p.current_seed());
@@ -145,7 +145,7 @@ void create_fission_sites(Particle& p)
     int dg;
     int gout;
     data::mg.macro_xs_[p.material()].sample_fission_energy(
-      p.g(), dg, gout, p.current_seed());
+      p.g(), dg, gout, p.current_seed(), p.mg_xs_cache().t, p.mg_xs_cache().a);
 
     // Store the energy and delayed groups on the fission bank
     site.E = gout;