Simulation.simulate() now accepts event.Vertex objects.

1 files changed, 7 insertions, 1 deletions

diff --git a/sim.py b/sim.py
index c0df524..931858f 100755
--- a/sim.py
+++ b/sim.py
@@ -61,12 +61,18 @@ class Simulation(object):
         self.pdf_config = None
 
     def simulate(self, iterable, keep_photons_beg=False, keep_photons_end=False, run_daq=True, max_steps=10):
-        first_element, iterable = peek(iterable)
+        try:
+            first_element, iterable = peek(iterable)
+        except TypeError:
+            first_element, iterable = iterable, [iterable]
 
         if isinstance(first_element, event.Event):
             iterable = self.photon_generator.generate_events(iterable)
         elif isinstance(first_element, event.Photons):
             iterable = (event.Event(photons_beg=x) for x in iterable)
+        elif isinstance(first_element, event.Vertex):
+            iterable = (event.Event(vertices=[vertex]) for vertex in iterable)
+            iterable = self.photon_generator.generate_events(iterable)
 
         for ev in iterable:
             gpu_photons = gpu.GPUPhotons(ev.photons_beg)