Document and slightly refactor Bethe-Heitler interactor (#1553)

* Precalculate some values
* Rewrite f in terms of phi
* Combine the branches
* Revert "Combine the branches"
* Additional documentation
* Make static functions static
* Rename screening threshold and fix no/full screening
* Update nomenclature

src/celeritas/em/interactor/BetheHeitlerInteractor.hh

@@ -11,12 +11,14 @@
 #include "corecel/Types.hh"
 #include "corecel/data/StackAllocator.hh"
 #include "corecel/math/Algorithms.hh"
+#include "corecel/math/ArrayOperators.hh"
 #include "corecel/math/ArrayUtils.hh"
 #include "celeritas/Constants.hh"
 #include "celeritas/Quantities.hh"
 #include "celeritas/em/data/BetheHeitlerData.hh"
 #include "celeritas/em/distribution/TsaiUrbanDistribution.hh"
 #include "celeritas/em/xs/LPMCalculator.hh"
+#include "celeritas/grid/PolyEvaluator.hh"
 #include "celeritas/mat/ElementView.hh"
 #include "celeritas/phys/Interaction.hh"
 #include "celeritas/phys/ParticleTrackView.hh"
@@ -41,6 +43,11 @@ namespace celeritas
  * G4PairProductionRelModel, as documented in sections 6.5 (gamma conversion)
  * and 10.2.2 (LPM effect) of the Geant4 Physics Reference Manual (release
  * 10.7)
+ *
+ * For additional context on the derivation see:
+ *     Butcher, J.C., and H. Messel. “Electron Number Distribution in
+ *     Electron-Photon Showers in Air and Aluminium Absorbers.” Nuclear Physics
+ *     20 (October 1960): 15–128. https://doi.org/10.1016/0029-5582(60)90162-0.
  */
 class BetheHeitlerInteractor
 {
@@ -68,34 +75,36 @@ class BetheHeitlerInteractor
   private:
     //// TYPES ////
 
-    //! Screening functions \f$ \Phi_1 \f$ and \f$ \Phi_2 \f$
-    struct ScreeningFunctions
-    {
-        real_type phi1;
-        real_type phi2;
-    };
+    using Real2 = Array<real_type, 2>;
 
     //// DATA ////
 
     // Shared model data
     BetheHeitlerData const& shared_;
     // Incident gamma energy
-    real_type const inc_energy_;
+    Energy const inc_energy_;
     // Incident direction
     Real3 const& inc_direction_;
     // Allocate space for a secondary particle
     StackAllocator<Secondary>& allocate_;
-    // Element properties for calculating screening functions and variables
-    ElementView const& element_;
     // Whether LPM supression is applied
     bool const enable_lpm_;
     // Used to calculate the LPM suppression functions
     LPMCalculator calc_lpm_functions_;
-    // Cached minimum epsilon, m_e*c^2/E_gamma; kinematical limit for Y -> e+e-
+    // Minimum epsilon, m_e*c^2/E_gamma; kinematical limit for Y -> e+e-
     real_type epsilon0_;
+    // 136/Z^1/3 factor on the screening variable, or zero for low energy
+    real_type screen_delta_;
+    real_type f_z_;
 
     //// CONSTANTS ////
 
+    //! Energy below which screening can be neglected
+    static CELER_CONSTEXPR_FUNCTION Energy no_screening_threshold()
+    {
+        return units::MevEnergy{2};
+    }
+
     //! Energy above which the Coulomb correction is applied [MeV]
     static CELER_CONSTEXPR_FUNCTION Energy coulomb_corr_threshold()
     {
@@ -114,14 +123,10 @@ class BetheHeitlerInteractor
     inline CELER_FUNCTION real_type impact_parameter(real_type eps) const;
 
     // Calculate the screening functions \f$ \Phi_1 \f$ and \f$ \Phi_2 \f$
-    inline CELER_FUNCTION ScreeningFunctions
-    screening_phi1_phi2(real_type delta) const;
-
-    // Calculate the auxiliary screening function \f$ F_1 \f$
-    inline CELER_FUNCTION real_type screening_f1(real_type delta) const;
+    static inline CELER_FUNCTION Real2 screening_phi(real_type delta);
 
-    // Calculate the auxiliary screening function \f$ F_2 \f$
-    inline CELER_FUNCTION real_type screening_f2(real_type delta) const;
+    // Calculate the auxiliary screening functions \f$ F_1 \f$ and \f$ F_2 \f$
+    static inline CELER_FUNCTION Real2 screening_f(real_type delta);
 };
 
 //---------------------------------------------------------------------------//
@@ -140,21 +145,36 @@ CELER_FUNCTION BetheHeitlerInteractor::BetheHeitlerInteractor(
     MaterialView const& material,
     ElementView const& element)
     : shared_(shared)
-    , inc_energy_(value_as<Energy>(particle.energy()))
+    , inc_energy_(particle.energy())
     , inc_direction_(inc_direction)
     , allocate_(allocate)
-    , element_(element)
-    , enable_lpm_(shared.enable_lpm
-                  && inc_energy_ > value_as<Energy>(lpm_threshold()))
-    , calc_lpm_functions_(material,
-                          element_,
-                          shared_.dielectric_suppression(),
-                          Energy{inc_energy_})
+    , enable_lpm_(shared.enable_lpm && inc_energy_ > lpm_threshold())
+    , calc_lpm_functions_(
+          material, element, shared_.dielectric_suppression(), inc_energy_)
 {
     CELER_EXPECT(particle.particle_id() == shared_.ids.gamma);
-    CELER_EXPECT(inc_energy_ >= 2 * value_as<Mass>(shared_.electron_mass));
+    CELER_EXPECT(value_as<Energy>(inc_energy_)
+                 >= 2 * value_as<Mass>(shared_.electron_mass));
 
-    epsilon0_ = value_as<Mass>(shared_.electron_mass) / inc_energy_;
+    epsilon0_ = value_as<Mass>(shared_.electron_mass)
+                / value_as<Energy>(inc_energy_);
+
+    if (inc_energy_ < no_screening_threshold())
+    {
+        // Don't reject: just sample uniformly
+        screen_delta_ = 0;
+        f_z_ = 0;
+    }
+    else
+    {
+        screen_delta_ = epsilon0_ * 136 / element.cbrt_z();
+        f_z_ = real_type(8) / real_type(3) * element.log_z();
+        if (inc_energy_ > coulomb_corr_threshold())
+        {
+            // Apply Coulomb correction function
+            f_z_ += 8 * element.coulomb_correction();
+        }
+    }
 }
 
 //---------------------------------------------------------------------------//
@@ -174,10 +194,11 @@ CELER_FUNCTION Interaction BetheHeitlerInteractor::operator()(Engine& rng)
 
     constexpr real_type half = 0.5;
 
-    // If E_gamma < 2 MeV, rejection not needed -- just sample uniformly
+    // Sample fraction of energy given to electron
     real_type epsilon;
-    if (inc_energy_ < value_as<Energy>(units::MevEnergy{2}))
+    if (screen_delta_ == 0)
     {
+        // If E_gamma < 2 MeV, rejection not needed -- just sample uniformly
         UniformRealDistribution<real_type> sample_eps(epsilon0_, half);
         epsilon = sample_eps(rng);
     }
@@ -186,14 +207,9 @@ CELER_FUNCTION Interaction BetheHeitlerInteractor::operator()(Engine& rng)
         // Calculate the minimum (when \epsilon = 1/2) and maximum (when
         // \epsilon = \epsilon_1) values of screening variable, \delta. Above
         // 50 MeV, a Coulomb correction function is introduced.
-        real_type const delta_min = 4 * 136 / element_.cbrt_z() * epsilon0_;
-        real_type f_z = real_type(8) / real_type(3) * element_.log_z();
-        if (inc_energy_ > value_as<Energy>(coulomb_corr_threshold()))
-        {
-            f_z += 8 * element_.coulomb_correction();
-        }
+        real_type const delta_min = 4 * screen_delta_;
         real_type const delta_max
-            = std::exp((real_type(42.038) - f_z) / real_type(8.29))
+            = std::exp((real_type(42.038) - f_z_) / real_type(8.29))
               - real_type(0.958);
         CELER_ASSERT(delta_min <= delta_max);
 
@@ -204,16 +220,15 @@ CELER_FUNCTION Interaction BetheHeitlerInteractor::operator()(Engine& rng)
         // introduced, where \epsilon_1 is the solution to
         // \Phi(\delta(\epsilon)) - F(Z)/2 = 0.
         real_type const epsilon1
-            = half - half * std::sqrt(1 - delta_min / delta_max);
+            = half * (1 - std::sqrt(1 - delta_min / delta_max));
         real_type const epsilon_min = celeritas::max(epsilon0_, epsilon1);
 
-        // Decide to choose f1, g1 or f2, g2 based on N1, N2 (factors from
-        // corrected Bethe-Heitler cross section; c.f. Eq. 6.6 of Geant4
-        // Physics Reference 10.6)
-        real_type const f10 = this->screening_f1(delta_min) - f_z;
-        real_type const f20 = this->screening_f2(delta_min) - f_z;
-        BernoulliDistribution choose_f1g1(ipow<2>(half - epsilon_min) * f10,
-                                          real_type(1.5) * f20);
+        // Decide to choose f1, g1 [brems] or f2, g2 [pair production]
+        // based on N1, N2 (factors from corrected Bethe-Heitler cross section;
+        // c.f. Eq. 6.6 of Geant4 Physics Reference 10.6)
+        Real2 const fmin = this->screening_f(delta_min) - f_z_;
+        BernoulliDistribution choose_f1g1(
+            ipow<2>(half - epsilon_min) * fmin[0], real_type(1.5) * fmin[1]);
 
         // Rejection function g_1 or g_2. Note the it's possible for g to be
         // greater than one
@@ -236,24 +251,23 @@ CELER_FUNCTION Interaction BetheHeitlerInteractor::operator()(Engine& rng)
                 // Calculate g_1 rejection function
                 if (enable_lpm_)
                 {
-                    auto screening = screening_phi1_phi2(delta);
+                    auto screening = screening_phi(delta);
                     auto lpm = calc_lpm_functions_(epsilon);
                     g = lpm.xi
-                        * ((2 * lpm.phi + lpm.g) * screening.phi1
-                           - lpm.g * screening.phi2 - lpm.phi * f_z)
-                        / f10;
+                        * ((2 * lpm.phi + lpm.g) * screening[0]
+                           - lpm.g * screening[1] - lpm.phi * f_z_);
                 }
                 else
                 {
-                    g = (this->screening_f1(delta) - f_z) / f10;
+                    g = this->screening_f(delta)[0] - f_z_;
                 }
+                g /= fmin[0];
                 CELER_ASSERT(g > 0);
             }
             else
             {
                 // Used to sample from f2
-                epsilon = epsilon_min
-                          + (half - epsilon_min) * generate_canonical(rng);
+                epsilon = UniformRealDistribution{epsilon_min, half}(rng);
                 CELER_ASSERT(epsilon >= epsilon_min && epsilon <= half);
 
                 // Calculate delta given the element atomic number and sampled
@@ -264,59 +278,58 @@ CELER_FUNCTION Interaction BetheHeitlerInteractor::operator()(Engine& rng)
                 // Calculate g_2 rejection function
                 if (enable_lpm_)
                 {
-                    auto screening = screening_phi1_phi2(delta);
+                    auto screening = screening_phi(delta);
                     auto lpm = calc_lpm_functions_(epsilon);
-                    g = lpm.xi
-                        * ((lpm.phi + half * lpm.g) * screening.phi1
-                           + half * lpm.g * screening.phi2
-                           - half * (lpm.g + lpm.phi) * f_z)
-                        / f20;
+                    g = half * lpm.xi
+                        * ((2 * lpm.phi + lpm.g) * screening[0]
+                           + lpm.g * screening[1] - (lpm.g + lpm.phi) * f_z_);
                 }
                 else
                 {
-                    g = (this->screening_f2(delta) - f_z) / f20;
+                    g = this->screening_f(delta)[1] - f_z_;
                 }
+                g /= fmin[1];
                 CELER_ASSERT(g > 0);
             }
+            // TODO: use rejection?
         } while (g < generate_canonical(rng));
     }
 
     // Construct interaction for change to primary (incident) particle (gamma)
     Interaction result = Interaction::from_absorption();
     result.secondaries = {secondaries, 2};
 
