Why `vectorize` is outperformed by `frompyfunc`?

The question is entirely moot. If speed is the question, then neither vectorize, nor frompyfunc, is the answer. Any speed difference between them pales into insignificance compared to faster ways of doing it.

I found this question wondering why frompyfunc broke my code (it returns objects), whereas vectorize worked (it returned what I told it to do), and found people talking about speed.

Now, in the 2020s, numba/jit is available, which blows any speed advantage of frompyfunc clean out of the water.

I coded a toy application, returning a large array of np.uint8 from another one, and got the following results.

pure python                      200 ms
vectorize                         58 ms
frompyfunc + cast back to uint8   53 ms   
np.empty + numba/njit             55 us (4 cores, 100 us single core)

So 1000x speedup over numpy, and 4000x over pure python

I can post the code if anyone is bothered. Coding the njit version involved little more than adding the line @njit before the pure python function, so you don’t need to be hardcore to do it.

It is less convenient than wrapping your function in vectorize, as you have to write the looping over the numpy array stuff manually, but it does avoid writing an external C function. You do need to write in a numpy/C-like subset of python, and avoid python objects.

Perhaps I’m being hard on numpy here, asking it to vectorise a pure python function. So, what if I benchmark a numpy native array function like min against numba?

Staggeringly, I got a 10x speedup using numba/jit over np.min on a 385×360 array of np.uint8. 230 us for np.min(array) was the baseline. Numba achieved 60 us using a single core, and 22 us with all four cores.

# minimum graphical reproducible case of difference between
# frompyfunc and vectorize

# apparently, from stack overflow,
# vectorize returns correct type, but is slow
# frompyfunc always returns an object

# let's see which is faster, casting frompyfunc, or plain vectorize
# and then compare those with plain python, and with njit

# spoiler
# python       200 ms
# vectorise     58 ms
# frompyfunc    53 ms
# njit parallel 55 us


from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
import sys
import time
from numba import njit, prange

THRESH = 128
FNAME = '3218_M.PNG' # monochrome screen grab of a sudoku puzzle
ROW = 200

def th_python(x, out):
    rows, cols = x.shape
    for row in range(rows):
        for col in range(cols):
            val = 250
            if x[row, col]<THRESH:
                val = 5
            out[row, col] = val

@njit(parallel=True)
def th_jit(x, out):
    rows, cols = x.shape
    for row in prange(rows):
        for col in prange(cols):
            val = 250
            if x[row, col]<THRESH:
                val = 5
            out[row, col] = val

@njit(parallel=True)
def min_jit(x):
    rows, cols = x.shape
    minval = 255
    for row in prange(rows):
        for col in prange(cols):
            val = x[row, col]
            if val<minval:
                minval = val
    return minval
           

def threshold(x):
    out = 250
    if x<THRESH:
        out = 5
    return np.uint8(out)

th_fpf = np.frompyfunc(threshold,1,1)
th_vec = np.vectorize(threshold, otypes=[np.uint8])

# load an image
image = Image.open(FNAME)
print(f'{image.mode=}')
npim = np.array(image)

# see what we've got
print(f'{type(npim)=}')
print(f'{type(npim[0,0])=}')
# print(npim[ROW,:])
print(f'{npim.shape=}')
print(f'{sys.getsizeof(npim)=}')

# plt.imshow(npim, cmap='gray', vmin=0, vmax=255)
# plt.show()

# threshold it with plain python
start = time.time()
npimpp = np.empty(npim.shape, dtype=np.uint8)
alloc = time.time()
th_python(npim, npimpp)
done = time.time()
print(f'nallocation took {alloc-start:g} seconds')
print(f'computation took {done-alloc:g} seconds')
print(f'total plain python took {done-start:g} seconds')
print(f'{sys.getsizeof(npimpp)=}')

# use vectorize
start = time.time()
npimv = th_vec(npim)
done = time.time()
print(f'nvectorize took {done-start:g} seconds')
print(f'{sys.getsizeof(npimv)=}')

# use fpf followed by cast
start = time.time()
npimf = th_fpf(npim)
comp = time.time()
npimfc = np.array(npimf, dtype=np.uint8)
done = time.time()
print(f'nfunction took {comp-start:g} seconds')
print(f'cast took {done-comp:g} seconds')
print(f'total was {done-start:g} seconds')
print(f'{sys.getsizeof(npimf)=}')

# threshold it with numba jit
for i in range(2):
    print(f'n go number {i}')
    start = time.time()
    npimjit = np.empty(npim.shape, dtype=np.uint8)
    alloc = time.time()
    th_jit(npim, npimjit)
    done = time.time()
    print(f'nallocation took {alloc-start:g} seconds')
    print(f'computation took {done-alloc:g} seconds')
    print(f'total with jit took {done-start:g} seconds')
    print(f'{sys.getsizeof(npimjit)=}')

# what about if we use a numpy native function?
start = time.time()
npmin = np.min(npim)
done = time.time()
print(f'ntotal for np.min was {done-start:g} seconds')

for i in range(2):
    print(f'n go number {i}')
    start = time.time()
    jit_min = min_jit(npim)
    done = time.time()
    print(f'total with min_jit took {done-start:g} seconds')