Output device usage to logger.

1 files changed, 22 insertions, 15 deletions

diff --git a/chroma/gpu/geometry.py b/chroma/gpu/geometry.py
index a223006..081563f 100644
--- a/chroma/gpu/geometry.py
+++ b/chroma/gpu/geometry.py
@@ -7,6 +7,7 @@ from chroma.geometry import standard_wavelengths
 from chroma.gpu.tools import get_cu_module, get_cu_source, cuda_options, \
     chunk_iterator, format_array, format_size, to_uint3, to_float3, \
     make_gpu_struct
+from chroma.log import logger
 
 class GPUGeometry(object):
     def __init__(self, geometry, wavelengths=None, print_usage=False):
@@ -134,23 +135,29 @@ class GPUGeometry(object):
 
         if print_usage:
             self.print_device_usage()
+        logger.info(self.device_usage_str())
+
+    def device_usage_str(self):
+        '''Returns a formatted string displaying the memory usage.'''
+        s = 'device usage:\n'
+        s += '-'*10 + '\n'
+        s += format_array('vertices', self.vertices) + '\n'
+        s += format_array('triangles', self.triangles) + '\n'
+        s += format_array('lower_bounds', self.lower_bounds) + '\n'
+        s += format_array('upper_bounds', self.upper_bounds) + '\n'
+        s += format_array('node_map', self.node_map) + '\n'
+        s += format_array('node_map_end', self.node_map_end) + '\n'
+        s += '%-15s %6s %6s' % ('total', '', format_size(self.vertices.nbytes + self.triangles.nbytes + self.lower_bounds.nbytes + self.upper_bounds.nbytes + self.node_map.nbytes + self.node_map_end.nbytes)) + '\n'
+        s += '-'*10 + '\n'
+        free, total = cuda.mem_get_info()
+        s += '%-15s %6s %6s' % ('device total', '', format_size(total)) + '\n'
+        s += '%-15s %6s %6s' % ('device used', '', format_size(total-free)) + '\n'
+        s += '%-15s %6s %6s' % ('device free', '', format_size(free)) + '\n'
+        return s
 
     def print_device_usage(self):
-        print 'device usage:'
-        print '-'*10
-        print format_array('vertices', self.vertices)
-        print format_array('triangles', self.triangles)
-        print format_array('lower_bounds', self.lower_bounds)
-        print format_array('upper_bounds', self.upper_bounds)
-        print format_array('node_map', self.node_map)
-        print format_array('node_map_end', self.node_map_end)
-        print '%-15s %6s %6s' % ('total', '', format_size(self.vertices.nbytes + self.triangles.nbytes + self.lower_bounds.nbytes + self.upper_bounds.nbytes + self.node_map.nbytes + self.node_map_end.nbytes))
-        print '-'*10
-        free, total = cuda.mem_get_info()
-        print '%-15s %6s %6s' % ('device total', '', format_size(total))
-        print '%-15s %6s %6s' % ('device used', '', format_size(total-free))
-        print '%-15s %6s %6s' % ('device free', '', format_size(free))
-        print
+        print self.device_usage_str()
+        print 
 
     def reset_colors(self):
         self.colors.set_async(self.geometry.colors.astype(np.uint32))