def communicateWavefront(self, wavefronts, distribution_plan, global_index,
root):
communicator = distribution_plan.communicator()
n_rank = distribution_plan.myRank()
if n_rank != root:
if global_index in distribution_plan.localRows():
local_index = distribution_plan.globalToLocalIndex(global_index)
wavefront = wavefronts[local_index]
as_array = wavefront.asNumpyArray()
for i_element, element in enumerate(as_array):
communicator.send(element, dest=root, tag=i_element)
return None
else:
if global_index in distribution_plan.localRows():
local_index = distribution_plan.globalToLocalIndex(global_index)
wavefront = wavefronts[local_index]
else:
as_array = list()
i_source = distribution_plan.rankByGlobalIndex(global_index)
for i_element in range(3):
as_array.append(
communicator.recv(source=i_source, tag=i_element))
wavefront = NumpyWavefront.fromNumpyArray(as_array[0], as_array[1],
as_array[2])
return wavefront
def communicateWavefront(self, wavefronts, distribution_plan, global_index, root):
communicator = distribution_plan.communicator()
n_rank = distribution_plan.myRank()
if n_rank != root:
if global_index in distribution_plan.localRows():
local_index = distribution_plan.globalToLocalIndex(global_index)
wavefront = wavefronts[local_index]
as_array = wavefront.asNumpyArray()
for i_element, element in enumerate(as_array):
communicator.send(element, dest=root, tag=i_element)
return None
else:
if global_index in distribution_plan.localRows():
local_index = distribution_plan.globalToLocalIndex(global_index)
wavefront = wavefronts[local_index]
else:
as_array = list()
i_source = distribution_plan.rankByGlobalIndex(global_index)
for i_element in range(3):
as_array.append(communicator.recv(source=i_source, tag=i_element))
wavefront = NumpyWavefront.fromNumpyArray(as_array[0],as_array[1],as_array[2])
return wavefront
def read(self, index):
#as_array = list()
# filename = self._file.replace(".npz", "") + str(index)+".wfs.npz"
# srio
filename = self._file.replace(".npz", "_") + "%04d"%(index)+".wfs.npz"
#for i in range(3):
# as_array.append(np.load(filename))
#wavefront = NumpyWavefront.fromNumpyArray(as_array[0], as_array[1], as_array[2])
wavefront = NumpyWavefront.load(filename)
return wavefront
def propagateWavefrontWofry(beamline,
wavefront,
i_mode,
python_to_be_used="python",
keep_temp_files=False):
s_id = str(mpi.COMM_WORLD.Get_rank()) + "_" + gethostname()
wavefront.save("./tmp/tmp%s_in" % s_id)
parameter_lines = "s_id=\"%s\"" % (s_id)
pickle.dump(beamline, open("./tmp/tmp%s_beamline.p" % s_id, "wb"))
lines = """
import pickle
from wofry.propagator.propagator import PropagationManager
from wofry.propagator.propagators2D.fresnel_zoom_xy import FresnelZoomXY2D
from comsyl.waveoptics.ComsylWofryBeamline import ComsylWofryBeamline
# initialize propagator
mypropagator = PropagationManager.Instance()
try:
mypropagator.add_propagator(FresnelZoomXY2D())
except:
print("May be you alreay initialized propagator and stored FresnelZoomXY2D")
beamline = pickle.load(open("./tmp/tmp%s_beamline.p"%s_id,"rb"))
ComsylWofryBeamline.propagate_numpy_wavefront(
"./tmp/tmp%s_in.npz"%s_id,
"./tmp/tmp%s_out.npz"%s_id,
beamline,mypropagator)
"""
file = open("./tmp/tmp%s.py" % s_id, "w")
file.writelines(parameter_lines)
file.writelines("\ni_mode=%d\n" % i_mode) # added srio
file.writelines(lines)
file.close()
os.system(python_to_be_used + " ./tmp/tmp%s.py" % s_id)
if keep_temp_files: # keep a copy of all files
logAll("cp %s %s" %
("./tmp/tmp%s_in.npz" % s_id, "./tmp/tmp%s_in_mode%d.npz" %
(s_id, i_mode)))
logAll("cp %s %s" %
("./tmp/tmp%s_out.npz" % s_id, "./tmp/tmp%s_out_mode%d.npz" %
(s_id, i_mode)))
os.system("cp %s %s" %
("./tmp/tmp%s_in.npz" % s_id, "./tmp/tmp%s_in_mode%d.npz" %
(s_id, i_mode)))
os.system("cp %s %s" %
("./tmp/tmp%s_out.npz" % s_id, "./tmp/tmp%s_out_mode%d.npz" %
(s_id, i_mode)))
return NumpyWavefront.load("./tmp/tmp%s_out.npz" % s_id)
def propagateWavefront(srw_beamline,
wavefront,
rx,
drx,
ry,
dry,
rescale_x,
rescale_y,
i_mode,
python_to_be_used="python"):
if isMaster():
if not os.path.exists("tmp"):
os.mkdir("tmp")
s_id = str(mpi.COMM_WORLD.Get_rank()) + "_" + gethostname()
wavefront.save("./tmp/tmp%s_in" % s_id)
parameter_lines = "rx=%f\ndrx=%f\nry=%f\ndry=%f\nrescale_x=%f\nrescale_y=%f\ns_id=\"%s\"" % (
rx, drx, ry, dry, rescale_x, rescale_y, s_id)
pickle.dump(srw_beamline, open("./tmp/tmp%s_beamline.p" % s_id, "wb"))
lines = """
import pickle
from srwlib import *
from comsyl.waveoptics.SRWAdapter import SRWAdapter
from comsyl.waveoptics.Wavefront import NumpyWavefront, SRWWavefront
wavefront = NumpyWavefront.load("./tmp/tmp%s_in.npz"%s_id)
adapter = SRWAdapter()
wfr = adapter.SRWWavefrontFromWavefront(wavefront,
rx,
drx,
ry,
dry,rescale_x,rescale_y)
#print("Doing propagation in external call")
srw_beamline = pickle.load(open("./tmp/tmp%s_beamline.p"%s_id,"rb"))
srwl.PropagElecField(wfr, srw_beamline)
tmp = SRWWavefront(wfr).toNumpyWavefront()
wfr = None
tmp.save("./tmp/tmp%s_out" % s_id)
"""
file = open("./tmp/tmp%s.py" % s_id, "w")
file.writelines(parameter_lines)
file.writelines("\ni_mode=%d\n" % i_mode) # added srio
file.writelines(lines)
file.close()
os.system(python_to_be_used + " ./tmp/tmp%s.py" % s_id)
return NumpyWavefront.load("./tmp/tmp%s_out.npz" % s_id)
def coherentModeAsWavefront(self, i_mode):
mode = self.coherentMode(i_mode)
e_field = np.zeros((1, mode.shape[0], mode.shape[1], 2),
dtype=mode.dtype)
e_field[0, :, :, 0] = mode[:, :]
e_field[abs(e_field) < 0.00001] = 0.0
wavefront = NumpyWavefront(
e_field,
self.xCoordinates().min(),
self.xCoordinates().max(),
self.yCoordinates().min(),
self.yCoordinates().max(),
self._wavefront.z(),
self._wavefront.energies(),
)
return wavefront
def propagateWavefrontWofry(beamline, wavefront, i_mode, python_to_be_used="python", keep_temp_files=False):
s_id=str(mpi.COMM_WORLD.Get_rank())+"_"+gethostname()
wavefront.save("./tmp/tmp%s_in"%s_id)
parameter_lines = "s_id=\"%s\"" % (s_id)
pickle.dump(beamline, open("./tmp/tmp%s_beamline.p"%s_id,"wb"))
lines ="""
import pickle
from wofry.propagator.propagator import PropagationManager
from wofry.propagator.propagators2D.fresnel_zoom_xy import FresnelZoomXY2D
from comsyl.waveoptics.ComsylWofryBeamline import ComsylWofryBeamline
# initialize propagator
mypropagator = PropagationManager.Instance()
try:
mypropagator.add_propagator(FresnelZoomXY2D())
except:
print("May be you alreay initialized propagator and stored FresnelZoomXY2D")
beamline = pickle.load(open("./tmp/tmp%s_beamline.p"%s_id,"rb"))
ComsylWofryBeamline.propagate_numpy_wavefront(
"./tmp/tmp%s_in.npz"%s_id,
"./tmp/tmp%s_out.npz"%s_id,
beamline,mypropagator)
"""
file = open("./tmp/tmp%s.py"%s_id, "w")
file.writelines(parameter_lines)
file.writelines("\ni_mode=%d\n"%i_mode) # added srio
file.writelines(lines)
file.close()
os.system(python_to_be_used+" ./tmp/tmp%s.py"%s_id)
if keep_temp_files: # keep a copy of all files
logAll("cp %s %s" % ("./tmp/tmp%s_in.npz" % s_id, "./tmp/tmp%s_in_mode%d.npz" % (s_id, i_mode)))
logAll("cp %s %s" % ("./tmp/tmp%s_out.npz" % s_id, "./tmp/tmp%s_out_mode%d.npz" % (s_id, i_mode)))
os.system("cp %s %s" % ("./tmp/tmp%s_in.npz" % s_id, "./tmp/tmp%s_in_mode%d.npz" % (s_id, i_mode)))
os.system("cp %s %s" % ("./tmp/tmp%s_out.npz" % s_id, "./tmp/tmp%s_out_mode%d.npz" % (s_id, i_mode)))
return NumpyWavefront.load("./tmp/tmp%s_out.npz"%s_id)
def propagateWavefront(srw_beamline, wavefront, rx, drx, ry, dry, rescale_x, rescale_y, i_mode, python_to_be_used="python"):
if isMaster():
if not os.path.exists("tmp"):
os.mkdir("tmp")
s_id=str(mpi.COMM_WORLD.Get_rank())+"_"+gethostname()
wavefront.save("./tmp/tmp%s_in"%s_id)
parameter_lines = "rx=%f\ndrx=%f\nry=%f\ndry=%f\nrescale_x=%f\nrescale_y=%f\ns_id=\"%s\"" % (rx, drx, ry, dry, rescale_x, rescale_y, s_id)
pickle.dump(srw_beamline, open("./tmp/tmp%s_beamline.p"%s_id,"wb"))
lines ="""
import pickle
from srwlib import *
from comsyl.waveoptics.SRWAdapter import SRWAdapter
from comsyl.waveoptics.Wavefront import NumpyWavefront, SRWWavefront
wavefront = NumpyWavefront.load("./tmp/tmp%s_in.npz"%s_id)
adapter = SRWAdapter()
wfr = adapter.SRWWavefrontFromWavefront(wavefront,
rx,
drx,
ry,
dry,rescale_x,rescale_y)
#print("Doing propagation in external call")
srw_beamline = pickle.load(open("./tmp/tmp%s_beamline.p"%s_id,"rb"))
srwl.PropagElecField(wfr, srw_beamline)
tmp = SRWWavefront(wfr).toNumpyWavefront()
wfr = None
tmp.save("./tmp/tmp%s_out" % s_id)
"""
file = open("./tmp/tmp%s.py"%s_id, "w")
file.writelines(parameter_lines)
file.writelines("\ni_mode=%d\n"%i_mode) # added srio
file.writelines(lines)
file.close()
os.system(python_to_be_used+" ./tmp/tmp%s.py"%s_id)
return NumpyWavefront.load("./tmp/tmp%s_out.npz"%s_id)
def _buildForXY(self, electron_beam, x_0, y_0, xp_0, yp_0, z, X, Y):
#TODO: X dimension equals Y dimension?
Y = X
beam = ElectronBeam(Electron_energy=electron_beam.energy(),
I_current=electron_beam.averageCurrent())
undulator = Undulator(K=self._undulator.K_vertical(),
period_length=self._undulator.periodLength(),
length=self._undulator.length())
initial_conditions = SourceUndulatorPlane(
undulator=undulator, electron_beam=beam,
magnetic_field=None).choose_initial_contidion_automatic()
v_z = initial_conditions[2]
initial_conditions[0] = xp_0 * speed_of_light
initial_conditions[1] = yp_0 * speed_of_light
initial_conditions[2] = np.sqrt(beam.electron_speed()**2 - xp_0**2 -
yp_0**2) * speed_of_light
initial_conditions[3] = x_0
initial_conditions[4] = y_0
if self._source_position == VIRTUAL_SOURCE_CENTER:
initial_conditions[5] = 0.0
print("initial cond:", initial_conditions)
simulation = create_simulation(
magnetic_structure=undulator,
electron_beam=beam,
traj_method=TRAJECTORY_METHOD_ODE,
rad_method=RADIATION_METHOD_NEAR_FIELD, #RADIATION_METHOD_FARFIELD,
distance=z,
X=X,
Y=Y,
photon_energy=self._photon_energy,
initial_condition=initial_conditions)
#simulation.trajectory.plot_3D()
#simulation.trajectory.plot()
#simulation.radiation.plot()
electrical_field = simulation.radiation_fact.calculate_electrical_field(
trajectory=simulation.trajectory,
source=simulation.source,
distance=simulation.radiation.distance,
X_array=simulation.radiation.X,
Y_array=simulation.radiation.Y)
efield = electrical_field.electrical_field()[np.newaxis, :, :, :]
efield = efield[:, :, :, 0:2]
calc_wavefront = NumpyWavefront(
e_field=efield,
x_start=X.min(),
x_end=X.max(),
y_start=Y.min(),
y_end=Y.max(),
z=z,
energies=np.array([self._photon_energy]),
)
#calc_wavefront.showEField()
self._last_simulation = simulation
self._last_initial_conditions = initial_conditions.copy()
return calc_wavefront
def fromDictionary(data_dict):
sigma_matrix = SigmaMatrix.fromNumpyArray(data_dict["sigma_matrix"])
undulator = undulator_from_numpy_array(data_dict["undulator"])
detuning_parameter = data_dict["detuning_parameter"][0]
energy = data_dict["energy"][0]
electron_beam_energy = data_dict["electron_beam_energy"][0]
np_wavefront_0=data_dict["wavefront_0"]
np_wavefront_1=data_dict["wavefront_1"]
np_wavefront_2=data_dict["wavefront_2"]
wavefront = NumpyWavefront.fromNumpyArray(np_wavefront_0, np_wavefront_1, np_wavefront_2)
try:
np_exit_slit_wavefront_0=data_dict["exit_slit_wavefront_0"]
np_exit_slit_wavefront_1=data_dict["exit_slit_wavefront_1"]
np_exit_slit_wavefront_2=data_dict["exit_slit_wavefront_2"]
exit_slit_wavefront = NumpyWavefront.fromNumpyArray(np_exit_slit_wavefront_0, np_exit_slit_wavefront_1, np_exit_slit_wavefront_2)
except:
exit_slit_wavefront = wavefront.clone()
try:
weighted_fields = data_dict["weighted_fields"]
except:
weighted_fields = None
srw_wavefront_rx=data_dict["srw_wavefront_rx"][0]
srw_wavefront_ry=data_dict["srw_wavefront_ry"][0]
srw_wavefront_drx = data_dict["srw_wavefront_drx"][0]
srw_wavefront_dry = data_dict["srw_wavefront_dry"][0]
info_string = str(data_dict["info"])
info = AutocorrelationInfo.fromString(info_string)
sampling_factor=data_dict["sampling_factor"][0]
minimal_size=data_dict["minimal_size"][0]
beam_energies = data_dict["beam_energies"]
static_electron_density = data_dict["static_electron_density"]
coordinates_x = data_dict["twoform_0"]
coordinates_y = data_dict["twoform_1"]
diagonal_elements = data_dict["twoform_2"]
eigenvalues = data_dict["twoform_3"]
twoform_vectors = data_dict["twoform_4"]
twoform = Twoform(coordinates_x, coordinates_y, diagonal_elements, eigenvalues, twoform_vectors)
eigenvector_errors = data_dict["twoform_5"]
twoform.setEigenvectorErrors(eigenvector_errors)
af = AutocorrelationFunction(sigma_matrix, undulator, detuning_parameter,energy,electron_beam_energy,
wavefront,exit_slit_wavefront,srw_wavefront_rx, srw_wavefront_drx, srw_wavefront_ry, srw_wavefront_dry,
sampling_factor,minimal_size, beam_energies, weighted_fields,
static_electron_density, twoform,
info)
return af
#
# # just to test
# a = ComsylWofryBeamline.initialize_from_lists(
# list_with_syned_optical_elements=[WOScreen(),WOScreen(),WOScreen()],
# list_with_syned_coordinates=[
# ElementCoordinates(p=0.0, q=28.3, angle_radial=0.0, angle_azimuthal=0.0),
# ElementCoordinates(p=0.0, q=28.3, angle_radial=0.0, angle_azimuthal=0.0),
# ElementCoordinates(p=0.0, q=28.3, angle_radial=0.0, angle_azimuthal=0.0),],
# list_with_wofry_propagator_handlers=['FRESNEL_ZOOM_XY_2D','FRESNEL_ZOOM_XY_2D','FRESNEL_ZOOM_XY_2D',],
# list_with_wofry_propagator_specific_parameters=[
# {'shift_half_pixel': 1, 'magnification_x': 8.0, 'magnification_y': 8.0},
# {'shift_half_pixel': 1, 'magnification_x': 8.0, 'magnification_y': 8.0},
# {'shift_half_pixel': 1, 'magnification_x': 8.0, 'magnification_y': 8.0},])
#
#
# print(a.info())
#
# print(a.to_json())
bl = ComsylWofryBeamline.load_pickle(
"/users/srio/Working/paper-hierarchical/CODE-COMSYL/tmp/tmp0_chac_beamline.p"
)
print(bl.info())
from comsyl.waveoptics.Wavefront import NumpyWavefront
w0 = NumpyWavefront.load(
"/users/srio/Working/paper-hierarchical/CODE-COMSYL/tmp/tmp0_chac_in.npz"
)
print(w0)
def fromDictionary(data_dict):
sigma_matrix = SigmaMatrix.fromNumpyArray(data_dict["sigma_matrix"])
undulator = undulator_from_numpy_array(data_dict["undulator"])
detuning_parameter = data_dict["detuning_parameter"][0]
energy = data_dict["energy"][0]
electron_beam_energy = data_dict["electron_beam_energy"][0]
np_wavefront_0 = data_dict["wavefront_0"]
np_wavefront_1 = data_dict["wavefront_1"]
np_wavefront_2 = data_dict["wavefront_2"]
wavefront = NumpyWavefront.fromNumpyArray(np_wavefront_0,
np_wavefront_1,
np_wavefront_2)
try:
np_exit_slit_wavefront_0 = data_dict["exit_slit_wavefront_0"]
np_exit_slit_wavefront_1 = data_dict["exit_slit_wavefront_1"]
np_exit_slit_wavefront_2 = data_dict["exit_slit_wavefront_2"]
exit_slit_wavefront = NumpyWavefront.fromNumpyArray(
np_exit_slit_wavefront_0, np_exit_slit_wavefront_1,
np_exit_slit_wavefront_2)
except:
exit_slit_wavefront = wavefront.clone()
try:
weighted_fields = data_dict["weighted_fields"]
except:
weighted_fields = None
srw_wavefront_rx = data_dict["srw_wavefront_rx"][0]
srw_wavefront_ry = data_dict["srw_wavefront_ry"][0]
srw_wavefront_drx = data_dict["srw_wavefront_drx"][0]
srw_wavefront_dry = data_dict["srw_wavefront_dry"][0]
info_string = str(data_dict["info"])
info = AutocorrelationInfo.fromString(info_string)
sampling_factor = data_dict["sampling_factor"][0]
minimal_size = data_dict["minimal_size"][0]
beam_energies = data_dict["beam_energies"]
static_electron_density = data_dict["static_electron_density"]
coordinates_x = data_dict["twoform_0"]
coordinates_y = data_dict["twoform_1"]
diagonal_elements = data_dict["twoform_2"]
eigenvalues = data_dict["twoform_3"]
twoform_vectors = data_dict["twoform_4"]
twoform = Twoform(coordinates_x, coordinates_y, diagonal_elements,
eigenvalues, twoform_vectors)
eigenvector_errors = data_dict["twoform_5"]
twoform.setEigenvectorErrors(eigenvector_errors)
af = AutocorrelationFunction(
sigma_matrix, undulator, detuning_parameter, energy,
electron_beam_energy, wavefront, exit_slit_wavefront,
srw_wavefront_rx, srw_wavefront_drx, srw_wavefront_ry,
srw_wavefront_dry, sampling_factor, minimal_size, beam_energies,
weighted_fields, static_electron_density, twoform, info)
return af
# from wofry.beamline.optical_elements.ideal_elements.screen import WOScreen
#
#
# # just to test
# a = ComsylWofryBeamline.initialize_from_lists(
# list_with_syned_optical_elements=[WOScreen(),WOScreen(),WOScreen()],
# list_with_syned_coordinates=[
# ElementCoordinates(p=0.0, q=28.3, angle_radial=0.0, angle_azimuthal=0.0),
# ElementCoordinates(p=0.0, q=28.3, angle_radial=0.0, angle_azimuthal=0.0),
# ElementCoordinates(p=0.0, q=28.3, angle_radial=0.0, angle_azimuthal=0.0),],
# list_with_wofry_propagator_handlers=['FRESNEL_ZOOM_XY_2D','FRESNEL_ZOOM_XY_2D','FRESNEL_ZOOM_XY_2D',],
# list_with_wofry_propagator_specific_parameters=[
# {'shift_half_pixel': 1, 'magnification_x': 8.0, 'magnification_y': 8.0},
# {'shift_half_pixel': 1, 'magnification_x': 8.0, 'magnification_y': 8.0},
# {'shift_half_pixel': 1, 'magnification_x': 8.0, 'magnification_y': 8.0},])
#
#
# print(a.info())
#
# print(a.to_json())
bl = ComsylWofryBeamline.load_pickle("/users/srio/Working/paper-hierarchical/CODE-COMSYL/tmp/tmp0_chac_beamline.p")
print(bl.info())
from comsyl.waveoptics.Wavefront import NumpyWavefront
w0 = NumpyWavefront.load("/users/srio/Working/paper-hierarchical/CODE-COMSYL/tmp/tmp0_chac_in.npz")
print(w0)