How to Create Custom NumPy Ufuncs in Python

Universal functions (ufuncs) in NumPy operate element-by-element on ndarray objects, providing native support for broadcasting, type casting, and methods like reduction and accumulation. This guide outlines the standard procedures for authoring custom vectorized ufuncs, spanning pure Python prototyping with numpy.frompyfunc, high-performance JIT compilation via Numba, and native extension development using the NumPy C API.

Method 1: Rapid Prototyping with numpy.frompyfunc

The standard library method for turning a standard Python function into a ufunc is numpy.frompyfunc. It wraps an arbitrary Python callable so that it supports NumPy broadcasting and array manipulation.

Step-by-Step Procedure

  1. Define a standard scalar function: Write a function that accepts individual elements as arguments and returns the computed result.
  2. Apply numpy.frompyfunc: Specify the function, the number of input arguments (nin), and the number of returned outputs (nout).
import numpy as np

# 1. Define scalar operation
def custom_add_subtract(x, y):
    return x + y, x - y

# 2. Convert to ufunc: 2 inputs, 2 outputs
my_ufunc = np.frompyfunc(custom_add_subtract, 2, 2)

# Usage with broadcasting
a = np.array([10, 20, 30])
b = np.array([1, 2, 3])
sum_res, diff_res = my_ufunc(a, b)

print(sum_res)   # array([11, 22, 33], dtype=object)
print(diff_res)  # array([9, 18, 27], dtype=object)

Note: numpy.frompyfunc always returns arrays of object dtype. If you require typed outputs, use numpy.vectorize with the otypes parameter, though neither approach bypasses Python interpreter overhead.

Method 2: High-Performance Ufuncs with Numba

To achieve C-like execution speeds without writing C extensions, use Numba’s @vectorize decorator. This compiles the Python code into machine code and creates a genuine NumPy ufunc that avoids Python GIL overhead.

Step-by-Step Procedure

  1. Define the scalar operation.
  2. Apply the @vectorize decorator: Pass the desired type signatures (e.g., 'float64(float64, float64)') and target architecture.
from numba import vectorize, float64
import numpy as np

# Define signature: return_type(input1_type, input2_type)
@vectorize([float64(float64, float64)], target='parallel')
def fast_power_diff(x, y):
    return (x ** 2) - (y ** 2)

a = np.linspace(0, 10, 1000000)
b = np.linspace(10, 20, 1000000)

result = fast_power_diff(a, b)

Functions created with Numba inherit full ufunc features, including methods such as .reduce():

# Sum-reduction over the array using the custom ufunc
total = fast_power_diff.reduce(a)

Method 3: Native C API Implementation

For absolute control and direct inclusion into compiled extensions, define the ufunc using NumPy's C API with PyUFunc_FromFuncAndData.

Step-by-Step Procedure

  1. Write the 1D loop function: Define a function matching the signature void loop(char **args, const npy_intp *dimensions, const npy_intp *steps, void *data).
  2. Define type arrays: Declare an array of char specifying input and output types using NumPy type characters (such as NPY_DOUBLE).
  3. Register the ufunc: Call PyUFunc_FromFuncAndData in the module initialization code.
#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
#include <Python.h>
#include <numpy/arrayobject.h>
#include <numpy/ufuncobject.h>
#include <math.h>

/* 1. The 1D loop */
static void double_logit(char **args, const npy_intp *dimensions,
                         const npy_intp *steps, void *data) {
    npy_intp n = dimensions[0];
    char *in = args[0];
    char *out = args[1];
    npy_intp in_step = steps[0];
    npy_intp out_step = steps[1];

    for (npy_intp i = 0; i < n; i++) {
        double val = *(double *)in;
        *(double *)out = log(val / (1.0 - val));
        in += in_step;
        out += out_step;
    }
}

/* 2. Type information */
static PyUFuncGenericFunction funcs[1] = {&double_logit};
static char types[2] = {NPY_DOUBLE, NPY_DOUBLE};
static void *data[1] = {NULL};

/* 3. Module Initialization */
static PyMethodDef LogitMethods[] = {
    {NULL, NULL, 0, NULL}
};

static struct PyModuleDef moduledef = {
    PyModuleDef_HEAD_INIT,
    "custom_ufuncs",
    NULL,
    -1,
    LogitMethods
};

PyMODINIT_FUNC PyInit_custom_ufuncs(void) {
    PyObject *m = PyModule_Create(&moduledef);
    if (!m) return NULL;

    import_array();
    import_ufunc();

    PyObject *logit_ufunc = PyUFunc_FromFuncAndData(
        funcs, data, types, 
        1, /* number of supported type variants */
        1, /* number of inputs */
        1, /* number of outputs */
        PyUFunc_None, "logit", "Computes the logit function", 0
    );

    PyModule_AddObject(m, "logit", logit_ufunc);
    return m;
}

Compile this file using setuptools with numpy.get_include() added to the extension's include directories to import and use the compiled ufunc in Python.