archiver_dir = '/afs/psi.ch/intranet/SF/Beamdynamics/Philipp/data/' magnet_file = archiver_dir + 'archiver_api_data/2021-03-16.h5' timestamp = elegant_matrix.get_timestamp(2021, 3, 16, 20, 41, 39) beam_offsets = [-0.0, -0.004724] median_screen = misc.image_to_screen(image_on.image, x_axis, False, x0) median_screen.cutoff2(3e-2) median_screen.crop() median_screen.reshape(len_profile) tracker = tracking.Tracker(magnet_file=magnet_file, timestamp=timestamp, n_particles=n_particles, n_emittances=n_emittances, screen_bins=screen_bins, screen_cutoff=screen_cutoff, smoothen=smoothen, profile_cutoff=profile_cutoff, len_screen=len_profile, bp_smoothen=bp_smoothen, quad_wake=False) energy_eV = tracker.energy_eV blmeas = data_dir+'113876237_bunch_length_meas.h5' profile_meas = iap.profile_from_blmeas(blmeas, 200e-15, charge, energy_eV) profile_meas.cutoff(5e-2) profile_meas.reshape(len_profile) profile_meas.crop() tt_range = (profile_meas.time.max() - profile_meas.time.min()) profile_gauss = iap.get_gaussian_profile(38e-15, tt_range, len_profile, charge, energy_eV) bp_back = tracker.track_backward2(median_screen, profile_gauss, gaps, beam_offsets, 1) bp_back2 = tracker.track_backward2(median_screen, bp_back, gaps, beam_offsets, 1) #bp_back3 = tracker.track_backward2(median_screen, bp_back2, gaps, beam_offsets, 1)
ny, nx = 2, 2
subplot = ms.subplot_factory(ny, nx)
sp_ctr = np.inf
chi_sq_centroid = np.zeros_like(gap_arr)
chi_sq_rms = np.zeros_like(gap_arr)
opt_values = np.zeros_like(gap_arr)
def chi_squared(meas, expected, weight=1):
return np.sum(((meas - expected) / weight)**2)
blmeas_profile = iap.profile_from_blmeas(blmeas_file, tt_halfrange,
gauss_kwargs['charge'],
tracker.energy_eV, True)
blmeas_profile.cutoff2(tracker.profile_cutoff)
blmeas_profile.crop()
blmeas_profile.reshape(tracker.len_screen)
streaker_offset = 376e-6
for gap_ctr, gap in enumerate(gap_arr):
sc.gap0 = gap
if sp_ctr > ny * nx:
fig = ms.figure('Screen distributions', figsize=(20, 16))
sp_ctr = 1
fignum = fig.number
struct_lengths,
n_particles,
n_emittances,
screen_bins,
screen_cutoff,
smoothen,
profile_cutoff,
len_profile,
quad_wake=quad_wake,
bp_smoothen=bp_smoothen)
blmeas = p_dict['blmeas']
flip_measured = p_dict['flipx']
profile_meas = iap.profile_from_blmeas(blmeas,
tt_halfrange,
charge,
tracker.energy_eV,
subtract_min=True)
profile_meas.reshape(len_profile)
profile_meas2 = iap.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()
struct_lengths,
n_particles,
n_emittances,
screen_bins,
screen_cutoff,
smoothen,
profile_cutoff,
len_profile,
quad_wake=quad_wake,
bp_smoothen=bp_smoothen)
tracker.wake2d = wake2d
tracker.split_streaker = split_streaker
profile_from_blmeas = profile_from_blmeas(blmeas,
tt_halfrange,
charge,
tracker.energy_eV,
subtract_min=True)
profile_from_blmeas.cutoff(3e-2)
profile_from_blmeas.crop()
profile_from_blmeas.reshape(len_profile)
emit_fit = tracker.fit_emittance(beamsize, smoothen, tt_halfrange)
print(emit_fit)
tracker.n_emittances = [emit_fit, emit_fit]
ms.figure('Forward propagation')
subplot = ms.subplot_factory(2, 2)
sp_ctr = 1
sp_pos = subplot(sp_ctr,
hostname = gethostname()
if hostname == 'desktop':
data_dir2 = '/storage/data_2021-05-19/'
elif hostname == 'pc11292.psi.ch':
data_dir2 = '/sf/data/measurements/2021/05/19/'
elif hostname == 'pubuntu':
data_dir2 = '/mnt/data/data_2021-05-19/'
data_dir1 = data_dir2.replace('19', '18')
lasing_off_file = data_dir1 + '2021_05_18-23_42_10_Lasing_True_SARBD02-DSCR050.h5'
lasing_on_file = data_dir1 + '2021_05_18-23_43_39_Lasing_False_SARBD02-DSCR050.h5'
blmeas_file = data_dir1 + '119325494_bunch_length_meas.h5'
blmeas_profile1 = iap.profile_from_blmeas(blmeas_file,
200e-15,
200e-12,
1,
zero_crossing=1)
blmeas_profile2 = iap.profile_from_blmeas(blmeas_file,
200e-15,
200e-12,
1,
zero_crossing=2)
for p in blmeas_profile1, blmeas_profile2:
p.cutoff2(5e-2)
p.crop()
p.reshape(int(1e3))
n_streaker = 1
screen_x0 = 4250e-6
def plot_rec_gauss(tracker, kwargs, gauss_dict, plot_handles=None, blmeas_file=None, do_plot=True, figsize=None, both_zero_crossings=True):
best_profile = gauss_dict['reconstructed_profile']
best_screen = gauss_dict['reconstructed_screen']
opt_func_values = gauss_dict['opt_func_values']
opt_func_screens = gauss_dict['opt_func_screens']
opt_func_profiles = gauss_dict['opt_func_profiles']
opt_func_sigmas = np.array(gauss_dict['opt_func_sigmas'])
meas_screen = gauss_dict['meas_screen']
gauss_sigma = gauss_dict['gauss_sigma']
if plot_handles is None:
fig, (sp_screen, sp_profile, sp_opt, sp_moments) = reconstruction_figure(figsize)
plt.suptitle('Optimization')
else:
sp_screen, sp_profile, sp_opt, sp_moments = plot_handles
meas_screen.plot_standard(sp_screen, color='black', lw=3)
rms_arr = np.zeros(len(opt_func_screens))
centroid_arr = rms_arr.copy()
for opt_ctr, (screen, profile, value, sigma) in enumerate(zip(opt_func_screens, opt_func_profiles, opt_func_values[:,1], opt_func_sigmas)):
screen.plot_standard(sp_screen, label='%i: %.1f fs' % (opt_ctr, sigma*1e15))
profile.plot_standard(sp_profile, label='%i: %.1f fs' % (opt_ctr, sigma*1e15), center='Mean')
rms_arr[opt_ctr] = screen.rms()
centroid_arr[opt_ctr] = screen.mean()
best_screen.plot_standard(sp_screen, color='red', lw=3, label='Final')
best_profile.plot_standard(sp_profile, color='red', lw=3, label='Final', center='Mean')
if blmeas_file is not None:
blmeas_profiles = []
for zero_crossing in (1, 2):
try:
blmeas_profile = iap.profile_from_blmeas(blmeas_file, kwargs['tt_halfrange'], kwargs['charge'], tracker.energy_eV, True, zero_crossing)
blmeas_profile.cutoff2(5e-2)
blmeas_profile.crop()
blmeas_profile.reshape(int(1e3))
blmeas_profiles.append(blmeas_profile)
except KeyError as e:
print(e)
print('No zero crossing %i in %s' % (zero_crossing, blmeas_file))
for blmeas_profile, ls, zero_crossing in zip(blmeas_profiles, ['--', 'dotted'], [1, 2]):
blmeas_profile.plot_standard(sp_profile, ls=ls, color='black', label='TDC %i %.1f fs' % (zero_crossing, blmeas_profile.rms()*1e15))
if not both_zero_crossings:
break
color = sp_moments.plot(opt_func_sigmas*1e15, np.abs(centroid_arr)*1e3, marker='.', label='Reconstructed centroid')[0].get_color()
sp_moments.axhline(np.abs(meas_screen.mean())*1e3, label='Measured centroid', color=color, ls='--')
color = sp_moments.plot(opt_func_sigmas*1e15, rms_arr*1e3, marker='.', label='Reconstructed rms')[0].get_color()
sp_moments.axhline(meas_screen.rms()*1e3, label='Measured rms', color=color, ls='--')
sp_moments.legend()
sp_screen.legend()
sp_profile.legend()
yy_opt = opt_func_values[:,1]
sp_opt.scatter(opt_func_sigmas*1e15, yy_opt)
sp_opt.set_ylim(0,1.1*yy_opt.max())
for sp_ in sp_opt, sp_moments:
sp_.axvline(gauss_sigma*1e15, color='black')
return gauss_dict
x_axis = dict0['x_axis'] * 1e-6
invert_x0 = (np.diff(x_axis)[0] < 0)
projx = dict0['projx']
projx0 = dict0['projx'][-1]
all_mean = []
for proj in projx0:
screen = get_screen_from_proj(proj, x_axis, invert_x0)
xx, yy = screen._xx, screen._yy
gf = gaussfit.GaussFit(xx, yy)
all_mean.append(gf.mean)
mean0 = np.mean(all_mean)
profile_meas = iap.profile_from_blmeas(blmeas38, tt_halfrange, charge,
energy_eV)
profile_meas.cutoff(profile_cutoff)
profile_meas.crop()
profile_meas.reshape(len_profile)
offset_arr = dict0['value'] * 1e-3 - mean_offset
if invert_offset:
offset_arr *= -1
ms.figure('Summary', figsize=(8, 6))
subplot = ms.subplot_factory(2, 2)
sp_summary_current = subplot(1,
xlabel='time [fs]',
ylabel='Current [kA]',
title='Deflector measurement')
profile_meas.plot_standard(sp_summary_current, color='black')
def reconstruct_lasing(file_or_dict_on, file_or_dict_off, screen_center, structure_center, structure_length, file_current, r12, disp, energy_eV, charge, streaker, plot_handles, pulse_energy, n_slices, len_profile):
if type(file_or_dict_on) is dict:
dict_on = file_or_dict_on
else:
dict_on = h5_storage.loadH5Recursive(file_or_dict_on)
dict_on_p = dict_on['pyscan_result']
dict_on_m = dict_on['meta_data_end']
if type(file_or_dict_off) is dict:
dict_off = file_or_dict_off
else:
dict_off = h5_storage.loadH5Recursive(file_or_dict_off)
dict_off_p = dict_off['pyscan_result']
dict_off_m = dict_off['meta_data_end']
if energy_eV == 'file':
if 'SARBD01-MBND100:ENERGY-OP' in dict_on_m:
energy_eV = dict_on_m['SARBD01-MBND100:ENERGY-OP']*1e6
elif 'SARBD01-MBND100:P-SET' in dict_on_m:
energy_eV = dict_on_m['SARBD01-MBND100:P-SET']*1e6
else:
raise ValueError('No energy saved!')
input_dict = {
'file_or_dict_on': file_or_dict_on,
'file_or_dict_off': file_or_dict_off,
'screen_center': screen_center,
'structure_center': structure_center,
'structure_length': structure_length,
'file_current': file_current,
'r12': r12,
'disp': disp,
'energy_eV': energy_eV,
'charge': charge,
'streaker': streaker,
'pulse_energy': pulse_energy,
}
gaps, beam_offsets = [], []
for dict_ in dict_off_m, dict_on_m:
gaps.append(dict_[streaker+':GAP']*1e-3)
beam_offsets.append(-(dict_[streaker+':CENTER']*1e-3-structure_center))
if abs(gaps[0] - gaps[1]) > 1e-6:
print('Gaps not the same!', gaps)
if abs(beam_offsets[0] - beam_offsets[1]) > 1e-6:
print('Beam offsets not the same!', beam_offsets)
gap = np.mean(gaps)
beam_offset = np.mean(beam_offsets)
if pulse_energy is None:
try:
pulse_energy = dict_on_m[config.gas_monitor_pvs['Aramis']]*1e-6
except KeyError:
print('No pulse energy found! Use 100 uJ')
pulse_energy = 100e-6
images0 = dict_off_p['image'].astype(np.float64)
images0_on = dict_on_p['image'].astype(np.float64)
if 'x_axis_m' in dict_off_p:
x_axis0 = dict_off_p['x_axis_m'].astype(np.float64)
y_axis0 = dict_off_p['y_axis_m'].astype(np.float64)
else:
x_axis0 = dict_off_p['x_axis'].astype(np.float64)
y_axis0 = dict_off_p['y_axis'].astype(np.float64)
projx0 = images0.sum(axis=-2)
median_index = misc.get_median(projx0, output='index')
median_image_off = iap.Image(images0[median_index], x_axis0, y_axis0, x_offset=screen_center)
projx0_on = images0_on.sum(axis=-2)
median_index = misc.get_median(projx0_on, output='index')
median_image_on = iap.Image(images0_on[median_index], x_axis0, y_axis0, x_offset=screen_center)
current_dict = h5_storage.loadH5Recursive(file_current)
if 'gaussian_reconstruction' in current_dict:
wake_profile_dict = current_dict['gaussian_reconstruction']['reconstructed_profile']
wake_profile = iap.BeamProfile.from_dict(wake_profile_dict)
else:
wake_profile = iap.profile_from_blmeas(current_dict, 200e-15, charge, energy_eV, True)
wake_profile.cutoff2(0.1)
wake_profile.crop()
wake_profile.reshape(len_profile)
wake_t, wake_x = wake_profile.get_x_t(gap, beam_offset, structure_length, r12)
lasing_dict = lasing.obtain_lasing(median_image_off, median_image_on, n_slices, wake_x, wake_t, len_profile, disp, energy_eV, charge, pulse_energy=pulse_energy, debug=False)
if plot_handles is None:
plot_handles = lasing_figures()
(fig, (sp_profile, sp_wake, sp_off, sp_on, sp_off_cut, sp_on_cut, sp_off_tE, sp_on_tE)) = plot_handles[0]
(fig, (sp_power, sp_current, sp_centroid, sp_slice_size)) = plot_handles[1]
slice_time = lasing_dict['slice_time']
all_slice_dict = lasing_dict['all_slice_dict']
power_from_Eloss = lasing_dict['power_Eloss']
power_from_Espread = lasing_dict['power_Espread']
sp_current.plot(slice_time*1e15, all_slice_dict['Lasing_off']['slice_current'], label='Off')
sp_current.plot(slice_time*1e15, all_slice_dict['Lasing_on']['slice_current'], label='On')
lasing_dict['all_images']['Lasing_off']['image_tE'].plot_img_and_proj(sp_off_tE)
lasing_dict['all_images']['Lasing_on']['image_tE'].plot_img_and_proj(sp_on_tE)
lasing_dict['all_images']['Lasing_off']['image_cut'].plot_img_and_proj(sp_off_cut)
lasing_dict['all_images']['Lasing_on']['image_cut'].plot_img_and_proj(sp_on_cut)
for key, sp_tE in [('Lasing_off',sp_off_tE), ('Lasing_on',sp_on_tE)]:
slice_sigma = all_slice_dict[key]['slice_sigma']
slice_centroid = all_slice_dict[key]['slice_mean']
sp_slice_size.plot(slice_time*1e15, slice_sigma/1e6, label=key)
sp_centroid.plot(slice_time*1e15, slice_centroid/1e6, label=key)
lims = sp_tE.get_ylim()
sp_tE.errorbar(slice_time*1e15, slice_centroid/1e6, yerr=slice_sigma/1e6, ls='None', marker='+', color='red')
sp_tE.set_ylim(*lims)
sp_power.plot(slice_time*1e15, power_from_Eloss/1e9, label='$\Delta$E')
sp_power.plot(slice_time*1e15, power_from_Espread/1e9, label='$\Delta\sigma_E$')
sp_current.legend()
sp_power.legend()
sp_slice_size.legend()
sp_centroid.legend()
median_image_off.plot_img_and_proj(sp_off)
median_image_on.plot_img_and_proj(sp_on)
sp_wake.plot(wake_t*1e15, wake_x*1e3)
wake_profile.plot_standard(sp_profile)
if type(input_dict['file_or_dict_on']) is dict:
del input_dict['file_or_dict_on']
if type(input_dict['file_or_dict_off']) is dict:
del input_dict['file_or_dict_off']
output = {
'input': input_dict,
'lasing_dict': lasing_dict,
}
#import pdb; pdb.set_trace()
return output
gap = 10e-3
beam_offset = 4.7e-3
struct_length = 1
gauss_kwargs = config.get_default_gauss_recon_settings()
tracker_kwargs = config.get_default_tracker_settings()
n_emittance = 300e-9
tracker_kwargs['n_emittances'] = [n_emittance, n_emittance]
tracker = tracking.Tracker(**tracker_kwargs)
blmeas_file = data_dir1 + '119325494_bunch_length_meas.h5'
blmeas_profile = iap.profile_from_blmeas(blmeas_file,
gauss_kwargs['tt_halfrange'],
gauss_kwargs['charge'], 0, True)
blmeas_profile.cutoff2(0.03)
blmeas_profile.crop()
blmeas_profile.reshape(1000)
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'
streaker_calib_files = [
data_dir +
'2021_05_18-23_07_20_Calibration_SARUN18-UDCP020.h5', # Affected but ok
data_dir +
'2021_05_18-23_32_12_Calibration_SARUN18-UDCP020.h5', # Affected but ok
]
calib2 = streaker_calibration.StreakerCalibration('Aramis', 1, 10e-3)
calib2.add_file(streaker_calib_files[0])
calib2.add_file(streaker_calib_files[1])
tracker_kwargs = config.get_default_tracker_settings()
tracker = tracking.Tracker(**tracker_kwargs)
tracker.set_simulator(calib2.meta_data)
beam_profile = iap.profile_from_blmeas(blmeas_file, tt_halfrange, charge,
tracker.energy_eV, True, 1)
beam_profile.reshape(tracker.len_screen)
beam_profile.cutoff2(5e-2)
beam_profile.crop()
beam_profile.reshape(tracker.len_screen)
order0_range = np.arange(2.6, 2.901, 0.05)
rec_gaps_rms = []
rec_gaps_centroid = []
rec_center_rms = []
rec_center_centroid = []
for order0 in order0_range:
calib2.order_rms = order0
calib2.order_centroid = order0
quad_wake=quad_wake)
ms.figure('Tracking')
subplot = ms.subplot_factory(1, 3)
sp_ctr = 1
sp_profile = subplot(sp_ctr,
title='Current profile',
xlabel='t (fs)',
ylabel='I (kA)')
sp_ctr += 1
for file_ctr, blmeas_file in enumerate(blmeas_files):
for zero_crossing in (1, 2):
blmeas = iap.profile_from_blmeas(blmeas_file, tt_halfrange, charge,
tracker.energy_eV, True,
zero_crossing)
blmeas.reshape(len_profile)
blmeas.cutoff2(profile_cutoff)
blmeas.crop()
blmeas.reshape(len_profile)
blmeas.plot_standard(sp_profile,
label='TDC %i zc %i %i fs' %
(file_ctr, zero_crossing, blmeas.rms() * 1e15),
ls='--',
center='Mean')
sp_screen_pos = subplot(sp_ctr,
title='Screen distributions',
xlabel='x (mm)',
ylabel='intensity (arb. units)')
all_offsets = np.array(sorted(offset_file_dict['ON'].keys()))
analysis_file = '/storage/tmp_reconstruction/2021_04_29-16_06_46_PassiveReconstruction.h5'
analysis_dict = loadH5Recursive(analysis_file)
#streaker_center = analysis_dict['input']['streaker_means'][1]
streaker_center = 309.66e-6 # + 40e-6
screen_x0 = analysis_dict['input']['screen_x0']
tracker_kwargs = analysis_dict['input']['tracker_kwargs']
tracker_kwargs['quad_wake'] = False
tracker_kwargs['override_quad_beamsize'] = False
#tracker_kwargs['n_emittances'][0] *= 2
gauss_kwargs = analysis_dict['input']['gaussian_reconstruction']
tracker = tracking.Tracker(**tracker_kwargs)
blmeas = iap.profile_from_blmeas(blmeas_file, 200e-15, 200e-12,
tracker.energy_eV, True)
blmeas.cutoff2(tracker.profile_cutoff)
ms.figure('Current profiles')
subplot = ms.subplot_factory(2, 2)
sp_ctr = 1
sp_profile0 = subplot(sp_ctr,
title='Reconstructed profiles',
xlabel='t (fs)',
ylabel='I (kA)')
sp_ctr += 1
sp_screen0 = subplot(sp_ctr,
title='Screens',
xlabel='x (mm)',