-    // Outgoing secondaries are electron and positron
+    // Outgoing secondaries are electron and positron with randomly selected
+    // charges
     secondaries[0].particle_id = shared_.ids.electron;
     secondaries[1].particle_id = shared_.ids.positron;
 
-    secondaries[0].energy = Energy{(1 - epsilon) * inc_energy_
+    // Incident energy split between the particles, with rest mass subtracted
+    secondaries[0].energy = Energy{(1 - epsilon) * value_as<Energy>(inc_energy_)
                                    - value_as<Mass>(shared_.electron_mass)};
-    secondaries[1].energy = Energy{epsilon * inc_energy_
+    secondaries[1].energy = Energy{epsilon * value_as<Energy>(inc_energy_)
                                    - value_as<Mass>(shared_.electron_mass)};
-
-    // Select charges for child particles (e-, e+) randomly
     if (BernoulliDistribution(half)(rng))
     {
         trivial_swap(secondaries[0].energy, secondaries[1].energy);
     }
 
-    // Sample secondary directions.  Note that momentum is not exactly
-    // conserved.
-    real_type phi
+    // Sample secondary directions: note that momentum is not exactly conserved
+    real_type const phi
         = UniformRealDistribution<real_type>(0, 2 * constants::pi)(rng);
+    auto sample_costheta = [&](Energy e) {
+        return TsaiUrbanDistribution{e, shared_.electron_mass}(rng);
+    };
 
-    // Electron
-    TsaiUrbanDistribution sample_electron_angle(secondaries[0].energy,
-                                                shared_.electron_mass);
-    real_type cost = sample_electron_angle(rng);
+    // Note that positron has opposite azimuthal angle
     secondaries[0].direction
-        = rotate(from_spherical(cost, phi), inc_direction_);
-    // Positron
-    TsaiUrbanDistribution sample_positron_angle(secondaries[1].energy,
-                                                shared_.electron_mass);
-    cost = sample_positron_angle(rng);
+        = rotate(from_spherical(sample_costheta(secondaries[0].energy), phi),
+                 inc_direction_);
     secondaries[1].direction
-        = rotate(from_spherical(cost, phi + constants::pi), inc_direction_);
+        = rotate(from_spherical(sample_costheta(secondaries[1].energy),
+                                phi + constants::pi),
+                 inc_direction_);
+
     return result;
 }
 
@@ -330,28 +343,33 @@ CELER_FUNCTION Interaction BetheHeitlerInteractor::operator()(Engine& rng)
 CELER_FUNCTION real_type
 BetheHeitlerInteractor::impact_parameter(real_type eps) const
 {
-    return 136 / element_.cbrt_z() * epsilon0_ / (eps * (1 - eps));
+    return screen_delta_ / (eps * (1 - eps));
 }
 
 //---------------------------------------------------------------------------//
 /*!
  * Screening functions \f$ \Phi_1(\delta) \f$ and \f$ \Phi_2(\delta) \f$.
+ *
+ * These correspond to \em f in Butcher: the first screening function is based
+ * on the bremsstrahlung cross sections, and the second is due to pair
+ * production. The values are improved function fits by M Novak (Geant4).
  */
-CELER_FUNCTION auto BetheHeitlerInteractor::screening_phi1_phi2(
-    real_type delta) const -> ScreeningFunctions
+CELER_FUNCTION auto
+BetheHeitlerInteractor::screening_phi(real_type delta) -> Real2
 {
     using R = real_type;
 
-    ScreeningFunctions result;
+    Real2 result;
     if (delta > R(1.4))
     {
-        result.phi1 = R(21.0190) - R(4.145) * std::log(delta + R(0.958));
-        result.phi2 = result.phi1;
+        result[0] = R(21.0190) - R(4.145) * std::log(delta + R(0.958));
+        result[1] = result[0];
     }
     else
     {
-        result.phi1 = R(20.806) - delta * (R(3.190) - R(0.5710) * delta);
-        result.phi2 = R(20.234) - delta * (R(2.126) - R(0.0903) * delta);
+        using PolyQuad = PolyEvaluator<real_type, 2>;
+        result[0] = PolyQuad{20.806, -3.190, 0.5710}(delta);
+        result[1] = PolyQuad{20.234, -2.126, 0.0903}(delta);
     }
     return result;
 }
@@ -361,27 +379,23 @@ CELER_FUNCTION auto BetheHeitlerInteractor::screening_phi1_phi2(
  * Auxiliary screening functions \f$ F_1(\delta) \f$ and \f$ F_2(\delta) \f$.
  *
  * The functions \f$ F_1 = 3 \Phi_1(\delta) - \Phi_2(\delta) \f$
- * and \f$ F_2 = 1.5\Phi_1(\delta) - 0.5\Phi_2(\delta) \f$
+ * and \f$ F_2 = 1.5\Phi_1(\delta) + 0.5\Phi_2(\delta) \f$
  * are decreasing functions of \f$ \delta \f$ for all \f$ \delta \f$
  * in \f$ [\delta_\textrm{min}, \delta_\textrm{max}] \f$.
  * They reach their maximum value at
  * \f$ \delta_\textrm{min} = \delta(\epsilon = 1/2)\f$. They are used in the
  * composition + rejection technique for sampling \f$ \epsilon \f$.
+ *
+ * Note that there's a typo in the Geant4 manual in the formula for F2:
+ * subtraction should be addition.
  */
-CELER_FUNCTION real_type BetheHeitlerInteractor::screening_f1(real_type delta) const
+CELER_FUNCTION auto
+BetheHeitlerInteractor::screening_f(real_type delta) -> Real2
 {
     using R = real_type;
-
-    return delta > R(1.4) ? R(42.038) - R(8.29) * std::log(delta + R(0.958))
-                          : R(42.184) - delta * (R(7.444) - R(1.623) * delta);
+    auto temp = screening_phi(delta);
+    return {3 * temp[0] - temp[1], R{1.5} * temp[0] + R{0.5} * temp[1]};
 }
 
-CELER_FUNCTION real_type BetheHeitlerInteractor::screening_f2(real_type delta) const
-{
-    using R = real_type;
-
-    return delta > R(1.4) ? R(42.038) - R(8.29) * std::log(delta + R(0.958))
-                          : R(41.326) - delta * (R(5.848) - R(0.902) * delta);
-}
 //---------------------------------------------------------------------------//
 }  // namespace celeritas