add some tests to gen-dark-matter

1 files changed, 68 insertions, 3 deletions

diff --git a/utils/gen-dark-matter b/utils/gen-dark-matter
index 06b146f..04ffeee 100755
--- a/utils/gen-dark-matter
+++ b/utils/gen-dark-matter
@@ -38,6 +38,7 @@ from __future__ import print_function, division
 import numpy as np
 import string
 import uuid
+from itertools import islice
 
 SNOMAN_MASS = {
     20: 0.511,
@@ -177,8 +178,9 @@ def rand_sphere():
 
 def lorentz_boost(x,n,beta):
     """
-    Performs a Lorentz boost on the 4-vector `x` along the direction `n` at a
-    velocity `beta` (in units of the speed of light).
+    Performs a Lorentz boost on the 4-vector `x`. Returns the 4-vector as would
+    be seen by a frame moving along the direction `n` at a velocity `beta` (in
+    units of the speed of light).
 
     Formula comes from 46.1 in the PDG "Kinematics" Review. See
     http://pdg.lbl.gov/2014/reviews/rpp2014-rev-kinematics.pdf.
@@ -236,6 +238,70 @@ def gen_decay(M, E, m1, m2):
 
         yield v1_new, v2_new
 
+def test_gen_decay1():
+    """
+    A super simple test that if we generate a decay with E = mass, the two
+    particles come out back to back.
+    """
+    for v1, v2 in islice(gen_decay(100,100,0.5,0.5),100):
+        dir1 = v1[1:]/np.linalg.norm(v1[1:])
+        dir2 = v2[1:]/np.linalg.norm(v2[1:])
+        assert np.isclose(np.dot(dir1,dir2),-1.0)
+
+def test_gen_decay2():
+    """
+    A super simple test that if we generate a decay with E >> mass, the two
+    particles come out in the same direction.
+    """
+    for v1, v2 in islice(gen_decay(100,100e3,0.5,0.5),100):
+        dir1 = v1[1:]/np.linalg.norm(v1[1:])
+        dir2 = v2[1:]/np.linalg.norm(v2[1:])
+        assert np.isclose(np.dot(dir1,dir2),1.0,atol=1e-3)
+
+def test_lorentz_boost2():
+    """
+    Simple test of lorentz_boost() using a problem from
+    http://electron6.phys.utk.edu/PhysicsProblems/Mechanics/8-Relativity/decay%20massless.html.
+    """
+    M = 140.0
+    P = 2e3
+    E = np.sqrt(P**2 + M**2)
+    m1 = 105.0
+    m2 = 0.0
+    # muon in lab frame is going in +z direction which means that the direction
+    # of the muon in the pion decay frame is also in the +z direction
+    v = np.array([0,0,1])
+    # calculate momentum of each daughter particle
+    # FIXME: need to double check my math
+    p1 = np.sqrt(((M**2-m1**2-m2**2)**2-4*m1**2*m2**2)/(4*M**2))
+    E1 = np.sqrt(m1**2 + p1**2)
+
+    if not np.isclose(E1,109.375):
+        print("Energy of muon in pion rest frame = %.2e but expected %.2e" % (E1,109.375))
+
+    assert np.isclose(E1,109.375)
+
+    E2 = np.sqrt(m2**2 + p1**2)
+    v1 = [E1] + (p1*v).tolist()
+    v2 = [E2] + (-p1*v).tolist()
+    # pion in lab frame is going in the +z direction, therefore the lab frame
+    # is going in the -z direction relative to the pion frame
+    n = np.array([0,0,-1])
+    beta = np.sqrt(E**2-M**2)/E
+
+    if not np.isclose(beta,0.99756,rtol=1e-3):
+        print("Speed of pion frame relative to lab frame = %.5f but expected %.5f" % (beta,0.99756))
+
+    assert np.isclose(beta,0.99756,rtol=1e-3)
+
+    v1_new = lorentz_boost(v1,n,beta)
+    v2_new = lorentz_boost(v2,n,beta)
+
+    if not np.isclose(v1_new[0],2003.8,atol=1e-2):
+        print("Energy of muon in lab frame = %.1f but expected %.1f" % (v1_new[0],2003.8))
+
+    assert np.isclose(v1_new[0],2003.8,atol=1e-2)
+
 def test_lorentz_boost():
     """
     Super simple test of the lorentz_boost() function to make sure that if we
@@ -253,7 +319,6 @@ def test_lorentz_boost():
 
 if __name__ == '__main__':
     import argparse
-    from itertools import islice
     import sys
     import os