CoveragePlotter.py

### Moyosore shittu ###
### mshittu@uni-bonn.de ###
### v0.1 ###

__usage__ = """
EITHER give an exact region:

python3 CoveragePlotter.py
--bam <BAM_FILE>[ <BAM_FILE2> <BAM_FILE3> ...]
--chr <CHROMOSOME_OF_INTEREST>
--start <START_POSITION>
--end <END_POSITION>
--out <FULL_PATH_TO_OUTPUT_FIGURE>

OR give a gene ID plus a GFF3 file, and let the script find the region for you:

python3 CoveragePlotter.py
--bam <BAM_FILE>[ <BAM_FILE2> <BAM_FILE3> ...]
--gff <FULL_PATH_TO_GFF3_FILE>
--gene <GENE_ID>
--out <FULL_PATH_TO_OUTPUT_FIGURE>

--bam accepts one or several BAM files, separated by spaces; each gets its own coverage panel.
Instead of --bam, you may give --bamdir <FOLDER> to automatically use every .bam file in that folder.

optional (usable with either mode above):
--bamdir <FOLDER_WITH_BAM_FILES>	use every .bam file found in this folder, instead of listing them with --bam
--gff <FULL_PATH_TO_GFF3_FILE>		also draws the exon/intron structure below the coverage track(s)
--gene <GENE_ID>			also draws the exon/intron structure below the coverage track(s) (needs --gff)
--padding <BASES>			extra bases to show around a --gene lookup [default=200]
--minqual <MIN_MAPPING_QUALITY>	minimal mapping quality of reads to consider [default=0]
--smooth <WINDOW_SIZE, INT>		smooths the coverage line with a moving average over this many bases [default=1, off]
--kde <BANDWIDTH, FLOAT>		smooths the coverage line with a Gaussian kernel (KDE-style) of this bandwidth instead of a moving average
					(--kde overrides --smooth if both are given)
--show-raw				also draws the unsmoothed raw coverage as a faint line behind the smoothed line
					(only useful together with --smooth or --kde)
--split					always save one separate figure per BAM file, even if there are few samples
					(with --out coverage.png and 3 BAM files, produces coverage_SAMPLE1.png, coverage_SAMPLE2.png, coverage_SAMPLE3.png)
--combine				always save one combined figure with all samples stacked, even if there are many samples
--group-size <N, INT>			if neither --split nor --combine is given, samples are auto-combined when there are
					N or fewer BAM files, and auto-split into separate figures when there are more than N [default=4]

if none of --split / --combine / --group-size is given, the script decides automatically using the default group size of 4.

bug reports and feature requests: mshittu@uni-bonn.de
"""

import sys, re, glob, os, math
import pysam
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

# --- end of imports --- #


def get_bam_files_from_flag( arguments ):
	"""! @brief collect one or more BAM files given after --bam, stopping at the next flag """
	
	bam_files = []
	idx = arguments.index( '--bam' ) + 1
	while idx < len( arguments ) and arguments[ idx ][0:2] != '--':
		bam_files.append( arguments[ idx ] )
		idx += 1
	
	return bam_files


def get_bam_files_from_dir( bam_dir ):
	"""! @brief collect all BAM files located in the given folder """
	
	if bam_dir[-1] != "/":
		bam_dir += "/"
	
	bam_files = sorted( glob.glob( bam_dir + "*.bam" ) )
	
	if len( bam_files ) == 0:
		sys.exit( "ERROR: no .bam files found in " + bam_dir )
	
	return bam_files


def get_coverage_from_bam( bam_file, chromosome, start, end, minqual ):
	"""! @brief extract per-base coverage for the given region from a BAM file """
	
	# --- open BAM file and check chromosome name --- #
	bam = pysam.AlignmentFile( bam_file, "rb" )
	
	if chromosome not in bam.references:
		if chromosome[0:3] == "chr":
			alt_chromosome = chromosome[3:]
		else:
			alt_chromosome = "chr" + chromosome
		if alt_chromosome in bam.references:
			chromosome = alt_chromosome
		else:
			sys.exit( "ERROR: chromosome '" + chromosome + "' not found in " + bam_file )
	
	# --- count per-base coverage across the region --- #
	cov_A, cov_C, cov_G, cov_T = bam.count_coverage( chromosome, start, end,
		quality_threshold=minqual )
	coverage = []
	pos = 0
	while pos < len( cov_A ):
		coverage.append( cov_A[ pos ] + cov_C[ pos ] + cov_G[ pos ] + cov_T[ pos ] )
		pos += 1
	
	bam.close()
	return coverage


