# -*- coding: utf-8 -*- """ The module provides visualization routines for displaying spatial and energy distributions of synchrotron sources in 2D and 3D.""" __author__ = "Konstantin Klementiev" __date__ = "12 Mar 2014" #import cmath import time import copy import numpy as np #import matplotlib as mpl import matplotlib.pyplot as plt from matplotlib.colors import LogNorm import os, sys; sys.path.append(os.path.join('..', '..')) # analysis:ignore #import xrt.backends.raycing as raycing import xrt.backends.raycing.sources as rs vmin = 2.5e-3 dpi = 82 xOrigin2d = 84 # all sizes are in pixels yOrigin2d = 48 space2dto1d = 8 height1d = 100 xSpaceExtra = 20 ySpaceExtra = 28 def visualize(source, data, title, saveName=None, sign=1): def one_fig(what, ts, tChar, otherChar): sh = what.shape xFigSize = float(xOrigin2d + sh[0] + space2dto1d + height1d + xSpaceExtra) yFigSize = float(yOrigin2d + sh[1] + space2dto1d + height1d + ySpaceExtra) fig = plt.figure(figsize=(xFigSize/dpi, yFigSize/dpi), dpi=dpi) rect_2D = [xOrigin2d / xFigSize, yOrigin2d / yFigSize, (sh[0]-1) / xFigSize, (sh[1]-1) / yFigSize] rect_1DE = copy.deepcopy(rect_2D) rect_1DE[1] = rect_2D[1] + rect_2D[3] + space2dto1d/yFigSize rect_1DE[3] = height1d / yFigSize rect_1Dx = copy.deepcopy(rect_2D) rect_1Dx[0] = rect_2D[0] + rect_2D[2] + space2dto1d/xFigSize rect_1Dx[2] = height1d / xFigSize extent = [source.eMin, source.eMax, ts[0], ts[-1]] ax2D = plt.axes(rect_2D) dataMax = what.max() ax2D.imshow( what.T, aspect='auto', cmap='hot', extent=extent, # interpolation='nearest', origin='lower', figure=fig, interpolation=None, origin='lower', figure=fig, norm=LogNorm(vmin=dataMax*vmin, vmax=dataMax)) ax2D.set_xlabel('energy (eV)') ax2D.set_ylabel(r"${0}'$ (mrad)".format(tChar)) ax1DE = plt.axes(rect_1DE, sharex=ax2D) ax1Dt = plt.axes(rect_1Dx, sharey=ax2D) plt.setp(ax1DE.get_xticklabels() + ax1Dt.get_yticklabels(), visible=False) dE = energies[1] - energies[0] dt = ts[1] - ts[0] ax1DE.plot(energies, np.sum(what, axis=1)*dt, 'r') # ax1DE.set_yscale('log') ax1Dt.plot(np.sum(what/energies[:, None]*dE, axis=0), ts, 'r') # ax1Dt.set_xscale('log') ax1DE.set_ylim(bottom=0) ax1DE.set_yticks(np.array([0, 2, 4, 6, 8])*1e16) ax2D.set_xlim(extent[0], extent[1]) ax2D.set_ylim(extent[2], extent[3]) ax1DE.text( 0.65, 1.0, r"Angular flux density {0} at ".format(title) + r"${0}'=0$".format(otherChar) + r" (ph/s/mrad$^2$/0.1%bw)", transform=ax1DE.transAxes, size=12, color='r', ha='center', va='bottom') ax1DE.text( 0.7, 0.95, "integrated\nover " + r"$d{0}'$".format(tChar), transform=ax1DE.transAxes, size=12, color='k', ha='center', va='top') ax1Dt.text( 0.25, 0.5, "integrated\nover " + r"$d$ln(energy)", rotation=-90, transform=ax1Dt.transAxes, size=12, color='k', ha='center', va='center') return fig, ax2D, ax1DE, ax1Dt if hasattr(source, 'xs'): xs = source.xs * 1e3 # from rad to mrad zs = source.zs * 1e3 # from rad to mrad energies = source.energies else: xs = np.mgrid[source.Theta_min:source.Theta_max + 0.5*source.dTheta: source.dTheta] * 1e3 zs = np.mgrid[source.Psi_min:source.Psi_max + 0.5*source.dPsi: source.dPsi] * 1e3 energies = np.mgrid[source.E_min:source.E_max + 0.5*source.dE: source.dE] xSlice = 0 zSlice = 0 if 'xrt' in source.prefix_save_name() or\ 'srw' in source.prefix_save_name(): pass else: data = np.concatenate((data[:, :0:-1, :], data), axis=1) data = np.concatenate((data[:, :, :0:-1], sign*data), axis=2) xs = np.concatenate((-xs[:0:-1], xs), axis=1) zs = np.concatenate((-zs[:0:-1], zs), axis=1) xSlice = (data.shape[1]-1) // 2 zSlice = (data.shape[2]-1) // 2 figX, ax2EX, ax1EX, ax1XX = one_fig(data[:, :, zSlice], xs, 'x', 'z') figZ, ax2EZ, ax1EZ, ax1ZZ = one_fig(data[:, xSlice, :], zs, 'z', 'x') dE = energies[1] - energies[0] integralEvsX = np.sum(data[:, :, zSlice]/energies[:, None], axis=0)*dE integralEvsZ = np.sum(data[:, xSlice, :]/energies[:, None], axis=0)*dE maxIntegral = max(np.max(integralEvsX), np.max(integralEvsZ)) ax1XX.set_xlim(maxIntegral*(sign-1)*0.5, maxIntegral) ax1ZZ.set_xlim(maxIntegral*(sign-1)*0.5, maxIntegral) if saveName is not None: fName = "{0}_{1}'E-" + source.prefix_save_name() + ".png" figX.savefig(fName.format(saveName, 'x')) figZ.savefig(fName.format(saveName, 'z')) def visualize3D(source, data, isZplane=True, saveName=None): # Enthought library imports from mayavi.scripts import mayavi2 from mayavi.sources.array_source import ArraySource # from mayavi.modules.outline import Outline # from mayavi.modules.volume import Volume from mayavi.modules.text3d import Text3D from mayavi.modules.image_plane_widget import ImagePlaneWidget from mayavi.tools.camera import view from mayavi.tools.animator import animate @mayavi2.standalone def view_data(data): """Example showing how to view a 3D numpy array in mayavi2. """ def set_labelE(ind): if hasattr(source, 'energies'): energies = source.energies else: energies = np.mgrid[source.E_min:source.E_max + 0.5*source.dE: source.dE] return '{0:.0f} eV'.format(energies[ind]) def move_view(obj, evt): labelE.text = set_labelE(ipwX.ipw.slice_index) pos = labelE.position labelE.position = ipwX.ipw.slice_index * src.spacing[0] + 1,\ pos[1], pos[2] labelE.vector_text.update() def set_lut(ipw): lutM = ipw.module_manager.scalar_lut_manager # lutM.show_scalar_bar = True # lutM.number_of_labels = 9 lutM.lut.scale = 'log10' lutM.lut.range = [dataMax*vmin, dataMax] lutM.lut_mode = 'hot' @animate() def anim(data, ipwX): scene.scene.off_screen_rendering = True scene.scene.anti_aliasing_frames = 0 for i in range(0, data.shape[0], 1): ipwX.ipw.slice_index = i move_view(None, None) if saveName is not None: scene.scene.save('{0}{1:04d}.png'.format(saveName, i)) yield # 'mayavi' is always defined on the interpreter. scene = mayavi.new_scene() # analysis:ignore scene.scene.background = (0, 0, 0) print(source.prefix_save_name()) src = ArraySource(transpose_input_array=True) sh = data.shape # print(sh) if 'xrt' in source.prefix_save_name() or\ 'srw' in source.prefix_save_name(): src.scalar_data = data[:, :sh[1]//2+1, :sh[2]//2+1].copy() else: src.scalar_data = data[:, ::-1, ::-1].copy() # src.spacing = np.array([-0.05, 1, 1]) # src.spacing = np.array([-0.25, 1, 1]) src.spacing = np.array([-0.25, 0.25, 0.25]) mayavi.add_source(src) # analysis:ignore # Visualize the data. # o = Outline() # mayavi.add_module(o) ipwY = ImagePlaneWidget() mayavi.add_module(ipwY) # analysis:ignore ipwY.ipw.plane_orientation = 'y_axes' # our x-axis ipwY.ipw.slice_index = int(data.shape[1] - 1) # if 'xrt' in source.prefix_save_name(): # ipwY.ipw.slice_index /= int(2) ipwY.ipw.left_button_action = 0 set_lut(ipwY) if isZplane: ipwZ = ImagePlaneWidget() mayavi.add_module(ipwZ) # analysis:ignore ipwZ.ipw.plane_orientation = 'z_axes' # our z-axis ipwZ.ipw.slice_index = int(data.shape[2] - 1) # if 'xrt' in source.prefix_save_name(): # ipwZ.ipw.slice_index /= int(2) ipwZ.ipw.left_button_action = 0 if 'xrt' in source.prefix_save_name() or\ 'srw' in source.prefix_save_name(): pass else: data = np.concatenate((data[:, :0:-1, :], data), axis=1) data = np.concatenate((data[:, :, :0:-1], data), axis=2) sh = data.shape print(sh) src = ArraySource(transpose_input_array=True) src.scalar_data = data.copy() # src.spacing = np.array([-0.05, 1, 1]) # src.spacing = np.array([-0.25, 1, 1]) src.spacing = np.array([-0.25, 0.25, 0.25]) mayavi.add_source(src) # analysis:ignore ipwX = ImagePlaneWidget() mayavi.add_module(ipwX) # analysis:ignore ipwX.ipw.plane_orientation = 'x_axes' # energy set_lut(ipwX) ipwX.ipw.add_observer('WindowLevelEvent', move_view) ipwX.ipw.add_observer('StartInteractionEvent', move_view) ipwX.ipw.add_observer('EndInteractionEvent', move_view) labelE = Text3D() mayavi.add_module(labelE) # analysis:ignore labelE.position = (1, data.shape[1]*0.73*src.spacing[1], data.shape[2]*0.85*src.spacing[2]) labelE.orientation = 90, 0, 90 labelE.text = 'Energy' labelE.scale = 3, 3, 1 labelE.actor.property.color = 0, 1, 1 labelE.orient_to_camera = False labelE.text = set_labelE(0) view(45, 70, 200) wantToAnimate = True if wantToAnimate: anim(data, ipwX) else: ipwX.ipw.slice_index = data.shape[0]-1 move_view(None, None) data[data < 1e-7] = 1e-7 dataMax = np.max(data) view_data(data) def test_synchrotron_source(SourceClass, **kwargs): tstart = time.time() source = SourceClass(**kwargs) # if source.prefix_save_name().startswith('srw'): # import pickle # pickleName = 'srw-und-non0em.pickle' # with open(pickleName, 'rb') as f: # I0, l1, l2, l3 = pickle.load(f)[0:4] print('started') I0, l1, l2, l3 = source.intensities_on_mesh() I0 *= 1e-6 # from /sr to /mrad² print('finished') tstop = time.time() print('calculations took {0:.1f} s'.format(tstop - tstart)) ##for long calculations like srw: # if source.prefix_save_name().startswith('srw'): # import pickle # pickleName = source.prefix_save_name()+'.pickle' # with open(pickleName, 'wb') as f: # pickle.dump((I0, l1, l2, l3, tstop-tstart), f, protocol=2) ##visualize in 2D: visualize(source, I0, r'$I_0$', 'I0') # visualize(source, I0*(1+l1)/2., r'$I_{\sigma\sigma}$', 'Is') # visualize(source, I0*(1-l1)/2., r'$I_{\pi\pi}$', 'Ip') # visualize(source, I0*l2/2., r'$\Re{I_{\sigma\pi}}$', 'IspRe') # sign = -1 # if hasattr(source, 'Kx'): # if source.Kx > 0: # sign = 1 # visualize(source, I0*l3/2., r'$\Im{I_{\sigma\pi}}$', 'IspIm', sign=sign) ##select only one visualize3D at a time: # visualize3D(source, I0, isZplane=False, saveName='Itot') # visualize3D(source, I0*(1+l1)/2., isZplane=False, saveName='IsPol') # visualize3D(source, I0*(1-l1)/2., isZplane=False, saveName='IpPol') # visualize3D(source, I0*l2/2., saveName='IspRe') # visualize3D(source, I0*l3/2., saveName='IspIm') # if __name__ == '__main__': """Uncomment the block you want to test.""" ##*********** Bending Magnet *************** # kwargs = dict(B0=1.7, eE=3., xPrimeMax=2.5, zPrimeMax=0.3, # eMin=1500, eMax=31500, eN=3000, nx=1, nz=10) ###by WS: ## Source = rs.BendingMagnetWS ##by xrt: # kwargs['distE'] = 'BW' # Source = rs.BendingMagnet ##*********** Wiggler *************** # kwargs = dict(period=80., K=13., n=12, eE=3., xPrimeMax=2.5, # zPrimeMax=0.3, eMin=1500, eMax=31500, eN=3000, nx=20, nz=20) ##by WS: # Source = rs.WigglerWS ##by xrt: # kwargs['distE'] = 'BW' # Source = rs.Wiggler #*********** undulator *************** # kwargs = dict( # period=31.4, K=2.7, n=63, eE=6.08, # xPrimeMax=0.3, zPrimeMax=0.3, # eMin=500, eMax=31500, eN=1000, nx=20, nz=20) # Kmax = 1.92 # thetaMax, psiMax = 100e-6, 50e-6 # kwargs = dict(name='IVU18.5', eE=3.0, eI=0.5, # eEpsilonX=0.263, eEpsilonZ=0.008, betaX=9., betaZ=2., # period=18.5, n=108, K=Kmax, # eMin=1500, eMax=31500, eN=1000, nx=40, nz=4, # xPrimeMax=thetaMax*1e3, zPrimeMax=psiMax*1e3, distE='BW') kwargs = dict( period=31.4, K=2.7, n=63, eE=6.08, eI=0.5, xPrimeMax=0.3, zPrimeMax=0.15, eSigmaX=134.2, eSigmaZ=6.325, eEpsilonX=1., eEpsilonZ=0.01, eMin=1500, eMax=5000, eN=350, nx=40*4, nz=20*4) ##by Urgent: # kwargs['icalc'] = 3 # 0 emittance # Source = rs.UndulatorUrgent ###by SRW: # import srw.xrtSRW as xrtSRW # kwargs['R0'] = 50000 ## 974 s - single electron ## 65501 s - zero spread ## 66180 s -nonzero spread # kwargs['eSigmaX'] = 0 # kwargs['eSigmaZ'] = 0 # kwargs['eEpsilonX'] = 0 # kwargs['eEpsilonZ'] = 0 ## kwargs['eEspread'] = 1e-3 # kwargs['harmonicStart'] = 1 # kwargs['harmonicFin'] = 4 # Source = xrtSRW.UndulatorSRW #by xrt: kwargs['R0'] = 50000 # kwargs['eSigmaX'] = 0 # kwargs['eSigmaZ'] = 0 # kwargs['eEpsilonX'] = 0 # kwargs['eEpsilonZ'] = 0 kwargs['eEspread'] = 1e-3*0 kwargs['distE'] = 'BW' kwargs['xPrimeMaxAutoReduce'] = False kwargs['zPrimeMaxAutoReduce'] = False # kwargs['targetOpenCL'] = "CPU" kwargs['filamentBeam'] = True Source = rs.Undulator ##*** helical undulator ************** # kwargs = dict( # period=31.4, Ky=2.7, Kx=2.7, n=63, eE=6.08, # xPrimeMax=0.3, zPrimeMax=0.3, # eMin=500, eMax=10500, eN=1000, nx=20, nz=20) ###by Urgent: ## Source = rs.UndulatorUrgent ##by xrt: # kwargs['phaseDeg'] = 90 # kwargs['distE'] = 'BW' # kwargs['xPrimeMaxAutoReduce'] = False # kwargs['zPrimeMaxAutoReduce'] = False # Source = rs.Undulator test_synchrotron_source(Source, **kwargs) plt.show()