Add optional weight to each photon being propagated.

For consistence, weights must be less than or equal to one at all
times.  When weight calculation is enabled by the likelihood
calculator, photons are prevented from being absorbed, and instead
their weight is decreased to reflect their probability of survival.
Once the photon's weight drops below a given threshold (0.01 for now),
the weighting procedure is stopped and the photon can be extinguished.

With weighting enabled, the detection efficiency of surfaces is also
applied to the weight, and the photon terminated with the DETECT bit
set in the history.  This is not completely accurate, as a photon
could pass through the surface, reflect, and reintersect the surface
later (or some other PMT) and be detected.  As a result, weighting
will slightly underestimate PMT efficiency compared to the true Monte
Carlo.  This is not intrinsic to the weighting procedure, but only
comes about because of the inclusion of detection efficiency into the
weight.

Without the detection efficiency included, weighting cuts in half the
number of evaluations required to achieve a given likelihood
uncertainty (at least for the hit probabilities).  Add in the
detection efficiency, and that factor becomes 1/5 or 1/6!

3 files changed, 89 insertions, 34 deletions

diff --git a/chroma/cuda/daq.cu b/chroma/cuda/daq.cu
index 3ed402d..719d928 100644
--- a/chroma/cuda/daq.cu
+++ b/chroma/cuda/daq.cu
@@ -36,6 +36,7 @@ __global__ void
 run_daq(curandState *s, unsigned int detection_state,
 	int first_photon, int nphotons, float *photon_times,
 	unsigned int *photon_histories, int *last_hit_triangles,
+	float *weights,
 	int *solid_map,
 	Detector *detector,
 	unsigned int *earliest_time_int,
@@ -55,22 +56,27 @@ run_daq(curandState *s, unsigned int detection_state,
 	    int channel_index = detector->solid_id_to_channel_index[solid_id];
 
 	    if (channel_index >= 0 && (history & detection_state)) {
-		float time = photon_times[photon_id] + 
-		    sample_cdf(&rng, detector->time_cdf_len,
-			       detector->time_cdf_x, detector->time_cdf_y);
-		unsigned int time_int = float_to_sortable_int(time);
+
+		float weight = weights[photon_id];
+		if (curand_uniform(&rng) < weight) {
+		    float time = photon_times[photon_id] + 
+			sample_cdf(&rng, detector->time_cdf_len,
+				   detector->time_cdf_x, detector->time_cdf_y);
+		    unsigned int time_int = float_to_sortable_int(time);
 			  
-		float charge = sample_cdf(&rng, detector->charge_cdf_len,
+		    float charge = sample_cdf(&rng, detector->charge_cdf_len,
 					  detector->charge_cdf_x,
 					  detector->charge_cdf_y);
-		unsigned int charge_int = roundf(charge / detector->charge_unit);
+		    unsigned int charge_int = roundf(charge / detector->charge_unit);
 
-		atomicMin(earliest_time_int + channel_index, time_int);
-		atomicAdd(channel_q_int + channel_index, charge_int);
-		atomicOr(channel_histories + channel_index, history);
-	    }
+		    atomicMin(earliest_time_int + channel_index, time_int);
+		    atomicAdd(channel_q_int + channel_index, charge_int);
+		    atomicOr(channel_histories + channel_index, history);
+		} // if weighted photon contributes
 
-	}
+	    } // if photon detected by a channel
+
+	} // if photon terminated on surface
 
 	s[id] = rng;
 	
