代码之家 › 专栏 › 技术社区 › Physicist

在cython中遍历指定轴

cython iterator pointers numpy python

Physicist · 技术社区 · 7 年前

tutorial 尝试进行数值微分:

import numpy as np
cimport numpy as np
import cython
np.import_array()

def test3(a, int order=2, int axis=-1):

    cdef int i

    if axis<0:
        axis = len(a.shape) + axis

    out = np.empty(a.shape, np.double)

    cdef np.flatiter ita = np.PyArray_IterAllButAxis(a, &axis)
    cdef np.flatiter ito = np.PyArray_IterAllButAxis(out, &axis)

    cdef int a_axis_stride = a.strides[axis]
    cdef int o_axis_stride = out.strides[axis]

    cdef int axis_length = out.shape[axis]

    cdef double value


    while np.PyArray_ITER_NOTDONE(ita):
        # first element
        pt1 = <double*>((<char*>np.PyArray_ITER_DATA(ita)))
        pt2 = (<double*>((<char*>np.PyArray_ITER_DATA(ita)) + 1*a_axis_stride))
        pt3 = (<double*>((<char*>np.PyArray_ITER_DATA(ita)) + 2*a_axis_stride))
        value = -1.5*pt1[0] + 2*pt2[0] - 0.5*pt3[0]
        (<double*>((<char*>np.PyArray_ITER_DATA(ito))))[0] = value

        for i in range(axis_length-2):
            pt1 = (<double*>((<char*>np.PyArray_ITER_DATA(ita)) + i*a_axis_stride))
            pt2 = (<double*>((<char*>np.PyArray_ITER_DATA(ita)) + (i+2)*a_axis_stride))
            value = -0.5*pt1[0] + 0.5*pt2[0]
            (<double*>((<char*>np.PyArray_ITER_DATA(ito)) + (i+1)*o_axis_stride))[0] = value

        # last element
        pt1 = (<double*>((<char*>np.PyArray_ITER_DATA(ita))+ (axis_length-3)*a_axis_stride))
        pt2 = (<double*>((<char*>np.PyArray_ITER_DATA(ita))+ (axis_length-2)*a_axis_stride))
        pt3 = (<double*>((<char*>np.PyArray_ITER_DATA(ita))+ (axis_length-1)*a_axis_stride))
        value = 1.5*pt3[0] - 2*pt2[0] + 0.5*pt1[0]
        (<double*>((<char*>np.PyArray_ITER_DATA(ito))+(axis_length-1)*o_axis_stride))[0] = value


        np.PyArray_ITER_NEXT(ita)
        np.PyArray_ITER_NEXT(ito)

    return out

代码生成正确的结果,而且确实比我自己的numpy实现(没有cython)要快。问题如下:

我只想过要一个 pt1 = (<double*>((<char*>np.PyArray_ITER_DATA(ita)) + i*a_axis_stride)) 语句,然后使用 pt1[0] , pt1[-1] , pt1[1] (<double*>((<char*>np.PyArray_ITER_DATA(ita)) + i*a_axis_stride)) 指向正确的,但是 pt[-1] 和 pt[1] pt[0] ,沿最后一个轴。目前的版本是可行的,但是如果我想实现更高阶的微分,这需要更多的点来评估,那么我最终会有很多这样的行,我不确定是否有更好/更有效的方法来使用它像这样的 pt[xxx]
有没有其他方法来加速这段代码?我正在寻找一些小细节,我可能忽略或微妙的事情,可以有很大的影响。

1 回复 | 直到 7 年前

Vladimir F ÐÐµÑÐ¾ÑÐ¼ ÑÐ»Ð°Ð²Ð° 7 年前

令我稍感意外的是,我实际上无法击败使用Cython类型的memoryview的您的版本-numpy迭代器看起来相当快。不过,我认为我可以显著提高可读性,让您使用Python切片语法。唯一的限制是,输入数组必须是C连续的,这样可以很容易地重新调整它的形状(我认为Fortran连续的也可以工作,但是我还没有测试)

基本技巧是展平选定轴前后的所有轴,使其成为已知的三维形状,此时可以使用Cython MemoryView。

@cython.boundscheck(False)
def test4(a,order=2,axis=-1):
    assert a.flags['C_CONTIGUOUS'] # otherwise the reshape doesn't work
    before = np.product(a.shape[:axis])
    after = np.product(a.shape[(axis+1):])
    cdef double[:,:,::1] a_new = a.reshape((before, a.shape[axis], after)) # this should not involve copying memory - it's just a new view
    cdef double[:] a_slice

    cdef double[:,:,::1] out = np.empty_like(a_new)

    assert a_new.shape[1] > 3

    cdef int m,n,i

    for m in range(a_new.shape[0]):
        for n in range(a_new.shape[2]):
            a_slice = a_new[m,:,n]

            out[m,0,n] = -1.5*a_slice[0] + 2*a_slice[1] - 0.5*a_slice[2]


            for i in range(a_slice.shape[0]-2):
                out[m,i+1,n] = -0.5*a_slice[i] + 0.5*a_slice[i+2]

            # last element
            out[m,-1,n] = 1.5*a_slice[-1] - 2*a_slice[-2] + 0.5*a_slice[-3]

    return np.asarray(out).reshape(a.shape)

速度比你的版本稍微慢一点。

在改进代码方面,可以用双倍而不是字节计算跨距( a_axis_stride_dbl = a_axis_stride/sizeof(double) )然后索引为 pt[i*a_axis_stride_dbl] ). 它可能不会获得太多的速度,但会更可读(这就是你在第1点中所问的问题)