add benchmarks for ray intersection and photon propagation

1 files changed, 21 insertions, 4 deletions

diff --git a/event.py b/event.py
index 5b93f17..657f414 100644
--- a/event.py
+++ b/event.py
@@ -1,13 +1,30 @@
 import numpy as np
 
+from chroma.sample import uniform_sphere
+
 class Photons(object):
-    def __init__(self, positions, directions, polarizations, times, wavelengths,
-                 last_hit_triangles=None, histories=None):
+    def __init__(self, positions, directions, wavelengths, polarizations=None, times=None, last_hit_triangles=None, histories=None):
         self.positions = positions
+
+        nphotons = len(positions)
+
+        assert len(directions) == nphotons
         self.directions = directions
-        self.polarizations = polarizations
-        self.times = times
+
+        assert len(wavelengths) == nphotons
         self.wavelengths = wavelengths
+
+        if polarizations is not None:
+            self.polarizations = polarizations
+        else:
+            # polarizations are isotropic in plane perpendicular
+            # to the direction vectors
+            polarizations = np.cross(directions, uniform_sphere(nphotons))
+            # normalize polarization vectors
+            polarizations /= np.tile(np.apply_along_axis(np.linalg.norm,1,polarizations),[3,1]).transpose()
+            self.polarizations = np.asarray(polarizations, order='C')
+
+        self.times = times
         self.last_hit_triangles = last_hit_triangles
         self.histories = histories