Numba njit并行实现局部秩变换出现非确定性错误输出问题

问题原因 #

• 核心竞态问题：padded_sig、result两个临时变量定义在并行循环for i in nb.prange(...)外部，所有并行线程共享这两块内存空间。不同线程会同时覆盖padded_sig的写入内容、同时修改result的元素值，导致计算时读取到其他线程的中间数据，最终结果非确定且不符合逻辑。
• 切片逻辑错误：多输入版本的邻域切片和已验证的单输入版本不一致，单输入版本切片为padded_sig[idx - pad:idx + pad]，多输入版本写为padded_sig[j - pad:j + pad + 1]，n=3时邻域大小多取1个元素，也是导致sum结果出现理论上不可能的3的原因之一。
• 嵌套并行冗余：外层已经使用prange做多输入并行，内层额外加prange会增加线程调度开销，还会加剧竞态问题的出现概率。

修复方案 #

方案1：独立临时变量+修正逻辑 #

将临时变量移到并行循环内部，保证每个线程拥有独立的临时空间，同时修正切片逻辑、删除内层多余的prange：

import numba as nb
import numpy as np

@nb.njit(["float64[:,:](float64[:,:], int32)",
          "float32[:,:](float32[:,:], int32)",
          "uint8[:,:](uint8[:,:], int32)"], parallel=True, nogil=True)
def local_rank_transform_many_jit(signals, n):
    pad = n//2
    sig_length = signals.shape[1]
    all_results = np.zeros_like(signals)
    
    for i in nb.prange(signals.shape[0]):
        # 每个线程独立创建临时变量，避免共享
        result = np.zeros(sig_length, dtype=signals.dtype)
        padded_sig = np.zeros(sig_length + n - 1, dtype=signals.dtype)
        padded_sig[pad:-pad] = signals[i]

        # 移除内层prange，避免嵌套并行
        for j in range(pad, len(result) + pad):
            # 邻域切片和单输入版本保持一致
            neighb = padded_sig[j - pad:j + pad]
            result[j - pad] = np.sum(neighb < padded_sig[j])
        
        all_results[i] = result
            
    return all_results

方案2：复用单输入函数（更易维护） #

直接在并行循环内调用已经验证正确的单输入版本函数，不需要重复编写逻辑，避免人为写错：

@nb.njit(["float64[:,:](float64[:,:], int32)",
          "float32[:,:](float32[:,:], int32)",
          "uint8[:,:](uint8[:,:], int32)"], parallel=True, nogil=True)
def local_rank_transform_many_jit(signals, n):
    all_results = np.zeros_like(signals)
    for i in nb.prange(len(signals)):
        all_results[i] = local_rank_transform_jit(signals[i], n)
    return all_results

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