add options and cumulative histogram

knodlseder
1 parent 211b3d95
Showing 1 changed file with 168 additions and 38 deletions
examples/show_ts_distribution.py
 #! /usr/bin/env python
 # ==========================================================================
-# This script displays pull histograms that have been generated using the
-# cspull script. The pull is defined by
-#
-#           (fitted value - simulated value) / estimated error
-#
-# and the pull histogram allows to assess possible biases in the fit
-# parameters and errors.
+# This script displays the TS distribution values generated using cstsdist.
 #
 # Required 3rd party modules:
 # - matplotlib
@@ -39,9 +33,9 @@ import math
 # ==================== #
 # Read TS distribution #
 # ==================== #
-def read_ts(filename):
+def read_ts(filename, tsname="TS"):
 	"""
-	Read TS distribution.
+	Read TS distribution values from CSV file.
 	"""
 	# Initialise list
 	values = []
@@ -57,9 +51,10 @@ def read_ts(filename):
 		# Get column index if first row
 		if first:
 			try:
-				index = row.index("TS")
+				index = row.index(tsname)
 			except:
 				print 'ERROR: Column "TS" not found in file'
+				print row
 				raise NameError("TS")
 			#print index
  
@@ -72,39 +67,108 @@ def read_ts(filename):
  
 	# Create numpy array
 	a = np.array(values)
-	#print a
  
 	# Return array
 	return a
  
  
-# ======================== #
-# Main routine entry point #
-# ======================== #
-if __name__ == '__main__':
+# =========== #
+# Compute erf #
+# =========== #
+def erf(x):
 	"""
-	Display pull histogram generated by cspull using matplotlib.
+	Compute error function.
 	"""
-	# Print usage information
-	usage = "Usage: show_ts_distribution filename [bins]"
-	if len(sys.argv) < 2:
-		print usage
-		sys.exit()
+	# save the sign of x
+	sign = 1
+	if x < 0: 
+		sign = -1
+	x = abs(x)
  
-	# Extract parameters
-	filename = sys.argv[1]
-	if len(sys.argv) == 3:
-		nbins = int(sys.argv[2])
-	else:
-		nbins = 30
+	# constants
+	a1 =  0.254829592
+	a2 = -0.284496736
+	a3 =  1.421413741
+	a4 = -1.453152027
+	a5 =  1.061405429
+	p  =  0.3275911
  
-	# Set lin/log
-	log = True
+	# 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)
  
-	# Read values from CSV file
-	values = read_ts(filename)
-	print len(values), "values read."
+	# 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)
@@ -117,7 +181,7 @@ if __name__ == &#39;__main__&#39;:
 	# Create histogram
 	n, bins, patches = plt.hist(values, nbins, range=[min_value, max_value], \
 	                            align='mid', facecolor='green', \
-								log=log)
+								log=True, label="simulations")
  
 	# Create expected distribution (onyl for positive TS). We compute the
 	# full and half of the distribution, as for positively constrained
@@ -139,17 +203,83 @@ if __name__ == &#39;__main__&#39;:
 			y_half.append(0.5*y)
  
 	# Show expected distribution (half)
-	if log:
-		l = plt.semilogy(x, y_half, 'ro', linewidth=2)
-	else:
-		l = plt.plot(x, y_half, 'r-', linewidth=2)
+	l = plt.semilogy(x, y_half, 'ro', linewidth=2, label="expectation")
  
 	# Set plot
 	plt.xlabel('TS')
 	plt.ylabel('Number of trials')
-	plt.title('TS distribution ('+filename+')')
+	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:
+		print usage
+		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:
+						print usage
+						sys.exit()
+				else:
+					print usage
+					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)
+	print len(values), "values read."
+	
+	# Show histogram
+	if plot == "pdf":
+		dist_pdf(values, nbins, title)
+	else:
+		dist_cdf(values, nbins, title)
+
  
 \ No newline at end of file