merge_and_convert_to_html_coloc2.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Nov 26 17:43:52 2019

@author: naim

This script is similar to merge_and_convert_to_html.py, but adds 
the additional columns required for coloc2 secondary datasets

Description file should contain the following information per row:
    1. File name to merge
    2. Description of the dataset
    3. Chromosome column name
    4. Basepair position column name
    5. SNP ID column name
    6. P-value column name
    7. Beta column name
    8. Standard error column name
    9. Number of samples column name
    10. Alternate or A1 allele column name
    11. Reference or A2 allele column name
    12. MAF column name
    13. ProbeID column name
"""

import argparse
import numpy as np
import pandas as pd
import gzip
from datetime import datetime
from tqdm import tqdm
import sys
#from bs4 import BeautifulSoup as bs

# Default column names:
CHROM = 'CHR'
BP = 'POS'
SNP = 'SNPID'
P = 'PVAL'
BETA = 'BETA'
SE = 'SE'
N = 'N'
ALT = 'A1'
REF = 'A2'
MAF = 'MAF'
ProbeID = 'ProbeID'

genomicWindowLimit = 2e6

##############################
## Helper functions
##############################
def parseRegionText(regiontext):
    regiontext = regiontext.strip().replace(' ','')
    chrom = regiontext.split(':')[0].replace('chr','').replace('Chr','')
    pos = regiontext.split(':')[1]
    startbp = pos.split('-')[0].replace(',','')
    endbp = pos.split('-')[1].replace(',','')
    #chromLengths = pd.read_csv(os.path.join(MYDIR, 'data/hg19_chrom_lengths.txt'), sep="\t", encoding='utf-8')
    #chromLengths.set_index('sequence',inplace=True)
    if chrom == 'X':
        chrom = 23
        #maxChromLength = chromLengths.loc['chrX', 'length']
        try:
            startbp = int(startbp)
            endbp = int(endbp)
        except:
            raise Exception("Invalid coordinates input")
    else:
        try:
            chrom = int(chrom)
            #maxChromLength = chromLengths.loc['chr'+str(chrom), 'length']
            startbp = int(startbp)
            endbp = int(endbp)
        except:
            raise Exception("Invalid coordinates input")
    if chrom < 1 or chrom > 23:
        raise Exception('Chromosome input must be between 1 and 23')
    elif startbp > endbp:
        raise Exception('Starting chromosome basepair position is greater than ending basepair position')
#    elif startbp > maxChromLength or endbp > maxChromLength:
#        raise Exception('Start or end coordinates are out of range')
    else:
        return chrom, startbp, endbp


def checkColnames(filename, cols):
    '''
    Checks if the 'filename' has the column names in cols (a list of column names); 
    returns False if any column is absent
    '''
    num_lines_to_check = 1000
    num_lines_checked = 0
    if filename.endswith('gz'):
        with gzip.open(filename,'rb') as f:
            nextline = f.readline().decode('utf-8').strip()
            num_lines_checked += 1
            while nextline[0:2] == "##" and num_lines_checked <= num_lines_to_check:
                nextline = f.readline().decode('utf-8').strip()
                num_lines_checked += 1
            headerline = nextline.split('\t')
            if all(col in headerline for col in cols):
                return True
            else:
                return False
    else:
        with open(filename,'r') as f:
            nextline = f.readline().strip()
            num_lines_checked += 1
            while nextline[0:2] == "##" and num_lines_checked <= num_lines_to_check:
                nextline = f.readline().strip()
                num_lines_checked += 1
            headerline = nextline.split('\t')
            if all(col in headerline for col in cols):
                return True
            else:
                return False
        
def getNumHeaderLines(file, num_lines_to_check = 1000):
    num_header_lines = 0
    num_lines_checked = 0
    if file.endswith('gz'):
        with gzip.open(file, 'rb') as f:
            nextline = f.readline().decode('utf-8')
            num_lines_checked += 1
            while nextline[0:2] == "##" and num_lines_checked <= num_lines_to_check:
                num_header_lines += 1
                nextline = f.readline().decode('utf-8')
                num_lines_checked += 1
    else:
        with open(file, 'r') as f:
            nextline = f.readline()
            num_lines_checked += 1
            while nextline[0:2] == "##" and num_lines_checked <= num_lines_to_check:
                num_header_lines += 1
                nextline = f.readline()
                num_lines_checked += 1
    return num_header_lines

class Logger(object):
    def __init__(self, outfilename):
        self.terminal = sys.stdout
        self.log = open(outfilename, "a")

    def write(self, message):
        self.terminal.write(message)
        self.log.write(message)  

    def flush(self):
        pass    


##############################
## MAIN
##############################
if __name__=='__main__':
    parser = argparse.ArgumentParser(description='Merge several datasets together into HTML tables separated by <h3> title tags')
    parser.add_argument('filelist_filename', 
                        help='Filename containing the list of files to be merged together.\n \
                        The second column (tab-delimited) may contain descriptions of the datasets.\n \
                        The remaining columns should contain the following information per row:\n \
                            1. File name to merge\n \
                            2. Description of the dataset\n \
                            3. Chromosome column name\n \
                            4. Basepair position column name\n \
                            5. SNP ID column name\n \
                            6. P-value column name\n \
                            7. Beta column name\n \
                            8. Standard error column name\n \
                            9. Number of samples column name\n \
                            10. Alternate or A1 allele column name\n \
                            11. Reference or A2 allele column name\n \
                            12. MAF column name\n \
                            13. ProbeID column name')
    parser.add_argument('coordinates', help="The region coordinates to subset from each file (e.g. 1:500,000-600,000")
    parser.add_argument('outfilename', help="Desired output filename for the merged file")
    args = parser.parse_args()  

    filelist_filename = str(args.filelist_filename)
    chrom, startbp, endbp = parseRegionText(args.coordinates)
    outfilename = str(args.outfilename.replace('.html','') + '.html')
    logfilename = str(outfilename.replace('.html','') + '.log')

#    filelist_filename = "files_test.txt"
#    outfilename = "slc9a3_gwas_mqtl.html"
#    coordinates = 'chr5:384,664-612803'
#    chrom, startbp, endbp = parseRegionText(coordinates)

    old_stdout = sys.stdout
    sys.stdout = Logger(logfilename)
    
    print('Start: ' + datetime.now().strftime('%c'))
    print('filelist_filename: ' + filelist_filename)
    print('Region: ' + str(chrom)+':'+str(startbp)+'-'+str(endbp))
    print('outfilename: ' + outfilename)
    print('Log file: ' + logfilename)

    print('Reading list of files file')
    filelist = pd.read_csv(filelist_filename, sep='\t', header=None)
    files = list(filelist.iloc[:,0])
    file_descriptions = list(filelist.iloc[:,1])
    chrom_colnames = list(filelist.iloc[:,2])
    bp_colnames = list(filelist.iloc[:,3])
    snp_colnames = list(filelist.iloc[:,4])
    pval_colnames = list(filelist.iloc[:,5])
    
    print('Verifying settings in filelist_filename')
    for i in np.arange(len(files)):
        if not checkColnames(files[i], list(filelist.iloc[i,2:12])):
            raise Exception('Column names given: ' + str(list(filelist.iloc[i,2:12])) + '\n Not all match for file ' + files[i])

    print('Merging files')
    final_merge = '<!DOCTYPE html>\n<html>'
    for i in tqdm(np.arange(len(files))):
        df = pd.read_csv(files[i], sep='\t', skiprows=getNumHeaderLines(files[i]))
        desired_cols = list(filelist.iloc[i,2:12])
        df = df[desired_cols].rename(columns={
                desired_cols[0]: CHROM
                ,desired_cols[1]: BP
                ,desired_cols[2]: SNP
                ,desired_cols[3]: P
                ,desired_cols[4]: BETA
                ,desired_cols[5]: SE
                ,desired_cols[6]: N
                ,desired_cols[7]: ALT
                ,desired_cols[8]: REF
                ,desired_cols[9]: MAF
                })
        if chrom == 23:
            df[CHROM] = np.repeat(chrom, df.shape[0])
        df[CHROM] = [ int(str(x).lower().replace('chr','').replace('chrom','')) for x in list(df[CHROM]) ]
        df = df.loc[ (df[CHROM]==chrom) & (df[BP]>=startbp) & (df[BP]<=endbp) ]
        probeidcol = pd.DataFrame({ProbeID: np.repeat(filelist.iloc[i,12], df.shape[0])})
        df = pd.concat([df, probeidcol], axis=1)
        h3tag = '<h3>'+file_descriptions[i]+'</h3>'
        df_html_str = df.to_html(index=False)
        final_merge += h3tag + df_html_str
    final_merge += '</html>'
    
    print('Writing merged output')
    with open(outfilename, 'w') as f_out:
        f_out.write(final_merge)
    
    # Quick check:
#    with open(outfilename, encoding='utf-8', errors='replace') as f:
#        html = f.read()
#    soup = bs(html, 'lxml')
#    table_titles = soup.find_all('h3')
#    table_titles = [x.text for x in table_titles]
#    tables = pd.read_html(outfilename)

    print('Done')
    print('End: ' + datetime.now().strftime('%c'))
    sys.stdout = old_stdout