diff options
-rwxr-xr-x | sim.py | 8 |
1 files changed, 7 insertions, 1 deletions
@@ -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) |