class DistributionFunctionUI(QtWidgets.QMainWindow):
def __init__(self, argv):
QtWidgets.QMainWindow.__init__(self)
self.ui = distfunc_design.Ui_DistfuncUI()
self.ui.setupUi(self)
self.plotWindow = PlotWindow(width=800, height=600)
self.ax = None
self.radialAx = None
self.colorbar = None
self.radialPlotWindow = PlotWindow(width=700, height=400)
self.anapitchWindow = DistributionFunctionAnalysePitch(self)
self.maxF = None # Maximum of distribution function
self.maxP = None # Maximum momentum on distribution grid
self.currentPlotHandle = None
self.plotHandles = []
self.linestyles = []
self.logarithmicPlot = False
self.generateLineStyles()
self.filename, self.disttype = self.parseCmdArgs(argv)
self.loadDistribution(self.filename, typehint=self.disttype)
self.setupFigure()
self.bindEvents()
def parseCmdArgs(self, argv):
"""
Parse the command-line arguments that were passed to
this program.
argv: List of command-line arguments.
"""
if len(argv) == 0:
raise Exception(
"At least the name of the distribution function to load must be specified on startup."
)
filename = None
typehint = None
def getarg(argv, i):
i += 1
if i >= len(argv):
raise Exception("Expected argument after '--type'.")
else:
return argv[i], i
i = 0
L = len(argv)
while i < L:
if argv[i].startswith('--'):
if argv[i] == '--type':
typehint, i = getarg(argv, i)
else:
raise Exception(
"Unrecognized command-line argument: '{}'.".format(
argv[i]))
else:
if filename is None:
filename = argv[i]
else:
raise Exception(
"Multiple input files specified on command-line.")
i += 1
return filename, typehint
def analysePitchDistribution(self):
r = self.ui.sliderRadius.value()
self.anapitchWindow.setDistribution(r, self.distfunc)
self.anapitchWindow.show()
def bindEvents(self):
self.ui.sliderRadius.valueChanged.connect(self.sliderRadiusChanged)
self.ui.gbMoments.toggled.connect(self.plotMomentsChanged)
self.ui.cbVolumeElement.toggled.connect(self.plotMomentsChanged)
self.ui.rbDistParPerp.toggled.connect(self.plotTypeChanged)
self.ui.rbDistPXi.toggled.connect(self.plotTypeChanged)
self.ui.rbDist1D.toggled.connect(self.plotTypeChanged)
self.ui.rbCumCurrent.toggled.connect(self.plotTypeChanged)
self.ui.rbSynchrotron.toggled.connect(self.rbSynchrotronChanged)
self.ui.rbRunaway.toggled.connect(self.rbRunawayChanged)
self.ui.btnAnalysePitch.clicked.connect(self.analysePitchDistribution)
self.ui.btnAutomaticY.clicked.connect(self.setAutomaticYLimit)
self.ui.btnClearKeptDistributions.clicked.connect(
self.keepDistributionClear)
self.ui.btnKeepDistribution.clicked.connect(self.keepDistribution)
self.ui.btnUpdateYAxis.clicked.connect(self.setYLimit)
self.ui.btnPlotCurrent.clicked.connect(self.plotCurrentDensity)
self.ui.btnPlotRadprof.clicked.connect(self.plotRadialDensity)
self.ui.btnMomentProfile.clicked.connect(self.plotMomentProfile)
self.ui.btnPlotNow.clicked.connect(self.plotNow)
self.ui.tbMinY.returnPressed.connect(self.setYLimit)
self.ui.tbMaxY.returnPressed.connect(self.setYLimit)
self.ui.tbRunawayPc.returnPressed.connect(self.plotNow)
def closeEvent(self, event):
self.exit()
def exit(self):
self.plotWindow.close()
self.radialPlotWindow.close()
self.close()
def generateLineStyles(self):
"""
Generate a list of line styles that can be used
when making 1D plots.
"""
NCOLORS = 4
ls = ['-', '--', ':']
cmap = GeriMap.get()
for i in range(0, len(ls) * NCOLORS):
ci = i % NCOLORS
li = i % len(ls)
clr = cmap(float(ci / NCOLORS))
self.linestyles.append((clr, ls[li]))
def generateMoment(self, P, XI, F):
"""
Calculate a moment of the distribution function according
to the settings in the 'Moments' group box in the GUI.
P: Momentum grid on which the distribution function is defined.
XI: Pitch grid on which the distribution function is defined.
F: Distribution function.
"""
rF = F
logged = False # Should the returned function be plotted on a logarithmic scale?
if self.ui.rbSynchrotron.isChecked():
wavelength = float(self.ui.sbSynchWavelength.value()) * 1e-9
magfield = self.ui.sbSynchMagneticField.value()
synch = Bekefi.synchrotron(P, XI, wavelength, magfield)
rF = synch * rF
elif self.ui.rbRunaway.isChecked():
try:
pc = float(self.ui.tbRunawayPc.text())
rF[np.where(P < pc)] = np.amin(rF)
logged = True
except ValueError:
QMessageBox.critical(
self, 'Invalid format for critical momentum',
"The specified value for the critical momentum has an invalid format."
)
else:
print(
'WARNING: A moment which has not been implemented seems to have been selected. Ignoring...'
)
return rF, logged
def getAngleAveragedDistribution(self, r):
"""
Returns the 1D function to plot at the given radius.
This may be just the distribution function, or it may
be the distribution weighted with some quantity.
"""
self.logarithmicPlot = False
p, fp = self.distfunc.getAngleAveragedDistribution(r)
# This is difficult for 1D distributions...
if self.ui.gbMoments.isChecked():
pass
return p, fp
def getParPerpDistribution(self, r):
"""
Returns the 2D function to plot in ppar/pperp coordinates
at the given radius. This may be just the distribution function,
or it may be the distribution weighted with some quantity.
"""
p = np.linspace(0, self.distfunc.getMaxP(), 1000)
P, XI, F = self.distfunc.eval(r, p)
PPAR = P * XI
PPERP = np.sqrt(P**2 - PPAR**2)
if self.ui.gbMoments.isChecked():
F, self.logarithmicPlot = self.generateMoment(P, XI, F)
if self.ui.cbVolumeElement.isChecked():
F = F * P**2
if self.logarithmicPlot:
F = np.log10(np.abs(F))
else:
F = np.log10(np.abs(F))
self.logarithmicPlot = True
return PPAR, PPERP, F
def getPXiDistribution(self, r):
"""
Returns the 2D function to plot in p/xi coordinates at the
given radius. This may be just the distribution function,
or it may be the distribution weighted with some quantity.
"""
p = np.linspace(0, self.distfunc.getMaxP(), 1000)
P, XI, F = self.distfunc.eval(r, p)
if self.ui.gbMoments.isChecked():
F, self.logarithmicPlot = self.generateMoment(P, XI, F)
if self.ui.cbVolumeElement.isChecked():
F = F * P**2
if self.logarithmicPlot:
F = np.log10(np.abs(F))
else:
F = np.log10(np.abs(F))
self.logarithmicPlot = True
return P, XI, F
def getRadius(self):
"""
The returns the currently selected radius.
"""
return self.distfunc.getRadius(self.ui.sliderRadius.value())
def keepDistribution(self):
self.plotHandles.append(self.currentPlotHandle)
self.currentPlotHandle = None
def keepDistributionClear(self):
for h in self.plotHandles:
h.remove()
self.plotHandles = []
self.currentPlotHandle.remove()
self.currentPlotHandle = None
self.plotSelection()
def loadDistribution(self, filename, typehint=None):
"""
Load the named distribution function.
filename: Name of file to load.
"""
self.ui.lblFileName.setText(os.path.basename(filename))
# Try to open the distribution using different methods
flist = [
self.loadGOCODEDistribution, self.loadCODEDistribution,
self.loadSOFTDistribution
]
if typehint is not None:
if typehint.lower() == 'code': flist = [self.loadCODEDistribution]
elif typehint.lower() == 'gocode':
flist = [self.loadGOCODEDistribution]
elif typehint.lower() == 'soft':
flist = [self.loadSOFTDistribution]
else:
print(
"WARNING: Unrecognized distribution function type '{}' specified. Ignoring..."
.format(typehint))
found = False
for f in flist:
try:
f(filename)
found = True
break
except Exception as ex:
if typehint is not None:
print(ex)
traceback.print_exc()
if not found:
QMessageBox.critical(
self, 'Unrecognized type',
"The specified file is either corrupt or not a SOFT compatible distribution function."
)
self.exit()
return
self.loadDistributionUI()
def loadCODEDistribution(self, filename):
self.distfunc = CODEDistribution(filename)
self.ui.lblDistType.setText('CODE')
def loadGOCODEDistribution(self, filename):
self.distfunc = GOCODEDistribution(filename)
self.ui.lblDistType.setText('GO+CODE')
def loadSOFTDistribution(self, filename):
self.distfunc = SOFTDistribution(filename)
self.ui.lblDistType.setText('SOFT')
def loadDistributionUI(self):
"""
This function sets up certain parts of the UI relating
to the distribution function, after the function has been
loaded.
"""
nr = self.distfunc.getNr()
nmom = self.distfunc.getNmomentum()
self.maxP = self.distfunc.getMaxP()
if nr is not None: self.ui.lblNRadii.setText(str(nr))
else:
raise ValueError(
"Internal error: Unable to determine the number of radial points in distribution."
)
if nmom is not None:
self.ui.lblNMomentum.setText(str(self.distfunc.getNmomentum()))
else:
self.ui.lblNMomentum.setText('[Several]')
if self.maxP is not None:
self.ui.lblMaxP.setText(str(self.maxP) + 'mc')
else:
self.ui.lblMaxP.setText('[Several]')
self.ui.sliderRadius.setMaximum(nr - 1)
def plotSelection(self, vmin=None, vmax=None):
first = False
r = self.getRadius()
# 2D Ppar/Pperp distribution
if self.ui.rbDistParPerp.isChecked():
if self.currentPlotHandle is not None:
self.ax.clear()
PPAR, PPERP, F = self.getParPerpDistribution(r)
self.maxF = np.amax(F)
if self.logarithmicPlot:
self.maxF = np.power(10, self.maxF)
if vmax is None:
vmax = np.amax(F)
levels = None
if vmin is None:
if self.logarithmicPlot:
levels = np.linspace(vmax - 20, vmax, 50)
else:
levels = np.linspace(vmax * 1e-20, vmax, 50)
else:
if self.logarithmicPlot:
levels = np.linspace(vmin, vmax, 50)
else:
levels = np.linspace(vmin, vmax, 50)
self.currentPlotHandle = self.ax.contourf(PPAR,
PPERP,
F,
cmap='GeriMap',
levels=levels)
self.ax.set_xlabel(r'$p_\parallel$')
self.ax.set_ylabel(r'$p_\perp / mc$')
self.ax.set_xlim([-self.maxP, self.maxP])
self.ax.set_ylim([0, self.maxP])
# D P/XI distribution
elif self.ui.rbDistPXi.isChecked():
if self.currentPlotHandle is not None:
self.ax.clear()
P, XI, F = self.getPXiDistribution(r)
self.maxF = np.amax(F)
if self.logarithmicPlot:
self.maxF = np.power(10, self.maxF)
if vmax is None:
vmax = np.amax(F)
levels = None
if vmin is None:
if self.logarithmicPlot:
levels = np.linspace(vmax - 20, vmax, 50)
else:
levels = np.linspace(vmax * 1e-20, vmax, 50)
else:
if self.logarithmicPlot:
levels = np.linspace(vmin, vmax, 50)
else:
levels = np.linspace(vmin, vmax, 50)
self.currentPlotHandle = self.ax.contourf(P,
XI,
F,
cmap='GeriMap',
levels=levels)
self.ax.set_xlim([0, self.maxP])
self.ax.set_ylim([-1, 1])
self.ax.set_xlabel(r'$p$')
self.ax.set_ylabel(r'$\xi$')
# 1D angle-averaged distribution
elif self.ui.rbDist1D.isChecked():
if self.currentPlotHandle is None:
lsc = self.linestyles[len(self.plotHandles)]
self.currentPlotHandle, = self.ax.semilogy([], [],
lsc[1],
color=lsc[0],
linewidth=3)
first = True
h = self.currentPlotHandle
p, fp = self.getAngleAveragedDistribution(r)
h.set_data(p, fp)
self.ax.set_xlim([0, self.maxP])
self.setYLimit()
self.maxF = np.amax(fp)
self.ax.set_xlabel(r'$p$')
self.ax.set_ylabel(r'$f(p)$')
elif self.ui.rbCumCurrent.isChecked():
if self.currentPlotHandle is None:
lsc = self.linestyles[len(self.plotHandles)]
self.currentPlotHandle, = self.ax.plot([], [],
lsc[1],
color=lsc[0],
linewidth=3)
first = True
h = self.currentPlotHandle
_, p, j = self.distfunc.getCurrentDensity(r=r, cumulative=True)
if j[0, -1] < 0:
j = -j
h.set_data(p, j)
self.ax.set_xlim([0, self.maxP])
self.maxF = np.amax(j)
self.ax.set_ylim([0, self.maxF * 1.1])
self.ax.set_xlabel(r'$p$')
self.ax.set_ylabel(
r"$\int_0^p j(p')\,\mathrm{d}p'$ ($\mathrm{A}\cdot\mathrm{s}\cdot\mathrm{kg}^{-1}\cdot\mathrm{m}^{-3}$)"
)
else:
raise Exception("Unrecognized or no plot type selected.")
if first:
self.setAutomaticYLimit()
self.plotWindow.drawSafe()
def rbRunawayChanged(self, checked):
self.ui.lblRunawayPc.setEnabled(checked)
self.ui.tbRunawayPc.setEnabled(checked)
self.plotMomentsChanged(checked)
def rbSynchrotronChanged(self, checked):
self.ui.lblSynchB0.setEnabled(checked)
self.ui.lblSynchWavelength.setEnabled(checked)
self.ui.sbSynchMagneticField.setEnabled(checked)
self.ui.sbSynchWavelength.setEnabled(checked)
self.plotMomentsChanged(checked)
def plotMomentsChanged(self, checked):
if isinstance(self, QtWidgets.QRadioButton) and not checked:
self.ui.btnMomentProfile.setEnabled(False)
return
self.ui.btnMomentProfile.setEnabled(True)
self.plotSelection()
def plotCurrentDensity(self):
r, nr = self.distfunc.getCurrentDensity()
self.plotRadialProfile(r, -nr, label=r'$j(r)$ (A/m$^2$)')
def plotRadialDensity(self):
r, nr = self.distfunc.getRadialDensity()
self.plotRadialProfile(r, nr, label=r'$n_e(r)$ (m$^{-3}$)')
def plotMomentProfile(self):
mc = 9.109e-31 * 299792458.0
r = np.copy(self.distfunc._radii)[:, 0]
nr = np.zeros(r.shape)
for i in range(0, r.size):
print('Evaluating moment at r = {:.3f}...'.format(r[i]))
p = np.linspace(0, self.distfunc.getMaxP(), 1000)
P, XI, F = self.distfunc.eval(r[i], p)
if self.ui.gbMoments.isChecked():
F, _ = self.generateMoment(P, XI, F)
dp = np.zeros(p.shape)
dp[:-1] = np.diff(p)
dp[-1] = dp[-2]
xi = XI[:, 0]
dxi = np.zeros(xi.shape)
dxi[:-1] = np.diff(xi)
dxi[-1] = dxi[-2]
DP, DXI = np.meshgrid(dp, dxi)
nr[i] = np.sum(F * P**2 * DP * DXI * mc**3)
label = 'UNKNOWN'
if self.ui.rbSynchrotron.isChecked():
label = r'$P$ (W/m$^{3?}$)'
elif self.ui.rbRunaway.isChecked():
label = r'$n_{\rm RE}$ (m$^{-3}$)'
self.plotRadialProfile(r, nr, label=label)
def plotRadialProfile(self, r, nr, label):
self.radialPlotAx = self.radialPlotWindow.figure.add_subplot(111)
self.radialPlotAx.plot(r, nr, linewidth=2, color='k')
self.radialPlotAx.set_xlabel(r'$r$ (m)')
self.radialPlotAx.set_ylabel(label)
self.radialPlotAx.set_xlim([0, np.amax(r)])
self.radialPlotAx.set_ylim([0, np.amax(nr) * 1.1])
self.radialPlotWindow.drawSafe()
if not self.radialPlotWindow.isVisible():
self.radialPlotWindow.show()
def plotNow(self):
self.plotTypeChanged(True)
def plotTypeChanged(self, checked):
if not checked: return
# Enable/disable moments group box
self.ui.gbMoments.setEnabled(not (
self.ui.rbDist1D.isChecked() or self.ui.rbCumCurrent.isChecked()))
self.setupFigure()
def setAutomaticYLimit(self):
if self.maxF is None: return
else:
if self.logarithmicPlot or self.ax.get_yaxis().get_scale(
) == 'log':
self.ui.tbMinY.setText('{:.7e}'.format(self.maxF * 1e-30))
self.ui.tbMaxY.setText('{:.7e}'.format(self.maxF * 10))
else:
self.ui.tbMinY.setText('0')
self.ui.tbMaxY.setText('{:.7e}'.format(self.maxF * 1.1))
self.setYLimit()
def setYLimit(self):
minY, maxY = 0, 0
try:
minY = float(self.ui.tbMinY.text())
maxY = float(self.ui.tbMaxY.text())
except Exception as ex:
print(ex)
QMessageBox.critical(
self, 'Invalid Y-axis limits',
"The Y-axis limits are specified in a invalid format.")
return
if self.ui.rbDist1D.isChecked() or self.ui.rbCumCurrent.isChecked():
if self.logarithmicPlot:
self.ax.set_ylim([np.log10(minY), np.log10(maxY)])
else:
self.ax.set_ylim([minY, maxY])
else:
if self.logarithmicPlot:
self.plotSelection(vmin=np.log10(minY), vmax=np.log10(maxY))
else:
self.plotSelection(vmin=minY, vmax=maxY)
self.plotWindow.drawSafe()
def setupFigure(self):
self.plotWindow.figure.clear()
self.ax = self.plotWindow.figure.add_subplot(111)
self.currentPlotHandle = None
if not self.plotWindow.isVisible():
self.plotWindow.show()
self.plotSelection()
def sliderRadiusChanged(self):
"""
Function called when the value of the radius slider is changed.
"""
self.ui.lblRadius.setText(str(self.getRadius()))
self.plotSelection()
class BeamsizeMeasurement(QtWidgets.QMainWindow):
def __init__(self, argv):
QtWidgets.QMainWindow.__init__(self)
self.ui = BeamsizeMeasurement_design.Ui_BeamsizeMeasurement()
self.ui.setupUi(self)
filename = None
if len(argv) > 1:
raise Exception("Too many input arguments given to 'beamsize'.")
elif len(argv) == 1:
filename = argv[0]
self.plotWindow = PlotWindow()
self.GF = None
self.beamHandle = None
self.radius = None
self.imageContoursHandle = None
self.imageHandle = None
self.overlay = None
self.overlayHandle = None
self.toggleEnabled(False)
self.bindEvents()
if filename is not None and os.path.isfile(filename):
self.loadFile(filename)
self.setupImage()
def bindEvents(self):
self.ui.btnBrowse.clicked.connect(self.openFile)
self.ui.btnBrowseOverlay.clicked.connect(self.openOverlay)
self.ui.btnSaveImage.clicked.connect(self.saveImage)
self.ui.btnSaveBoth.clicked.connect(self.saveBoth)
self.ui.btnSaveProfile.clicked.connect(self.saveProfile)
self.ui.sliderBeamsize.valueChanged.connect(self.sliderBeamsizeChanged)
self.ui.sliderIntensity.valueChanged.connect(
self.sliderIntensityChanged)
self.ui.tbRadialProfile.textChanged.connect(self.radialProfileChanged)
self.ui.gbRadialProfile.toggled.connect(self.fullProfileToggled)
self.ui.sliderIntensity.valueChanged.connect(self.intensityChanged)
self.ui.sliderOverlay.valueChanged.connect(self.updateOverlay)
self.ui.cbContour.toggled.connect(self.contourToggled)
self.ui.actionSave.triggered.connect(self.saveImage)
self.ui.actionExit.triggered.connect(self.exit)
def closeEvent(self, event):
self.exit()
def contourToggled(self):
enabled = self.ui.cbContour.isChecked()
self.ui.sliderIntensity.setEnabled(enabled)
self.ui.lblIntensity_desc.setEnabled(enabled)
self.ui.lblIntensity.setEnabled(enabled)
self.ui.lblIntensity0.setEnabled(enabled)
self.ui.lblIntensity20.setEnabled(enabled)
self.ui.lblIntensity40.setEnabled(enabled)
self.ui.lblIntensity60.setEnabled(enabled)
self.ui.lblIntensity80.setEnabled(enabled)
self.ui.lblIntensity100.setEnabled(enabled)
self.updateImage()
def exit(self):
self.plotWindow.close()
self.close()
def toggleEnabled(self, enabled):
self.ui.lblBeamRadius.setEnabled(enabled)
self.ui.sliderBeamsize.setEnabled(enabled)
self.ui.lblBeamsize.setEnabled(enabled)
self.ui.lblBeamsize_desc.setEnabled(enabled)
self.ui.lblBeamsize0.setEnabled(enabled)
self.ui.lblBeamsize20.setEnabled(enabled)
self.ui.lblBeamsize40.setEnabled(enabled)
self.ui.lblBeamsize60.setEnabled(enabled)
self.ui.lblBeamsize80.setEnabled(enabled)
self.ui.lblBeamsize100.setEnabled(enabled)
self.ui.sliderIntensity.setEnabled(enabled)
self.ui.lblIntensity.setEnabled(enabled)
self.ui.lblIntensity_desc.setEnabled(enabled)
self.ui.lblIntensity0.setEnabled(enabled)
self.ui.lblIntensity20.setEnabled(enabled)
self.ui.lblIntensity40.setEnabled(enabled)
self.ui.lblIntensity60.setEnabled(enabled)
self.ui.lblIntensity80.setEnabled(enabled)
self.ui.lblIntensity100.setEnabled(enabled)
self.ui.sliderOverlay.setEnabled(enabled)
self.ui.lblOverlay_desc.setEnabled(enabled)
self.ui.tbOverlay.setEnabled(enabled)
self.ui.btnBrowseOverlay.setEnabled(enabled)
self.ui.lblOverlay0.setEnabled(enabled)
self.ui.lblOverlay20.setEnabled(enabled)
self.ui.lblOverlay40.setEnabled(enabled)
self.ui.lblOverlay60.setEnabled(enabled)
self.ui.lblOverlay80.setEnabled(enabled)
self.ui.lblOverlay100.setEnabled(enabled)
self.ui.gbRadialProfile.setEnabled(enabled)
def loadFile(self, filename):
self.ui.tbGreensFunction.setText(filename)
self.GF = Green(filename)
if not self.validateGreensFunction(self.GF):
return
# Store radii in centimeters
self.radius = (self.GF._r - self.GF._r[0]) * 100.0
self.toggleEnabled(True)
self.ui.sliderBeamsize.setMaximum(self.GF.nr - 1)
self.ui.sliderBeamsize.setSliderPosition(self.GF.nr - 1)
self.updateBeamRadiusLabel()
self.setupRadialProfile()
def openFile(self):
filename, _ = QFileDialog.getOpenFileName(
parent=self,
caption="Open SOFT Green's function file",
filter="SOFT Green's function (*.mat *.h5 *.hdf5);;All files (*.*)"
)
if filename:
self.loadFile(filename)
def loadOverlay(self, filename):
self.ui.tbOverlay.setText(filename)
self.overlay = mpimg.imread(filename)
if self.overlayHandle is not None:
self.overlayHandle.remove()
a = (self.ui.sliderOverlay.value()) / 100.0
self.overlayHandle = self.imageAx.imshow(self.overlay,
alpha=a,
extent=[-1, 1, -1, 1])
self.plotWindow.drawSafe()
def openOverlay(self):
filename, _ = QFileDialog.getOpenFileName(
parent=self,
caption="Open overlay image",
filter="Portable Network Graphics (*.png)")
if filename:
self.loadOverlay(filename)
def updateBeamRadiusLabel(self):
v = self.ui.sliderBeamsize.value()
r = self.radius[v]
p = int(np.round((v / self.radius.size) * 100.0))
self.ui.lblBeamRadius.setText('{0:.1f} cm'.format(r))
self.ui.lblBeamsize.setText('{0}%'.format(p))
def validateGreensFunction(self, gf):
if gf.getFormat() != 'rij':
QMessageBox.critical(
self, 'Invalid input file',
"The specified Green's function is not of the appropriate format. Expected 'rij', got {0}."
.format(gf.getFormat()))
return False
return True
def sliderBeamsizeChanged(self):
self.updateBeamRadiusLabel()
if self.imageHandle is not None:
self.updateRadialProfile()
def sliderIntensityChanged(self):
v = self.ui.sliderIntensity.value()
self.ui.lblIntensity.setText('{0}%'.format(v))
i = float(v) / 100.0
def setupRadialProfile(self):
self.radialProfileLayout = QtWidgets.QVBoxLayout(
self.ui.widgetRadialProfile)
self.radialProfileCanvas = FigureCanvas(Figure())
self.radialProfileLayout.addWidget(self.radialProfileCanvas)
f = self.getRadialProfile()
self.radialProfileAx = self.radialProfileCanvas.figure.subplots()
self.radialProfileHandle, = self.radialProfileAx.plot(self.radius, f)
self.radialProfileAx.set_xlim([0, self.radius[-1]])
self.radialProfileAx.set_ylim([0, 1.2])
self.radialProfileAx.set_xlabel(r'$r$ (cm)')
self.radialProfileAx.set_ylabel(r'Radial density')
self.radialProfileAx.figure.tight_layout(pad=4.5)
def updateRadialProfile(self):
f = self.getRadialProfile()
self.radialProfileHandle.set_ydata(f)
maxf = np.amax(f)
if maxf == 0:
self.radialProfileAx.set_ylim([0, 1])
else:
self.radialProfileAx.set_ylim([0, np.amax(f) * 1.2])
self.updateImage()
self.radialProfileCanvas.draw()
def getRadialProfile(self):
s = self.ui.tbRadialProfile.toPlainText().strip()
x = self.radius / self.radius[-1]
f0 = np.zeros(self.radius.shape)
a = self.radius[self.ui.sliderBeamsize.value()] / self.radius[-1]
if not s:
f = np.ones(self.radius.shape)
else:
# Parse string
f = None
lcls = {'a': a}
try:
f = evaluateExpression(s, x, lcls=lcls)
except Exception as ex:
return np.zeros(self.radius.shape)
# Set negative values to zero
f = np.where(f < 0, f0, f)
# Apply step function
f = np.where(x < a, f, f0)
return f
def radialProfileChanged(self):
self.updateRadialProfile()
def setupImage(self):
self.imageAx = self.plotWindow.figure.add_subplot(111)
dummy = np.zeros(self.GF._pixels)
self.imageHandle = self.imageAx.imshow(dummy,
cmap='GeriMap',
interpolation=None,
clim=(0, 1),
extent=[-1, 1, -1, 1])
self.imageAx.get_xaxis().set_visible(False)
self.imageAx.get_yaxis().set_visible(False)
if not self.plotWindow.isVisible():
self.plotWindow.show()
self.updateImage()
def updateImage(self):
# Generate image
f = self.getRadialProfile()
img = np.zeros(self.GF._pixels)
for i in range(0, len(self.radius)):
img += self.GF[i, :, :] * f[i]
img = img.T / np.amax(img)
if self.ui.gbRadialProfile.isChecked():
self.imageHandle.set_data(img)
else:
self.imageHandle.set_data(np.zeros(self.GF._pixels))
if self.imageContoursHandle is not None:
self.imageContoursHandle.remove()
self.imageContoursHandle = None
if self.ui.cbContour.isChecked():
threshold = self.ui.sliderIntensity.value() / 100.0
cntr = skimage.measure.find_contours(img.T, threshold)[0]
ipix, jpix = img.shape
i, j = cntr[:, 0], cntr[:, 1]
i = (i - ipix / 2) / ipix * 2
j = (-j + jpix / 2) / jpix * 2
self.imageContoursHandle, = self.imageAx.plot(i, j, 'w--')
self.plotWindow.drawSafe()
def fullProfileToggled(self):
self.updateImage()
def intensityChanged(self):
self.updateImage()
def updateOverlay(self):
if self.overlayHandle is not None:
a = float(self.ui.sliderOverlay.value()) / 100.0
self.overlayHandle.set_alpha(a)
self.plotWindow.drawSafe()
def saveImagePNG(self, filename=False):
"""
Save the currently displayed SOFT image to a PNG file.
"""
if filename is False:
filename, _ = QFileDialog.getSaveFileName(
self,
caption='Save PNG image',
filter='Portable Network Graphics (*.png)')
if filename:
f = self.getRadialProfile()
img = np.zeros(self.GF._pixels)
for i in range(0, len(self.radius)):
img += self.GF[i, :, :] * f[i]
img = img.T / np.amax(img)
cmap = plt.get_cmap('GeriMap')
im = Image.fromarray(np.uint8(cmap(img) * 255))
im.save(filename)
def saveImage(self, filename=False):
"""
Save the currently displayed SOFT image to a PNG file.
"""
if filename is False:
filename, _ = QFileDialog.getSaveFileName(
self,
caption='Save image',
filter=
'Portable Document Format (*.pdf);;Portable Network Graphics (*.png)'
)
if not filename: return
if filename.endswith('.png'):
self.saveImagePNG(filename=filename)
return
self.imageAx.set_axis_off()
self.plotWindow.figure.subplots_adjust(top=1,
bottom=0,
right=1,
left=0,
hspace=0,
wspace=0)
self.imageAx.get_xaxis().set_major_locator(
matplotlib.ticker.NullLocator())
self.imageAx.get_yaxis().set_major_locator(
matplotlib.ticker.NullLocator())
fcolor = self.plotWindow.figure.patch.get_facecolor()
self.plotWindow.canvas.print_figure(filename,
bbox_inches='tight',
pad_inches=0,
dpi=300)
def saveProfile(self, filename=False):
"""
Saves the current radial profile.
"""
if filename is False:
filename, _ = QFileDialog.getSaveFileName(
self,
caption='Save image',
filter='Portable Network Graphics (*.png)')
if not filename:
return
self.radialProfileCanvas.figure.canvas.print_figure(
filename, bbox_inches='tight')
def saveBoth(self):
filename, _ = QFileDialog.getSaveFileName(
self,
caption='Save both figures',
filter=
'Portable Document Form (*.pdf);;Portable Network Graphics (*.png);;Encapsulated Post-Script (*.eps);;Scalable Vector Graphics (*.svg)'
)
if not filename:
return
f = filename.split('.')
filename = str.join('.', f[:-1])
ext = f[-1]
if filename.endswith('_image') or filename.endswith('_super'):
filename = filename[:-6]
imgname = filename + '_image.' + ext
supname = filename + '_profile.' + ext
self.saveImage(filename=imgname)
self.saveProfile(filename=supname)
class DetectorCalibration(QtWidgets.QMainWindow):
def __init__(self, argv):
global COLORS
QtWidgets.QMainWindow.__init__(self)
self.ui = detcal_design.Ui_DetectorCalibration()
self.ui.setupUi(self)
self.magfield = None
self.image = None
self.plotWindow = PlotWindow()
self.toggleEnabled(False)
if len(argv) > 2:
QMessageBox.critical(
self, 'Too many input arguments',
'Too many input arguments were given. Expected at most 2 arguments.'
)
self.exit()
imagefile = None
meqfile = None
for arg in argv:
if arg.endswith('.png'):
imagefile = arg
elif arg.endswith('.h5') or arg.endswith('.mat') or args.endswith(
'.hdf5'):
meqfile = arg
else:
QMessageBox.critical(
self, 'Unrecognized input file',
'The given input file is of an unrecognized type: {0}'.
format(arg))
self.exit()
i, selindex = 0, 0
for clr, _ in COLORS.items():
self.ui.cbColor.addItem(clr)
if clr == 'white':
selindex = i
i += 1
self.ui.cbColor.setCurrentIndex(selindex)
if imagefile is not None:
self.loadImage(imagefile)
if meqfile is not None:
self.loadEquilibrium(meqfile)
self.bindEvents()
def bindEvents(self):
self.ui.btnBrowseEq.clicked.connect(self.openEquilibrium)
self.ui.btnBrowseImage.clicked.connect(self.openImage)
self.ui.btnRedraw.clicked.connect(self.updateWall)
self.ui.cbColor.currentIndexChanged.connect(self.toroidalChanged)
self.ui.sliderToroidal.valueChanged.connect(self.toroidalChanged)
self.ui.dsbLinewidth.valueChanged.connect(self.toroidalChanged)
def toggleEnabled(self, enabled=False):
self.ui.gbDetector.setEnabled(enabled)
self.ui.gbOverlay.setEnabled(enabled)
"""
self.ui.sliderToroidal.setEnabled(enabled)
self.ui.lblToroidal.setEnabled(enabled)
self.ui.lblTor0.setEnabled(enabled)
self.ui.lblTor90.setEnabled(enabled)
self.ui.lblTor180.setEnabled(enabled)
self.ui.lblTor270.setEnabled(enabled)
self.ui.lblTor360.setEnabled(enabled)
"""
def closeEvent(self, event):
self.exit()
def exit(self):
self.plotWindow.close()
self.close()
def openEquilibrium(self):
filename, _ = QFileDialog.getOpenFileName(
parent=self,
caption="Open SOFT magnetic equilibrium",
filter="SOFT Equilibrium Data (*.h5 *.mat)")
if filename:
self.loadEquilibrium(filename)
def loadEquilibrium(self, filename):
self.ui.tbEquilibrium.setText(filename)
self.magfield = MagneticField(filename)
self.toggleEnabled(True)
self.updateWall()
def openImage(self):
filename, _ = QFileDialog.getOpenFileName(
parent=self,
caption="Open SOFT magnetic equilibrium",
filter="Image (*.png)")
if filename:
self.loadImage(filename)
def loadImage(self, filename):
self.ui.tbImage.setText(filename)
self.image = mpimg.imread(filename)
if not self.plotWindow.isVisible():
self.plotWindow.show()
self.setImage(self.image)
def updateWall(self):
if not self.plotWindow.isVisible():
self.plotWindow.show()
try:
detpos = np.array([
float(self.ui.tbPosX.text()),
float(self.ui.tbPosY.text()),
float(self.ui.tbPosZ.text())
])
detdir = np.array([
float(self.ui.tbDirX.text()),
float(self.ui.tbDirY.text()),
float(self.ui.tbDirZ.text())
])
detdir = detdir / np.linalg.norm(detdir)
visang = float(self.ui.tbVisang.text())
tiltAngle = float(self.ui.tbTilt.text())
self.setWall(detpos, detdir, visang, tiltAngle)
except ValueError as e:
msg = QMessageBox()
msg.setIcon(QMessageBox.Critical)
msg.setText(e.strerror)
msg.setWindowTitle('Runtime Error')
msg.setStandardButtons(QMessageBox.Ok)
msg.exec_()
def toroidalChanged(self):
self.plotWindow.ax = self.gen()
self.plotWindow.drawSafe()
def gen(self):
fig = self.plotWindow.figure
fig.clear()
ax = fig.add_subplot(111)
pixelscale = 1
if self.magfield is not None:
pixelscale = np.tan(self.visang) / 2
if self.image is not None:
h = self.image.shape[0]
w = self.image.shape[1]
extent = []
if h >= w:
extent = [
-(w / h) * pixelscale, (w / h) * pixelscale, -pixelscale,
pixelscale
]
else:
extent = [
-pixelscale, pixelscale, -(h / w) * pixelscale,
(h / w) * pixelscale
]
ax.imshow(self.image, extent=extent)
if self.magfield is not None:
toffset = self.ui.sliderToroidal.value()
clr = self.ui.cbColor.currentText()
linewidth = self.ui.dsbLinewidth.value()
plotwall(ax,
self.magfield.wall,
self.detpos,
self.detdir,
degreesStart=[toffset - 30, toffset + 210],
degreesEnd=[toffset - 29, toffset + 211],
tiltAngle=self.tiltAngle,
color=COLORS[clr],
linewidth=linewidth)
#plotwall(ax, self.magfield.wall, self.detpos, self.detdir, degreesStart=[toffset+190], degreesEnd=[toffset+350], rlim=0.46, spacing=5, tiltAngle=self.tiltAngle, color=COLORS[clr], linewidth=linewidth)
#plotwall(ax, self.magfield.wall, self.detpos, self.detdir, degreesStart=[toffset+190], degreesEnd=[toffset+350], zuplim=-0.4, spacing=3, tiltAngle=self.tiltAngle, color=COLORS[clr], linewidth=linewidth)
ax.set_xlim([-pixelscale, pixelscale])
ax.set_ylim([-pixelscale, pixelscale])
fig.subplots_adjust(top=1,
bottom=0,
right=1,
left=0,
hspace=0,
wspace=0)
ax.set_axis_off()
return ax
def plot(self):
self.plotWindow.ax = self.gen()
self.plotWindow.drawSafe()
def setImage(self, image):
self.image = image
self.plot()
def setWall(self, detpos, detdir, visang, tiltAngle):
self.detpos = detpos
self.detdir = detdir
self.visang = visang
self.tiltAngle = tiltAngle
self.plot()
class GreensFunctionR12(QtWidgets.QMainWindow):
labels = {
'gamma': r'$\gamma$',
'p': r'$p / mc$',
'ppar': r'$p_{\parallel} / mc$',
'pperp': r'$p_{\perp} / mc$',
'thetap': r'$\theta_{\rm p}$ (rad)',
'ithetap': r'$\theta_{\rm p}$ (rad)',
'xi': r'$\xi$'
}
poltypes = {
'Intensity': lambda f: (f[0], 0, 1),
'Horizontal polarization': lambda f: (0.5 * (f[0] + f[1]), 0, 1),
'Vertical polarization': lambda f: (0.5 * (f[0] - f[1]), 0, 1),
'Polarization angle': calcpolangle,
'Polarization fraction': calcpolfrac,
'Stokes I': lambda f: (f[0], 0, 1),
'Stokes Q': lambda f: (f[1], -1, 1),
'Stokes U': lambda f: (f[2], -1, 1),
'Stokes V': lambda f: (f[3], -1, 1)
}
def __init__(self, argv):
QtWidgets.QMainWindow.__init__(self)
self.ui = r12_design.Ui_R12()
self.ui.setupUi(self)
if len(argv) != 1:
raise Exception(
"The Green's function must be specified at startup")
self.filename = argv[0]
self.plotWindow = PlotWindow(width=800, height=600)
self.radialPlotWindow = PlotWindow(width=700, height=400)
self.ax = None
self.radialAx = None
self.colorbar = None
self.hasStokesParameters = False
for key in self.poltypes:
self.ui.cbRadiationType.addItem(key)
self.loadGreensFunction(self.filename)
self.setupFigure()
self.bindEvents()
def closeEvent(self, event):
self.exit()
def exit(self):
self.plotWindow.close()
self.radialPlotWindow.close()
self.close()
def bindEvents(self):
self.ui.rbSingleR.toggled.connect(self.toggleSingleSum)
self.ui.sliderRadius.valueChanged.connect(self.sliderRadiusChanged)
self.ui.sliderWavelength.valueChanged.connect(
self.sliderWavelengthChanged)
self.ui.cbRadiationType.currentTextChanged.connect(
self.cbRadiationTypeChanged)
self.ui.btnMark.clicked.connect(self.markSuperParticle)
self.ui.btnPlotRadialProfile.clicked.connect(self.plotRadialProfile)
self.ui.btnSave.clicked.connect(self.saveFigure)
def getGF(self):
F = None
FUNC = self.gf.FUNC
if self.format[0] == 'w':
idx = self.ui.sliderWavelength.value()
FUNC = FUNC[idx, :]
fmin, fmax = 0, 1
if self.hasStokesParameters:
FUNC, fmin, fmax = self.getPolFunction()
if self.format in ['12', 'w12']:
F = np.copy(FUNC).T
mx = np.amax(np.abs(F))
if mx != 0: F /= mx
elif self.ui.rbSingleR.isChecked():
idx = self.ui.sliderRadius.value()
F = np.copy(FUNC[idx]).T
mx = np.amax(np.abs(F))
if mx != 0: F /= mx
else:
F = np.sum(FUNC, axis=0).T
mx = np.amax(np.abs(F))
if mx != 0: F /= mx
# Locate super particle
self.getSuperParticle(F)
return F, fmin, fmax
def getPolFunction(self):
tp = self.ui.cbRadiationType.currentText()
fun = self.poltypes[tp]
return fun(self.gf.FUNC)
def getSuperParticle(self, F):
i, j = np.unravel_index(np.argmax(F), F.shape)
param1max = self.gf._param1[j]
param2max = self.gf._param2[i]
self.ui.lblParam1.setText('{:.4}'.format(param1max))
self.ui.lblParam2.setText('{:.4}'.format(param2max))
self.ui.lblSuperEnergy.setText('{:.4} mc²'.format(self.gf.GAMMA[i, j]))
self.ui.lblSuperPitch.setText('{:.4} rad'.format(self.gf.THETAP[i, j]))
return param1max, param2max, i, j
def loadGreensFunction(self, filename):
self.gf = Green(filename=filename)
fmt = self.gf.getFormat()
if fmt == '12':
self.ui.rbSingleR.setEnabled(False)
elif fmt not in ['r12', 'wr12']:
raise Exception(
"The Green's function has an invalid format: '{}'. Expected '(s)12' or '(s)r12'."
.format(fmt))
self.format = fmt
self.hasStokesParameters = self.gf.stokesparams
self.ui.cbRadiationType.setEnabled(self.gf.stokesparams)
self.ui.sliderRadius.setMaximum(self.gf.nr - 1)
self.ui.sliderRadius.setTickInterval(
max(1, int(np.round(self.gf.nr / 20))))
if fmt == 'wr12':
self.ui.sliderWavelength.setMaximum(self.gf._wavelengths.size - 1)
self.sliderWavelengthChanged()
else:
self.ui.sliderWavelength.hide()
self.ui.lblWavelengthLbl.hide()
self.ui.lblWavelength.hide()
self.ui.lblParam1Name.setText('Parameter 1 ({}):'.format(
self.gf._param1name))
self.ui.lblParam2Name.setText('Parameter 2 ({}):'.format(
self.gf._param2name))
self.sliderRadiusChanged()
def markSuperParticle(self):
F, _, _ = self.getGF()
m1, m2, _, _ = self.getSuperParticle(F)
self.ax.plot(m1,
m2,
'x',
color=(0, 1, 0),
markersize=10,
markeredgewidth=3)
self.plotWindow.drawSafe()
def plotRadialProfile(self):
mc = 9.109e-31 * 299792458
J = self.gf.getJacobian().T * mc**3
r = self.gf._r
fr = np.zeros((r.size, ))
F = self.gf.FUNC
for i in range(0, r.size):
rrr = np.sum(F[i, :, :] * J)
fr[i] = rrr
self.radialAx = self.radialPlotWindow.figure.add_subplot(111)
self.radialAx.plot(r, fr / (r - r[0])**2, linewidth=2)
self.radialAx.set_xlabel(r'$\mathrm{Major\ radius}\ \rho$')
self.radialAx.set_ylabel(
r'$\mathrm{Radial\ intensity}\ \partial I/\partial\rho$')
self.radialPlotWindow.drawSafe()
if not self.radialPlotWindow.isVisible():
self.radialPlotWindow.show()
def redrawFigure(self):
F, fmin, fmax = self.getGF()
r = np.linspace(fmin, fmax, 20)
self.ax.clear()
cmap = 'GeriMap'
if fmin < 0:
cmap = 'RdBu'
cntr = self.ax.contourf(self.gf._param1,
self.gf._param2,
F,
levels=r,
cmap=cmap,
vmin=fmin,
vmax=fmax)
if self.colorbar is None:
self.colorbar = self.plotWindow.figure.colorbar(cntr)
else:
self.colorbar.ax.clear()
self.colorbar = self.plotWindow.figure.colorbar(
cntr, cax=self.colorbar.ax)
self.ax.set_xlabel(self.labels[self.gf._param1name])
self.ax.set_ylabel(self.labels[self.gf._param2name])
self.plotWindow.drawSafe()
def saveFigure(self):
filename, _ = QFileDialog.getSaveFileName(
self,
caption='Save figure',
filter=
'Portable Document Format (*.pdf);;Encapsulated PostScript (*.eps);;Portable Network Graphics (*.png);;All files (*.*)'
)
if filename:
self.plotWindow.canvas.print_figure(filename, bbox_inches='tight')
def setupFigure(self):
self.ax = self.plotWindow.figure.add_subplot(111)
if not self.plotWindow.isVisible():
self.plotWindow.show()
self.redrawFigure()
def cbRadiationTypeChanged(self):
self.redrawFigure()
def sliderRadiusChanged(self):
idx = self.ui.sliderRadius.value()
self.ui.lblRIndex.setText(str(idx))
self.ui.lblR.setText('r = {:.4}'.format(self.gf._r[idx]))
if self.ax is not None:
self.redrawFigure()
def sliderWavelengthChanged(self):
idx = self.ui.sliderWavelength.value()
self.ui.lblWavelength.setText('{} nm'.format(
self.gf._wavelengths[idx] / 1e-9))
if self.ax is not None:
self.redrawFigure()
def toggleSingleSum(self):
enbl = self.ui.rbSingleR.isChecked()
self.ui.sliderRadius.setEnabled(enbl)
self.ui.lblRIndex.setEnabled(enbl)
self.ui.lblR.setEnabled(enbl)
self.redrawFigure()
class SingleEnergyPitchIJ(QtWidgets.QMainWindow):
def __init__(self, argv):
QtWidgets.QMainWindow.__init__(self)
self.ui = SingleEnergyPitchIJ_design.Ui_SingleEnergyPitchIJ()
self.ui.setupUi(self)
filename = None
if len(argv) > 1:
raise Exception("Too many input arguments given to 's12ij'.")
elif len(argv) == 1:
filename = argv[0]
self.plotWindow = PlotWindow()
self.GF = None
self.GFintensity = None
self.superPlotCanvas = None
self.superPlotLayout = None
self.superPlotAx = None
self.superPlotHandle = None
self.superPlotDomHandle = None
self.pitchAngles = None
self.overlay = None
self.overlayHandle = None
self.toggleEnabled(False)
self.bindEvents()
self.gerimap, _ = registerGeriMap(None)
if filename is not None and os.path.isfile(filename):
self.loadFile(filename)
def bindEvents(self):
self.ui.btnBrowse.clicked.connect(self.openFile)
self.ui.btnBrowseOverlay.clicked.connect(self.openOverlay)
self.ui.btnSaveImage.clicked.connect(self.saveImage)
self.ui.btnSaveSuper.clicked.connect(self.saveSuper)
self.ui.btnSaveBoth.clicked.connect(self.saveBoth)
self.ui.sliderEnergy.valueChanged.connect(self.energyChanged)
self.ui.sliderPitchAngle.valueChanged.connect(
self.pitchAngleParameterChanged)
self.ui.sliderOverlay.valueChanged.connect(self.overlaySliderChanged)
self.ui.cbUnderlay.toggled.connect(self.toggleOverlayType)
def closeEvent(self, event):
self.exit()
def exit(self):
self.plotWindow.close()
self.close()
def setupSuperPlot(self):
ymax = 1.1
z = np.zeros(self.pitchAngles.shape)
self.superPlotLayout = QtWidgets.QVBoxLayout(self.ui.widgetDistPlot)
self.superPlotCanvas = FigureCanvas(Figure())
self.superPlotLayout.addWidget(self.superPlotCanvas)
self.superPlotAx = self.superPlotCanvas.figure.subplots()
self.superPlotHandle, = self.superPlotAx.plot(self.pitchAngles,
z,
linewidth=2)
self.superPlotDomHandle, = self.superPlotAx.plot([0, 0], [0, ymax],
'k--')
self.superPlotAx.set_xlim([self.pitchAngles[0], self.pitchAngles[-1]])
self.superPlotAx.set_ylim([0, ymax])
self.superPlotAx.get_yaxis().set_ticks([])
self.superPlotAx.set_xlabel(r'$\theta_{\rm p}\ \mathrm{(rad)}$')
self.superPlotAx.set_ylabel(r'$f(\theta_{\rm p}) / f_{\rm max}$')
self.superPlotCanvas.figure.tight_layout(pad=2.5)
def updateSuperPlot(self, f=None):
ei = self.getEnergyIndex()
if f is None:
f = self.getDistributionFunction()
superParticle = self.GFintensity[ei, :] * f
superParticle = superParticle / np.amax(superParticle)
maxpitch = self.pitchAngles[np.argmax(superParticle)]
self.ui.lblDomPitch.setText('{0:.3f} rad'.format(maxpitch))
self.superPlotHandle.set_ydata(superParticle)
self.superPlotDomHandle.set_xdata([maxpitch, maxpitch])
self.superPlotCanvas.draw()
self.superPlotCanvas.flush_events()
def getDistributionFunction(self):
C = self.ui.sliderPitchAngle.value()
f = np.exp(C * self.cosPitchAngles)
return f
def energyChanged(self):
ei = self.ui.sliderEnergy.value()
if len(self.GF._param1.shape) == 2:
self.ui.lblEnergy.setText('{0:.2f}'.format(self.GF._param1[0][ei]))
else:
self.ui.lblEnergy.setText('{0:.2f}'.format(self.GF._param1[ei]))
f = self.getDistributionFunction()
self.updateSuperPlot(f=f)
self.updateImage(f=f)
def pitchAngleParameterChanged(self):
self.ui.lblPitchAngle.setText(str(self.ui.sliderPitchAngle.value()))
f = self.getDistributionFunction()
self.updateSuperPlot(f=f)
self.updateImage(f=f)
def toggleEnabled(self, enabled=False):
self.ui.lblREEnergy.setEnabled(enabled)
self.ui.lblREPitchAngle.setEnabled(enabled)
self.ui.lblEnergy.setEnabled(enabled)
self.ui.lblPitchAngle.setEnabled(enabled)
self.ui.sliderEnergy.setEnabled(enabled)
self.ui.sliderPitchAngle.setEnabled(enabled)
self.ui.lblEnergyMin.setEnabled(enabled)
self.ui.lblEnergyMax.setEnabled(enabled)
self.ui.lblPitchAngleMin.setEnabled(enabled)
self.ui.lblPitchAngleMax.setEnabled(enabled)
self.ui.lblOverlay.setEnabled(enabled)
self.ui.lblOverlayMin.setEnabled(enabled)
self.ui.lblOverlayMax.setEnabled(enabled)
self.ui.lblOverlay25.setEnabled(enabled)
self.ui.lblOverlay50.setEnabled(enabled)
self.ui.lblOverlay75.setEnabled(enabled)
self.ui.btnBrowseOverlay.setEnabled(enabled)
self.ui.tbOverlay.setEnabled(enabled)
self.ui.sliderOverlay.setEnabled(enabled)
self.ui.widgetDistPlot.setEnabled(enabled)
def loadFile(self, filename):
self.ui.tbFilename.setText(filename)
self.GF = Green(filename)
if not self.validateGreensFunction(self.GF):
return
self.toggleEnabled(True)
if len(self.GF._param2.shape) == 2:
self.pitchAngles = np.abs(self.GF._param2[0])
else:
self.pitchAngles = np.abs(self.GF._param2)
self.cosPitchAngles = np.cos(self.pitchAngles)
# Sum all pixels of each image
self.GFintensity = np.sum(self.GF.FUNC, axis=(2, 3)) * np.sin(
self.pitchAngles)
self.setupEnergySlider()
self.setupSuperPlot()
self.setupImage()
f = self.getDistributionFunction()
self.updateSuperPlot(f=f)
self.updateImage(f=f)
def loadOverlay(self, filename):
self.ui.tbOverlay.setText(filename)
self.overlay = mpimg.imread(filename)
self.setupOverlay()
def openFile(self):
filename, _ = QFileDialog.getOpenFileName(
parent=self,
caption="Open SOFT Green's function file",
filter="SOFT Green's function (*.mat *.h5 *.hdf5);;All files (*.*)"
)
if filename:
self.loadFile(filename)
def openOverlay(self):
filename, _ = QFileDialog.getOpenFileName(
parent=self,
caption="Open image overlay",
filter="Portable Network Graphics (*.png)")
if filename:
self.loadOverlay(filename)
def validateGreensFunction(self, gf):
if gf.getFormat() != '12ij':
QMessageBox.critical(
self, 'Invalid input file',
"The specified Green's function is not of the appropriate format. Expected '12ij', got '{0}'."
.format(gf.getFormat()))
return False
pn = gf.getParameterName('1')
if pn != 'gamma' and pn != 'p':
QMessageBox.critical(
self, 'Invalid input file',
"The first momentum parameter has an invalid type: '{0}'. Expected either 'gamma' or 'p'."
.format(pn))
return False
return True
def saveImage(self):
self.doSaveImage()
def saveSuper(self):
self.doSaveSuper()
def doSaveImage(self, filename=None):
if filename is None:
filename, _ = QFileDialog.getSaveFileName(
self,
caption='Save synthetic image',
filter=
'Portable Document Form (*.pdf);;Portable Network Graphics (*.png);;Encapsulated Post-Script (*.eps);;Scalable Vector Graphics (*.svg)'
)
if not filename:
return
self.imageAx.set_axis_off()
self.plotWindow.figure.subplots_adjust(top=1,
bottom=0,
right=1,
left=0,
hspace=0,
wspace=0)
self.imageAx.get_xaxis().set_major_locator(
matplotlib.ticker.NullLocator())
self.imageAx.get_yaxis().set_major_locator(
matplotlib.ticker.NullLocator())
fcolor = self.plotWindow.figure.patch.get_facecolor()
self.plotWindow.canvas.print_figure(filename,
bbox_inches='tight',
pad_inches=0,
dpi=300)
def doSaveSuper(self, filename=None):
if filename is None:
filename, _ = QFileDialog.getSaveFileName(
self,
caption='Save super particle',
filter=
'Portable Document Form (*.pdf);;Portable Network Graphics (*.png);;Encapsulated Post-Script (*.eps);;Scalable Vector Graphics (*.svg)'
)
if not filename:
return
self.superPlotCanvas.print_figure(filename, bbox_inches='tight')
def saveBoth(self):
filename, _ = QFileDialog.getSaveFileName(
self,
caption='Save both figures',
filter=
'Portable Document Form (*.pdf);;Portable Network Graphics (*.png);;Encapsulated Post-Script (*.eps);;Scalable Vector Graphics (*.svg)'
)
if not filename:
return
f = filename.split('.')
filename = str.join('.', f[:-1])
ext = f[-1]
if filename.endswith('_image') or filename.endswith('_super'):
filename = filename[:-6]
imgname = filename + '_image.' + ext
supname = filename + '_super.' + ext
self.doSaveImage(filename=imgname)
self.doSaveSuper(filename=supname)
def setupEnergySlider(self):
if len(self.GF._param1.shape) == 2:
vmin, vmax, vn = self.GF._param1[0][0], self.GF._param1[0][
-1], self.GF._param1[0].size
else:
vmin, vmax, vn = self.GF._param1[0], self.GF._param1[
-1], self.GF._param1.size
if self.GF.getParameterName('1') == 'gamma':
self.ui.lblREEnergy.setText('Runaway energy (mc²)')
else:
self.ui.lblREEnergy.setText('Runaway momentum (mc)')
self.ui.lblEnergyMin.setText('{0:.2f}'.format(vmin))
self.ui.lblEnergyMax.setText('{0:.2f}'.format(vmax))
self.ui.lblEnergy.setText('{0:.2f}'.format(vmin))
self.ui.sliderEnergy.setMinimum(0)
self.ui.sliderEnergy.setMaximum(vn - 1)
self.ui.sliderEnergy.setSingleStep(1)
def getEnergyIndex(self):
return self.ui.sliderEnergy.value()
def setupImage(self):
lbl = ''.join(random.choices(string.ascii_uppercase, k=4))
self.imageAx = self.plotWindow.figure.add_subplot(111, label=lbl)
a = 1
if self.overlayHandle is not None:
self.setupOverlay()
if not self.ui.cbUnderlay.isChecked():
a = 1 - (self.ui.sliderOverlay.value()) / 100.0
dummy = np.zeros(self.GF._pixels)
self.imageHandle = self.imageAx.imshow(dummy,
cmap=self.gerimap,
alpha=a,
interpolation=None,
clim=(0, 1),
extent=[-1, 1, -1, 1],
zorder=1)
self.imageAx.get_xaxis().set_visible(False)
self.imageAx.get_yaxis().set_visible(False)
if not self.plotWindow.isVisible():
self.plotWindow.show()
self.plotWindow.drawSafe()
def setupOverlay(self):
if self.overlayHandle is not None:
self.overlayHandle.remove()
self.overlayHandle = self.imageAx.imshow(self.overlay,
extent=[-1, 1, -1, 1],
zorder=0)
self.updateImage()
self.plotWindow.drawSafe()
def toggleOverlayType(self):
ic = self.ui.cbUnderlay.isChecked()
if not ic:
self.gerimap, _ = registerGeriMap(None)
else:
self.gerimap, _ = registerGeriMap(transparencyThreshold=0.25)
self.ui.sliderOverlay.setEnabled(not ic)
self.ui.lblOverlayMin.setEnabled(not ic)
self.ui.lblOverlay25.setEnabled(not ic)
self.ui.lblOverlay50.setEnabled(not ic)
self.ui.lblOverlay75.setEnabled(not ic)
self.ui.lblOverlayMax.setEnabled(not ic)
self.setupImage()
f = self.getDistributionFunction()
self.updateSuperPlot(f=f)
self.updateImage(f=f)
def overlaySliderChanged(self):
self.updateImage()
def updateImage(self, f=None):
ei = self.getEnergyIndex()
if f is None:
f = self.getDistributionFunction()
g = self.GF[ei, :, :, :]
I = 0
for i in range(0, f.size):
I += g[i, :, :] * f[i]
I = I.T / np.amax(I)
self.imageHandle.set_data(I)
if self.ui.cbUnderlay.isChecked() or self.overlayHandle is None:
self.imageHandle.set_alpha(1)
else:
a = float(self.ui.sliderOverlay.value()) / 100.0
self.imageHandle.set_alpha(1 - a)
self.plotWindow.drawSafe()
class GreensFunctionIJ(QtWidgets.QMainWindow):
WIDTH = 550
HEIGHT = 450
def __init__(self, argv):
QtWidgets.QMainWindow.__init__(self)
self.setupUi()
if len(argv) != 1:
raise Exception(
"The Green's function must be specified at startup")
self.filename = argv[0]
# Create plot window
self.plotWindow = PlotWindow()
self.imageAx = None
self.overlayHandle = None
# Combobox used for select polarization quantity to plot
self.stokesbox = None
# Overlay controls
self.tbOverlay = None
self.btnOverlay = None
self.lblOverlaySlider = None
self.sliderOverlay = None
# Load Green's function
self.gf = Green(filename=self.filename)
nparams, dims = self.classifyGreensFunction(self.gf)
self.buildControls(dims, self.gf)
self.setDetails()
self.setWindowTitle("Green's function with image")
self.setupImage()
def closeEvent(self, event):
self.exit()
def exit(self):
self.plotWindow.close()
self.close()
def buildControls(self, dims, gf):
i = 0
for d in dims:
if d == '1':
self.buildControl(index=i,
coordname=self.getCoordinateName(
gf._param1name),
vals=gf._param1)
elif d == '2':
self.buildControl(index=i,
coordname=self.getCoordinateName(
gf._param2name),
vals=gf._param2)
elif d == 'r':
self.buildControl(index=i, coordname='Radius', vals=gf._r)
elif d == 's':
self.buildStokes(index=i, coordname='Polarization quantity')
elif d == 'w':
self.buildControl(index=i,
coordname='Wavelength',
vals=gf._wavelengths)
else:
raise Exception(
"Unrecognized or unsupported Green's function format: '{}'."
.format(d))
i += 1
self.buildOverlay(index=i)
def buildControl(self, index, coordname, vals):
vmin = np.amin(vals)
vmax = np.amax(vals)
vn = vals.size
self.paramValues.append(vals)
gb = QtWidgets.QGroupBox(self.centralwidget)
gb.setTitle(coordname)
self.controlGroupboxes.append(gb)
vl = QtWidgets.QVBoxLayout(gb)
lbl = QtWidgets.QLabel(gb)
font = QtGui.QFont()
font.setPointSize(14)
lbl.setFont(font)
lbl.setText('{}'.format(vmin))
lbl.setAlignment(QtCore.Qt.AlignRight | QtCore.Qt.AlignTrailing
| QtCore.Qt.AlignVCenter)
self.paramLabels.append(lbl)
hs = QtWidgets.QSlider(gb)
hs.setOrientation(QtCore.Qt.Horizontal)
hs.setMinimum(0)
hs.setMaximum(vn - 1)
hs.setTickPosition(QtWidgets.QSlider.TicksBelow)
hs.setTickInterval(1)
self.paramSliders.append(hs)
hs.valueChanged.connect(self.sliderChanged)
vl.addWidget(lbl)
vl.addWidget(hs)
# Insert groupbox into window
idx = self.verticalLayout.count() - 1
self.verticalLayout.insertWidget(idx, gb)
self.HEIGHT += gb.height()
self.resize(self.WIDTH, self.HEIGHT)
def buildOverlay(self, index):
hl = QtWidgets.QHBoxLayout()
lbl = QtWidgets.QLabel(self.centralwidget)
lbl.setText("Overlay:")
tb = QtWidgets.QLineEdit(self.centralwidget)
tb.setReadOnly(True)
self.tbOverlay = tb
btn = QtWidgets.QPushButton(self.centralwidget)
btn.setText('Browse...')
self.btnOverlay = btn
self.btnOverlay.clicked.connect(self.openOverlay)
lblOverlaySlider = QtWidgets.QLabel(self.centralwidget)
lblOverlaySlider.setText('50%')
lblOverlaySlider.setAlignment(QtCore.Qt.AlignRight)
self.lblOverlaySlider = lblOverlaySlider
slider = QtWidgets.QSlider(self.centralwidget)
slider.setOrientation(QtCore.Qt.Horizontal)
slider.setMinimum(0)
slider.setMaximum(100)
slider.setValue(50)
slider.setTickPosition(QtWidgets.QSlider.TicksBelow)
slider.setTickInterval(5)
self.sliderOverlay = slider
slider.valueChanged.connect(self.sliderOverlayChanged)
hl.addWidget(tb)
hl.addWidget(btn)
self.verticalLayout.addWidget(lbl)
self.verticalLayout.addLayout(hl)
self.verticalLayout.addWidget(lblOverlaySlider)
self.verticalLayout.addWidget(slider)
self.HEIGHT += tb.height() + lbl.height() + slider.height(
) + lblOverlaySlider.height()
self.resize(self.WIDTH, self.HEIGHT)
def buildStokes(self, index, coordname):
cb = QtWidgets.QComboBox(self.centralwidget)
self.stokesbox = cb
cb.addItem("Polarization angle")
cb.addItem("Polarization fraction")
cb.addItem("Stokes I")
cb.addItem("Stokes Q")
cb.addItem("Stokes U")
cb.addItem("Stokes V")
cb.setCurrentIndex(2)
cb.currentIndexChanged.connect(self.redrawFigure)
self.verticalLayout.insertWidget(index, cb)
self.HEIGHT += cb.height()
self.resize(self.WIDTH, self.HEIGHT)
def classifyGreensFunction(self, gf):
dims = gf.format
if dims[-2] != 'i' or dims[-1] != 'j':
raise Exception(
"Invalid format of Green's function. Format string must end in 'ij'."
)
# Just pick out the interesting dimensions
dims = dims[:-2]
nparams = len(dims)
# Append Stoke's dimension?
if gf.stokesparams:
dims = 's' + dims
nparams += 1
return nparams, dims
def getCoordinateName(self, s):
"""
Converts a SOFT parameter name to a proper parameter label.
"""
if s == "gamma": return "Energy (γ)"
elif s == "p": return "Momentum (p)"
elif s == "ppar": return "Parallel momentum"
elif s == "pperp": return "Perpendicular momentum"
elif s == "thetap": return "Pitch angle (θ)"
elif s == "ithetap": return "Pitch angle (θ)"
elif s == "xi": return "Pitch (ξ)"
else: return "<UNKNOWN>"
def getSelectedGreensFunction(self):
"""
Returns the appropriate image to draw based on
how the GUI controls are set.
"""
Fmax, Fmin = None, 0
F = self.gf.FUNC
colormap = 'GeriMap'
# Compute select polarization quantity
if self.stokesbox is not None:
val = self.stokesbox.currentText()
if val == "Polarization angle":
F = 0.5 * np.arctan2(F[2], F[1]) * 180 / np.pi
F[np.where(F < -45)] += 180
Fmax, Fmin = 135, -45
colormap = 'RdBu'
elif val == "Polarization fraction":
F = np.sqrt(F[1]**2 + F[2]**2) / F[0]
F[np.where(np.isnan(F))] = 0
Fmax = 1
elif val == "Stokes I":
F = F[0]
elif val == "Stokes Q":
F = F[1]
fmax = max(abs(np.amax(F)), abs(np.amin(F)))
Fmax, Fmin = fmax, -fmax
colormap = 'RdBu'
elif val == "Stokes U":
F = F[2]
fmax = max(abs(np.amax(F)), abs(np.amin(F)))
Fmax, Fmin = fmax, -fmax
colormap = 'RdBu'
elif val == "Stokes V":
F = F[3]
fmax = max(abs(np.amax(F)), abs(np.amin(F)))
Fmax, Fmin = fmax, -fmax
colormap = 'RdBu'
else:
raise Exception(
"INTERNAL ERROR: Unrecognized polarization quantity select: '{}'."
.format(val))
if self.wfgb.isChecked(): # Draw with weight function
print('Not supported yet...')
return np.zeros(self.gf._pixels)
else:
indices = list()
for s in self.paramSliders:
F = F[s.value()]
if len(F.shape) != 2:
raise Exception(
'INTERNAL ERROR: F does not have the expected shape.')
if Fmax is None:
F /= np.amax(F)
Fmax = 1
return F.T, Fmax, Fmin, colormap
def loadOverlay(self, filename):
self.tbOverlay.setText(filename)
self.overlay = mpimg.imread(filename)
self.setupOverlay()
def openOverlay(self):
filename, _ = QFileDialog.getOpenFileName(
parent=self,
caption="Open image overlay",
filter="Portable Network Graphics (*.png)")
if filename:
self.loadOverlay(filename)
def redrawFigure(self):
F, Fmax, Fmin, cmap = self.getSelectedGreensFunction()
self.imageHandle.set_data(F)
self.imageHandle.set_clim(vmin=Fmin, vmax=Fmax)
self.imageHandle.set_cmap(cmap)
self.plotWindow.drawSafe()
def setDetails(self):
self.lblFilename.setText(self.filename)
# Green's function size
totsize = self.gf.FUNC.size * 8
fsize = totsize
prefixes = ['ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi', 'Yi']
pfi = -1
while fsize > 1024:
fsize /= 1024.0
pfi += 1
prefix = ''
if pfi >= 0:
prefix = prefixes[pfi]
self.lblGFSize.setText('{:.1f} {:s}B ({:d} bytes)'.format(
fsize, prefix, totsize))
# Format
self.lblFormat.setText(self.gf.format)
# Number of pixels
hpix, vpix = self.gf._pixels[0], self.gf._pixels[1]
self.lblPixels.setText('{} × {} pixels'.format(hpix, vpix))
def setupImage(self):
self.imageAx = self.plotWindow.figure.add_subplot(111)
F, Fmax, Fmin, cmap = self.getSelectedGreensFunction()
self.imageHandle = self.imageAx.imshow(F,
cmap=cmap,
interpolation=None,
clim=(Fmin, Fmax),
extent=[-1, 1, -1, 1])
self.imageAx.get_xaxis().set_visible(False)
self.imageAx.get_yaxis().set_visible(False)
self.colorbar = self.plotWindow.figure.colorbar(self.imageHandle,
ax=self.imageAx)
if not self.plotWindow.isVisible():
self.plotWindow.show()
self.plotWindow.drawSafe()
def setupOverlay(self):
if self.overlayHandle is not None:
self.overlayHandle.remove()
val = self.sliderOverlay.value() / 100.0
self.overlayHandle = self.imageAx.imshow(self.overlay,
alpha=val,
extent=[-1, 1, -1, 1],
zorder=100)
self.plotWindow.drawSafe()
def setupUi(self):
self.resize(self.WIDTH, self.HEIGHT)
self.centralwidget = QtWidgets.QWidget(self)
self.centralwidget.setObjectName("centralwidget")
self.verticalLayout = QtWidgets.QVBoxLayout(self.centralwidget)
# Control groupboxes
self.controlGroupboxes = list()
self.paramLabels = list()
self.paramSliders = list()
self.paramValues = list()
# Information groupbox
gb = QtWidgets.QGroupBox(self.centralwidget)
gb.setTitle("Green's function details")
# (declare labels)
self.lblFilename = QtWidgets.QLabel(gb)
self.lblGFSize = QtWidgets.QLabel(gb)
self.lblFormat = QtWidgets.QLabel(gb)
self.lblPixels = QtWidgets.QLabel(gb)
vli = QtWidgets.QVBoxLayout(gb)
h1 = QtWidgets.QHBoxLayout()
h1.addWidget(QtWidgets.QLabel('File name:', gb))
h1.addWidget(self.lblFilename)
h2 = QtWidgets.QHBoxLayout()
h2.addWidget(QtWidgets.QLabel('Function size: ', gb))
h2.addWidget(self.lblGFSize)
h3 = QtWidgets.QHBoxLayout()
h3.addWidget(QtWidgets.QLabel('Pixels:', gb))
h3.addWidget(self.lblPixels)
h4 = QtWidgets.QHBoxLayout()
h4.addWidget(QtWidgets.QLabel('Format: ', gb))
h4.addWidget(self.lblFormat)
vli.addLayout(h1)
vli.addLayout(h2)
vli.addLayout(h3)
vli.addLayout(h4)
# Distribution function window
self.wfgb = QtWidgets.QGroupBox(self.centralwidget)
self.wfgb.setCheckable(True)
self.wfgb.setChecked(False)
self.wfgb.setTitle('Weight function')
vl = QtWidgets.QVBoxLayout(self.wfgb)
font = QtGui.QFont()
font.setFamily("Droid Sans Mono")
txt = QtWidgets.QPlainTextEdit(self.wfgb)
txt.setFont(font)
txt.setMaximumSize(QtCore.QSize(16777215, 300))
frm = QtWidgets.QFrame(self.wfgb)
frm.setMinimumSize(QtCore.QSize(0, 200))
frm.setFrameShape(QtWidgets.QFrame.StyledPanel)
frm.setFrameShadow(QtWidgets.QFrame.Raised)
vl.addWidget(txt)
vl.addWidget(frm)
self.verticalLayout.addWidget(gb)
self.verticalLayout.addWidget(self.wfgb)
self.setCentralWidget(self.centralwidget)
self.wfgb.toggled.connect(self.toggleWeightFunction)
def sliderChanged(self):
for i in range(0, len(self.paramSliders)):
idx = self.paramSliders[i].value()
val = self.paramValues[i][idx]
self.paramLabels[i].setText('{}'.format(val))
self.redrawFigure()
def sliderOverlayChanged(self):
val = self.sliderOverlay.value() / 100.0
self.lblOverlaySlider.setText('{}%'.format(self.sliderOverlay.value()))
if self.overlayHandle is not None:
self.overlayHandle.set_alpha(val)
self.plotWindow.drawSafe()
def toggleWeightFunction(self):
for gb in self.controlGroupboxes:
gb.setEnabled(not self.wfgb.isChecked())
