def __init__(self, args, tempdir, scaling=4, overwrite=True, verbose=False, summary=False): self.cons = constraintsClass() self.beam = rbf.beam() self.twiss = rtf.twiss() self.tmpdir = tempdir self.verbose = verbose self.summary = summary self.parameters = list(args) self.dirname = os.path.relpath(self.tmpdir) self.framework = fw.Framework(self.dirname, clean=True, verbose=False) self.framework.loadSettings('CLA10-BA1_TOMP.def') self.framework.change_Lattice_Code('All', 'elegant') if not os.name == 'nt': self.framework.defineASTRACommand([ 'mpiexec', '-np', str(3 * scaling), '/opt/ASTRA/astra_MPICH2.sh' ]) self.framework.defineGeneratorCommand(['/opt/ASTRA/generator.sh']) self.framework.defineCSRTrackCommand([ '/opt/OpenMPI-1.4.3/bin/mpiexec', '-n', str(3 * scaling), '/opt/CSRTrack/csrtrack_openmpi.sh' ]) self.framework.generator.number_of_particles = 2**(3 * scaling) else: self.framework.generator.number_of_particles = 2**(3 * 3) self.framework.defineElegantCommand(['elegant']) self.framework.setElementType('quadrupole', 'k1l', self.parameters)
def __init__(self, lattice='Lattices/clara400_v12_v3_elegant_jkj.def'): self.CLARA_dir = os.path.relpath( os.path.dirname(os.path.dirname(os.path.abspath( __file__)))) # if CLARA_dir is None else CLARA_dir self.cons = constraintsClass() self.beam = rbf.beam() self.twiss = rtf.twiss() self.lattice_file = lattice self.base_files = None self.verbose = False self.start_lattice = None self.scaling = 6 self.overwrite = True self.post_injector = True self.startcharge = None self.changes = None self.sample_interval = 1
def load_elegant_file(self, filename, floor_filename=None): interpret = elegantInterpret() lattice = interpret.readElegantFile(filename) twiss = rtf.twiss() if floor_filename is None: dot_index = filename.rfind('.') floor_filename = filename[:dot_index]+'.flr' xyz = twiss.read_elegant_floor_file(floor_filename) # print(xyz) # exit() if len(lattice.keys()) == 1: latt = lattice[list(lattice.keys())[0]] for n, start, end, rot in xyz: if n.lower() in latt: latt[n.lower()]['position_start'] = start latt[n.lower()]['position_end'] = end latt[n.lower()]['global_rotation'] = rot print (n.lower(), start, rot)
def __init__(self, args, tempdir, scaling=4, overwrite=True, verbose=False, summary=False, clean=False): self.cons = constraintsClass() self.beam = rbf.beam() self.twiss = rtf.twiss() self.tmpdir = tempdir self.scaling = scaling self.verbose = verbose self.summary = summary self.parameters = list(args) self.dirname = os.path.basename(self.tmpdir) self.framework = Framework(self.dirname, clean=clean, verbose=False) self.framework.loadSettings('Lattices/claraX400_v12_80MVm.def') if not os.name == 'nt': self.framework.defineASTRACommand(['mpiexec','-np',str(3*scaling),'/opt/ASTRA/astra_MPICH2.sh']) self.framework.defineGeneratorCommand(['/opt/ASTRA/generator.sh']) self.framework.defineCSRTrackCommand(['/opt/OpenMPI-1.4.3/bin/mpiexec','-n',str(3*scaling),'/opt/CSRTrack/csrtrack_openmpi.sh']) # self.framework.defineElegantCommand(['mpiexec','-np',str(3*scaling),'Pelegant']) # else: self.framework.defineElegantCommand(['elegant']) self.framework.setElementType('quadrupole','k1', self.parameters)
def __init__(self, args, tempdir, scaling=4, overwrite=True, verbose=False, summary=False, post_injector=True, parameterfile=[]): self.cons = constraintsClass() self.beam = rbf.beam() self.twiss = rtf.twiss() self.scaling = scaling self.tmpdir = tempdir self.verbose = verbose self.summary = summary self.overwrite = overwrite self.post_injector = post_injector self.parameters = list(args) self.dirname = os.path.basename(self.tmpdir) self.framework = Framework(self.dirname, clean=False) self.framework.loadSettings('Lattices/clara400_v12_v3.def') if isinstance(parameterfile, (list, tuple)): for p in parameterfile: self.framework.loadParametersFile(p) elif parameterfile is not None: self.framework.loadParametersFile(parameterfile) if not os.name == 'nt': self.framework.defineASTRACommand([ 'mpiexec', '-np', str(3 * scaling), '/opt/ASTRA/astra_MPICH2.sh' ]) self.framework.defineGeneratorCommand(['/opt/ASTRA/generator.sh']) self.framework.defineCSRTrackCommand([ '/opt/OpenMPI-1.4.3/bin/mpiexec', '-n', str(3 * scaling), '/opt/CSRTrack/csrtrack_openmpi.sh' ]) self.framework.generator.number_of_particles = 2**(3 * scaling) else: self.framework.generator.number_of_particles = 2**(3 * 3) self.framework.setElementType('quadrupole', 'k1', self.parameters)
def __init__(self, lattice='Lattices/clara400_v12_v3_elegant_jkj.def'): self.CLARA_dir = os.path.relpath(os.path.dirname(os.path.dirname( os.path.abspath(__file__))))# if CLARA_dir is None else CLARA_dir self.cons = constraintsClass() self.beam = rbf.beam() self.twiss = rtf.twiss() self.lattice_file = lattice self.base_files = None self.verbose = False self.start_lattice = None self.scaling = 6 self.overwrite = True self.clean = False self.post_injector = True self.startcharge = None self.changes = None self.sample_interval = 1 self.npart=2**(3*self.scaling) self.ncpu = self.scaling*3 self.elegant_ncpu = 1 self.genesis_ncpu = 2 self.doTracking = True self.change_to_elegant = True self.framework = fw.Framework('.', clean=False, overwrite=False, verbose=False)
import os import sys sys.path.append(os.path.abspath(__file__ + '/../../../../../')) from SimulationFramework.Framework import * import SimulationFramework.Modules.read_beam_file as rbf import SimulationFramework.Modules.read_twiss_file as rtf import numpy as np ncpu = 10 beam = rbf.beam() twiss = rtf.twiss() import csv ################################ CSRTrack ##################################### # framework = Framework('VBC_CSRTrack') # if not os.name == 'nt': # framework.defineGeneratorCommand(['/opt/ASTRA/generator']) # framework.defineASTRACommand(['mpiexec','-np',str(ncpu),'/opt/ASTRA/astra_MPICH2.sh']) # framework.defineCSRTrackCommand(['/opt/OpenMPI-1.4.3/bin/mpiexec','-n',str(ncpu),'/opt/CSRTrack/csrtrack_openmpi.sh']) # framework.defineElegantCommand(['elegant']) # # framework.loadSettings('Lattices/clara400_v12_v3.def') # framework['VBC'].file_block['input']['prefix'] = '../basefiles_5/' # framework.track(startfile='VBC', endfile='S07') ################################ ELEGANT ###################################### # framework = Framework('VBC_Elegant') # if not os.name == 'nt': # framework.defineGeneratorCommand(['/opt/ASTRA/generator']) # framework.defineASTRACommand(['mpiexec','-np',str(ncpu),'/opt/ASTRA/astra_MPICH2.sh']) # framework.defineCSRTrackCommand(['/opt/OpenMPI-1.4.3/bin/mpiexec','-n',str(ncpu),'/opt/CSRTrack/csrtrack_openmpi.sh'])
def __init__(self, args, tempdir, scaling=4, overwrite=True, verbose=False, summary=False, post_injector=True): self.cons = constraintsClass() self.beam = rbf.beam() self.twiss = rtf.twiss() self.scaling = scaling self.tmpdir = tempdir self.verbose = verbose self.summary = summary self.overwrite = overwrite self.post_injector = post_injector # print 'self.post_injector = ', self.post_injector ''' if only post-injector optimisation''' if self.post_injector: linac2field, linac2phase, linac3field, linac3phase, fhcfield, fhcphase, linac4field, linac4phase, bcangle = args self.parameters = dict( zip([ 'linac2field', 'linac2phase', 'linac3field', 'linac3phase', 'fhcfield', 'fhcphase', 'linac4field', 'linac4phase', 'bcangle' ], args)) else: ''' including injector parameters ''' gunphase, gunsol, linac1field, linac1phase, linac1sol1, linac1sol2, linac2field, linac2phase, linac3field, linac3phase, fhcfield, fhcphase, linac4field, linac4phase, bcangle = args self.parameters = dict( zip([ 'gunphase', 'gunsol', 'linac1field', 'linac1phase', 'linac1sol1', 'linac1sol2', 'linac2field', 'linac2phase', 'linac3field', 'linac3phase', 'fhcfield', 'fhcphase', 'linac4field', 'linac4phase', 'bcangle' ], args)) self.npart = 2**(3 * scaling) ncpu = scaling * 3 if self.post_injector: self.sbandlinacfields = np.array( [linac2field, linac3field, linac4field]) else: self.sbandlinacfields = np.array( [linac1field, linac2field, linac3field, linac4field]) self.dirname = os.path.basename(self.tmpdir) self.framework = Framework(self.dirname, overwrite=overwrite, verbose=verbose) if not os.name == 'nt': self.framework.defineGeneratorCommand(['/opt/ASTRA/generator']) self.framework.defineASTRACommand( ['mpiexec', '-np', str(ncpu), '/opt/ASTRA/astra_MPICH2.sh']) self.framework.defineCSRTrackCommand([ '/opt/OpenMPI-1.4.3/bin/mpiexec', '-n', str(ncpu), '/opt/CSRTrack/csrtrack_openmpi.sh' ]) self.framework.defineElegantCommand(['elegant']) self.framework.loadSettings('Lattices/clara400_v12_v3.def') if not self.post_injector: self.framework.generator.particles = self.npart self.framework.modifyElement('CLA-HRG1-GUN-CAV', 'phase', gunphase) self.framework.modifyElement('CLA-HRG1-GUN-SOL', 'field_amplitude', gunsol) self.framework.modifyElement('CLA-L01-CAV', 'field_amplitude', abs(linac1field)) self.framework.modifyElement('CLA-L01-CAV', 'phase', linac1phase) self.framework.modifyElement('CLA-L01-CAV-SOL-01', 'field_amplitude', linac1sol1) self.framework.modifyElement('CLA-L01-CAV-SOL-02', 'field_amplitude', linac1sol2) self.framework.modifyElement('CLA-L02-CAV', 'field_amplitude', abs(linac2field)) self.framework.modifyElement('CLA-L02-CAV', 'phase', linac2phase) self.framework.modifyElement('CLA-L03-CAV', 'field_amplitude', abs(linac3field)) self.framework.modifyElement('CLA-L03-CAV', 'phase', linac3phase) self.framework.modifyElement('CLA-L4H-CAV', 'field_amplitude', abs(fhcfield)) self.framework.modifyElement('CLA-L4H-CAV', 'phase', fhcphase) self.framework.modifyElement('CLA-L04-CAV', 'field_amplitude', abs(linac4field)) self.framework.modifyElement('CLA-L04-CAV', 'phase', linac4phase) self.framework['bunch_compressor'].set_angle(abs(bcangle))
def analyse_image(mw, dir=None): dir = str(dir) if dir is not None: if os.path.isfile(dir + '/' + 'CLA-S07-APER-01.hdf5'): args = parser.parse_args() beam = rbf.beam() twiss = rtf.twiss() beam.read_HDF5_beam_file(dir + '/CLA-S07-APER-01.hdf5') twiss.read_elegant_twiss_files(dir + '/CLARAX.twi') beam.slice_length = 0.01e-12 t = 1e12 * (beam.t - np.mean(beam.t)) t_grid = np.linspace(min(t), max(t), 2**8) # t_grid = np.arange(min(t), max(t), 0.01) peakIPDF = beam.PDFI(t, t_grid, bandwidth=beam.rms(t) / (2**3)) * 250 peakICDF = beam.CDF(t, t_grid, bandwidth=beam.rms(t) / (2**3)) peakIFWHM, indexes = beam.FWHM(t_grid, peakIPDF, frac=0.5) peakIFWHM2, indexes2 = beam.FWHM(t_grid, peakIPDF, frac=2) stdpeakIPDF = np.std( peakIPDF[indexes2] ) #(max(peakIPDF[indexes2]) - min(peakIPDF[indexes2]))/np.mean(peakIPDF[indexes2]) # Flat-top in the distribution! # print('stdpeakIPDF = ', stdpeakIPDF) # print 'Peak Fraction = ', 100*peakICDF[indexes][-1]-peakICDF[indexes][0], stdpeakIPDF beam.bin_time() t = 1e12 * (beam.t - np.mean(beam.t)) dt = beam.slice_length * (max(t) - min(t)) sigmat = np.std(t) sigmap = np.std(beam.p) meanp = np.mean(beam.p) fitp = 100 * sigmap / meanp peakI, peakIstd, peakIMomentumSpread, peakIEmittanceX, peakIEmittanceY, peakIMomentum, peakIDensity = beam.sliceAnalysis( ) peakI = max(peakIPDF) # chirp = beam.chirp chirp = 1e-6 * (max(beam.cp) - min(beam.cp)) print('bunchlength = ', 1e6 * beam.sigma_z, 'um chirp = ', -1 * beam.chirp, 'MeV/ps') plt.subplots_adjust(left=0.1, bottom=0.1, right=0.9, top=0.9, wspace=0.25, hspace=0.15) ax1 = mw.getFigure().add_subplot(321) ax2 = mw.getFigure().add_subplot(323) ax3 = mw.getFigure().add_subplot(325) ax4 = mw.getFigure().add_subplot(322) ax5 = mw.getFigure().add_subplot(324) ax6 = mw.getFigure().add_subplot(326) ax = [[ax1, ax2, ax3], [ax4, ax5, ax6]] exponent = np.floor(np.log10(np.abs(beam.slice_length))) x = 10**(12) * np.array((beam.slice_bins - np.mean(beam.t))) ax[0][0].plot(x, beam.slice_peak_current) ax[0][0].set(xlabel='t (ps)', ylabel='I [A]') # ax[0][1].plot(t_grid, peakIPDF, color='blue', alpha=0.5, lw=3) # ax[0][1].fill_between(t_grid[indexes], peakIPDF[indexes], 0, facecolor='gray', edgecolor='gray', alpha=0.4) beam.read_HDF5_beam_file(dir + '/CLA-S07-CAV-04-SCR.hdf5') t = 1e12 * (beam.t - np.mean(beam.t)) dt = beam.slice_length * (max(t) - min(t)) p = 1e-6 * beam.cp ymax = max(p) + 1 ymin = min(p) - 1 if ymax - ymin < 5: ymax = np.mean(p) + 2.5 ymin = np.mean(p) - 2.5 ax[0][1].set_xlim(min(t) - dt, max(t) + dt) ax[0][1].hist2d(t, p, bins=(250, 250), cmap=plt.cm.jet, range=[[min(t), max(t)], [ymin, ymax]]) # ax[0][2].set_ylim(top=ymax, bottom=ymin) ax[0][1].set(ylabel='cp [Mev]') beam.read_HDF5_beam_file(dir + '/CLA-S07-APER-01.hdf5') t = 1e12 * (beam.t - np.mean(beam.t)) dt = beam.slice_length * (max(t) - min(t)) p = 1e-6 * beam.cp ymax = max(p) + 1 ymin = min(p) - 1 if ymax - ymin < 5: ymax = np.mean(p) + 2.5 ymin = np.mean(p) - 2.5 ax[0][2].set_xlim(min(t) - dt, max(t) + dt) ax[0][2].hist2d(t, p, bins=(250, 250), cmap=plt.cm.jet, range=[[min(t), max(t)], [ymin, ymax]]) # ax[0][2].set_ylim(top=ymax, bottom=ymin) ax[0][2].set(ylabel='cp [Mev]') ax[1][0].plot(x, 1e6 * beam.slice_normalized_horizontal_emittance) ax[1][0].plot(x, 1e6 * beam.slice_normalized_vertical_emittance) ax[1][0].set_ylim(top=3, bottom=0) ax[1][0].set(ylabel='emit_x / emit_y [m]') ax[1][0].grid() ax[1][1].plot(twiss.elegant['s'], 0.511 * twiss.elegant['pCentral0']) # ax[1][1].set_ylim(top=1100, bottom=0) ax[1][1].set(ylabel='Momentum [MeV/c]') ax[1][1].grid() ax[1][2].plot(twiss.elegant['s'], 1e3 * twiss['sigma_x']) ax[1][2].plot(twiss.elegant['s'], 1e3 * twiss['sigma_y']) ax[1][2].set(ylabel='sigma_x / sigma_y [m]') ax[1][2].grid()
def __init__(self, directory='.', **kwargs): super(astraPlotWidget, self).__init__(**kwargs) self.beam = raf.beam() self.twiss = rtf.twiss() self.directory = directory ''' twissPlotWidget ''' self.twissPlotView = GraphicsView(useOpenGL=True) self.twissPlotWidget = GraphicsLayout() self.twissPlotView.setCentralItem(self.twissPlotWidget) # self.latticePlotData = imageio.imread(os.path.dirname( os.path.abspath(__file__))+'/lattice_plot.png') self.latticePlots = {} self.twissPlots = {} i = -1 for entry in self.twissplotLayout: if entry == 'next_row': self.twissPlotWidget.nextRow() else: i += 1 p = self.twissPlotWidget.addPlot(title=entry['name']) p.showGrid(x=True, y=True) vb = p.vb vb.setYRange(*entry['range']) # latticePlot = ImageItem(self.latticePlotData) # latticePlot.setOpts(axisOrder='row-major') # vb.addItem(latticePlot) # latticePlot.setZValue(-1) # make sure this image is on top # latticePlot.setOpacity(0.5) self.twissPlots[entry['name']] = p.plot( pen=mkPen('b', width=3)) # self.latticePlots[p.vb] = latticePlot p.vb.sigRangeChanged.connect(self.scaleLattice) ''' beamPlotWidget ''' self.beamPlotWidget = QWidget() self.beamPlotLayout = QVBoxLayout() # self.item = ImageItem() self.beamPlotWidget.setLayout(self.beamPlotLayout) # self.beamPlotView = ImageView(imageItem=self.item) # self.rainbow = rainbow() # self.item.setLookupTable(self.rainbow) # self.item.setLevels([0,1]) # # self.beamPlotWidgetGraphicsLayout = GraphicsLayout() # # p = self.beamPlotWidgetGraphicsLayout.addPlot(title='beam') # # p.showGrid(x=True, y=True) # # self.beamPlot = p.plot(pen=None, symbol='+') # # self.beamPlotView.setCentralItem(self.beamPlotWidgetGraphicsLayout) self.beamPlotXAxisCombo = QComboBox() self.beamPlotXAxisDict = OrderedDict() self.beamPlotXAxisDict['x'] = {'scale': 1e3, 'axis': 'x [mm]'} self.beamPlotXAxisDict['y'] = {'scale': 1e3, 'axis': 'y [mm]'} self.beamPlotXAxisDict['z'] = { 'scale': 1e6, 'axis': 'z [micron]', 'norm': True } self.beamPlotXAxisDict['t'] = { 'scale': 1e12, 'axis': 't [ps]', 'norm': True } self.beamPlotXAxisDict['cpx'] = {'scale': 1e3, 'axis': 'cpx [keV]'} self.beamPlotXAxisDict['cpy'] = {'scale': 1e3, 'axis': 'cpy [keV]'} self.beamPlotXAxisDict['BetaGamma'] = { 'scale': 0.511, 'axis': 'cp [MeV]' } self.beamPlotXAxisCombo.addItems(list(self.beamPlotXAxisDict.keys())) self.beamPlotXAxisCombo.setCurrentIndex(2) self.beamPlotYAxisCombo = QComboBox() self.beamPlotYAxisCombo.addItems(list(self.beamPlotXAxisDict.keys())) self.beamPlotXAxisCombo.setCurrentIndex(6) self.beamPlotNumberBins = QSpinBox() self.beamPlotNumberBins.setRange(10, 500) self.beamPlotNumberBins.setSingleStep(10) self.histogramBins = 100 self.beamPlotNumberBins.setValue(self.histogramBins) self.beamPlotAxisWidget = QWidget() self.beamPlotAxisWidget.setMaximumHeight(100) self.beamPlotAxisLayout = QHBoxLayout() self.beamPlotAxisWidget.setLayout(self.beamPlotAxisLayout) self.beamPlotAxisLayout.addWidget(self.beamPlotXAxisCombo) self.beamPlotAxisLayout.addWidget(self.beamPlotYAxisCombo) self.beamPlotAxisLayout.addWidget(self.beamPlotNumberBins) self.beamPlotXAxisCombo.currentIndexChanged.connect(self.plotDataBeam) self.beamPlotYAxisCombo.currentIndexChanged.connect(self.plotDataBeam) self.beamPlotNumberBins.valueChanged.connect(self.plotDataBeam) self.beamPlotXAxisCombo.setCurrentIndex(2) self.beamPlotYAxisCombo.setCurrentIndex(6) self.canvasWidget = QWidget() l = QVBoxLayout(self.canvasWidget) self.sc = MyStaticMplCanvas(self.canvasWidget, width=1, height=1, dpi=150) l.addWidget(self.sc) self.beamPlotLayout.addWidget(self.beamPlotAxisWidget) self.beamPlotLayout.addWidget(self.canvasWidget) ''' slicePlotWidget ''' self.sliceParams = [ { 'name': 'slice_normalized_horizontal_emittance', 'units': 'm-rad', 'text': 'enx' }, { 'name': 'slice_normalized_vertical_emittance', 'units': 'm-rad', 'text': 'eny' }, { 'name': 'slice_peak_current', 'units': 'A', 'text': 'PeakI' }, { 'name': 'slice_relative_momentum_spread', 'units': '%', 'text': 'sigma-p' }, ] self.slicePlotWidget = QWidget() self.slicePlotLayout = QVBoxLayout() self.slicePlotWidget.setLayout(self.slicePlotLayout) # self.slicePlotView = GraphicsView(useOpenGL=True) self.slicePlotWidgetGraphicsLayout = GraphicsLayoutWidget() # self.slicePlots = {} self.slicePlotCheckbox = {} self.curve = {} self.sliceaxis = {} self.slicePlotCheckboxWidget = QWidget() self.slicePlotCheckboxLayout = QVBoxLayout() self.slicePlotCheckboxWidget.setLayout(self.slicePlotCheckboxLayout) self.slicePlot = self.slicePlotWidgetGraphicsLayout.addPlot( title='Slice', row=0, col=50) self.slicePlot.showAxis('left', False) self.slicePlot.showGrid(x=True, y=True) i = -1 colors = ['b', 'r', 'g', 'k'] for param in self.sliceParams: i += 1 axis = AxisItem("left") labelStyle = {'color': '#' + colorStr(mkColor(colors[i]))[0:-2]} axis.setLabel(text=param['text'], units=param['units'], **labelStyle) viewbox = ViewBox() axis.linkToView(viewbox) viewbox.setXLink(self.slicePlot.vb) self.sliceaxis[param['name']] = [axis, viewbox] self.curve[param['name']] = PlotDataItem(pen=colors[i], symbol='+') viewbox.addItem(self.curve[param['name']]) col = self.findFirstEmptyColumnInGraphicsLayout() self.slicePlotWidgetGraphicsLayout.ci.addItem(axis, row=0, col=col, rowspan=1, colspan=1) self.slicePlotWidgetGraphicsLayout.ci.addItem(viewbox, row=0, col=50) p.showGrid(x=True, y=True) # self.slicePlots[param] = self.slicePlot.plot(pen=colors[i], symbol='+') self.slicePlotCheckbox[param['name']] = QCheckBox(param['text']) self.slicePlotCheckboxLayout.addWidget( self.slicePlotCheckbox[param['name']]) self.slicePlotCheckbox[param['name']].stateChanged.connect( self.plotDataSlice) # self.slicePlotView.setCentralItem(self.slicePlotWidgetGraphicsLayout) self.slicePlotSliceWidthWidget = QSpinBox() self.slicePlotSliceWidthWidget.setMaximum(100) self.slicePlotSliceWidthWidget.setValue(20) self.slicePlotSliceWidthWidget.setSingleStep(1) self.slicePlotSliceWidthWidget.setSuffix(" slices") self.slicePlotSliceWidthWidget.setSpecialValueText('Automatic') self.slicePlotAxisWidget = QWidget() self.slicePlotAxisLayout = QHBoxLayout() self.slicePlotAxisWidget.setLayout(self.slicePlotAxisLayout) self.slicePlotAxisLayout.addWidget(self.slicePlotCheckboxWidget) self.slicePlotAxisLayout.addWidget(self.slicePlotSliceWidthWidget) # self.slicePlotXAxisCombo.currentIndexChanged.connect(self.plotDataSlice) self.slicePlotSliceWidthWidget.valueChanged.connect( self.changeSliceLength) # self.beamPlotXAxisCombo.setCurrentIndex(2) # self.beamPlotYAxisCombo.setCurrentIndex(5) self.slicePlotLayout.addWidget(self.slicePlotAxisWidget) self.slicePlotLayout.addWidget(self.slicePlotWidgetGraphicsLayout) self.layout = QVBoxLayout() self.setLayout(self.layout) self.tabWidget = QTabWidget() self.folderButton = QPushButton('Select Directory') self.folderLineEdit = QLineEdit() self.folderLineEdit.setReadOnly(True) self.folderLineEdit.setText(self.directory) self.reloadButton = QPushButton() self.reloadButton.setIcon(qApp.style().standardIcon( QStyle.SP_BrowserReload)) self.folderWidget = QGroupBox() self.folderLayout = QHBoxLayout() self.folderLayout.addWidget(self.folderButton) self.folderLayout.addWidget(self.folderLineEdit) self.folderLayout.addWidget(self.reloadButton) self.folderWidget.setLayout(self.folderLayout) self.folderWidget.setMaximumWidth(800) self.reloadButton.clicked.connect( lambda: self.changeDirectory(self.directory)) self.folderButton.clicked.connect(self.changeDirectory) self.fileSelector = QComboBox() self.fileSelector.currentIndexChanged.connect(self.updateScreenCombo) self.screenSelector = QComboBox() self.screenSelector.currentIndexChanged.connect(self.changeScreen) self.beamWidget = QGroupBox() self.beamLayout = QHBoxLayout() self.beamLayout.addWidget(self.fileSelector) self.beamLayout.addWidget(self.screenSelector) self.beamWidget.setLayout(self.beamLayout) self.beamWidget.setMaximumWidth(800) self.beamWidget.setVisible(False) self.folderBeamWidget = QWidget() self.folderBeamLayout = QHBoxLayout() self.folderBeamLayout.setAlignment(Qt.AlignLeft) self.folderBeamWidget.setLayout(self.folderBeamLayout) self.folderBeamLayout.addWidget(self.folderWidget) self.folderBeamLayout.addWidget(self.beamWidget) self.tabWidget.addTab(self.twissPlotView, 'Twiss Plots') self.tabWidget.addTab(self.beamPlotWidget, 'Beam Plots') self.tabWidget.addTab(self.slicePlotWidget, 'Slice Beam Plots') self.tabWidget.currentChanged.connect(self.changeTab) self.layout.addWidget(self.folderBeamWidget) self.layout.addWidget(self.tabWidget) self.plotType = 'Slice'
def main(): args = parser.parse_args() dirs = args.directory if isinstance(args.directory, (list, tuple)) else [args.directory] twiss = rtf.twiss() # fig = plt.figure(figsize=(15,15)) # fig.subplots_adjust(wspace=0, hspace=0) # gs = gridspec.GridSpec(3, 3) # ax_main = plt.subplot(gs[1:3, :2]) # ax_xDist = plt.subplot(gs[0, :2],sharex=ax_main) # ax_yDist = plt.subplot(gs[1:3, 2],sharey=ax_main) for d in dirs: beam = beam_analysis(d, args.file) twiss.read_elegant_twiss_files(d + '/CLARAX.twi') x, y = 1e12 * (beam.t - np.mean(beam.t)), 1e-6 * beam.cp t = 1e12 * (beam.t - np.mean(beam.t)) dt = 0.05 * (max(t) - min(t)) # ax_main.scatter(x,y,marker='.') # ax_main.set(xlabel="t [ps]", ylabel="cp [MeV/c]") beam.slices = 100 beam.bin_time() x = beam.slice['t_Bins'] bin_centers = 10**12 * np.array( (0.5 * (x[1:] + x[:-1]) - np.mean(beam.t))) # ax_xDist.plot(bin_centers, beam.slice_peak_current) # ax_xDist.set(ylabel='I [A]') fig2, ax = plt.subplots(3, 2, figsize=(10, 10)) fig2.subplots_adjust(hspace=0.01) t = 1e12 * (beam.t - np.mean(beam.t)) t_grid = np.linspace(min(t), max(t), 2**8) peakIPDF = beam.PDF(t, t_grid, bandwidth=beam.rms(t) / (2**4)) peakICDF = beam.CDF(t, t_grid, bandwidth=beam.rms(t) / (2**4)) peakIFWHM, indexes = beam.FWHM(t_grid, peakIPDF, frac=0.1) peakIFWHM2, indexes2 = beam.FWHM(t_grid, peakIPDF, frac=20) stdpeakIPDF = (max(peakIPDF[indexes2]) - min(peakIPDF[indexes2])) / np.mean( peakIPDF[indexes2]) # Flat-top in the distribution! # print('stdpeakIPDF = ', stdpeakIPDF) # print 'Peak Fraction = ', 100*peakICDF[indexes][-1]-peakICDF[indexes][0], stdpeakIPDF beam.bin_time() t = 1e12 * (beam.t - np.mean(beam.t)) dt = 0.05 * (max(t) - min(t)) ax[0][0].clear() ax[0][1].clear() ax[1][0].clear() ax[1][1].clear() ax[2][0].clear() ax[2][1].clear() exponent = np.floor(np.log10(np.abs(beam.slice_length))) x = 10**(12) * np.array((beam.slice_bins - np.mean(beam.t))) ax[0][0].plot(x, beam.slice_peak_current) ax[0][0].set(ylabel='I [A]') ax[1][0].set_xlim(min(t) - dt, max(t) + dt) t = 1e12 * (beam.t - np.mean(beam.t)) p = 1e-6 * beam.cp ymax = max(p) + 1 ymin = min(p) - 1 if ymax - ymin < 5: ymax = np.mean(p) + 2.5 ymin = np.mean(p) - 2.5 ax[1][0].hist2d(t, p, bins=(50, 50), cmap=plt.cm.jet, range=[[min(t), max(t)], [ymin, ymax]]) # ax[0][2].set_ylim(top=ymax, bottom=ymin) ax[1][0].set(ylabel='cp [Mev]') ax[2][0].plot(x, 1e6 * beam.slice_normalized_horizontal_emittance) ax[2][0].plot(x, 1e6 * beam.slice_normalized_vertical_emittance) ax[2][0].set_ylim(top=3, bottom=0) ax[2][0].set(ylabel='emit_x / emit_y [m]', xlabel='t [ps]') ax[2][0].grid() ax[0][1].plot(twiss.elegant['s'], 1e12 * twiss['sigma_z'] / 2.99e8) ax[0][1].set(ylabel='sigma_l [ps]') ax[0][1].set_ylim(bottom=0) print('bunch length = ', 1e12 * twiss['sigma_z'][-1] / 2.99e8) print('chirp = ', 1e-6 * (max(beam.cp) - min(beam.cp))) # ax[0][1].fill_between(t_grid[indexes], peakICDF[indexes], 0, facecolor='gray', edgecolor='gray', alpha=0.4) ax[1][1].plot(twiss.elegant['s'], 0.511 * twiss.elegant['pCentral0']) # ax[1][1].set_ylim(top=1100, bottom=0) ax[1][1].set(ylabel='Momentum [MeV/c]') ax[1][1].grid() ax[2][1].plot(twiss.elegant['s'], 1e3 * twiss['sigma_x']) ax[2][1].plot(twiss.elegant['s'], 1e3 * twiss['sigma_y']) ax[2][1].set(ylabel='sigma_x / sigma_y [m]', xlabel='s [m]') ax[2][1].grid() i = 0 j = 0 while i < 2: j = 0 while j < 2: plt.setp(ax[i][j].get_xticklabels(), visible=False) j = j + 1 i = i + 1 # plt.setp(ax_yDist.get_yticklabels(), visible=False) plt.show()