def plot_images(self, type_, title='', **kwargs):
if type_ == 'raw':
images = self.raw_image_objs
elif type_ == 'cut':
images = self.cut_images
elif type_ == 'tE':
images = self.images_tE
elif type_ == 'slice':
images = self.images_sliced
sp_ctr = np.inf
ny, nx = 3, 3
subplot = ms.subplot_factory(ny, nx, grid=False)
figs = []
subplots = []
for n_image, image in enumerate(images):
if sp_ctr > ny*nx:
fig = ms.figure('%s Images %s' % (title, type_))
figs.append(fig)
this_subplots = []
subplots.append(this_subplots)
sp_ctr = 1
sp = subplot(sp_ctr, title='Image %i' % n_image, xlabel=image.xlabel, ylabel=image.ylabel)
sp_ctr += 1
this_subplots.append(sp)
slice_dict = None
if type_ in ('tE', 'slice') and hasattr(self, 'slice_dicts'):
slice_dict = self.slice_dicts[n_image]
image.plot_img_and_proj(sp, slice_dict=slice_dict, **kwargs)
return figs, subplots
def x_to_t(self, wake_x, wake_time, debug=False, print_=False):
if wake_time[1] < wake_time[0]:
wake_x = wake_x[::-1]
wake_time = wake_time[::-1]
new_img0 = np.zeros_like(self.image)
new_t_axis = np.linspace(wake_time.min(), wake_time.max(), self.image.shape[1])
x_interp = np.interp(new_t_axis, wake_time, wake_x)
to_print = []
for t_index, (t, x) in enumerate(zip(new_t_axis, x_interp)):
x_index = np.argmin((self.x_axis - x)**2)
new_img0[:,t_index] = self.image[:,x_index]
if print_:
to_print.append('%i %i %.1f %.1f' % (t_index, x_index, t*1e15, x*1e6))
if print_:
print('\n'.join(to_print))
diff_x = np.concatenate([np.diff(x_interp), [0]])
new_img = new_img0 * np.abs(diff_x)
new_img = new_img / new_img.sum() * self.image.sum()
output = self.child(new_img, new_t_axis, self.y_axis, x_unit='s', xlabel='t (fs)')
if debug:
ms.figure('Debug x_to_t')
subplot = ms.subplot_factory(2,3)
sp_ctr = 1
sp = subplot(sp_ctr, title='Wake', xlabel='time [fs]', ylabel='Screen x [mm]')
sp_ctr += 1
sp.plot(wake_time*1e15, wake_x*1e3)
sp = subplot(sp_ctr, title='Image projection X', xlabel='x [mm]', ylabel='Intensity (arb. units)')
sp_ctr += 1
sp.plot(self.x_axis*1e3, self.image.sum(axis=-2))
sp = subplot(sp_ctr, title='Image projection T', xlabel='t [fs]', ylabel='Intensity (arb. units)')
sp_ctr += 1
sp.plot(output.x_axis*1e15, output.image.sum(axis=-2))
sp = subplot(sp_ctr, title='Image old', xlabel='x [mm]', ylabel='y [mm]', grid=False)
sp_ctr += 1
self.plot_img_and_proj(sp)
sp = subplot(sp_ctr, title='Image new', xlabel='t [fs]', ylabel='y [mm]', grid=False)
sp_ctr += 1
output.plot_img_and_proj(sp)
sp = subplot(sp_ctr, title='Image new 0', xlabel='t [fs]', ylabel=' y [mm]', grid=False)
sp_ctr += 1
new_obj0 = self.child(new_img0, new_t_axis, self.y_axis, x_unit='s', xlabel='t (fs)')
new_obj0.plot_img_and_proj(sp)
#ms.plt.show()
#import pdb; pdb.set_trace()
return output
def lasing_figures():
output = []
fig = plt.figure()
fig.canvas.set_window_title('Lasing reconstruction')
subplot = ms.subplot_factory(3,3, grid=False)
plot_handles = tuple((subplot(sp_ctr) for sp_ctr in range(1, 1+8)))
output.append((fig, plot_handles))
fig.subplots_adjust(hspace=0.5, wspace=0.3)
fig = plt.figure()
fig.subplots_adjust(hspace=0.4)
subplot = ms.subplot_factory(2,2, grid=False)
plot_handles = tuple((subplot(sp_ctr) for sp_ctr in range(1, 1+4)))
output.append((fig, plot_handles))
clear_lasing(output)
return output
def lasing_figure(figsize=None):
fig = plt.figure(figsize=figsize)
fig.canvas.set_window_title('Lasing reconstruction')
fig.subplots_adjust(hspace=0.4)
subplot = ms.subplot_factory(2,3)
subplots = [subplot(sp_ctr) for sp_ctr in range(1, 1+6)]
clear_lasing_figure(*subplots)
return fig, subplots
def reconstruction_figure(figsize=None):
fig = plt.figure(figsize=figsize)
fig.canvas.set_window_title('Current reconstruction')
fig.subplots_adjust(hspace=0.4)
subplot = ms.subplot_factory(2,2)
subplots = [subplot(sp_ctr) for sp_ctr in range(1, 1+4)]
clear_reconstruction(*subplots)
return fig, subplots
def screen_calibration_figure():
fig = plt.figure()
fig.canvas.set_window_title('Screen center calibration')
fig.subplots_adjust(hspace=0.35)
sp_ctr = 1
subplot = ms.subplot_factory(1, 1)
sp_proj = subplot(sp_ctr, xlabel='x (mm)', ylabel='Intensity (arb. units)', sciy=True)
sp_ctr += 1
clear_screen_calibration(sp_proj)
return fig, (sp_proj, )
def plot_slice_dict(slice_dict):
subplot = ms.subplot_factory(3, 3)
sp_ctr = np.inf
for n_slice, slice_gf in enumerate(slice_dict['slice_gf']):
slice_sigma = slice_dict['slice_sigma_sq'][n_slice]
slice_rms = slice_dict['slice_rms_sq'][n_slice]
slice_cut = slice_dict['slice_cut_rms_sq'][n_slice]
if sp_ctr > 9:
ms.figure('Investigate slice')
sp_ctr = 1
sp = subplot(sp_ctr, title='Slice %i, $\sigma$=%.1e, rms=%.1e, cut=%.1e' % (n_slice, slice_sigma, slice_rms, slice_cut))
sp_ctr += 1
slice_gf.plot_data_and_fit(sp)
sp.legend()
def inspect(watcher, bins=(100, 100), show=True, title=None, charge=200e-12):
dimensions = 'x', 'xp', 'y', 'yp', 't', 'p'
if title is None:
title = watcher
fig = ms.figure(title)
plt.subplots_adjust(hspace=0.35, wspace=0.25, bottom=0.1)
subplot = ms.subplot_factory(4, 4)
sp_ctr = 1
for dim_ctr1, dim1 in enumerate(dimensions):
for dim_ctr2, dim2 in enumerate(dimensions[dim_ctr1 + 1:],
dim_ctr1 + 1):
sp = subplot(sp_ctr,
title='%s-%s' % (dim1, dim2),
xlabel=dim1,
ylabel=dim2,
scix=True,
sciy=True,
grid=False)
sp_ctr += 1
x_arr = watcher[dim1]
y_arr = watcher[dim2]
if dim1 == 't':
x_arr = x_arr - x_arr.mean()
if dim2 == 't':
y_arr = y_arr - y_arr.mean()
sp.hist2d(x_arr, y_arr, bins=bins)
sp = subplot(sp_ctr,
title='Beam current',
xlabel='t',
ylabel='I [A]',
scix=True,
sciy=True,
grid=False)
xx, yy = watcher.get_current('t', bins=bins[0], charge=charge)
sp.step(xx, yy)
if show:
plt.show()
return fig
final_profile = d['final_profile']
xx = d['xx']
tt = d['tt']
meas_screen = d['meas_screen']
if xx[1] < xx[0]:
xx = xx[::-1]
tt = tt[::-1]
image_obj = iap.Image(image, x_axis, y_axis)
image_cut = image_obj.cut(xx.min(), xx.max())
image2 = image_cut.reshape_x(len(final_profile))
figure = ms.figure('Backtrack image')
ms.plt.subplots_adjust(hspace=0.3)
subplot = ms.subplot_factory(2, 3, grid=False)
sp_ctr = 1
sp = subplot(sp_ctr, title='X space 1', xlabel='x [mm]', ylabel='y [mm]')
sp_ctr += 1
image_cut.plot_img_and_proj(sp)
sp = subplot(sp_ctr, title='X space 2', xlabel='x [mm]', ylabel='y [mm]')
sp_ctr += 1
image2.plot_img_and_proj(sp)
new_img = image2.x_to_t(xx, tt, debug=True, print_=False)
ms.plt.figure(figure.number)
sp = subplot(sp_ctr, title='T space', xlabel='t [fs]', ylabel='y [mm]')
sp_ctr += 1
def main():
fig_paper = ms.figure('Comparison plots')
subplot = ms.subplot_factory(2, 2)
sp_ctr_paper = 1
#images0 = dict0['projx'][-1]
#x_axis = dict0['x_axis']*1e-6
#if np.diff(x_axis)[0] < 0:
# x_axis = x_axis[::-1]
# invert_x = True
#else:
# invert_x = False
process_dict = {
'Long': {
'filename': file38,
'main_dict': dict38,
'proj0': dict0['projx'][-1],
'x_axis0': dict0['x_axis'] * 1e-6,
'n_offset': None,
'filename0': file0,
'blmeas': blmeas38,
'flipx': False,
},
'Medium': {
'filename': file25,
'main_dict': dict25,
'proj0': dict25['projx'][7],
'x_axis0': dict25['x_axis'] * 1e-6,
'n_offset': 0,
'filename0': file25,
'blmeas': blmeas25,
'flipx': False,
},
}
for main_label, p_dict in process_dict.items():
#if main_label != 'Medium':
# continue
projx0 = p_dict['proj0']
x_axis0 = p_dict['x_axis0']
if np.diff(x_axis0)[0] < 0:
x_axis0 = x_axis0[::-1]
invert_x0 = True
all_mean = []
for proj in projx0:
screen = get_screen_from_proj(proj, x_axis0, invert_x0)
xx, yy = screen._xx, screen._yy
gf = gaussfit.GaussFit(xx, yy)
all_mean.append(gf.mean)
mean0 = np.mean(all_mean)
timestamp0 = misc.get_timestamp(os.path.basename(p_dict['filename0']))
tracker0 = tracking.Tracker(
archiver_dir + 'archiver_api_data/2020-10-03.h5', timestamp0,
struct_lengths, n_particles, n_emittances, screen_bins,
screen_cutoff, smoothen, profile_cutoff, len_profile)
bp_test = tracking.get_gaussian_profile(40e-15, tt_halfrange,
len_profile, charge,
tracker0.energy_eV)
screen_sim = tracker0.matrix_forward(bp_test, [10e-3, 10e-3],
[0, 0])['screen']
all_emittances = []
for proj in projx0:
screen_meas = get_screen_from_proj(proj, x_axis0, invert_x0)
emittance_fit = misc.fit_nat_beamsize(screen_meas, screen_sim,
n_emittances[0])
all_emittances.append(emittance_fit)
new_emittance = np.mean(all_emittances)
print(main_label, 'Emittance [nm]', new_emittance * 1e9)
n_emittances[0] = new_emittance
dict_ = p_dict['main_dict']
file_ = p_dict['filename']
x_axis = dict_['x_axis'] * 1e-6
y_axis = dict_['y_axis'] * 1e-6
n_offset = p_dict['n_offset']
if np.diff(x_axis)[0] < 0:
x_axis = x_axis[::-1]
invert_x = True
else:
invert_x = False
if np.diff(y_axis)[0] < 0:
y_axis = y_axis[::-1]
invert_y = True
else:
invert_y = False
timestamp = misc.get_timestamp(os.path.basename(file_))
tracker = tracking.Tracker(
archiver_dir + 'archiver_api_data/2020-10-03.h5', timestamp,
struct_lengths, n_particles, n_emittances, screen_bins,
screen_cutoff, smoothen, profile_cutoff, len_profile)
blmeas = p_dict['blmeas']
flip_measured = p_dict['flipx']
profile_meas = tracking.profile_from_blmeas(blmeas,
tt_halfrange,
charge,
tracker.energy_eV,
subtract_min=True)
profile_meas.reshape(len_profile)
profile_meas2 = tracking.profile_from_blmeas(blmeas,
tt_halfrange,
charge,
tracker.energy_eV,
subtract_min=True,
zero_crossing=2)
profile_meas2.reshape(len_profile)
if flip_measured:
profile_meas.flipx()
else:
profile_meas2.flipx()
profile_meas.cutoff(1e-2)
profile_meas2.cutoff(1e-2)
beam_offsets = [0., -(dict_['value'] * 1e-3 - mean_struct2)]
distance_um = (gaps[n_streaker] / 2. - beam_offsets[n_streaker]) * 1e6
if n_offset is not None:
distance_um = distance_um[n_offset]
beam_offsets = [beam_offsets[0], beam_offsets[1][n_offset]]
tdc_screen1 = tracker.matrix_forward(profile_meas, gaps,
beam_offsets)['screen']
tdc_screen2 = tracker.matrix_forward(profile_meas, gaps,
beam_offsets)['screen']
plt.figure(fig_paper.number)
sp_profile_comp = subplot(sp_ctr_paper,
title=main_label,
xlabel='t [fs]',
ylabel='Intensity (arb. units)')
sp_ctr_paper += 1
profile_meas.plot_standard(sp_profile_comp,
norm=True,
color='black',
label='TDC',
center='Right')
ny, nx = 2, 4
subplot = ms.subplot_factory(ny, nx)
sp_ctr = np.inf
all_profiles, all_screens = [], []
if n_offset is None:
projections = dict_['projx']
else:
projections = dict_['projx'][n_offset]
for n_image in range(len(projections)):
screen = get_screen_from_proj(projections[n_image], x_axis,
invert_x)
screen.crop()
screen._xx = screen._xx - mean0
gauss_dict = tracker.find_best_gauss(
sig_t_range,
tt_halfrange,
screen,
gaps,
beam_offsets,
n_streaker,
charge,
self_consistent=self_consistent)
best_screen = gauss_dict['reconstructed_screen']
best_screen.cutoff(1e-3)
best_screen.crop()
best_profile = gauss_dict['reconstructed_profile']
if n_image == 0:
screen00 = screen
bp00 = best_profile
best_screen00 = best_screen
best_gauss = gauss_dict['best_gauss']
if sp_ctr > (ny * nx):
ms.figure('All reconstructions Distance %i %s' %
(distance_um, main_label))
sp_ctr = 1
if n_image % 2 == 0:
sp_profile = subplot(sp_ctr, title='Reconstructions')
sp_ctr += 1
sp_screen = subplot(sp_ctr, title='Screens')
sp_ctr += 1
profile_meas.plot_standard(sp_profile,
color='black',
label='Measured',
norm=True,
center='Right')
tdc_screen1.plot_standard(sp_screen, color='black')
color = screen.plot_standard(sp_screen,
label=n_image)[0].get_color()
best_screen.plot_standard(sp_screen, color=color, ls='--')
best_profile.plot_standard(sp_profile,
label=n_image,
norm=True,
center='Right')
sp_profile.legend()
sp_screen.legend()
all_profiles.append(best_profile)
# Averaging the reconstructed profiles
all_profiles_time, all_profiles_current = [], []
for profile in all_profiles:
profile.shift('Right')
#all_profiles_time.append(profile.time - profile.time[np.argmax(profile.current)])
all_profiles_time.append(profile.time)
new_time = np.linspace(min(x.min() for x in all_profiles_time),
max(x.max() for x in all_profiles_time),
len_profile)
for tt, profile in zip(all_profiles_time, all_profiles):
new_current = np.interp(new_time,
tt,
profile.current,
left=0,
right=0)
new_current *= charge / new_current.sum()
all_profiles_current.append(new_current)
all_profiles_current = np.array(all_profiles_current)
mean_profile = np.mean(all_profiles_current, axis=0)
std_profile = np.std(all_profiles_current, axis=0)
average_profile = tracking.BeamProfile(new_time, mean_profile,
tracker.energy_eV, charge)
average_profile.plot_standard(sp_profile_comp,
label='Reconstructed',
norm=True,
center='Right')
ms.figure('Test averaging %s' % main_label)
sp = plt.subplot(1, 1, 1)
for yy in all_profiles_current:
sp.plot(new_time, yy / np.trapz(yy, new_time), lw=0.5)
to_plot = [
('Average', new_time, mean_profile, 'black', 3),
('+1 STD', new_time, mean_profile + std_profile, 'black', 1),
('-1 STD', new_time, mean_profile - std_profile, 'black', 1),
]
integral = np.trapz(mean_profile, new_time)
for pm, ctr, color in [(profile_meas, 1, 'red'),
(profile_meas2, 2, 'green')]:
#factor = integral/np.trapz(pm.current, pm.time)
#t_meas = pm.time-pm.time[np.argmax(pm.current)]
i_meas = np.interp(new_time, pm.time, pm.current)
bp = tracking.BeamProfile(new_time,
i_meas,
energy_eV=tracker.energy_eV,
charge=charge)
bp.shift('Right')
to_plot.append(('TDC %i' % ctr, bp.time, bp.current, color, 3))
for label, tt, profile, color, lw in to_plot:
gf = gaussfit.GaussFit(tt, profile)
width_fs = gf.sigma * 1e15
if label is None:
label = ''
label = (label + ' %i fs' % width_fs).strip()
factor = np.trapz(profile, tt)
sp.plot(tt, profile / factor, color=color, lw=lw, label=label)
sp.legend(title='Gaussian fit $\sigma$')
plt.show()
for ctr, (data_file, label) in enumerate([(lasing_off_file, 'Lasing Off'),
(data_files[-1], 'Lasing On')]):
data_dict = loadH5Recursive(data_file)
images = data_dict['pyscan_result']['image'].astype(np.float64)
x_axis = data_dict['pyscan_result']['x_axis'].astype(np.float64)
y_axis = data_dict['pyscan_result']['y_axis'].astype(np.float64)
projx = images.sum(axis=-2)
median_index = misc.get_median(projx, output='index')
median_indices[label] = median_index
print(label, median_index)
sp_ctr = np.inf
ny, nx = 3, 3
subplot = ms.subplot_factory(ny, nx)
full_slice_dict[label] = np.zeros([len(images), 5, n_slices])
if ctr == 1:
ref_y = full_slice_dict['Lasing Off'][:, 4, 0].mean()
for n_image, image in enumerate(images):
image_obj = iap.Image(image,
x_axis,
y_axis,
subtract_median=True,
x_offset=screen_center)
image_cut = image_obj.cut(wake_x.min(), wake_x.max())
image_reshaped = image_cut.reshape_x(len_profile)
image_t = image_reshaped.x_to_t(wake_x, wake_t, debug=False)
tracker = tracking.Tracker(**tracker_kwargs)
print('R12', tracker.calcR12()[1])
print('disp', tracker.calcDisp()[1])
images = data_dict['pyscan_result']['image'].astype(np.float64)
x_axis = data_dict['pyscan_result']['x_axis_m'].astype(np.float64)
y_axis = data_dict['pyscan_result']['y_axis_m'].astype(np.float64)
projx = images.sum(axis=-2)
median_index = misc.get_median(projx, output='index')
median_indices[label] = median_index
print(label, median_index)
sp_ctr = np.inf
ny, nx = 3, 3
subplot = ms.subplot_factory(ny, nx)
full_slice_dict[label] = np.zeros([len(images), 5, n_slices])
if ctr == 1:
ref_y = full_slice_dict['Lasing Off'][:, 4, 0].mean()
for n_image, image in enumerate(images):
image_obj = iap.Image(image,
x_axis,
y_axis,
subtract_median=True,
x_offset=screen_center)
image_cut = image_obj.cut(wake_x.min(), wake_x.max())
image_reshaped = image_cut.reshape_x(len_profile)
image_t = image_reshaped.x_to_t(wake_x, wake_t, debug=False)
streaker_gap = streaker_data.get_prev_datapoint(streaker + ':GAP',
timestamp) * 1e-3
streaker_center = streaker_data.get_prev_datapoint(streaker + ':CENTER',
timestamp) * 1e-3
print('Streaker properties [mm]', timestamp, streaker_gap * 1e3,
streaker_center * 1e3)
lasing_on = loadH5Recursive(lasing_on_file)
lasing_off = loadH5Recursive(lasing_off_file)
image_on = lasing_on['camera1']['image'].astype(float)
image_off = lasing_off['camera1']['image'].astype(float)
fig = ms.figure('See effect of lasing')
fig.subplots_adjust(hspace=0.3)
subplot = ms.subplot_factory(2, 2, grid=True)
sp_ctr = 1
x_axis = lasing_on['camera1']['x_axis'].astype(float) * 1e-6 - x0
y_axis = lasing_on['camera1']['y_axis'].astype(float) * 1e-6
if x_axis[1] < x_axis[0]:
x_axis = x_axis[::-1]
image_on = image_on[:, ::-1]
image_off = image_off[:, ::-1]
if y_axis[1] < y_axis[0]:
y_axis = y_axis[::-1]
image_on = image_on[::-1, :]
image_off = image_off[::-1, :]
empty_snapshot = dirname2 + '20210519_121718_SARBD02-DSCR050_camera_snapshot.h5'
image_dict = h5_storage.loadH5Recursive(empty_snapshot)
x_axis0 = image_dict['camera1']['x_axis']
y_axis0 = image_dict['camera1']['y_axis']
lasing_snapshot = dirname1+'20210518_233946_SARBD02-DSCR050_camera_snapshot.h5'
image_dict = h5_storage.loadH5Recursive(lasing_snapshot)
x_axis = image_dict['camera1']['x_axis']
y_axis = image_dict['camera1']['y_axis']
image = image_dict['camera1']['image']
ms.figure('Saved')
subplot = ms.subplot_factory(1,2, False)
sp_ctr = 1
with open('./bytes.pkl', 'rb') as f:
bytes = pickle.load(f)
arr0 = np.frombuffer(bytes, dtype=np.uint16)
bg = arr0.reshape([2160, 2560])
mask_bg = bg > 150
min_index_x = np.argwhere(x_axis0 == x_axis[0]).squeeze()
max_index_x = np.argwhere(x_axis0 == x_axis[-1]).squeeze()
min_index_y = np.argwhere(y_axis0 == y_axis[0]).squeeze()
max_index_y = np.argwhere(y_axis0 == y_axis[-1]).squeeze()
if y_axis[1] < y_axis[0]:
y_axis = y_axis[::-1]
image_on = image_on[::-1,:]
image_off = image_off[::-1,:]
image_on -= np.median(image_on)
image_off -= np.median(image_off)
np.clip(image_on, 0, None, out=image_on)
np.clip(image_off, 0, None, out=image_off)
image_on = iap.Image(image_on, x_axis, y_axis, x_offset=x0)
image_off = iap.Image(image_off, x_axis, y_axis, x_offset=x0)
ms.figure('')
subplot = ms.subplot_factory(2,2, grid=False)
sp_ctr = 1
all_slice_dict = OrderedDict()
for image, label in [(image_on, 'Lasing On'), (image_off, 'Lasing Off')][::-1]:
image_cut = image.cut(xx.min(), 0.5e-3)
sp = subplot(sp_ctr, xlabel='x [mm]', ylabel='y [mm]', title=label)
sp_ctr += 1
image_cut.plot_img_and_proj(sp)
n_slices = 20
image_slice = image_cut.slice_x(n_slices)
slice_dict = image_slice.fit_slice()
slice_mean = slice_dict['slice_mean']
disp_factor = disp_dict['SARBD02.DSCR050'] / energy_eV
data_dict = loadmat(data_file)
x_axis = data_dict['x_axis'].squeeze() * 1e-6
y_axis = data_dict['y_axis'].squeeze() * 1e-6
n_images, n_gaps = data_dict['Image'].shape
if transpose:
x_axis, y_axis = y_axis, x_axis
gap_list = result_dict['gap_list']
nx, ny = 3, 3
subplot = ms.subplot_factory(ny, nx)
subplot0 = ms.subplot_factory(2, 2)
fig0 = ms.figure('Energy change')
sp_ctr = 1
sp_current = subplot0(sp_ctr, title='Current')
sp_ctr += 1
sp_current2 = subplot0(sp_ctr, title='Current')
sp_ctr += 1
sp_ene = subplot0(sp_ctr, title='Energy change')
sp_ctr += 1
sp_ene2 = subplot0(sp_ctr,
print('Delta gap %i um' % (delta_gap * 1e6))
tracker_kwargs = config.get_default_tracker_settings()
recon_kwargs = config.get_default_gauss_recon_settings()
tracker = tracking.Tracker(**tracker_kwargs)
tracker.set_simulator(sc.meta_data)
recon_kwargs['gaps'] = [10e-3, 10e-3 + delta_gap]
recon_kwargs['beam_offsets'] = [0., -(sc.offsets[index] - streaker_offset)]
recon_kwargs['n_streaker'] = 1
recon_kwargs['meas_screen'] = meas_screen
hspace, wspace = 0.5, 0.4
fig = ms.figure('Current profile reconstruction', figsize=(6, 12))
ms.plt.subplots_adjust(hspace=hspace, wspace=wspace)
subplot = ms.subplot_factory(4, 2, grid=False)
sp_ctr = 1
where0 = np.argwhere(sc.offsets == 0).squeeze()
xlim = -3e-3, 1e-3
ylim = 1e-3, 5e-3
for img_index, title in [(index, '(b) Streaked'),
(where0, '(a) Unstreaked')][::-1]:
raw_image = sc.images[img_index][0]
x_axis = sc.plot_list_x[img_index]
y_axis = sc.y_axis_list[img_index]
img = iap.Image(raw_image, x_axis, y_axis)
sp_img = subplot(sp_ctr,
title=title,
def obtain_lasing(image_off, image_on, n_slices, wake_x, wake_t, len_profile, dispersion, energy_eV, charge, pulse_energy, debug=False):
all_slice_dict = OrderedDict()
all_images = OrderedDict()
for ctr, (image_obj, label) in enumerate([(image_off, 'Lasing_off'), (image_on, 'Lasing_on')]):
image_cut = image_obj.cut(wake_x.min(), wake_x.max())
image_reshaped = image_cut.reshape_x(len_profile)
image_t = image_reshaped.x_to_t(wake_x, wake_t)
if ctr == 0:
ref_y = None
image_tE, ref_y = image_t.y_to_eV(dispersion, energy_eV, ref_y)
image_t_reduced = image_tE.slice_x(n_slices)
slice_dict = image_t_reduced.fit_slice(charge=charge, smoothen_first=True, smoothen=1e6)
all_slice_dict[label] = slice_dict
all_images[label] = {
'image_xy': image_obj,
'image_tE': image_tE,
'image_cut': image_cut,
'image_t': image_t,
'image_t_reduced': image_t_reduced,
}
slice_time = all_slice_dict['Lasing_off']['slice_x']
mean_current = (all_slice_dict['Lasing_off']['slice_current']+all_slice_dict['Lasing_on']['slice_current'])/2.
delta_E = all_slice_dict['Lasing_off']['slice_mean'] - all_slice_dict['Lasing_on']['slice_mean']
delta_std_sq = all_slice_dict['Lasing_on']['slice_sigma'] - all_slice_dict['Lasing_off']['slice_sigma']
np.clip(delta_std_sq, 0, None, out=delta_std_sq)
power_from_Eloss = power_Eloss(mean_current, delta_E)
E_total = np.trapz(power_from_Eloss, slice_time)
power_from_Espread = power_Espread(slice_time, mean_current, delta_std_sq, pulse_energy)
if debug:
ms.figure('Lasing')
subplot = ms.subplot_factory(2,2)
sp_ctr = 1
sp_power = subplot(sp_ctr, title='Power')
sp_ctr += 1
sp_current = subplot(sp_ctr, title='Current')
sp_ctr += 1
sp_current.plot(slice_time, all_slice_dict['Lasing_off']['slice_current'], label='Off')
sp_current.plot(slice_time, all_slice_dict['Lasing_on']['slice_current'], label='On')
sp_power.plot(slice_time, power_from_Eloss)
sp_power.plot(slice_time, power_from_Espread)
plt.show()
output = {
'all_slice_dict': all_slice_dict,
'power_Eloss': power_from_Eloss,
'energy_Eloss': E_total,
'power_Espread': power_from_Espread,
'all_images': all_images,
'current': mean_current,
'slice_time': slice_time,
}
return output
(screen_sim, 'Initial')]:
color = screen.plot_standard(sp, label=label)[0].get_color()
xx, yy = screen.gaussfit.xx, screen.gaussfit.reconstruction
sp.plot(xx * 1e3,
yy / np.trapz(yy, xx),
color=color,
ls='--',
label='%i' % (screen.gaussfit.sigma * 1e6))
sp.legend()
else:
new_emittance = n_emittances[0]
fig_paper = ms.figure('For paper %s ($\epsilon$=%i nm)' %
(main_label, new_emittance * 1e9))
fig_paper.subplots_adjust(hspace=0.3)
subplot = ms.subplot_factory(2, 3)
sp_ctr_p = 1
sp_tdc_meas = subplot(sp_ctr_p,
title='Current profiles',
xlabel='time [fs]',
ylabel='Current (A)')
sp_ctr_p += 1
sp_backtrack_tdc_screen = subplot(sp_ctr_p,
title='TDC forward',
xlabel='x [mm]',
ylabel='Screen projection')
sp_ctr_p += 1
sp_back_forward = subplot(sp_ctr_p,
calib_dict = h5_storage.loadH5Recursive(calib_file)
raw_data = calib_dict['raw_data']
result0 = analysis.analyze_streaker_calibration(raw_data, True)
print('Module streaker center', result0['meta_data']['streaker_offset']*1e6, 'um')
mask = np.logical_and(raw_data['streaker_offsets'] < 5.04e-3, raw_data['streaker_offsets'] > -0.0044)
#mask = np.ones_like(raw_data['streaker_offsets'], dtype=bool)
raw_data['streaker_offsets'] = raw_data['streaker_offsets'][mask]
raw_data['pyscan_result']['image'] = raw_data['pyscan_result']['image'][mask]
print(raw_data['streaker_offsets'])
ms.figure('Screen projections')
subplot = ms.subplot_factory(1,1)
sp_proj = subplot(1, title='Screen projections', xlabel='x (mm)', ylabel='Intensity (arb. units)')
screen_data = h5_storage.loadH5Recursive(screen_file)
image = screen_data['pyscan_result']['image'][0]
screen = misc.image_to_screen(image, screen_data['pyscan_result']['x_axis'], True, 0)
screen.plot_standard(sp_proj, color='black', label='Screen calib')
sp_proj.axvline(screen_analysis['x0']*1e3, color='black')
#for res in result0, result1:
# print(int(res['meta_data']['streaker_offset']*1e6))
def streaker_calibration_fit_func(offsets, streaker_offset, strength, order=3, const=0, semigap=0):
wall0, wall1 = -semigap, semigap
import h5_storage
import myplotstyle as ms
plt.close('all')
snapshot_file = '/sf/data/measurements/2021/05/19/20210519_121718_SARBD02-DSCR050_camera_snapshot.h5'
image_dict = h5_storage.loadH5Recursive(snapshot_file)
x_axis = image_dict['camera1']['x_axis']
y_axis = image_dict['camera1']['y_axis']
image = image_dict['camera1']['image']
with open('./bytes.pkl', 'rb') as f:
bytes = pickle.load(f)
arr0 = np.frombuffer(bytes, dtype=np.uint16)
arr = arr0.reshape([2160, 2560])
ms.figure('Background')
sp = ms.subplot_factory(1, 1, False)(1)
sp.imshow(arr, aspect='auto')
ms.figure('Saved')
sp = ms.subplot_factory(1, 1, False)(1)
sp.imshow(image, aspect='auto')
plt.show()
profile_meas = tracking.profile_from_blmeas(p_dict['blmeas'], tt_halfrange,
charge, energy_eV)
profile_dhf = tracking.dhf_profile(profile_meas)
profile_dhf.cutoff(screen_cutoff)
profile_dhf.crop()
#profile_meas.flipx()
for n_proj in range(3):
screen0 = get_screen_from_proj(projections[n_proj], x_axis, invert_x0)
screen0._xx = screen0._xx - mean0
screen0.cutoff(3e-2)
screen0.crop()
ms.figure('Fit quad effect %s' % main_label)
subplot = ms.subplot_factory(1, 2)
sp_ctr = 1
sp_profile = subplot(sp_ctr,
title='Current profile',
xlabel='time [fs]',
ylabel='I (arb. units)')
sp_ctr += 1
profile_meas.plot_standard(sp_profile,
label='$\sigma$ %i fs' %
(profile_meas.gaussfit.sigma * 1e15))
sp_profile.legend()
sp_forward = subplot(sp_ctr,
title='Screen distribution',
#elif hostname == 'pubuntu':
# dirname = '/home/work/data_2020-10-03/'
#
#
#if all_files:
# files = sorted(glob.glob(dirname+'Passive_alignment*.mat')[2:])
#else:
# files = [dirname + f for f in [
# 'Passive_alignment_20201003T221023.mat',
# 'Passive_alignment_20201003T214629.mat',
# 'Passive_alignment_20201003T222227.mat',
# ]]
fig0 = ms.figure(title='a', figsize=(12, 12))
#fig0.subplots_adjust(wspace=0.4)
subplot = ms.subplot_factory(2, 2)
bpm_sp_dict = {}
bpm_sp_dict2 = {}
sp_ctr = 1
for bpm in bpms_plot:
bpm_sp_dict[bpm] = subplot(
sp_ctr,
title='',
xlabel='Center (mm)',
ylabel='Beam position (mm)',
grid=False) # , title_fs=fs_title, label_fs=fs_label)
sp_ctr += 1
fig0 = ms.figure(title='b', figsize=(2, 2))
#fig0.subplots_adjust(wspace=0.4)
long_proj = np.interp(long_x_axis, x_axis, proj)
t_interp0 = np.interp(long_x_axis, xx, tt)
intensity, bins = np.histogram(t_interp0, bins=100, weights=long_proj)
new_axis = (bins[1:] + bins[:-1]) / 2.
intensity[0] = 0
intensity[-1] = 0
manual_profile = iap.BeamProfile(new_axis, intensity, final_profile.energy_eV,
200e-12)
manual_profile_rev = iap.BeamProfile(new_axis, intensity[::-1],
final_profile.energy_eV, 200e-12)
ms.figure('Debug')
subplot = ms.subplot_factory(2, 2)
sp_ctr = 1
sp_screen = subplot(sp_ctr, title='Screen', xlabel='x [mm]')
sp_ctr += 1
sp_screen.plot(meas_screen.x * 1e3, meas_screen.intensity)
sp_wake = subplot(sp_ctr, title='Wake', xlabel='t [fs]', ylabel='x [mm]')
sp_ctr += 1
sp_wake.plot(tt * 1e15, xx * 1e3)
sp_profile = subplot(sp_ctr, title='Profile')
sp_ctr += 1
def track_backward(self,
screen,
wake_effect,
n_streaker,
plot_details=False):
screen = copy.deepcopy(screen)
wake_time = wake_effect['t']
wake_x = wake_effect['x']
q_wake_x = wake_effect['quad']
charge = wake_effect['charge']
diff_x = np.diff(wake_x)
if np.all(diff_x <= 0):
wake_x = wake_x[::-1]
q_wake_x = q_wake_x[::-1]
wake_time = wake_time[::-1]
elif np.all(diff_x >= 0):
pass
else:
raise ValueError('Wake x is not monotonous')
if abs(wake_x.max()) > abs(wake_x.min()):
mask_negative = screen.x < 0
else:
mask_negative = screen.x > 0
if self.compensate_negative_screen and np.any(mask_negative):
x_positive = -screen.x[mask_negative][::-1]
y_positive = screen.intensity[mask_negative][::-1]
if np.all(np.diff(x_positive) < 0):
x_positive = x_positive[::-1]
y_positive = y_positive[::-1]
positive_interp = np.interp(screen.x,
x_positive,
y_positive,
left=0,
right=0)
screen_intensity = screen.intensity + positive_interp
screen_intensity[mask_negative] = 0
screen._yy = screen_intensity
if self.quad_wake_back:
if self.override_quad_beamsize:
bs_at_streaker = self.quad_x_beamsize[n_streaker]
else:
if self.bs_at_streaker is None:
self.set_bs_at_streaker()
bs_at_streaker = self.bs_at_streaker[n_streaker]
screen.reshape(self.n_particles)
rand0 = np.random.randn(len(screen.x))
np.clip(rand0, -3, 3, out=rand0)
randx = rand0 * bs_at_streaker
t_interp0 = np.zeros_like(screen.x)
for n_x, (x, rx) in enumerate(zip(screen.x, randx)):
t_interp0[n_x] = np.interp(x, wake_x + rx * q_wake_x,
wake_time)
charge_interp, hist_edges = np.histogram(t_interp0,
bins=self.n_particles //
100,
weights=screen.intensity,
density=True)
charge_interp[0] = 0
charge_interp[-1] = 0
t_interp = (hist_edges[1:] + hist_edges[:-1]) / 2.
else:
screen.reshape(self.n_particles)
t_interp0 = np.interp(screen.x, wake_x, wake_time)
charge_interp, hist_edges = np.histogram(t_interp0,
bins=self.n_particles //
100,
weights=screen.intensity,
density=True)
charge_interp[0] = 0
charge_interp[-1] = 0
t_interp = (hist_edges[1:] + hist_edges[:-1]) / 2.
try:
if np.any(np.diff(t_interp) < 0):
t_interp = t_interp[::-1]
charge_interp = charge_interp[::-1]
assert np.all(np.diff(t_interp) >= 0)
bp = BeamProfile(t_interp, charge_interp, self.energy_eV, charge)
except (ValueError, AssertionError) as e:
print(e)
ms.figure('')
self.set_bs_at_streaker()
subplot = ms.subplot_factory(2, 2)
sp = subplot(1, title='Wake', xlabel='t', ylabel='$\Delta$ x')
sp.plot(wake_time, wake_x, label='Dipole')
sp.plot(wake_time,
q_wake_x * self.bs_at_streaker[n_streaker],
label='Quad')
sp.legend()
sp = subplot(2, title='Screen', xlabel='x')
sp.plot(screen.x, screen.intensity)
sp = subplot(3,
title='Current profile',
xlabel='time',
ylabel='Current')
sp.plot(t_interp, charge_interp)
plt.show()
raise
bp.reshape(self.len_screen)
bp.cutoff2(self.profile_cutoff)
bp.crop()
if np.any(np.isnan(bp.time)) or np.any(np.isnan(bp.current)):
raise ValueError('NaNs in beam profile')
if self.bp_smoothen:
#bp0 = copy.deepcopy(bp)
bp.smoothen(self.bp_smoothen)
#bp1 = copy.deepcopy(bp)
bp.reshape(self.len_screen)
#import pdb; pdb.set_trace()
if plot_details:
ms.figure('track_backward')
subplot = ms.subplot_factory(2, 2)
sp_wake = subplot(1,
title='Wake effect',
xlabel='t [fs]',
ylabel='$\Delta$ x [mm]')
sp_wake.plot(wake_time * 1e15, wake_x * 1e3)
sp_screen = subplot(2,
title='Screen dist',
xlabel='x [mm]',
ylabel='Intensity (arb. units)')
screen.plot_standard(sp_screen)
sp_profile = subplot(3,
title='Interpolated profile',
xlabel='t [fs]',
ylabel='Current [kA]')
bp.plot_standard(sp_profile)
return bp
charge_xx -= charge_xx.min()
current_profile = bl_meas['current1']
charge_profile = current_profile * charge / np.sum(current_profile)
energy_eV = bl_meas['energy_eV']
mask_charge = np.logical_and(charge_xx > lims[0] * 1e-6,
charge_xx < lims[1] * 1e-6)
assert np.any(mask_charge)
charge_xx = charge_xx[mask_charge]
charge_xx -= charge_xx.min()
charge_profile = charge_profile[mask_charge]
fig = ms.figure('Surface variable scan')
sp_ctr = 1
subplot = ms.subplot_factory(3, 3)
sp_charge = subplot(sp_ctr, title='Current profile', xlabel='s [$\mu$m]')
sp_ctr += 1
sp_charge.plot(charge_xx * 1e6, charge_profile)
sp_surface = subplot(sp_ctr, title='Surface WF', xlabel='s [$\mu$m]')
sp_ctr += 1
h_arr = [50e-9, 100e-9, 200e-9]
kappa_factor_arr = [0.5, 1, 2]
for h, kappa_factor in itertools.product(h_arr, kappa_factor_arr):
kappa = uwf.aramis_kappa * kappa_factor
result, err = uwf.surface_round_tube(charge_xx,
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
import numpy as np
import wf_model
import myplotstyle as ms
plt.close('all')
gap = 10e-3
distance_arr = np.linspace(250e-6, 500e-6, int(1e4))
ms.figure('Scaling')
sp = ms.subplot_factory(1, 1)(1, xlabel='Distance [um]')
for s in [1e-6, 10e-6, 25e-6, 50e-6, 75e-6]:
wake = wf_model.wxd(s, gap / 2., gap / 2. - distance_arr)
color = sp.plot(distance_arr * 1e6, wake,
label='s=%i um' % (s * 1e6))[0].get_color()
def fit_func(x, scale, order):
return scale * 1 / x**order
def fit_func3(x, scale):
return scale * 1 / x**3
fit, _ = curve_fit(fit_func, distance_arr, wake, p0=[1e12, 1])
fit2, _ = curve_fit(fit_func3, distance_arr, wake, p0=[1e12])
reconstruction = fit_func(distance_arr, *fit)
reconstruction2 = fit_func3(distance_arr, *fit2)
sp.plot(distance_arr * 1e6,