@@ -82,6 +88,7 @@ __global__ void
 run_daq_many(curandState *s, unsigned int detection_state,
 	     int first_photon, int nphotons, float *photon_times,
 	     unsigned int *photon_histories, int *last_hit_triangles,
+	     float *weights,
 	     int *solid_map,
 	     Detector *detector,
 	     unsigned int *earliest_time_int,
@@ -94,7 +101,7 @@ run_daq_many(curandState *s, unsigned int detection_state,
     __shared__ int channel_index;
     __shared__ unsigned int history;
     __shared__ float photon_time;
-    
+    __shared__ float weight;
 
     if (threadIdx.x == 0) {
 	photon_id = first_photon + blockIdx.x;
@@ -105,6 +112,7 @@ run_daq_many(curandState *s, unsigned int detection_state,
 	    history = photon_histories[photon_id];
 	    channel_index = detector->solid_id_to_channel_index[solid_id];
 	    photon_time = photon_times[photon_id];
+	    weight = weights[photon_id];
 	}
     }
 
@@ -119,19 +127,21 @@ run_daq_many(curandState *s, unsigned int detection_state,
     for (int i = threadIdx.x; i < ndaq; i += blockDim.x) {
 	int channel_offset = channel_index + i * channel_stride;
 
-	float time = photon_time + curand_normal(&rng) * 1.2f;// + 
-	//sample_cdf(&rng, detector->time_cdf_len,
-	//	       detector->time_cdf_x, detector->time_cdf_y);
-	unsigned int time_int = float_to_sortable_int(time);
-	
-	float charge = 1.0f; //sample_cdf(&rng, detector->charge_cdf_len,
-	//detector->charge_cdf_x,
-	//detector->charge_cdf_y);
-	unsigned int charge_int = roundf(charge / detector->charge_unit);
+	if (curand_uniform(&rng) < weight) {
+	    float time = photon_time + curand_normal(&rng) * 1.2f;// + 
+	    //sample_cdf(&rng, detector->time_cdf_len,
+	    //	       detector->time_cdf_x, detector->time_cdf_y);
+	    unsigned int time_int = float_to_sortable_int(time);
 	
-	atomicMin(earliest_time_int + channel_offset, time_int);
-	atomicAdd(channel_q_int + channel_offset, charge_int);
-	atomicOr(channel_histories + channel_offset, history);
+	    float charge = 1.0f; //sample_cdf(&rng, detector->charge_cdf_len,
+	    //detector->charge_cdf_x,
+	    //detector->charge_cdf_y);
+	    unsigned int charge_int = roundf(charge / detector->charge_unit);
+	    
+	    atomicMin(earliest_time_int + channel_offset, time_int);
+	    atomicAdd(channel_q_int + channel_offset, charge_int);
+	    atomicOr(channel_histories + channel_offset, history);
+	}
     }
 
     s[id] = rng;
diff --git a/chroma/cuda/photon.h b/chroma/cuda/photon.h
index 8b7e588..79a87f0 100644
--- a/chroma/cuda/photon.h
+++ b/chroma/cuda/photon.h
@@ -9,6 +9,8 @@
 #include "mesh.h"
 #include "geometry.h"
 
+#define WEIGHT_LOWER_THRESHOLD 0.01f
+
 struct Photon
 {
     float3 position;
@@ -16,7 +18,9 @@ struct Photon
     float3 polarization;
     float wavelength;
     float time;
-
+  
+    float weight;
+  
     unsigned int history;
 
     int last_hit_triangle;
@@ -175,11 +179,17 @@ rayleigh_scatter(Photon &p, curandState &rng)
     p.polarization /= norm(p.polarization);
 } // scatter
 
-__device__ int propagate_to_boundary(Photon &p, State &s, curandState &rng)
+__device__ int propagate_to_boundary(Photon &p, State &s, curandState &rng,
+                                     bool use_weights=false)
 {
     float absorption_distance = -s.absorption_length*logf(curand_uniform(&rng));
     float scattering_distance = -s.scattering_length*logf(curand_uniform(&rng));
 
+    if (use_weights && p.weight > WEIGHT_LOWER_THRESHOLD) // Prevent absorption
+	absorption_distance = 1e30;
+    else
+	use_weights = false;
+
     if (absorption_distance <= scattering_distance) {
 	if (absorption_distance <= s.distance_to_boundary) {
 	    p.time += absorption_distance/(SPEED_OF_LIGHT/s.refractive_index1);
@@ -193,6 +203,11 @@ __device__ int propagate_to_boundary(Photon &p, State &s, curandState &rng)
     }
     else {
 	if (scattering_distance <= s.distance_to_boundary) {
+
+	    // Scale weight by absorption probability along this distance
+	    if (use_weights)
+		p.weight *= expf(-scattering_distance/s.absorption_length);
+
 	    p.time += scattering_distance/(SPEED_OF_LIGHT/s.refractive_index1);
 	    p.position += scattering_distance*p.direction;
 
@@ -206,6 +221,10 @@ __device__ int propagate_to_boundary(Photon &p, State &s, curandState &rng)
 	} // photon is scattered in material1
     } // if scattering_distance < absorption_distance
 
+    // Scale weight by absorption probability along this distance
+    if (use_weights)
+	p.weight *= expf(-s.distance_to_boundary/s.absorption_length);
+    
     p.position += s.distance_to_boundary*p.direction;
     p.time += s.distance_to_boundary/(SPEED_OF_LIGHT/s.refractive_index1);
 
@@ -269,7 +288,8 @@ propagate_at_boundary(Photon &p, State &s, curandState &rng)
 } // propagate_at_boundary
 
 __device__ int
-propagate_at_surface(Photon &p, State &s, curandState &rng, Geometry *geometry)
+propagate_at_surface(Photon &p, State &s, curandState &rng, Geometry *geometry,
+                     bool use_weights=false)
 {
     Surface *surface = geometry->surfaces[s.surface_index];
 
@@ -283,6 +303,25 @@ propagate_at_surface(Photon &p, State &s, curandState &rng, Geometry *geometry)
 
     float uniform_sample = curand_uniform(&rng);
 
+    if (use_weights && p.weight > WEIGHT_LOWER_THRESHOLD 
+	&& absorb < (1.0f - WEIGHT_LOWER_THRESHOLD)) {
+	// Prevent absorption and reweight accordingly
+	float survive = 1.0f - absorb;
+	absorb = 0.0f;
+	p.weight *= survive;
+
+	// Renormalize remaining probabilities
+	detect /= survive;
+	reflect_diffuse /= survive;
+	reflect_specular /= survive;
+    }
+
+    if (use_weights && detect > 0.0f) {
+	p.history |= SURFACE_DETECT;
+	p.weight *= detect;
+	return BREAK;
+    }
+
     if (uniform_sample < absorb) {
 	p.history |= SURFACE_ABSORB;
 	return BREAK;
diff --git a/chroma/cuda/propagate.cu b/chroma/cuda/propagate.cu
index 87ba9f5..29f8aea 100644
--- a/chroma/cuda/propagate.cu
+++ b/chroma/cuda/propagate.cu
@@ -14,7 +14,7 @@ photon_duplicate(int first_photon, int nthreads,
 		 float3 *positions, float3 *directions,
 		 float *wavelengths, float3 *polarizations,
 		 float *times, unsigned int *histories,
-		 int *last_hit_triangles, 
+		 int *last_hit_triangles, float *weights,
 		 int copies, int stride)
 {
     int id = blockIdx.x*blockDim.x + threadIdx.x;
@@ -32,6 +32,7 @@ photon_duplicate(int first_photon, int nthreads,
     p.time = times[photon_id];
     p.last_hit_triangle = last_hit_triangles[photon_id];
     p.history = histories[photon_id];
+    p.weight = weights[photon_id];
 
     for (int i=1; i <= copies; i++) {
       int target_photon_id = photon_id + stride * i;
@@ -43,6 +44,7 @@ photon_duplicate(int first_photon, int nthreads,
       times[target_photon_id] = p.time;
       last_hit_triangles[target_photon_id] = p.last_hit_triangle;
       histories[target_photon_id] = p.history;
+      weights[target_photon_id] = p.weight;
     }
 }
 
@@ -79,11 +81,11 @@ copy_photons(int first_photon, int nthreads, unsigned int target_flag,
 	     float3 *positions, float3 *directions,
 	     float *wavelengths, float3 *polarizations,
 	     float *times, unsigned int *histories,
-	     int *last_hit_triangles, 
+	     int *last_hit_triangles, float *weights,
 	     float3 *new_positions, float3 *new_directions,
 	     float *new_wavelengths, float3 *new_polarizations,
 	     float *new_times, unsigned int *new_histories,
-	     int *new_last_hit_triangles)
+	     int *new_last_hit_triangles, float *new_weights)
 {
     int id = blockIdx.x*blockDim.x + threadIdx.x;
     
@@ -102,6 +104,7 @@ copy_photons(int first_photon, int nthreads, unsigned int target_flag,
 	new_times[offset] = times[photon_id];
 	new_histories[offset] = histories[photon_id];
 	new_last_hit_triangles[offset] = last_hit_triangles[photon_id];
+	new_weights[offset] = weights[photon_id];
     }
 }
 
@@ -112,7 +115,8 @@ propagate(int first_photon, int nthreads, unsigned int *input_queue,
 	  float3 *positions, float3 *directions,
 	  float *wavelengths, float3 *polarizations,
 	  float *times, unsigned int *histories,
-	  int *last_hit_triangles, int max_steps,
+	  int *last_hit_triangles, float *weights,
+	  int max_steps, int use_weights,
 	  Geometry *g)
 {
     __shared__ Geometry sg;
@@ -143,6 +147,7 @@ propagate(int first_photon, int nthreads, unsigned int *input_queue,
     p.time = times[photon_id];
     p.last_hit_triangle = last_hit_triangles[photon_id];
     p.history = histories[photon_id];
+    p.weight = weights[photon_id];
 
     if (p.history & (NO_HIT | BULK_ABSORB | SURFACE_DETECT | SURFACE_ABSORB))
 	return;
@@ -166,7 +171,7 @@ propagate(int first_photon, int nthreads, unsigned int *input_queue,
 	if (p.last_hit_triangle == -1)
 	    break;
 
-	command = propagate_to_boundary(p, s, rng);
+	command = propagate_to_boundary(p, s, rng, use_weights);
 
 	if (command == BREAK)
 	    break;
@@ -175,7 +180,7 @@ propagate(int first_photon, int nthreads, unsigned int *input_queue,
 	    continue;
 
 	if (s.surface_index != -1) {
-	    command = propagate_at_surface(p, s, rng, g);
+	  command = propagate_at_surface(p, s, rng, g, use_weights);
 
 	    if (command == BREAK)
 		break;
@@ -196,7 +201,8 @@ propagate(int first_photon, int nthreads, unsigned int *input_queue,
     times[photon_id] = p.time;
     histories[photon_id] = p.history;
     last_hit_triangles[photon_id] = p.last_hit_triangle;
-	
+    weights[photon_id] = p.weight;
+
     // Not done, put photon in output queue
     if ((p.history & (NO_HIT | BULK_ABSORB | SURFACE_DETECT | SURFACE_ABSORB)) == 0) {
 	int out_idx = atomicAdd(output_queue, 1);