Source code for fudge.vis.matplotlib.plot_matrix

# <<BEGIN-copyright>>
# <<END-copyright>>

import numpy
from matplotlib import pyplot
from matplotlib import colors
from matplotlib import ticker

[docs]def plot_matrix( matrix, energyBoundariesX=None, energyBoundariesY=None, title="Matrix", xyTitle='Energy (MeV)', xylog=False, zlog=False, zRange=(), switchY=True, colorMap='jet', subplot=None ): """Plot a matrix. Arguments are: @matrix (required): 2-dimensional numpy array (or equivalent) @energyBoundariesX: list of group boundary edges. If omitted, use matrix indices for axis @energyBoundariesY: "". If omitted, use same boundaries as x-axis @title @xyTitle: specify x (and optionally y) titles. May be string or tuple of two strings @xylog: if True, energyBoundariesX (and Y if different) must also be specified @zlog @zRange: for plotting a sub-range of z @switchY: if True, plot matrix[0,0] at upper left rather than lower left @colorMap: select color range for z-axis. Options include 'jet', 'jet_r', 'correlation', etc. @subplot: matplotlib.axes.AxesSubplot instance (for drawing matrix inside larger figure) This prepares a 2-d matrix plot. Use pyplot.show() to view the result.""" matrix = numpy.array( matrix ) if matrix.ndim != 2: raise Exception("plot_matrix can only handle 2-d arrays!") if not subplot: fig = pyplot.figure() subplot = fig.add_subplot(111) if type(xyTitle) is str: xtitle = ytitle = xyTitle elif len(xyTitle)==2: xtitle, ytitle = xyTitle else: raise Exception("Can't understand xyTitle argument: %s" % xyTitle) pyplot.xlabel(xtitle) pyplot.ylabel(ytitle) if energyBoundariesY is None: energyBoundariesY = energyBoundariesX if energyBoundariesX is not None: lowX, highX = energyBoundariesX[0], energyBoundariesX[-1] lowY, highY = energyBoundariesY[0], energyBoundariesY[-1] else: lowX = lowY = 0 highX, highY = matrix.shape[::-1] pyplot.xlim( lowX, highX ) if switchY: # switch y-axis: put low energy at top pyplot.ylim( highY, lowY ) else: pyplot.ylim( lowY, highY ) if zRange: vmin, vmax = zRange else: vmin, vmax = matrix.min(), matrix.max() if zlog: try: vmin = min( matrix[matrix>0] ) except ValueError: print "WARNING: no matrix elements > 0. Switching to linear scale for z-axis" zlog=False colorBarOpts = {} if zlog: # z-axis log scale if vmin < 0: zopts = {'norm': colors.SymLogNorm( linthresh=vmax*1e-9, vmin=vmin, vmax=vmax )} ticks = [] # FIXME future versions of matplotlib may do this automatically for SymLogNorm zmax = 10**numpy.floor( numpy.log10( vmax ) ) for i in range(4): ticks.append( -zmax * 10**(-3*i) ) ticks.append(0) for i in range(3,-1,-1): ticks.append( zmax * 10**(-3*i) ) colorBarOpts['ticks'] = ticks colorBarOpts['format'] = ticker.LogFormatterMathtext() else: zopts = {'norm': colors.LogNorm( vmin=vmin, vmax=vmax ) } else: zopts = {'vmin':vmin, 'vmax':vmax } # color map: zopts['cmap'] = colorMap if colorMap=='correlation': zopts['cmap'] = generateColorMap() if xylog: if energyBoundariesX is None: raise Exception("xylog option requires energy group boundary info") X,Y = numpy.lib.function_base.meshgrid(energyBoundariesX,energyBoundariesY) pyplot.pcolormesh( X,Y, matrix, **zopts ) pyplot.setp( subplot, xscale='log', yscale='log' ) elif energyBoundariesX is not None: pyplot.pcolormesh( energyBoundariesX, energyBoundariesY, matrix, **zopts ) else: pyplot.imshow( matrix, interpolation='none', **zopts ) pyplot.title(title) pyplot.colorbar( **colorBarOpts ) pyplot.matplotlib.cm.jet.set_under('w')
[docs]def generateColorMap(): # based on jet_r, but the center is white (for 0 correlation) import matplotlib cdict = {'blue': [(0.0, 0, 0), (0.34999999999999998, 0, 0), (0.4995, 0.4828709677, 1), (0.5005, 1, 0.4828709677), (0.65999999999999992, 1, 1), (0.89000000000000001, 1, 1), (1.0, 0.5, 0.5)], 'green': [(0.0, 0, 0), (0.089999999999999969, 0, 0), (0.35999999999999999, 1, 1), (0.4995, 1, 1), (0.5005, 1, 1), (0.625, 1, 1), (0.875, 0, 0), (1.0, 0, 0)], 'red': [(0.0, 0.5, 0.5), (0.10999999999999999, 1, 1), (0.33999999999999997, 1, 1), (0.4995, 0.516129, 1), (0.5005, 1, 0.516129), (0.65000000000000002, 0, 0), (1.0, 0, 0)]} return matplotlib.colors.LinearSegmentedColormap('correlation',cdict,1024)