Использовать многопроцессорность для цикла for, Python

Question

У меня есть цикл for, который использует некоторые двоичные условия и, наконец, записывает файл соответствующим образом. Проблема у меня есть, условия верны для многих файлов (иногда требуется около 1000 файлов). Поэтому писать их занимает много времени (около 10 минут). Я знаю, что могу как-то использовать Python multiprocessing и использовать некоторые из ядер.

Это код, который работает, но использует только одно ядро.

for i,n in enumerate(halo_param.strip()): 
    mask = var1['halo_id'] == n 
    newtbdata = tbdata1[mask] 
    hdu = pyfits.BinTableHDU(newtbdata) 
    hdu.writeto(('/home/Documments/file_{0}.fits').format(i)) 

я наткнулся, что это может быть сделано с помощью Pool из multiprocessing.

if __name__ == '__main__': pool = Pool(processes=4)

Я хотел бы знать, как сделать это и использовать по крайней мере 4 из моих ядер.

Answer 1

Реструктурируйте корпус петли в качестве функции и используйте функцию Pool.map с функцией.

def work(arg): 
    i, n = arg 
    mask = var1['halo_id'] == n 
    newtbdata = tbdata1[mask] 
    hdu = pyfits.BinTableHDU(newtbdata) 
    hdu.writeto(('/home/Documments/file_{0}.fits').format(i)) 

if __name__ == '__main__': 
    pool = Pool(processes=4) 
    pool.map(work, enumerate(halo_param.strip())) 
    pool.close() 
    pool.join()