Implement jacobian calculation and energy and angle transformation

include/openmc/angle_energy.h

@@ -1,8 +1,11 @@
 #ifndef OPENMC_ANGLE_ENERGY_H
 #define OPENMC_ANGLE_ENERGY_H
 
+#include <cmath> // for sqrt
 #include <cstdint>
 
+#include "openmc/random_lcg.h"
+
 namespace openmc {
 
 //==============================================================================
@@ -27,12 +30,51 @@ class AngleEnergy {
   //! \param[in] mu Scattering cosine with respect to current direction
   //! \param[out] E_out Outgoing energy in [eV]
   //! \param[inout] seed Pseudorandom seed pointer
-  //! \return Probability density for the scattering cosine
-  virtual double sample_energy_and_pdf(
-    double E_in, double mu, double& E_out, uint64_t* seed) const = 0;
+  //! \param[in] is_com Is scattering cosine is given in center of mass
+  //! coordinates \param[in] awr Weight of nucleus in neutron masses \return
+  //! Probability density for the scattering cosine
+  virtual double sample_energy_and_pdf(double E_in, double mu, double& E_out,
+    uint64_t* seed, bool is_com, double awr) const = 0;
   virtual ~AngleEnergy() = default;
 };
 
+inline double get_jac_and_transform(
+  double E_in, double& mu, double& E_out, uint64_t* seed, double awr)
+{
+  double E_cm = E_out;
+  double mu_lab = mu;
+  double D = mu_lab * mu_lab - 1.0 + (awr + 1.0) * (awr + 1.0) * E_cm / E_in;
+  if (D < 0.0)
+    return 0.0;
+  D = std::sqrt(D);
+  double E_out1 =
+    E_in * ((mu_lab + D) / (awr + 1.0)) * ((mu_lab + D) / (awr + 1.0));
+  double E_out2 =
+    E_in * ((mu_lab - D) / (awr + 1.0)) * ((mu_lab - D) / (awr + 1.0));
+  double mu_cm1 =
+    mu_lab * std::sqrt(E_out1 / E_cm) - std::sqrt(E_in / E_cm) / (awr + 1.0);
+  double mu_cm2 =
+    mu_lab * std::sqrt(E_out1 / E_cm) - std::sqrt(E_in / E_cm) / (awr + 1.0);
+  double mult = 1.0;
+  if (std::abs(mu_cm1) > 1.0) {
+    mu = mu_cm2;
+    E_out = E_out2;
+  } else if (std::abs(mu_cm2) > 1.0) {
+    mu = mu_cm1;
+    E_out = E_out1;
+  } else {
+    mult = 2.0;
+    if (prn(seed) < 0.5) {
+      mu = mu_cm2;
+      E_out = E_out2;
+    } else {
+      mu = mu_cm2;
+      E_out = E_out2;
+    }
+  }
+  return mult * E_out * (awr + 1.0) / (D * std::sqrt(E_cm * E_in));
+}
+
 } // namespace openmc
 
 #endif // OPENMC_ANGLE_ENERGY_H

include/openmc/chain.h

