如何进一步优化Wasserstein距离计算的Cython代码？

我正在将计算Wasserstein距离的Python程序转为Cython以提升速度，目前已实现初步提速，但作为Cython新手，想挖掘更多优化空间，尤其想知道如何用prange实现并行优化。输入是形状为(a, s, s)的NumPy数组，其中a是直方图数量，s是直方图的二维维度。

现有Cython代码 #

#cython: language_level=3
cimport cython
import numpy as np
from numpy cimport ndarray, float64_t, uint32_t
from libc.math cimport sqrt, fabs

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.nonecheck(False)
cdef float64_t move_supply(
        ndarray[float64_t, ndim=2] supply,
        ndarray[float64_t, ndim=2] capacity,
        (uint32_t, uint32_t) from_index,
        (uint32_t, uint32_t) to_index,
):
    cdef float64_t flow, distance
    if supply[from_index] <= capacity[to_index]:
        flow = supply[from_index]
        capacity[to_index] = capacity[to_index]-flow
        supply[from_index] = <float64_t>0.0
    else:
        flow = capacity[to_index]
        supply[from_index] = supply[from_index]-flow
        capacity[to_index] = <float64_t>0.0
    distance = sqrt(fabs(from_index[0]-to_index[0])**2 + fabs(from_index[1]-to_index[1])**2)
    return flow*distance

@cython.boundscheck(False)
@cython.wraparound(False)
cdef (uint32_t, uint32_t) find_first_candidate(ndarray[float64_t, ndim=2] array, int hist_shape):
    cdef uint32_t ie, je
    cdef (uint32_t, uint32_t) neg_return = (-1, -1)
    for ie in range(hist_shape):
        for je in range(hist_shape):
            if array[ie][je] > <float64_t>0.0:
                return ie, je
    return neg_return

@cython.boundscheck(False)
@cython.wraparound(False)
cdef float64_t wasserstein_distance(
        ndarray[float64_t, ndim=2] supply,
        ndarray[float64_t, ndim=2] capacity,
        int hist_shape,
):
    cdef float64_t score = 0.0
    cdef uint32_t i, j
    cdef (uint32_t, uint32_t) from_index, to_index
    cdef (uint32_t, uint32_t) neg = (-1, -1)


    for i in range(hist_shape):
        from_index = find_first_candidate(supply, hist_shape)
        for j in range(hist_shape):
            to_index = find_first_candidate(capacity, hist_shape)
            if from_index == neg or to_index  == neg:
                return score
            work = move_supply(supply, capacity, from_index, to_index)
            score += work
    return score

@cython.boundscheck(False)
@cython.wraparound(False)
cpdef ndarray[float64_t, ndim=2] calculate_novelty(ndarray[float64_t, ndim=3] histograms):
    cdef uint32_t i, j
    cdef uint32_t amount_instances = histograms.shape[0], histogram_size = histograms.shape[1]
    cdef ndarray[float64_t, ndim=1] novelty_scores = np.zeros(amount_instances, dtype=np.float64)

    for i in range(amount_instances-1):
        for j in range(i+1, amount_instances):
            score = wasserstein_distance(histograms[i].copy(), histograms[j].copy(), histogram_size)
            novelty_scores[i] += score
            novelty_scores[j] += score
    return novelty_scores

关键优化点 #

1. 修复核心算法逻辑缺陷 #

现有wasserstein_distance的循环次数固定为hist_shape，会导致计算不完整或冗余；find_first_candidate每次从头遍历数组找非零元素，效率极低。建议改用队列跟踪非零索引，避免重复遍历：

