可以使用numba.jit修饰f2py编译的函数吗?

时间:2016-02-04 08:31:52

标签: python numpy numba f2py

假设我有一个适用于一维数组的功能:

def opfunc(a, b):
    if len(a.shape) != 1 and len(b.shape) != 1:
        raise ValueError('opfunc: both inputs must be 1-d arrays.')
    return 2 * b + 2 * (a[0] - a[-1])

我希望在循环或嵌套循环中重复使用它,比如

nx1, nx2, nx3 = a.shape
c = np.zeros(a.shape)
for ix1 in range(nx1):
    for ix2 in range(nx2):
        avec = a[ix1, ix2, :]
        bvec = b[ix1, ix2, :]
        cvec = opfunc(avec, bvec)
            c[ix1, ix2, :] = cvec

我可以通过将这些语句放在一个函数中并使用opfunc numba.jit进行装饰来加快速度:

@numba.jit(nopython=True)
def opfunc(a, b):
    if len(a.shape) != 1 and len(b.shape) != 1:
        raise ValueError('opfunc: both inputs must be 1-d arrays.')
    return 2 * b + 2 * (a[0] - a[-1])

@numba.jit(nopython=True)
def opfunc1(a, b):
    nx1, nx2, nx3 = a.shape
    c = np.zeros(a.shape)
    for ix1 in range(nx1):
        for ix2 in range(nx2):
            avec = a[ix1, ix2, :]
            bvec = b[ix1, ix2, :]
            cvec = opfunc(avec, bvec)
            c[ix1, ix2, :] = cvec
    return c

这里确定加速。

如果在Fortran中编写opfunc

,是否可以做类似的事情
$ !cat fortran_test.f90

module test_fortran

  implicit none

contains

  subroutine opfunc(a, b, c)
    implicit none

    real (kind=8), dimension(6), intent(in) :: a, b
    real (kind=8), dimension(6), intent(out) :: c

    c = 2 * b + 2 !* (a(1) - a(ubound(a)))
  end subroutine opfunc

end module test_fortran

并已使用f2py导入Python?

import fortran_test

opfuncx = fortran_test.test_fortran.opfunc

@numba.jit(nopython=True)
def opfunc1x(a, b):
    nx1, nx2, nx3 = a.shape
    c = np.zeros(a.shape)
    for ix1 in range(nx1):
        for ix2 in range(nx2):
            avec = a[ix1, ix2, :]
            bvec = b[ix1, ix2, :]
            cvec = opfuncx(avec, bvec)
            c[ix1, ix2, :] = cvec
    return c

实际上,我正在

TypingError: Failed at nopython (nopython frontend)
Untyped global name 'opfuncx'
File "<ipython-input-48-c51d5b4fdad5>", line 13

运行该功能时。我猜这是因为opfuncx没有用numba.jit装饰。有可能这样做吗?

0 个答案:

没有答案