Cython不会显着提高速度:如何处理Cython中的NumPy广播?

时间:2019-07-05 19:32:48

标签: python numpy cython

我是Cython的初学者。我正在尝试加快一个被多次调用的函数的速度,而大幅提高速度将非常有帮助。

此功能的原始版本大量使用了大型多维数组上的NumPy广播。当尝试Cythonize函数时,我最初尝试将这些数组转换为Cython memoryviews。但是,由于函数对这些数组执行算术运算,因此我遇到了编译错误,此后我了解到内存视图不支持此操作。

因此,我修改了Cythonized代码,以将数组声明为NumPy ndarrays,而不是memoryviews。 Cythonized函数现在可以使用,但是与原始的纯Python / NumPy版本相比,速度上没有任何明显的差异。

因此,我现在可以:a。)接受这段代码不适合Cython化,然后在其他地方寻找提高速度的方法,或者b。)恢复为将大数组作为内存视图来处理,以某种方式克服了需要执行的算法。

继续执行选项b。),我希望解决此函数的for循环中的算法。我可以潜在地编写其他Cython函数,这些函数在逐个元素的基础上执行数组乘法和加法,据我所知,这样我便可以使用memoryviews。但是,鉴于此代码的NumPy广播很复杂,我希望这可能需要付出很大的努力(而且我不一定知道如何开始...)。此外,我不确定这种努力是否会取得成果,因为在逐个元素的基础上做事(诚然在C语言中)可能实际上并不比广播NumPy操作快。

我非常欢迎任何建议或支持。谢谢。

from numpy import pi, exp, sin, cos
import numpy as np
cimport numpy as np
cimport cython

cdef np.ndarray[np.double_t, ndim=3] foo(bar, double dt, f, xi):
    cdef int nt, nh, nj
    cdef Py_ssize_t t

    nt = bar.shape[0]
    nh = bar.shape[1]

    if len(bar.shape) < 3:
        bar = bar[:, np.newaxis, :]
    cdef np.ndarray[np.double_t, ndim=3] bar_c = bar

    nj = len(f)
    k = (2 * pi * f) ** 2
    wn = k ** 0.5
    wd = (wn * (1 - xi ** 2) ** 0.5)

    cdef np.ndarray[np.double_t, ndim=3] u = np.zeros((nt, nj, nh))
    cdef np.ndarray[np.double_t, ndim=3] v = np.zeros((nt, nj, nh))

    C1 = exp(-xi * wn * dt)
    C2 = sin(wd * dt)
    C3 = cos(wd * dt)
    cdef np.ndarray[np.double_t, ndim=2] A11 = C1 * (C3 + (xi / (1 - xi ** 2) ** 0.5) * C2)
    cdef np.ndarray[np.double_t, ndim=2] A12 = (C1 / wd) * C2
    cdef np.ndarray[np.double_t, ndim=2] A21 = (-wn / (1 - xi ** 2) ** 0.5) * C1 * C2
    cdef np.ndarray[np.double_t, ndim=2] A22 = C1 * (C3 - (xi / (1 - xi ** 2) ** 0.5) * C2)
    cdef np.ndarray[np.double_t, ndim=2] B11 = C1 * (
        (((2 * xi ** 2 - 1) / (dt * wn ** 2)) + xi / wn) * C2 / wd
        + ((2 * xi / (dt * wn ** 3)) + (1 / wn ** 2)) * C3
    ) - 2 * xi / (dt * wn ** 3)
    cdef np.ndarray[np.double_t, ndim=2] B12 = (
        -C1
        * (
            ((2 * xi ** 2 - 1) / (dt * wn ** 2)) * C2 / wd
            + ((2 * xi) / (dt * wn ** 3)) * C3
        )
        - (1 / wn ** 2)
        + 2 * xi / (dt * wn ** 3)
    )
    cdef np.ndarray[np.double_t, ndim=2] B21 = -A12 - ((A11 - 1) / (dt * wn ** 2))
    cdef np.ndarray[np.double_t, ndim=2] B22 = -B21 - A12
    for t in range(0, nt - 1):
        u[t + 1, :, :] = (
            A11 * u[t, :, :]
            + A12 * v[t, :, :]
            + B11 * bar_c[t, :, :]
            + B12 * bar_c[t + 1, :, :]
        )
        v[t + 1, :, :] = (
            A21 * u[t, :, :]
            + A22 * v[t, :, :]
            + B21 * bar_c[t, :, :]
            + B22 * bar_c[t + 1, :, :]
        )

    cdef np.ndarray[np.double_t, ndim=3] out = -2 * xi * wn * v - (wn ** 2) * u - bar_c
    return out

不减少代码大小的歉意。考虑到查询的性质,我很难确定一个最小的,可重复的示例。

1 个答案:

答案 0 :(得分:1)

感谢@CodeSurgeon和@ 9000给出的答案。您已经确认/建议,在某些情况下,C语言确实可以比NumPy操作提高速度-即使考虑到Cython内存视图要求C操作必须在逐个元素的基础上执行的事实,而我的原始代码是使用NumPy操作在大型阵列上广播。

这促使我进一步探索。现在,在内存视图上逐个元素地运行,代码速度提高了3到40倍(强烈取决于输入数组的大小,后者会有所不同)。

修改基本上是这样的:

    # Memoryview declarations
    cdef double[:] u_mv = u
    cdef double[:] v_mv = v
    cdef double[:, :, :] out_mv = out
    cdef double[:, :, :] bar_mv = bar
    for j in range(nj):

        ...same calculation of constants C1, C2 etc. as before...

        for h in range(nh):
            for t in range(0, nt - 1):
                u_mv[t + 1] = (
                    A11 * u_mv[t]
                    + A12 * v_mv[t]
                    + B11 * bar_mv[t, 0, h]
                    + B12 * bar_mv[t + 1, 0, h]
                )
                v_mv[t + 1] = (
                    A21 * u_mv[t]
                    + A22 * v_mv[t]
                    + B21 * bar_mv[t, 0, h]
                    + B22 * bar_mv[t + 1, 0, h]
                )
                out_mv[t + 1, j, h] = (
                    -2 * xi_j * wn * v_mv[t + 1] - (wn ** 2) * u_mv[t + 1]
                    - bar_mv[t + 1, 0, h]
                )
            out_mv[0, j, h] = -bar_mv[0, 0, h]
    return out

再次感谢。