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

#!/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


# Numeric version
def Numeric_ensure_numeric(A, typecode=None):
    """Ensure that sequence is a numeric array.

    Inputs:
        A: Sequence. If A is already a Numeric array it will be returned
                     unaltered
                     If not, an attempt is made to convert it to a Numeric
                     array
        A: Scalar.   Return 0-dimensional array of length 1, containing that value
        A: String.   Array of ASCII values
        typecode: Numeric type. If specified, use this in the conversion.
                                If not, let Numeric decide

    This function is necessary as array(A) can cause memory overflow.
    """

    if typecode is None:
        if type(A) == Numeric.ArrayType:
            return A
        else:
            return Numeric.array(A)
    else:
        if type(A) == Numeric.ArrayType:
            if A.typecode() == typecode:
                return A
            else:
                return Numeric.array(A, typecode)
        else:
            return Numeric.array(A, typecode)


# numpy version
def numpy_ensure_numeric(A, typecode=None):
    """Ensure that sequence is a numeric array.

    Inputs:
        A: Sequence. If A is already a numeric array it will be returned
                     unaltered
                     If not, an attempt is made to convert it to a numeric
                     array
        A: Scalar.   Return 0-dimensional array containing that value. Note
                     that a 0-dim array DOES NOT HAVE A LENGTH UNDER numpy.
        A: String.   Array of ASCII values (numpy can't handle this)

        typecode:    numeric type. If specified, use this in the conversion.
                     If not, let numeric package decide.
                     typecode will always be one of num.float, num.int, etc.

    Note that num.array(A, dtype) will sometimes copy.  Use 'copy=False' to
    copy only when required.

    This function is necessary as array(A) can cause memory overflow.
    """

#    if isinstance(A, basestring):
#        msg = 'Sorry, cannot handle strings in ensure_numeric()'
#        raise Exception, msg

    if typecode is None:
        if isinstance(A, numpy.ndarray):
            return A
        else:
            return numpy.array(A)
    else:
        return numpy.array(A, dtype=typecode, copy=False)


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, f, en, ls):
    start_time = time.time()
    start_mem = mem_usage()

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

    # do some numeric calculations
    A = en(A)
    B = en(B)
    C = module.ones(ARRAY_SIZE, f)
    C = en(C)
    for i in xrange(ls):
        C = en(2.6*A + B + C + i)

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

    delta_time = stop_time - start_time
    delta_mem = stop_mem - start_mem

    del A, B, C
    gc.collect()

    return (delta_time, delta_mem)

for loop_size in (10, 100, 1000):
    # Do numpy work
    (t, m) = test_usage(numpy, numpy.float, numpy_ensure_numeric, loop_size)
    print('  numpy %4d loops: %5.1f s, %d KiB' % (loop_size, t, m))

    # Do Numeric work
    (t, m) = test_usage(Numeric, Numeric.Float, Numeric_ensure_numeric,
                        loop_size)
    print('Numeric %4d loops: %5.1f s, %d KiB' % (loop_size, t, m))