Python嵌套For循环的PyCUDA并行化及HPC适配技术问询

首先，32万×700万的嵌套循环确实是天文数字，单线程Python跑起来肯定遥遥无期。既然你没有Python CUDA编程经验，下面给你几个完全不用手动写CUDA代码的最简改造方案，直接适配你的HPC多节点/多GPU资源，按上手难度从易到难排序：

你的原始核心代码 #

for i in range(0, len(longitude_aq)):
    center = Coordinates(latitude_aq[i], longitude_aq[i])
    currentAq = aq[i, :]
    for j in range(0, len(longitude_meo)):
        currentMeo = meo[j, :]
        grid_point = Coordinates(latitude_meo[j], longitude_meo[j])
        if is_in_circle(center, RADIUS, grid_point):
            if currentAq[TIME_AQ] == currentMeo[TIME_MEO]:
                humidity += currentMeo[HUMIDITY_MEO]
                pressure += currentMeo[PRESSURE_MEO]
                temperature += currentMeo[TEMPERATURE_MEO]
                wind_speed += currentMeo[WIND_SPEED_MEO]
                wind_direction += currentMeo[WIND_DIRECTION_MEO]
                count += 1.0
    if count != 0.0:
        final_tmp[i, HUMIDITY_FINAL] = humidity/count
        final_tmp[i, PRESSURE_FINAL] = pressure/count
        final_tmp[i, TEMPERATURE_FINAL] = temperature/count
        final_tmp[i, WIND_SPEED_FINAL] = wind_speed/count
        final_tmp[i, WIND_DIRECTION_FINAL] = wind_direction/count
    humidity, pressure, temperature, wind_speed, wind_direction, count = 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
final.loc[:, :] = final_tmp[:, :]

方案1：Numba自动CPU/GPU并行（上手最快，改动最小） #

Numba是Python的即时编译库，能自动把你的Python代码编译成高效的机器码（甚至CUDA代码），完全不需要你懂CUDA语法。这是最适合你的入门方案，改动极小就能获得几十倍的性能提升。

改造步骤： #

1. 适配Numba的nopython模式：Numba对自定义类（比如你的Coordinates）支持有限，所以我们把类改成直接传递经纬度数值，同时把is_in_circle函数用Numba编译。
2. 并行标记外层循环：用numba.prange代替普通range，Numba会自动把外层循环分配到所有CPU核心上；如果要用到GPU，只需要换个装饰器，Numba会自动生成CUDA内核。

CPU并行版代码（几乎和原逻辑一致）

import numba
import numpy as np

# 先把距离判断函数用Numba编译，改成直接传数值
@numba.jit(nopython=True)
def is_in_circle(center_lat, center_lon, radius, point_lat, point_lon):
    # 替换成你实际的距离判断逻辑（比如球面距离），示例用平面欧氏距离简化
    dx = center_lon - point_lon
    dy = center_lat - point_lat
    return dx*dx + dy*dy <= radius*radius

# 标记并行编译，nopython=True确保完全编译成机器码
@numba.jit(nopython=True, parallel=True)
def compute_aq_meo_mean(latitude_aq, longitude_aq, aq, latitude_meo, longitude_meo, meo, RADIUS, TIME_AQ, TIME_MEO, HUMIDITY_MEO, PRESSURE_MEO, TEMPERATURE_MEO, WIND_SPEED_MEO, WIND_DIRECTION_MEO):
    n_aq = latitude_aq.shape[0]
    n_meo = latitude_meo.shape[0]
    final_tmp = np.zeros((n_aq, 5))  # 对应5个气象要素的结果

    # prange告诉Numba并行处理这个循环
    for i in numba.prange(n_aq):
        center_lat = latitude_aq[i]
        center_lon = longitude_aq[i]
        aq_time = aq[i, TIME_AQ]
        
        # 初始化累加变量
        humidity, pressure, temperature, wind_speed, wind_direction, count = 0.0, 0.0, 0.0, 0.0, 0.0, 0.0

        # 内层循环：先判断时间，减少不必要的距离计算
        for j in range(n_meo):
            meo_time = meo[j, TIME_MEO]
            if aq_time != meo_time:
                continue
            # 判断是否在半径范围内
            if is_in_circle(center_lat, center_lon, RADIUS, latitude_meo[j], longitude_meo[j]):
                humidity += meo[j, HUMIDITY_MEO]
                pressure += meo[j, PRESSURE_MEO]
                temperature += meo[j, TEMPERATURE_MEO]
                wind_speed += meo[j, WIND_SPEED_MEO]
                wind_direction += meo[j, WIND_DIRECTION_MEO]
                count += 1.0
        
        # 计算均值
        if count != 0.0:
            final_tmp[i, 0] = humidity / count
            final_tmp[i, 1] = pressure / count
            final_tmp[i, 2] = temperature / count
            final_tmp[i, 3] = wind_speed / count
            final_tmp[i, 4] = wind_direction / count

    return final_tmp

