parsing.py

"""
Utils for reading flat files that are loaded into database
"""
import copy
import re
import traceback
from tqdm import tqdm
from utils import *

POPS = {
    'AFR': 'African',
    'AMR': 'Latino',
    'ASJ': 'Ashkenazi Jewish',
    'EAS': 'East Asian',
    'FIN': 'European (Finnish)',
    'NFE': 'European (Non-Finnish)',
    'SAS': 'South Asian',
    'OTH': 'Other'
}

def par(chrom, pos):
    return (chrom == 'X' and (60001 <= pos <= 2699520) or (154931044 <= pos <= 155260560))

def get_base_coverage_from_file(base_coverage_file):
    """
    Read a base coverage file and return iter of dicts that look like:
    {
        'xpos': 1e9+1,
        'mean': 0.0,
        'median': 0.0,
        '1': 0.0,
        '5': 0.0,
        '10': 0.0,
        '15': 0.0,
        '20': 0.0,
        '25': 0.0,
        '30': 0.0,
        '50': 0.0,
        '100': 0.0,
    }
    """

    float_header_fields = ['mean', 'median', '1', '5', '10', '15', '20', '25', '30', '50', '100']
    for line in tqdm(base_coverage_file, unit=" coverage"):
        if line.startswith('#'):
            continue
        fields = line.strip('\n').split('\t')
        d = {
            'xpos': get_xpos(fields[0], int(fields[1])),
            'pos': int(fields[1]),
        }
        for i, k in enumerate(float_header_fields):
            try:
                d[k] = float(fields[i+2])
            except IndexError:
                print 'Index error at file:', base_coverage_file, 'index:', i, 'field:', k, 'line:', line

        yield d


