def multi_electron(undulator, electron_beam, screen, file):
"""
Calculate wavefront from multi-electron source to screens.
Args: undulator - the parameters of undualtor.
electron_beam - the parameters of electron_beam.
screen - the parameters of screen.
file - file name.
Return: None.
"""
rank = _multi._get_rank()
_create_source_propagate(undulator, electron_beam, screen, file)
# Choose different CMD methods for different size matrix. Depend on the
# experience, 150**2 is setted as a boundary.
if screen["nx"] * screen["ny"] < 150**2:
if rank == 0:
# scipy.sparse.linalg.eigsh is arnoldi method, single process
_source_utils._scipy_sparse_cmd(n_vector=screen["n_vector"],
file_name=file)
else:
# if matrix is too large, multi-process arnoldi method is used.
_source_utils._arnoldi_cmd(screen["nx"] * screen["ny"],
n_vector=screen["n_vector"],
file_name=file)
def _cal_wfrs(undulator, electron_beam, screen, file_name=None, method='srw'):
"""
---------------------------------------------------------------------------
calculate wavefronts from electron source to screens.
args: undulator - the parameters of undualtor.
electron_beam - the parameters of electron_beam.
screen - the parameters of screen.
file - file name.
return: none.
---------------------------------------------------------------------------
"""
# The multi_process parameters
c_rank = _multi._get_rank()
n_rank = _multi._get_size()
# The construction of undulator
und = _undulator(undulator)
und.wave_length()
und.cal_k(electron_beam_energy=electron_beam["energy"])
und_magnetic_structure = und.magnetic_structure()
# The construction of electron beam
e_beam = _srw_electron_beam(electron_beam, und.n_period, und.period_length)
screen['screen'] = screen['screen'] - e_beam.initial_z / 2
e_beam.monte_carlo()
# for 'vib'
vir_part_beam, wavelength, resonance_energy = e_beam.after_monte_carlo(
[0, 0, 0, 0, 0], und.period_length, und.k_vertical, und.k_horizontal)
vir_wavefront = _propagate_wave_front(screen, resonance_energy)
vir_wavefront._cal_wave_front(vir_part_beam, und_magnetic_structure)
vir_wavefront = np.reshape(vir_wavefront.wfr.arEx,
(screen['nx'], screen['ny'], 2))
vir_wavefront = vir_wavefront[:, :, 0] + 1j * vir_wavefront[:, :, 1]
# The parameters of undulator and beam
# reset number of electrons for the multi_layer_svd method
n_electron = electron_beam['n_electron']
if n_electron > _N_SVD_TOL * (n_rank - 1) and n_electron < _N_SVD_TOP * (
n_rank - 1):
pass
else:
n_electron = n_electron + (_N_SVD_OPT - n_electron % _N_SVD_OPT)
n_points = screen['nx'] * screen['ny']
electron_count, electron_index = _support._cal_rank_part(
n_electron, n_rank)
xtick = np.linspace(screen['xstart'], screen['xfin'], screen['nx'])
ytick = np.linspace(screen['ystart'], screen['yfin'], screen['ny'])
gridx, gridy = np.meshgrid(xtick, ytick)
# Set monte carlo random seed
random.seed(c_rank * 123)
newSeed = random.randint(0, 1000000)
random.seed(newSeed)
#------------------------------------------------
# wavefront calculation
if c_rank > 0:
n_electron_per_process = electron_count[c_rank - 1]
e_crank_count, e_crank_index = _support._cal_part(
n_electron_per_process, 500)
_file_utils._create_multi_wfrs(c_rank,
n_electron_per_process,
n_points,
file_name=file_name)
flag = 0
for i_part in e_crank_count:
wfr_arrays = np.zeros((i_part, n_points), dtype=complex)
for ie in range(i_part):
rand_array = [random.gauss(0, 1) for ir in range(5)]
e_beam.monte_carlo()
part_beam, wavelength, resonance_energy = e_beam.after_monte_carlo(
rand_array, und.period_length, und.k_vertical,
und.k_horizontal)
if method is 'srw':
# The magnetic structure setting
wavefront = _propagate_wave_front(screen, resonance_energy)
wavefront._cal_wave_front(part_beam,
und_magnetic_structure)
wfr_arex = np.reshape(np.array(wavefront.wfr.arEx),
(screen['nx'], screen['ny'], 2))
wfr_arex = wfr_arex[:, :, 0] + 1j * wfr_arex[:, :, 1]
ie_wfr = np.reshape(wfr_arex, (1, n_points))
elif method is 'vib':
offset_x = part_beam.partStatMom1.x
angle_x = part_beam.partStatMom1.xp
offset_y = part_beam.partStatMom1.y
angle_y = part_beam.partStatMom1.yp
# phase shift
rot_x_phase = np.exp(
-1j * (2 * np.pi / wavelength) *
(angle_x * gridx -
(1 - np.cos(angle_x)) * screen['screen']))
rot_y_phase = np.exp(
-1j * (2 * np.pi / wavelength) *
(angle_y * gridy -
(1 - np.cos(angle_y)) * screen['screen']))
# offset range
offset_x = offset_x + np.sin(angle_x) * screen['screen']
offset_y = offset_y + np.sin(angle_y) * screen['screen']
# calculate range
x_00, x_01, x_10, x_11 = _support._shift_plane(
offset_x, xtick, screen['xstart'], screen['xfin'],
screen['nx'])
y_00, y_01, y_10, y_11 = _support._shift_plane(
offset_y, ytick, screen['ystart'], screen['yfin'],
screen['ny'])
shift_wfr = np.zeros((screen['nx'], screen['ny']),
dtype=np.complex128)
shift_wfr[y_10:y_11,
x_10:x_11] = ((vir_wavefront * rot_x_phase *
rot_y_phase)[y_00:y_01,
x_00:x_01])
ie_wfr = np.reshape(shift_wfr, (1, n_points))
wfr_arrays[ie, :] = ie_wfr
# save wavefront
_file_utils._save_multi_wfrs(c_rank, wfr_arrays,
e_crank_index[flag][0],
e_crank_index[flag][-1] + 1,
file_name)
flag += 1
# source_file.close()
_multi._send(np.array(1, dtype='i'), dtype=_multi.i, tag=c_rank)
#------------------------------------------------
if c_rank == 0:
import os
import h5py as h
if os.path.isfile(file_name): os.remove(file_name)
with h.File(file_name, 'a') as f:
descript = f.create_group("description")
_support._dict_to_h5(descript, electron_beam)
_support._dict_to_h5(descript, undulator)
_support._dict_to_h5(descript, screen)
descript.create_dataset("wavelength", data=wavelength)
print("wavefront calculation start..... ", flush=True)
for i in range(n_rank - 1):
flag = _multi._recv(1,
np_dtype=np.int,
dtype=_multi.i,
source=i + 1,
tag=i + 1)
print("wavefront calculation finished.", flush=True)
def _arnoldi_cmd(n_points, n_vector=500, file_name="test.h5"):
"""
CSD CMD arnoldi method.
Args: n_points - the size (one dimension) of matrix
n_vector - the number of vectors to approxmite.
file_name -
Return: None.
"""
# The multi process parameters
n_rank = _multi._get_size()
rank = _multi._get_rank()
# multi-process distribution plan
count, index = _cal_rank_part(n_points, n_rank)
order = 2 * n_vector + 1
# construct guess_vector, shcur and hess matrix. loading csd matrix and
# distrub csd parts to different multi-process
if rank == 0:
print("| coherent mode calculation start.|", flush=True)
# construct guess_vector and normalise.
guess_vector = np.random.random(n_points)
guess_vector = guess_vector / np.linalg.norm(guess_vector)
schur = np.zeros((order, n_points), dtype=np.complex128)
schur_conj = np.copy(schur)
schur[0, :] = guess_vector
schur_conj[0, :] = schur[0, :].conj()
hess = np.zeros((order, order), dtype=np.complex128)
with h.File("_cache.h5", 'a') as f:
csdx = np.array(f["coherence/csdx"])
csdy = np.array(f["coherence/csdy"])
miux = np.array(f["coherence/miux"])
miuy = np.array(f["coherence/miuy"])
print("| csd data loading... |", flush=True)
for ic in range(n_rank - 1):
csd_part = np.array(f["coherence/csd"][index[ic], :])
_multi._send(csd_part, dtype=_multi.c, dest=ic + 1, tag=0)
del csd_part
# Recving csd parts
elif rank > 0:
csd_part = _multi._recv((count[rank - 1], n_points),
np_dtype=np.complex128,
dtype=_multi.c,
source=0,
tag=0)
if rank == n_rank - 1:
print("| all csd data loaded. |", flush=True)
print("___________________________________", flush=True)
else:
print("| rank%.2d csd loaded. |" % (rank), flush=True)
# start arnoldi loop
for i in range(1, order):
# for rank >0 process, A*q, A2*q, A3q .... is calculated.
if rank > 0:
vector = _multi._recv(n_points,
np_dtype=np.complex128,
dtype=_multi.f,
source=0,
tag=10 * i + 1)
vector_to_send = np.dot(csd_part, vector)
_multi._send(vector_to_send, dtype=_multi.c, dest=0, tag=i)
# GSM process
if rank == 0:
if (10 * i) % (10 * int(order / 10)) == 0:
print("| %.2d percent |" %
(int(100 * i / order)),
flush=True)
for ir in range(n_rank - 1):
_multi._send(schur[i - 1, :],
dtype=_multi.c,
dest=ir + 1,
tag=10 * i + 1)
idx = 0
for ir in range(n_rank - 1):
vector_part = _multi._recv(count[ir],
np_dtype=np.complex128,
dtype=_multi.c,
source=ir + 1,
tag=i)
schur[i, int(idx):int(idx + count[ir])] = vector_part
idx += count[ir]
for ih in range(i):
hess[ih, i - 1] = schur_conj[ih, :].dot(schur[i, :])
schur[i, :] = schur[i, :] - hess[ih, i - 1] * schur[ih, :]
hess[i, i - 1] = np.linalg.norm(schur[i, :])
schur[i, :] = schur[i, :] / hess[i, i - 1]
schur_conj[i, :] = schur[i, :].conj()
# calculate approximate value and vector
if rank == 0:
hess = hess[0:order, 0:order]
schur = schur[0:order, :]
ratio = np.linalg.eigh(hess)
eig_values = ratio[0][::-1]
vector = ratio[1].transpose()[::-1, :]
vector = vector.transpose()
schur_vector = np.zeros((n_points, n_vector), dtype=np.complex128)
for it in range(n_vector):
t = vector[:, it]
schur_vector[:, it] = schur.transpose().dot(t)
with h.File(file_name, 'a') as f:
if os.path.isfile("_cache.h5"): os.remove("_cache.h5")
coherence = f.create_group("coherence")
f["coherence"].create_dataset("eig_vector", data=schur_vector)
f["coherence"].create_dataset("eig_value", data=eig_values)
f["coherence"].create_dataset("csdx", data=csdx)
f["coherence"].create_dataset("csdy", data=csdy)
f["coherence"].create_dataset("miux", data=miux)
f["coherence"].create_dataset("miuy", data=miuy)
print("| coherent mode calculated. |", flush=True)
print("___________________________________", flush=True)
def _create_source_propagate(undulator, electron_beam, screen, file):
"""
Calculate wavefronts from electron source to screens.
Args: undulator - the parameters of undualtor.
electron_beam - the parameters of electron_beam.
screen - the parameters of screen.
file - file name.
Return: None.
"""
# The multi_process parameters
rank = _multi._get_rank()
n_rank = _multi._get_size()
# The construction of undulator
und = _undulator(undulator)
und.wave_length()
und.cal_k(electron_beam_energy=electron_beam["energy"])
und_magnetic_structure = und.magnetic_structure()
# The construction of electron beam
e_beam = _srw_electron_beam(electron_beam, und.n_period, und.period_length)
e_beam.monte_carlo()
# The parameters of undulator and beam
n_electron = electron_beam['n_electron']
n_points = screen['nx'] * screen['ny']
elec_count, elec_index = _source_utils._cal_rank_part(n_electron, n_rank)
# Set monte carlo random seed
random.seed(rank * 123)
newSeed = random.randint(0, 1000000)
random.seed(newSeed)
#------------------------------------------------
# wavefront calculation
if rank > 0:
n_electron_per_process = elec_count[rank - 1]
_source_utils._create_multi_source_file(rank,
n_electron_per_process,
n_points,
file_name=file)
for i in range(n_electron_per_process):
# random parameters
rand_array = [random.gauss(0, 1) for ir in range(5)]
# monte carlo process
e_beam.monte_carlo()
part_beam, wavelength, resonance_energy = e_beam.after_monte_carlo(
rand_array, und.period_length, und.k_vertical,
und.k_horizontal)
# The magnetic structure setting
wavefront = _propagate_wave_front(screen, resonance_energy)
wavefront._cal_wave_front(part_beam, und_magnetic_structure)
# send wavefront
_source_utils._save_source_file(i, rank, n_points, file, part_beam,
wavefront, screen['nx'],
screen['ny'], resonance_energy)
if rank == 1 and (10 * i) % (
10 * int(n_electron_per_process / 10)) == 0:
print("| %.2d percent |" %
(int(100 * i / n_electron_per_process)),
flush=True)
# source_file.close()
_multi._send(np.array(1, dtype='i'), dtype=_multi.i, tag=rank)
#------------------------------------------------
if rank == 0:
print("___________________________________", flush=True)
print("| wavefront calculation start. |", flush=True)
for i in range(n_rank - 1):
flag = _multi._recv(1,
np_dtype=np.int,
dtype=_multi.i,
source=i + 1,
tag=i + 1)
_source_utils._reconstruct_source_file(file, n_electron, n_points,
electron_beam, undulator,
screen, und.wavelength, n_rank)
print("| wavefront calculation finished. |", flush=True)
print("___________________________________", flush=True)
_source_utils._cal_csd(n_electron,
n_points,
screen["nx"],
screen["ny"],
file_name=file)
def _multi_layer_svd_exceed_0(
n_electron, xcount, ycount, k_cut_off, file_name = "test.h5"
):
"""
---------------------------------------------------------------------------
description: perform coherent mode decompostion by multi layer svd methods.
self._exceed = 0
args: xcount - the pixel number of screen (x).
ycount - the pixel number of screen (y).
k_cut_off - the cut off index of coherent mode.
file_name - file name of the saved source.
return: none.
---------------------------------------------------------------------------
"""
# use the scipy svds algorthm
import scipy.sparse.linalg as ssl
# multi-process parameters
n_rank = _multi._get_size()
c_rank = _multi._get_rank()
# root process
if c_rank == 0:
pre_cmodes = np.zeros(
(xcount * ycount, k_cut_off * (n_rank - 1)),
dtype = complex
)
flag = 0
for i in range(1, n_rank):
icore_cmode = _multi._recv(
(xcount * ycount, k_cut_off), np_dtype = np.complex128,
dtype = _multi.c, source = i, tag = i
)
pre_cmodes[:, flag : flag + k_cut_off] = icore_cmode
flag += k_cut_off
# further calcuated the final svd results
svd_matrix = pre_cmodes
vector, value, evolution = ssl.svds(svd_matrix, k = int(2*k_cut_off))
eig_vector = np.copy(vector[:, ::-1], order = 'C')
value = np.copy(np.abs(value[::-1]), order = 'C')
with h.File(file_name, 'a') as f:
coherence_dict = f.create_group("coherence")
coherence_dict.create_dataset("eig_vector", data = eig_vector)
coherence_dict.create_dataset("eig_value", data = value)
elif c_rank > 0:
crank_file_name = ('_' + file_name.split('.')[0] + '_%.2d.h5') % (c_rank)
with h.File(crank_file_name, 'a') as crank_file:
crank_wfrs = np.array(crank_file['wave_front/arex'])
os.remove(crank_file_name)
vectors, values, evolution = ssl.svds(
crank_wfrs.T, k = int(2*k_cut_off)
)
crank_vectors = np.copy(vectors[:, ::-1], order = 'C')
crank_value = np.copy(np.abs(values[::-1], order = 'C'))
crank_vectors = crank_vectors * crank_value
crank_vectors = np.array(crank_vectors[:, 0 : k_cut_off])
_multi._send(
crank_vectors, dtype = _multi.c,
dest = 0, tag = c_rank
)
def _multi_layer_svd_exceed_1(
n_electron, xcount, ycount, k_cut_off, file_name = "test.h5"
):
"""
---------------------------------------------------------------------------
description: perform coherent mode decompostion by multi layer svd methods.
self._exceed = 1
args: xcount - the pixel number of screen (x).
ycount - the pixel number of screen (y).
k_cut_off - the cut off index of coherent mode.
file_name - file name of the saved source.
return: none.
---------------------------------------------------------------------------
"""
# use the scipy svds algorthm
import scipy.sparse.linalg as ssl
# multi-process parameters
n_rank = _multi._get_size()
c_rank = _multi._get_rank()
r_rank = int(n_electron / _N_SVD_OPT)
# reset electrons
if n_electron > _N_SVD_TOL * (n_rank - 1) and n_electron < _N_SVD_TOP * (n_rank - 1):
pass
else:
n_electron = n_electron + (_N_SVD_OPT - n_electron % _N_SVD_OPT)
# root process
if c_rank == 0:
wfr_arrays = np.zeros((n_electron, int(xcount * ycount)), dtype = complex)
start_index = 0
end_index = 0
for i in range(1, n_rank):
crank_file_name = ('_' + file_name.split('.')[0] + '_%.2d.h5') % (i)
with h.File(crank_file_name, 'a') as crank_file:
crank_wfrs = np.array(crank_file['wave_front/arex'])
print(np.shape(crank_wfrs), flush = True)
end_index += np.shape(crank_wfrs)[0]
print(start_index, flush = True)
print(end_index, flush = True)
wfr_arrays[start_index : end_index, :] = crank_wfrs
start_index += np.shape(crank_wfrs)[0]
os.remove(crank_file_name)
start_index = 0
end_index = 0
for i in range(1, r_rank + 1):
end_index += _N_SVD_OPT
_multi._send(wfr_arrays[start_index : end_index, :], dtype = _multi.c, dest = i, tag = i)
start_index += _N_SVD_OPT
pre_cmodes = np.zeros(
(xcount * ycount, k_cut_off * r_rank),
dtype = complex
)
flag = 0
for i in range(1, r_rank + 1):
icore_cmode = _multi._recv(
(xcount * ycount, k_cut_off), np_dtype = np.complex128,
dtype = _multi.c, source = i, tag = i
)
pre_cmodes[:, flag : flag + k_cut_off] = icore_cmode
flag += k_cut_off
# further calcuated the final svd results
svd_matrix = pre_cmodes
vector, value, evolution = ssl.svds(svd_matrix, k = int(2*k_cut_off))
eig_vector = np.copy(vector[:, ::-1], order = 'C')
value = np.copy(np.abs(value[::-1]), order = 'C')
with h.File(file_name, 'a') as f:
coherence_dict = f.create_group("coherence")
coherence_dict.create_dataset("eig_vector", data = eig_vector)
coherence_dict.create_dataset("eig_value", data = value)
elif c_rank > 0:
if c_rank in list(range(1, r_rank + 1)):
crank_wfrs = _multi._recv(
(_N_SVD_OPT, int(xcount * ycount)), np_dtype = np.complex128,
dtype = _multi.c, source = 0, tag = c_rank
)
vectors, values, evolution = ssl.svds(
crank_wfrs.T, k = int(2*k_cut_off)
)
crank_vectors = np.copy(vectors[:, ::-1], order = 'C')
crank_value = np.copy(np.abs(values[::-1], order = 'C'))
crank_vectors = crank_vectors * crank_value
crank_vectors = np.array(crank_vectors[:, 0 : k_cut_off])
print(np.shape(crank_vectors), flush = True)
_multi._send(
crank_vectors, dtype = _multi.c,
dest = 0, tag = c_rank
)
else:
pass