merge

5 files changed, 161 insertions, 22 deletions

diff --git a/gpu.py b/gpu.py
index 519cf0f..eeb759e 100644
--- a/gpu.py
+++ b/gpu.py
@@ -306,7 +306,7 @@ class GPU(object):
                 input_queue_gpu = output_queue_gpu
                 output_queue_gpu = temp
                 output_queue_gpu[:1].set(np.uint32(1))
-                nphotons = input_queue_gpu[:1].get()[0]
+                nphotons = input_queue_gpu[:1].get()[0] - 1
         
         if gpuarray.max(self.histories_gpu).get() & (1 << 31):
             print >>sys.stderr, "WARNING: ABORTED PHOTONS IN THIS EVENT"
diff --git a/sim.py b/sim.py
index fce559e..c52a071 100755
--- a/sim.py
+++ b/sim.py
@@ -85,7 +85,7 @@ class Simulation(object):
 
     def propagate_photons(self, photons, max_steps=10):
         self.gpu_worker.load_photons(photons)
-        self.gpu_worker.propagate()
+        self.gpu_worker.propagate(max_steps=max_steps)
         return self.gpu_worker.get_photons()
 
 @profile_if_possible
diff --git a/src/photon.h b/src/photon.h
index f76ca54..e781864 100644
--- a/src/photon.h
+++ b/src/photon.h
@@ -105,30 +105,52 @@ __device__ void fill_state(State &s, Photon &p)
 
 } // fill_state
 
