1 files changed, 50 insertions, 20 deletions

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;