merge

1 files changed, 34 insertions, 40 deletions

diff --git a/gpu.py b/gpu.py
index 7e45c3d..327e692 100644
--- a/gpu.py
+++ b/gpu.py
@@ -13,6 +13,7 @@ import src
 from geometry import standard_wavelengths
 from color import map_to_color
 import sys
+import event
 
 cuda.init()
 
@@ -232,40 +233,31 @@ class GPU(object):
                                  grid=(self.max_blocks,1))
         #self.context.synchronize()
 
-    def load_photons(self, pos, dir, pol, wavelength, t0, 
-                     histories=None, last_hit_triangles=None):
-        '''Load N photons onto the GPU.
-
-           pos: numpy.array(shape=(N, 3)) of photon starting positions (meters)
-           dir: numpy.array(shape=(N, 3)) of photon starting directions (unit vectors)
-           pol: numpy.array(shape=(N, 3)) of photon polarization directions (unit vectors)
-           wavelength: numpy.array(shape=N) of photon wavelengths (nm)
-           t0: numpy.array(shape=N) of photon start times (s)
-           
-           Optional args will be loaded with defaults on GPU if not set:
-           histories: Bitmask of interactions that have occurred over history of photon
-           last_hit_triangles: The triangle ID number that the photon last interacted with,
-                               if any.  -1 if no triangle was hit in the last step
+    def load_photons(self, photons):
+        '''Load photons onto the GPU from an event.Photons object.
+
+           If photon.histories or photon.last_hit_triangles are set to none,
+           they will be initialized to 0 and -1 on the GPU, respectively.
         '''
-        self.nphotons = len(pos)
-        assert len(dir) == self.nphotons
-        assert len(pol) == self.nphotons
-        assert len(wavelength) == self.nphotons
-        assert len(t0) == self.nphotons
-
-        self.positions_gpu = gpuarray.to_gpu(pos.astype(np.float32).view(gpuarray.vec.float3))
-        self.directions_gpu = gpuarray.to_gpu(dir.astype(np.float32).view(gpuarray.vec.float3))
-        self.polarizations_gpu = gpuarray.to_gpu(pol.astype(np.float32).view(gpuarray.vec.float3))
-        self.wavelengths_gpu = gpuarray.to_gpu(wavelength.astype(np.float32))
-        self.times_gpu = gpuarray.to_gpu(t0.astype(np.float32))
-
-        if histories is not None:
-            self.histories_gpu = gpuarray.to_gpu(histories.astype(np.uint32))
+        self.nphotons = len(photons.positions)
+        assert len(photons.directions) == self.nphotons
+        assert len(photons.polarizations) == self.nphotons
+        assert len(photons.wavelengths) == self.nphotons
+        assert len(photons.times) == self.nphotons
+
+        self.positions_gpu = gpuarray.to_gpu(photons.positions.astype(np.float32).view(gpuarray.vec.float3))
+        self.directions_gpu = gpuarray.to_gpu(photons.directions.astype(np.float32).view(gpuarray.vec.float3))
+        self.polarizations_gpu = gpuarray.to_gpu(photons.polarizations.astype(np.float32).view(gpuarray.vec.float3))
+        self.wavelengths_gpu = gpuarray.to_gpu(photons.wavelengths.astype(np.float32))
+        self.times_gpu = gpuarray.to_gpu(photons.times.astype(np.float32))
+
+        if photons.histories is not None:
+            self.histories_gpu = gpuarray.to_gpu(photons.histories.astype(np.uint32))
         else:
             self.histories_gpu = gpuarray.zeros(self.nphotons, dtype=np.uint32)
 
-        if last_hit_triangles is not None:
-            self.last_hit_triangles_gpu = gpuarray.to_gpu(last_hit_triangles.astype(np.int32))
+        if photons.last_hit_triangles is not None:
+            self.last_hit_triangles_gpu = gpuarray.to_gpu(photons.last_hit_triangles.astype(np.int32))
         else:
             self.last_hit_triangles_gpu = gpuarray.empty(self.nphotons, dtype=np.int32)
             self.last_hit_triangles_gpu.fill(-1)
@@ -328,13 +320,13 @@ class GPU(object):
 
         Contents of dictionary have the same names as the parameters to load_photons().
         '''
-        return { 'pos' : self.positions_gpu.get().view(np.float32).reshape(self.positions_gpu.size, 3),
-                 'dir' : self.directions_gpu.get().view(np.float32).reshape(self.directions_gpu.size, 3),
-                 'pol' : self.polarizations_gpu.get().view(np.float32).reshape(self.polarizations_gpu.size, 3),
-                 'wavelength' : self.wavelengths_gpu.get(),
-                 't0' : self.times_gpu.get(),
-                 'histories' : self.histories_gpu.get(),
-                 'last_hit_triangles' : self.last_hit_triangles_gpu.get()}
+        return event.Photons(positions=self.positions_gpu.get().view(np.float32).reshape(self.positions_gpu.size, 3),
+                             directions=self.directions_gpu.get().view(np.float32).reshape(self.directions_gpu.size, 3),
+                             polarizations=self.polarizations_gpu.get().view(np.float32).reshape(self.polarizations_gpu.size, 3),
+                             wavelengths=self.wavelengths_gpu.get(),
+                             times=self.times_gpu.get(),
+                             histories=self.histories_gpu.get(),
+                             last_hit_triangles=self.last_hit_triangles_gpu.get())
     
     def setup_daq(self, max_pmt_id, pmt_rms=1.2e-9):
         self.earliest_time_gpu = gpuarray.GPUArray(shape=(max_pmt_id+1,), dtype=np.float32)
@@ -378,9 +370,11 @@ class GPU(object):
         
 
     def get_hits(self):
-        return { 't': self.earliest_time_gpu.get(), 
-                 'q': self.channel_q_gpu.get().astype(np.float32),
-                 'history': self.channel_history_gpu.get()}
+        t = self.earliest_time_gpu.get()
+        # For now, assume all channels with small enough hit time were hit
+        return event.Channels(hit=t<1e8, t=t, 
+                        q=self.channel_q_gpu.get().astype(np.float32),
+                        histories=self.channel_history_gpu.get())
 
     def __del__(self):
         self.context.pop()