Python：并行化嵌套的 for 循环

Question

我在 python 中嵌套了 for 循环来创建一个 netCDF 文件。 for 循环采用带有时间、纬度、地段和参数的 pandas 数据帧，并将 netCDF 文件中的信息替换为正确位置和时间的参数。这花费了太长时间，因为 pandas 数据帧有超过 80000 行，而 netCDF 文件有大约 8000 个时间步长。我一直在寻找使用xargs 或multiprocessing，但在第一种情况下使用文件作为输入，在第二种情况下，它产生与我使用的进程一样多的输出。我没有并行处理的经验，所以我的断言可能完全错误。这是我正在使用的代码：

with Dataset(os.path.join('Downloads', inv, 'observations.nc'), 'w') as dset:
    
    dset.createDimension('time_components', 6)
    
    groups = ['obs', 'mix_apri', 'mix_apos', 'mix_background']
    
    for group in groups:
        
        dset.createGroup(group)
        dset[group].createDimension('nt', 8760)
        dset[group].createDimension('nlat', 80)
        dset[group].createDimension('nlon', 100)
        
        times_start = dset[group].createVariable('times_start', 'i4', ('nt', 'time_components'))
        times_end = dset[group].createVariable('times_end', 'i4', ('nt', 'time_components'))
        lats = dset[group].createVariable('lats', 'f4', ('nlat'))
        lons = dset[group].createVariable('lons', 'f4', ('nlon'))
        
        times_start[:,:] = list(emis_apri['biosphere']['times_start'])
        times_end[:,:] = list(emis_apri['biosphere']['times_end'])
        lats[:] = list(emis_apri['biosphere']['lats'])
        lons[:] = list(emis_apri['biosphere']['lons'])
        
    conc_obs = dset['obs'].createVariable('conc', 'f8', ('nt', 'nlat', 'nlon'))
    conc_mix_apri = dset['mix_apri'].createVariable('conc', 'f8', ('nt', 'nlat', 'nlon'))
    conc_mix_apos = dset['mix_apos'].createVariable('conc', 'f8', ('nt', 'nlat', 'nlon'))
    conc_mix_background = dset['mix_background'].createVariable('conc', 'f8', ('nt', 'nlat', 'nlon'))
    
    for i in range(8760):
        conc_obs[i,:,:] = emis_apri['biosphere']['emis'][i][:,:]*0
        
    conc_mix_apri[:,:,:] = list(conc_obs)
    conc_mix_apos[:,:,:] = list(conc_obs)
    conc_mix_background[:,:,:] = list(conc_obs)
    
    db = obsdb(os.path.join('Downloads', inv, 'observations.apos.tar.gz'))
    nsites = db.sites.shape[0]
    for isite, site in enumerate(db.sites.itertuples()):
        dbs = db.observations.loc[db.observations.site == site.Index]
        lat = where((array(emis_apri['biosphere']['lats']) >= list(dbs.lat)[0]-0.25) & (array(emis_apri['biosphere']['lats']) <= list(dbs.lat)[0]+0.25))[0][0]
        lon = where((array(emis_apri['biosphere']['lons']) >= list(dbs.lon)[0]-0.25) & (array(emis_apri['biosphere']['lons']) <= list(dbs.lon)[0]+0.25))[0][0]
        for i in range(len(list(dbs.time))):
            for j in range(len(times_start)):
                if datetime(*times_start[j,:].data) >= Timestamp.to_pydatetime(list(dbs.time)[i]) and datetime(*times_end[j,:].data) >= Timestamp.to_pydatetime(list(dbs.time)[i]):
                    conc_obs[i,lat,lon] = list(dbs.obs)[i]
                    conc_mix_apri[i,lat,lon] = list(dbs.mix_apri)[i]
                    conc_mix_apos[i,lat,lon] = list(dbs.mix_apos)[i]
                    conc_mix_background[i,lat,lon] = list(dbs.mix_background)[i]

来自for isite, site in enumerate(db.sites.itertuples()): 是我需要并行化的代码部分。我非常感谢您对此的任何见解。

Answer 1

将以下内容视为伪代码，因为我无法在没有任何样本的情况下运行任何测试等。我通常使用 mpi4py 并行化我的代码，在你的情况下，你可以在一开始就这样做：

from mpi4py import MPI
comm = MPI.COMM_WORLD
size = comm.Get_size(); # let your program know how many processors you are using
rank = comm.Get_rank() # let the running program know, which processor it is

现在，在代码的开头，让其中一个进程成为所谓的主任务，它可以完成所有任务无法同时完成的所有基本/重要工作。例如，为输出打开/初始化一些文件。因此，在您的代码中，对于这些部分，您可以使用：

if rank==0:
    # do some important stuff
else:
    # do something not important (for example a = 5)
comm.barrier() # this is important to synchronize the  processes

现在，要并行化您的代码，您可以在分布式 db.sites 上执行循环，即将 db.sites.itertuples() 除以您要使用的处理器数量：

allsites = db.sites.itertuples() # all the processor have to know all the sites
sites = allsites[rank::size] # each starts from it's current rank and jumps with the size

for isite, site in enumerate(sites):
    dbs = db.observations.loc[db.observations.site == site.Index]
    lat = where((array(emis_apri['biosphere']['lats']) >= list(dbs.lat)[0]-0.25) & (array(emis_apri['biosphere']['lats']) <= list(dbs.lat)[0]+0.25))[0][0]
    lon = where((array(emis_apri['biosphere']['lons']) >= list(dbs.lon)[0]-0.25) & (array(emis_apri['biosphere']['lons']) <= list(dbs.lon)[0]+0.25))[0][0]
    for i in range(len(list(dbs.time))):
        for j in range(len(times_start)):
            if datetime(*times_start[j,:].data) >= Timestamp.to_pydatetime(list(dbs.time)[i]) and datetime(*times_end[j,:].data) >= Timestamp.to_pydatetime(list(dbs.time)[i]):
                conc_obs[i,lat,lon] = list(dbs.obs)[i]
                conc_mix_apri[i,lat,lon] = list(dbs.mix_apri)[i]
                conc_mix_apos[i,lat,lon] = list(dbs.mix_apos)[i]
                conc_mix_background[i,lat,lon] = list(dbs.mix_background)[i]

comm.barrier() # do not forget to synchronize

尽管如此，在这种情况下，“isite”现在具有基于列表大小的值，您正在放弃。因此，它不是 0...len(allsites)，而是 0...len(allsites ）/尺寸。如果“isite”对于具有从 0 到 len(allsites) 的值很重要，则您必须以某种方式重新计算。也许 isite_global = isite*size+rank 可以得到处理器正在做的实际数字。

那么，到底如何运行代码，我通常是这样做的：

mpiexec -np 10 ipython script_name

在终端上在 10 个处理器上运行代码。

但是，无论如何，最困难的部分是在没有库的特定支持的情况下并行化 I/O 操作。我不确定 netCDF4 是否支持并行 I/O，这意味着如果您的处理器等级为 0...X 为 X 处理器同时打开文件，将某些内容写入文件中的特定位置并关闭文件，然后来自所有处理器写在那里。

因此，最安全的想法是让一个处理器（主）负责输出并交换/收集需要在写入之前从所有子处理器写入的数据。

希望这会有所帮助，祝代码好运！