减少Python dict的RAM消耗

时间:2015-04-09 15:45:54

标签: python memory memory-management dictionary

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

2 个答案:

答案 0 :(得分:0)

对于这类事情,你可能最好使用其他数据结构。大熊猫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]

但结果会略有不同。它将是一个pandas 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]

答案 1 :(得分:0)

我最终采用了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