def load_exons_from_gff( gff_file, gene_id ):
	"""! @brief load all exon positions of the given gene from a GFF3/GTF file """
	
	exons = []
	
	with open( gff_file, "r" ) as f:
		line = f.readline()
		while line:
			if line[0] != "#":
				parts = line.strip().split( '\t' )
				if len( parts ) == 9:
					feature = parts[2]
					attributes = parts[8]
					if feature == "exon" and gene_id in attributes:
						exons.append( ( int( parts[3] ), int( parts[4] ) ) )
			line = f.readline()
	
	if len( exons ) == 0:
		sys.exit( "ERROR: no exons found for gene '" + gene_id + "' in " + gff_file )
	
	return sorted( exons )


def find_gene_region_in_gff( gff_file, gene_id, padding ):
	"""! @brief look up chromosome/start/end of a gene from a GFF3/GTF file """
	
	with open( gff_file, "r" ) as f:
		line = f.readline()
		while line:
			if line[0] != "#":
				parts = line.strip().split( '\t' )
				if len( parts ) == 9:
					feature = parts[2]
					attributes = parts[8]
					if feature in ( "gene", "mRNA", "transcript" ) and gene_id in attributes:
						chromosome = parts[0]
						start = max( 1, int( parts[3] ) - padding )
						end = int( parts[4] ) + padding
						return chromosome, start, end
			line = f.readline()
	
	sys.exit( "ERROR: gene '" + gene_id + "' not found in " + gff_file )


def smooth_coverage( coverage, window ):
	"""! @brief apply a simple moving average to the coverage track """
	
	if window <= 1:
		return coverage
	
	smoothed = []
	half = window // 2
	n = len( coverage )
	pos = 0
	while pos < n:
		lo = max( 0, pos - half )
		hi = min( n, pos + half + 1 )
		chunk = coverage[ lo:hi ]
		smoothed.append( sum( chunk ) / len( chunk ) )
		pos += 1
	
	return smoothed


def kde_smooth_coverage( coverage, bandwidth ):
	"""! @brief apply a Gaussian kernel (KDE-style) smoothing to the coverage track """
	
	if bandwidth <= 0:
		return coverage
	
	smoothed = []
	radius = int( round( 3 * bandwidth ) )
	n = len( coverage )
	pos = 0
	while pos < n:
		lo = max( 0, pos - radius )
		hi = min( n, pos + radius + 1 )
		weight_sum = 0.0
		value_sum = 0.0
		j = lo
		while j < hi:
			distance = pos - j
			weight = math.exp( -0.5 * ( distance / bandwidth ) ** 2 )
			value_sum += weight * coverage[ j ]
			weight_sum += weight
			j += 1
		smoothed.append( value_sum / weight_sum )
		pos += 1
	
	return smoothed


def build_split_filename( fig_file, sample_label ):
	"""! @brief build a per-sample output filename, e.g. coverage.png + ERR123 -> coverage_ERR123.png """
	
	sample_name = sample_label.split( '.' )[0]
	
	if '.' in fig_file:
		prefix = fig_file.rsplit( '.', 1 )[0]
		suffix = fig_file.rsplit( '.', 1 )[1]
		return prefix + "_" + sample_name + "." + suffix
	else:
		return fig_file + "_" + sample_name + ".png"


def decide_split( arguments, number_of_samples, group_size ):
	"""! @brief decide whether to save one combined figure or one figure per sample """
	
	if '--split' in arguments:
		return True, "one figure per sample - forced with --split"
	elif '--combine' in arguments:
		return False, "one combined figure - forced with --combine"
	elif number_of_samples > group_size:
		return True, "one figure per sample - auto-split, since " + str( number_of_samples ) + " samples is more than the group size of " + str( group_size )
	else:
		return False, "one combined figure - auto-combined, since " + str( number_of_samples ) + " samples is within the group size of " + str( group_size )


