This post/notebook is the follow-up to a recent one : Heapsort with Numba and Cython, where we implemented heapsort in Python/Numba and Cython and compared the execution time with NumPy heapsort. However, heapsort in NumPy is written in C++ [wrapped in Python] and not exactly the same implementation as the one we used in the previous post. So in the following, we present two different heapsort implementations translated into Cython, which are a little bit more refined than the previous one:

  • the first one [Cython_1] is taken from the book Numerical recipes in C [1]
  • the second one [Cython_2] is taken from NumPy source code: heapsort.cpp

In both implementations, indices are 1-based and left bitwise shift [<<] and right bitwise shift [>>] operators are used. In order to use 1-based indices with NumPy arrays, we add an extra element at the beginning:

def heapsort(A_in):
    A = np.copy(A_in)
    A = np.concatenate((np.array([np.nan]), A))
    _heapsort(A)
    return A[1:]

Also, the algorithms are implemented in a single function, and without recursion.

Imports

import gc
from time import perf_counter
import cython
import numpy as np
import pandas as pd
import perfplot
%load_ext cython
SD = 124  # random seed
rng = np.random.default_rng(SD)  # random number generator

Float array creation

We still assume that we want to sort an NumPy array of 64-bit floating-point numbers between 0 and 1:

n = 10
A = rng.random(n, dtype=np.float64)
A
array([0.78525311, 0.78585936, 0.96913602, 0.74805977, 0.65555081,
       0.93888454, 0.17861445, 0.58864721, 0.44279917, 0.34884712])

Cython_1

The authors of Numerical recipes in C [1] use a "corporate hierarchy metaphor":

[...] a heap has every “supervisor” greater than or equal to its two “supervisees,” down through the levels of the hierarchy.

If you have managed to rearrange your array into an order that forms a heap, then sorting it is very easy: You pull off the “top of the heap,” which will be the largest element yet unsorted. Then you “promote” to the top of the heap its largest underling. Then you promote its largest underling, and so on. The process is like what happens [or is supposed to happen] in a large corporation when the chairman of the board retires. You then repeat the whole process by retiring the new chairman of the board.

Imagine that the corporation starts out with N/2 employees on the production line, but with no supervisors. Now a supervisor is hired to supervise two workers. If he is less capable than one of his workers, that one is promoted in his place, and he joins the production line. After supervisors are hired, then supervisors of supervisors are hired, and so on up the corporate ladder. Each employee is brought in at the top of the tree, but then immediately sifted down, with more capable workers promoted until their proper corporate level has been reached.

One execution of the Heapsort function represents the entire life-cycle of a giant corporation: N/2 workers are hired; N/2 potential supervisors are hired; there is a sifting up in the ranks, a sort of super Peter Principle: in due course, each of the original employees gets promoted to chairman of the board.

So we have two phases:

  • hiring
  • retirement-and-promotion

This is very similar to what we described in the previous post: a first step where we build the heap in a bottom-up approach, a second step where we keep on extracting the root until the heap is empty. Both steps rely on a sift-down process. However, we have here a single main loop for both phases.

%%cython -3 --compile-args=-Ofast
# cython: boundscheck=False, cdivision=True, initializedcheck=False, wraparound=False
import cython
import numpy as np
from cython import double, ssize_t

@cython.exceptval(check=False)
@cython.nogil
@cython.cfunc
def _heapsort_1(ra: double[::1]) -> cython.void:
    i: ssize_t
    ir: ssize_t
    j: ssize_t
    l: ssize_t
    n: ssize_t = ra.shape[0] - 1
    rra: double
    if n < 2:
        return
    l = (n >> 1) + 1
    ir = n
    # The index l will be decremented from its initial value down to 1 during the “hiring” (heap
    # creation) phase. Once it reaches 1, the index ir will be decremented from its initial value
    # down to 1 during the “retirement-and-promotion” (heap selection) phase.
    while 1:
        if l > 1:  # Still in hiring phase.
            l -= 1
            rra = ra[l]
        else:  # In retirement-and-promotion phase.
            rra = ra[ir]  # Clear a space at end of array.
            ra[ir] = ra[1]  # Retire the top of the heap into it.
            ir -= 1
            if ir == 1:  # Done with the last promotion
                ra[1] = rra  # The least competent worker of all!
                break
        i = l  # Whether in the hiring phase or promotion phase, we
        j = l + l  # here set up to sift down element rra to its proper level.
        while j <= ir:
            if (j < ir) & (ra[j] < ra[j + 1]):  # Compare to the better underling.
                j += 1
            if rra < ra[j]:  # Demote rra.
                ra[i] = ra[j]
                i = j
                j = j << 1
            else:
                # j = ir + 1  # This is rra's level. Set j to terminate the sift-down.
                break
        ra[i] = rra  # Put rra into its slot.

