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
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'] all_emittances = [] all_beamsizes = [] for proj in projx0: screen_meas = get_screen_from_proj(proj, x_axis0, invert_x0) all_beamsizes.append(screen_meas.gaussfit.sigma) emittance_fit = misc.fit_nat_beamsize(screen_meas, screen_sim, n_emittances[0], smoothen, print_=False) #print(screen_meas.gaussfit.sigma) all_emittances.append(emittance_fit)
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()
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 screen.plot_standard(sp_proj, label='Streaking') screen0.plot_standard(sp_proj, label='No streaking') 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) r12 = tracker.calcR12()[1] bp_gauss = tracking.get_gaussian_profile(30e-15, 200e-15, 1e3, charge, tracker.energy_eV) offset_arr = dict_['value']*1e-3 - mean_struct2 wf = bp_gauss.calc_wake(gaps[1], offset_arr[n_offset], struct_lengths[1]) dipole_wake = wf['dipole']['wake_potential'] dipole_t = wf['input']['charge_xx']/c sp_wf = subplot(sp_ctr, title='Wakefield', xlabel='Time [fs]', ylabel='Wake [V/m]') sp_ctr += 1 sp_wf.plot(dipole_t, dipole_wake) sp_deltax = subplot(sp_ctr, title='Wake effect', xlabel='Time [fs]', ylabel='$\Delta$ x') sp_ctr += 1
def opt_func(sig_t_fs, count_nfev, profile_cutoff, screen_cutoff, smoothen): a = np.array(opt_func_values) if len(a) > 0 and np.any(a[:, 0] == sig_t_fs): index = np.argwhere(sig_t_fs == a[:, 0])[0] return a[index, 1] global opt_ctr, sp_ctr, plot_ctr, sp, nfev_ctr, sp2 sig_t = sig_t_fs / 1e15 bp_wake = tracking.get_gaussian_profile(sig_t, tt_halfrange, len_profile, charge, tracker.energy_eV) baf = tracker.back_and_forward(meas_screen, bp_wake, gaps, beam_offsets, n_streaker) if self_consistent: baf_self = tracker.back_and_forward(meas_screen, baf['beam_profile'], gaps, beam_offsets, n_streaker) screen_self = baf_self['screen'] else: screen_self = baf['screen'] profile = baf['beam_profile'] diff = screen_self.compare(meas_screen) print(opt_ctr, '%f fs' % sig_t_fs, '%.1e' % diff) opt_ctr += 1 if opt_plot: if plot_ctr == 5: plot_ctr = 0 if sp_ctr == ny * nx + 1: ms.figure('Optimization %s' % label) sp_ctr = 1 sp = subplot(sp_ctr, title='Screen') sp_ctr += 1 sp.plot(meas_screen_shift.x * 1e3, meas_screen_shift.intensity / meas_screen_shift.integral, label='Original') sp2 = subplot(sp_ctr, title='Profile') sp_ctr += 1 sp2.plot(profile_meas.time * 1e15, profile_meas.current / profile_meas.integral, label='Original') plot_ctr += 1 sp.plot(screen_self.x * 1e3, screen_self.intensity / screen_self.integral, label='%i: %.1f fs %.3e' % (opt_ctr, sig_t_fs, diff)) sp2.plot(profile.time * 1e15, profile.current / profile.integral, label='%i: %.1f fs %.3e' % (opt_ctr, sig_t_fs, diff)) sp.legend() sp2.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_fs), diff)) opt_func_screens.append(screen_self) opt_func_profiles.append(baf['beam_profile']) return diff
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') energy_eV = tracker.energy_eV profile_meas = tracking.profile_from_blmeas(bl_meas_file, tt_halfrange, charge, energy_eV, subtract_min=False) profile_gauss = tracking.get_gaussian_profile(sig_t, tt_halfrange, len_profile, charge, energy_eV) meas_screen_dict = {} for beam_offset0 in beam_offsets_s0_real: beam_offsets = [beam_offset0, 0.] meas_screen = tracker.elegant_forward(profile_meas, gaps, beam_offsets, n_bins)['screen'] meas_screen.smoothen(smoothen) meas_screen.cut(backtrack_cutoff) for beam_offset0 in beam_offsets_s0: beam_offsets = [beam_offset0, 0.] fab_dict_real = tracker.forward_and_back(profile_meas, profile_meas, gaps, beam_offsets,
hostname = gethostname() if hostname == 'desktop': dirname1 = '/storage/data_2020-10-03/' dirname2 = '/storage/data_2020-10-04/' archiver_dir = '/storage/Philipp_data_folder/' elif hostname == 'pc11292.psi.ch': dirname1 = '/sf/data/measurements/2020/10/03/' dirname2 = '/sf/data/measurements/2020/10/04/' elif hostname == 'pubuntu': dirname1 = '/home/work/data_2020-10-03/' dirname2 = '/home/work/data_2020-10-04/' archiver_dir = '/home/work/' bp_gauss = tracking.get_gaussian_profile(40e-15, 200e-15, len_profile, charge, energy_eV) flat_current = np.zeros_like(bp_gauss.current) flat_time = bp_gauss.time flat_current[np.logical_and(flat_time > -40e-15, flat_time < 40e-15)] = 1 bp_flat = tracking.BeamProfile(flat_time, flat_current, energy_eV, charge) sig = 5e-15 dhf_current = doublehornfit.DoublehornFit.fit_func(flat_time, -20e-15, 20e-15, sig, sig, sig, sig, 0.5, 1, 1) dhf_current *= charge / dhf_current.sum() bp_dhf = tracking.BeamProfile(flat_time, dhf_current, energy_eV, charge) bp_dhf.cutoff(1e-3)
screen_bins=screen_bins, smoothen=smoothen, profile_cutoff=profile_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() profile_gauss = tracking.get_gaussian_profile(sig_t, tt_halfrange, len_profile, charge, tracker.energy_eV, cutoff=profile_cutoff) screen0 = tracker.matrix_forward(profile_meas, gaps, [0, 0])['screen'] meas_screen = tracker.matrix_forward(profile_meas, gaps, beam_offsets)['screen'] ms.figure('') subplot = ms.subplot_factory(2, 2) sp_ctr = 1 sp_profile0 = subplot(sp_ctr, title='Beam profiles', xlabel='t [fs]', ylabel='Current (arb. units)')