def construct_figure( coverage_tracks, labels, chromosome, start, end, fig_file, exons, raw_tracks=None, color_offset=0 ):
	"""! @brief construct coverage figure for the region of interest, with one panel per sample """
	
	positions = list( range( start, end ) )
	number_of_tracks = len( coverage_tracks )
	track_colors = [ "green", "blue", "red", "orange", "purple", "brown" ]
	
	# --- set up figure layout: one panel per sample, plus one for the gene model if given --- #
	height_ratios = [ 4 ] * number_of_tracks
	if exons is not None:
		height_ratios.append( 1 )
	
	fig, axes = plt.subplots( len( height_ratios ), 1, figsize=(10, 2.2 * number_of_tracks + 0.8),
		gridspec_kw={ "height_ratios": height_ratios }, sharex=True )
	if len( height_ratios ) == 1:
		axes = [ axes ]
	
	# --- plot one coverage track per sample --- #
	track_index = 0
	while track_index < number_of_tracks:
		cov_ax = axes[ track_index ]
		color = track_colors[ ( track_index + color_offset ) % len( track_colors ) ]
		
		# --- draw the raw (unsmoothed) coverage faintly behind the smoothed line, if requested --- #
		if raw_tracks is not None:
			cov_ax.plot( positions, raw_tracks[ track_index ], color="gray", linewidth=0.5, alpha=0.5, label="raw" )
		
		cov_ax.fill_between( positions, coverage_tracks[ track_index ], color=color, alpha=0.4 )
		cov_ax.plot( positions, coverage_tracks[ track_index ], color=color, linewidth=0.8, label=labels[ track_index ] )
		cov_ax.set_ylabel( "coverage" )
		cov_ax.legend( loc="upper right", fontsize=8, frameon=False )
		track_index += 1
	
	axes[0].set_title( chromosome + ":" + str( start ) + "-" + str( end ) )
	
	# --- plot gene model track, if provided --- #
	if exons is not None:
		gene_ax = axes[-1]
		gene_ax.axhline( 0, color="black", linewidth=1 )
		for exon_start, exon_end in exons:
			es, ee = max( exon_start, start ), min( exon_end, end )
			if ee > es:
				gene_ax.add_patch( plt.Rectangle( ( es, -0.4 ), ee - es, 0.8,
					facecolor="steelblue", edgecolor="black" ) )
		gene_ax.set_yticks( [] )
	
	axes[-1].set_xlabel( chromosome + " position (bp)" )
	
	fig.tight_layout()
	
	# --- make sure the output folder exists before saving, so --out into a new folder just works --- #
	out_dir = os.path.dirname( fig_file )
	if out_dir != "" and not os.path.isdir( out_dir ):
		os.makedirs( out_dir )
	
	fig.savefig( fig_file, dpi=300 )


