def get_SRWLMagFldC(self):
return SRWLMagFldC(_arMagFld=[self.get_SRWMagneticStructure()],
_arXc=array('d',
[self.horizontal_central_position]),
_arYc=array('d', [self.vertical_central_position]),
_arZc=array('d',
[self.longitudinal_central_position]))
def as_srw(self, gap="min", energy=None, harmonic=1, **kwargs):
use_srw = kwargs.get("use_srw", False)
if energy is not None:
pars = self.find_harmonic_and_gap(energy,
sort_harmonics=True,
use_srw=use_srw)[0]
gap = pars["gap"]
harmonic = pars["harmonic"]
if isinstance(gap, str) and gap == "min":
gap = self.min_gap
import srwlib
ebeam = beam.srw_ebeam(self.ebeam)
harmB = srwlib.SRWLMagFldH() # magnetic field harmonic
harmB.n = harmonic # harmonic number
harmB.h_or_v = "v" # magnetic field plane: horzontal ('h') or vertical ('v')
harmB.B = self.field(gap=gap) # magnetic field amplitude [T]
und = srwlib.SRWLMagFldU([harmB])
und.per = self.period / 1e3 # period length [m]
und.nPer = self.N # number of periods (will be rounded to integer)
magFldCnt = srwlib.SRWLMagFldC(
[und],
srwlib.array("d", [0]),
srwlib.array("d", [0]),
srwlib.array("d", [0]),
) # Container of all magnetic field elements
return ds(ebeam=ebeam, und=und, field=magFldCnt)
def _gaussian_beam_intensity(GsnBm, _mesh, **kwargs):
wfr = srwlib.SRWLWfr()
wfr.allocate(_mesh.ne, _mesh.nx, _mesh.ny) #Numbers of points vs Photon Energy, Horizontal and Vertical Positions
wfr.mesh = srwlib.deepcopy(_mesh)
wfr.partBeam.partStatMom1.x = GsnBm.x #Some information about the source in the Wavefront structure
wfr.partBeam.partStatMom1.y = GsnBm.y
wfr.partBeam.partStatMom1.z = GsnBm.z
wfr.partBeam.partStatMom1.xp = GsnBm.xp
wfr.partBeam.partStatMom1.yp = GsnBm.yp
arPrecPar = [kwargs['_samp_fact']]
srwlpy.CalcElecFieldGaussian(wfr, GsnBm, arPrecPar)
depType = -1
if((_mesh.ne >= 1) and (_mesh.nx == 1) and (_mesh.ny == 1)): depType = 0
elif((_mesh.ne == 1) and (_mesh.nx > 1) and (_mesh.ny == 1)): depType = 1
elif((_mesh.ne == 1) and (_mesh.nx == 1) and (_mesh.ny > 1)): depType = 2
elif((_mesh.ne == 1) and (_mesh.nx > 1) and (_mesh.ny > 1)): depType = 3
elif((_mesh.ne > 1) and (_mesh.nx > 1) and (_mesh.ny == 1)): depType = 4
elif((_mesh.ne > 1) and (_mesh.nx == 1) and (_mesh.ny > 1)): depType = 5
elif((_mesh.ne > 1) and (_mesh.nx > 1) and (_mesh.ny > 1)): depType = 6
if(depType < 0): Exception('Incorrect numbers of points in the mesh structure')
sNumTypeInt = 'f'
if(kwargs['_int_type'] == 4): sNumTypeInt = 'd'
arI = srwlib.array(sNumTypeInt, [0]*wfr.mesh.ne*wfr.mesh.nx*wfr.mesh.ny)
srwlpy.CalcIntFromElecField(arI, wfr, kwargs['_pol'], kwargs['_int_type'], depType, wfr.mesh.eStart, wfr.mesh.xStart, wfr.mesh.yStart) #extracts intensity
_fname = kwargs['_fname']
if(len(_fname) > 0): srwlib.srwl_uti_save_intens_ascii(arI, wfr.mesh, _fname, 0, ['Photon Energy', 'Horizontal Position', 'Vertical Position', ''], _arUnits=['eV', 'm', 'm', 'ph/s/.1%bw/mm^2'])
return wfr, arI
def wavefront_to_intensity(wavefront):
"""
Computes the intensity from a wavefront.
"""
intensity = srwlib.array('f', [0] * wavefront.mesh.nx * wavefront.mesh.ny)
srwlib.srwl.CalcIntFromElecField(intensity, wavefront, 6, 0, 3,
wavefront.mesh.eStart, 0, 0)
return np.array(intensity).reshape(wavefront.mesh.ny, wavefront.mesh.nx)
def get_SRWMagneticStructure(self):
if self._number_of_periods <= 3:
n_points = 100
arBx = array('d', [0] *
n_points) if self._K_horizontal > 0.0 else None
arBy = array('d', [0] *
n_points) if self._K_vertical > 0.0 else None
longitudinal_mesh = numpy.linspace(-self._period_length / 2.,
self._period_length / 2.,
n_points)
phases = 2.0 * numpy.pi * (longitudinal_mesh / self._period_length)
for i in range(0, n_points):
if self._K_horizontal > 0.0:
arBx[i] = -self.magnetic_field_horizontal() * numpy.cos(
phases[i])
if self._K_vertical > 0.0:
arBy[i] = self.magnetic_field_vertical() * numpy.cos(
phases[i])
return SRWLMagFld3D(_arBx=arBx,
_arBy=arBy,
_nx=1,
_ny=1,
_nz=n_points,
_rz=self._period_length,
_nRep=int(self._number_of_periods),
_interp=1)
else:
magnetic_fields = []
if self._K_vertical > 0.0:
magnetic_fields.append(
SRWLMagFldH(1, 'v', self.magnetic_field_vertical(), 0, 1,
1))
if self._K_horizontal > 0.0:
magnetic_fields.append(
SRWLMagFldH(1, 'h', self.magnetic_field_horizontal(), 0,
-1, 1))
return SRWLMagFldU(magnetic_fields, self._period_length,
int(self._number_of_periods))
def calc2d_srw(self, photon_energy, photon_energy_step, scanning_data):
Kv = scanning_data.get_additional_parameter("Kv")
Kh = scanning_data.get_additional_parameter("Kh")
period_id = scanning_data.get_additional_parameter("period_id")
n_periods = scanning_data.get_additional_parameter("n_periods")
B0v = Kv/period_id/(codata.e/(2*numpy.pi*codata.electron_mass*codata.c))
B0h = Kh/period_id/(codata.e/(2*numpy.pi*codata.electron_mass*codata.c))
eBeam = srwlib.SRWLPartBeam()
eBeam.Iavg = scanning_data.get_additional_parameter("electron_current")
eBeam.partStatMom1.gamma = scanning_data.get_additional_parameter("electron_energy") / (codata_mee * 1e-3)
eBeam.partStatMom1.relE0 = 1.0
eBeam.partStatMom1.nq = -1
eBeam.partStatMom1.x = 0.0
eBeam.partStatMom1.y = 0.0
eBeam.partStatMom1.z = -0.5*period_id*n_periods + 4
eBeam.partStatMom1.xp = 0.0
eBeam.partStatMom1.yp = 0.0
eBeam.arStatMom2[ 0] = scanning_data.get_additional_parameter("electron_beam_size_h") ** 2
eBeam.arStatMom2[ 1] = 0.0
eBeam.arStatMom2[ 2] = scanning_data.get_additional_parameter("electron_beam_divergence_h") ** 2
eBeam.arStatMom2[ 3] = scanning_data.get_additional_parameter("electron_beam_size_v") ** 2
eBeam.arStatMom2[ 4] = 0.0
eBeam.arStatMom2[ 5] = scanning_data.get_additional_parameter("electron_beam_divergence_v") ** 2
eBeam.arStatMom2[10] = scanning_data.get_additional_parameter("electron_energy_spread") ** 2
gap_h = scanning_data.get_additional_parameter("gap_h")
gap_v = scanning_data.get_additional_parameter("gap_v")
mesh = srwlib.SRWLRadMesh(photon_energy,
photon_energy,
1,
-gap_h / 2, gap_h / 2, scanning_data.get_additional_parameter("h_slits_points"),
-gap_v / 2, gap_v / 2, scanning_data.get_additional_parameter("v_slits_points"),
scanning_data.get_additional_parameter("distance"))
srw_magnetic_fields = []
if B0v > 0: srw_magnetic_fields.append(srwlib.SRWLMagFldH(1, "v", B0v))
if B0h > 0: srw_magnetic_fields.append(srwlib.SRWLMagFldH(1, "h", B0h))
magnetic_structure = srwlib.SRWLMagFldC([srwlib.SRWLMagFldU(srw_magnetic_fields, period_id, n_periods)],
srwlib.array("d", [0]), srwlib.array("d", [0]), srwlib.array("d", [0]))
wfr = srwlib.SRWLWfr()
wfr.mesh = mesh
wfr.partBeam = eBeam
wfr.allocate(mesh.ne, mesh.nx, mesh.ny)
srwlib.srwl.CalcElecFieldSR(wfr, 0, magnetic_structure, [1, 0.01, 0, 0, 50000, 1, 0])
mesh_out = wfr.mesh
h_array=numpy.linspace(mesh_out.xStart, mesh_out.xFin, mesh_out.nx)*1e3 # in mm
v_array=numpy.linspace(mesh_out.yStart, mesh_out.yFin, mesh_out.ny)*1e3 # in mm
intensity_array = numpy.zeros((h_array.size, v_array.size))
arI0 = srwlib.array("f", [0]*mesh_out.nx*mesh_out.ny) #"flat" array to take 2D intensity data
srwlib.srwl.CalcIntFromElecField(arI0, wfr, 6, 1, 3, photon_energy, 0, 0)
data = numpy.ndarray(buffer=arI0, shape=(mesh_out.ny, mesh_out.nx),dtype=arI0.typecode)
for ix in range(h_array.size):
for iy in range(v_array.size):
intensity_array[ix, iy] = data[iy,ix]
return self.calculate_power(h_array, v_array, intensity_array, photon_energy_step)
def _plotTest(self):
import pylab
from srwlib import srwl, array
from copy import deepcopy
wfr = self.SRWWavefront()
y = int(self.dim_y() / 2)
inten = self.intensity_at_y(y)
cor = self.absolute_x_coordinates()
# I_y(x)
s = self.interpolate(0)
e = s([0.0], self.absolute_x_coordinates())
i = e[0].real**2 + e[0].imag**2 + e[1].real**2 + e[1].imag**2
print(i)
mesh0 = deepcopy(wfr.mesh)
arI0 = array('f', [0] * mesh0.nx *
mesh0.ny) #"flat" array to take 2D intensity data
srwl.CalcIntFromElecField(arI0, wfr, 6, 0, 3, mesh0.eStart, 0, 0)
arI0x = array('f',
[0] * mesh0.nx) #array to take 1D intensity data (vs X)
srwl.CalcIntFromElecField(arI0x, wfr, 6, 0, 1, mesh0.eStart, 0, 0)
arI0y = array('f',
[0] * mesh0.ny) #array to take 1D intensity data (vs Y)
srwl.CalcIntFromElecField(arI0y, wfr, 6, 0, 2, mesh0.eStart, 0, 0)
print('done')
pylab.plot(cor, i[0, :], cor, inten, cor, arI0x)
pylab.show()
e = s(self.absolute_x_coordinates(), self.absolute_y_coordinates())
# I_x(y)
cor = self.absolute_y_coordinates()
e = s(cor, [0.0])
i = e[0].real**2 + e[0].imag**2 + e[1].real**2 + e[1].imag**2
print(i)
mesh0 = deepcopy(wfr.mesh)
arI0 = array('f', [0] * mesh0.nx *
mesh0.ny) #"flat" array to take 2D intensity data
srwl.CalcIntFromElecField(arI0, wfr, 6, 0, 3, mesh0.eStart, 0, 0)
arI0x = array('f',
[0] * mesh0.nx) #array to take 1D intensity data (vs X)
srwl.CalcIntFromElecField(arI0x, wfr, 6, 0, 1, mesh0.eStart, 0, 0)
arI0y = array('f',
[0] * mesh0.ny) #array to take 1D intensity data (vs Y)
srwl.CalcIntFromElecField(arI0y, wfr, 6, 0, 2, mesh0.eStart, 0, 0)
print('done')
pylab.plot(cor, i[:, 0], cor, arI0y)
pylab.show()
# I(x,y)
#e = s(my_wavefront.absolute_x_coordinates(),my_wavefront.absolute_y_coordinates())
e = self.E_field_as_numpy()
i = e[:, :, 0, 0].real**2 + e[:, :, 0,
0].imag**2 + e[:, :, 0,
1].real**2 + e[:, :, 0,
1].imag**2
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from scipy import meshgrid, array
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
X, Y = meshgrid(self.absolute_y_coordinates(),
self.absolute_x_coordinates())
zs = array(i)
print(zs.shape)
Z = zs.reshape(X.shape)
ax.plot_surface(X, Y, Z)
ax.set_xlabel('X in plane')
ax.set_ylabel('Y in plane')
ax.set_zlabel('Intensity')
plt.show()