减少Python dict的RAM消耗

Question

我有一个python脚本，可以处理几个千兆字节的文件。使用下面的代码，我将一些数据存储到一个列表中，该列表存储在字典snp_dict中。RAM消耗巨大。看看我的代码，你能建议一些减少RAM消耗的方法吗？

def extractAF(files_vcf):
    z=0
    snp_dict=dict()
    for infile_name in sorted(files_vcf):
        print '      * ' + infile_name
        ###single files
        vcf_reader = vcf.Reader(open(infile_name, 'r'))
        for record in vcf_reader:
            snp_position='_'.join([record.CHROM, str(record.POS)])
            ref_F = float(record.INFO['DP4'][0])
            ref_R = float(record.INFO['DP4'][1])
            alt_F = float(record.INFO['DP4'][2])
            alt_R = float(record.INFO['DP4'][3])
            AF = (alt_F+alt_R)/(alt_F+alt_R+ref_F+ref_R)
            if not snp_position in snp_dict:
                snp_dict[snp_position]=list((0) for _ in range(len(files_vcf)))
            snp_dict[snp_position][z] = round(AF, 3) #record.INFO['DP4']
        z+=1
    return snp_dict

Answer 1

我最终采用了以下MySQL实现：

for infile_name in sorted(files_vcf):
    print infile_name
    ###single files
    vcf_reader = vcf.Reader(open(infile_name, 'r'))
    for record in vcf_reader:
        snp_position='_'.join([record.CHROM, str(record.POS)])
        ref_F = float(record.INFO['DP4'][0])
        ref_R = float(record.INFO['DP4'][1])
        alt_F = float(record.INFO['DP4'][2])
        alt_R = float(record.INFO['DP4'][3])
        AF = (alt_F+alt_R)/(alt_F+alt_R+ref_F+ref_R)
        if not snp_position in snp_dict:
            sql_insert_table = "INSERT INTO snps VALUES ('" + snp_position + "'," + ",".join(list(('0') for _ in range(len(files_vcf)))) + ")"
            cursor = db1.cursor()
            cursor.execute(sql_insert_table)
            db1.commit()
            snp_dict.append(snp_position)
        sql_update = "UPDATE snps SET " + str(z) + "g=" + str(AF) + " WHERE snp_pos='" + snp_position + "'";
        cursor = db1.cursor()
        cursor.execute(sql_update)
        db1.commit()
    z+=1
return snp_dict

Answer 2

对于这类事情，您最好使用另一种数据结构。熊猫DataFrame在你的情况下会很好用。

最简单的解决方案是使用现有的库，而不是编写自己的解析器。vcfnp可以将vcf文件读取为易于转换为pandas DataFrame的格式。像这样的东西应该起作用：

import pandas as pd
    def extractAF(files_vcf):
    dfs = []
    for fname in sorted(files_vcf):
        vars = vcfnp.variants(fname, fields=['CHROM', 'POS', 'DP4'])
        snp_pos = np.char.add(np.char.add(vars.CHROM, '_'), record.POS.astype('S'))
        dp4 = vars.DP4.astype('float')
        AF = dp4[2:].sum(axis=0)/dp4.sum(axis=0)
        dfs.append(pd.DataFrame(AF, index=snp_pos, columns=[fname]).T)
    return pd.concat(dfs).fillna(0.0)

如果你绝对必须使用PyVCF，它会更慢，但希望这至少会比你现有的实现更快，并且应该产生与上面代码相同的结果：

def extractAF(files_vcf):
    files_vcf = sorted(files_vcf)
    dfs = []
    for fname in files_vcf:
        print '      * ' + fname
        vcf_reader = vcf.Reader(open(fname, 'r'))
        vars = ((rec.CHROM, rec.POS) + tuple(rec.INFO['DP4']) for rec in vcf_reader)
        df = pd.DataFrame(vars, columns=['CHROMS', 'POS', 'ref_F', 'ref_R', 'alt_F', 'alt_R'])
        df['snp_position'] = df['CHROMS'] + '_' + df['POS'].astype('S')
        df_alt = df.loc[:, ('alt_F', 'alt_R')]
        df_dp4 = df.loc[:, ('alt_F', 'alt_R', 'ref_F', 'ref_R')]
        df[fname] = df_alt.sum(axis=1)/df_dp4.sum(axis=1)
        df = df.set_index('snp_position', drop=True).loc[:, fname:fname].T
        dfs.append(df)
    return pd.concat(dfs).fillna(0.0)

现在假设你想读取一个特定的snp_position，比如说包含在变量snp_pos中，它可能在那里，也可能不在那里（根据你的评论），你实际上不需要更改任何内容：

all_vcf = extractAF(files_vcf)
if snp_pos in all_vcf:
     linea_di_AF = all_vcf[snp_pos]

不过，结果会略有不同。它将是熊猫Series，它就像一个数组，但也可以像字典一样访问：

all_vcf = extractAF(files_vcf)
if snp_pos in all_vcf:
     linea_di_AF = all_vcf[snp_pos]
     f_di_AF = linea_di_AF[files_vcf[0]]

这允许您直接访问特定的文件/snp_pos对：

all_vcf = extractAF(files_vcf)
if snp_pos in all_vcf:
     f_di_AF = linea_di_AF[snp_pos][files_vcf[0]]

或者，更好的是：

all_vcf = extractAF(files_vcf)
if snp_pos in all_vcf:
     f_di_AF = linea_di_AF.loc[files_vcf[0], snp_pos]

或者，您可以获取给定文件的所有snp_pos值：

all_vcf = extractAF(files_vcf)
fpos = linea_di_AF.loc[fname]