def __init__(self, parent=None): # initialization super(LaserTab, self).__init__() self.ui = Ui_LaserInterface() self.ui.setupUi(self) # set default values self.numTicks = 5 self.plotTitles = True self.lambda0 = 1.0e-06 self.w0 = 20.0e-06 self.w0_z = 0.0 self.zR = math.pi * self.w0 ** 2 / self.lambda0 # Rayleigh range [m] # Define the maximum order(s) of the Hermite expansion self.mMax = 24 # Create two instances of the Hermite expansion class self.hS1 = hermite.RbGaussHermiteMN(self.lambda0, self.w0, self.w0, 0.0) self.hS2 = hermite.RbGaussHermiteMN(self.lambda0, self.w0, self.w0, 0.0) # the laser pulse (to be instantiated later) self.myPulse = 0 self.pulseInitialized = False # the external fields (used to find GH coefficients of laser pulse) self.externalFields = False # plotting flags self.plotTitles = True # link the simple push buttons to appropriate methods self.ui.generateCoeffs.clicked.connect(self.generateCoeffs) self.ui.noTitles.clicked.connect(self.togglePlotTitles) # create a menu for saving files exportMenu = QtGui.QMenu(self) saveToCSV = QtGui.QAction("RadTrack CSV format", self) exportMenu.addAction(saveToCSV) saveToSDDS = QtGui.QAction("SRW SDDS format", self) exportMenu.addAction(saveToSDDS) # associate these actions with class methods saveToCSV.triggered.connect(self.saveToCSV) saveToSDDS.triggered.connect(self.saveToSDDS) # grab an existing button & insert the menu saveToFileButton = self.ui.saveToFile saveToFileButton.setMenu(exportMenu) saveToFileButton.setPopupMode(QtGui.QToolButton.InstantPopup) # create a menu for importing particle data importMenu = QtGui.QMenu(self) readFromCSV = QtGui.QAction("RadTrack CSV format", self) importMenu.addAction(readFromCSV) readFromSDDS = QtGui.QAction("SRW SDDS format", self) importMenu.addAction(readFromSDDS) # associate these actions with class methods readFromCSV.triggered.connect(self.readFromCSV) readFromSDDS.triggered.connect(self.readFromSDDS) # grab an existing button & insert the menu importFileButton = self.ui.importFile importFileButton.setMenu(importMenu) importFileButton.setPopupMode(QtGui.QToolButton.InstantPopup) # create a menu for laser pulse generation pulseMenu = QtGui.QMenu(self) paraxialGaussian = QtGui.QAction("Paraxial approx - Gaussian", self) pulseMenu.addAction(paraxialGaussian) # associate these actions with class methods paraxialGaussian.triggered.connect(self.paraxialGaussian) # grab an existing button & insert the menu pulseButton = self.ui.generatePulse pulseButton.setMenu(pulseMenu) pulseButton.setPopupMode(QtGui.QToolButton.InstantPopup) # create a menu for external fields extFieldsMenu = QtGui.QMenu(self) mirrorWithHole = QtGui.QAction("Mirror with hole", self) extFieldsMenu.addAction(mirrorWithHole) # associate these actions with class methods mirrorWithHole.triggered.connect(self.mirrorWithHole) # grab an existing button & insert the menu extFieldsButton = self.ui.externalFields extFieldsButton.setMenu(extFieldsMenu) extFieldsButton.setPopupMode(QtGui.QToolButton.InstantPopup) # specify physical constants self.c = 299792458.0 # speed of light [m/s] self.cSq = self.c ** 2 # speed of light squared self.cInv = 1.0 / self.c # one over the speed of light self.mu0 = 4.0e-07 * math.pi # permeability of free space self.eps0 = 1.0 / self.mu0 / self.cSq # permittivity of free space self.eMass = 9.10938215e-31 # electron mass [kG] self.eCharge = 1.602176487e-19 # elementary charge [C] self.eMassEV = self.eMass * self.cSq / self.eCharge # eMass [eV] # specify default values for all input fields self.ui.wavelength.setText("{:.0f}".format(self.lambda0 * 1.0e06) + " um") self.ui.waistSize.setText("{:.0f}".format(self.w0 * 1.0e06) + " um") self.ui.waistPosition.setText("{:.0f}".format(self.w0_z * 1.0e3) + " mm") self.unitsXY = "um" self.unitsZ = "mm" self.ui.unitsXY.setText(self.unitsXY) self.ui.unitsZ.setText(self.unitsZ) self.ui.numTicks.setText(str(self.numTicks)) # load up the table of coefficients self.ui.ghTable.setEditTriggers(QtGui.QAbstractItemView.CurrentChanged) self.ui.ghTable.setItem(0, 0, QtGui.QTableWidgetItem("1")) self.ui.ghTable.setItem(0, 1, QtGui.QTableWidgetItem("1")) maxLoop = max(self.mMax, 100) for iLoop in range(1, maxLoop): self.ui.ghTable.setItem(iLoop, 0, QtGui.QTableWidgetItem("0")) self.ui.ghTable.setItem(iLoop, 1, QtGui.QTableWidgetItem("0")) # file directories self.parent = parent if self.parent is None: self.parent = self self.parent.lastUsedDirectory = os.path.expanduser("~") self.fileExtension = ".sdds" self.exportToFile = self.saveToSDDS self.importFile = self.readFromSDDS # try to make the blank plotting regions look nice self.erasePlots() self.container = QtGui.QScrollArea(parent) self.container.setWidget(self)
class LaserTab(QtGui.QWidget): acceptsFileTypes = ["sdds", "csv"] defaultTitle = "Laser" task = "Create a laser beam" category = "beams" def __init__(self, parent=None): # initialization super(LaserTab, self).__init__() self.ui = Ui_LaserInterface() self.ui.setupUi(self) # set default values self.numTicks = 5 self.plotTitles = True self.lambda0 = 1.0e-06 self.w0 = 20.0e-06 self.w0_z = 0.0 self.zR = math.pi * self.w0 ** 2 / self.lambda0 # Rayleigh range [m] # Define the maximum order(s) of the Hermite expansion self.mMax = 24 # Create two instances of the Hermite expansion class self.hS1 = hermite.RbGaussHermiteMN(self.lambda0, self.w0, self.w0, 0.0) self.hS2 = hermite.RbGaussHermiteMN(self.lambda0, self.w0, self.w0, 0.0) # the laser pulse (to be instantiated later) self.myPulse = 0 self.pulseInitialized = False # the external fields (used to find GH coefficients of laser pulse) self.externalFields = False # plotting flags self.plotTitles = True # link the simple push buttons to appropriate methods self.ui.generateCoeffs.clicked.connect(self.generateCoeffs) self.ui.noTitles.clicked.connect(self.togglePlotTitles) # create a menu for saving files exportMenu = QtGui.QMenu(self) saveToCSV = QtGui.QAction("RadTrack CSV format", self) exportMenu.addAction(saveToCSV) saveToSDDS = QtGui.QAction("SRW SDDS format", self) exportMenu.addAction(saveToSDDS) # associate these actions with class methods saveToCSV.triggered.connect(self.saveToCSV) saveToSDDS.triggered.connect(self.saveToSDDS) # grab an existing button & insert the menu saveToFileButton = self.ui.saveToFile saveToFileButton.setMenu(exportMenu) saveToFileButton.setPopupMode(QtGui.QToolButton.InstantPopup) # create a menu for importing particle data importMenu = QtGui.QMenu(self) readFromCSV = QtGui.QAction("RadTrack CSV format", self) importMenu.addAction(readFromCSV) readFromSDDS = QtGui.QAction("SRW SDDS format", self) importMenu.addAction(readFromSDDS) # associate these actions with class methods readFromCSV.triggered.connect(self.readFromCSV) readFromSDDS.triggered.connect(self.readFromSDDS) # grab an existing button & insert the menu importFileButton = self.ui.importFile importFileButton.setMenu(importMenu) importFileButton.setPopupMode(QtGui.QToolButton.InstantPopup) # create a menu for laser pulse generation pulseMenu = QtGui.QMenu(self) paraxialGaussian = QtGui.QAction("Paraxial approx - Gaussian", self) pulseMenu.addAction(paraxialGaussian) # associate these actions with class methods paraxialGaussian.triggered.connect(self.paraxialGaussian) # grab an existing button & insert the menu pulseButton = self.ui.generatePulse pulseButton.setMenu(pulseMenu) pulseButton.setPopupMode(QtGui.QToolButton.InstantPopup) # create a menu for external fields extFieldsMenu = QtGui.QMenu(self) mirrorWithHole = QtGui.QAction("Mirror with hole", self) extFieldsMenu.addAction(mirrorWithHole) # associate these actions with class methods mirrorWithHole.triggered.connect(self.mirrorWithHole) # grab an existing button & insert the menu extFieldsButton = self.ui.externalFields extFieldsButton.setMenu(extFieldsMenu) extFieldsButton.setPopupMode(QtGui.QToolButton.InstantPopup) # specify physical constants self.c = 299792458.0 # speed of light [m/s] self.cSq = self.c ** 2 # speed of light squared self.cInv = 1.0 / self.c # one over the speed of light self.mu0 = 4.0e-07 * math.pi # permeability of free space self.eps0 = 1.0 / self.mu0 / self.cSq # permittivity of free space self.eMass = 9.10938215e-31 # electron mass [kG] self.eCharge = 1.602176487e-19 # elementary charge [C] self.eMassEV = self.eMass * self.cSq / self.eCharge # eMass [eV] # specify default values for all input fields self.ui.wavelength.setText("{:.0f}".format(self.lambda0 * 1.0e06) + " um") self.ui.waistSize.setText("{:.0f}".format(self.w0 * 1.0e06) + " um") self.ui.waistPosition.setText("{:.0f}".format(self.w0_z * 1.0e3) + " mm") self.unitsXY = "um" self.unitsZ = "mm" self.ui.unitsXY.setText(self.unitsXY) self.ui.unitsZ.setText(self.unitsZ) self.ui.numTicks.setText(str(self.numTicks)) # load up the table of coefficients self.ui.ghTable.setEditTriggers(QtGui.QAbstractItemView.CurrentChanged) self.ui.ghTable.setItem(0, 0, QtGui.QTableWidgetItem("1")) self.ui.ghTable.setItem(0, 1, QtGui.QTableWidgetItem("1")) maxLoop = max(self.mMax, 100) for iLoop in range(1, maxLoop): self.ui.ghTable.setItem(iLoop, 0, QtGui.QTableWidgetItem("0")) self.ui.ghTable.setItem(iLoop, 1, QtGui.QTableWidgetItem("0")) # file directories self.parent = parent if self.parent is None: self.parent = self self.parent.lastUsedDirectory = os.path.expanduser("~") self.fileExtension = ".sdds" self.exportToFile = self.saveToSDDS self.importFile = self.readFromSDDS # try to make the blank plotting regions look nice self.erasePlots() self.container = QtGui.QScrollArea(parent) self.container.setWidget(self) def paraxialGaussian(self): # get input from text boxes self.lambda0 = convertUnitsStringToNumber(self.ui.wavelength.text(), "m") self.w0 = convertUnitsStringToNumber(self.ui.waistSize.text(), "m") self.w0_z = convertUnitsStringToNumber(self.ui.waistPosition.text(), "m") self.zR = math.pi * self.w0 ** 2 / self.lambda0 # horiz. Rayleigh range [m] # instantiate the laser pulse self.myPulse = hermite.RbGaussHermiteMN(self.lambda0, self.w0, self.w0, 0.0) # load up the coefficients mCoefs = np.zeros(8) nCoefs = np.zeros(8) for iLoop in range(8): mCoefs[iLoop] = float(self.ui.ghTable.item(iLoop, 0).text()) nCoefs[iLoop] = float(self.ui.ghTable.item(iLoop, 1).text()) self.myPulse.setMCoef(mCoefs) self.myPulse.setNCoef(nCoefs) # update flag accordingly self.pulseInitialized = True # generate the plots self.refreshPlots() def togglePlotTitles(self): if self.plotTitles == True: self.plotTitles = False else: self.plotTitles = True self.refreshPlots() def refreshPlots(self): # nothing to plot, if beam hasn't been initialized if self.pulseInitialized == False: return # get the specified units for plotting self.unitsXY = self.ui.unitsXY.text() self.unitsZ = self.ui.unitsZ.text() # get the number of tick marks self.numTicks = int(self.ui.numTicks.text()) # Specify the desired grid size self.numZ = 64 self.numX = 64 self.numCellsZX = self.numZ * self.numX # specify the min's and max's self.minZ = self.w0_z - 4.0 * self.zR self.maxZ = self.w0_z + 4.0 * self.zR self.minX = -8.0 * self.w0 self.maxX = 8.0 * self.w0 self.plotXY() self.plotZY() self.plotZX() def plotZX(self): # instance of the plot utility class myPlotUtils = radtrack.plot.RbPlotUtils.RbPlotUtils() zArr = np.zeros(self.numZ) xArr = np.zeros(self.numX) for iLoop in range(self.numZ): zArr[iLoop] = self.minZ + iLoop * (self.maxZ - self.minZ) / (self.numZ - 1) # print('zArr[', iLoop, '] = ', zArr[iLoop]) for jLoop in range(self.numX): xArr[jLoop] = self.minX + jLoop * (self.maxX - self.minX) / (self.numX - 1) # print('xArr[', jLoop, '] = ', xArr[jLoop]) # specify y position for plot yValue = 0.0 # Calculate Ex at the 2D array of x,y values zxEData = np.zeros((self.numX, self.numZ)) for iLoop in range(self.numZ): for jLoop in range(self.numX): zxEData[jLoop, iLoop] = np.real(self.myPulse.evalEnvelopeEx(xArr[jLoop], yValue, zArr[iLoop])) # print('zxEData[', iLoop, jLoop, '] = ', zxEData[iLoop, jLoop] # generate the xy plot canvas = self.ui.zxPlot.canvas canvas.ax.clear() levels = myPlotUtils.generateContourLevels(zxEData) canvas.ax.contourf( zArr * convertUnitsNumber(1, "m", self.unitsZ), xArr * convertUnitsNumber(1, "m", self.unitsXY), zxEData, levels, extent="none", aspect="equal", ) canvas.ax.axis( [ self.minZ * convertUnitsNumber(1, "m", self.unitsZ), self.maxZ * convertUnitsNumber(1, "m", self.unitsZ), self.minX * convertUnitsNumber(1, "m", self.unitsXY), self.maxX * convertUnitsNumber(1, "m", self.unitsXY), ] ) canvas.ax.xaxis.set_major_locator(plt.MaxNLocator(self.numTicks)) canvas.ax.yaxis.set_major_locator(plt.MaxNLocator(self.numTicks)) canvas.ax.set_xlabel("z [" + self.unitsZ + "]") canvas.ax.set_ylabel("x [" + self.unitsXY + "]") if self.plotTitles == True: canvas.ax.set_title( "ZX slice, at y={0:4.2f} [{1}]".format(yValue * convertUnitsNumber(1, "m", self.unitsXY), self.unitsXY) ) canvas.fig.tight_layout() canvas.fig.set_facecolor("w") canvas.draw() def plotZY(self): # instance of the plot utility class myPlotUtils = radtrack.plot.RbPlotUtils.RbPlotUtils() freq0 = self.c / self.lambda0 yLoc = 0.0 # Specify the desired grid size self.zyNumH = 64 self.zyNumV = 64 self.zyNumCells = self.zyNumH * self.zyNumV # specify the min's and max's self.zyMinH = -4.0 * self.lambda0 self.zyMaxH = 4 * self.lambda0 self.zyMinV = -4.0 * self.w0 self.zyMaxV = 4.0 * self.w0 zArr = np.zeros(self.zyNumH) yArr = np.zeros(self.zyNumV) for iLoop in range(self.zyNumH): zArr[iLoop] = self.zyMinH + iLoop * (self.zyMaxH - self.zyMinH) / (self.zyNumH - 1) # print('zArr[', iLoop, '] = ', zArr[iLoop]) for jLoop in range(self.zyNumV): yArr[jLoop] = self.zyMinV + jLoop * (self.zyMaxV - self.zyMinV) / (self.zyNumV - 1) # print('xArr[', jLoop, '] = ', xArr[jLoop]) # Choose values of x,t for plot xValue = 0.0 tValue = 0.0 # Calculate Ex at the 2D array of x,y values zyEData = np.zeros((self.zyNumV, self.zyNumH)) for iLoop in range(self.zyNumH): for jLoop in range(self.zyNumV): zyEData[jLoop, iLoop] = self.myPulse.evaluateEx(xValue, yArr[jLoop], zArr[iLoop], tValue) # print('zyEData[', iLoop, jLoop, '] = ', zyEData[iLoop, jLoop]) # generate the xy plot canvas = self.ui.zyPlot.canvas canvas.ax.clear() levels = myPlotUtils.generateContourLevels(zyEData) canvas.ax.contourf( zArr * convertUnitsNumber(1, "m", self.unitsZ), yArr * convertUnitsNumber(1, "m", self.unitsXY), zyEData, levels, extent="none", aspect="equal", ) # plt.colorbar(cs1, format='%3.2e') # plt.gcf().colorbar(contours, format='%3.2e') canvas.ax.axis( [ self.zyMinH * convertUnitsNumber(1, "m", self.unitsZ), self.zyMaxH * convertUnitsNumber(1, "m", self.unitsZ), self.zyMinV * convertUnitsNumber(1, "m", self.unitsXY), self.zyMaxV * convertUnitsNumber(1, "m", self.unitsXY), ] ) canvas.ax.xaxis.set_major_locator(plt.MaxNLocator(self.numTicks)) canvas.ax.yaxis.set_major_locator(plt.MaxNLocator(self.numTicks)) canvas.ax.set_xlabel("z [" + self.unitsZ + "]") canvas.ax.set_ylabel("y [" + self.unitsXY + "]") if self.plotTitles == True: canvas.ax.set_title( "ZY slice, at x={0:4.2f} [{1}]".format(xValue * convertUnitsNumber(1, "m", self.unitsXY), self.unitsXY) ) canvas.fig.tight_layout() canvas.fig.set_facecolor("w") canvas.draw() def plotXY(self): # instance of the plot utility class myPlotUtils = radtrack.plot.RbPlotUtils.RbPlotUtils() # Specify the desired grid size self.xyNumH = 64 self.xyNumV = 64 self.xyNumCells = self.xyNumH * self.xyNumV # specify the min's and max's self.xyMinH = -4.0 * self.w0 self.xyMaxH = 4.0 * self.w0 self.xyMinV = -4.0 * self.w0 self.xyMaxV = 4.0 * self.w0 xArr = np.zeros(self.xyNumH) yArr = np.zeros(self.xyNumV) xTmp = np.zeros(self.xyNumV) for iLoop in range(self.xyNumH): xArr[iLoop] = self.xyMinH + iLoop * (self.xyMaxH - self.xyMinH) / (self.xyNumH - 1) for jLoop in range(self.xyNumV): yArr[jLoop] = self.xyMinV + jLoop * (self.xyMaxV - self.xyMinV) / (self.xyNumV - 1) # Calculate Ex at the 2D array of x,y values xyEData = np.zeros((self.xyNumV, self.xyNumH)) # interchange of V/H is weirdly necessary!? for iLoop in range(self.xyNumH): for jLoop in range(self.xyNumV): xTmp[jLoop] = xArr[iLoop] xyEData[0 : self.xyNumV, iLoop] = np.real(self.myPulse.evalEnvelopeEx(xTmp, yArr, self.w0_z)) # generate the xy plot canvas = self.ui.xyPlot.canvas canvas.ax.clear() levels = myPlotUtils.generateContourLevels(xyEData) canvas.ax.contourf( xArr * convertUnitsNumber(1, "m", self.unitsXY), yArr * convertUnitsNumber(1, "m", self.unitsXY), xyEData, levels, extent="none", aspect="equal", ) # For some reason, I can't get the color bar to appear... # contours = canvas.ax.contourf(xArr, yArr, xyEData, levels, extent='none', aspect='equal') # plt.gcf().colorbar(contours, format='%3.2e') canvas.ax.axis( [ self.xyMinH * convertUnitsNumber(1, "m", self.unitsXY), self.xyMaxH * convertUnitsNumber(1, "m", self.unitsXY), self.xyMinV * convertUnitsNumber(1, "m", self.unitsXY), self.xyMaxV * convertUnitsNumber(1, "m", self.unitsXY), ] ) canvas.ax.xaxis.set_major_locator(plt.MaxNLocator(self.numTicks)) canvas.ax.yaxis.set_major_locator(plt.MaxNLocator(self.numTicks)) canvas.ax.set_xlabel("x [" + self.unitsXY + "]") canvas.ax.set_ylabel("y [" + self.unitsXY + "]") if self.plotTitles == True: canvas.ax.set_title( "XY slice, at z={0:4.2f} [{1}]".format( self.w0_z * convertUnitsNumber(1, "m", self.unitsZ), self.unitsZ ) ) canvas.fig.tight_layout() canvas.fig.set_facecolor("w") canvas.draw() def mirrorWithHole(self): # toggle the corresponding flag self.externalFields = True # get input from text boxes self.lambda0 = convertUnitsStringToNumber(self.ui.wavelength.text(), "m") self.w0 = convertUnitsStringToNumber(self.ui.waistSize.text(), "m") self.w0_z = convertUnitsStringToNumber(self.ui.waistPosition.text(), "m") self.zR = math.pi * self.w0 ** 2 / self.lambda0 # horiz. Rayleigh range [m] # load up the x,y locations of the mesh self.xMin = -4.0 * self.w0 self.xMax = 4.0 * self.w0 self.yMin = self.xMin self.yMax = self.xMax self.numPts = 64 self.nCells = self.numPts ** 2 self.xArr = np.zeros(self.numPts) for iLoop in range(self.numPts): self.xArr[iLoop] = self.xMin + iLoop * (self.xMax - self.xMin) / (self.numPts - 1) self.yArr = np.zeros(self.numPts) for jLoop in range(self.numPts): self.yArr[jLoop] = self.yMin + jLoop * (self.yMax - self.yMin) / (self.numPts - 1) self.xGrid = np.zeros((self.numPts, self.numPts)) self.yGrid = np.zeros((self.numPts, self.numPts)) for iLoop in range(self.numPts): for jLoop in range(self.numPts): self.xGrid[iLoop, jLoop] = self.xMin + iLoop * (self.xMax - self.xMin) / (self.numPts - 1) self.yGrid[iLoop, jLoop] = self.yMin + jLoop * (self.yMax - self.yMin) / (self.numPts - 1) # Create transverse field profile (#3 elliptical Gaussian donut) self.ExGridExternal = np.zeros((self.numPts, self.numPts)) self.wx3 = 2.0 * self.w0 self.rad1 = 1.0 * self.w0 self.rad2 = 2.0 * self.w0 for iLoop in range(self.numPts): for jLoop in range(self.numPts): xArg = self.xArr[iLoop] yArg = self.yArr[jLoop] rArg = math.sqrt(xArg ** 2 + yArg ** 2) rFactor = 1.0 if rArg <= self.rad2: rFactor = 0.5 + 0.5 * math.cos(math.pi * ((rArg - self.rad1) / (self.rad2 - self.rad1) - 1.0)) if rArg <= self.rad1: rFactor = 0.0 self.ExGridExternal[iLoop, jLoop] = ( rFactor * math.exp(-(xArg / self.wx3) ** 2) * math.exp(-(yArg / self.wx3) ** 2) ) # generate the xy plot canvas = self.ui.xyPlotExtFields.canvas canvas.ax.clear() myPlotUtils = radtrack.plot.RbPlotUtils.RbPlotUtils() levels = myPlotUtils.generateContourLevels(self.ExGridExternal) canvas.ax.contourf( self.xGrid * convertUnitsNumber(1, "m", self.unitsXY), self.yGrid * convertUnitsNumber(1, "m", self.unitsXY), self.ExGridExternal, levels, extent="none", aspect="equal", ) canvas.ax.xaxis.set_major_locator(plt.MaxNLocator(self.numTicks)) canvas.ax.yaxis.set_major_locator(plt.MaxNLocator(self.numTicks)) canvas.ax.set_xlabel("x [" + self.unitsXY + "]") canvas.ax.set_ylabel("y [" + self.unitsXY + "]") canvas.ax.axis( [ self.xMin * convertUnitsNumber(1, "m", self.unitsXY), self.xMax * convertUnitsNumber(1, "m", self.unitsXY), self.yMin * convertUnitsNumber(1, "m", self.unitsXY), self.yMax * convertUnitsNumber(1, "m", self.unitsXY), ] ) if self.plotTitles == True: canvas.ax.set_title("slice: quadratic square; at z={0:4.2f} [{1}]".format(0.0, self.unitsZ)) canvas.fig.tight_layout() canvas.fig.set_facecolor("w") canvas.draw() def generateCoeffs(self): # check whether the external fields exist if self.externalFields == False: return # do nothing (should throw an exception) # set default values self.numFuncCalls = 0 # choose initial guesses for all fitting parameters # also, specify the scale of variations for each paramGuess = np.zeros(2 * self.mMax + 2) paramGuess[0] = self.w0 # horizontal waist paramGuess[1] = 1.0 for ii in range(self.mMax - 1): paramGuess[ii + 2] = 1.0e-5 # horiz. coeff's paramGuess[self.mMax + 1] = 1.0 for ii in range(self.mMax): paramGuess[self.mMax + 1 + ii] = 1.0e-5 # vertical coeff's # invoke the least squares algorithm result = leastsq( self.residuals, paramGuess, args=( np.reshape(self.ExGridExternal, self.nCells), np.reshape(self.xGrid, self.nCells), np.reshape(self.yGrid, self.nCells), ), full_output=True, ftol=1e-5, maxfev=1200, ) parFit = result[0] nEvals = result[2]["nfev"] resVals = result[2]["fvec"] message = result[3] iError = result[4] print(" ") print(" iError = ", iError) print(" message = ", message) print(" nEvals = ", nEvals) print(" resVals = ", resVals) # load the results into named variables (for clarity) wxFit = parFit[0] mCFit = np.zeros(self.mMax + 1) for ii in range(self.mMax): mCFit[ii + 1] = parFit[1 + ii] nCFit = np.zeros(self.mMax + 1) for ii in range(self.mMax + 1): nCFit[ii] = parFit[self.mMax + 1 + ii] # check the results print(" ") print("The least squares minimization has completed:") print(" wx = ", self.wx3, "; ", wxFit) self.ui.ghTable.setItem(0, 0, QTableWidgetItem(str(mCFit[0]))) self.ui.ghTable.setItem(0, 1, QTableWidgetItem(nCFit[0])) # load the coefficient table in the GUI minLoop = min(self.mMax, 101) for iLoop in range(0, minLoop): self.ui.ghTable.setItem(iLoop, 0, QTableWidgetItem(str(mCFit[iLoop]))) self.ui.ghTable.setItem(iLoop, 1, QTableWidgetItem(str(nCFit[iLoop]))) maxLoop = max(self.mMax, 101) for iLoop in range(minLoop + 1, maxLoop): self.ui.ghTable.setItem(iLoop, 0, QTableWidgetItem("0")) self.ui.ghTable.setItem(iLoop, 1, QTableWidgetItem("0")) # modify the laser pulse object, using these GH coefficients self.myPulse.setMCoef(mCFit) self.myPulse.setNCoef(nCFit) # regenerate the standard plots, using this new info self.refreshPlots() # Calculate residuals for the least squares analysis # params - array of fitting parameters def residuals(self, _params, _e, _x, _y): self.hS1.setWaistX(_params[0]) self.hS1.setWaistY(_params[0]) self.hS2.setWaistX(_params[0]) self.hS2.setWaistY(_params[0]) hCoefs = np.zeros(self.mMax + 1) for ii in range(self.mMax): hCoefs[ii + 1] = _params[1 + ii] self.hS1.setMCoef(hCoefs) self.hS2.setNCoef(hCoefs) vCoefs = np.zeros(self.mMax + 1) for ii in range(self.mMax + 1): vCoefs[ii] = _params[self.mMax + 1 + ii] self.hS1.setNCoef(vCoefs) self.hS2.setMCoef(vCoefs) # let the user know what's going on if many function calls are required if self.numFuncCalls == 0: print(" ") print("Number of calls to method residual():") self.numFuncCalls += 1 if 100 * int(self.numFuncCalls / 100.0) == self.numFuncCalls: print(" ", self.numFuncCalls) # xDum = np.real(self.myPulse.evalEnvelopeEx(xTmp, yArr, self.w0_z)) return ( _e - np.real(self.hS1.evalEnvelopeEx(_x, _y, self.w0_z)) - np.real(self.hS2.evalEnvelopeEx(_x, _y, self.w0_z)) ) def erasePlots(self): self.ui.xyPlot.canvas.ax.clear() self.ui.xyPlot.canvas.ax.axis([-1.0, 1.0, -1.0, 1.0]) self.ui.xyPlot.canvas.ax.set_xlabel("x [" + self.unitsXY + "]") self.ui.xyPlot.canvas.ax.set_ylabel("y [" + self.unitsXY + "]") # if self.plotTitles == True: # self.ui.xyPlot.canvas.ax.set_title('cross-section') self.ui.xyPlot.canvas.fig.tight_layout() self.ui.xyPlot.canvas.fig.set_facecolor("w") self.ui.zxPlot.canvas.ax.clear() self.ui.zxPlot.canvas.ax.axis([-1.0, 1.0, -1.0, 1.0]) self.ui.zxPlot.canvas.ax.set_xlabel("z [" + self.unitsZ + "]") self.ui.zxPlot.canvas.ax.set_ylabel("x [" + self.unitsXY + "]") # if self.plotTitles == True: # self.ui.zxPlot.canvas.ax.set_title('cross-section') self.ui.zxPlot.canvas.fig.tight_layout() self.ui.zxPlot.canvas.fig.set_facecolor("w") # self.ui.zxPlot.canvas.draw() self.ui.zyPlot.canvas.ax.clear() self.ui.zyPlot.canvas.ax.axis([-1.0, 1.0, -1.0, 1.0]) self.ui.zyPlot.canvas.ax.set_xlabel("z [" + self.unitsZ + "]") self.ui.zyPlot.canvas.ax.set_ylabel("y [" + self.unitsXY + "]") # if self.plotTitles == True: # self.ui.zyPlot.canvas.ax.set_title('cross-section') self.ui.zyPlot.canvas.fig.tight_layout() self.ui.zyPlot.canvas.fig.set_facecolor("w") # self.ui.zyPlot.canvas.draw() def calculateLimits(self, _arr): # nothing to do, if beam hasn't been initialized if self.pulseInitialized == False: return def readFromSDDS(self, fileName=None): # use Qt file dialog if not fileName: fileName = QFileDialog.getOpenFileName( self, "Import Elegant/SDDS particle file -- ", self.parent.lastUsedDirectory, "*.sdds" ) # if user cancels out, do nothing if not fileName: return self.parent.lastUsedDirectory = os.path.dirname(fileName) if False: print(" ") print(" File to be parsed: ", fileName) # index is always zero...? sddsIndex = 0 # initialize sdds.sddsdata.pyd library (Windows only) with data file if sdds.sddsdata.InitializeInput(sddsIndex, fileName) != 1: sdds.sddsdata.PrintErrors(1) # get data storage mode...? sddsStorageMode = sdds.sddsdata.GetMode(sddsIndex) if False: print(" Storage mode for index ", sddsIndex, ": ", sddsStorageMode) # get description text...? sddsDescription = sdds.sddsdata.GetDescription(sddsIndex) if False: print(" Description for index ", sddsIndex, ": ", sddsDescription) # get parameter names paramNames = sdds.sddsdata.GetParameterNames(sddsIndex) numParams = len(paramNames) if False: print(" numParams = ", numParams) print(" Parameter names for index ", sddsIndex, ": \n", paramNames) # get parameter definitions paramDefs = range(numParams) for iLoop in range(numParams): paramDefs[iLoop] = sdds.sddsdata.GetParameterDefinition(sddsIndex, paramNames[iLoop]) if False: print(" paramDefs[", iLoop, "] = ", paramDefs[iLoop]) # give the user a look at the parameters (if any) msgBox = QtGui.QMessageBox() if numParams == 0: message = "WARNING --\n\n" message += "No parameters were found in your selected SDDS file!!\n\n" message += "It will be impossible to set the design momentum, or\n" message += " the total beam charge, etc." else: message = "The parameter names in your selected SDDS file are: \n" for iLoop in range(numParams): message += " " + paramNames[iLoop] + "\n" message += "The parameter definitions in the file are: \n" for iLoop in range(numParams): message += " " + str(paramDefs[iLoop]) + "\n\n" message += "WARNING --\n" message += " Logic for extracting the design momentum, total beam\n" message += " charge, etc. has not yet been implemented!" msgBox.setText(message) msgBox.exec_() # get column names columnNames = sdds.sddsdata.GetColumnNames(sddsIndex) numColumns = len(columnNames) if False: print(" numColumns = ", numColumns) print(" Column names for index ", sddsIndex, ": \n", columnNames) # initialize the parameter arrays paramData = range(numParams) for iLoop in range(numParams): paramData[iLoop] = [] # column data has to be handled differently; # it will be a 6D python array of N-D NumPy arrays columnData = range(numColumns) # read parameter data from the SDDS file # mus read particle data at the same time errorCode = sdds.sddsdata.ReadPage(sddsIndex) # print(' ') # print(' errorCode = ', errorCode) if errorCode != 1: sdds.sddsdata.PrintErrors(1) while errorCode > 0: for iLoop in range(numParams): paramData[iLoop].append(sdds.sddsdata.GetParameter(sddsIndex, iLoop)) for jLoop in range(numColumns): tmpData = [] tmpData.append(sdds.sddsdata.GetColumn(sddsIndex, jLoop)) if False: print(" ") print(" jLoop = ", jLoop) print(" tmpData = ", tmpData) columnData[jLoop] = np.array(tmpData[0]) if False: print(" ") print(" columnData[", jLoop, "] = ", columnData[jLoop]) errorCode = sdds.sddsdata.ReadPage(sddsIndex) # logic for deciphering and making use of parameter data goes here! # check whether the particle data is 6D if numColumns != 6: msgBox = QtGui.QMessageBox() message = "ERROR --\n\n" message += " Particle data in the selected SDDS file is not 6D!\n\n" message += " Column names are: \n" message += " " + str(columnNames) + "\n\n" message += "Please select another file.\n" message += "Thanks!" msgBox.setText(message) msgBox.exec_() return # get column definitions # units are in the 2nd column columnDefs = range(numColumns) unitStrings = range(numColumns) for iLoop in range(numColumns): columnDefs[iLoop] = sdds.sddsdata.GetColumnDefinition(sddsIndex, columnNames[iLoop]) unitStrings[iLoop] = columnDefs[iLoop][1] if False: print(" columnDefs[", iLoop, "] = ", columnDefs[iLoop]) print(" unitStrings[", iLoop, "] = ", unitStrings[iLoop]) # begin deciphering the column data dataRead = [False, False, False, False, False, False] dataIndex = [-1, -1, -1, -1, -1, -1] for iLoop in range(6): if columnNames[iLoop] == "x" or columnNames[iLoop] == "X": if dataRead[0] == True: message = "Error -- \n\n" message += " X column appears twice, for iLoop = " message += str(dataIndex[0]) + " and " + str(iLoop) dataRead[0] = True dataIndex[0] = iLoop if columnNames[iLoop] == "xp" or columnNames[iLoop] == "px" or columnNames[iLoop] == "x'": if dataRead[1] == True: message = "Error -- \n\n" message += " XP column appears twice, for iLoop = " message += str(dataIndex[1]) + " and " + str(iLoop) dataRead[1] = True dataIndex[1] = iLoop if columnNames[iLoop] == "y" or columnNames[iLoop] == "Y": if dataRead[2] == True: message = "Error -- \n\n" message += " Y column appears twice, for iLoop = " message += str(dataIndex[2]) + " and " + str(iLoop) dataRead[2] = True dataIndex[2] = iLoop if columnNames[iLoop] == "yp" or columnNames[iLoop] == "py" or columnNames[iLoop] == "y'": if dataRead[3] == True: message = "Error -- \n\n" message += " YP column appears twice, for iLoop = " message += str(dataIndex[3]) + " and " + str(iLoop) dataRead[3] = True dataIndex[3] = iLoop if columnNames[iLoop] == "s" or columnNames[iLoop] == "ct" or columnNames[iLoop] == "t": if dataRead[4] == True: message = "Error -- \n\n" message += " S column appears twice, for iLoop = " message += str(dataIndex[4]) + " and " + str(iLoop) dataRead[4] = True dataIndex[4] = iLoop if columnNames[iLoop] == "p" or columnNames[iLoop] == "pt" or columnNames[iLoop] == "dp": if dataRead[5] == True: message = "Error -- \n\n" message += " DP column appears twice, for iLoop = " message += str(dataIndex[5]) + " and " + str(iLoop) dataRead[5] = True dataIndex[5] = iLoop # initial validation of the column data for iLoop in range(6): if dataRead[iLoop] == False: msgBox = QtGui.QMessageBox() message = "ERROR --\n\n" message += " Not all of the data columns could be correctly interpreted!\n" message += " These are the column headings that were parsed from the file:\n" message += " " + str(columnNames) + "\n\n" message += "The parsing logic failed on: " + columnNames[iLoop] + "\n" message += "The code is looking for [x, xp, y, yp, s, dp] or something similar." msgBox.setText(message) msgBox.exec_() return # check for unspecified units, and set them to default value # if the units are specified, but incorrect, the problem is detected below defaultUnits = ["m", "rad", "m", "rad", "m", "rad"] for iLoop in range(6): # print(' before: unitStrings[', iLoop, '] = ', unitStrings[iLoop]) if not unitStrings[iLoop]: unitStrings[iLoop] = defaultUnits[dataIndex[iLoop]] # print(' after: unitStrings[', iLoop, '] = ', unitStrings[iLoop]) if False: print(" ") print(" Here is columnData[:]:") print(columnData) # check that all data columns are the same length numElements = [0, 0, 0, 0, 0, 0] for iLoop in range(6): numElements[iLoop] = len(columnData[iLoop]) # print(' size of column # ', iLoop, ' = ', numElements[iLoop]) for iLoop in range(5): if numElements[iLoop + 1] != numElements[0]: msgBox = QtGui.QMessageBox() message = "ERROR --\n\n" message += " Not all of the data columns have the same length!\n" message += " Here is the number of elements found in each column:\n" message += " " + str(numElements) + "\n\n" message += "Please try again with a valid particle file..." message += "Thanks!" msgBox.setText(message) msgBox.exec_() return # now we know the number of macro-particles numParticles = numElements[0] # print(' ') # print(' numParticles = ', numParticles) # all seems to be well, so load particle data into local array, # accounting for any non-standard physical units tmp6 = np.zeros((6, numParticles)) for iLoop in range(6): tmp6[dataIndex[iLoop], :] = columnData[iLoop] # another sanity check # myShape = np.shape(tmp6) # print(' ') # print(' myShape = ', myShape) # close the SDDS particle file if sdds.sddsdata.Terminate(sddsIndex) != 1: sdds.sddsdata.PrintErrors(1) # instantiate the particle bunch self.myBunch = beam.RbParticleBeam6D(numParticles) self.myBunch.setDesignMomentumEV(self.designMomentumEV) self.myBunch.setMassEV(self.eMassEV) # assume electrons self.pulseInitialized = True # load particle array into the phase space object self.myBunch.getDistribution6D().getPhaseSpace6D().setArray6D(tmp6) # post top-level parameters to GUI self.ui.numPtcls.setText("{:d}".format(numParticles)) self.ui.designMomentum.setText("{:.0f}".format(self.designMomentumEV * 1.0e-6) + " MeV") self.ui.totalCharge.setText("{:.0f}".format(self.totalCharge * 1.0e9) + " nC") # calculate bunch statistics and populate text boxes self.calculateTwiss() # plot the results self.refreshPlots() def readFromCSV(self, fileName=None): if not fileName: fileName = QtGui.QFileDialog.getOpenFileName( self, "Import RadTrack particle file -- ", self.parent.lastUsedDirectory, "*.csv" ) if not fileName: return self.parent.lastUsedDirectory = os.path.dirname(fileName) # check whether this is a RadTrack generated CSV file fileObject = open(fileName) csvReader = csv.reader(fileObject, delimiter=",") lineNumber = 0 for rawData in csvReader: lineNumber += 1 # for testing purposes if False: print(" ") print(" lineNumber = ", lineNumber) print(" rawData = ", rawData) # make sure this file follows the RadTrack format if lineNumber == 1: if rawData[0] != "RadTrack": fileObject.close() msgBox = QtGui.QMessageBox() message = "ERROR --\n\n" message += " The selected CSV file was not generated by RadTrack.\n" message += " Please select another file.\n\n" message += "Thanks!" msgBox.setText(message) msgBox.exec_() # ignore the 2nd line elif lineNumber == 2: continue # 3rd line contains the parametic data elif lineNumber == 3: self.designMomentumEV = float(rawData[0]) self.totalCharge = float(rawData[1]) # for testing only if False: print(" ") print(" p0 = ", self.designMomentumEV) print(" Q = ", self.totalCharge) # don't read beyond the first three lines elif lineNumber > 3: break # close the file fileObject.close() # load file into temporary data array tmp6 = np.loadtxt(fileName, dtype=float, skiprows=5, delimiter=",", unpack=True) # check whether the particle data is 6D arrayShape = np.shape(tmp6) numDimensions = arrayShape[0] if numDimensions != 6: msgBox = QtGui.QMessageBox() message = "ERROR --\n\n" message += " Particle data in the selected CSV file is not 6D!\n" message += " Please select another file.\n\n" message += "Thanks!" msgBox.setText(message) msgBox.exec_() return # store the number of particles read from the file numParticles = arrayShape[1] # instantiate the particle bunch self.myBunch = beam.RbParticleBeam6D(numParticles) self.myBunch.setDesignMomentumEV(self.designMomentumEV) self.myBunch.setMassEV(self.eMassEV) # assume electrons self.pulseInitialized = True # load particle array into the phase space object self.myBunch.getDistribution6D().getPhaseSpace6D().setArray6D(tmp6) # for testing purposes only if False: print(" ") print(" numParticles = ", numParticles) q6 = self.myBunch.getDistribution6D().getPhaseSpace6D().getArray6D() print(" 1st particle: ", q6[:, 0]) # post top-level parameters to GUI self.ui.numPtcls.setText("{:d}".format(numParticles)) self.ui.designMomentum.setText("{:.0f}".format(self.designMomentumEV * 1.0e-6) + " MeV") self.ui.totalCharge.setText("{:.0f}".format(self.totalCharge * 1.0e9) + " nC") # calculate bunch statistics and populate text boxes self.calculateTwiss() # plot the results self.refreshPlots() def saveToCSV(self, fileName=None): if not fileName: fileName = getSaveFileName(self, "csv") if not fileName: return with open(fileName, "w"): return # make sure the top-level parameters are up-to-date self.designMomentumEV = convertUnitsStringToNumber(self.ui.designMomentum.text(), "eV") self.totalCharge = convertUnitsStringToNumber(self.ui.totalCharge.text(), "C") # create local pointer to particle array tmp6 = self.myBunch.getDistribution6D().getPhaseSpace6D().getArray6D() numParticles = tmp6.shape[1] # create a header to identify this as a RadTrack file h1 = "RadTrack,Copyright 2012-2014 by RadiaBeam Technologies LLC - All rights reserved (C)\n " # names of the top-level parameters h2 = "p0 [eV],Q [C],mass [eV]\n " # values of the top-level parameters h3 = str(self.designMomentumEV) + "," + str(self.totalCharge) + "," + str(self.eMassEV) + "\n " # label the columns h4 = "x,xp,y,yp,s,dp\n " # specify the units h5 = "[m],[rad],[m],[rad],[m],[rad]" # assemble the full header myHeader = h1 + h2 + h3 + h4 + h5 # write particle data into the file # The following ugliness is used to accommodate savetxt() # There is probably a better way... f6 = np.zeros((numParticles, 6)) f6[:, 0] = tmp6[0, :] f6[:, 1] = tmp6[1, :] f6[:, 2] = tmp6[2, :] f6[:, 3] = tmp6[3, :] f6[:, 4] = tmp6[4, :] f6[:, 5] = tmp6[5, :] np.savetxt(fileName, f6, fmt="%.12e", delimiter=",", comments="", header=myHeader) def saveToSDDS(self, sddsFileName=None): if not sddsFileName: sddsFileName = getSaveFileName(self, "sdds") if not sddsFileName: return with open(sddsFileName, "w"): return # make sure the top-level parameters are up-to-date self.designMomentumEV = convertUnitsStringToNumber(self.ui.designMomentum.text(), "eV") self.totalCharge = convertUnitsStringToNumber(self.ui.totalCharge.text(), "C") # create local pointer to particle array tmp6 = self.myBunch.getDistribution6D().getPhaseSpace6D().getArray6D() mySDDS = sdds.SDDS(0) mySDDS.description[0] = "RadTrack" mySDDS.description[1] = "Copyright 2013-2014 by RadiaBeam Technologies. All rights reserved." mySDDS.parameterName = ["designMomentumEV", "totalCharge", "eMassEV"] mySDDS.parameterData = [[self.designMomentumEV], [self.totalCharge], [self.eMassEV]] mySDDS.parameterDefinition = [ ["", "", "", "", mySDDS.SDDS_DOUBLE, ""], ["", "", "", "", mySDDS.SDDS_DOUBLE, ""], ["", "", "", "", mySDDS.SDDS_DOUBLE, ""], ] mySDDS.columnName = ["x", "xp", "y", "yp", "s", "dp"] mySDDS.columnData = [ [list(tmp6[0, :])], [list(tmp6[1, :])], [list(tmp6[2, :])], [list(tmp6[3, :])], [list(tmp6[4, :])], [list(tmp6[5, :])], ] if False: print(" ") print(" Here is mySDDS.columnData[:]:") print(mySDDS.columnData) mySDDS.columnDefinition = [ ["", "m", "", "", mySDDS.SDDS_DOUBLE, 0], ["", "rad", "", "", mySDDS.SDDS_DOUBLE, 0], ["", "m", "", "", mySDDS.SDDS_DOUBLE, 0], ["", "rad", "", "", mySDDS.SDDS_DOUBLE, 0], ["", "m", "", "", mySDDS.SDDS_DOUBLE, 0], ["", "rad", "", "", mySDDS.SDDS_DOUBLE, 0], ] mySDDS.save(sddsFileName)