def analyze_screen_calibration(filename_or_dict, do_plot=True, plot_handles=None):
if type(filename_or_dict) is dict:
data_dict = filename_or_dict
elif type(filename_or_dict) is str:
data_dict = h5_storage.loadH5Recursive(filename_or_dict)
else:
raise ValueError(type(filename_or_dict))
screen_data = data_dict['pyscan_result']
if 'x_axis_m' in screen_data:
x_axis = screen_data['x_axis_m']
else:
#print(screen_data['x_axis'].shape)
x_axis = screen_data['x_axis'].squeeze()*1e-6
if len(x_axis.shape) == 2:
x_axis = x_axis[0]
assert len(x_axis.squeeze().shape) == 1
if 'projx' in screen_data:
projx = screen_data['projx']
else:
images = screen_data['image'].astype(float).squeeze()
projx = images.sum(axis=-2)
all_mean = []
all_std = []
for proj in projx:
gf = gaussfit.GaussFit(x_axis, proj)
all_mean.append(gf.mean)
all_std.append(gf.sigma)
index_median = np.argsort(all_mean)[len(all_mean)//2]
projx_median = projx[index_median]
beamsize = np.mean(all_std)
projx_median -= projx_median.min()
x0 = gaussfit.GaussFit(x_axis, projx_median).mean
output = {
'raw_data': data_dict,
'x0': x0,
'beamsize': beamsize,
}
if not do_plot:
return output
if plot_handles is None:
fig, (sp_proj,) = screen_calibration_figure()
else:
(sp_proj, ) = plot_handles
for proj in projx:
sp_proj.plot(x_axis*1e3, proj-proj.min())
sp_proj.plot(x_axis*1e3, projx_median-projx_median.min(), lw=3)
sp_proj.axvline(x0*1e3)
return output
def get_median(projx, method='gf_mean', output='proj'):
"""
From list of projections, return the median one
Methods: gf_mean, gf_sigma, mean, rms
Output: proj, index
"""
x_axis = np.arange(len(projx[0]))
all_mean = []
for proj in projx:
proj = proj - proj.min()
if method == 'gf_mean':
gf = gaussfit.GaussFit(x_axis, proj)
all_mean.append(gf.mean)
elif method == 'gf_sigma':
gf = gaussfit.GaussFit(x_axis, proj)
all_mean.append(gf.sigma)
elif method == 'mean':
mean = np.sum(x_axis * proj) / np.sum(proj)
all_mean.append(mean)
elif method == 'std':
mean = np.sum(x_axis * proj) / np.sum(proj)
rms = np.sqrt(np.sum((x_axis - mean)**2 * proj) / np.sum(proj))
all_mean.append(rms)
else:
raise ValueError(method)
index_median = np.argsort(all_mean)[len(all_mean) // 2]
projx_median = projx[index_median]
#import matplotlib.pyplot as plt
#plt.figure()
#for proj in projx:
# plt.plot(proj)
#plt.plot(projx_median, color='black', lw=3)
#plt.show()
#import pdb; pdb.set_trace()
if output == 'proj':
return projx_median
elif output == 'index':
return index_median
#sp_ctr_paper = 1
if type(p_dict['proj0']) is str:
mean0 = -0.08e-3
else:
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)
print('%s: Mean0: %.3e mm' % (main_label, mean0 * 1e3))
if fit_emittance and main_label != 'Short':
timestamp0 = misc.get_timestamp(os.path.basename(p_dict['filename0']))
tracker0 = tracking.Tracker(magnet_file,
timestamp0,
struct_lengths,
n_particles,
n_emittances,
screen_bins,
screen_cutoff,
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)
if False:
timestamp0 = misc.get_timestamp(os.path.basename(p_dict['filename0']))
tracker0 = tracking.Tracker(magnet_file, timestamp0, struct_lengths, n_particles, n_emittances, screen_bins, screen_cutoff, smoothen, profile_cutoff, len_profile, quad_wake=quad_wake)
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']
#screen_sim.smoothen(25e-6)
all_emittances = []
all_beamsizes = []
for proj in projx0:
sp_proj = subplot(sp_ctr, title='Projections X')
sp_ctr += 1
sp_projY = subplot(sp_ctr, title='Projections Y')
sp_ctr += 1
histx, xedges = np.histogram(beam_forward[0], bins=200)
histy, yedges = np.histogram(beam_forward[2], bins=200)
for arr, arry, axis, axisy, label in [
(image0.sum(axis=0), image0.sum(axis=1), x_axis, y_axis, 'Measured'),
(histx, histy, xedges[1:], yedges[1:], 'Simulated')
]:
for sp, arr_, ax in [(sp_proj, arr, axis), (sp_projY, arry, axisy)]:
gf = gaussfit.GaussFit(ax, arr_)
sigma = gf.sigma
x_max = ax[np.argmax(arr_)]
sp.plot(ax - x_max,
arr_ / arr_.max(),
label=label + ' %i' % (sigma * 1e6))
sp_proj.legend()
sp_projY.legend()
sp_ctr = np.inf
n_proj = np.inf
ny, nx = 3, 3
subplot = ms.subplot_factory(ny, nx)
images0 = dict_['Image'][-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()
ms.figure('Resolution', figsize=(10, 8))
ms.plt.subplots_adjust(hspace=0.4, wspace=0.8)
subplot = ms.subplot_factory(2, 3, grid=False)
sp_ctr = 1
image_file = data_dir1 + '2021_05_18-21_02_13_Lasing_False_SARBD02-DSCR050.h5'
image_dict = h5_storage.loadH5Recursive(image_file)
meta_data1 = image_dict['meta_data_begin']
screen_calib_file = data_dir1 + '2021_05_18-16_39_27_Screen_Calibration_SARBD02-DSCR050.h5'
screen_calib_dict = h5_storage.loadH5Recursive(screen_calib_file)
screen_calib_raw_image = screen_calib_dict['pyscan_result']['image'][0].astype(
float)
x_axis_calib = screen_calib_dict['pyscan_result']['x_axis_m']
screen_x0 = gaussfit.GaussFit(x_axis_calib,
screen_calib_raw_image.sum(axis=0)).mean
x_axis_calib -= screen_x0
y_axis_calib = screen_calib_dict['pyscan_result']['y_axis_m']
screen_calib_raw_image -= np.median(screen_calib_raw_image)
screen_calib_image = iap.Image(screen_calib_raw_image, x_axis_calib,
y_axis_calib)
images = image_dict['pyscan_result']['image'].astype(float)
x_axis = image_dict['pyscan_result']['x_axis_m'] - screen_x0
y_axis = image_dict['pyscan_result']['y_axis_m']
projx = images.sum(axis=-2)
median_index = misc.get_median(projx, method='mean', output='index')
raw_image1 = images[median_index]
raw_image1 -= np.median(raw_image1)
image1 = iap.Image(raw_image1, x_axis, y_axis)
x_axis = dict_['pyscan_result']['x_axis'].astype(float)
y_axis = dict_['pyscan_result']['y_axis'].astype(float)
def fit_func(xx, scale, mean, sig, const):
return scale*np.exp(-(xx-mean)**2/(2*sig**2))+const
for n_img in range(4):
image0 = dict_['pyscan_result']['image'][n_img].astype(float)
if n_img == 0:
projx = image0.sum(axis=-2)
projx -= projx.min()
sp_proj_x.plot(x_axis, projx)
gf = gaussfit.GaussFit(x_axis, projx, fit_const=True)
p0 = list(gf.p0)
p0[2] *= 5
p_opt, p_cov = curve_fit(fit_func, x_axis, projx, p0=p0)
yy = fit_func(x_axis, *p_opt)
sp_proj_x.plot(x_axis, yy)
sp_proj_x.plot(gf.xx, gf.reconstruction)
sp_proj_y.plot(y_axis, image0.sum(axis=-1))
sp_img = subplot(sp_ctr)
sp_ctr += 1
sp_img.imshow(image0, aspect='auto')
slice_sigma = arr[n_image, 2, mask]
slice_eloss = lasing_input['Lasing Off']['mean'] - slice_mean
slice_Espread_sqr_increase = slice_sigma**2 - lasing_input['Lasing Off']['sigma']**2
p_eloss = lasing.power_Eloss(curr, slice_eloss)
E_eloss = np.trapz(p_eloss, slice_time2)
p_espread = lasing.power_Espread(slice_time2, curr, slice_Espread_sqr_increase, E_norm)
sp_mean.plot(slice_time2*1e15, p_eloss/1e9)
sp_sigma.plot(slice_time2*1e15, p_espread/1e9)
for ctr, (power, error, label, color, sp2) in enumerate([
(power_mean, power_mean_err, '$P_m$', 'red', sp_mean),
(power_sigma, power_sigma_err, '$P_\sigma$', 'blue', sp_sigma),
]):
gf = gaussfit.GaussFit(slice_time2, power, fit_const=False)
label2 = label+ ' $\sigma$=%.1f fs' % (gf.sigma*1e15)
for sp_, color2, lw in [(sp_lasing, color, None), (sp2, 'black', 3)]:
sp_.errorbar(slice_time2*1e15, power/1e9, yerr=error/1e9, label=label2, color=color2, lw=lw)
sp_lasing.legend()
ms.saveall('/tmp/lasing_analysis', ending='.pdf', hspace=.35, wspace=.35)
plt.show()
