NumPy: Parallelization with Numba

NumPy’s vectorized operations are highly efficient, but for computationally intensive tasks, further performance gains can be achieved by parallelizing code with Numba. Numba is a JIT (Just-In-Time) compiler that optimizes Python code, including NumPy arrays, by compiling it to machine code and enabling parallel execution. This tutorial explores NumPy parallelization with Numba, covering its integration, key techniques, and practical approaches for accelerating numerical computations in scientific computing, machine learning, and data analysis.

---

01. What Is Numba and Its Role with NumPy?

Numba is a Python library that accelerates numerical computations by compiling Python functions to optimized machine code using LLVM. When combined with NumPy, Numba enhances performance by optimizing loops, enabling parallel execution, and reducing Python overhead. Built on NumPy Array Operations, Numba’s JIT compilation and parallelization capabilities make it ideal for tasks like matrix operations, simulations, and large-scale data processing that benefit from multi-core CPU utilization.

Example: Basic Numba JIT with NumPy

import numpy as np
from numba import jit

# Define a function with Numba JIT
@jit(nopython=True)
def matrix_sum(a, b):
    return a + b

# Create NumPy arrays
a = np.random.rand(1000, 1000)
b = np.random.rand(1000, 1000)
# Call JIT-compiled function
result = matrix_sum(a, b)
print("First element:", result[0, 0])

Output:

First element: [sum value]

Explanation:

• @jit(nopython=True) - Compiles the function to machine code, ensuring no Python runtime overhead.
• NumPy arrays are passed directly, leveraging Numba’s compatibility with NumPy.

---

02. Parallelization Techniques with Numba and NumPy

Numba provides decorators and features like @jit, @njit, and prange to parallelize NumPy computations across multiple CPU cores. These techniques are particularly effective for loop-heavy operations that NumPy’s vectorization alone cannot fully optimize. The table below summarizes key Numba parallelization techniques for NumPy:

Technique	Description	Example
@jit	Compiles function to machine code	@jit(nopython=True)
@njit	Shortcut for @jit(nopython=True)	@njit
prange	Parallelizes loops across CPU cores	from numba import prange
Explicit Parallelization	Enable parallel execution	@jit(parallel=True)
Vectorization with @vectorize	Create NumPy ufuncs	@vectorize

2.1 Basic JIT Compilation

Example: Optimizing a Loop with Numba

import numpy as np
from numba import jit

# JIT-compiled function
@jit(nopython=True)
def element_wise_product(a, b, result):
    for i in range(a.shape[0]):
        for j in range(a.shape[1]):
            result[i, j] = a[i, j] * b[i, j]

# Create NumPy arrays
a = np.random.rand(1000, 1000)
b = np.random.rand(1000, 1000)
result = np.zeros((1000, 1000))
# Call function
element_wise_product(a, b, result)
print("First element:", result[0, 0])

Output:

First element: [product value]

Explanation:

• Numba compiles the nested loop to machine code, making it much faster than a Python loop.
• Compare to NumPy’s vectorized a * b, which is simpler but less flexible for custom logic.

2.2 Parallel Loops with prange

Example: Parallel Matrix Addition

import numpy as np
from numba import jit, prange

# Parallel JIT-compiled function
@jit(nopython=True, parallel=True)
def parallel_add(a, b, result):
    for i in prange(a.shape[0]):
        for j in range(a.shape[1]):
            result[i, j] = a[i, j] + b[i, j]

# Create NumPy arrays
a = np.random.rand(1000, 1000)
b = np.random.rand(1000, 1000)
result = np.zeros((1000, 1000))
# Call function
parallel_add(a, b, result)
print("First element:", result[0, 0])

Output:

First element: [sum value]

Explanation:

• prange - Parallelizes the outer loop across CPU cores.
• parallel=True - Enables Numba’s parallel execution mode.

2.3 Vectorization with @vectorize

Example: Custom Vectorized Function

import numpy as np
from numba import vectorize

# Define vectorized function
@vectorize(['float64(float64, float64)'])
def custom_op(x, y):
    return x * y + x

# Create NumPy arrays
a = np.random.rand(10000)
b = np.random.rand(10000)
# Apply vectorized function
result = custom_op(a, b)
print("First element:", result[0])

Output:

First element: [computed value]

Explanation:

• @vectorize - Creates a NumPy ufunc, applying the function element-wise with Numba optimization.

2.4 Parallel Matrix Multiplication

Example: Parallel Matrix Multiplication

import numpy as np
from numba import jit, prange

# Parallel matrix multiplication
@jit(nopython=True, parallel=True)
def parallel_matmul(a, b, result):
    m, n = a.shape
    n, p = b.shape
    for i in prange(m):
        for j in range(p):
            for k in range(n):
                result[i, j] += a[i, k] * b[k, j]