def get_variants_from_sites_vcf(sites_vcf):
    """
    Parse exac sites VCF file and return iter of variant dicts
    sites_vcf is a file (gzipped), not file path
    """
    vep_field_names = None

    # hard-code some VCF header fields until vcf gets updated.
    line = '##INFO=<ID=DP_HIST,Number=R,Type=String,Description="Histogram for DP; Mids: 2.5|7.5|12.5|17.5|22.5|27.5|32.5|37.5|42.5|47.5|52.5|57.5|62.5|67.5|72.5|77.5|82.5|87.5|92.5|97.5">'
    dp_mids = map(float, line.split('Mids: ')[-1].strip('">').split('|'))
    line = '##INFO=<ID=GQ_HIST,Number=R,Type=String,Description="Histogram for GQ; Mids: 2.5|7.5|12.5|17.5|22.5|27.5|32.5|37.5|42.5|47.5|52.5|57.5|62.5|67.5|72.5|77.5|82.5|87.5|92.5|97.5">'
    gq_mids = map(float, line.split('Mids: ')[-1].strip('">').split('|'))
    line = '##INFO=<ID=GQ_HIST,Number=R,Type=String,Description="Histogram for GQ; Mids: 2.5|7.5|12.5|17.5|22.5|27.5|32.5|37.5|42.5|47.5|52.5|57.5|62.5|67.5|72.5|77.5|82.5|87.5|92.5|97.5">'
    ab_mids = map(float, line.split('Mids: ')[-1].strip('">').split('|'))
    line = '##INFO=<ID=CSQ,Number=.,Type=String,Description="Consequence annotations from Ensembl VEP. Format: Allele|Consequence|IMPACT|SYMBOL|Gene|Feature_type|Feature|BIOTYPE|EXON|INTRON|HGVSc|HGVSp|cDNA_position|CDS_position|Protein_position|Amino_acids|Codons|Existing_variation|ALLELE_NUM|DISTANCE|STRAND|FLAGS|VARIANT_CLASS|MINIMISED|SYMBOL_SOURCE|HGNC_ID|CANONICAL|TSL|APPRIS|CCDS|ENSP|SWISSPROT|TREMBL|UNIPARC|GENE_PHENO|SIFT|PolyPhen|DOMAINS|HGVS_OFFSET|GMAF|AFR_MAF|AMR_MAF|EAS_MAF|EUR_MAF|SAS_MAF|AA_MAF|EA_MAF|ExAC_MAF|ExAC_Adj_MAF|ExAC_AFR_MAF|ExAC_AMR_MAF|ExAC_EAS_MAF|ExAC_FIN_MAF|ExAC_NFE_MAF|ExAC_OTH_MAF|ExAC_SAS_MAF|CLIN_SIG|SOMATIC|PHENO|PUBMED|MOTIF_NAME|MOTIF_POS|HIGH_INF_POS|MOTIF_SCORE_CHANGE|LoF|LoF_filter|LoF_flags|LoF_info">'

    vep_field_names = line.split('Format: ')[-1].strip('">').split('|')

    for line in sites_vcf:
        try:
            line = line.strip('\n')
            # if line.startswith('##INFO=<ID=CSQ'):
            #     vep_field_names = line.split('Format: ')[-1].strip('">').split('|')
                #print(vep_field_names)

            #if line.startswith('##INFO=<ID=DP_HIST'):
            #dp_mids = map(float, line.split('Mids: ')[-1].strip('">').split('|'))
            #if line.startswith('##INFO=<ID=GQ_HIST'):
            #    gq_mids = map(float, line.split('Mids: ')[-1].strip('">').split('|'))
            if line.startswith('#'):
                continue

            # If we get here, it's a variant line
            if vep_field_names is None:
                raise Exception("VEP_field_names is None. Make sure VCF header is present.")

            # This elegant parsing code below is copied from https://github.com/konradjk/loftee
            fields = line.split('\t')
            info_field = dict([(x.split('=', 1)) if '=' in x else (x, x) for x in re.split(';(?=\w)', fields[7])])
            consequence_array = info_field['CSQ'].split(',') if 'CSQ' in info_field else []
            annotations = [dict(zip(vep_field_names, x.split('|'))) for x in consequence_array if len(vep_field_names) == len(x.split('|'))]
            coding_annotations = [ann for ann in annotations if ann['Feature'].startswith('ENST')]

            #if not coding_annotations:
            #    continue

            alt_alleles = fields[4].split(',')

            # different variant for each alt allele
            # print ' '
            # print 'number of alternate alleles:', len(alt_alleles)
            # print '--------------------'
            for i, alt_allele in enumerate(alt_alleles):
                if alt_allele == "*":
                    continue

                vep_annotations = [ann for ann in coding_annotations if int(ann['ALLELE_NUM']) == i + 1]

                # Variant is just a dict
                # Make a copy of the info_field dict - so all the original data remains
                # Add some new keys that are allele-specific
                pos, ref, alt = get_minimal_representation(fields[1], fields[3], alt_allele)

                variant = {}
                variant['chrom'] = fields[0]
                variant['pos'] = pos
                variant['rsid'] = fields[2]
                variant['xpos'] = get_xpos(variant['chrom'], variant['pos'])
                variant['ref'] = ref
                variant['alt'] = alt
                variant['xstart'] = variant['xpos']
                variant['xstop'] = variant['xpos'] + len(variant['alt']) - len(variant['ref'])
                variant['variant_id'] = '{}-{}-{}-{}'.format(variant['chrom'], variant['pos'], variant['ref'], variant['alt'])
                variant['orig_alt_alleles'] = [
                    '{}-{}-{}-{}'.format(variant['chrom'], *get_minimal_representation(fields[1], fields[3], x))
                    for x in alt_alleles
                ]
                variant['site_quality'] = float(fields[5])

                filter_status = info_field['AS_FilterStatus'].split(',')[i]
                filter_field_list = fields[6].split(';')
                fail_filters = ('LCR', 'SEGDUP', 'InbreedingCoeff')
                failed = set(filter_field_list).intersection(fail_filters)

                if filter_status == 'PASS' and len(failed) > 0:
                    variant['filter'] = '|'.join(failed)
                elif filter_status is not 'PASS' and len(failed) > 0:
                    variant['filter'] = filter_status + '|' + '|'.join(failed)
                else:
                    variant['filter'] = filter_status

                if 'lcr' in info_field:
                    variant['lcr'] = True
                if 'segdup' in info_field:
                    variant['segdup'] = True

                variant['vep_annotations'] = vep_annotations

                variant['allele_count'] = int(info_field['AC'].split(',')[i])
                if not variant['allele_count'] and variant['filter'] == 'PASS': variant['filter'] = 'AC0' # Temporary filter
                variant['allele_num'] = int(info_field['AN'])

                if variant['allele_num'] > 0:
                    variant['allele_freq'] = variant['allele_count']/float(info_field['AN'])
                else:
                    variant['allele_freq'] = None

                variant['pop_acs'] = dict([(POPS[x], int(info_field['AC_%s' % x].split(',')[i]) if ('AC_%s' % x) in info_field else 0) for x in POPS])
                variant['pop_ans'] = dict([(POPS[x], int(info_field.get('AN_%s' % x, 0))) for x in POPS])
                variant['pop_homs'] = dict([(POPS[x], int(info_field['Hom_%s' % x].split(',')[i] if ('Hom_%s' % x) in info_field else 0)) for x in POPS])
                if variant['chrom'] not in ('X', 'Y'):
                    if not info_field['AC_Male'].split(',')[i] == ".":
                        variant['ac_male'] = int(info_field['AC_Male'].split(',')[i])
                    if not info_field['AC_Female'].split(',')[i] == ".":
                        variant['ac_female'] = int(info_field['AC_Female'].split(',')[i])
                    variant['an_male'] = int(info_field['AN_Male'])
                    variant['an_female'] = int(info_field['AN_Female'])
                variant['hom_count'] = sum(variant['pop_homs'].values())
                if variant['chrom'] == 'X':
                    if not par(variant['chrom'], variant['pos']):
                        variant['pop_hemis'] = dict([(POPS[x], int(info_field['Hemi_%s' % x].split(',')[i] if ('Hemi_%s' % x) in info_field else 0)) for x in POPS])
                        variant['hemi_count'] = sum(variant['pop_hemis'].values())
                if variant['chrom'] == 'Y':
                    variant['pop_hemis'] = variant['pop_acs']
                variant['quality_metrics'] = dict([(x, info_field[x]) for x in METRICS if x in info_field])
                for metric in AS_METRICS:
                    if metric in info_field and info_field[metric] != '.':
                        variant['quality_metrics'][metric] = info_field[metric].split(',')[i]
                variant['genes'] = list({annotation['Gene'] for annotation in vep_annotations})
                variant['transcripts'] = list({annotation['Feature'] for annotation in vep_annotations})
                if 'DP_HIST_ALL' in info_field:
                    # hists_all = [info_field['DP_HIST'].split(',')[0], info_field['DP_HIST'].split(',')[i+1]]
                    hists_all = [info_field['DP_HIST_ALL'], info_field['DP_HIST_ALT'].split(',')[i]]
                    # print hists_all
                    variant['genotype_depths'] = [zip(dp_mids, map(int, x.split('|'))) for x in hists_all]
                if 'GQ_HIST_ALL' in info_field:
                    # hists_all = [info_field['GQ_HIST'].split(',')[0], info_field['GQ_HIST'].split(',')[i+1]]
                    hists_all = [info_field['GQ_HIST_ALL'], info_field['GQ_HIST_ALT'].split(',')[i]]
                    variant['genotype_qualities'] = [zip(gq_mids, map(int, x.split('|'))) for x in hists_all]
                if 'AB_HIST_ALL' in info_field:
                    # hists_all = [info_field['GQ_HIST'].split(',')[0], info_field['GQ_HIST'].split(',')[i+1]]
                    hists_all = [info_field['AB_HIST_ALL'], info_field['AB_HIST_ALT'].split(',')[i]]
                    variant['allele_balance'] = [zip(ab_mids, map(int, x.split('|'))) for x in hists_all]

                yield variant
        except Exception:
            print("Error parsing vcf line: " + line)
            traceback.print_exc()
            break


