Fix energy-momentum conservation in Compton secondaries production

src/PROPOSAL/detail/PROPOSAL/secondaries/parametrization/compton/NaivCompton.cxx

@@ -15,12 +15,10 @@ double secondaries::NaivCompton::CalculateRho(
 std::tuple<Cartesian3D, Cartesian3D> secondaries::NaivCompton::CalculateDirections(
     const Vector3D& primary_dir, double energy, double v, double rnd)
 {
-    auto com_energy = energy + ME; // center of mass energy
-    auto kin_energy = energy;      // kinetic energy
-    auto cosphi_gamma = (com_energy * (1. - v) - ME)
-        / ((1. - v) * std::sqrt(com_energy * kin_energy));
-    auto cosphi_electron
-        = (com_energy * v - ME) / (v * std::sqrt(com_energy * kin_energy));
+    // use energy-momentum conservation formula
+    auto cosphi_gamma = 1 - (v * ME) / (energy * (1 - v));
+    auto cosphi_electron = v * (energy + ME) / std::sqrt(2 * v * energy * ME + v * v * energy * energy);
+
     auto rnd_theta = rnd * 2. * PI;
     auto dir_gamma = Cartesian3D(primary_dir);
     dir_gamma.deflect(cosphi_gamma, rnd_theta);
@@ -32,15 +30,15 @@ std::tuple<Cartesian3D, Cartesian3D> secondaries::NaivCompton::CalculateDirectio
 std::tuple<double, double> secondaries::NaivCompton::CalculateEnergy(
     double energy, double v)
 {
-    return std::make_tuple(energy * (1 - v), energy * v);
+    return std::make_tuple(energy * (1 - v), ME + energy * v);
 }
 
 std::vector<ParticleState> secondaries::NaivCompton::CalculateSecondaries(
         StochasticLoss loss, const Component&, std::vector<double> &rnd)
 {
     auto v = loss.energy /  loss.parent_particle_energy;
     auto secondary_energies = CalculateEnergy(loss.parent_particle_energy, v);
-    auto secondary_dir = CalculateDirections(loss.direction, loss.energy, v,
+    auto secondary_dir = CalculateDirections(loss.direction, loss.parent_particle_energy, v,
                                              rnd[0]);
 
     auto sec = std::vector<ParticleState>();

tests/SecondaryProduction_TEST.cxx

@@ -1,6 +1,7 @@
 #include "gtest/gtest.h"
 #include <cmath>
 
+#include "PROPOSAL/secondaries/parametrization/compton/NaivCompton.h"
 #include "PROPOSAL/secondaries/parametrization/ionization/NaivIonization.h"
 #include "PROPOSAL/secondaries/parametrization/photomupairproduction/PhotoMuPairProductionBurkhardtKelnerKokoulin.h"
 #include "PROPOSAL/secondaries/parametrization/photopairproduction/PhotoPairProductionKochMotz.h"
@@ -208,6 +209,47 @@ TEST(HeitlerAnnihilation, EnergyMomentumConservation)
     }
 }
 
+TEST(Compton, EnergyMomentumConservation)
+{
+    RandomGenerator::Get().SetSeed(1909);
+
+    auto particle = GammaDef();
+    auto medium = Air();
+    auto target = medium.GetComponents().front();
+
+    auto sec_calculator = secondaries::NaivCompton(particle, medium);
+
+    double E = 1e2;
+    double v = 0.4;
+    auto direction = Cartesian3D(0, 0, 1);
+
+    auto stochastic_loss = StochasticLoss((int)InteractionType::Compton, E * v, Cartesian3D(0, 0, 0),
+                                          direction, 0, 0, E);
+
+    // calculate random numbers
+    std::vector<double> rnd;
+    for (int i=0; i < sec_calculator.RequiredRandomNumbers(); i++) {
+        rnd.push_back(RandomGenerator::Get().RandomDouble());
+    }
+    auto secondaries = sec_calculator.CalculateSecondaries(stochastic_loss, target, rnd);
+
+    // energy conservation
+    EXPECT_NEAR(E + ME, secondaries.at(0).energy + secondaries.at(1).energy, COMPUTER_PRECISION);
+
+    // momentum conservation
+    ASSERT_EQ(secondaries.at(0).type, 22);
+    auto p_photon = secondaries.at(0).direction * secondaries.at(0).energy;
+
+    ASSERT_EQ(secondaries.at(1).type, 11);
+    auto E_electron = secondaries.at(1).energy;
+    auto p_electron = secondaries.at(1).direction * std::sqrt(E_electron * E_electron - ME * ME);
+
+    auto p_init = direction * E;
+
+    for (unsigned int i = 0; i<=2; i++)
+        EXPECT_NEAR(p_init[i], p_photon[i] + p_electron[i], COMPUTER_PRECISION);
+}
+
 int main(int argc, char** argv)
 {
     ::testing::InitGoogleTest(&argc, argv);