# 调用函数（确保所有输入都是NumPy数组，Pandas列可以用.values转成NumPy数组）
final_tmp = compute_aq_meo_mean(
    latitude_aq=np.array(latitude_aq),
    longitude_aq=np.array(longitude_aq),
    aq=aq.values,  # 如果aq是Pandas DataFrame
    latitude_meo=np.array(latitude_meo),
    longitude_meo=np.array(longitude_meo),
    meo=meo.values,
    RADIUS=RADIUS,
    TIME_AQ=TIME_AQ,
    TIME_MEO=TIME_MEO,
    HUMIDITY_MEO=HUMIDITY_MEO,
    PRESSURE_MEO=PRESSURE_MEO,
    TEMPERATURE_MEO=TEMPERATURE_MEO,
    WIND_SPEED_MEO=WIND_SPEED_MEO,
    WIND_DIRECTION_MEO=WIND_DIRECTION_MEO
)

# 把结果写回Pandas DataFrame
final.loc[:, [HUMIDITY_FINAL, PRESSURE_FINAL, TEMPERATURE_FINAL, WIND_SPEED_FINAL, WIND_DIRECTION_FINAL]] = final_tmp

GPU版本（依然不用写CUDA）

如果要用到HPC的GPU，只需要把装饰器换成@numba.cuda.jit，然后调整循环为CUDA的网格结构，或者用Numba提供的简化接口。不过GPU版本需要确保你的数据能加载到GPU内存里（700万条MEO数据应该没问题），核心逻辑和CPU版完全一致，不需要手动写CUDA内核。

方案2：Dask分布式并行（适合超大规模数据，多节点适配） #

如果你的数据大到单节点内存放不下，Dask可以把数据分片到HPC的多个节点上，自动处理分布式计算，同样不需要写CUDA代码。

改造步骤： #

1. 把所有数据转换成Dask数组（或Dask DataFrame），设置合适的分片大小。
2. 用map_blocks并行处理每个AQ站点分片，最后汇总结果。

简化示例代码

import dask.array as da
import numpy as np
from dask.diagnostics import ProgressBar

# 把NumPy数组/Pandas DataFrame转换成Dask数组，设置分片大小（根据你的内存调整）
latitude_aq_da = da.from_array(latitude_aq, chunks=1000)  # 每块1000个AQ站点
longitude_aq_da = da.from_array(longitude_aq, chunks=1000)
aq_da = da.from_array(aq.values, chunks=(1000, aq.shape[1]))
latitude_meo_da = da.from_array(latitude_meo, chunks=10000)
longitude_meo_da = da.from_array(longitude_meo, chunks=10000)
meo_da = da.from_array(meo.values, chunks=(10000, meo.shape[1]))

# 定义每个分片的处理函数（和Numba版的核心逻辑一致）
def process_chunk(lat_aq_chunk, lon_aq_chunk, aq_chunk, lat_meo, lon_meo, meo, RADIUS, TIME_AQ, TIME_MEO):
    n_chunk = lat_aq_chunk.shape[0]
    result = np.zeros((n_chunk, 5))
    for i in range(n_chunk):
        center_lat = lat_aq_chunk[i]
        center_lon = lon_aq_chunk[i]
        aq_time = aq_chunk[i, TIME_AQ]
        humidity, pressure, temperature, wind_speed, wind_direction, count = 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
        
        for j in range(lat_meo.shape[0]):
            meo_time = meo[j, TIME_MEO]
            if aq_time != meo_time:
                continue
            if is_in_circle(center_lat, center_lon, RADIUS, lat_meo[j], lon_meo[j]):
                humidity += meo[j, HUMIDITY_MEO]
                pressure += meo[j, PRESSURE_MEO]
                temperature += meo[j, TEMPERATURE_MEO]
                wind_speed += meo[j, WIND_SPEED_MEO]
                wind_direction += meo[j, WIND_DIRECTION_MEO]
                count += 1.0
        
        if count != 0.0:
            result[i] = [humidity/count, pressure/count, temperature/count, wind_speed/count, wind_direction/count]
    return result

# 用map_blocks并行处理所有AQ分片
final_tmp_da = da.map_blocks(
    process_chunk,
    latitude_aq_da,
    longitude_aq_da,
    aq_da,
    latitude_meo_da,
    longitude_meo_da,
    meo_da,
    RADIUS=RADIUS,
    TIME_AQ=TIME_AQ,
    TIME_MEO=TIME_MEO,
    chunks=(latitude_aq_da.chunks[0], 5)
)

