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()