python - 如何使用多处理模块迭代列表并将其与字典中的键匹配？

Question

我有一个名为 master_lst 的列表，使用以下代码从 CSV 文件创建

infile= open(sys.argv[1], "r")
lines = infile.readlines()[1:]
master_lst = ["read"]
for line in lines:
 line= line.strip().split(',')
 fourth_field = line [3]
 master_lst.append(fourth_field)

此主列表具有唯一的序列集。现在我必须循环 30 个折叠的 FASTA 文件来计算主列表中每个序列的出现次数。这30个文件的文件格式如下:

>AAAAAAAAAAAAAAA
7451
>AAAAAAAAAAAAAAAA
4133
>AAAAAAAAAAAAAAAAA
2783

为了计算出现次数，我循环遍历了 30 个文件中的每一个，并创建了一个字典，其中序列作为键，出现次数作为值。然后，我迭代了 master_lst 的每个元素，并将其与上一步创建的字典中的键进行匹配。如果存在匹配项，我会将键的值附加到新列表 (ind_lst)。如果不是，我将 0 附加到 ind_lst 中。其代码如下:

for file in files:
 ind_lst = []
 if file.endswith('.fa'):
  first = file.split(".")
  first_field = first [0]
  ind_lst.append(first_field)
  fasta= open(file)
  individual_dict= {}
  for line in fasta:
   line= line.strip()
   if line == '':
    continue
   if line.startswith('>'):
    header = line.lstrip('>')
    individual_dict[header]= ''
   else:
    individual_dict[header] += line
  for key in master_lst[1:]:
   a = 0
   if key in individual_dict.keys():
     a = individual_dict[key]
   else:
     a = 0
   ind_lst.append(a)

然后我使用此处解释的代码将 master_lst 写入 CSV 文件和 ind_lst:How to append a new list to an existing CSV file?

最终输出应如下所示:

Read                           file1     file2 so on until file 30
AAAAAAAAAAAAAAA                 7451      4456
AAAAAAAAAAAAAAAA                4133      3624
AAAAAAAAAAAAAAAAA               2783      7012

当我使用较小的 master_lst 时，此代码工作得非常好。但是，当 master_lst 的大小增加时，执行时间会增加太多。我现在正在使用的 master_lst 有 35,718,501 个序列(元素)。当我对 50 个序列进行子集化并运行代码时，脚本需要 2 小时才能执行。因此对于 35,718,501 个序列来说，需要很长时间才能完成。

现在我不知道如何加快脚本速度。我不太确定是否可以对该脚本进行一些改进以使其在更短的时间内执行。我正在具有 16 个 CPU 核心的 Linux 服务器上运行我的脚本。当我使用 top 命令时，我可以看到该脚本仅使用一个 CPU。但我不是 python 专家，我不知道如何使用多处理模块使其在所有可用的 CPU 核心上运行。我检查了这个网页:Learning Python's Multiprocessing Module。

但是，我不太确定 def 和 if __name__ == '__main__': 下应该包含什么。我也不太确定应该将哪些参数传递给该函数。当我尝试 Douglas 的第一个代码时，没有传递任何参数，出现错误，如下所示:

  File "/usr/lib/python2.7/multiprocessing/process.py", line 114, in run

self._target(*self._args, **self._kwargs)

过去几天我一直在研究这个问题，但没有成功地产生我想要的输出。如果有人可以建议一个可以快速运行的替代代码，或者如果有人可以建议如何在多个 CPU 上运行此代码，那就太棒了。任何解决此问题的帮助将不胜感激。

Answer 1

这是一个多处理版本。它使用的方法与您在代码中使用的方法略有不同，无需创建 ind_lst。

区别的本质在于它首先生成所需数据的转置，然后将其转置为所需结果。

换句话说，不是直接创建它:

Read,file1,file2
AAAAAAAAAAAAAAA,7451,4456
AAAAAAAAAAAAAAAA,4133,3624
AAAAAAAAAAAAAAAAA,2783,7012

它首先产生:

Read,AAAAAAAAAAAAAAA,AAAAAAAAAAAAAAAA,AAAAAAAAAAAAAAAAA 
file1,7451,4133,2783
file2,4456,3624,7012

...然后使用内置的 zip() 函数转置它以获得所需的格式。

除了不需要创建ind_lst之外，它还允许为每个文件创建一行数据，而不是其中的一列(这更容易，并且可以更有效地完成) .

代码如下:

from __future__ import print_function

import csv
from functools import partial
from glob import glob
from itertools import izip  # Python 2
import operator
import os
from multiprocessing import cpu_count, Pool, Queue
import sys

def get_master_list(filename):
    with open(filename, "rb") as csvfile:
        reader = csv.reader(csvfile)
        next(reader)  # ignore first row
        sequence_getter = operator.itemgetter(3)  # retrieves fourth column of each row
        return map(sequence_getter, reader)

def process_fa_file(master_list, filename):
    fa_dict = {}
    with open(filename) as fa_file:
        for line in fa_file:
            if line and line[0] != '>':
                fa_dict[sequence] = int(line)
            elif line:
                sequence = line[1:-1]

    get = fa_dict.get  # local var to expedite access
    basename = os.path.basename(os.path.splitext(filename)[0])
    return [basename] + [get(key, 0) for key in master_list]

def process_fa_files(master_list, filenames):
    pool = Pool(processes=4)  # "processes" is the number of worker processes to
                              # use. If processes is None then the number returned
                              # by cpu_count() is used.
    # Only one argument can be passed to the target function using Pool.map(),
    # so create a partial to pass first argument, which doesn't vary.
    results = pool.map(partial(process_fa_file, master_list), filenames)
    header_row = ['Read'] + master_list
    return [header_row] + results

if __name__ == '__main__':
    master_list = get_master_list('master_list.csv')

    fa_files_dir = '.'  # current directory
    filenames = glob(os.path.join(fa_files_dir, '*.fa'))

    data = process_fa_files(master_list, filenames)

    rows = zip(*data)  # transpose
    with open('output.csv', 'wb') as outfile:
        writer = csv.writer(outfile)
        writer.writerows(rows)

    # show data written to file
    for row in rows:
        print(','.join(map(str, row)))