1 files changed, 29 insertions, 15 deletions
diff --git a/benchmark.py b/benchmark.py
index 8c75d35..c2f7ac9 100755
--- a/benchmark.py
+++ b/benchmark.py
@@ -228,26 +228,40 @@ if __name__ == '__main__':
     from chroma import detectors
     import gc
 
+    # Default to run all tests
+    tests = ['ray', 'load', 'propagate', 'pdf', 'pdf_eval']
+    if len(sys.argv) > 1:
+        tests = sys.argv[1:] # unless test names given on command line
+
     lbne = detectors.lbne()
     lbne.build(bits=11)
 
     context = gpu.create_cuda_context()
     gpu_geometry = gpu.GPUGeometry(lbne)
     
-    print '%s ray intersections/sec.' % \
-        tools.ufloat_to_str(intersect(gpu_geometry))
-    # run garbage collection since there is a reference loop
-    # in the GPUArray class.
-    gc.collect()
-    print '%s photons loaded/sec.' % tools.ufloat_to_str(load_photons())
-    gc.collect()
-    print '%s photons propagated/sec.' % \
-        tools.ufloat_to_str(propagate(gpu_geometry))
-    gc.collect()
-    print '%s 100 MeV events histogrammed/s' % \
-        tools.ufloat_to_str(pdf(gpu_geometry, max(lbne.pmtids)))
-
-    print '%s 100 MeV events/s accumulated in PDF evaluation data structure (100 GEANT4 x 16 Chroma x 128 DAQ)' % \
-        tools.ufloat_to_str(pdf_eval(gpu_geometry, max(lbne.pmtids)))
+    if 'ray' in tests:
+        print '%s ray intersections/sec.' % \
+            tools.ufloat_to_str(intersect(gpu_geometry))
+        # run garbage collection since there is a reference loop
+        # in the GPUArray class.
+        gc.collect()
+
+    if 'load' in tests:
+        print '%s photons loaded/sec.' % tools.ufloat_to_str(load_photons())
+        gc.collect()
+
+    if 'propagate' in tests:
+        print '%s photons propagated/sec.' % \
+            tools.ufloat_to_str(propagate(gpu_geometry))
+        gc.collect()
+
+    if 'pdf' in tests:
+        print '%s 100 MeV events histogrammed/s' % \
+            tools.ufloat_to_str(pdf(gpu_geometry, max(lbne.pmtids)))
+        gc.collect()
+
+    if 'pdf_eval' in tests:
+        print '%s 100 MeV events/s accumulated in PDF evaluation data structure (100 GEANT4 x 16 Chroma x 128 DAQ)' % \
+            tools.ufloat_to_str(pdf_eval(gpu_geometry, max(lbne.pmtids)))
 
     context.pop()