Adding memory profiling demos and an exercise.

Author: Jonathan Rocher

memory_profile/mem_profile_test.py

@@ -0,0 +1,14 @@
+""" File to test the memory profiler package. To run it, 
+
+python -m memory_profiler mem_profile_test.py 
+"""
+
+@profile
+def my_func():
+    a = [1] * (10 ** 6)
+    b = [2] * (2 * 10 ** 7)
+    del b
+    return a
+
+if __name__ == '__main__':
+    my_func()

memory_profile/memory_footprint.py

@@ -0,0 +1,10 @@
+"""
+Write a function that takes an object and computes it total memory footprint. 
+Test your function on 
+- simple datastructures, 
+- custom classes
+- numpy arrays.
+"""
+
+from sys import getsizeof
+from gc import get_referents

memory_profile/memory_profiler_support.py

@@ -0,0 +1,11 @@
+""" Supporting function for memory profiler testing
+"""
+
+from numpy import array, int64, float32
+
+def my_func2(b):
+    a = 1.0 # Negligible footprint on the memory usage
+    c = array(b, dtype = int64)
+    c += a
+    c = c.astype(float32)
+    return c

memory_profile/memory_profiler_test_file.py

@@ -0,0 +1,36 @@
+""" Demo file for testing memory profiling with memory_profiler and heapy. 
+
+"""
+# Line below is necessary to run the code without -m memory_profiler
+from memory_profiler import profile
+from memory_profiler_support import my_func2
+
+    
+@profile
+def my_func():
+    result = {}
+    a = [1] * (10 ** 6)
+    b = [2] * (2 * 10 ** 7)
+    c = my_func2(b)
+    del b
+    result["a"] = a
+    result["c"] = c
+    return result
+
+class MyClass(object):
+    def __init__(self, x, y):
+        self.x = x
+        self.y = y
+        self.a = linspace(self.y)
+
+if __name__ == '__main__':
+    a = my_func()
+    
+    try:
+        from guppy import hpy
+        hp = hpy()
+        h1 = hp.heap()
+        print h1
+        print h1.bymodule
+    except ImportError:
+        print "heapy is not installed on your machine"

memory_profile/simple_heap.py

@@ -0,0 +1,76 @@
+""" Simple class to track down the memory usage of an object, including the 
+objects it points to that are likely to be part of the object.
+
+Based on an original version by Robert Kern
+"""
+from collections import deque
+import gc
+import sys
+import types
+
+from traits.api import Any, HasTraits, Int, Set
+from numpy import ndarray, dtype
+
+class SimpleHeap(HasTraits):
+    """ A simple heap memory tracker.
+    """
+
+    # The root object we are trying to ascertain the size of.
+    root = Any()
+
+    # The total size of the rooted object graph.
+    size = Int()
+
+    # The IDs of objects that have been seen.
+    _seen = Set(Int)
+
+    def __init__(self, root, **traits):
+        super(SimpleHeap, self).__init__(root=root, **traits)
+
+        self.compute_size()
+
+    def compute_size(self):
+        queue = deque([self.root])
+        while queue:
+            self.size += self._pop_and_size(queue)
+        print "Size of the object and what it contains is %s bytes." % self.size
+
+    def _pop_and_size(self, queue):
+        """ Pop an item off the queue and return its size. Add its referents to 
+        the queue.
+        """
+        obj = queue.popleft()
+        if id(obj) in self._seen:
+            return 0
+        else:
+            self._seen.add(id(obj))
+        if isinstance(obj, ndarray):
+            # The elements of a numpy array are not seen by the gc as a referent
+            container_size = sys.getsizeof(obj)
+            if obj.dtype != dtype('object'): 
+                return container_size + obj.nbytes
+            else:
+                # In this case, the element is a pointer (4 bytes). 
+                # Convert to a list to find the referents.
+                referents = self._filtered_referents(obj.tolist())
+                queue.extend(referents)
+                return container_size
+        else:
+            size = sys.getsizeof(obj)
+            referents = self._filtered_referents(obj)
+            queue.extend(referents)
+        return size
+
+    @staticmethod
+    def _filtered_referents(obj):
+        """ Return the referents of an object that we consider as "part"
+        of the object using the garbage collector.
+        """
+        referents = gc.get_referents(obj)
+        badtypes = (types.CodeType, types.ModuleType, types.FunctionType,
+            types.ClassType, types.TypeType, types.FrameType,
+            types.GeneratorType, types.GetSetDescriptorType, types.LambdaType,
+            types.MemberDescriptorType, types.MethodType, types.TracebackType,
+            types.UnboundMethodType)
+        good = [x for x in referents if not isinstance(x, badtypes)]
+        return good