@cython.ccall
def heapsort_cython_1(A_in):
    A = np.copy(A_in)
    A = np.concatenate((np.array([np.nan]), A))
    _heapsort_1(A)
    return A[1:]
A_1 = heapsort_cython_1(A)
A_1
array([0.17861445, 0.34884712, 0.44279917, 0.58864721, 0.65555081,
       0.74805977, 0.78525311, 0.78585936, 0.93888454, 0.96913602])
np.testing.assert_array_equal(A_1, np.sort(A))

Cython_2

This is a straight translation of NumPy Source code from C++ to Cython. Here is what we can read in the source code about its origin:

 * These sorting functions are copied almost directly from numarray
 * with a few modifications (complex comparisons compare the imaginary
 * part if the real parts are equal, for example), and the names
 * are changed.
 *
 * The original sorting code is due to Charles R. Harris who wrote
 * it for numarray.

This time we have two distinct loops in the _heapsort_2() function. But this is very similar to the code above [Cython_1].

%%cython -3 --compile-args=-Ofast
# cython: boundscheck=False, cdivision=True, initializedcheck=False, wraparound=False
import cython
import numpy as np
from cython import double, ssize_t

@cython.exceptval(check=False)
@cython.nogil
@cython.cfunc
def _heapsort_2(a: double[::1]) -> cython.void:
    i: ssize_t
    j: ssize_t
    l: ssize_t
    n: ssize_t = a.shape[0] - 1
    tmp: double
    for l in range(n >> 1, 0, -1):
        tmp = a[l]
        i = l
        j = l << 1
        while j <= n:
            if (j < n) & (a[j] < a[j + 1]):
                j += 1
            if tmp < a[j]:
                a[i] = a[j]
                i = j
                j += j
            else:
                break
        a[i] = tmp
    while n > 1:
        tmp = a[n]
        a[n] = a[1]
        n -= 1
        i = 1
        j = 2
        while j <= n:
            if (j < n) & (a[j] < a[j + 1]):
                j += 1
            if tmp < a[j]:
                a[i] = a[j]
                i = j
                j += j
            else:
                break
        a[i] = tmp

@cython.ccall
def heapsort_cython_2(A_in):
    A = np.copy(A_in)
    A = np.concatenate((np.array([np.nan]), A))
    _heapsort_2(A)
    return A[1:]
A_2 = heapsort_cython_2(A)
np.testing.assert_array_equal(A_2, np.sort(A))

Cython_3

For the sake of completeness, we also included the implementation from the previous post, with the same compiler directives as above:

%%cython -3 --compile-args=-Ofast
# cython: boundscheck=False, cdivision=True, initializedcheck=False, wraparound=False
import cython
import numpy as np
from cython import double, ssize_t

@cython.exceptval(check=False)
@cython.nogil
@cython.cfunc
def _max_heapify(A: double[::1], size: ssize_t, node_idx: ssize_t) -> cython.void:
    largest: ssize_t = node_idx
    left_child = 2 * node_idx + 1
    right_child = 2 * (node_idx + 1)
    if left_child < size and A[left_child] > A[largest]:
        largest = left_child
    if right_child < size and A[right_child] > A[largest]:
        largest = right_child
    if largest != node_idx:
        A[node_idx], A[largest] = A[largest], A[node_idx]
        _max_heapify(A, size, largest)

@cython.exceptval(check=False)
@cython.nogil
@cython.cfunc
def _heapsort_3(A: double[::1]) -> cython.void:
    i: cython.int
    size: ssize_t = len(A)
    node_idx: ssize_t = size // 2 - 1
    for i in range(node_idx, -1, -1):
        _max_heapify(A, size, i)
    for i in range(size - 1, 0, -1):
        A[i], A[0] = A[0], A[i]
        size -= 1
        _max_heapify(A, size, 0)

@cython.ccall
def heapsort_cython_3(A_in):
    A = np.copy(A_in)
    _heapsort_3(A)
    return A
A_3 = heapsort_cython_3(A)
np.testing.assert_array_equal(A_3, np.sort(A))

Performance comparison

Again, we compare the different versions against NumPy heapsort.

Perfplot

