Numba加速嵌套循环调用Cartopy投影转换失败问题排查

我刚接触Numba，正在学习其使用方法。需要转换大气模型输出的4D风场矢量（维度为time, height, latitude, longitude）的投影。由于Cartopy的矢量投影转换函数仅支持2D（水平）矢量场，因此必须遍历时间和高度维度。但面对尺寸为[48, 90, 1200, 1200]的大型数组，使用Numba并行化嵌套循环后并未提速，耗时长达2小时！怀疑是Numba无法加速Cartopy函数，附上精简代码求指出问题：

import cartopy.crs as ccrs
import numpy, xarray
from numba import jit, prange

def coord_ref_systems(cube):
     global source_crs, target_crs
     source_crs = ccrs.RotatedPole(pole_longitude=140,
                      pole_latitude=36,
                      globe=ccrs.Globe(semimajor_axis=6370000, semiminor_axis=6370000))
     target_crs = ccrs.PlateCarree()

def vector_transform(u, v, lon, lat):
    xx, yy = numpy.meshgrid(lon, lat)
    u_t, v_t = target_crs.transform_vectors(source_crs, xx, yy, u, v)
    return u_t, v_t

@jit(parallel=True)
def transformer(u, v, lon, lat):
    u_new, v_new = numpy.empty(u.shape), numpy.empty(u.shape)
    for ti in prange(len(u.shape[0])):
        for zi in prange(len(u.shape[1])):
            u_t, v_t = vector_transform(u[ti, zi, :, :].squeeze(), v[ti, zi, :, :].squeeze(), lon, lat)
            u_new[ti, zi, :, :] = u_t
            v_new[ti, zi, :, :] = v_t
    return u_new, v_new

def main():
    ds = xarray.open_mfdataset('wind_vectors.nc', chunks={'time':6, 'height':-1, 'latitude': 200, 'longitude': 200})
    u = ds['U']
    v = ds['V']
    _u_, _v_ = transformer(u.values, v.values, u.longitude.values, u.latitude.values)

if __name__ == '__main__':
    main()

核心问题：Numba无法加速Cartopy函数 #

你猜的没错，Numba对Cartopy的transform_vectors函数完全起不到加速作用。原因很简单：Numba的JIT加速只对纯Python/NumPy代码有效，而Cartopy的底层实现是C++/GEOS库，属于Python扩展模块，Numba无法对这类外部函数进行编译优化。你的transformer函数里，绝大多数耗时都在调用vector_transform（本质是调用Cartopy的函数），Numba只能优化循环的框架部分，这部分在整个流程里占比极低，所以并行化后完全看不到提速效果，甚至可能因为线程调度开销变慢。

具体代码问题 #

• 嵌套prange浪费资源：Numba的prange嵌套两层会导致线程数爆炸（比如8核CPU会生成64个线程），增加调度成本，完全没必要，只需要并行化最外层循环即可。
• 重复生成网格：vector_transform里每次循环都调用numpy.meshgrid(lon, lat)，生成完全相同的网格数据，属于无效重复计算。
• 未利用Xarray分块：用chunks参数做了分块，但最后直接调用u.values把整个数组加载到内存，完全浪费了Dask的分块并行能力。

优化后的代码 #

方案一：提前生成网格+合理并行化 #

import cartopy.crs as ccrs
import numpy, xarray
from numba import jit, prange

# 提前定义坐标系（避免全局变量带来的副作用）
source_crs = ccrs.RotatedPole(pole_longitude=140,
                              pole_latitude=36,
                              globe=ccrs.Globe(semimajor_axis=6370000, semiminor_axis=6370000))
target_crs = ccrs.PlateCarree()

def main():
    ds = xarray.open_mfdataset('wind_vectors.nc', chunks={'time':6, 'height':-1, 'latitude': 200, 'longitude': 200})
    u = ds['U']
    v = ds['V']
    
    # 提前生成一次网格，避免重复计算
    lon = u.longitude.values
    lat = u.latitude.values
    xx, yy = numpy.meshgrid(lon, lat)
    
    # 只并行化最外层时间维度，减少线程调度开销
    @jit(parallel=True, nogil=True)
    def transformer(u_arr, v_arr):
        u_new = numpy.empty_like(u_arr)
        v_new = numpy.empty_like(v_arr)
        for ti in prange(u_arr.shape[0]):
            for zi in range(u_arr.shape[1]):
                u_t, v_t = target_crs.transform_vectors(source_crs, xx, yy, 
                                                        u_arr[ti, zi, :, :], 
                                                        v_arr[ti, zi, :, :])
                u_new[ti, zi, :, :] = u_t
                v_new[ti, zi, :, :] = v_t
        return u_new, v_new
    
    _u_, _v_ = transformer(u.values, v.values)

if __name__ == '__main__':
    main()

方案二：利用Dask并行化（适合超大型数组） #

如果数组内存装不下，直接用Dask的map_blocks分块处理，自动实现并行，完全不需要Numba：

import cartopy.crs as ccrs
import numpy, xarray

source_crs = ccrs.RotatedPole(pole_longitude=140,
                              pole_latitude=36,
                              globe=ccrs.Globe(semimajor_axis=6370000, semiminor_axis=6370000))
target_crs = ccrs.PlateCarree()

def transform_block(u_block, v_block, lon, lat):
    xx, yy = numpy.meshgrid(lon, lat)
    u_t, v_t = target_crs.transform_vectors(source_crs, xx, yy, u_block, v_block)
    return u_t, v_t

def main():
    ds = xarray.open_mfdataset('wind_vectors.nc', chunks={'time':6, 'height':-1, 'latitude': 200, 'longitude': 200})
    u = ds['U']
    v = ds['V']
    
    lon = u.longitude.values
    lat = u.latitude.values
    
    # 对每个时间-高度块应用转换
    u_transformed = u.map_blocks(lambda x: transform_block(x, v.sel(time=x.time, height=x.height), lon, lat),
                                 template=u)
    v_transformed = v.map_blocks(lambda x: transform_block(u.sel(time=x.time, height=x.height), x, lon, lat)[1],
                                 template=v)
    
    # 计算并保存结果
    u_transformed.to_netcdf('u_transformed.nc')
    v_transformed.to_netcdf('v_transformed.nc')

if __name__ == '__main__':
    main()

关键优化点总结 #

• 移除Numba对Cartopy函数的无效包裹，仅在纯Python循环部分合理使用，或直接用Dask并行。
• 提前生成网格数据，避免重复计算。
• 利用Xarray+Dask的分块能力，处理超大型数组时避免内存溢出，同时实现高效并行。

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