"""

NAME

CornerPlotter.py

PURPOSE

Plot 2D projections of ND parameter space as triangular

array of panels. Include 1D marginalised distributions

on the diagonal.

COMMENTS

Expected data format is plain text, header marked by # in column 1, with

header lines listing:

1) Parameter names separated by commas. These names will

be used as labels on the plots, and should be in latex.

2) Parameter ranges [min,max,] for plot axes.

USAGE

CornerPlotter.py [flags]

FLAGS

-h Print this message [0]

-v Verbose operation [0]

--eps Postscript output

--just-2D Don't bother plotting 1D PDFs.

-c --conditional Plot 2D conditional distributions only

INPUTS

file1,color1 Name of textfile containing data, and color to use

file2,color2 etc.

OPTIONAL INPUTS

-n --columns List of columns to plot [all] NB. inputs are one-indexed!

-w iw Index of column containing weight of sample

-L iL Index of column containing likelihood of sample

--plot-points Plot the samples themselves

-s --smooth x Plot contours that are smoother by factor x [def=1]

-o --output Name of output file

--test Write out a set of test commands and exit

OUTPUTS

stdout Useful information

pngfile Output plot in png format

EXAMPLES

CornerPlotter.py --test

BUGS

- Only works from command line at the moment!

- Tick labels overlap, cannot remove first and last tick mark

- Figure has no legend

- Overlay of 1D estimates not enabled

- No overlay of 1-D priors

- Lines are too thin

HISTORY

2010-05-06 started Marshall/Auger (UCSB)

2011-06-24 generalized Marshall (Courmayeur/Bologna)

2012-05-03 split into pappy module Marshall (Oxford)

"""

# --------------------------------------------------------------------

try:

opts, args = getopt.getopt(argv, "hvcew:L:n:o:s:",["help","verbose","test","conditional","eps","smooth","plot-points","just-2D","columns","output"])

except getopt.GetoptError, err:

# print help information and exit:

print str(err) # will print something like "option -a not recognized"

print CornerPlotter.__doc__

return

vb = False

wcol = -1

Lcol = -1

plotpoints = False

eps = False

plotlegend = True

plot2D = True

plot1D = True

conditional = False

columns = 'All'

outfile = None

test = False

smoothscale = 1.0

# NB. wcol and Lcol are assumed to be entered indexed to 1!

for o,a in opts:

if o in ("-v", "--verbose"):

vb = True

elif o in ("--test"):

test = True

elif o in ("-n", "--columns"):

columns = a

elif o in ("--plot-points"):

plotpoints = True

elif o in ("-s","--smooth"):

smoothscale = float(a)

elif o in ("--just-2D"):

plot1D = False

elif o in ("-c", "--conditional"):

conditional = True

plot1D = False

elif o in ("-w"):

wcol = int(a) - 1

elif o in ("-L"):

Lcol = int(a) - 1

elif o in ("--eps"):

eps = True

elif o in ("-o", "--output"):

outfile = a

elif o in ("-h", "--help"):

print CornerPlotter.__doc__

return

else:

print "Couldn't understand option,argument: ",o,a

assert False

# Test cases:

if test:

print 'Try these:'

print ' CornerPlotter.py -o test1.png examples/thetas.cpt,blue,shaded'

print ' CornerPlotter.py -o test2.png -w 1 -n 2,3,4 examples/localgroup.cpt,red,shaded'

print ' CornerPlotter.py -o test3.png --just-2D -w 1 -n 3,4 examples/localgroup.cpt,purple,shaded'

print ' CornerPlotter.py -o test4.png --conditional --plot-points -n 2,3,4 examples/localgroup.cpt,green,shaded'

return

# Sort out output filename:

print outfile

if outfile == None:

if eps:

outfile = "cornerplot.eps"

else:

outfile = "cornerplot.png"

# Check for datafiles in array args:

# BUG: calling CornerPlotter from python, this line is needed for args

# to be treated as a list of strings, not a list of characters

# args = [args]

# But then the command line version does not work...

# Partial solution - split useful functions into pappy module. PJM

# args = numpy.array([args]) might help?

if len(args) > 0:

datafiles = []

colors = []

styles = []

legends = []

for i in range(len(args)):

pieces = args[i].split(',')

if len(pieces) == 1: pieces = pieces + ['black']

if len(pieces) == 2: pieces = pieces + ['outlined']

if len(pieces) == 3:

pieces = pieces + [' ']

plotlegend = False

datafiles = datafiles + [pieces[0]]

colors = colors + [pieces[1]]

styles = styles + [pieces[2]]

legends = legends + [pieces[3]]

if vb:

print "Making corner plot of data in following files:",datafiles

print "using following colors:",colors

if not plot1D: print "Leaving out 1D plots"

print "Output will be stored in "+output

if eps: print "and will be in postscript format"

else :

print CornerPlotter.__doc__

return

# --------------------------------------------------------------------

# Start figure, set up viewing area:

figprops = dict(figsize=(8.0, 8.0), dpi=128) # Figure properties

fig = pylab.figure(**figprops)

