def correct_image(img, x_axis, y_axis, limits):
image = iap.Image(img, x_axis, y_axis)
index1 = np.argmin((image.x_axis - limits[0])**2)
index2 = np.argmin((image.x_axis - limits[1])**2) + 1
new_image = image.image.copy()
new_image[:, index1:index2 - 1] = new_image[:, index1][:, np.newaxis]
proj = np.sum(new_image, axis=0)
interp_proj = np.interp(image.x_axis[index1:index2],
[image.x_axis[index1], image.x_axis[index2]],
[proj[index1], proj[index2]])
new_image[:, index1:index2] *= interp_proj / proj[index1:index2]
image_corrected = iap.Image(new_image, image.x_axis, image.y_axis)
return image, image_corrected
def add_images(self, meta_data, images, x_axis, y_axis, max_index=None):
self.meta_data = meta_data
self.epics_gap = meta_data[self.streaker+':GAP']*1e-3
self.epics_center = meta_data[self.streaker+':CENTER']*1e-3
self.gap = self.epics_gap + self.delta_gap
self.streaker_center = self.epics_center - self.streaker_offset
self.x_axis0 = x_axis
self.x_axis = x_axis - self.screen_x0
self.y_axis = y_axis
self.raw_images = images
self.raw_image_objs = []
self.meas_screens = []
rms_arr = []
for n_image, img in enumerate(images):
if max_index is not None and n_image >= max_index:
break
if self.subtract_median:
img = img - np.median(img)
img[img<0] = 0
image = iap.Image(img, self.x_axis, y_axis)
self.raw_image_objs.append(image)
screen = iap.ScreenDistribution(image.x_axis, image.image.sum(axis=-2))
self.meas_screens.append(screen)
rms_arr.append(screen.rms())
self.median_meas_screen_index = np.argsort(np.array(rms_arr))[len(self.meas_screens)//2]
def correct_image(img, x_axis, y_axis, limits):
image = iap.Image(img, x_axis, y_axis)
index1 = np.argmin((image.x_axis - limits[0])**2)
index2 = np.argmin((image.x_axis - limits[1])**2) + 1
new_image = image.image.copy()
for index_y in range(new_image.shape[0]):
new_image[index_y, index1:index2] = np.interp(
image.x_axis[index1:index2],
[image.x_axis[index1], image.x_axis[index2]], [
image.image[index_y, index1], image.image[index_y,
index2]
])
image_corrected = iap.Image(new_image, image.x_axis,
image.y_axis)
return image, image_corrected
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,:]
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
with open('./backtrack_image_no_compensate.pkl', 'rb') as f:
d = pickle.load(f)
image = d['image']
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)
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)
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)
full_slice_dict[label][n_image, 0] = slice_dict['slice_x']
full_slice_dict[label][n_image, 1] = slice_dict['slice_mean']
full_slice_dict[label][n_image, 2] = slice_dict['slice_sigma_sq']
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, :]
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)
image_off = iap.Image(image_off, x_axis, y_axis)
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
index_x - 1]
return new_image
old_pyscan_image = data_dict['pyscan_result']['image'].astype(
float)
new_pyscan_image = np.zeros_like(old_pyscan_image)
for n_image, img in enumerate(old_pyscan_image):
image, image_corrected = correct_image(img, x_axis, y_axis,
limits)
img_sat = correct_saturation(image_corrected.image,
image_corrected.x_axis,
image_corrected.y_axis)
image_sat = iap.Image(img_sat, image_corrected.x_axis,
image_corrected.y_axis)
new_pyscan_image[n_image] = image_sat.image
if plot_repair:
ms.figure('Repair data')
subplot = ms.subplot_factory(2, 2, grid=False)
sp_ctr = 1
sp_raw = subplot(sp_ctr, title='Image raw')
sp_ctr += 1
sp_corrected = subplot(sp_ctr, title='Image corrected')
sp_ctr += 1
sp_sat = subplot(sp_ctr, title='Saturation')
sp_ctr += 1
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
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)
strong_streaking_file = data_dir1 + '2021_05_18-23_43_39_Lasing_False_SARBD02-DSCR050.h5'
strong_streaking_dict = h5_storage.loadH5Recursive(strong_streaking_file)
meta_data2 = strong_streaking_dict['meta_data_begin']
strong_calib_file = data_dir1 + '2021_05_18-23_32_12_Calibration_SARUN18-UDCP020.h5'
fig1 = ms.figure('Backtracked images 2', figsize=(9, 7))
ms.plt.subplots_adjust(hspace=0.3)
sp_ctr1 = 1
dispersion = tracker.calcDisp()[1]
all_slice_dict = OrderedDict()
for ctr, (image, label) in enumerate([(image_off, 'Lasing_off'),
(image_on, 'Lasing_on')]):
image_reduced = image - np.median(image)
np.clip(image_reduced, 0, None, out=image_reduced)
image_obj = iap.Image(image_reduced, x_axis, y_axis)
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)
if ctr == 0:
ref_y = None
image_tE, ref_y = image_t.y_to_eV(dispersion, tracker.energy_eV, ref_y)
if ctr == 1:
image_tE.y_axis -= 250e3
image_t_reduced = image_tE.slice_x(n_slices)
slice_dict = image_t_reduced.fit_slice(intensity_cutoff=None,
smoothen_first=True,
smoothen=1e6)
all_slice_dict[label] = slice_dict
#ms.plt.figure(fig.number)
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,
xlabel='x (mm)',
ylabel='y (mm)',
title_fs=title_fs)
sp_ctr += 1
img.plot_img_and_proj(sp_img, xlim=xlim, ylim=ylim, plot_gauss=False)
sumx = raw_image.sum(axis=0)
prof = iap.AnyProfile(x_axis, sumx - np.min(sumx))
prof.cutoff2(0.05)
prof.crop()
prof.reshape(1000)
x_rms = prof.rms()
x_gf = prof.gaussfit.sigma
print('%s RMS: %i um; Gauss sigma: %i um' %
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)
x_axis = lasing_on['camera1']['x_axis'].astype(float) * 1e-6
y_axis = lasing_on['camera1']['y_axis'].astype(float) * 1e-6
image_obj_on = iap.Image(image_on,
x_axis,
y_axis,
subtract_median=True,
x_offset=x0)
image_obj_off = iap.Image(image_off,
x_axis,
y_axis,
subtract_median=True,
x_offset=x0)
tt_halfrange = 200e-15
charge = 200e-12
screen_cutoff = 2e-3
profile_cutoff = 2e-2
len_profile = int(2e3)
struct_lengths = [1., 1.]
screen_bins = 400