source: branches/numpy_misc/tools/test_numpy_numeric/test_numpy_numeric.py @ 7244

#!/usr/bin/env python

'''
A program to test differences in time and memory usage between Numeric
and numpy code.

Due to the above, this will run only on a machine that has both of the
above packages installed.
'''


import sys
import os
import time
import gc
import Numeric
import numpy


ARRAY_SIZE = 1000 * 1000
LOOP_SIZE = 1000


def mem_usage():
    '''Get memory usage (virtual) in KiB.'''

    _scale = {'KB': 1024, 'MB': 1024*1024, 'GB': 1024*1024*1024,
              'kB': 1024, 'mB': 1024*1024, 'gB': 1024*1024*1024}

    if sys.platform != 'win32':
        _proc_status = '/proc/%d/status' % os.getpid()
        
        def _VmB(VmKey):
            '''Get number of virtual bytes used.'''

            # get pseudo file /proc/<pid>/status
            try:
                t = open(_proc_status)
                v = t.read()
                t.close()
            except IOError:
                return 0.0

            # get VmKey line, eg: 'VmRSS: 999 kB\n ...
            i = v.index(VmKey)
            v = v[i:].split(None, 3)
            if len(v) < 3:
                return 0.0

            # convert Vm value to bytes
            return float(v[1]) * _scale[v[2]]

        return int(_VmB('VmSize:')/_scale['KB'])
    else:
        # Windows code from: http://code.activestate.com/recipes/511491/
        try:
            import ctypes
            import _winreg
        except:
            log(level, 'Windows resource usage not available')
            return

        kernel32 = ctypes.windll.kernel32
        c_ulong = ctypes.c_ulong
        c_ulonglong = ctypes.c_ulonglong
        class MEMORYSTATUSEX(ctypes.Structure):
            _fields_ = [('dwLength', c_ulong),
                        ('dwMemoryLoad', c_ulong),
                        ('ullTotalPhys', c_ulonglong),
                        ('ullAvailPhys', c_ulonglong),
                        ('ullTotalPageFile', c_ulonglong),
                        ('ullAvailPageFile', c_ulonglong),
                        ('ullTotalVirtual', c_ulonglong),
                        ('ullAvailVirtual', c_ulonglong),
                        ('ullAvailExtendedVirtual', c_ulonglong)
                       ]

        memoryStatusEx = MEMORYSTATUSEX()
        memoryStatusEx.dwLength = ctypes.sizeof(MEMORYSTATUSEX)
        kernel32.GlobalMemoryStatusEx(ctypes.byref(memoryStatusEx))

        return int(memoryStatusEx.ullTotalPhys/_scale['KB'])

def test_usage(module, A, B):
    # do some numeric calculations
    C = B
    for i in xrange(LOOP_SIZE):
        C = 2.6*A + B + C + i

# Do numpy work
start_time = time.time()
start_mem = mem_usage()

A = numpy.ones(ARRAY_SIZE, numpy.float)
B = numpy.ones(ARRAY_SIZE, numpy.float)
for i in xrange(ARRAY_SIZE):
    A[i] *= float(i)
    B[i] *= float(i)

test_usage(numpy, A, B)

stop_mem = mem_usage()
stop_time = time.time()

del A, B
gc.collect()

delta_time = stop_time - start_time
delta_mem = stop_mem - start_mem
print('  numpy: %.1f s, %d KiB' % (delta_time, delta_mem))

# Do Numeric work
start_time = time.time()
start_mem = mem_usage()

A = Numeric.ones(ARRAY_SIZE, Numeric.Float)
B = Numeric.ones(ARRAY_SIZE, Numeric.Float)
for i in xrange(ARRAY_SIZE):
    A[i] *= float(i)
    B[i] *= float(i)

test_usage(Numeric, A, B)

stop_mem = mem_usage()
stop_time = time.time()

del A, B
gc.collect()

delta_time = stop_time - start_time
delta_mem = stop_mem - start_mem
print('Numeric: %.1f s, %d KiB' % (delta_time, delta_mem))