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 _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 _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_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