Replace Rayleigh scattering implementation with that from SNOMAN. The

angular distribution is slightly different, and now fits with the
distribution given in the GEANT4 physics reference manual.

Unit test is now included to verify the correctness of the scattering.

2 files changed, 97 insertions, 19 deletions

diff --git a/src/photon.h b/src/photon.h
index 45ad79b..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)
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)
+