A simple unit test of the constant particle gun

1 files changed, 23 insertions, 0 deletions

diff --git a/tests/test_generator_vertex.py b/tests/test_generator_vertex.py
new file mode 100644
index 0000000..d773579
--- /dev/null
+++ b/tests/test_generator_vertex.py
@@ -0,0 +1,23 @@
+import unittest
+import itertools
+import numpy as np
+import chroma.generator.vertex
+
+class TestParticleGun(unittest.TestCase):
+    def test_constant_particle_gun_center(self):
+        '''Generate electron vertices at the center of the world volume.'''
+        vertex = chroma.generator.vertex.constant_particle_gun('e-', (0,0,0), (1,0,0), 100)
+        for ev in itertools.islice(vertex, 100):
+            self.assertEquals(ev.particle_name, 'e-')
+            self.assertTrue(np.allclose(ev.gen_position, [0,0,0]))
+            self.assertTrue(np.allclose(ev.gen_direction, [1,0,0]))
+            self.assertTrue(np.allclose(ev.gen_total_energy, 100))
+
+    def test_off_center(self):
+        '''Generate electron vertices at (1,0,0) in the world volume.'''
+        vertex = chroma.generator.vertex.constant_particle_gun('e-', (1,0,0), (1,0,0), 100)
+        for ev in itertools.islice(vertex, 100):
+            self.assertEquals(ev.particle_name, 'e-')
+            self.assertTrue(np.allclose(ev.gen_position, [1,0,0]))
+            self.assertTrue(np.allclose(ev.gen_direction, [1,0,0]))
+            self.assertTrue(np.allclose(ev.gen_total_energy, 100))