def get_mnp_data(mnp_file):
    header = mnp_file.readline().strip().split('\t')
    for line in tqdm(mnp_file, unit=" mnps"):
        data = dict(zip(header, line.split('\t')))
        if any(map(lambda x: x == 'True', data['QUESTIONABLE_PHASING'])): continue
        chroms = data['CHROM'].split(',')
        chrom = chroms[0]
        sites = data['SITES'].split(',')
        refs = data['REF'].split(',')
        alts = data['ALT'].split(',')
        for i, site in enumerate(sites):
            all_sites = zip(chroms, sites, refs, alts)
            all_sites.remove(all_sites[i])
            mnp = {}
            mnp['xpos'] = get_xpos(chrom, site)
            mnp['ref'] = refs[i]
            mnp['alt'] = alts[i]
            mnp['site2'] = '-'.join(all_sites[0])
            if len(all_sites) > 1:
                mnp['site3'] = all_sites[1]
            mnp['combined_codon_change'] = data['COMBINED_CODON_CHANGE']
            mnp['category'] = data['CATEGORY']
            mnp['number_samples'] = data['NSAMPS']
            yield mnp


def get_constraint_information(constraint_file):
    _, _, _, header = constraint_file.readline().strip().split(None, 3)
    header = header.split()
    for line in constraint_file:
        transcript, gene, chrom, info = line.strip().split(None, 3)
        transcript_info = dict(zip(header, map(float, info.split())))
        transcript_info['transcript'] = transcript.split('.')[0]
        yield transcript_info


