Hi Martyn,
Thanks very much, this is very helpful. I have slightly modified your script to allow the z- (colour-)scale to be log. For this I had to carry out some checks as many of the detectorpixels will have 0.0 intensity.
In addition I added a modification that allows to specify the limits of the three axes:
import mantid
from mantid.simpleapi import *
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as colors
try:
mplcmd = gui_cmd
except NameError:
def mplcmd(callable, *args, **kwargs):
return callable(*args, **kwargs)
#-----------------------------------------------------------------------------
# Reset data suitable for logscale plotting
#-----------------------------------------------------------------------------
def ResetDataForLogscale(data,min=1e-20):
# logconversion will fail for 0.0 and negative numbers
#find out where logconversion will fail
isokforlogscale = np.log(data)
isokforlogscale = np.isfinite(isokforlogscale) # True if log worked, False otherwise
# replace data points where logconversion fails with positive, finite minimum value
dataforlogscale = np.where(isokforlogscale,data,[min])
return dataforlogscale
# ------------------------------------------------------------------------------
# Save a figure as filename
# ------------------------------------------------------------------------------
def CreateFigure3D(wksp, xlabel, ylabel,
rasterized=False, aspect='auto', xmin = None, xmax = None, ymin = None, ymax = None, zmin = None, zmax = None, logz = False, logmin = 1e-20):
# prepare data for treatment by matplotlib
# Signal values (really the Z axis)
intensity = wksp.extractY()
if logz: intensity = ResetDataForLogscale(intensity,logmin)
if not zmin: zmin = intensity.min()
if not zmax: zmax = intensity.max()
# x- and y-axes
data_xmin, data_xmax = wksp.readX(0)[0], wksp.readX(0)[-1]
#xmin, xmax = wksp.readX(0)[0], wksp.readX(0)[-1]
xstep = wksp.readX(0)[1] - data_xmin
x = np.arange(data_xmin, data_xmax + xstep/2.,xstep)
data_ymin, data_ymax = wksp.getAxis(1).getMin(), wksp.getAxis(1).getMax()
ystep = wksp.getAxis(1).getValue(1) - data_ymin
y = np.arange(data_ymin, data_ymax + ystep/2.,ystep)
xx, yy = np.meshgrid(x, y)
fig, ax1 = mplcmd(plt.subplots, 1, 1)
ax1.set_aspect(aspect)
if logz:
norm = colors.LogNorm(vmin=zmin, vmax=zmax)
else:
norm = colors.Normalize(vmin=zmin, vmax=zmax)
cmesh = mplcmd(ax1.pcolormesh, xx, yy, intensity,norm = norm)
cmesh.set_rasterized(rasterized)
#ax1.set_title(title)
ax1.set_xlabel(xlabel)
ax1.set_ylabel(ylabel)
for limit,data_limit in ((xmin, data_xmin), (xmax, data_xmax), (ymin, data_ymin), (ymax,data_ymax)):
if not limit: limit = data_limit
ax1.set_xlim(xmin, xmax)
ax1.set_ylim(ymin, ymax)
return fig
# ------------------------------------------------------------------------------
# Show a figure compatible with MantidPlot
# ------------------------------------------------------------------------------
def ShowFigure(figure):
mplcmd(figure.show)
# ------------------------------------------------------------------------------
# Save a figure as filename
# ------------------------------------------------------------------------------
def SaveFigure(figure, filename, dpi=600):
figure.savefig(filename, dpi=dpi, format="pdf")
# ------------------------------------------------------------------------------
# Create a test workspace for demo - replace with actual workspace
# ------------------------------------------------------------------------------
#ws = CreateSimulationWorkspace(Instrument='MAR',BinParams='-20,0.5,30')
#ws *= 0
#ws += 1
#sofqw = SofQW(ws, QAxisBinning='0,0.1,10',Emode='Direct',EFixed=60)
# load the example workspace somehow
#example = LoadISISNexus("example_ws.nxs")
example = mtd['example_ws']
# ------------------------------------------------------------------------------
# Create figure
fig = CreateFigure3D(example, xlabel='$Q_x\, [\AA^{-1}]$',
ylabel='$Q_z \,[\AA^{-1}]$',
rasterized=False, aspect='auto', logz = True,xmin = -0.00055, xmax = 0.00055, zmax=5e-4,zmin=5e-8)
# To display the plot:
ShowFigure(fig)
# To save to file
SaveFigure(fig, filename='test_save3d.pdf', dpi=600)
# to periodically clean the memory, matplotlib retains a reference to the figures even after they have been closed
plt.close('all')