def main( arguments ):
	"""! @brief run all parts of this script """
	
	if '--bam' in arguments:
		bam_files = get_bam_files_from_flag( arguments )
	else:
		bam_dir = arguments[ arguments.index( '--bamdir' )+1 ]
		bam_files = get_bam_files_from_dir( bam_dir )
	
	fig_file = arguments[ arguments.index( '--out' )+1 ]
	
	if '--minqual' in arguments:
		minqual = int( arguments[ arguments.index( '--minqual' )+1 ] )
	else:
		minqual = 0
	
	if '--padding' in arguments:
		padding = int( arguments[ arguments.index( '--padding' )+1 ] )
	else:
		padding = 200
	
	if '--smooth' in arguments:
		smooth_window = int( arguments[ arguments.index( '--smooth' )+1 ] )
	else:
		smooth_window = 1
	
	if '--kde' in arguments:
		kde_bandwidth = float( arguments[ arguments.index( '--kde' )+1 ] )
	else:
		kde_bandwidth = None
	
	if kde_bandwidth is not None and smooth_window > 1:
		print( "NOTE: both --smooth and --kde were given - using --kde (Gaussian kernel) and ignoring --smooth, since only one smoothing method is applied at a time" )
	
	show_raw = '--show-raw' in arguments
	
	if '--group-size' in arguments:
		group_size = int( arguments[ arguments.index( '--group-size' )+1 ] )
	else:
		group_size = 4
	
	if '--gff' in arguments:
		gff_file = arguments[ arguments.index( '--gff' )+1 ]
	else:
		gff_file = None
	
	if '--gene' in arguments:
		gene_id = arguments[ arguments.index( '--gene' )+1 ]
	else:
		gene_id = None
	
	if '--chr' in arguments and '--start' in arguments and '--end' in arguments:
		# region given explicitly by the user
		chromosome = arguments[ arguments.index( '--chr' )+1 ]
		start = int( arguments[ arguments.index( '--start' )+1 ] )
		end = int( arguments[ arguments.index( '--end' )+1 ] )
	elif gff_file is not None and gene_id is not None:
		# region derived automatically from the gene ID in the GFF3 file
		chromosome, start, end = find_gene_region_in_gff( gff_file, gene_id, padding )
		print( "gene '" + gene_id + "' found at " + chromosome + ":" + str( start ) + "-" + str( end ) )
	else:
		sys.exit( __usage__ )
	
	exons = None
	if gff_file is not None and gene_id is not None:
		exons = load_exons_from_gff( gff_file, gene_id )
	
	# --- extract and smooth coverage for every BAM file --- #
	coverage_tracks = []
	raw_tracks = []
	labels = []
	skipped_files = []
	for bam_file in bam_files:
		print( "processing " + bam_file + " ..." )
		try:
			raw_coverage = get_coverage_from_bam( bam_file, chromosome, start, end, minqual )
			if kde_bandwidth is not None:
				coverage = kde_smooth_coverage( raw_coverage, kde_bandwidth )
			else:
				coverage = smooth_coverage( raw_coverage, smooth_window )
			coverage_tracks.append( coverage )
			raw_tracks.append( raw_coverage )
			labels.append( os.path.basename( bam_file ) )
		except Exception as error:
			print( "WARNING: skipping " + bam_file + " - could not read this BAM file (" + str( error ) + ")" )
			skipped_files.append( bam_file )
	
	if len( coverage_tracks ) == 0:
		sys.exit( "ERROR: no BAM file could be read - nothing to plot" )
	
	do_split, split_reason = decide_split( arguments, len( coverage_tracks ), group_size )
	print( split_reason )
	
	if do_split:
		# --- save one separate figure per sample --- #
		track_index = 0
		while track_index < len( coverage_tracks ):
			sample_fig_file = build_split_filename( fig_file, labels[ track_index ] )
			if show_raw:
				sample_raw_tracks = [ raw_tracks[ track_index ] ]
			else:
				sample_raw_tracks = None
			construct_figure( [ coverage_tracks[ track_index ] ], [ labels[ track_index ] ],
				chromosome, start, end, sample_fig_file, exons, raw_tracks=sample_raw_tracks, color_offset=track_index )
			print( "coverage figure saved to " + sample_fig_file )
			track_index += 1
	else:
		# --- save one combined figure with all samples stacked --- #
		if show_raw:
			combined_raw_tracks = raw_tracks
		else:
			combined_raw_tracks = None
		construct_figure( coverage_tracks, labels, chromosome, start, end, fig_file, exons, raw_tracks=combined_raw_tracks )
		print( "coverage figure saved to " + fig_file )
	
	if len( skipped_files ) > 0:
		print( str( len( skipped_files ) ) + " file(s) were skipped due to errors: " + ", ".join( skipped_files ) )


if __name__ == '__main__':
	
	has_region = '--chr' in sys.argv and '--start' in sys.argv and '--end' in sys.argv
	has_gene_lookup = '--gff' in sys.argv and '--gene' in sys.argv
	has_bam_input = '--bam' in sys.argv or '--bamdir' in sys.argv
	
	if has_bam_input and '--out' in sys.argv and ( has_region or has_gene_lookup ):
		main( sys.argv )
	else:
		sys.exit( __usage__ )