def rangeUnits(textBox, array):
# For some reason, max() doesn't work on numpy arrays containing non-finite numbers
maxValue = float('-inf')
for value in array:
if isfinite(value) and value > maxValue:
maxValue = value
if textBox.unit != '':
unit = util.displayWithUnitsNumber(maxValue, textBox.unit).split()[1]
rangeInUnits = [util.convertUnitsNumber(value, textBox.unit, unit) for value in array]
else:
power = int(round(floor(log10(abs(maxValue)))/3)*3) # round power to nearest multiple of 3
unit = '10^' + str(power) if power != 0 else ''
rangeInUnits = [x/(10**power) for x in array]
axisLabel = textBox.objectName() + (' (' + unit + ')' if unit else '')
return unit, axisLabel, rangeInUnits
def refreshPlots(self):
self.parent.ui.statusbar.showMessage('Redrawing plots ...')
self.erasePlots()
# nothing to plot, if beam hasn't been initialized
if not self.myBunch:
return
# get the specified units for plotting
self.unitsPos = self.ui.unitsPos.text()
self.unitsAngle = self.ui.unitsAngle.text()
# get the number of tick marks
self.numTicks = int(self.ui.numTicks.text())
# create local pointer to particle array
tmp6 = randomSampleOfBunch(
self.myBunch.getDistribution6D().getPhaseSpace6D().getArray6D(),
min(self.maxParticles, int(self.ui.numPtcls.text())))
numParticles = tmp6.shape[1]
self.calculateLimits(tmp6)
nLevels = 5 + int(math.pow(numParticles, 0.33333333))
nDivs = 10 + int(math.pow(numParticles, 0.2))
# generate the four plots
self.plotXY( util.convertUnitsNumber(tmp6[0,:], 'm', self.unitsPos),
util.convertUnitsNumber(tmp6[2,:], 'm', self.unitsPos),
nDivs, nLevels)
self.plotXPX(util.convertUnitsNumber(tmp6[0,:], 'm', self.unitsPos),
util.convertUnitsNumber(tmp6[1,:], 'rad', self.unitsAngle),
nDivs, nLevels)
self.plotYPY(util.convertUnitsNumber(tmp6[2,:], 'm', self.unitsPos),
util.convertUnitsNumber(tmp6[3,:], 'rad', self.unitsAngle),
nDivs, nLevels)
self.plotSDP(util.convertUnitsNumber(tmp6[4,:], 'm', self.unitsPos),
util.convertUnitsNumber(tmp6[5,:], 'rad', self.unitsAngle),
nDivs, nLevels)
self.parent.ui.statusbar.clearMessage()
def plot(self):
try:
xTextBox = self.textBox[self.ui.x.currentText()]
xOrginalValue = xTextBox.text()
yTextBox = self.textBox[self.ui.y.currentText()]
yOrginalValue = yTextBox.text()
zTextBox = self.textBox[self.ui.z.currentText()]
except KeyError:
return # Combo box choice was left blank
try:
xmin = util.convertUnitsStringToNumber(self.ui.xmin.text(), xTextBox.unit)
xmax = util.convertUnitsStringToNumber(self.ui.xmax.text(), xTextBox.unit)
if xmin > xmax:
xmin, xmax = xmax, xmin
ymin = util.convertUnitsStringToNumber(self.ui.ymin.text(), yTextBox.unit)
ymax = util.convertUnitsStringToNumber(self.ui.ymax.text(), yTextBox.unit)
if ymin > ymax:
ymin, ymax = ymax, ymin
except ValueError: # text box left blank
return
numPoints = 32 # per variable, 32*32 = 1,024 points total
xRange = numpy.linspace(xmin, xmax, numPoints)
_, xAxisLabel, xRangeUnits = rangeUnits(xTextBox, xRange)
yRange = numpy.linspace(ymin, ymax, numPoints)
_, yAxisLabel, yRangeUnits = rangeUnits(yTextBox, yRange)
Z = numpy.zeros((numPoints, numPoints))
plotProgress = QtGui.QProgressDialog("Plotting ...", None, 0, (numPoints**2)-1)
plotProgress.setMinimumDuration(0)
try:
for j, x in enumerate(xRange):
for i, y in enumerate(yRange):
plotProgress.setValue(plotProgress.value()+1)
xTextBox.setText(str(x))
yTextBox.setText(str(y))
results = calculate(self.userInputDict())
try:
Z[i,j] = results[zTextBox.dictName]
except KeyError:
Z[i,j] = float('nan')
zUnit, zAxisLabel, _ = rangeUnits(zTextBox, Z.flat)
for z in numpy.nditer(Z, op_flags=['readwrite']):
z[...] = util.convertUnitsNumber(z, zTextBox.unit, zUnit)
# Plotting
self.ui.plotWidget.canvas.fig.clear()
self.ui.plotWidget.canvas.ax = self.ui.plotWidget.canvas.fig.add_subplot(111)
c = self.ui.plotWidget.canvas.ax.imshow(numpy.flipud(Z), cmap = 'hot',
extent = [min(xRangeUnits), max(xRangeUnits),
min(yRangeUnits), max(yRangeUnits)],
aspect = 'auto')
self.ui.plotWidget.canvas.ax.set_xlabel(xAxisLabel)
self.ui.plotWidget.canvas.ax.set_ylabel(yAxisLabel)
cb = self.ui.plotWidget.canvas.fig.colorbar(c)
cb.set_label(zAxisLabel)
self.ui.plotWidget.canvas.ax.set_xlim(min(xRangeUnits), max(xRangeUnits))
self.ui.plotWidget.canvas.ax.set_ylim(min(yRangeUnits), max(yRangeUnits))
self.ui.plotWidget.canvas.fig.tight_layout()
self.ui.plotWidget.canvas.draw()
finally:
# Restore text boxes to original state
xTextBox.setText(xOrginalValue)
xTextBox.setCursorPosition(0)
yTextBox.setText(yOrginalValue)
yTextBox.setCursorPosition(0)
self.calculateAll()
def calculateLimits(self, _arr):
# nothing to do, if beam hasn't been initialized
if not self.myBunch:
return
# get average, RMS, min, max values and diffs
avgArray = self.myStat.calcAverages6D(_arr)
rmsArray = self.myStat.calcRmsValues6D(_arr)
minArray = self.myStat.calcMinValues6D(_arr)
maxArray = self.myStat.calcMaxValues6D(_arr)
# calculate the differences, imposing symmetry
diffZero = np.zeros(6)
for iLoop in range(6):
diffZero[iLoop] = max( (avgArray[iLoop]-minArray[iLoop]),
(maxArray[iLoop]-avgArray[iLoop]) )
# now switch based on the specified axis flag
# specify plot limits, symmetric around the zero axis
if self.axisFlag == 'symmetric':
self.xMax = (abs(avgArray[0])+diffZero[0])*util.convertUnitsNumber(1, 'm', self.unitsPos)
self.xpMax = (abs(avgArray[1])+diffZero[1])*util.convertUnitsNumber(1, 'rad', self.unitsAngle)
self.yMax = (abs(avgArray[2])+diffZero[2])*util.convertUnitsNumber(1, 'm', self.unitsPos)
self.ypMax = (abs(avgArray[3])+diffZero[3])*util.convertUnitsNumber(1, 'rad', self.unitsAngle)
self.ptMax = (abs(avgArray[5])+diffZero[5])*util.convertUnitsNumber(1, 'rad', self.unitsAngle)
self.xMin = -self.xMax
self.xpMin = -self.xpMax
self.yMin = -self.yMax
self.ypMin = -self.ypMax
self.ptMin = -self.ptMax
# specify plot limits, symmetric around the bunch (confined to 3 rms)
elif self.axisFlag == 'compact':
self.xMin = (avgArray[0]-3.*rmsArray[0])*util.convertUnitsNumber(1, 'm', self.unitsPos)
self.xpMin = (avgArray[1]-3.*rmsArray[1])*util.convertUnitsNumber(1, 'rad', self.unitsAngle)
self.yMin = (avgArray[2]-3.*rmsArray[2])*util.convertUnitsNumber(1, 'm', self.unitsPos)
self.ypMin = (avgArray[3]-3.*rmsArray[3])*util.convertUnitsNumber(1, 'rad', self.unitsAngle)
self.ptMin = (avgArray[5]-3.*rmsArray[5])*util.convertUnitsNumber(1, 'rad', self.unitsAngle)
self.xMax = (avgArray[0]+3.*rmsArray[0])*util.convertUnitsNumber(1, 'm', self.unitsPos)
self.xpMax = (avgArray[1]+3.*rmsArray[1])*util.convertUnitsNumber(1, 'rad', self.unitsAngle)
self.yMax = (avgArray[2]+3.*rmsArray[2])*util.convertUnitsNumber(1, 'm', self.unitsPos)
self.ypMax = (avgArray[3]+3.*rmsArray[3])*util.convertUnitsNumber(1, 'rad', self.unitsAngle)
self.ptMax = (avgArray[5]+3.*rmsArray[5])*util.convertUnitsNumber(1, 'rad', self.unitsAngle)
# symmetric around the bunch
elif self.axisFlag == 'bunch-centered':
self.xMin = (avgArray[0]-diffZero[0])*util.convertUnitsNumber(1, 'm', self.unitsPos)
self.xpMin = (avgArray[1]-diffZero[1])*util.convertUnitsNumber(1, 'rad', self.unitsAngle)
self.yMin = (avgArray[2]-diffZero[2])*util.convertUnitsNumber(1, 'm', self.unitsPos)
self.ypMin = (avgArray[3]-diffZero[3])*util.convertUnitsNumber(1, 'rad', self.unitsAngle)
self.ptMin = (avgArray[5]-diffZero[5])*util.convertUnitsNumber(1, 'rad', self.unitsAngle)
self.xMax = (avgArray[0]+diffZero[0])*util.convertUnitsNumber(1, 'm', self.unitsPos)
self.xpMax = (avgArray[1]+diffZero[1])*util.convertUnitsNumber(1, 'rad', self.unitsAngle)
self.yMax = (avgArray[2]+diffZero[2])*util.convertUnitsNumber(1, 'm', self.unitsPos)
self.ypMax = (avgArray[3]+diffZero[3])*util.convertUnitsNumber(1, 'rad', self.unitsAngle)
self.ptMax = (avgArray[5]+diffZero[5])*util.convertUnitsNumber(1, 'rad', self.unitsAngle)
if self.axisFlag=='compact' or self.axisFlag=='symmetric-compact':
# sMin / sMax always have to be 'bunch centered'
self.sMin = (avgArray[4]-3.*rmsArray[4])*util.convertUnitsNumber(1, 'm', self.unitsPos)
self.sMax = (avgArray[4]+3.*rmsArray[4])*util.convertUnitsNumber(1, 'm', self.unitsPos)
if self.axisFlag=='symmetric' or self.axisFlag=='bunch-centered':
# sMin / sMax always have to be 'bunch centered'
self.sMin = (avgArray[4]-diffZero[4])*util.convertUnitsNumber(1, 'm', self.unitsPos)
self.sMax = (avgArray[4]+diffZero[4])*util.convertUnitsNumber(1, 'm', self.unitsPos)