out = perfplot.bench(
    setup=lambda n: rng.random(n, dtype=np.float64),
    kernels=[
        lambda A: heapsort_cython_1(A),
        lambda A: heapsort_cython_2(A),
        lambda A: heapsort_cython_3(A),
        lambda A: np.sort(A, kind="heapsort"),
    ],
    labels=["Cython_1", "Cython_2", "Cython_3", "NumPy"],
    n_range=[10**k for k in range(1, 9)],
)
t_df = pd.DataFrame(out.timings_s.T, columns=out.labels, index=out.n_range)
t_df
Cython_1Cython_2Cython_3NumPy
103.915e-063.908e-061.663e-061.737e-06
1005.111e-064.890e-065.054e-062.977e-06
10005.326e-055.319e-057.090e-056.266e-05
100008.942e-049.227e-048.871e-049.634e-04
1000001.202e-021.200e-021.272e-021.315e-02
10000001.779e-011.735e-012.202e-011.973e-01
100000002.666e+002.654e+004.641e+003.066e+00
1000000003.740e+013.695e+017.699e+014.161e+01

So is seems that the Cython_1 and Cython_2 implementations achieve similar speed as NumPy's version. However, the Cython_3 implementation is clearly slower when the array gets larger.

A Larger array

It is possible to run the algorithms on arrays with 1 billion elements, but without using perfplot for reasons related to memory capacity of the computer used to run the tests.

%%time
k = 9
A = rng.random(10**k, dtype=np.float64)
CPU times: user 4.05 s, sys: 679 ms, total: 4.72 s
Wall time: 4.72 s
timings = []
for _ in range(3):
    start = perf_counter()
    A_1 = heapsort_cython_1(A)
    end = perf_counter()
    elapsed_time = end - start
    del A_1
    _ = gc.collect()
    timings.append(elapsed_time)
elapsed_time_1 = np.amin(timings)
timings = []
for _ in range(3):
    start = perf_counter()
    A_2 = heapsort_cython_2(A)
    end = perf_counter()
    elapsed_time = end - start
    del A_2
    _ = gc.collect()
    timings.append(elapsed_time)
elapsed_time_2 = np.amin(timings)
timings = []
for _ in range(3):
    start = perf_counter()
    A_3 = heapsort_cython_3(A)
    end = perf_counter()
    elapsed_time = end - start
    del A_3
    _ = gc.collect()
    timings.append(elapsed_time)
elapsed_time_3 = np.amin(timings)
timings = []
for _ in range(3):
    start = perf_counter()
    A_4 = np.sort(A, kind="heapsort")
    end = perf_counter()
    elapsed_time = end - start
    del A_4
    _ = gc.collect()
    timings.append(elapsed_time)
elapsed_time_4 = np.amin(timings)
t_df = pd.concat(
    [
        t_df,
        pd.DataFrame(
            [[elapsed_time_1, elapsed_time_2, elapsed_time_3, elapsed_time_4]],
            columns=out.labels,
            index=[10**k],
        ),
    ],
    axis=0,
)
t_df["n log(n)"] = 5.0e-8 * t_df.index * np.log(t_df.index)

Note that we also included the curve corresponding to $O(n\log n)$ in the following figure. However, we are aware that this is an asymptotic behavior on an abstract machine, but do not take into account some hardware effects, such has cache optimization.

t_df[['Cython_1', 'Cython_2', 'Cython_3', 'NumPy']]
Cython_1Cython_2Cython_3NumPy
103.915e-063.908e-061.663e-061.737e-06
1005.111e-064.890e-065.054e-062.977e-06
10005.326e-055.319e-057.090e-056.266e-05
100008.942e-049.227e-048.871e-049.634e-04
1000001.202e-021.200e-021.272e-021.315e-02
10000001.779e-011.735e-012.202e-011.973e-01
100000002.666e+002.654e+004.641e+003.066e+00
1000000003.740e+013.695e+017.699e+014.161e+01
10000000004.547e+024.536e+021.059e+034.142e+02
ax = t_df.plot(loglog=True, figsize=(12, 12), ms=10, alpha=0.6)
markers = ("o", "v", "^", "<", ">", "s", "p", "P", "*", "h", "X", "D", ".")
for i, line in enumerate(ax.get_lines()):
    marker = markers[i]
    line.set_marker(marker)
plt.legend()
plt.grid("on")
_ = ax.set(
    title="Timing comparison between Cython and NumPy heapsort",
    xlabel="Array length (log scale)",
    ylabel="Elapsed_time [s] (log scale)",
)

Timings

Reference

[1] William H. Press, Saul A. Teukolsky, William T. Vetterling, and Brian P. Flannery. 1992. Numerical recipes in C [2nd ed.]: the art of scientific computing. Cambridge University Press, USA.