-__device__ void rayleigh_scatter(Photon &p, curandState &rng)
+__device__ float3 pick_new_direction(float3 axis, float theta, float phi)
 {
-	float theta, y;
-
-	while (true)
-	{
-		y = curand_uniform(&rng);
-		theta = uniform(&rng, 0, 2*PI);
-
-		if (y < powf(cosf(theta),2))
-			break;
+	// Taken from SNOMAN rayscatter.for
+	float cos_theta = cosf(theta);
+	float sin_theta = sinf(theta);
+	float cos_phi   = cosf(phi);
+	float sin_phi   = sinf(phi);
+	
+	float sin_axis_theta = sqrt(1.0f - axis.z*axis.z);
+	float cos_axis_phi, sin_axis_phi;
+	
+	if (isnan(sin_axis_theta) || sin_axis_theta < 0.00001f) {
+		cos_axis_phi = 1.0f;
+		sin_axis_phi = 0.0f;
+	} else {
+		cos_axis_phi = axis.x / sin_axis_theta;
+		sin_axis_phi = axis.y / sin_axis_theta;
 	}
 
-	float phi = uniform(&rng, 0, 2*PI);
-
-	float3 b = cross(p.polarization, p.direction);
-	float3 c = p.polarization;
-
-	p.direction = rotate(p.direction, theta, b);
-	p.direction = rotate(p.direction, phi, c);
+	return make_float3(cos_theta*axis.x + sin_theta*(axis.z*cos_phi*cos_axis_phi - sin_phi*sin_axis_phi),
+			   cos_theta*axis.y + sin_theta*(cos_phi*axis.z*sin_axis_phi - sin_phi*cos_axis_phi),
+			   cos_theta*axis.z - sin_theta*cos_phi*sin_axis_theta);
+}
 
-	p.polarization = rotate(p.polarization, theta, b);
-	p.polarization = rotate(p.polarization, phi, c);
+__device__ void rayleigh_scatter(Photon &p, curandState &rng)
+{
+	float cos_theta = 2.0f*cosf((acosf(1.0f - 2.0f*curand_uniform(&rng))-2*PI)/3.0f);
+	if (cos_theta > 1.0f)
+	  cos_theta = 1.0f;
+	else if (cos_theta < -1.0f)
+	  cos_theta = -1.0f;
+
+	float theta = acosf(cos_theta);
+	float phi = uniform(&rng, 0.0f, 2.0f * PI);
+
+	p.direction = pick_new_direction(p.polarization, theta, phi);
+
+	if (1.0f - fabsf(cos_theta) < 1e-6f) {
+		p.polarization = pick_new_direction(p.polarization, PI/2.0f, phi);
+	} else {
+		// linear combination of old polarization and new direction
+		p.polarization = p.polarization - cos_theta * p.direction;
+	}
 
+	p.direction /= norm(p.direction);
+	p.polarization /= norm(p.polarization);
 } // scatter
 
 __device__ int propagate_to_boundary(Photon &p, State &s, curandState &rng)
@@ -179,7 +201,15 @@ __device__ void propagate_at_boundary(Photon &p, State &s, curandState &rng)
 	float refracted_angle = asinf(sinf(incident_angle)*s.refractive_index1/s.refractive_index2);
 
 	float3 incident_plane_normal = cross(p.direction, s.surface_normal);
-	incident_plane_normal /= norm(incident_plane_normal);
+	float incident_plane_normal_length = norm(incident_plane_normal);
+
+	// Photons at normal incidence do not have a unique plane of incidence,
+	// so we have to pick the plane normal to be the polarization vector
+	// to get the correct logic below
+	if (incident_plane_normal_length < 1e-6f)
+	  incident_plane_normal = p.polarization;
+	else
+	  incident_plane_normal /= incident_plane_normal_length;
 
 	float normal_coefficient = dot(p.polarization, incident_plane_normal);
 	float normal_probability = normal_coefficient*normal_coefficient;
diff --git a/tests/propagation.py b/tests/propagation.py
new file mode 100644
index 0000000..331242b
--- /dev/null
+++ b/tests/propagation.py
@@ -0,0 +1,53 @@
+import unittest
+import numpy as np
+
+from chroma.geometry import Solid, Geometry
+from chroma.make import box
+from chroma.sim import Simulation
+from chroma.optics import vacuum
+from chroma.event import Photons
+
+class TestPropagation(unittest.TestCase):
+    def testAbort(self):
+        '''Photons that hit a triangle at normal incidence should not abort.
+
+        Photons that hit a triangle at exactly normal incidence can sometimes produce a dot product
+        that is outside the range allowed by acos().  Trigger these with axis aligned photons in a box
+        '''
+
+        # Setup geometry
+        cube = Geometry()
+        cube.add_solid(Solid(box(100,100,100), vacuum, vacuum))
+        cube.pmtids = [0]
+        sim = Simulation(cube, vacuum, bvh_bits=4, geant4_processes=0,
+                              use_cache=False)
+
+        # Create initial photons
+        nphotons = 10000
+        positions = np.tile([0,0,0], (nphotons,1)).astype(np.float32)
+        directions = np.tile([0,0,1], (nphotons,1)).astype(np.float32)
+        polarizations = np.zeros_like(positions)
+        phi = np.random.uniform(0, 2*np.pi, nphotons).astype(np.float32)
+        polarizations[:,0] = np.cos(phi)
+        polarizations[:,1] = np.sin(phi)
+        times = np.zeros(nphotons, dtype=np.float32)
+        wavelengths = np.empty(nphotons, np.float32)
+        wavelengths.fill(400.0)
+
+        photons = Photons(positions=positions, directions=directions, polarizations=polarizations,
+                          times=times, wavelengths=wavelengths)
+
+        # First make one step to check for strangeness
+        photon_stop = sim.propagate_photons(photons, max_steps=1)
+        self.assertFalse(np.isnan(photon_stop.positions).any())
+        self.assertFalse(np.isnan(photon_stop.directions).any())
+        self.assertFalse(np.isnan(photon_stop.polarizations).any())
+        self.assertFalse(np.isnan(photon_stop.times).any())
+        self.assertFalse(np.isnan(photon_stop.wavelengths).any())
+
+        # Now let it run the usual ten steps
+        photon_stop = sim.propagate_photons(photons, max_steps=10)
+        aborted = (photon_stop.histories & (1 << 31)) > 0
+        print 'aborted photons: %1.1f' % (float(np.count_nonzero(aborted)) / nphotons)
+        self.assertFalse(aborted.any())
+        
diff --git a/tests/rayleigh.py b/tests/rayleigh.py
new file mode 100644
index 0000000..4394ada
--- /dev/null
+++ b/tests/rayleigh.py
@@ -0,0 +1,56 @@
+import unittest
+import numpy as np
+
+from chroma.geometry import Solid, Geometry
+from chroma.make import box
+from chroma.sim import Simulation
+from chroma.optics import water_wcsim
+from chroma.event import Photons
+import histogram
+from histogram.root import rootify
+import ROOT
+ROOT.gROOT.SetBatch(1)
+
+class TestRayleigh(unittest.TestCase):
+    def setUp(self):
+        self.cube = Geometry()
+        self.cube.add_solid(Solid(box(100,100,100), water_wcsim, water_wcsim))
+        self.cube.pmtids = [0]
+        self.sim = Simulation(self.cube, water_wcsim, bvh_bits=4, geant4_processes=0,
+                              use_cache=False)
+        nphotons = 100000
+        positions = np.tile([0,0,0], (nphotons,1)).astype(np.float32)
+        directions = np.tile([0,0,1], (nphotons,1)).astype(np.float32)
+        polarizations = np.zeros_like(positions)
+        phi = np.random.uniform(0, 2*np.pi, nphotons).astype(np.float32)
+        polarizations[:,0] = np.cos(phi)
+        polarizations[:,1] = np.sin(phi)
+        times = np.zeros(nphotons, dtype=np.float32)
+        wavelengths = np.empty(nphotons, np.float32)
+        wavelengths.fill(400.0)
+
+        self.photons = Photons(positions=positions, directions=directions, polarizations=polarizations,
+                               times=times, wavelengths=wavelengths)
+
+    def testAngularDistributionPolarized(self):
+        # Fully polarized photons
+        self.photons.polarizations[:] = [1.0, 0.0, 0.0]
+
+        photon_stop = self.sim.propagate_photons(self.photons, max_steps=1)
+        aborted = (photon_stop.histories & (1 << 31)) > 0
+        self.assertFalse(aborted.any())
+
+        # Compute the dot product between initial and final directions
+        rayleigh_scatters = (photon_stop.histories & (1 << 4)) > 0
+        cos_scatter = (self.photons.directions[rayleigh_scatters] * photon_stop.directions[rayleigh_scatters]).sum(axis=1)
+        theta_scatter = np.arccos(cos_scatter)
+        h = histogram.Histogram(bins=100, range=(0, np.pi))
+        h.fill(theta_scatter)
+        h = rootify(h)
+
+        # The functional form for polarized light should be (1 + \cos^2 \theta)\sin \theta
+        # according to GEANT4 physics reference manual
+        f = ROOT.TF1("pol_func", "[0]*(1+cos(x)**2)*sin(x)", 0, np.pi)
+        h.Fit(f)
+        self.assertGreater(f.GetProb(), 1e-3)
+