# Need small space between subplots to avoid deletion due to overlap...

adjustprops = dict(\

left=0.1,\

bottom=0.1,\

right=0.95,\

top=0.95,\

wspace=0.04,\

hspace=0.04)

fig.subplots_adjust(**adjustprops)

# Need to set global linewidth thicker here!

# No. of bins used:

nbins = 161

# Small tweak to prevent too many numeric tick marks

tiny = 0.01

# --------------------------------------------------------------------

# Plot files in turn, using specified color scheme.

for k in range(len(datafiles)):

datafile = datafiles[k]

color = colors[k]

# if (color == 'black' or eps):

# style = 'outlines'

# else:

# style = 'shaded'

style = styles[k]

legend = legends[k]

print "\nPlotting PDFs given in "+datafile+" as "+color+" "+style+" contours, and labelled as '"+legend+"'"

# Read in data:

data = numpy.loadtxt(datafile)

# Start figuring out how many parameters we have - index will be a

# list of column numbers containg the parameter values.

# NB. ALL files must have same parameters/weights structure!

if ( k == 0 ):

npars = data.shape[1]

index = numpy.arange(npars)

hmax = numpy.zeros([npars])

# Some cols may not be data, adjust index accordingly.

# Weights/importances:

if wcol >= 0 and wcol < data.shape[1]:

if vb: print "using weight in column: ",wcol

wht = data[:,wcol].copy()

if ( k == 0 ):

keep = (index != wcol)

index = index[keep]

npars = npars - 1

if vb: print " index = ",index

else:

if (wcol >= data.shape[1]):

print "WARNING: only ",data.shape[1]," columns are available"

print "Setting weights to 1"

wht = 0.0*data[:,0].copy() + 1.0

# Likelihood values:

if Lcol >= 0:

Lhood = data[:,Lcol].copy()

if ( k == 0 ):

keep = (index != Lcol)

index = index[keep]

npars = npars - 1

if vb: print "using lhood in column: ",Lcol

else:

Lhood = 0.0*data[:,0].copy() + 1.0

# Having done all that, optionally overwrite index with specified list

# of column numbers. Note conversion to zero-indexed python:

if columns != 'All':

pieces = columns.split(',')

index = []

for piece in pieces:

index.append(int(piece) - 1)

npars = len(index)

print "Only using data in",npars,"columns ",index,": "

# Now parameter list is in index - which is fixed for other datasets

# Font sizes - can only be set when no. of panels in grid is known:

# Big font sizes: {npars,bfs}={1,14},{2,13},{3,12},{4,10}

if plot1D:

ngrid = npars

else:

ngrid = npars - 1

if ngrid < 4:

bfs = 20 - 2*ngrid

else:

bfs = 10

sfs = bfs - 2

params = { 'axes.labelsize': bfs,

'text.fontsize': bfs,

'legend.fontsize': bfs,

'xtick.labelsize': sfs,

'ytick.labelsize': sfs}

pylab.rcParams.update(params)

# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

# 1st data file defines labels and parameter ranges:

if k == 0:

alllabels,alllimits,usedylimits = pappy.read_header(datafiles[k])

if vb: print 'No axis limits found, using 5-sigma ranges'

# Now pull out just the labels and limits we need:

limits = numpy.zeros([npars,2])

ii = 0

for i in index:

limits[ii,:] = alllimits[i,:]

ii = ii + 1

labels = alllabels[:]

# Set up dynamic axis limits, and smoothing scales:

dylimits = numpy.zeros([npars,2])

smooth = numpy.zeros(npars)

for i in range(npars):

col = index[i]

# Get data subarray, and measure its mean and stdev:

d = data[:,col].copy()

mean,stdev,Neff,N95 = pappy.meansd(d,wht=wht)

if vb: print "col = ",col," mean,stdev,nd = ",mean,stdev,Neff,N95

# Set smoothing scale for this parameter, in physical units.

smooth[i] = smoothscale*0.5*4.0*stdev/(N95**0.33)

# Cf Jullo et al 2007, who use a bin size given by

# w = 2*IQR/N^(1/3) for N samples, interquartile range IQR

# For a Gaussian, IQR is not too different from 2sigma. 4sigma/N^1/3?

# Also need N to be the effective number of parameters - return

# form meansd as sum of weights!

# Set 10 sigma limits:

dylimits[i,0] = mean - 10*stdev

dylimits[i,1] = mean + 10*stdev

# Now set up bin arrays based on dynamic limits,

# and convert smoothing to pixels:

bins = numpy.zeros([npars,nbins])

for i in range(npars):

col = index[i]

bins[i] = numpy.linspace(dylimits[i,0],dylimits[i,1],nbins)

smooth[i] = smooth[i]/((dylimits[i,1]-dylimits[i,0])/float(nbins))

if vb: print "col = ",col," smooth = ",smooth[i]

if vb: print "binning limits:",dylimits[i,0],dylimits[i,1]

if (k == 0):

# Finalise limits, again at 1st datafile:

if (usedylimits == 1): limits = dylimits

