Python 中 Numba @njit 编译失败及复数数组 dtype 报错如何解决？

2026-11-04数据小陈

结论

这两个错误发生在不同位置，需要分别处理：

让 res() 调用的 Ll() 和 dLl() 能被 Numba 编译，例如为它们添加 @njit。
创建复数数组时，使用明确的 NumPy 类型对象 np.complex128，不要传入字符串 "complex128"。
将 solve_ivp 的初值显式转换为复数数组。右端函数包含 1j，后续解通常会产生虚部。

还应移除右端函数中的频繁打印，并检查空间网格与差分公式使用的 h 是否一致。编译通过并不代表离散方程已经设置正确。

为什么会出现普通函数类型错误？

@njit 使用 nopython 模式编译函数。编译 res() 时，Numba 也需要识别其中调用的函数。

原代码中的：

def Ll(a, b, t):
    return a + b*t

def dLl(b):
    return b

是普通 Python 函数。在这条调用路径下，Numba 无法直接为它们生成所需的本地代码，因此可能报错：

Cannot determine Numba type of <class 'function'>

可以将它们改为：

@njit
def Ll(a, b, t):
    return a + b*t

@njit
def dLl(b):
    return b

也可以直接把这两条简单表达式写进 res()。

当前线性边界长度的计算路径没有调用 Ls() 和 dLs()，它们不必因为定义在同一个文件里就添加 @njit。f() 只在求解前计算一次初值，也不需要为解决这个错误而编译。

为什么 np.zeros_like 仍然报错？

错误日志中有两处信息：

dtype=Literal[str](complex128)

以及：

dtype=unicode_type

这说明此次编译收到的 dtype 是字符串。NumPy 能接受某些字符串形式的 dtype，Numba 对应的实现不一定支持。

可改为：

dphi = np.zeros_like(phi, dtype=np.complex128)

或者：

dphi = np.zeros(phi.shape, dtype=np.complex128)

这里的 np.complex128 是类型对象，没有引号。如果 phi 已经确定为复数数组，也可以写成：

dphi = np.zeros_like(phi)

问题展示的代码使用的是 dtype=complex，补充日志显示的却是字符串 complex128。需要确认实际执行的代码是否已经修改，以及交互环境中的函数定义是否已经重新运行。仅凭这些材料，不能断定最初使用 complex 的调用与补充日志来自同一次编译。

日志中的 readonly array(float64, 1d, C) 表示输入数组在编译器看来是只读的，这常见于作为全局变量使用的数组。zeros_like 只读取它的形状和类型，不会修改它。当前日志明确指出的匹配问题是字符串形式的 dtype，不能仅凭 readonly 就把原因归到数组只读上。

可参考的修改示例

下面的示例补全了题面代码的缩进，保留原右端表达式中 i 作为循环索引的含义，并显式传入 h、a 和 b。

示例假设 phi 表示区间内部的 n 个节点，左右端点值为零。这与原代码在首尾节点外侧补零的写法对应；如果实际边界条件不同，需要调整。

import numpy as np
from scipy.integrate import solve_ivp
from numba import njit


@njit
def Ll(a, b, t):
    return a + b*t


@njit
def dLl(b):
    return b


@njit
def res(t, phi, h, a, b):
    n = phi.size
    dphi = np.zeros(phi.shape, dtype=np.complex128)

    length = Ll(a, b, t)
    kinetic = 1j / (h**2 * length**2)
    drift = dLl(b) / (2.0 * h * length)

    for i in range(n):
        left = 0.0j
        right = 0.0j

        if i > 0:
            left = phi[i - 1]
        if i < n - 1:
            right = phi[i + 1]

        dphi[i] = (
            kinetic * (right - 2.0*phi[i] + left)
            + drift * i * (right - left)
        )

    return dphi


z0 = 0.0
zf = 1.0
n = 50

t0 = 0.0
tf = 1.0
a = 1.0
b = 3.0

# n 个内部节点，另有两个边界节点。
h = (zf - z0) / (n + 1)
z = np.linspace(z0, zf, n + 2)[1:-1]

y0 = (np.sqrt(2.0) * np.sin(2.0*np.pi*z)).astype(np.complex128)
t_eval = np.linspace(t0, tf, n)

# 提前触发 JIT 编译，便于将编译开销与后续求解区分开。
res(t0, y0, h, a, b)

sol = solve_ivp(
    res,
    (t0, tf),
    y0,
    args=(h, a, b),
    t_eval=t_eval,
    method="RK45",
)

if not sol.success:
    raise RuntimeError(sol.message)

print(sol.y)

这里使用支持复数状态的 RK45。如果更换求解方法，需要确认该方法是否支持复数问题；也可以将复数状态拆成实部和虚部后求解。

示例没有使用全局 z 来分配导数数组，因为导数的长度应与当前状态 phi 一致。显式传入参数便于后续改变网格或模型参数，也能避免依赖 Numba 对全局值的编译期处理。

还需要注意哪些数值问题？

初值为实数，不代表解始终为实数

原代码的：

y0 = f(z)

会生成实数数组，但右端函数包含 1j。要按复数状态积分，应在调用 solve_ivp 前转换初值：

y0 = np.asarray(f(z), dtype=np.complex128)

求解器接收到的初值需要是复数数组，不能只将 dphi 设为复数数组。程序能够执行，也不足以证明虚部得到了正确保留。

原网格的实际间距与 h 不一致

原代码设置：

h = 1/m
z = np.linspace(z0+h, zf-h, m)

在 z0=0、zf=1、m=50 时，h 为 0.02，但 linspace 生成的相邻点间距为：

(0.98 - 0.02) / 49 = 0.96 / 49

这个间距不等于 0.02。差分公式中的 h 应与实际节点间距一致。上面的示例使用 n 个内部节点，因此取 h=(zf-z0)/(n+1)。

循环索引 i 与虚数单位 1j 不同

原式中的：

1/2/h*i*dLl(b)/Ll(a,b,t)

使用的是整数索引 i。示例保留了这一含义，没有把它改成 1j。它是否应当表示索引、空间坐标或其他系数，需要依据实际离散方程核对。

另外，计算期间必须保证 Ll(a,b,t) 不为零，因为表达式中存在它的倒数和平方倒数。

如何评估加速效果？

先删除 res() 中的 print()。solve_ivp 会多次调用右端函数，打印边界值和整个导数数组会增加大量输出开销。

第一次调用 @njit 函数通常包含编译耗时。比较计算速度时，应在预热后使用相同参数、初值和求解容差计时。t_eval 指定输出时间点，求解器的计算并不限于这些时间点；可以结合 sol.nfev 查看右端函数调用次数。

这里由 Numba 加速的是右端函数，solve_ivp 仍在 Python 层调度。现有材料不足以预测加速倍数，也无法确定原先约 31 秒的耗时是否主要来自差分循环。