add benchmarks for ray intersection and photon propagation

1 files changed, 99 insertions, 0 deletions

diff --git a/benchmark.py b/benchmark.py
new file mode 100644
index 0000000..0f72811
--- /dev/null
+++ b/benchmark.py
@@ -0,0 +1,99 @@
+import numpy as np
+from pycuda import gpuarray as ga
+import time
+from uncertainties import ufloat
+import sys
+
+from chroma.gpu import GPU, to_float3
+from chroma.camera import get_rays
+from chroma.event import Photons
+from chroma.sample import uniform_sphere
+
+def progress(seq):
+    "Print progress while iterating over `seq`."
+    n = len(seq)
+    print '[' + ' '*21 + ']\r[',
+    sys.stdout.flush()
+    for i, item in enumerate(seq):
+        if i % (n//10) == 0:
+            print '.',
+            sys.stdout.flush()
+        yield item
+    print ']'
+    sys.stdout.flush()
+
+def ray_trace(gpu, number=1000):
+    """
+    Return the number of mean and standard deviation of the number of ray
+    intersections per second as a ufloat for the geometry loaded onto `gpu`.
+
+    .. note::
+        The rays are thrown from a camera sitting *outside* of the geometry.
+
+    Args:
+        - gpu, chroma.gpu.GPU
+            The GPU object with a geometry already loaded.
+        - number, int
+            The number of kernel calls to average.
+    """
+    lb, ub = gpu.geometry.mesh.get_bounds()
+    scale = np.linalg.norm(ub-lb)
+    point = [0,scale,(lb[2]+ub[2])/2]
+
+    size = (800,600)
+    width, height = size
+
+    origins, directions = get_rays(point, size, 0.035, focal_length=0.018)
+
+    origins_gpu = ga.to_gpu(to_float3(origins))
+    directions_gpu = ga.to_gpu(to_float3(directions))
+    pixels_gpu = ga.zeros(width*height, dtype=np.int32)
+
+    run_times = []
+    for i in progress(range(number)):
+        t0 = time.time()
+        gpu.kernels.ray_trace(np.int32(pixels_gpu.size), origins_gpu, directions_gpu, pixels_gpu, block=(gpu.nthreads_per_block,1,1), grid=(pixels_gpu.size//gpu.nthreads_per_block+1,1))
+        gpu.context.synchronize()
+        elapsed = time.time() - t0
+        run_times.append(elapsed)
+
+    return pixels_gpu.size/ufloat((np.mean(run_times),np.std(run_times)))
+
+def propagate(gpu, number=10, nphotons=500000):
+    """
+    Return the mean and standard deviation of the number of photons propagated
+    per second as a ufloat for the geometry loaded onto `gpu`.
+
+    Args:
+        - gpu, chroma.gpu.GPU
+            The GPU object with a geometry already loaded.
+        - number, int
+            The number of kernel calls to average.
+        - nphotons, int
+            The number of photons to propagate per kernel call.
+    """
+    gpu.setup_propagate()
+
+    run_times = []
+    for i in progress(range(number)):
+        photons = Photons(np.zeros((nphotons,3)), uniform_sphere(nphotons), np.random.uniform(400,800,size=nphotons))
+        gpu.load_photons(photons)
+        t0 = time.time()
+        gpu.propagate()
+        gpu.context.synchronize()
+        elapsed = time.time() - t0
+        run_times.append(elapsed)
+
+    return nphotons/ufloat((np.mean(run_times),np.std(run_times)))
+
+if __name__ == '__main__':
+    from chroma.detectors import build_lbne_200kton, build_minilbne
+
+    lbne = build_lbne_200kton()
+    lbne.build(bits=11)
+
+    gpu = GPU()
+    gpu.load_geometry(lbne, print_usage=False)
+
+    print '%s track steps/s' % ray_trace(gpu)
+    print '%s photons/s' % propagate(gpu)