add a pi0 decay generator

1 files changed, 102 insertions, 0 deletions

diff --git a/pi0.py b/pi0.py
new file mode 100644
index 0000000..fa624d0
--- /dev/null
+++ b/pi0.py
@@ -0,0 +1,102 @@
+import numpy as np
+
+#c = 2.99792458e8
+#h = 6.6260689633e-34
+#joules_per_MeV = 1.602176487e-19/1e-6
+kg_per_MeV = 1.782661758e-36/1e-6
+pi0_mass = 134.9766*kg_per_MeV
+
+def rocket_to_lab(energy, momentum, v):
+    """
+    Return the energy and momentum of a particle in the lab frame from its
+    energy and momentum in a rocket frame traveling at a velocity `v` with
+    respect to the lab frame.
+
+    Args:
+        - energy: float
+            The energy of the particle in the rocket frame in kg.
+        - momentum: array
+            The momentum of the particle in the rocket frame in kg.
+        - v: array
+            The velocity of the rocket frame as seen in the lab frame in units
+            of c.
+    """
+    e0 = float(energy)#/c**2
+    p0 = np.asarray(momentum, float)#/c
+    v = np.asarray(v, float)#/c
+
+    try:
+        assert e0**2 - p0.dot(p0) >= -1.0e-70
+    except AssertionError:
+        print e0**2 - p0.dot(p0)
+        raise
+
+    g = 1.0/np.sqrt(1.0-v.dot(v))
+
+    x = np.dot(p0, v)/np.linalg.norm(v)
+    p = p0 + ((g-1.0)*x + g*np.linalg.norm(v)*e0)*v/np.linalg.norm(v)
+    e = np.sqrt(e0**2 - p0.dot(p0) + p.dot(p))
+
+    #return e*c**2, p*c
+    return e, p
+
+def pi0_decay(energy, direction, theta, phi):
+    """
+    Return the energy and directions for two photons produced from a pi0
+    decay in the lab frame given that one of the photons decayed at polar
+    angles `theta` and `phi` in the pi0 rest frame.
+
+    Args:
+        - energy: float
+            The total energy of the pi0 in MeV.
+        - direction: array
+            The direction of the pi0's velocity in the lab frame.
+    """
+    direction = np.asarray(direction)/np.linalg.norm(direction)
+    pi0_e = float(energy)*kg_per_MeV#/c**2
+    pi0_p = np.sqrt(pi0_e**2-pi0_mass**2)*direction
+    pi0_v = pi0_p/pi0_e
+
+    photon_e0 = pi0_mass/2.0
+    photon_p0 = photon_e0*np.array([np.cos(theta)*np.sin(phi), np.sin(theta)*np.sin(phi), np.cos(phi)])
+
+    #print photon_p0/np.linalg.norm(photon_p0)
+
+    e1, p1 = rocket_to_lab(photon_e0, photon_p0, pi0_v)
+    v1 = p1/np.linalg.norm(p1)
+    e2, p2 = rocket_to_lab(photon_e0, -photon_p0, pi0_v)
+    v2 = p2/np.linalg.norm(p2)
+
+    return (e1/kg_per_MeV, v1), (e2/kg_per_MeV, v2)
+
+if __name__ == '__main__':
+    import sys
+
+    npi0s = 10000
+
+    pi0_e = 300.0*kg_per_MeV
+    pi0_p = np.sqrt(pi0_e**2 - pi0_mass**2)
+
+    print 'pi0 e: %f MeV' % (pi0_e/kg_per_MeV)
+    print 'pi0 p: %f MeV' % (pi0_p/kg_per_MeV)
+
+    e, cos_theta = [], []
+    for i, (theta, phi) in enumerate(zip(np.random.rand(npi0s), np.arccos(np.random.uniform(-1.0, 1.0, size=npi0s)))):
+        print '\r%i' % (i+1),
+        sys.stdout.flush()
+
+        (e1, v1), (e2, v2) = pi0_decay(pi0_e/kg_per_MeV, [0,0,1], theta, phi)
+        e += [e1, e2]
+        cos_theta += [v1.dot(v2)]
+
+    import matplotlib.pyplot as plt
+
+    plt.figure()
+    plt.title('Energy Distribution of Photons from a %.0f MeV Pi0 Decay' % (pi0_e/kg_per_MeV))
+    plt.hist(e, 100)
+    plt.xlabel('Energy (MeV)')
+    plt.figure()
+    plt.title('Opening Angle between Photons in Lab from a %.0f MeV Pi0 Decay' % (pi0_e/kg_per_MeV))
+    plt.hist(cos_theta, 100)
+    plt.xlabel('Cos($\theta$)')
+    plt.show()