# 计算结果（自动利用HPC多节点资源）
with ProgressBar():
    final_tmp = final_tmp_da.compute()

# 写回结果
final.loc[:, [HUMIDITY_FINAL, PRESSURE_FINAL, TEMPERATURE_FINAL, WIND_SPEED_FINAL, WIND_DIRECTION_FINAL]] = final_tmp

HPC上运行Dask的注意事项： #

• 用dask-jobqueue提交任务到HPC的调度系统（比如Slurm），它会自动创建Dask集群。
• 或者用dask-mpi连接多节点，适合MPI环境的HPC。

方案3：mpi4py多节点并行（贴合HPC原生MPI环境） #

如果你的HPC主要用MPI调度，mpi4py可以让你把AQ站点分配到不同的节点上，每个节点独立处理一部分站点，最后汇总结果，完全不需要复杂的分布式框架。

示例代码

from mpi4py import MPI
import numpy as np
import pandas as pd

# 初始化MPI环境
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

# 主进程（rank=0）加载所有数据
if rank == 0:
    # 替换成你的数据加载逻辑
    latitude_aq = np.array(pd.read_csv("aq_data.csv")["latitude"])
    longitude_aq = np.array(pd.read_csv("aq_data.csv")["longitude"])
    aq = pd.read_csv("aq_data.csv").values
    latitude_meo = np.array(pd.read_csv("meo_data.csv")["latitude"])
    longitude_meo = np.array(pd.read_csv("meo_data.csv")["longitude"])
    meo = pd.read_csv("meo_data.csv").values
else:
    latitude_aq = None
    longitude_aq = None
    aq = None
    latitude_meo = None
    longitude_meo = None
    meo = None

# 把MEO数据广播到所有进程（每个进程都需要所有MEO数据）
latitude_meo = comm.bcast(latitude_meo, root=0)
longitude_meo = comm.bcast(longitude_meo, root=0)
meo = comm.bcast(meo, root=0)

# 分配AQ站点到各个进程
n_aq = latitude_aq.shape[0] if rank == 0 else 0
n_aq = comm.bcast(n_aq, root=0)
chunk_size = n_aq // size
start = rank * chunk_size
end = start + chunk_size if rank != size-1 else n_aq

# 每个进程获取自己的AQ分片
lat_aq_chunk = comm.scatter(np.array_split(latitude_aq, size), root=0)
lon_aq_chunk = comm.scatter(np.array_split(longitude_aq, size), root=0)
aq_chunk = comm.scatter(np.array_split(aq, size), root=0)

# 进程独立计算自己的分片
result_chunk = np.zeros((len(lat_aq_chunk), 5))
for i in range(len(lat_aq_chunk)):
    center_lat = lat_aq_chunk[i]
    center_lon = lon_aq_chunk[i]
    aq_time = aq_chunk[i, TIME_AQ]
    humidity, pressure, temperature, wind_speed, wind_direction, count = 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
    
    for j in range(len(latitude_meo)):
        meo_time = meo[j, TIME_MEO]
        if aq_time != meo_time:
            continue
        if is_in_circle(center_lat, center_lon, RADIUS, latitude_meo[j], longitude_meo[j]):
            humidity += meo[j, HUMIDITY_MEO]
            pressure += meo[j, PRESSURE_MEO]
            temperature += meo[j, TEMPERATURE_MEO]
            wind_speed += meo[j, WIND_SPEED_MEO]
            wind_direction += meo[j, WIND_DIRECTION_MEO]
            count += 1.0
    
    if count != 0.0:
        result_chunk[i] = [humidity/count, pressure/count, temperature/count, wind_speed/count, wind_direction/count]

# 汇总所有进程的结果到主进程
final_tmp = comm.gather(result_chunk, root=0)

# 主进程写回结果
if rank == 0:
    final_tmp = np.concatenate(final_tmp, axis=0)
    final.loc[:, [HUMIDITY_FINAL, PRESSURE_FINAL, TEMPERATURE_FINAL, WIND_SPEED_FINAL, WIND_DIRECTION_FINAL]] = final_tmp

HPC上运行方式： #

用MPI命令提交任务，比如：

mpiexec -n 16 python your_script.py

其中16是你要用到的进程数（对应HPC的核心数/节点数）。

选择建议 #

1. 优先选Numba CPU并行版：上手最快，改动最小，不需要了解HPC调度细节，直接在多核心CPU上就能获得几十倍的性能提升，完全满足你的需求。
2. 如果数据太大单节点内存放不下：选Dask，自动处理分布式内存和计算，适配多节点。
3. 如果HPC用MPI调度：选mpi4py，贴合原生环境，性能稳定。

所有这些方案都不需要你写CUDA代码，完美适配你的情况。

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