使用Cython包装LAPACKE函数

时间:2014-04-21 15:09:13

标签: python numpy linear-algebra cython lapack

我试图使用Cython包装LAPACK函数dgtsv(三对角方程组的求解器)。

我遇到this previous answer,但由于dgtsv不是scipy.linalg中包含的LAPACK函数之一,我不认为我可以使用这种特殊方法。相反,我一直试图关注this example

这是我lapacke.pxd文件的内容:

ctypedef int lapack_int

cdef extern from "lapacke.h" nogil:

    int LAPACK_ROW_MAJOR
    int LAPACK_COL_MAJOR

    lapack_int LAPACKE_dgtsv(int matrix_order,
                             lapack_int n,
                             lapack_int nrhs,
                             double * dl,
                             double * d,
                             double * du,
                             double * b,
                             lapack_int ldb)

...这是我在_solvers.pyx中的瘦Cython包装器:

#!python

cimport cython
from lapacke cimport *

cpdef TDMA_lapacke(double[::1] DL, double[::1] D, double[::1] DU,
                   double[:, ::1] B):

    cdef:
        lapack_int n = D.shape[0]
        lapack_int nrhs = B.shape[1]
        lapack_int ldb = B.shape[0]
        double * dl = &DL[0]
        double * d = &D[0]
        double * du = &DU[0]
        double * b = &B[0, 0]
        lapack_int info

    info = LAPACKE_dgtsv(LAPACK_ROW_MAJOR, n, nrhs, dl, d, du, b, ldb)

    return info

...这里是一个Python包装器和测试脚本:

import numpy as np
from scipy import sparse
from cymodules import _solvers


def trisolve_lapacke(dl, d, du, b, inplace=False):

    if (dl.shape[0] != du.shape[0] or dl.shape[0] != d.shape[0] - 1
            or b.shape != d.shape):
        raise ValueError('Invalid diagonal shapes')

    if b.ndim == 1:
        # b is (LDB, NRHS)
        b = b[:, None]

    # be sure to force a copy of d and b if we're not solving in place
    if not inplace:
        d = d.copy()
        b = b.copy()

    # this may also force copies if arrays are improperly typed/noncontiguous
    dl, d, du, b = (np.ascontiguousarray(v, dtype=np.float64)
                    for v in (dl, d, du, b))

    # b will now be modified in place to contain the solution
    info = _solvers.TDMA_lapacke(dl, d, du, b)
    print info

    return b.ravel()


def test_trisolve(n=20000):

    dl = np.random.randn(n - 1)
    d = np.random.randn(n)
    du = np.random.randn(n - 1)

    M = sparse.diags((dl, d, du), (-1, 0, 1), format='csc')
    x = np.random.randn(n)
    b = M.dot(x)

    x_hat = trisolve_lapacke(dl, d, du, b)

    print "||x - x_hat|| = ", np.linalg.norm(x - x_hat)

不幸的是,test_trisolve只是在调用_solvers.TDMA_lapacke时出现了段错误。 我非常确定我的setup.py是正确的 - ldd _solvers.so表明_solvers.so在运行时正在链接到正确的共享库。

我不确定如何从这里开始 - 任何想法?


简要更新

对于较小的n值我不会立即得到段错误,但我确实得到了无意义的结果( || x - x_hat || 应该非常接近0):

In [28]: test_trisolve2.test_trisolve(10)
0
||x - x_hat|| =  6.23202576396

In [29]: test_trisolve2.test_trisolve(10)
-7
||x - x_hat|| =  3.88623414288

In [30]: test_trisolve2.test_trisolve(10)
0
||x - x_hat|| =  2.60190676562

In [31]: test_trisolve2.test_trisolve(10)
0
||x - x_hat|| =  3.86631743386

In [32]: test_trisolve2.test_trisolve(10)
Segmentation fault

通常LAPACKE_dgtsv返回代码0(应该表示成功),但偶尔会得到-7,这意味着参数7(b)具有非法值。发生的事情是b的第一个值实际上正在被修改。如果我继续致电test_trisolve,即使n很小,我也会最终遇到段错误。

2 个答案:

答案 0 :(得分:3)

好吧,我最终弄明白了 - 似乎我误解了在这种情况下行和列专业所指的是什么。

由于C连续数组遵循行主顺序,我认为我应该指定LAPACK_ROW_MAJOR作为LAPACKE_dgtsv的第一个参数。

事实上,如果我改变了

info = LAPACKE_dgtsv(LAPACK_ROW_MAJOR, ...)

info = LAPACKE_dgtsv(LAPACK_COL_MAJOR, ...)

然后我的功能起作用:

test_trisolve2.test_trisolve()
0
||x - x_hat|| =  6.67064747632e-12

这对我来说似乎非常直观 - 任何人都能解释为什么会这样吗?

答案 1 :(得分:1)

虽然相当陈旧但这个问题似乎仍然具有相关性。 观察到的行为是参数LDB错误解释的结果:

  • Fortran阵列是主要的,并且阵列B的主要维度对应于N.因此LDB> = max(1,N)。
  • 行主要LDB对应于NRHS,因此必须满足条件LDB> = max(1,NRHS)。

注释 #b是(LDB,NRHS)不正确,因为b具有维度(LDB,N),在这种情况下LDB应为1。

只要NRHS等于1,从LAPACK_ROW_MAJOR切换到LAPACK_COL_MAJOR就可以解决问题.col major(N,1)的内存布局与行major(1,N)相同。但是,如果NRHS大于1,它将失败。