def get_screen_from_image(image, invert=False):
projX = np.sum(image, axis=0)
if invert:
screen = tracking.ScreenDistribution(-x_axis[::-1], projX[::-1])
else:
screen = tracking.ScreenDistribution(x_axis, projX)
screen.normalize()
screen.cutoff(screen_cutoff)
screen.reshape(len_profile)
return screen
def get_screen_from_image(image):
projX = np.sum(image, axis=0)
if invert_x:
screen = tracking.ScreenDistribution((-x_axis[::-1]).copy(), (projX[::-1]).copy())
else:
screen = tracking.ScreenDistribution(x_axis.copy(), projX.copy())
screen.normalize()
screen.cutoff(screen_cutoff)
screen.reshape(len_profile)
return screen
def get_screen_from_proj(projX, x_axis, invert_x):
if invert_x:
xx, yy = (-x_axis[::-1]).copy(), (projX[::-1]).copy()
else:
xx, yy = x_axis.copy(), projX.copy()
if subtract_min:
yy -= yy.min()
screen = tracking.ScreenDistribution(xx, yy)
screen.normalize()
screen.cutoff(screen_cutoff)
screen.reshape(len_profile)
return screen
def opt_func(sig_t_fs, count_nfev):
global ctr, sp_ctr, plot_ctr, sp, nfev_ctr
sig_t = sig_t_fs / 1e15
bp_wake = tracking.get_gaussian_profile(sig_t, tt_halfrange,
len_profile, charge,
tracker.energy_eV)
screen_recon = tracker.back_and_forward(meas_screen,
meas_screen0,
bp_wake,
gaps,
beam_offsets,
n_streaker,
n_bins,
back_cutoff=0.1)
screen_max_x = screen_recon.x[np.argmax(screen_recon.intensity)]
screen_shift = tracking.ScreenDistribution(
screen_recon.x - screen_max_x, screen_recon.intensity.copy())
screen_shift.cutoff(0.1)
diff = screen_shift.compare(meas_screen_shift)
print(ctr, '%f fs' % sig_t_fs, '%.1e' % diff)
ctr += 1
if plot_ctr == 5:
plot_ctr = 0
if sp_ctr == 7:
ms.figure('Optimization bo %.1e' % beam_offset0)
sp_ctr = 1
sp = subplot(sp_ctr, title='Profile')
sp_ctr += 1
sp.plot(meas_screen_shift.x * 1e3,
meas_screen_shift.intensity,
label='Original')
plot_ctr += 1
sp.plot(screen_shift.x * 1e3,
screen_shift.intensity,
label='%i: %.1f fs %.3e' % (ctr, sig_t_fs, diff))
sp.legend()
plt.show()
plt.pause(0.01)
if count_nfev:
nfev_ctr += 1
if nfev_ctr > max_nfev:
raise StopIteration(sig_t_fs)
opt_func_values.append((float(sig_t), diff))
return diff
convoluted_screen = convolve(screen_meas00.intensity, cut_gauss)
diff = np.diff(xx_arr).mean()
convolved_xx = np.arange(0, len(convoluted_screen)*diff, diff) - len(cut_gauss)/2*diff
zero_arr = np.zeros([len(cut_gauss)//2])
intensity_padded = np.concatenate([zero_arr, screen_meas.intensity, zero_arr[:-1]])
deconvolved, remainder = deconvolve(intensity_padded, cut_gauss)
#random_distortion = np.random.randn(len(convoluted_screen))*0.01*convoluted_screen.max()
#random_distortion = 0
#deconvolved, remainder = deconvolve(intensity_padded, cut_gauss)
screen_decon = tracking.ScreenDistribution(screen_meas.x, deconvolved)
ms.figure('Screens')
subplot = ms.subplot_factory(2,2)
sp_ctr = 1
sp0 = subplot(1, title='Screens')
sp0.plot(screen_meas.x, screen_meas.intensity, label='Streak 300 nm')
sp0.plot(screen_meas0.x, screen_meas0.intensity/screen_meas0.intensity.max()*screen_meas.intensity.max(), label='No Streak 300 nm')
sp0.plot(screen_meas00.x, screen_meas00.intensity, label='Streak 1 nm')
sp0.plot(convolved_xx, convoluted_screen, label='1 nm streak x 300 nm no streak')
sp0.plot(screen_decon.x, screen_decon.intensity, label='Deconvoluted Streak 300 nm')
color = sp.plot(x_axis, projX)[0].get_color()
sp.axvline(gf.mean, color=color, ls='--')
sp.plot(gf.xx, gf.reconstruction, ls='--', color=color)
all_means.append(gf.mean)
n_proj += 1
sp.set_xlim(-0.001, 0.001)
all_means = np.array(all_means)
sp_all = subplot(sp_ctr, title='All means')
sp_ctr += 1
sp_all.hist(all_means * 1e6)
bp_test = tracking.get_gaussian_profile(40e-15, tt_halfrange, len_profile,
charge, energy_eV)
screen_sim = tracker.matrix_forward(bp_test, [10e-3, 10e-3], [0, 0])['screen']
emittances_fit = []
for n_image, image in enumerate(images0):
screen_meas = tracking.ScreenDistribution(x_axis, image.T.sum(axis=0))
emittance_fit = misc.fit_nat_beamsize(screen_meas, screen_sim,
n_emittances[0])
emittances_fit.append(emittance_fit)
emittances_fit = np.array(emittances_fit)
mean_emittance = emittances_fit.mean()
plt.show()
image = image[:,::-1]
if np.diff(y_axis0)[0] < 0:
image = image[::-1,:]
return image
n_offset = 0
image = get_image(n_offset, 0)
image0 = get_image(-1, 0)
ms.figure('Investigate')
subplot = ms.subplot_factory(2,2)
sp_ctr = 1
screen = tracking.ScreenDistribution(x_axis.copy(), image.sum(axis=0))
screen0 = tracking.ScreenDistribution(x_axis.copy(), image0.sum(axis=0))
shift0 = float(screen0.gaussfit.mean)
screen._xx -= shift0
screen0._xx -= shift0
screen._yy -= screen._yy.min()
screen.reshape(1e3)
screen.cutoff(0.05)
screen.remove0()
screen.reshape(1e3)
sp_proj = subplot(sp_ctr, title='Screen', xlabel='Position [mm]', ylabel='Intensity (arb. units)')
sp_ctr += 1
def reconstruct_current(data_file_or_dict, n_streaker, beamline, tracker_kwargs_or_tracker, rec_mode, kwargs_recon, screen_x0, streaker_centers, blmeas_file=None, plot_handles=None, do_plot=True):
if type(tracker_kwargs_or_tracker) is dict:
tracker = tracking.Tracker(**tracker_kwargs_or_tracker)
elif type(tracker_kwargs_or_tracker) is tracking.Tracker:
tracker = tracker_kwargs_or_tracker
else:
raise ValueError(type(tracker_kwargs_or_tracker))
if type(data_file_or_dict) is dict:
screen_data = data_file_or_dict
else:
screen_data = h5_storage.loadH5Recursive(data_file_or_dict)
if 'meta_data' in screen_data:
meta_data = screen_data['meta_data']
elif 'meta_data_begin' in screen_data:
meta_data = screen_data['meta_data_begin']
else:
print(screen_data.keys())
raise ValueError
if 'pyscan_result' in screen_data:
pyscan_data = screen_data['pyscan_result']
else:
pyscan_data = screen_data
x_axis = pyscan_data['x_axis_m']
projx = pyscan_data['image'].sum(axis=-2)
if rec_mode == 'Median':
median_projx = misc.get_median(projx)
proj_list = [median_projx]
elif rec_mode == 'All':
proj_list = projx
tracker.set_simulator(meta_data)
if x_axis[1] < x_axis[0]:
revert = True
x_axis = x_axis[::-1]
else:
revert = False
output_dicts = []
for proj in proj_list:
if revert:
proj = proj[::-1]
meas_screen = tracking.ScreenDistribution(x_axis, proj)
kwargs_recon['meas_screen'] = meas_screen
#print('Analysing reconstruction')
kwargs = copy.deepcopy(kwargs_recon)
gaps, streaker_offsets = get_gap_and_offset(meta_data, beamline)
kwargs['meas_screen']._xx = kwargs['meas_screen']._xx - screen_x0
kwargs['beam_offsets'] = -(streaker_offsets - streaker_centers)
kwargs['gaps'] = gaps
kwargs['meas_screen'].cutoff2(tracker.screen_cutoff)
kwargs['meas_screen'].crop()
kwargs['meas_screen'].reshape(tracker.len_screen)
kwargs['n_streaker'] = n_streaker
# Only allow one streaker at the moment
for n in (0,1):
if n != kwargs['n_streaker']:
kwargs['beam_offsets'][n] = 0
gauss_dict = current_profile_rec_gauss(tracker, kwargs, plot_handles, blmeas_file, do_plot=do_plot)
output_dict = {
'input': {
'data_file_or_dict': data_file_or_dict,
'n_streaker': n_streaker,
'beamline': beamline,
'tracker_kwargs': tracker_kwargs_or_tracker,
'rec_mode': rec_mode,
'kwargs_recon': kwargs_recon,
'screen_x0': screen_x0,
'streaker_centers': streaker_centers,
'blmeas_file': blmeas_file,
},
'gauss_dict': gauss_dict,
}
output_dicts.append(output_dict)
if rec_mode == 'Median':
return output_dict
elif rec_mode == 'All':
return output_dicts
else:
print(rec_mode)
import tracking
import myplotstyle as ms
plt.close('all')
with open('./example_screen.pkl', 'rb') as f:
screen_old = pickle.load(f)
ms.figure('Investigate Smoothing and binning')
subplot = ms.subplot_factory(2, 2)
sp_ctr = 1
sp = subplot(sp_ctr, title='Histtograms')
sp_ctr += 1
sp2 = subplot(sp_ctr, title='Histograms smooth')
sp_ctr += 1
xx = screen_old.real_x
for n_bins in (150, 300, 450):
hist, bin_edges = np.histogram(xx, bins=n_bins, density=True)
screen = tracking.ScreenDistribution(bin_edges[1:], hist)
screen.smoothen(30e-6)
sp.plot(bin_edges[1:], hist, label=n_bins)
sp2.plot(screen.x, screen.intensity, label=n_bins)
plt.show()
(meas_screen2, 'Streaking back and forward real'),
(meas_screen3, 'Streaking back and forward gauss'),
#(forward_dict_gauss['screen'], 'Forward gauss'),
#(forward_dict_real['screen'], 'Forward real'),
]:
sp.plot(sd.x * 1e3, sd.intensity, label=label)
sp.legend()
ctr = 0
plot_ctr = 5
sp_ctr = 7
subplot = ms.subplot_factory(2, 3)
meas_screen_max_x = meas_screen.x[np.argmax(meas_screen.intensity)]
meas_screen_shift = tracking.ScreenDistribution(
meas_screen.x - meas_screen_max_x, meas_screen.intensity.copy())
meas_screen_shift.cutoff(0.1)
nfev_ctr = 0
max_nfev = 15
opt_func_values = []
@functools.lru_cache()
def opt_func(sig_t_fs, count_nfev):
global ctr, sp_ctr, plot_ctr, sp, nfev_ctr
sig_t = sig_t_fs / 1e15
bp_wake = tracking.get_gaussian_profile(sig_t, tt_halfrange,
len_profile, charge,
tracker.energy_eV)
