Vectorizing Python Functions

Reviewed & published by Brayan K

np.vectorize wraps an ordinary Python function so it accepts whole arrays and broadcasts its arguments, giving you clean element-wise syntax — though it is a convenience for readability, not a real speed-up.

Part of the free NumPy course at LearnCodingFast — hands-on lessons with examples you run in your browser, plus practice exercises and a quick quiz.

You'll wrap functions with np.vectorize and np.frompyfunc, apply functions along an axis with np.apply_along_axis, and see why true vectorization with ufuncs and broadcasting is dramatically faster.

What You'll Learn in This Lesson

• Wrap a Python function with np.vectorize
• Why vectorize is convenience, not speed
• Turn a function into a ufunc with np.frompyfunc
• Apply a function along an axis with apply_along_axis
• How true ufunc vectorization beats these helpers
• When these convenience wrappers are acceptable

1 np.vectorize: Clean Syntax for Element-wise Logic

Suppose you have a plain Python function with some branching that does not map neatly onto array operations. np.vectorize(func) returns a new callable that accepts arrays, applies func to each element, and broadcasts multiple arguments. The code reads cleanly, with no explicit loop in sight.

import numpy as np

def label(score):
    return "pass" if score >= 60 else "fail"

vlabel = np.vectorize(label)
vlabel(np.array([55, 70, 60, 40]))
# array(['fail', 'pass', 'pass', 'fail'], dtype='<U4')

import numpy as np

def label(score):
    # arbitrary per-element Python logic
    if score >= 90:
        return "A"
    elif score >= 60:
        return "pass"
    else:
        return "fail"

vlabel = np.vectorize(label)

scores = np.array([95, 72, 58, 60])
print(vlabel(scores))
# ['A' 'pass' 'fail' 'pass']

# It also broadcasts multiple arguments
def clamp(x, lo):
    return max(x, lo)

vclamp = np.vectorize(clamp)
print(vclamp(np.array([1, 5, 9]), 4))
# [4 5 9]

# ✅ Expected output:
# ['A' 'pass' 'fail' 'pass']
# [4 5 9]

2 frompyfunc and apply_along_axis

np.frompyfunc(func, nin, nout) is a lower-level cousin that turns a function of nin inputs and nout outputs into an element-wise ufunc — but it returns object arrays, so you often cast the result. Separately, np.apply_along_axis(func, axis, arr) runs a function over whole 1D slices (rows or columns) of a 2D array.

# frompyfunc: 1 input, 1 output, returns object dtype
square = np.frompyfunc(lambda v: v * v, 1, 1)
square(np.array([1, 2, 3]))          # array([1, 4, 9], dtype=object)

# apply_along_axis: run on each row
def span(row):
    return row.max() - row.min()
np.apply_along_axis(span, 1, m)      # one value per row

import numpy as np

# frompyfunc returns an OBJECT array; cast to int to do math
square = np.frompyfunc(lambda v: v * v, 1, 1)
result = square(np.array([1, 2, 3, 4]))
print(result)                        # [1 4 9 16]
print(result.dtype)                  # object
print(result.astype(int).sum())      # 30

# apply_along_axis: reduce each row to one number
m = np.array([[1, 9, 4],
              [3, 3, 3],
              [0, 5, 2]])

def span(row):
    return row.max() - row.min()

print(np.apply_along_axis(span, 1, m))   # [8 0 5]

# ✅ Expected output:
# [1 4 9 16]
# object
# 30
# [8 0 5]

3 True Vectorization Wins

Whenever a calculation can be written with NumPy's built-in ufuncs and array operations, prefer that. Real vectorization runs in fast compiled C over the whole array at once, while np.vectorize loops in Python. The same logic often rewrites cleanly using np.where for branching or plain arithmetic.

scores = np.array([55, 70, 60, 40])

# Convenience: per-element Python loop
np.vectorize(lambda s: "pass" if s >= 60 else "fail")(scores)

# True vectorization: one compiled operation
np.where(scores >= 60, "pass", "fail")    # same result, much faster

import numpy as np

scores = np.array([55, 70, 60, 40, 88])

# Convenience wrapper (Python loop under the hood)
slow = np.vectorize(lambda s: "pass" if s >= 60 else "fail")
print(slow(scores))
# ['fail' 'pass' 'pass' 'fail' 'pass']

# True vectorization with np.where -> same answer, compiled speed
fast = np.where(scores >= 60, "pass", "fail")
print(fast)
# ['fail' 'pass' 'pass' 'fail' 'pass']