def get_canonical_transcripts(canonical_transcript_file):
    for line in canonical_transcript_file:
        gene, transcript = line.strip().split()
        yield gene, transcript


def get_omim_associations(omim_file):
    for line in omim_file:
        fields = line.strip().split('\t')
        if len(fields) == 4:
            yield fields
        else:
            yield None


def get_genes_from_gencode_gtf(gtf_file):
    """
    Parse gencode GTF file;
    Returns iter of gene dicts
    """
    for line in tqdm(gtf_file, unit=" genes"):
        if line.startswith('#'):
            continue
        fields = line.strip('\n').split('\t')

        if fields[2] != 'gene':
            continue

        chrom = fields[0][3:]
        start = int(fields[3]) + 1  # bed files are 0-indexed
        stop = int(fields[4]) + 1
        info = dict(x.strip().split() for x in fields[8].split(';') if x != '')
        info = {k: v.strip('"') for k, v in info.items()}
        gene_id = info['gene_id'].split('.')[0]

        gene = {
            'gene_id': gene_id,
            'gene_name': info['gene_name'],
            'gene_name_upper': info['gene_name'].upper(),
            'chrom': chrom,
            'start': start,
            'stop': stop,
            'strand': fields[6],
            'xstart': get_xpos(chrom, start),
            'xstop': get_xpos(chrom, stop),
        }
        yield gene


def get_transcripts_from_gencode_gtf(gtf_file):
    """
    Parse gencode GTF file;
    Returns iter of transcript dicts
    """
    for line in tqdm(gtf_file, unit=" transcripts"):
        if line.startswith('#'):
            continue
        fields = line.strip('\n').split('\t')

        if fields[2] != 'transcript':
            continue

        chrom = fields[0][3:]
        start = int(fields[3]) + 1  # bed files are 0-indexed
        stop = int(fields[4]) + 1
        info = dict(x.strip().split() for x in fields[8].split(';') if x != '')
        info = {k: v.strip('"') for k, v in info.items()}
        transcript_id = info['transcript_id'].split('.')[0]
        gene_id = info['gene_id'].split('.')[0]

        gene = {
            'transcript_id': transcript_id,
            'gene_id': gene_id,
            'chrom': chrom,
            'start': start,
            'stop': stop,
            'strand': fields[6],
            'xstart': get_xpos(chrom, start),
            'xstop': get_xpos(chrom, stop),
        }
        yield gene


