Andreas Specovius / ctools

#! /usr/bin/env python
# ==========================================================================
# This script displays the TS distribution values generated using cstsdist.
#
# Required 3rd party modules:
# - matplotlib
# - numpy
#
# Copyright (C) 2011 Jurgen Knodlseder
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# ==========================================================================
import matplotlib.mlab as mlab
import matplotlib.pyplot as plt
import numpy as np
import sys
import csv
import math


# ==================== #
# Read TS distribution #
# ==================== #
def read_ts(filename, tsname="TS"):
	"""
	Read TS distribution values from CSV file.
	"""
	# Initialise list
	values = []
	
	# Open reader
	reader = csv.reader(open(filename, 'r'), delimiter=',')
	
	# Read rows
	first = True
	index = -1
	for row in reader:
		
		# Get column index if first row
		if first:
			try:
				index = row.index(tsname)
			except:
				sys.stdout.write('ERROR: Column "TS" not found in file\n')
				sys.stdout.write(row+"\n")
				raise NameError("TS")

		# Handle data rows
		else:
			values.append(float(row[index]))
		
		# Flag that first row has been passed
		first = False
	
	# Create numpy array
	a = np.array(values)
	
	# Return array
	return a


# =========== #
# Compute erf #
# =========== #
def erf(x):
	"""
	Compute error function.
	"""
	# save the sign of x
	sign = 1
	if x < 0: 
		sign = -1
	x = abs(x)

	# constants
	a1 =  0.254829592
	a2 = -0.284496736
	a3 =  1.421413741
	a4 = -1.453152027
	a5 =  1.061405429
	p  =  0.3275911
	
	# A&S formula 7.1.26
	t = 1.0/(1.0 + p*x)
	y = 1.0 - (((((a5*t + a4)*t) + a3)*t + a2)*t + a1)*t*math.exp(-x*x)
	
	# Return result
	return sign*y # erf(-x) = -erf(x)


# =============================== #
# Display cumulative distribution #
# =============================== #
def dist_cdf(values, nbins, title):
	"""
	Displays cumulative TS distribution and overplots expectation.
	"""
	# Set range and adapt number of bins. We make sure that we have a
	# histogram bin centred on 0 that will capture the TS=0 part.
	max_value = max(values)
	binsize   = max_value/nbins
	min_value = 0.5*binsize
	while min_value > min(values):
		min_value -= binsize
		nbins     += 1
	
	# Create histogram
	hist_x      = []
	hist_y      = []
	hist_expect = []
	norm        = 0.5*len(values)
	for i in range(nbins):
		x_min = min_value+i*binsize
		x_max = min_value+(i+1)*binsize
		x_val = 0.5*(x_min+x_max)
		y_val = 0.0
		for ts in values:
			if ts > x_min:
				y_val += 1.0
		hist_x.append(x_min)
		hist_y.append(y_val)
		hist_x.append(x_max)
		hist_y.append(y_val)
		if x_min > 0.0:
			value = norm*erf(math.sqrt(0.5*x_min)) + norm
			hist_expect.append(len(values)-value)
			hist_expect.append(len(values)-value)
		else:
			hist_expect.append(norm)
			hist_expect.append(norm)

	# Show expected distribution (half)
	l = plt.semilogy(hist_x, hist_y,      'r-', linewidth=2, label="simulations")
	l = plt.semilogy(hist_x, hist_expect, 'k-', linewidth=2, label="expectation")
	l = plt.semilogy(hist_x, hist_y,      'r-', linewidth=2)

	# Set plot
	plt.xlabel('TS')
	plt.ylabel('Number of trials > TS')
	plt.title(title)
	plt.grid(True)
	plt.legend(loc="upper right")

	# Show histogram
	plt.show()
	
	# Return
	return


# ======================================== #
# Display probability density distribution #
# ======================================== #
def dist_pdf(values, nbins, title):
	"""
	Display probability density TS distribution and overplots
	expectation.
	"""
	# Set range and adapt number of bins. We make sure that we have a
	# histogram bin centred on 0 that will capture the TS=0 part.
	max_value = max(values)
	binsize   = max_value/nbins
	min_value = 0.5*binsize
	while min_value > min(values):
		min_value -= binsize
		nbins     += 1
	
	# Create histogram
	n, bins, patches = plt.hist(values, nbins, range=[min_value, max_value], \
	                            align='mid', facecolor='green', \
								log=True, label="simulations")

	# Create expected distribution (onyl for positive TS). We compute the
	# full and half of the distribution, as for positively constrained
	# quantities half of the values have TS=0
	x      = []
	y_full = []
	y_half = []
	width  = bins[1]-bins[0]
	norm   = len(values)/(math.sqrt(2.0)*math.sqrt(math.pi))*width
	x.append(0.0)
	y_full.append(len(values))
	y_half.append(0.5*len(values))
	for bin in bins:
		ts = bin + 0.5*width
		if ts > 0:
			y  = norm*math.pow(ts,-0.5)*math.exp(-0.5*ts)
			x.append(ts)
			y_full.append(y)
			y_half.append(0.5*y)
	
	# Show expected distribution (half)
	l = plt.semilogy(x, y_half, 'ro', linewidth=2, label="expectation")

	# Set plot
	plt.xlabel('TS')
	plt.ylabel('Number of trials')
	plt.title(title)
	plt.grid(True)
	plt.legend(loc="upper right")

	# Show histogram
	plt.show()
		
	
# ======================== #
# Main routine entry point #
# ======================== #
if __name__ == '__main__':
	"""
	Display TS distribution generated using cstsdist using matplotlib.
	"""
	# Print usage information
	usage = "Usage: show_ts_distribution filename" \
	        " [-n bins] [-c column] [-t title] [-p plot]"
	if len(sys.argv) < 2:
		sys.stdout.write(usage+"\n")
		sys.exit()

	# Extract parameters
	filename = sys.argv[1]

	# Set default parameters
	nbins  = 30
	tsname = "TS"
	title  = "TS distribution"
	plot   = "pdf"

	# Parameter dictionnary
	pars = [{'option': '-n', 'value': nbins}, \
	        {'option': '-c', 'value': tsname}, \
	        {'option': '-t', 'value': title}, \
			{'option': '-p', 'value': plot}]
	
	# Gather parameters from command line
	i = 2
	while i < len(sys.argv):
		
		# Search for options
		for par in pars:
			if sys.argv[i] == par['option']:
				if len(sys.argv) > i+1:
					i      += 1
					try:
						par['value'] = sys.argv[i]
					except:
						sys.stdout.write(usage+"\n")
						sys.exit()
				else:
					sys.stdout.write(usage+"\n")
					sys.exit()
		
		# Next item
		i += 1
	
	# Recover parameters
	nbins  = int(pars[0]['value'])
	tsname = pars[1]['value']
	title  = pars[2]['value']
	plot   = pars[3]['value']
	
	# Read values from CSV file
	values = read_ts(filename, tsname=tsname)
	sys.stdout.write(len(values), "values read.\n")
	
	# Show histogram
	if plot == "pdf":
		dist_pdf(values, nbins, title)
	else:
		dist_cdf(values, nbins, title)