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

时间:2017-03-04 22:50:23

标签: python python-2.7 dictionary multiprocessing

我有一个名为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_lstHow 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

但是,我不太确定defif __name__ == '__main__':下应该有什么。我也不太确定我应该将哪些参数传递给函数。当我尝试Douglas的第一个代码时,我收到了一个错误,没有传递任何参数如下:

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

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

过去几天我一直在努力工作,但我没有成功地产生我想要的输出。如果任何人都可以建议一个可以快速运行的替代代码,或者如果有人可以建议如何在多个CPU上运行此代码,那将是非常棒的。我们非常感谢您解决此问题的任何帮助。

1 个答案:

答案 0 :(得分: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)))