无法解决 Numba TypingError

Question

有一个过程涉及搜索（最终）大量的参数组合，我正试图使用 numba 来加快速度。我试图更改我的初始代码以使其兼容，但仍然出现我不知道如何解决的 TypingError。我在使用 'np.asarray()' 时遇到了很多问题，所以我删除了它并且现在主要处理标准列表，但它仍然不起作用。我只会提到代码中与错误消息相关的部分，如果需要更多，我也会提供（只是为了保持简短）。

错误消息中看似关键的部分（据我所知，其余部分只是由于在另一个函数中调用此函数而在另一个函数中再次调用）如下：

Failed in nopython mode pipeline (step: nopython frontend) No implementation of function Function(<function dot at 0x0000021490EA2400>) found for signature:

dot(list(list(float64)<iv=None>)<iv=None>, reflected list(float64)<iv=None>)

有 4 个候选实现： - 其中 4 个不匹配，原因是：函数“_OverloadWrapper._build..ol_generated”中的重载：文件：numba\core\overload_glue.py：第 131 行。带有参数：'(list(list(float64)), reflected list(float64))' : 由于实施引发了特定错误而被拒绝： TypingError：在 nopython 模式管道中失败（步骤：nopython 前端）未找到签名的函数 Function() 的实现：

  >>> stub(list(list(float64)<iv=None>)<iv=None>, reflected list(float64)<iv=None>)

好像来源于函数'nvaxes_lab'定义如下：

# reference position of NV axis vectors, each row holds the orientation (vector) of one NV axis    
nvaxes_ref = [[1.0, 1.0, 1.0],
     [-1.0, 1.0, -1.0], [1.0, -1.0, -1.0], [-1.0, -1.0, 1.0]]

@njit()
def rot_euler(a, b, g):
    """
    :param a: float, angle of third rotation (around z-axis) in radians
    :param b: float, angle of second rotation (around x-axis) in radians
    :param g: float, angle of first rotation (around z-axis) in radians
    :return: rotation matrix for a general rotation around the Euler angles a (=alpha), b (=beta), g (=gamma) defined
        as follows: rot_euler(a,b,g) = rotZ(a).rotX(b).rotZ(g)
    """
    return [[np.cos(a) * np.cos(g) - np.cos(b) * np.sin(a) * np.sin(g),
             -np.cos(b) * np.cos(g) * np.sin(a) - np.cos(a) * np.sin(g), np.sin(a) * np.sin(b)],
            [np.cos(g) * np.sin(a) + np.cos(a) * np.cos(b) * np.sin(g),
             np.cos(a) * np.cos(b) * np.cos(g) - np.sin(a) * np.sin(g), -np.cos(a) * np.sin(b)],
            [np.sin(b) * np.sin(g), np.cos(g) * np.sin(b), np.cos(b)]]

@njit()    
def nvaxes_lab(a, b, g):
    """
    :param a: float, angle of third Euler rotation (around z-axis) in radians
    :param b: float, angle of second Euler rotation (around x-axis) in radians
    :param g: float, angle of first Euler rotation (around z-axis) in radians
    :return: numpy array of dimensions 3x4, each row holds the orientation (vector) of the corresponding NV-axis
        after rotation by the Euler angles a,b,g
    """
    return [np.dot(rot_euler(a, b, g), nvaxes_ref[0]).astype(float), np.dot(rot_euler(a, b, g), nvaxes_ref[1]).astype(float),
            np.dot(rot_euler(a, b, g).astype(float), nvaxes_ref[2]), np.dot(rot_euler(a, b, g), nvaxes_ref[3]).astype(float)]

问题似乎出在 'nvaxes_lab' 中返回的数组的每个分量中的点积......不幸的是，我不知道如何定义数组（我尝试了多个版本）或其他什么因为这是我第一次使用 numba，所以我可以让这个错误消失。

非常感谢您的帮助！

Answer 1

你有几个问题。首先，您看到的实际错误来自于 numba 编译您的函数时，它使用了 numpy 函数的 C 实现。这些实现需要 numpy 数组，而不是列表。所以错误：

dot(list(list(float64)<iv=None>)<iv=None>, reflected list(float64)<iv=None>)

基本上是想告诉您 np.dot 函数不知道如何处理列表。在 python 版本中，它会将列表强制转换为一个 numpy 数组，但我想在 C/numba 版本中不会。您可以将 np.array 环绕在每个列表周围，或者更好的选择是输出一个数组 rot_euler.

接下来，最好不要引用在您尝试进行 JIT 的函数范围之外定义的可变变量。在 nvaxes_lab 中，您引用了 nvaxes_ref 的每个元素。您最好将它们作为函数的参数传递，并且为了保持一致，也将它们设为数组。

最后，您的文档字符串与 nvaxes_lab 的 return 类型不匹配。您可以将整个列表 vstack 以使其成为一个数组。

在将东西转换为浮点数方面，因为 nvaxes_ref 数组已经是一个浮点数数组，所有输出也将是浮点数，所以你可以放弃 astype 方法。

import numpy as np
from numba import njit


nvaxes_ref = np.array([
     [1.0, 1.0, 1.0],
     [-1.0, 1.0, -1.0], 
     [1.0, -1.0, -1.0], 
     [-1.0, -1.0, 1.0]
])

@njit
def rot_euler(a, b, g):
    """
    :param a: float, angle of third rotation (around z-axis) in radians
    :param b: float, angle of second rotation (around x-axis) in radians
    :param g: float, angle of first rotation (around z-axis) in radians
    :return: rotation matrix for a general rotation around the Euler angles
        a (=alpha), b (=beta), g (=gamma) defined as follows: 
        rot_euler(a,b,g) = rotZ(a).rotX(b).rotZ(g)
    """
    return np.array([
        [
            np.cos(a) * np.cos(g) - np.cos(b) * np.sin(a) * np.sin(g), 
            -np.cos(b) * np.cos(g) * np.sin(a) - np.cos(a) * np.sin(g), 
            np.sin(a) * np.sin(b)
        ],
        [
            np.cos(g) * np.sin(a) + np.cos(a) * np.cos(b) * np.sin(g),
            np.cos(a) * np.cos(b) * np.cos(g) - np.sin(a) * np.sin(g), 
            -np.cos(a) * np.sin(b)
        ],
        [
            np.sin(b) * np.sin(g), 
            np.cos(g) * np.sin(b), 
            np.cos(b)
        ]
    ])

@njit
def nvaxes_lab(a, b, g, refs):
    """
    :param a: float, angle of third Euler rotation (around z-axis) in radians
    :param b: float, angle of second Euler rotation (around x-axis) in radians
    :param g: float, angle of first Euler rotation (around z-axis) in radians
    :return: numpy array of dimensions 3x4, each row holds the orientation (vector) 
             of the corresponding NV-axis after rotation by the Euler angles a,b,g
    """
    return [
            np.dot(rot_euler(a, b, g), refs[0]),
            np.dot(rot_euler(a, b, g), refs[1]),
            np.dot(rot_euler(a, b, g), refs[2]),
            np.dot(rot_euler(a, b, g), refs[3])
    ]

无法解决 Numba TypingError

Cannot resolve Numba TypingError

python

numpy

list

numba