n50

#!/bin/env python

import sys
import os
import pandas as pd
import numpy as np
import argparse
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
from matplotlib.ticker import LogLocator, LogFormatter
import logging


BIN_SIZE = 10000000

def get_n50(vals):
    
    vals = vals.sort_values(ascending=False)
    vals_csum = np.cumsum(vals)
    
    return vals.iloc[np.sum(vals_csum <= (vals_csum.iloc[-1] // 2))]/1000000

def bp_auto_formatter(x,pos):
    units = {
        1000:"kbp",
        1000**2:"Mbp",
        1000**3:"Gbp"
    }
    scaled_x = x
    unit_text = "bp"
    for unit_max in sorted(units.keys(),reverse=True):
        if x >= unit_max:
            scaled_x = x / unit_max
            unit_text = units[unit_max]
            break
    return f"{int(scaled_x)} {unit_text}"    
        
def get_log_max(max_len, digit=0):
    if max_len >= 10**digit:
        return get_log_max(max_len, digit+1)
    else:
        return digit


if __name__=="__main__":

    parser = argparse.ArgumentParser() 

    parser.add_argument('faidx', type=str, help='Index file of the assembly fasta file')
    parser.add_argument('--rows', '-r',required=False, action='store_true', help='Output information in a transposed format')
    parser.add_argument('--plot', '-p', type=str, required=False, help='Plot a scatter plot of read length and save to the provided argument')
    parser.add_argument('--len_dist', '-d', type=str, required=False, help='Plot the read length distribution and save to the provided argument')
    parser.add_argument('--log_scale', '-l', required=False, action='store_true', help='Plot in log-scale')
    parser.add_argument('--title', '-t', type=str, required=False, help='Title text for the plot')
    parser.add_argument('--no_output', '-n',required=False, action='store_true', help='No output information')
    
    
    args = parser.parse_args()

    ASM_FAIDX = args.faidx
    if not ASM_FAIDX.endswith('.fai'):
        ASM_FAIDX = "%s.fai"%ASM_FAIDX
    if not os.path.isfile("%s"%ASM_FAIDX):
        logging.error(f"{ASM_FAIDX} Not found.")
        sys.exit(1)

    with open(ASM_FAIDX, 'r') as in_file:
        len_list = [int(record.split('\t')[1]) for record in in_file]

    len_list = pd.Series(len_list)
    len_list.sort_values(ascending=False, inplace=True)
    log_max = get_log_max(max(len_list))

    aun = np.sum(len_list*len_list)/np.sum(len_list)/1000000

    n50 = get_n50(len_list)
    n50_realnum = n50*1000000

    if not args.no_output:
        if args.rows:
            print('Bases (Gbp)\tContigs\tN50 (Mbp)\t100k+ (Gbp)\tAuN (Mbp)')
            print('{:,.3f}\t{:,d}\t{:,.2f}\t{:,.3f}\t{:,.3f}'.format(
                np.sum(len_list)/1000000000.000,
                len(len_list),
                n50,
                np.sum(len_list.loc[len_list >= 100000])/1000000000.000,
                aun
            ))
        else:
            print(
            (
                'Bases (Gbp): {:,.3f}\n'
                'Contigs:     {:,d}\n'
                'N50 (Mbp):   {:,.2f}\n'
                '100k+ (Gbp): {:,.3f}\n'
                'AuN (Mbp):   {:,.3f}\n'
            ).format(
                np.sum(len_list)/1000000000.000,
                len(len_list),
                n50,
                np.sum(len_list.loc[len_list >= 100000])/1000000000.000,
                aun
            )
            )

    if args.plot:
        if not args.plot.endswith("png"):
            logging.error("-p/--plot must end with png")
            sys.exit(1)

        fig_p, ax_p = plt.subplots()
        ax_p.scatter(range(len(len_list)),len_list.sort_values(ascending=True))
        if args.log_scale:
            ax_p.set_yscale('log')
            ax_p.set_ylabel("Contig Length(Log-Scaled)")
            ax_p.text(int(len(len_list)/3),(n50_realnum*1.1),'N50 = {:,.2f} Mbp'.format(n50))
            ax_p.set_ylim(1000,10**log_max)
        else:
            ax_p.set_ylabel("Contig Length")
            ax_p.text(int(len(len_list)/3),(n50_realnum*1.1),'N50 = {:,.2f} Mbp'.format(n50))
            
        ax_p.yaxis.set_major_formatter(ticker.FuncFormatter(bp_auto_formatter))
        ax_p.axhline(y=n50_realnum, color="r", linestyle="-")
        ax_p.set_xlabel("Contig")
        

        if args.title:
            plt.title(f"{args.title}")	
        plt.savefig(args.plot, bbox_inches='tight')


    if args.len_dist:
        if not args.len_dist.endswith("png"):
            logging.error("-d/--len_dist must end with png")
            sys.exit(1)
        fig_l, ax_l = plt.subplots()
        if args.log_scale:
            bin_num = int((10**log_max)/BIN_SIZE/2)
            bins = np.logspace(0, log_max,)
            hist, edges = np.histogram(len_list,bins=bins)
            ax_l.bar(edges[:-1], hist, width=np.diff(edges), edgecolor='black', align='edge')
            ax_l.set_xscale("log")
            ax_l.set_xlabel("Contig Length(Log-Scaled)")
            ax_l.text((n50_realnum*0.002),max(hist)*3/4,'N50 = {:,.2f} Mbp'.format(n50))
            ax_l.set_xlim(1000,10**log_max)

        else:
            bins = np.arange(0, max(len_list) + BIN_SIZE, BIN_SIZE)
            hist, edges = np.histogram(len_list,bins=bins)
            ax_l.hist(len_list, bins=bins, edgecolor='black', align='left')
            ax_l.set_xlabel("Contig Length")
            ax_l.text((n50_realnum*1.1),max(hist)*2/3,'N50 = {:,.2f} Mbp'.format(n50))

        
        ax_l.set_ylabel("Frequency")
        ax_l.xaxis.set_major_formatter(ticker.FuncFormatter(bp_auto_formatter))
        ax_l.axvline(x=n50_realnum, color="r", linestyle="-")
        
        
        if args.title:
            plt.title(f"{args.title}")	
        plt.savefig(args.len_dist, bbox_inches='tight')