from collections import deque

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.nonecheck(False)
cdef float64_t wasserstein_distance(
        ndarray[float64_t, ndim=2] supply,
        ndarray[float64_t, ndim=2] capacity,
        uint32_t hist_shape,
        float64_t[:, :] dist_matrix,
):
    cdef float64_t score = 0.0
    cdef uint32_t x, y
    cdef (uint32_t, uint32_t) from_idx, to_idx
    cdef float64_t flow
    cdef deque from_queue = deque()
    cdef deque to_queue = deque()

    # 初始化队列，存储非零元素索引
    for x in range(hist_shape):
        for y in range(hist_shape):
            if supply[x][y] > 1e-9:  # 用小阈值规避浮点精度问题
                from_queue.append((x, y))
            if capacity[x][y] > 1e-9:
                to_queue.append((x, y))

    while from_queue and to_queue:
        from_idx = from_queue.popleft()
        to_idx = to_queue.popleft()

        # 计算流动量
        if supply[from_idx] <= capacity[to_idx]:
            flow = supply[from_idx]
            capacity[to_idx] -= flow
            supply[from_idx] = 0.0
        else:
            flow = capacity[to_idx]
            supply[from_idx] -= flow
            capacity[to_idx] = 0.0

        # 复用预计算的距离矩阵，避免重复计算sqrt
        score += flow * dist_matrix[from_idx[0]*hist_shape + from_idx[1]][to_idx[0]*hist_shape + to_idx[1]]

        # 剩余量放回队列
        if supply[from_idx] > 1e-9:
            from_queue.appendleft(from_idx)
        if capacity[to_idx] > 1e-9:
            to_queue.appendleft(to_idx)

    return score

2. 预计算距离矩阵 #

距离计算是重复操作，提前预计算所有点对的距离并缓存，大幅减少计算开销：

@cython.boundscheck(False)
@cython.wraparound(False)
cdef float64_t[:, :] precompute_distances(uint32_t hist_shape):
    cdef uint32_t x1, y1, x2, y2
    cdef uint32_t flat_size = hist_shape * hist_shape
    cdef float64_t[:, :] dist_matrix = np.zeros((flat_size, flat_size), dtype=np.float64)

    for x1 in range(hist_shape):
        for y1 in range(hist_shape):
            for x2 in range(hist_shape):
                for y2 in range(hist_shape):
                    dist_matrix[x1*hist_shape + y1][x2*hist_shape + y2] = sqrt(fabs(x1-x2)**2 + fabs(y1-y2)**2)
    return dist_matrix

3. 用prange实现并行优化 #

calculate_novelty中每个(i,j)对的计算完全独立，适合用OpenMP并行加速。通过局部数组暂存线程结果，避免全局竞争：

from cython.parallel cimport prange, parallel

@cython.boundscheck(False)
@cython.wraparound(False)
cpdef ndarray[float64_t, ndim=1] calculate_novelty(ndarray[float64_t, ndim=3] histograms):
    cdef uint32_t i, j, k
    cdef uint32_t amount_instances = histograms.shape[0], histogram_size = histograms.shape[1]
    cdef ndarray[float64_t, ndim=1] novelty_scores = np.zeros(amount_instances, dtype=np.float64)
    cdef float64_t[:, :] dist_matrix = precompute_distances(histogram_size)

    with nogil, parallel():
        # 每个线程维护局部分数数组，避免全局写竞争
        cdef float64_t[:] local_scores = np.zeros(amount_instances, dtype=np.float64)
        # 并行遍历外层i
        for i in prange(amount_instances-1, nogil=True):
            for j in range(i+1, amount_instances):
                cdef ndarray[float64_t, ndim=2] supply = histograms[i].copy()
                cdef ndarray[float64_t, ndim=2] capacity = histograms[j].copy()
                score = wasserstein_distance(supply, capacity, histogram_size, dist_matrix)
                local_scores[i] += score
                local_scores[j] += score
        # 合并局部结果到全局数组
        with gil:
            for k in range(amount_instances):
                novelty_scores[k] += local_scores[k]
    return novelty_scores

编译注意事项

启用OpenMP需在编译时添加参数（以setup.py为例）：

from setuptools import setup
from Cython.Build import cythonize
import numpy as np

setup(
    ext_modules=cythonize("your_module.pyx"),
    include_dirs=[np.get_include()],
    extra_compile_args=['-fopenmp'],
    extra_link_args=['-fopenmp']
)

4. 其他细节优化 #

• 统一类型：将hist_shape从int改为uint32_t，避免类型转换开销；
• 浮点精度：用>1e-9替代>0判断非零值，规避浮点误差；
• 内存视图：对NumPy数组使用内存视图（如float64_t[:]），比ndarray访问更高效；
• 减少拷贝：在Cython内部完成数组拷贝，避免Python层额外开销。

内容的提问来源于stack exchange，提问作者Oliver