@@ -76,9 +76,11 @@ class DecayPhotonAngleEnergy : public AngleEnergy {
   //! \param[in] mu Scattering cosine with respect to current direction
   //! \param[out] E_out Outgoing energy in [eV]
   //! \param[inout] seed Pseudorandom seed pointer
-  //! \return Probability density for the scattering cosine
-  double sample_energy_and_pdf(
-    double E_in, double mu, double& E_out, uint64_t* seed) const override;
+  //! \param[in] is_com Is scattering cosine is given in center of mass
+  //! coordinates \param[in] awr Weight of nucleus in neutron masses \return
+  //! Probability density for the scattering cosine
+  double sample_energy_and_pdf(double E_in, double mu, double& E_out,
+    uint64_t* seed, bool is_com, double awr) const override;
 
 private:
   const Distribution* photon_energy_;

include/openmc/reaction_product.h

@@ -60,9 +60,11 @@ class ReactionProduct {
   //! \param[in] mu Scattering cosine with respect to current direction
   //! \param[out] E_out Outgoing energy in [eV]
   //! \param[inout] seed Pseudorandom seed pointer
-  //! \return Probability density for the scattering cosine
-  double sample_energy_and_pdf(
-    double E_in, double mu, double& E_out, uint64_t* seed) const;
+  //! \param[in] is_com Is scattering cosine is given in center of mass
+  //! coordinates \param[in] awr Weight of nucleus in neutron masses \return
+  //! Probability density for the scattering cosine
+  double sample_energy_and_pdf(double E_in, double mu, double& E_out,
+    uint64_t* seed, bool is_com, double awr) const;
 
   ParticleType particle_;      //!< Particle type
   EmissionMode emission_mode_; //!< Emission mode

include/openmc/secondary_correlated.h

@@ -53,9 +53,11 @@ class CorrelatedAngleEnergy : public AngleEnergy {
   //! \param[in] mu Scattering cosine with respect to current direction
   //! \param[out] E_out Outgoing energy in [eV]
   //! \param[inout] seed Pseudorandom seed pointer
-  //! \return Probability density for the scattering cosine
-  double sample_energy_and_pdf(
-    double E_in, double mu, double& E_out, uint64_t* seed) const override;
+  //! \param[in] is_com Is scattering cosine is given in center of mass
+  //! coordinates \param[in] awr Weight of nucleus in neutron masses \return
+  //! Probability density for the scattering cosine
+  double sample_energy_and_pdf(double E_in, double mu, double& E_out,
+    uint64_t* seed, bool is_com, double awr) const override;
 
   // energy property
   vector<double>& energy() { return energy_; }

include/openmc/secondary_kalbach.h

@@ -46,9 +46,11 @@ class KalbachMann : public AngleEnergy {
   //! \param[in] mu Scattering cosine with respect to current direction
   //! \param[out] E_out Outgoing energy in [eV]
   //! \param[inout] seed Pseudorandom seed pointer
-  //! \return Probability density for the scattering cosine
-  double sample_energy_and_pdf(
-    double E_in, double mu, double& E_out, uint64_t* seed) const override;
+  //! \param[in] is_com Is scattering cosine is given in center of mass
+  //! coordinates \param[in] awr Weight of nucleus in neutron masses \return
+  //! Probability density for the scattering cosine
+  double sample_energy_and_pdf(double E_in, double mu, double& E_out,
+    uint64_t* seed, bool is_com, double awr) const override;
 
 private:
   //! Outgoing energy/angle at a single incoming energy

include/openmc/secondary_nbody.h

@@ -39,9 +39,11 @@ class NBodyPhaseSpace : public AngleEnergy {
   //! \param[in] mu Scattering cosine with respect to current direction
   //! \param[out] E_out Outgoing energy in [eV]
   //! \param[inout] seed Pseudorandom seed pointer
-  //! \return Probability density for the scattering cosine
-  double sample_energy_and_pdf(
-    double E_in, double mu, double& E_out, uint64_t* seed) const override;
+  //! \param[in] is_com Is scattering cosine is given in center of mass
+  //! coordinates \param[in] awr Weight of nucleus in neutron masses \return
+  //! Probability density for the scattering cosine
+  double sample_energy_and_pdf(double E_in, double mu, double& E_out,
+    uint64_t* seed, bool is_com, double awr) const override;
 
 private:
   int n_bodies_;      //!< Number of particles distributed

include/openmc/secondary_thermal.h

@@ -39,9 +39,11 @@ class CoherentElasticAE : public AngleEnergy {
   //! \param[in] mu Scattering cosine with respect to current direction
   //! \param[out] E_out Outgoing energy in [eV]
   //! \param[inout] seed Pseudorandom seed pointer
-  //! \return Probability density for the scattering cosine
-  double sample_energy_and_pdf(
-    double E_in, double mu, double& E_out, uint64_t* seed) const override;
+  //! \param[in] is_com Is scattering cosine is given in center of mass
+  //! coordinates \param[in] awr Weight of nucleus in neutron masses \return
+  //! Probability density for the scattering cosine
+  double sample_energy_and_pdf(double E_in, double mu, double& E_out,
+    uint64_t* seed, bool is_com, double awr) const override;
 
 private:
   const CoherentElasticXS& xs_; //!< Coherent elastic scattering cross section
@@ -74,9 +76,11 @@ class IncoherentElasticAE : public AngleEnergy {
   //! \param[in] mu Scattering cosine with respect to current direction
   //! \param[out] E_out Outgoing energy in [eV]
   //! \param[inout] seed Pseudorandom seed pointer
-  //! \return Probability density for the scattering cosine
-  double sample_energy_and_pdf(
-    double E_in, double mu, double& E_out, uint64_t* seed) const override;
+  //! \param[in] is_com Is scattering cosine is given in center of mass
+  //! coordinates \param[in] awr Weight of nucleus in neutron masses \return
+  //! Probability density for the scattering cosine
+  double sample_energy_and_pdf(double E_in, double mu, double& E_out,
+    uint64_t* seed, bool is_com, double awr) const override;
 
 private:
   double debye_waller_;
@@ -108,9 +112,11 @@ class IncoherentElasticAEDiscrete : public AngleEnergy {
   //! \param[in] mu Scattering cosine with respect to current direction
   //! \param[out] E_out Outgoing energy in [eV]
   //! \param[inout] seed Pseudorandom seed pointer
-  //! \return Probability density for the scattering cosine
-  double sample_energy_and_pdf(
-    double E_in, double mu, double& E_out, uint64_t* seed) const override;
+  //! \param[in] is_com Is scattering cosine is given in center of mass
+  //! coordinates \param[in] awr Weight of nucleus in neutron masses \return
+  //! Probability density for the scattering cosine
+  double sample_energy_and_pdf(double E_in, double mu, double& E_out,
+    uint64_t* seed, bool is_com, double awr) const override;
 
 private:
   const vector<double>& energy_;  //!< Energies at which cosines are tabulated
@@ -149,9 +155,11 @@ class IncoherentInelasticAEDiscrete : public AngleEnergy {
   //! \param[in] mu Scattering cosine with respect to current direction
   //! \param[out] E_out Outgoing energy in [eV]
   //! \param[inout] seed Pseudorandom seed pointer
-  //! \return Probability density for the scattering cosine
-  double sample_energy_and_pdf(
-    double E_in, double mu, double& E_out, uint64_t* seed) const override;
+  //! \param[in] is_com Is scattering cosine is given in center of mass
+  //! coordinates \param[in] awr Weight of nucleus in neutron masses \return
+  //! Probability density for the scattering cosine
+  double sample_energy_and_pdf(double E_in, double mu, double& E_out,
+    uint64_t* seed, bool is_com, double awr) const override;
 
 private:
   const vector<double>& energy_; //!< Incident energies
@@ -196,9 +204,11 @@ class IncoherentInelasticAE : public AngleEnergy {
   //! \param[in] mu Scattering cosine with respect to current direction
   //! \param[out] E_out Outgoing energy in [eV]
   //! \param[inout] seed Pseudorandom seed pointer
-  //! \return Probability density for the scattering cosine
-  double sample_energy_and_pdf(
-    double E_in, double mu, double& E_out, uint64_t* seed) const override;
+  //! \param[in] is_com Is scattering cosine is given in center of mass
+  //! coordinates \param[in] awr Weight of nucleus in neutron masses \return
+  //! Probability density for the scattering cosine
+  double sample_energy_and_pdf(double E_in, double mu, double& E_out,
+    uint64_t* seed, bool is_com, double awr) const override;
 
 private:
   //! Secondary energy/angle distribution
@@ -246,9 +256,11 @@ class MixedElasticAE : public AngleEnergy {
   //! \param[in] mu Scattering cosine with respect to current direction
   //! \param[out] E_out Outgoing energy in [eV]
   //! \param[inout] seed Pseudorandom seed pointer
-  //! \return Probability density for the scattering cosine
-  double sample_energy_and_pdf(
-    double E_in, double mu, double& E_out, uint64_t* seed) const override;
+  //! \param[in] is_com Is scattering cosine is given in center of mass
+  //! coordinates \param[in] awr Weight of nucleus in neutron masses \return
+  //! Probability density for the scattering cosine
+  double sample_energy_and_pdf(double E_in, double mu, double& E_out,
+    uint64_t* seed, bool is_com, double awr) const override;
 
 private:
   CoherentElasticAE coherent_dist_;         //!< Coherent distribution

include/openmc/secondary_uncorrelated.h

@@ -37,9 +37,11 @@ class UncorrelatedAngleEnergy : public AngleEnergy {
   //! \param[in] mu Scattering cosine with respect to current direction
   //! \param[out] E_out Outgoing energy in [eV]
   //! \param[inout] seed Pseudorandom seed pointer
-  //! \return Probability density for the scattering cosine
-  double sample_energy_and_pdf(
-    double E_in, double mu, double& E_out, uint64_t* seed) const override;
+  //! \param[in] is_com Is scattering cosine is given in center of mass
+  //! coordinates \param[in] awr Weight of nucleus in neutron masses \return
+  //! Probability density for the scattering cosine
+  double sample_energy_and_pdf(double E_in, double mu, double& E_out,
+    uint64_t* seed, bool is_com, double awr) const override;
 
   // Accessors
   AngleDistribution& angle() { return angle_; }

include/openmc/thermal.h

@@ -67,9 +67,11 @@ class ThermalData {
   //! \param[in] mu Scattering cosine with respect to current direction
   //! \param[out] E_out Outgoing energy in [eV]
   //! \param[inout] seed Pseudorandom seed pointer
-  //! \return Probability density for the scattering cosine
+  //! \param[in] is_com Is scattering cosine is given in center of mass
+  //! coordinates \param[in] awr Weight of nucleus in neutron masses \return
+  //! Probability density for the scattering cosine
   double sample_energy_and_pdf(const NuclideMicroXS& micro_xs, double E_in,
-    double mu, double& E_out, uint64_t* seed) const;
+    double mu, double& E_out, uint64_t* seed, bool is_com, double awr) const;
 
 private:
   struct Reaction {

src/chain.cpp

@@ -74,9 +74,13 @@ void DecayPhotonAngleEnergy::sample(
   mu = Uniform(-1., 1.).sample(seed).first;
 }
 
-double DecayPhotonAngleEnergy::sample_energy_and_pdf(
-  double E_in, double mu, double& E_out, uint64_t* seed) const
+double DecayPhotonAngleEnergy::sample_energy_and_pdf(double E_in, double mu,
+  double& E_out, uint64_t* seed, bool is_com, double awr) const
 {
+  if (is_com)
+    fatal_error(
+      "DecayPhotonAngleEnergy distribution must be in lab coordinates");
+
   E_out = photon_energy_->sample(seed).first;
   return 0.5;
 }

src/reaction_product.cpp

@@ -134,10 +134,11 @@ void ReactionProduct::sample(
   sample_dist(E_in, seed).sample(E_in, E_out, mu, seed);
 }
 
-double ReactionProduct::sample_energy_and_pdf(
-  double E_in, double mu, double& E_out, uint64_t* seed) const
+double ReactionProduct::sample_energy_and_pdf(double E_in, double mu,
+  double& E_out, uint64_t* seed, bool is_com, double awr) const
 {
-  return sample_dist(E_in, seed).sample_energy_and_pdf(E_in, mu, E_out, seed);
+  return sample_dist(E_in, seed)
+    .sample_energy_and_pdf(E_in, mu, E_out, seed, is_com, awr);
 }
 
 } // namespace openmc

src/secondary_correlated.cpp

@@ -260,10 +260,15 @@ void CorrelatedAngleEnergy::sample(
   mu = sample_dist(E_in, E_out, seed).sample(seed).first;
 }
 
-double CorrelatedAngleEnergy::sample_energy_and_pdf(
-  double E_in, double mu, double& E_out, uint64_t* seed) const
+double CorrelatedAngleEnergy::sample_energy_and_pdf(double E_in, double mu,
+  double& E_out, uint64_t* seed, bool is_com, double awr) const
 {
-  return sample_dist(E_in, E_out, seed).evaluate(mu);
+  auto& dist = sample_dist(E_in, E_out, seed);
+  double jac = 1.0;
+  if (is_com)
+    jac = get_jac_and_transform(E_in, mu, E_out, seed, awr);
+
+  return jac * dist.evaluate(mu);
 }
 
 } // namespace openmc

src/secondary_kalbach.cpp

@@ -232,14 +232,18 @@ void KalbachMann::sample(
     mu = std::log(r1 * std::exp(km_a) + (1.0 - r1) * std::exp(-km_a)) / km_a;
   }
 }
-double KalbachMann::sample_energy_and_pdf(
-  double E_in, double mu, double& E_out, uint64_t* seed) const
+double KalbachMann::sample_energy_and_pdf(double E_in, double mu, double& E_out,
+  uint64_t* seed, bool is_com, double awr) const
 {
   double km_r, km_a;
   sample_params(E_in, E_out, km_a, km_r, seed);
 
+  double jac = 1.0;
+  if (is_com)
+    jac = get_jac_and_transform(E_in, mu, E_out, seed, awr);
+
   // https://docs.openmc.org/en/latest/methods/neutron_physics.html#equation-KM-pdf-angle
-  return km_a / (2 * std::sinh(km_a)) *
+  return jac * km_a / (2 * std::sinh(km_a)) *
          (std::cosh(km_a * mu) + km_r * std::sinh(km_a * mu));
 }
 

src/secondary_nbody.cpp

@@ -72,11 +72,14 @@ void NBodyPhaseSpace::sample(
   E_out = sample_energy(E_in, seed);
 }
 
-double NBodyPhaseSpace::sample_energy_and_pdf(
-  double E_in, double mu, double& E_out, uint64_t* seed) const
+double NBodyPhaseSpace::sample_energy_and_pdf(double E_in, double mu,
+  double& E_out, uint64_t* seed, bool is_com, double awr) const
 {
   E_out = sample_energy(E_in, seed);
-  return 0.5;
+  double jac = 1.0;
+  if (is_com)
+    jac = get_jac_and_transform(E_in, mu, E_out, seed, awr);
+  return jac * 0.5;
 }
 
 } // namespace openmc