# Create NumPy arrays
a = np.random.rand(500, 300)
b = np.random.rand(300, 400)
result = np.zeros((500, 400))
# Call function
parallel_matmul(a, b, result)
print("First element:", result[0, 0])

Output:

First element: [product value]

Explanation:

• Parallelizes the outer loop with prange, distributing work across cores.
• Note: For standard matrix multiplication, np.dot with BLAS is often faster, but Numba is useful for custom logic.

2.5 Incorrect Numba Usage

Example: Non-Numba-Compatible Code

import numpy as np
from numba import jit

# Incorrect: Using Python objects
@jit(nopython=True)
def bad_function(a):
    return [x * 2 for x in a]  # List comprehension not supported

# Create NumPy array
a = np.array([1, 2, 3])
# Call function
bad_function(a)  # CompilationError

Output:

TypingError: list comprehension not supported in nopython mode

Explanation:

• Numba’s nopython=True mode doesn’t support Python objects like lists; use NumPy arrays and loops instead.

---

03. Effective Usage

3.1 Recommended Practices

• Use @njit and prange for parallelizing loop-heavy NumPy operations.

Example: Parallel Sum of Squares

import numpy as np
from numba import njit, prange

# Parallel sum of squares
@njit(parallel=True)
def sum_squares(a, result):
    for i in prange(a.shape[0]):
        result[i] = a[i] ** 2

# Create NumPy array
a = np.random.rand(10000)
result = np.zeros(10000)
# Call function
sum_squares(a, result)
print("First element:", result[0])

Output:

First element: [squared value]

• Combine Numba with NumPy’s vectorized operations for hybrid optimization.
• Use @vectorize for simple element-wise operations requiring custom logic.

3.2 Practices to Avoid

• Avoid using Numba for operations already optimized by NumPy/BLAS (e.g., np.dot).

Example: Redundant Numba for Matrix Multiplication

import numpy as np
from numba import jit

# Inefficient: Numba for simple matrix multiplication
@jit(nopython=True)
def redundant_matmul(a, b):
    return a @ b  # Better to use np.matmul directly

# Create NumPy arrays
a = np.random.rand(100, 100)
b = np.random.rand(100, 100)
# Call function
result = redundant_matmul(a, b)
print("Result shape:", result.shape)

Output:

Result shape: (100, 100)

• Use NumPy’s BLAS-backed np.matmul instead of Numba for standard matrix operations.

---

04. Common Use Cases

4.1 Scientific Simulations

Numba parallelizes custom numerical simulations with NumPy arrays.

Example: Parallel Particle Simulation

import numpy as np
from numba import njit, prange

# Parallel distance computation
@njit(parallel=True)
def particle_distances(positions, distances):
    n = positions.shape[0]
    for i in prange(n):
        for j in range(i + 1, n):
            distances[i, j] = np.sqrt(np.sum((positions[i] - positions[j])**2))
            distances[j, i] = distances[i, j]

# Create NumPy arrays
positions = np.random.rand(1000, 3)
distances = np.zeros((1000, 1000))
# Call function
particle_distances(positions, distances)
print("First distance:", distances[0, 1])

Output:

First distance: [distance value]

4.2 Machine Learning

Numba accelerates custom algorithms, like distance metrics in clustering.

Example: Parallel K-Means Distance

import numpy as np
from numba import njit, prange

# Parallel distance to centroids
@njit(parallel=True)
def compute_distances(points, centroids, distances):
    n, d = points.shape
    k = centroids.shape[0]
    for i in prange(n):
        for j in range(k):
            distances[i, j] = np.sqrt(np.sum((points[i] - centroids[j])**2))

# Create NumPy arrays
points = np.random.rand(10000, 2)
centroids = np.random.rand(5, 2)
distances = np.zeros((10000, 5))
# Call function
compute_distances(points, centroids, distances)
print("First distance:", distances[0, 0])

Output:

First distance: [distance value]

---

Conclusion

Numba enhances NumPy’s performance by compiling and parallelizing numerical computations, particularly for custom algorithms and loop-heavy tasks. Using decorators like @njit, prange, and @vectorize, you can leverage multi-core CPUs for significant speedups. Key takeaways:

• Use Numba for custom computations not fully optimized by NumPy.
• Parallelize loops with prange for multi-core performance.
• Avoid Numba for operations already optimized by NumPy/BLAS.
• Apply these techniques in scientific simulations and machine learning.

With these strategies, you’re equipped to accelerate your NumPy Array Operations using Numba’s parallelization capabilities!