How to Create Custom NumPy Functions with Cython

NumPy is an essential library for numerical computing in Python, highly optimized for performance. However, when you need even more speed, you can use Cython, a superset of the Python language that enables C-level performance boosts. In this tutorial, we’ll walk through how to create custom NumPy functions with Cython, starting from the basics and working up to more advanced concepts.

Introduction to Cython

Cython is a programming language that makes writing C extensions for the Python language as easy as Python itself. Cython is a compiled language that produces very efficient C code, which can then be used by Python programs. To make the most of Cython, you should have a basic understanding of both Python and C.

Getting Started with Cython

Before you can write your first Cython program, you need to set up your environment:

• Ensure you have Python installed on your system.
• Install Cython using pip install Cython.
• Additionally, it’s good to have a C compiler available on your system.

Simple Cython Function

Here’s a simple example of how to write a Cython function:

cdef int add(int a, int b):
    return a + b

To use this function from Python code, you would declare it in a .pyx file and then use a setup.py script to compile it.

from setuptools import setup
from Cython.Build import cythonize
setup(ext_modules = cythonize('example.pyx'))

Enhancing NumPy Operations with Cython

To use NumPy with Cython, you can directly import it and type your numpy arrays:

cimport numpy as cnp
def operate_on_array(cnp.ndarray[cnp.float64_t, ndim=1] arr):
    cdef Py_ssize_t i
    for i in range(arr.size):
        arr[i] *= 2  # Will multiply each element by 2
    return arr

Compiling and Running the Cython Code

Once you’ve written your Cython file, you’ll need to compile it into a shared library that Python can import:

setup(ext_modules = cythonize('my_numpy_extension.pyx'))

After compiling, you can import and use your Cython function in Python:

import my_numpy_extension
import numpy as np

arr = np.array([1.0, 2.0, 3.0])
my_numpy_extension.operate_on_array(arr)
print(arr)

Working with Multiple Dimensions

NumPy arrays aren’t limited to a single dimension. Here’s how you would write and compile a Cython function that operates on a two-dimensional array:

cdef operate_on_2d_array(cnp.ndarray[cnp.float64_t, ndim=2] arr):
    cdef Py_ssize_t i, j
    for i in range(arr.shape[0]):
        for j in range(arr.shape[1]):
            arr[i, j] += 10  # Adds 10 to each element
    return arr

After compilation, using the function looks very similar to the one-dimensional case.

Improving Performance with Typed Memoryviews

Typed memoryviews are one of Cython’s mechanisms for accessing memory buffers such as arrays and buffer-like objects without incurring Python overhead:

def operate_on_array_with_memoryview(double[:] arr):
    cdef Py_ssize_t i
    for i in range(arr.shape[0]):
        arr[i] *= 2
    return arr

Typed memoryviews are more flexible and efficient, they often result in faster code compared to using the ndarray directly.

Optimizing with Cython Directives

Cython provides several directives that can be used to optimize your code. For example, you can enable/disable bounds checking and wraparound, which can speed up array access:

from cython cimport boundscheck, wraparound

@boundscheck(False)
@wraparound(False)
cdef operate_on_array_optimized(double[:] arr):
    ...

Parallelization

Cython also has support for parallelizing code using OpenMP, a parallel programming interface. This can be very effective when working with large data sets that can be divided across multiple CPU cores:

from cython.parallel import prange

def parallel_operation(double[:] arr):
    cdef Py_ssize_t i
    for i in prange(len(arr), nogil=True):
        arr[i] *= 2

Do not forget to tell the compiler to enable OpenMP in your setup.py file:

setup(
    ext_modules=cythonize('my_parallel_extension.pyx', compiler_directives={'language_level': 3}),
    extra_compile_args=['-fopenmp'],
    extra_link_args=['-fopenmp']
)

Conclusion

By now, you should have a solid foundation for optimizing your NumPy operations using Cython. Experimenting with these examples and directives will help you to harness the full potential of your code’s performance.