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 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
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 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 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
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,
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()
for l in ll:
l_len = l.shape[-1]
bpm_data[key][:, ctr:ctr + l_len] = l
ctr += l_len
# filter out non-unique bpm data
for a in bpm_data.values():
for n_col in range(a.shape[0]):
old = a[n_col].copy()
a[n_col] = np.nan
arr2 = np.array(list(set(old)))
a[n_col][:len(arr2)] = arr2
if sp_ctr > ny * nx:
sp_ctr = 1
fig_raw = ms.figure('Streaker gap scan (raw)', figsize=figsize)
fig_rescaled = ms.figure('Streaker gap scan (normalized by R12)',
figsize=figsize)
plt.figure(fig_raw.number)
sp_raw = subplot(sp_ctr,
title='Streaker %i Gap %.1f mm' % (n_streaker, gap),
xlabel='Offset [mm]',
ylabel='BPM reading [mm]')
plt.figure(fig_rescaled.number)
sp_rescaled = subplot(sp_ctr,
title='Streaker %i Gap %.1f mm' %
(n_streaker, gap),
xlabel='Offset [mm]',
ylabel='BPM reading / R12',
sciy=True)
import wf_model
import data_loader
import elegant_matrix
data_dir = '/storage/data_2020-02-03/'
bpms = ['SARBD02.DBPM010', 'SARBD02.DBPM040']
reverse_current_profile = True
linear_fit_details = True
quadratic_fit_details = True
plt.close('all')
figsize = (16, 12)
fig = ms.figure('BPM trajectories in SARBD02', figsize=figsize)
subplot = ms.subplot_factory(2, 3)
sp_ctr = 1
xlabel = 'Offset [mm]'
ylabel = 'BPM reading [mm]'
sp_x = subplot(sp_ctr, title='DEH1 BPM 010 X', xlabel=xlabel, ylabel=ylabel)
sp_ctr += 1
sp_x2 = subplot(sp_ctr, title='DEH1 BPM 040 X', xlabel=xlabel, ylabel=ylabel)
sp_ctr += 1
sp_x3 = subplot(sp_ctr, title='DEH2 BPM 010 X', xlabel=xlabel, ylabel=ylabel)
sp_ctr += 1
def get_image(i, j):
image = dict_['Image'][i][j].T.astype(np.float64)
if np.diff(x_axis0)[0] < 0:
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()
plt.close('all')
pi = np.pi
sigma = uwf.sigma_cu
ctau = uwf.ctau_cu
#integral_re = uwf.impedance_flat_re(1, 3e-3)
a = 2.5e-3
s0 = uwf.s0(a, uwf.sigma_cu)
potency_2 = 13
n_s0 = 10
n_samples = int(1e3)
fig = ms.figure('Details of flat wake calc')
subplot = ms.subplot_factory(2, 3)
sp_ctr = 1
for s0_factor in range(0, 5):
s0_factor *= 2
k = s0_factor / s0
sigma_tilde = sigma / (1 - 1j * k * ctau)
lamb = np.sqrt((2 * pi * sigma * np.abs(k)) / c) * (1j + np.sign(k))
lamb_tilde = np.sqrt(
(2 * pi * sigma_tilde * np.abs(k)) / c) * (1j + np.sign(k))
sp = subplot(sp_ctr, title='k*s0=%i' % s0_factor)
sp_ctr += 1
xx = np.linspace(-10000, 10000, n_samples)
profile_meas = tracking.profile_from_blmeas(bl_meas_file,
tt_halfrange,
energy_eV,
subtract_min=False)
profile_back = tracker.forward_and_back(profile_meas, profile_meas, gaps,
beam_offsets, 0)
#track_dict_forward = tracker.elegant_forward(profile_meas, gaps, beam_offsets, [1, 1])
#track_dict_forward0 = tracker.elegant_forward(profile_meas, gaps, [0,0], [1, 1])
#
#wf_dict = profile_meas.calc_wake(gaps[0], beam_offsets[0], 1.)
#wake_effect = profile_meas.wake_effect_on_screen(wf_dict, track_dict_forward0['r12_dict'][0])
#profile_back = tracker.track_backward(track_dict_forward, track_dict_forward0, wake_effect)
#profile_back.reshape(len_profile)
ms.figure('Back and forward with real profile')
subplot = ms.subplot_factory(2, 2)
sp_ctr = 1
sp = subplot(sp_ctr, title='Profiles')
for bp, label in [(profile_meas, 'Measured'), (profile_back, 'Reconstructed')]:
norm = np.trapz(bp.current, bp.time)
sp.plot(bp.time * 1e15, bp.current / norm, label=label)
sp.legend()
plt.show()
'Passive_data_20201003T231958.mat',
'Passive_data_20201003T233852.mat',]
files2 = [
'Passive_data_20201004T161118.mat',
'Passive_data_20201004T172425.mat',
'Passive_data_20201004T223859.mat',
'Passive_data_20201004T163828.mat',
'Passive_data_20201004T221502.mat',
#'Passive_money_20201004T012247.mat',
]
#blmeas_1 = dirname1+'Bunch_length_meas_2020-10-03_15-43-29.h5'
blmeas_1 = dirname1+'129833611_bunch_length_meas.h5'
blmeas_2 = dirname2+'129858802_bunch_length_meas.h5'
ms.figure('Measured beam profiles')
subplot = ms.subplot_factory(2,2)
sp_ctr = 1
sp = subplot(sp_ctr, title='Beam profiles', xlabel='time [fs]', ylabel='Current (arb. units)')
sp_ctr += 1
for blmeas in (blmeas_1, blmeas_2):
for zero_crossing in (1, 2):
bp = tracking.profile_from_blmeas(blmeas, tt_halfrange, charge, energy_eV=1, subtract_min=True, zero_crossing=zero_crossing)
if zero_crossing == 2:
bp.flipx()
label = '%s %i' % (blmeas, zero_crossing)
bp.plot_standard(sp, label=label)
sp.legend()
x_axis_cut = x_axis[x_mask]
image2 = np.zeros([len(y_axis), len_profile])
x_axis2 = np.linspace(x_axis_cut.min(), x_axis_cut.max(), len_profile)
delta_x = np.zeros_like(image_cut)
delta_x[:, :-1] = image_cut[:, 1:] - image_cut[:, :-1]
grid_points, points = x_axis_cut, x_axis2
index_float = (points - grid_points[0]) / (grid_points[1] - grid_points[0])
index = index_float.astype(int)
index_delta = index_float - index
np.clip(index, 0, len(grid_points) - 1, out=index)
image2 = image_cut[:, index] + index_delta * delta_x[:, index]
#image2 *= image_cut.sum()/image2.sum()
figure = ms.figure('Backtrack image')
ms.plt.subplots_adjust(hspace=0.3)
subplot = ms.subplot_factory(2, 3)
sp_ctr = 1
sp = subplot(sp_ctr, title='X space 1', xlabel='x [mm]', ylabel='y [mm]')
sp_ctr += 1
plot_img_and_proj(sp, image_cut, x_axis_cut, y_axis)
sp = subplot(sp_ctr, title='X space 2', xlabel='x [mm]', ylabel='y [mm]')
sp_ctr += 1
plot_img_and_proj(sp, image2, x_axis2, y_axis)
new_img0 = np.zeros_like(image2)
new_t_axis = np.linspace(t_axis.min(), t_axis.max(), new_img0.shape[1])
diff_t = np.concatenate([[0], np.diff(tt)])
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
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)
forward_dict = tracker.matrix_forward(profile_meas, [10e-3, 10e-3], [0., 0.])
forward_dict_ele = tracker.elegant_forward(profile_meas, [10e-3, 10e-3],
[0., 0.])
beam_forward = forward_dict['beam0_at_screen']
image_sim, xedges, yedges = np.histogram2d(beam_forward[0],
beam_forward[2],
bins=(200, 100),
normed=True)
fig = ms.figure('Compare images')
subplot = ms.subplot_factory(2, 2)
sp_ctr = 1
sp_img = subplot(sp_ctr, title='Measured', grid=False)
sp_ctr += 1
sp_img.imshow(image0,
aspect='auto',
extent=(x_axis[0], x_axis[-1], y_axis[-1], y_axis[0]))
sp_img = subplot(sp_ctr, title='Simulated', grid=False)
sp_ctr += 1
sp_img.imshow(image_sim,
aspect='auto',
extent=(xedges[0], xedges[-1], yedges[-1], yedges[0]))
x_axis = d['x_axis']
y_axis = d['y_axis']
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)
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()
magnet_file = '/storage/Philipp_data_folder/archiver_api_data/2020-07-26.h5'
bl_meas_file = '/storage/data_2020-02-03/Bunch_length_meas_2020-02-03_15-59-13.h5'
else:
magnet_file = '/afs/psi.ch/intranet/SF/Beamdynamics/Philipp/data/archiver_api_data/2020-07-26.h5'
bl_meas_file = '/sf/data/measurements/2020/02/03/Bunch_length_meas_2020-02-03_15-59-13.h5'
tracker = tracking.Tracker(magnet_file, timestamp, struct_lengths, energy_eV='file', n_emittances=n_emittances, screen_bins=screen_bins, n_particles=n_particles, smoothen=smoothen, profile_cutoff=profile_cutoff, screen_cutoff=screen_cutoff, len_screen=len_profile)
energy_eV = tracker.energy_eV
profile_meas = tracking.profile_from_blmeas(bl_meas_file, tt_halfrange, charge, energy_eV, subtract_min=False)
profile_meas.center()
meas_screen_dict = {}
ms.figure('Compare offset errors')
sp_ctr = 1
subplot = ms.subplot_factory(2,3)
sp_profile0 = subplot(sp_ctr, title='Beam profiles', xlabel='t [fs]', ylabel='Current (arb. units)')
sp_ctr += 1
sp_profile0.plot(profile_meas.time*1e15, profile_meas.current/profile_meas.integral, label='Real / %i' % (profile_meas.gaussfit.sigma*1e15), color='black', lw=3)
sp_error = subplot(sp_ctr, title='Screen errors', xlabel='Offset error [$\mu$m]', ylabel='Difference at screen (arb. units)')
sp_ctr += 1
sp_screen_dict = OrderedDict()
for n_bo, beam_offsets0 in enumerate(beam_offset_list):
sp = subplot(sp_ctr, title='Measured screen dist offset %.2f mm' % (beam_offsets0[0]*1e3), xlabel='x [mm]', ylabel='Intensity (arb. units)')
long_x_axis = np.linspace(x_axis[0], x_axis[-1], int(100e3))
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')
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')
sp0.legend()
]
L = 58.8
L0 = 1.
L_factor = L / L0
for mat_file, _, total_charge in mat_files_current_charge:
result_dict = loadH5Recursive(os.path.basename(mat_file) + '_wake.h5')
dd = loadmat(data_dir + mat_file)
charge_profile = result_dict['charge_profile']
charge_xx = result_dict['charge_xx']
energy_eV = result_dict['energy_eV']
gap_list = result_dict['gap_list']
#result_dict = {str(i): result_list[i] for i, gap in enumerate(gap_list)}
result_list = [result_dict[str(i)] for i, gap in enumerate(gap_list)]
ms.figure('Undulator wakefield measurements %s' % mat_file)
plt.subplots_adjust(hspace=0.5, wspace=0.4)
subplot = ms.subplot_factory(3, 3)
sp_ctr = 1
sp_charge = subplot(sp_ctr, title='Current_profile')
sp_ctr += 1
sp_charge.plot(charge_xx * 1e6, charge_profile)
xlabel = 's [$\mu$m]'
ylabel = 'w [kV/(pC$\cdot$m)]'
ylabel2 = '$\Delta$ E [MeV]'
ylabel3 = '$\Delta$ E [keV/m]'
sp_wf_surface = subplot(sp_ctr,
title='Surface wake',
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
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': {
for grid_points, points in [(quad_s, beam_s), (quad_x, beam_x + beam_offset)]:
index_float = (points - grid_points[0]) / (grid_points[1] - grid_points[0])
index = np.round(index_float).astype(int)
np.clip(index, 0, len(grid_points) - 1, out=index)
indices.append(index)
vals = quad_wake[indices[0], indices[1]]
vals2 = wake2d_dict_quad['wake_on_particles']
interp_quad = interpolate.interp2d(quad_s, quad_x, quad_wake.T)
quad_effect = np.zeros_like(vals)
for n, (s_, x) in enumerate(zip(beam_s, beam_x + beam_offset)):
quad_effect[n] = interp_quad(s_, x)
ms.figure('Test interp')
subplot = ms.subplot_factory(1, 2)
sp = subplot(1)
sp.hist((vals - quad_effect) / np.mean(np.abs(quad_effect)), bins=100)
sp.set_yscale('log')
sp = subplot(2)
sp.hist((vals2 - quad_effect) / np.mean(np.abs(quad_effect)), bins=100)
sp.set_yscale('log')
ms.plt.show()
},
}
data = {}
for key1, subdict in name_dict.items():
data[key1] = {}
for key2, alternate_key in subdict.items():
data[key1][key2] = {
'xx': mat[alternate_key + '_XData'].squeeze(),
'yy': mat[alternate_key + '_YData'].squeeze(),
'err1': mat[alternate_key + '_YPositiveDelta'].squeeze(),
'err2': mat[alternate_key + '_YNegativeDelta'].squeeze(),
}
full_title = 'Longitudinal measurements %s' % title
fig = ms.figure(title=full_title)
subplot = ms.subplot_factory(2, 3)
sp_ctr = 1
# Plot Raw_data
sp_eloss = subplot(sp_ctr,
title='Energy loss',
xlabel='Gap [mm]',
ylabel='Rel. energy loss',
sciy=True)
sp_ctr += 1
sp_espread = subplot(sp_ctr,
title='Energy spread',
xlabel='Gap [mm]',
dscr0 = 'SARBD01.DSCR050'
dscr1 = 'SARBD02.DSCR050'
#print(mat0[dscr0])
#print(mat1[dscr0])
#print(mat0[dscr1])
print(mat1[dscr1])
sys.exit()
sim0 = ElegantSimulation(
'./for_alex//elegant_5521_2020-06-19_17-17-32_0/SwissFEL_in0.ele')
sim1 = ElegantSimulation(
'./for_alex//elegant_5521_2020-06-19_17-17-33_1/SwissFEL_in0.ele')
ms.figure()
subplot = ms.subplot_factory(2, 2)
sp_ctr = 1
for dim in 'x', 'y':
sp = subplot(sp_ctr, title='Beta_%s' % dim)
sp_ctr += 1
for n_sim, sim in enumerate([sim0, sim1]):
twi = sim.twi
sp.plot(twi['s'], twi['beta%s' % dim], label=n_sim)
sp.set_xlim(None, 25)
sp.legend()
sp = subplot(sp_ctr, title='R12')
sp_ctr += 1
n_emittances[0],
smoothen,
print_=False)
#print(screen_meas.gaussfit.sigma)
all_emittances.append(emittance_fit)
new_emittance = np.mean(all_emittances)
print(main_label, 'Emittance [nm]', new_emittance * 1e9)
n_emittances[0] = new_emittance
tracker0.n_emittances[0] = new_emittance
new_screen0 = tracker0.matrix_forward(bp_test, [10e-3, 10e-3],
[0, 0])['screen']
ms.figure('Test nat bs')
sp = plt.subplot(1, 1, 1)
sp.set_title('New emittance %i nm' % (new_emittance * 1e9))
screen_meas.center()
for screen, label in [(new_screen0, 'New'), (screen_meas, 'Meas'),
(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]
saveH5Recursive('./example_current_profile.h5', {
'time_profile': time_profile,
'current': current,
'energy_eV': energy_eV
})
rescale_factor = 200e-12 / np.sum(current)
yy_charge = current * rescale_factor
xx_space = time_profile * c
xx_space -= xx_space.min()
#cutoff
#yy_charge[yy_charge < 0.1*yy_charge.max()] = 0
fig = ms.figure('Wake effect')
subplot = ms.subplot_factory(2, 3)
sp_ctr = 1
sp_charge = subplot(sp_ctr,
title='Charge',
xlabel='s [m]',
ylabel='Charge [C]',
scix=True,
sciy=True)
sp_ctr += 1
sp_charge.plot(xx_space, yy_charge)
sp_single = subplot(sp_ctr,
title='Single particle wake',
charge_xx = bl_meas['time_profile1'] * c
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):
mean_x_arr = mean_s_arr.copy()
size_x_arr = mean_s_arr.copy()
for n_slice in range(n_slices):
mask = np.logical_and(beam_s > borders[n_slice],
beam_s < borders[n_slice + 1])
mean_x_arr[n_slice] = np.mean(beam_x[mask])
size_x_arr[n_slice] = np.std(beam_x[mask])
mean_s_arr[n_slice] = np.mean(beam_s[mask])
return mean_s_arr, mean_x_arr, size_x_arr
for n_split in range(n_splits):
fig = ms.figure('Test 2d interpolation step %i' % n_split)
subplot = ms.subplot_factory(2, 2)
sp_ctr = 1
ss = np.linspace(0, (beam[4, :].max() - beam[4, :].min()) * c * 1.01,
len_s)
xx = np.linspace(beam[0, :].min(), beam[0, :].max(), len_x) + beam_offset
distance = semigap - beam_offset
sp = sp0 = subplot(sp_ctr,
title='Single particle wake',
xlabel='s [$\mu$m]',
ylabel='Wake [kV/pC/m]')
sp_ctr += 1