for i in range(npars):

limits[i,0] = limits[i,0] + tiny*abs(limits[i,0])

limits[i,1] = limits[i,1] - tiny*abs(limits[i,1])

# Good - limits are set.

# Report stats:

print "Effective number of samples: Neff,N95,N =",Neff,N95,len(d)

# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

# Loop over plotting panels. Arrangement is bottom left hand corner,

# most important parameter as x in first column, (row=0,col=0) is

# top left hand corner panel.

# panel = (col,row):

# 1=0,0 | 2=1,0 | 3=2,0 | 4=3,0 | 5=4,0

# ---------------------------------------------

# 6=0,1 | 7=1,1 | 8=2,1 | 9=3,1 | 10=4,1

# ---------------------------------------------

#11=0,2 | 12=1,2 | 13=2,2 | 14=3,2 | 15=4,2

# ---------------------------------------------

#16=0,3 | 17=1,3 | 18=2,3 | 19=3,3 | 20=4,3

# ---------------------------------------------

#21=0,4 | 22=1,4 | 23=2,4 | 24=3,4 | 25=4,4

for i in range(npars):

col = index[i]

# Get data subarray:

d1 = data[:,col].copy()

if vb: print "Read in ",col,"th column of data: min,max = ",min(d1),max(d1)

for j in range(i,npars):

row = index[j]

if plot1D and j==i:

# Move to next subplot:

panel = j*ngrid+i+1

pylab.subplot(ngrid,ngrid,panel)

# # Percentiles etc are too slow - get PDF1D to do it?

# # Report some statistcs:

# if vb:

# pct = percentiles(d1,wht)

# print " Percentiles (16,50,84th) =",pct

# median = pct[1]

# errplus = pct[2] - pct[1]

# errminus = pct[1] - pct[0]

# print " -> ",labels[col]+" = $",median,"^{",errplus,"}_{",errminus,"}$"

# Plot 1D PDF, defined in subroutine below

dummy,estimate = pappy.pdf1d(d1,wht,bins[i],smooth[i],color)

print " Par no.",col+1,",",labels[col],"=",estimate

if k == 0: hmax[i] = dummy

# Write the estimate on the plot?

# Force axes to obey limits:

pylab.axis([limits[i,0],limits[i,1],0.0,1.2*hmax[i]])

# Adjust axes of 1D plots:

ax = pylab.gca()

# Turn off the y axis tick labels for all 1D panels:

ax.yaxis.set_major_formatter(pylab.NullFormatter())

# Turn off the x axis tick labels for all but the last 1D panel:

if j<(npars-1):

ax.xaxis.set_major_formatter(pylab.NullFormatter())

pylab.xticks(rotation=45)

# Label x axis, only on the bottom panels:

if j==npars-1:

pylab.xlabel(labels[col])

plottedsomething = True

elif plot2D and j!=i:

# Move to next subplot:

if plot1D:

panel = j*ngrid+i+1

else:

# i j panel

# 0 1 1 = (j-1)*ngrid+i+1

# 0 2 3

# 1 2 4

panel = (j-1)*ngrid+i+1

# print "XXXXX i,j,panel = ",i,j,panel

pylab.subplot(ngrid,ngrid,panel)

# Get 2nd data set:

d2 = data[:,row].copy()

if vb: print "Read in ",row,"th column of data: min,max = ",min(d2),max(d2)

# Plot 2D PDF, defined in subroutine below

if vb: print "Calling pdf2d for col,row = ",col,row

fwhm = 0.5*(smooth[i]+smooth[j])

pappy.pdf2d(d1,d2,wht,bins[i],bins[j],fwhm,color,style,conditional=conditional)

# If we are just plotting one file, overlay samples:

if (len(datafiles) == 1 and plotpoints):

pylab.plot(d1,d2,'ko',ms=0.1)

# Force axes to obey limits:

pylab.axis([limits[i,0],limits[i,1],limits[j,0],limits[j,1]])

# Adjust axes of 2D plots:

ax = pylab.gca()

if i>0:

# Turn off the y axis tick labels

ax.yaxis.set_major_formatter(pylab.NullFormatter())

if j<npars-1:

# Turn off the x axis tick labels

ax.xaxis.set_major_formatter(pylab.NullFormatter())

# Rotate ticks so that axis labels don't overlap

pylab.xticks(rotation=45)

# Label x axes, only on the bottom panels:

if j==npars-1:

pylab.xlabel(labels[col])

if i==0 and j>0:

# Label y axes in the left-hand panels

pylab.ylabel(labels[row])

plottedsomething = True

else:

plottedsomething = False

if plottedsomething:

if vb: print "Done subplot", panel,"= (", i, j,")"

if vb: print " - plotted",labels[col],"vs",labels[row]

if vb: print "--------------------------------------------------"

# endfor

# endfor

# endfor

# Plot legend? In opposite corner if ngrid > 1, otherwise?

if plotlegend:

pass

# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

# Plot graph to file:

pylab.savefig(outfile,dpi=300)

print "\nFigure saved to file:",outfile