如何在Numba（Python）中实现1D与2D数组的行级NaN检查？

你遇到的问题本质是Numba在nopython模式下对变量类型的严格要求：同一个变量不能同时是1D和2D数组，而np.isnan对不同维度输入的返回值形状不一致，直接导致类型推导失败。

针对这个问题，拆分出专门处理1D和2D数组的子函数是最稳妥且高效的方案，当然也可以利用Numba的重载机制来实现更透明的调用。下面详细说明两种解决思路，并修正你的代码：

方案一：拆分主函数与维度专属子函数 #

这种方式逻辑清晰，每个子函数只处理单一维度的数组，Numba可以完美推导类型，避免冲突。

修正后的代码 #

import numpy as np
import numba as nb

# 专门处理1D数组的子函数
@nb.njit(nb.float64[:](nb.float64[:]))
def _check_1d(x):
    result = np.full_like(x, np.nan)
    for i in range(x.shape[0]):
        if not np.isnan(x[i]):
            # 替换为你实际需要执行的操作
            result[i] = 1.0
    return result

# 专门处理2D数组的子函数
@nb.njit(nb.float64[:,:](nb.float64[:,:]))
def _check_2d(x):
    result = np.full_like(x, np.nan)
    # 遍历每行每列（如果需要检查整行非NaN，可以调整逻辑）
    for i in range(x.shape[0]):
        for j in range(x.shape[1]):
            if not np.isnan(x[i, j]):
                # 替换为你实际需要执行的操作
                result[i, j] = 1.0
    return result

# 主函数：根据数组维度分发到对应子函数
@nb.njit
def is_not_nan_in_array_or_scalar(x):
    if x.ndim == 1:
        return _check_1d(x)
    elif x.ndim == 2:
        return _check_2d(x)
    else:
        raise ValueError("Only 1D or 2D float64 arrays are supported")

# 测试
print(is_not_nan_in_array_or_scalar(np.array([1, 2.5, np.NaN])))
print(is_not_nan_in_array_or_scalar(np.array([[1], [2.5], [np.NaN]])))

输出结果 #

[ 1.  1. nan]
[[ 1.]
 [ 1.]
 [nan]]

方案二：利用Numba的函数重载机制 #

如果你希望调用时更透明（不需要显式的条件判断），可以使用Numba的@nb.overload装饰器，让Numba根据输入数组的维度自动选择对应的实现：

import numpy as np
import numba as nb

@nb.overload(is_not_nan_in_array_or_scalar)
def ol_is_not_nan(x):
    # 匹配1D float64数组
    if x.ndim == 1 and x.dtype == nb.float64:
        def impl(x):
            result = np.full_like(x, np.nan)
            for i in range(x.shape[0]):
                if not np.isnan(x[i]):
                    result[i] = 1.0
            return result
        return impl
    # 匹配2D float64数组
    elif x.ndim == 2 and x.dtype == nb.float64:
        def impl(x):
            result = np.full_like(x, np.nan)
            for i in range(x.shape[0]):
                for j in range(x.shape[1]):
                    if not np.isnan(x[i, j]):
                        result[i, j] = 1.0
            return result
        return impl
    # 不支持的输入类型
    else:
        raise ValueError("Only 1D or 2D float64 arrays are supported")

# 空的主函数，实际实现由overload提供
@nb.njit
def is_not_nan_in_array_or_scalar(x):
    pass

# 测试
print(is_not_nan_in_array_or_scalar(np.array([1, 2.5, np.NaN])))
print(is_not_nan_in_array_or_scalar(np.array([[1], [2.5], [np.NaN]])))

原代码的问题分析 #

1. 类型冲突：is_not_nan变量在1D场景下是1D布尔数组，在2D场景下是flatten后的1D数组（原代码中~np.isnan(x)对2D数组返回2D布尔数组，而else分支又flatten成1D，直接导致类型无法统一），Numba在nopython模式下不允许这种变量类型的不确定性。
2. 逻辑错误：循环使用len(x)，当输入是2D数组时，len(x)是行数，而is_not_nan的长度是数组总元素数，会导致索引越界。

总结 #

无论是拆分子函数还是使用重载，核心思路都是让每个处理逻辑只对应单一维度的输入，避免Numba进行类型推导时遇到冲突。这种方式不仅解决了类型问题，还能让代码的性能和可读性更优。

内容的提问来源于stack exchange，提问作者Olibarer