# They match
print(np.array_equal(slow(scores), fast))   # True

# ✅ Expected output:
# ['fail' 'pass' 'pass' 'fail' 'pass']
# ['fail' 'pass' 'pass' 'fail' 'pass']
# True

🎯 YOUR TURN: Fill in the Blank

Wrap the function so it can take an array of inputs by filling in the wrapper name.

import numpy as np

def grade(score):
    return "pass" if score >= 50 else "fail"

# Wrap grade so it works element-wise on arrays
vgrade = np.___(grade)

print(vgrade(np.array([40, 60, 50])))
# Expected output:
# ['fail' 'pass' 'pass']

Answer: vectorize (so np.vectorize(grade)). Remember it is for convenience; here np.where(arr >= 50, "pass", "fail") would be the faster equivalent.

! Common Errors (And How to Fix Them)

❌ Expecting np.vectorize to speed up a big array

It loops in Python, so a million-element array is still slow:

np.vectorize(lambda x: x * 2)(big_array)   # ❌ no faster than a loop

✅ Fix: use a real array operation:

big_array * 2   # ✅ compiled, fast

❌ Doing math on a frompyfunc result without casting

np.frompyfunc returns an object array, which can behave unexpectedly in later math.

✅ Fix: cast it, e.g. result.astype(int) or result.astype(float), before numeric operations.

❌ Wrong axis in apply_along_axis

Using axis=0 when you meant rows applies the function down columns instead.

✅ Fix: axis=1 runs the function on each row, axis=0 on each column — pick deliberately.

🎯 Mini Challenge: Categorize and Then Optimize

Categorize temperatures with np.vectorize, then rewrite the same logic the fast way with np.select.

import numpy as np

temps = np.array([-5, 12, 25, 38])

# Convenience version
def category(t):
    if t < 0:
        return "freezing"
    elif t < 20:
        return "cool"
    elif t < 35:
        return "warm"
    return "hot"

vcat = np.vectorize(category)
print(vcat(temps))
# ['freezing' 'cool' 'warm' 'hot']

# Faster, fully vectorized version with np.select
conditions = [temps < 0, temps < 20, temps < 35]
choices = ["freezing", "cool", "warm"]
fast = np.select(conditions, choices, default="hot")
print(fast)
# ['freezing' 'cool' 'warm' 'hot']

print(np.array_equal(vcat(temps), fast))   # True

# ✅ Expected output:
# ['freezing' 'cool' 'warm' 'hot']
# ['freezing' 'cool' 'warm' 'hot']
# True

❓ Frequently Asked Questions

Lesson complete — convenience and speed in balance!

You can wrap functions with np.vectorize and np.frompyfunc, run logic along an axis with np.apply_along_axis, and you understand why true ufunc vectorization is the fast path.

Practice quiz

Continue this course

• Previous: Polynomials: polyfit & poly1d
• Next: Einstein Summation: einsum — Express dot products, matmuls, transposes, and traces with one notation
• Quick reference: NumPy cheat sheet

Frequently asked questions

Does np.vectorize make my code faster?

No, and this is the most important thing to know about it. np.vectorize is a convenience wrapper that still calls your Python function once per element behind the scenes, so it runs at roughly Python loop speed. Its value is cleaner syntax and automatic broadcasting of arguments, not performance. For real speed you need true vectorization with NumPy ufuncs and array operations.

When is np.vectorize actually a good choice?

Use it when a function genuinely cannot be expressed with array operations, for example one that contains complex branching, calls an external library per element, or applies arbitrary Python logic, and when the array is small enough that speed does not matter. It also gives you free broadcasting of multiple arguments, which can be tidier than writing the loop yourself.

What is np.frompyfunc and how does it differ from np.vectorize?

np.frompyfunc(func, nin, nout) also turns a Python function into something that works element-wise, but it always returns arrays of Python objects (dtype=object) rather than inferring a numeric dtype. np.vectorize is friendlier because it infers an output dtype and supports keyword options, so frompyfunc is mostly useful for advanced cases or when you specifically want object output.

What does np.apply_along_axis do?

np.apply_along_axis(func, axis, arr) applies a function that takes a 1D slice and returns a result, running it along the chosen axis of a multi-dimensional array. It is useful when your function naturally operates on whole rows or columns rather than single elements, but like vectorize it loops in Python and is not a performance tool.