Enthought / ~ischnell

mkufunc

Ilan Schnell, 2008-07-01

Download:

Example:

>>> from numpy import arange
>>> from mkufunc.api import mkufunc
>>> @mkufunc
... def foo(x):
...     return 4.2 * x*x - x + 6.3
...
>>> a = arange(5)
>>> a
array([0, 1, 2, 3, 4])
>>> foo(a)
array([  6.3,   9.5,  21.1,  41.1,  69.5])

Benchmark (example above but array of size 10 Million):

Method Runtime (sec)Speed vs.
numpy.vectorize 8.674 69.9
x as numpy.array 0.467 3.8
weave.blitz 0.135 1.1
mkufunc 0.124 1.0
hand weave (inline) 0.076 0.61

Count prime example:

def count_primes(N):
    res = 0
    for n in xrange(2, N):
        for i in xrange(2, min(n, int(sqrt(n)+2.0))):
            if n % i == 0:
                break
        else:
            res += 1
    return res

Translated python function:

Program overview (pseudo Python code):

def func_hash(f, salt=None):
    return md5sum(bytecode(f)+repr(f))

def translate(f, argtypes):
    try:
        return open(cache_file_name).read()
    except IOError:
        import pypy
        make the pypy translation and store output in c_source_filename
        os.rename(c_source_filename, cache_file_name)
        return translate(f, argtypes)

class Cfunc(object):
    f           -- the Python function object
    n           -- id number
    sig         -- signature
    nin         -- number of input arguments
    nout        -- number of output arguments
    cname       -- name of the C function

    cfunc()     -- returns the C function (as string)
    ufunc_support_code()
                -- generate the C support code to make this
                   function part work with PyUFuncGenericFunction

def support_code(cfuncs):
    return header + support_code

def code(f, signatures):
    return '/* %s */' % func_hash(f) + \
           'return_val = PyUFunc_FromFuncAndData(...)';

def genufunc(f, signatures):
    cfuncs = [Cfunc(f, sig, n) for n, sig in enumerate(signatures)]
    return weave.inline(code(f, signatures),
                        support_code=support_code(cfuncs),
                        customize=ufunc_info)

def mkufunc(arg0=[float]):
    class UFunc(object):
        def __init__(self, f):
            check(signatures)

        self.ufunc = genufunc(f, signatures)

        def __call__(self, *args):
            return self.ufunc(*args)

    if isinstance(arg0, FunctionType):
        f = arg0
        signatures = [float]
        return UFunc(f)

    elif isinstance(arg0, list):
        signatures = arg0
        return UFunc

    elif arg0 in typedict.keys():
        signatures = [arg0]
        return UFunc