def get_exons_from_gencode_gtf(gtf_file):
    """
    Parse gencode GTF file;
    Returns iter of transcript dicts
    """
    for line in tqdm(gtf_file, unit=" exons"):
        if line.startswith('#'):
            continue
        fields = line.strip('\n').split('\t')

        if fields[2] not in ['exon', 'CDS', 'UTR']:
            continue

        chrom = fields[0][3:]
        feature_type = fields[2]
        start = int(fields[3]) + 1  # bed files are 0-indexed
        stop = int(fields[4]) + 1
        info = dict(x.strip().split() for x in fields[8].split(';') if x != '')
        info = {k: v.strip('"') for k, v in info.items()}
        transcript_id = info['transcript_id'].split('.')[0]
        gene_id = info['gene_id'].split('.')[0]

        exon = {
            'feature_type': feature_type,
            'transcript_id': transcript_id,
            'gene_id': gene_id,
            'chrom': chrom,
            'start': start,
            'stop': stop,
            'strand': fields[6],
            'xstart': get_xpos(chrom, start),
            'xstop': get_xpos(chrom, stop),
        }
        yield exon


def get_cnvs_from_txt(cnv_txt_file):
    """
    Parse gencode txt file;
    Returns iter of gene dicts
    """
    header = cnv_txt_file.next() # gets rid of the header
    #print header
    for line in cnv_txt_file:

        fields = line.rsplit()
        transcript = fields[0]
        gene = fields[1]
        chrom = fields[2]
        start = int(fields[3])
        stop = int(fields[4])
        del0 = int(fields[5])
        del60 = int(fields[6])
        dup0 = int(fields[7])
        dup60 = int(fields[8])
        delpop0 = fields[9]
        delpop60 = fields[10]
        duppop0 = fields[11]
        duppop60 = fields[12]


        #find gene from DB.genes, get ID
        #find exon of that gene that this CNV referes to from db.exons, get ID
        #add the object reference to the cnv dict.
        cnv = {
            'transcript': transcript,
            'gene': gene,
            'chrom': chrom,
            'start': start,
            'stop': stop,
            'del0': del0,
            'dup0': dup0,
            'dup60': dup60,
            'del60' : del60,
            'delpop0' : delpop0,
            'delpop60' : delpop60,
            'duppop0' : duppop0,
            'duppop60' : duppop60,
            'xstart': get_xpos(chrom, start),
            'xstop': get_xpos(chrom, stop),
        }
        yield cnv


def get_cnvs_per_gene(cnv_gene_file):
    header = cnv_gene_file.next() # gets rid of the header
    for line in cnv_gene_file:

        fields = line.rsplit()
        gene = fields[0]
        symbol = fields[1]
        del0 = int(fields[2])
        dup0 = int(fields[3])
        cnv0 = int(fields[4])
        del60 = int(fields[5])
        dup60 = int(fields[6])
        cnv60 = int(fields[7])
        del_score = float(fields[8])
        dup_score = float(fields[9])
        cnv_score = float(fields[10])
        rank = int(fields[11])

        cnv_gene = {
            'gene': gene,
            'symbol': symbol,
            'del0': del0,
            'dup0': dup0,
            'cnv0': cnv0,
            'del60': del60,
            'dup60': dup60,
            'cnv60' : cnv60,
            'del_score': del_score,
            'dup_score': dup_score,
            'cnv_score': cnv_score,
            'rank': rank,
            }
        yield cnv_gene


def get_dbnsfp_info(dbnsfp_file):
    """
    Parse dbNSFP_gene file;
    Returns iter of transcript dicts
    """
    header = dbnsfp_file.next().split('\t')
    fields = dict(zip(header, range(len(header))))
    for line in dbnsfp_file:
        line = line.split('\t')
        other_names = line[fields["Gene_old_names"]].split(';') if line[fields["Gene_old_names"]] != '.' else []
        if line[fields["Gene_other_names"]] != '.':
            other_names.extend(line[fields["Gene_other_names"]].split(';'))
        gene_info = {
            'gene_name': line[fields["Gene_name"]],
            'ensembl_gene': line[fields["Ensembl_gene"]],
            'gene_full_name': line[fields["Gene_full_name"]],
            'gene_other_names': other_names
        }
        yield gene_info


def get_snp_from_dbsnp_file(dbsnp_file):
    for line in dbsnp_file:
        fields = line.split('\t')
        if len(fields) < 3: continue
        rsid = int(fields[0])
        chrom = fields[1].rstrip('T')
        if chrom == 'PAR': continue
        start = int(fields[2]) + 1
        snp = {
            'xpos': get_xpos(chrom, start),
            'rsid': rsid
        }
        yield snp