def plot_save_ene_ene_spread(slm, savepath, h5plot=True, t_is_z=True): if t_is_z: xlabel_str = r'$k_p z$' else: xlabel_str = r'$\omega_p t$' t = slm.get_time_array() tot_charge = np.sum(slm.get_charge(), axis=1) avg_gamma = slm.get_proj(pz=1) sigma_gamma = np.sqrt(slm.get_proj(pz=2)) sigma_gamma_per_gamma = np.divide( sigma_gamma, avg_gamma) fig_sg = plt.figure() plt.plot(t, sigma_gamma_per_gamma) ax = plt.gca() if magn_check(sigma_gamma_per_gamma): ax.yaxis.set_major_formatter(FormatStrFormatter('%.1e')) plt.gcf().subplots_adjust(left=0.18) else: plt.gcf().subplots_adjust(left=0.15) ax.set_xlabel(xlabel_str, fontsize=14) ax.set_ylabel(r'$\sigma_\gamma/\gamma$', fontsize=14) saveas_eps_pdf(fig_sg, savepath, 'sigma_gamma_per_gamma', h5plot=h5plot) plt.close(fig_sg) fig_g = plt.figure() plt.plot(t, avg_gamma) ax = plt.gca() ax.set_xlabel(xlabel_str, fontsize=14) ax.set_ylabel(r'$\gamma$', fontsize=14) saveas_eps_pdf(fig_g, savepath, 'avg_gamma', h5plot=h5plot) plt.close(fig_g)
def plot_save_slice_exkurtosis_lines(slm, savepath, time = None, axdir='x', h5plot=True): if time == None: tidx = [0,-1]; else: tidx = [(np.abs(slm.get_time_array() - time)).argmin()] for i in tidx: if axdir == 'y': var_xy = slm.get(y=2)[i,:] var_xy_nonzero_idx = np.nonzero(var_xy)[0] exkurtosis_xy = np.divide( slm.get(y=4)[i,var_xy_nonzero_idx],\ np.power(var_xy[var_xy_nonzero_idx],2)) - 3 exkurtosis_xy_lab = r'$\left \langle y^4 \right\rangle -3$' exkurtosis_xy_savename = 'exkurtosis_y' var_pxy = slm.get(py=2)[i,:] var_pxy_nonzero_idx = np.nonzero(var_pxy)[0] exkurtosis_pxy = np.divide( slm.get(py=4)[i,var_pxy_nonzero_idx],\ np.power(var_pxy[var_pxy_nonzero_idx],2)) - 3 exkurtosis_pxy_lab = r'$\left \langle p_y^4/\sigma_{py}^4 \right\rangle -3$' exkurtosis_pxy_savename = 'exkurtosis_py' else: # xdir == 'x' or xdir == 'r' var_xy = slm.get(x=2)[i,:] var_xy_nonzero_idx = np.nonzero(var_xy)[0] exkurtosis_xy = np.divide( slm.get(x=4)[i,var_xy_nonzero_idx],\ np.power(var_xy[var_xy_nonzero_idx],2)) - 3 exkurtosis_xy_lab = r'$\left \langle x^4/\sigma_x^4 \right\rangle - 3$' exkurtosis_xy_savename = 'exkurtosis_x' var_pxy = slm.get(px=2)[i,:] var_pxy_nonzero_idx = np.nonzero(var_pxy)[0] exkurtosis_pxy = np.divide( slm.get(px=4)[i,var_pxy_nonzero_idx],\ np.power(var_pxy[var_pxy_nonzero_idx],2)) - 3 exkurtosis_pxy_lab = r'$\left \langle p_x^4/\sigma_{px}^4 \right\rangle -3$' exkurtosis_pxy_savename = 'exkurtosis_px' # Also define labels and savenames here! fig_exkurtosis_xy = plt.figure() plt.plot( slm.get_zeta_array()[var_xy_nonzero_idx], exkurtosis_xy) ax = plt.gca() ax.set_xlabel(r'$k_p \zeta$', fontsize=14) ax.set_ylabel(exkurtosis_xy_lab, fontsize=14) saveas_eps_pdf(fig_exkurtosis_xy, savepath, ('%s_time_%0.1f' % (exkurtosis_xy_savename, slm.get_time_array()[i])), h5plot=h5plot) plt.close(fig_exkurtosis_xy) fig_exkurtosis_pxy = plt.figure() plt.plot( slm.get_zeta_array()[var_pxy_nonzero_idx], exkurtosis_pxy) ax = plt.gca() ax.set_xlabel(r'$k_p \zeta$', fontsize=14) ax.set_ylabel(exkurtosis_pxy_lab, fontsize=14) saveas_eps_pdf(fig_exkurtosis_pxy, savepath, ('%s_time_%0.1f' % (exkurtosis_pxy_savename, slm.get_time_array()[i])), h5plot=h5plot) plt.close(fig_exkurtosis_pxy)
def oneD(raw, args): appstr = '' psv, xlabel, savename = get_props(raw, args.psv) if args.nbins == None: nbins = np.int(10 + np.power(np.log(raw.get_npart()), 3)) else: nbins = args.nbins hist, bin_edges = np.histogram( psv,\ bins=nbins,\ weights=raw.get('q'),\ density=True,\ range=args.range) x_array = bin_edges[0:-1] + (bin_edges[1] - bin_edges[0]) / 2 if args.psv == 'r': ## Scaling dr with 1/r # rmax = np.max(psv) # r_range = [0, rmax] # r_lsp = np.linspace(0,rmax,nbins) # lsp_centers = r_lsp[:-1] + np.diff(r_lsp) # dr_unnormed = np.divide(np.diff(r_lsp),lsp_centers) # dr_normed = rmax * dr_unnormed/np.sum(dr_unnormed) # r_bin_edges = np.insert(np.cumsum(dr_normed),0,0) #hist, bin_edges = np.histogram(psv,bins=r_bin_edges,weights=raw.get('q'),density=True,range=args.range) ## Scaling number of hits with 1/r hist = np.divide(hist, x_array) / (2 * math.pi) fig = plt.figure(figsize=(6, 5)) plt.plot(x_array, hist) ax = plt.gca() ax.set_xlabel(xlabel, fontsize=14) plt.gcf().subplots_adjust(left=0.15, bottom=0.15) if args.part_selection_criterion != '' or args.tagfile_in != None: appstr += '_sel' raw.read_attrs() timestamp = '_%08.1f' % (raw.get_time()) savename += get_zeta_range_str(args.zeta_range) \ + appstr \ + timestamp if args.file_format == 'png': saveas_png(fig, args.savepath, savename, verbose=True) else: saveas_eps_pdf(fig, args.savepath, savename, h5plot=True, verbose=True, fformat=args.file_format) plt.close(fig)
def plot_save_slice_rms_lines(slm, savepath, time = None, axdir='x', h5plot=True): if time == None: tidx = [0,-1]; else: tidx = [(np.abs(slm.get_time_array() - time)).argmin()] for i in tidx: if axdir == 'x': sigma_xy = np.sqrt(slm.get(x=2)[i,:]) sigma_xy_lab = r'$k_p \sigma_{x}$' sigma_xy_savename = 'sigma_x' sigma_pxy = np.sqrt(slm.get(px=2)[i,:]) sigma_pxy_lab = r'$\sigma_{p_x}/mc$' sigma_pxy_savename = 'sigma_px' # Also define labels and savenames here! elif axdir == 'y': sigma_xy = np.sqrt(slm.get(y=2)[i,:]) sigma_xy_lab = r'$k_p \sigma_{y}$' sigma_xy_savename = 'sigma_y' sigma_pxy = np.sqrt(slm.get(py=2)[i,:]) sigma_pxy_lab = r'$\sigma_{p_y}/mc$' sigma_pxy_savename = 'sigma_py' elif axdir == 'r': sigma_xy = np.sqrt(slm.get(x=2)[i,:] + slm.get(y=2)[i,:]) sigma_xy_lab = r'$k_p \sigma_{r}$' sigma_xy_savename = 'sigma_r' sigma_pxy = np.sqrt(slm.get(px=2)[i,:] + slm.get(py=2)[i,:]) sigma_pxy_lab = r'$(\sigma_{p_x}^2+\sigma_{p_y}^2)^{1/2}/mc$' sigma_pxy_savename = 'sigma_pr' fig_sigma_xy = plt.figure() plt.plot( slm.get_zeta_array(), sigma_xy) ax = plt.gca() ax.set_xlabel(r'$k_p \zeta$', fontsize=14) ax.set_ylabel(sigma_xy_lab, fontsize=14) saveas_eps_pdf(fig_sigma_xy, savepath, ('%s_time_%0.1f' % (sigma_xy_savename, slm.get_time_array()[i])), h5plot=h5plot) plt.close(fig_sigma_xy) fig_sigma_pxy = plt.figure() plt.plot( slm.get_zeta_array(), sigma_pxy) ax = plt.gca() ax.set_xlabel(r'$k_p \zeta$', fontsize=14) ax.set_ylabel(sigma_pxy_lab, fontsize=14) saveas_eps_pdf(fig_sigma_pxy, savepath, ('%s_time_%0.1f' % (sigma_pxy_savename, slm.get_time_array()[i])), h5plot=h5plot) plt.close(fig_sigma_pxy)
def plot_curr_profile(slm, savepath, time = None, h5plot=True, t_is_z=True): dzeta = abs(slm.get_zeta_array()[1] - slm.get_zeta_array()[0]); curr = slm.get_charge() / dzeta if t_is_z: xlabel_str = r'$k_p z$' else: xlabel_str = r'$\omega_p t$' # print('Q = %0.3e' % (np.sum(curr) * dzeta )) # I_A = 4 * pi * epsilon_0 * m * c^3 / e # [curr] = epsilon_0 * m * c^3 / e # curr * 4 * pi = I_b/I_A Ib_per_IA = curr/(4 * math.pi) if time == None: tidx = [0,-1]; else: tidx = [(np.abs(slm.get_time_array() - time)).argmin()] for i in tidx: fig = plt.figure() plt.plot(slm.get_zeta_array(), Ib_per_IA[i,:]) ax = plt.gca() ax.set_xlabel(r'$k_p \zeta$', fontsize=14) ax.set_ylabel(r'$I_{b}/I_A$', fontsize=14) saveas_eps_pdf(fig, savepath, ('Ib_time_%0.1f' % (slm.get_time_array()[i])), h5plot=h5plot) plt.close(fig) figIb = plt.figure() cax = plt.pcolormesh( slm.get_zeta_array(), slm.get_time_array(), Ib_per_IA, cmap=cm.Greys, vmin=0, vmax=np.amax(abs(Ib_per_IA))) cbar = figIb.colorbar(cax) cbar.ax.set_ylabel('$I_b/I_A$', fontsize=14) ax = plt.gca() ax.set_xlabel(r'$k_p \zeta$', fontsize=14) ax.set_ylabel(xlabel_str, fontsize=14) saveas_png(figIb, savepath, 'Ib') plt.close(figIb)
def plot_save_slice_ene(slm, savepath, time = None, h5plot=True, t_is_z=True): if t_is_z: xlabel_str = r'$k_p z$' else: xlabel_str = r'$\omega_p t$' gamma = np.sqrt( 1 + np.power(slm.get(pz=1),2) + np.power(slm.get(px=1),2) + np.power(slm.get(py=1),2) ) if time == None: tidx = [0,-1]; else: tidx = [(np.abs(slm.get_time_array() - time)).argmin()] for i in tidx: fig = plt.figure() plt.plot(slm.get_zeta_array(), gamma[i,:]) ax = plt.gca() ax.set_xlabel(r'$k_p \zeta$', fontsize=14) ax.set_ylabel(r'$\gamma$', fontsize=14) saveas_eps_pdf(fig, savepath, ('gamma_time_%0.1f' % (slm.get_time_array()[i])), h5plot=h5plot) plt.close(fig) figG = plt.figure() cax = plt.pcolormesh( slm.get_zeta_array(), slm.get_time_array(), gamma, cmap=cm.GnBu, vmin=np.amin(gamma), vmax=np.amax(gamma)) ax = plt.gca() ax.set_xlabel(r'$k_p \zeta$', fontsize=14) ax.set_ylabel(xlabel_str, fontsize=14) cbar = figG.colorbar( cax ) cbar.ax.set_ylabel(r'$\gamma$', fontsize=14) saveas_png(figG, savepath, 'gamma') plt.close(figG)
def plot_xterms(slm, savepath, time = None, h5plot=True, t_is_z=True): dzeta = abs(slm.get_zeta_array()[1] - slm.get_zeta_array()[0]); curr = slm.get_charge() / dzeta if t_is_z: xlabel_str = r'$k_p z$' else: xlabel_str = r'$\omega_p t$' if time == None: tidx = [0,-1]; else: tidx = [(np.abs(slm.get_time_array() - time)).argmin()] for i in tidx: fig = plt.figure() plt.plot(slm.get_zeta_array(), slm.get(x=1,y=1)[i,:]) ax = plt.gca() ax.set_xlabel(r'$k_p \zeta$', fontsize=14) ax.set_ylabel(r'$<xy>$', fontsize=14) saveas_eps_pdf(fig, savepath, ('xy_%0.1f' % (slm.get_time_array()[i])), h5plot=h5plot) plt.close(fig)
def plot_save_slice_quad_corr_lines(slm, savepath, time = None, axdir='x', h5plot=True): if time == None: tidx = [0,-1]; else: tidx = [(np.abs(slm.get_time_array() - time)).argmin()] for i in tidx: if axdir == 'y': var_xy = slm.get(y=2)[i,:] var_pxy = slm.get(py=2)[i,:] nonzero_idx = np.nonzero(np.multiply(var_xy, var_pxy))[0] quad_corr_xy = np.divide( slm.get(y=2,py=2)[i,nonzero_idx],\ np.multiply(var_xy[nonzero_idx],var_pxy[nonzero_idx]) ) - 1.0 quad_corr_xy_lab = r'$\left \langle y^2 p_y^2\right\rangle/(\sigma_y^2\sigma_{p_y}^2) - 1$' quad_corr_xy_savename = 'quad_corr_y' else: # axdir == 'x' or axdir == 'r' var_xy = slm.get(x=2)[i,:] var_pxy = slm.get(px=2)[i,:] nonzero_idx = np.nonzero(np.multiply(var_xy, var_pxy))[0] quad_corr_xy = np.divide( slm.get(x=2,px=2)[i,nonzero_idx],\ np.multiply(var_xy[nonzero_idx],var_pxy[nonzero_idx]) ) - 1.0 quad_corr_xy_lab = r'$\left \langle x^2 p_x^2\right\rangle/(\sigma_x^2\sigma_{p_x}^2) - 1$' quad_corr_xy_savename = 'quad_corr_x' fig_quad_corr_xy = plt.figure() plt.plot( slm.get_zeta_array()[nonzero_idx], quad_corr_xy) ax = plt.gca() ax.set_xlabel(r'$k_p \zeta$', fontsize=14) ax.set_ylabel(quad_corr_xy_lab, fontsize=14) saveas_eps_pdf(fig_quad_corr_xy, savepath, ('%s_time_%0.1f' % (quad_corr_xy_savename, slm.get_time_array()[i])), h5plot=h5plot) plt.close(fig_quad_corr_xy)
def adiabatic_matching_plots(savepath, z_array, emittance, gamma): log_rel_emittance = np.log(emittance/emittance[0]) kbeta0 = np.power(2*gamma,-1/2) dz_per_dgamma = np.divide(np.gradient(z_array),np.gradient(gamma)) dlog_rel_emittance_per_dgamma = np.divide(np.gradient(log_rel_emittance),np.gradient(gamma)) eta = np.multiply( np.divide(dz_per_dgamma,kbeta0), dlog_rel_emittance_per_dgamma) fig_eta = plt.figure() plt.plot(gamma, eta) ax = plt.gca() ax.set_xlabel(r'$\overline{\gamma}$', fontsize=14) ax.set_ylabel(r'$\eta$', fontsize=14) #ylim = ax.get_ylim() ax.set_ylim([0, 1e-2]) if magn_check(eta): ax.yaxis.set_major_formatter(FormatStrFormatter('%.1e')) plt.gcf().subplots_adjust(left=0.18) else: plt.gcf().subplots_adjust(left=0.15) saveas_eps_pdf(fig_eta, savepath, 'eta', h5plot=True) plt.close(fig_eta)
def density_inversion(r_nb, nb, r_psi, psi, zeta_pos, savepath='./plots/equilib', ifplot=True): psi_norm = psi - psi[0] psi_interp = np.interp(r_nb, r_psi, psi_norm) dnb = np.diff(nb) dpsi = np.diff(psi_interp) x = psi_interp[1:] F = -1 * np.divide(dnb, dpsi) / (2 * math.pi) popt, pcov = curve_fit(exp_func, x, F) if ifplot: mkdirs_if_nexist(savepath) fig = plt.figure() plt.plot(r_nb, nb) ax = plt.gca() ax.set_xlabel(r'$k_p r$', fontsize=14) ax.set_ylabel(r'$n_b/n_0$', fontsize=14) saveas_eps_pdf(fig, savepath, 'nb_zeta_%0.2f' % zeta_pos) plt.close(fig) fig = plt.figure() plt.plot(r_nb[:-1], np.divide(dnb, np.diff(r_nb))) ax = plt.gca() ax.set_xlabel(r'$k_p r$', fontsize=14) ax.set_ylabel(r'$d(n_b/n_0)/(k_p dr)$', fontsize=14) saveas_eps_pdf(fig, savepath, 'dnb_dr_zeta_%0.2f' % zeta_pos) plt.close(fig) fig = plt.figure() plt.plot(r_nb[:-1], np.divide(dpsi, np.diff(r_nb))) ax = plt.gca() ax.set_xlabel(r'$k_p r$', fontsize=14) ax.set_ylabel(r'$d(e\psi/mc^2)/(k_p dr)$', fontsize=14) saveas_eps_pdf(fig, savepath, 'dpsi_dr_zeta_%0.2f' % zeta_pos) plt.close(fig) fig = plt.figure() plt.plot(psi_norm, r_nb) ax = plt.gca() ax.set_ylabel(r'$k_p r$', fontsize=14) ax.set_xlabel(r'$e\psi/mc^2$', fontsize=14) saveas_eps_pdf(fig, savepath, 'r_vs_psi_zeta_%0.2f' % zeta_pos) plt.close(fig) fig = plt.figure() plt.plot(r_nb, psi_norm) ax = plt.gca() ax.set_xlabel(r'$k_p r$', fontsize=14) ax.set_ylabel(r'$e\psi/mc^2$', fontsize=14) saveas_eps_pdf(fig, savepath, 'psi_zeta_%0.2f' % zeta_pos) plt.close(fig) fig = plt.figure() plt.plot(x, F, '-', label=r'F(x)') plt.plot(x, exp_func(x, *popt), '--', label=r'fit: $%0.2f \cdot exp(-%0.2f \cdot x)$' % tuple(popt)) ax = plt.gca() ax.set_xlabel(r'$x$', fontsize=14) ax.set_ylabel(r'$F$', fontsize=14) plt.legend(frameon=False) saveas_eps_pdf(fig, savepath, 'F_zeta_%0.2f' % zeta_pos) plt.close(fig) return x, F
def cmp_plot_Wr(args, g3d, plasma, beam): if args.zeta_pos == None: zeta_pos_list = [0.0] else: zeta_pos_list = args.zeta_pos x_array = g3d.get_x_arr(1) Nsimlines = len(zeta_pos_list) Nzeta = g3d.get_nx(1) Wr_sim = np.zeros((Nzeta, Nsimlines)) Wr_theo = np.zeros((Nzeta, Nsimlines)) for i in range(0, Nsimlines): zeta_pos = zeta_pos_list[i] Wr_sim[:, i] = g3d.read(x0=zeta_pos, x2=0.0) Wr_theo[:, i] = lin_Wr_theo(plasma=plasma, beam=beam, x_array=x_array, zeta_pos=zeta_pos) cmap = cm.tab10 fig = plt.figure(num=None, figsize=(9, 7), dpi=80, facecolor='w', edgecolor='k') for i in range(0, Nsimlines): zeta_pos = zeta_pos_list[i] if args.rel_to_hom: Wr_sim_vals = Wr_sim[:, i] - x_array * 0.5 Wr_theo_vals = Wr_theo[:, i] - x_array * 0.5 ylab = r'$W_r/E_0 - r/2$' else: Wr_sim_vals = Wr_sim[:, i] Wr_theo_vals = Wr_theo[:, i] ylab = r'$W_r/E_0$' label_sim = r'PIC: $k_p \zeta = %0.1f$' % zeta_pos ax_sim = plt.plot(x_array, Wr_sim_vals, linestyle='-', color=cmap(i), label=label_sim) label_theo = r'Theo: $k_p \zeta = %0.1f$' % zeta_pos ax_theo = plt.plot(x_array, Wr_theo_vals, linestyle='--', color=cmap(i), label=label_theo) if not args.rel_to_hom: label_half = r'$k_p r/2$' ax_half = plt.plot(x_array, x_array * 0.5, linestyle='-', label=label_half, color=[0.5, 0.5, 0.5]) ax = plt.gca() handles, labels = ax.get_legend_handles_labels() plt.legend(flip(handles, 2), flip(labels, 2), ncol=2) if args.rmax != None: xmax = args.rmax ymax = 1.2 * np.amax(Wr_sim_vals[np.logical_and( x_array < xmax, x_array > 0.0)]) else: max_idx = np.where(Wr_sim_vals == np.amax(Wr_sim_vals))[0][0] xmax = x_array[max_idx] ymax = 1.2 * np.amax(Wr_sim_vals) ax.set_xlim([0, xmax]) ax.set_ylim([0, ymax]) ax.set_ylabel(ylab, fontsize=14) ax.set_xlabel(r'$k_p r$', fontsize=14) if not (-3.0 < math.log(np.max(abs(Wr_sim)), 10) < 3.0): ax.yaxis.set_major_formatter(FormatStrFormatter('%.1e')) plt.gcf().subplots_adjust(left=0.18) else: plt.gcf().subplots_adjust(left=0.15) filesuffix = '_t_%06.f' % (np.floor(g3d.get_time())) fileprefix = 'Wr_' + 'zeta_%0.1f' % zeta_pos saveformat = 'pdf' savepath = args.savepath savename = fileprefix + filesuffix + '.' + saveformat saveas_eps_pdf(fig, savepath, savename, h5plot=args.h5plot, fformat=saveformat) plt.close(fig)
def plot_save_slice_ene_spread(slm, savepath, h5plot=True, time=None, zeta_pos=None, t_is_z=True): if t_is_z: xlabel_str = r'$k_p z$' else: xlabel_str = r'$\omega_p t$' sigma_gamma_per_gamma = np.divide( np.sqrt( slm.get(pz=2) ),slm.get(pz=1)) if time == None: time = slm.get_time_array()[-1] tidx = [-1]; else: tidx = [(np.abs(slm.get_time_array() - time)).argmin()] if len(tidx) == 1: time = [time] for i in range(0,len(tidx)): sig_gam_per_gam_savename = 'sigma_gamma_per_gamma_time_%0.2f' % time[i] sig_gam_per_gam_ylab = r'$\sigma_{\gamma}/\gamma(t=%0.1f)$' % time[i] fig_sig_gam_per_gam = plt.figure() plt.plot(slm.get_zeta_array(), sigma_gamma_per_gamma[tidx[i],:] ) #/ # np.mean(np.sqrt(1 + np.power(slm.get(px=1)[tidx[i],:],2) # + np.power(slm.get(py=1)[tidx[i],:],2) # + np.power(slm.get(pz=1)[tidx[i],:],2) )) ) ax = plt.gca() ymin, ymax = ax.get_ylim() if ymin > 0 and ymax > 0: plt.ylim(0, ymax*1.2) elif ymin < 0 and ymax < 0: plt.ylim(ymin*1.2, 0) ax.set_xlabel(r'$k_p \zeta$', fontsize=14) ax.set_ylabel(sig_gam_per_gam_ylab, fontsize=14) saveas_eps_pdf(fig_sig_gam_per_gam, savepath, sig_gam_per_gam_savename, h5plot=h5plot) plt.close(fig_sig_gam_per_gam) figsigma_gamma_per_gamma = plt.figure() cax = plt.pcolormesh( slm.get_zeta_array(), slm.get_time_array(), sigma_gamma_per_gamma, cmap=cm.GnBu, vmin=np.amin(abs(sigma_gamma_per_gamma)), vmax=np.amax(abs(sigma_gamma_per_gamma))) cbar = figsigma_gamma_per_gamma.colorbar(cax) cbar.ax.set_ylabel(r'$\sigma_{\gamma} / \gamma$', fontsize=14) ax = plt.gca() ax.set_xlabel(r'$k_p \zeta$', fontsize=14) ax.set_ylabel(xlabel_str, fontsize=14) saveas_png(figsigma_gamma_per_gamma, savepath, 'sigma_gamma_per_gamma') plt.close(figsigma_gamma_per_gamma) figXb0 = plt.figure() plt.plot(slm.get_zeta_array(), sigma_gamma_per_gamma[0,:]) ax = plt.gca() ymin, ymax = ax.get_ylim() if ymin > 0 and ymax > 0: plt.ylim(0, ymax*1.2) elif ymin < 0 and ymax < 0: plt.ylim(ymin*1.2, 0) ax.set_xlabel(r'$k_p \zeta$', fontsize=14) ax.set_ylabel(r'$\sigma_{\gamma} / \gamma$', fontsize=14) saveas_eps_pdf(figXb0, savepath, 'sigma_gamma_per_gamma_time_0', h5plot=h5plot) plt.close(figXb0) if zeta_pos != None: zetaidx = [(np.abs(slm.get_zeta_array() - zeta_pos)).argmin()] sg_per_g_zeta_savename = 'sigma_gamma_per_gamma_zeta_%0.2f' % zeta_pos # figsigma_gamma_per_gamma_zetapos = plt.figure() plt.plot(slm.get_time_array(), sigma_gamma_per_gamma[:,zetaidx]) ax = plt.gca() ax.set_xlabel(xlabel_str, fontsize=14) ax.set_ylabel(r'$\sigma_{\gamma} / \gamma$', fontsize=14) if magn_check(sigma_gamma_per_gamma[:,zetaidx]): ax.yaxis.set_major_formatter(FormatStrFormatter('%.1e')) plt.gcf().subplots_adjust(left=0.18) else: plt.gcf().subplots_adjust(left=0.15) saveas_eps_pdf(figsigma_gamma_per_gamma_zetapos, savepath, sg_per_g_zeta_savename, h5plot=h5plot) plt.close(figsigma_gamma_per_gamma_zetapos)
def plot_save_slice_centroids(slm, savepath, h5plot=True, time=None, zeta_pos=None, t_is_z=True): if t_is_z: xlabel_str = r'$k_p z$' else: xlabel_str = r'$\omega_p t$' Xb0_for_norm = np.ones(slm.get(x=1)[0,:].shape) Yb0_for_norm = np.ones(slm.get(y=1)[0,:].shape) for i in range(0,len(slm.get_zeta_array())): if (slm.get(x=1)[0,i] != 0.0): Xb0_for_norm[i] = slm.get(x=1)[0,i] Yb0_for_norm[i] = slm.get(y=1)[0,i] if time == None: time = slm.get_time_array()[-1] tidx = [-1]; else: tidx = [(np.abs(slm.get_time_array() - time)).argmin()] if len(tidx) == 1: time = [time] for i in range(0,len(tidx)): Xb_lout_savename = 'Xb_time_%0.2f' % time[i] Yb_lout_savename = 'Yb_time_%0.2f' % time[i] Xb_lout_ylab = r'$k_p X_{b}(t=%0.1f)$' % time[i] Yb_lout_ylab = r'$k_p Y_{b}(t=%0.1f)$' % time[i] figXb_lout = plt.figure() plt.plot(slm.get_zeta_array(), slm.get(x=1)[tidx[i],:]) ax = plt.gca() ymin, ymax = ax.get_ylim() if ymin > 0 and ymax > 0: plt.ylim(0, ymax*1.2) elif ymin < 0 and ymax < 0: plt.ylim(ymin*1.2, 0) ax.set_xlabel(r'$k_p \zeta$', fontsize=14) ax.set_ylabel(Xb_lout_ylab, fontsize=14) saveas_eps_pdf(figXb_lout, savepath, Xb_lout_savename, h5plot=h5plot) plt.close(figXb_lout) figYb_lout = plt.figure() plt.plot(slm.get_zeta_array(), slm.get(y=1)[tidx[i],:]) ax = plt.gca() ymin, ymax = ax.get_ylim() if ymin > 0 and ymax > 0: plt.ylim(0, ymax*1.2) elif ymin < 0 and ymax < 0: plt.ylim(ymin*1.2, 0) ax.set_xlabel(r'$k_p \zeta$', fontsize=14) ax.set_ylabel(Yb_lout_ylab, fontsize=14) saveas_eps_pdf(figYb_lout, savepath, Yb_lout_savename, h5plot=h5plot) plt.close(figYb_lout) Xb_norm = np.zeros( slm.get(x=1).shape ) Yb_norm = np.zeros( slm.get(y=1).shape ) # zeta_hseed = np.min(slm.get_zeta_array()) # idx_hseed = (np.abs(slm.get_zeta_array()-zeta_hseed)).argmin() for i in range(0,len(slm.get_zeta_array())): # if (slm.get_zeta_array()[i] <= zeta_hseed): Xb_norm[:,i] = np.absolute( slm.get(x=1)[:,i]/Xb0_for_norm[i] ) Yb_norm[:,i] = np.absolute( slm.get(y=1)[:,i]/Yb0_for_norm[i] ) figXb = plt.figure() cax = plt.pcolormesh( slm.get_zeta_array(), slm.get_time_array(), slm.get(x=1), cmap=cm.PuOr, vmin=-np.amax(abs(slm.get(x=1))), vmax=np.amax(abs(slm.get(x=1)))) cbar = figXb.colorbar(cax) cbar.ax.set_ylabel('$k_p X_b$', fontsize=14) ax = plt.gca() ax.set_xlabel(r'$k_p \zeta$', fontsize=14) ax.set_ylabel(xlabel_str, fontsize=14) saveas_png(figXb, savepath, 'Xb') plt.close(figXb) figYb = plt.figure() cax = plt.pcolormesh( slm.get_zeta_array(), slm.get_time_array(), slm.get(y=1), cmap=cm.PuOr, vmin=-np.amax(abs(slm.get(y=1))), vmax=np.amax(abs(slm.get(y=1)))) cbar = figYb.colorbar(cax) cbar.ax.set_ylabel('$k_p Y_b$', fontsize=14) ax = plt.gca() ax.set_xlabel(r'$k_p \zeta$', fontsize=14) ax.set_ylabel(xlabel_str, fontsize=14) saveas_png(figYb, savepath, 'Yb') plt.close(figYb) figXbnorm = plt.figure() cax = plt.pcolormesh( slm.get_zeta_array(), slm.get_time_array(), Xb_norm, cmap=cm.Blues, vmin=0, vmax=np.amax(Xb_norm)) cbar = figXbnorm.colorbar(cax) cbar.ax.set_ylabel('$|X_b/X_{b,0}|$', fontsize=14) ax = plt.gca() ax.set_xlabel(r'$k_p \zeta$', fontsize=14) ax.set_ylabel(xlabel_str, fontsize=14) saveas_png(figXbnorm, savepath, 'Xb_rel') plt.close(figXbnorm) figYbnorm = plt.figure() cax = plt.pcolormesh( slm.get_zeta_array(), slm.get_time_array(), Yb_norm, cmap=cm.Blues, vmin=0, vmax=np.amax(Yb_norm)) cbar = figYbnorm.colorbar(cax) cbar.ax.set_ylabel('$|Y_b/Y_{b,0}|$', fontsize=14) ax = plt.gca() ax.set_xlabel(r'$k_p \zeta$', fontsize=14) ax.set_ylabel(xlabel_str, fontsize=14) saveas_png(figYbnorm, savepath, 'Yb_rel') plt.close(figYbnorm) figXb0 = plt.figure() plt.plot(slm.get_zeta_array(), slm.get(x=1)[0,:]) ax = plt.gca() ymin, ymax = ax.get_ylim() if ymin > 0 and ymax > 0: plt.ylim(0, ymax*1.2) elif ymin < 0 and ymax < 0: plt.ylim(ymin*1.2, 0) ax.set_xlabel(r'$k_p \zeta$', fontsize=14) ax.set_ylabel(r'$k_p X_{b,0}$', fontsize=14) saveas_eps_pdf(figXb0, savepath, 'Xb0', h5plot=h5plot) plt.close(figXb0) figYb0 = plt.figure() plt.plot(slm.get_zeta_array(), slm.get(y=1)[0,:]) ax = plt.gca() ymin, ymax = ax.get_ylim() if ymin > 0 and ymax > 0: plt.ylim(0, ymax*1.2) elif ymin < 0 and ymax < 0: plt.ylim(ymin*1.2, 0) ax.set_xlabel(r'$k_p \zeta$', fontsize=14) ax.set_ylabel(r'$k_p Y_{b,0}$', fontsize=14) saveas_eps_pdf(figYb0, savepath, 'Yb0', h5plot=h5plot) plt.close(figYb0) if zeta_pos != None: zetaidx = [(np.abs(slm.get_zeta_array() - zeta_pos)).argmin()] Xb_savename = 'Xb_zeta_%0.2f' % zeta_pos Yb_savename = 'Yb_zeta_%0.2f' % zeta_pos else: zetaidx = 0 Xb_savename = 'Xb_tail' Yb_savename = 'Yb_tail' figXbtail = plt.figure() plt.plot(slm.get_time_array(), slm.get(x=1)[:,zetaidx]) ax = plt.gca() ax.set_xlabel(xlabel_str, fontsize=14) ax.set_ylabel(r'$k_p X_{b,\mathrm{tail}}$', fontsize=14) if magn_check(slm.get(x=1)[:,zetaidx]): ax.yaxis.set_major_formatter(FormatStrFormatter('%.1e')) plt.gcf().subplots_adjust(left=0.18) else: plt.gcf().subplots_adjust(left=0.15) saveas_eps_pdf(figXbtail, savepath, Xb_savename, h5plot=h5plot) plt.close(figXbtail) figYbtail = plt.figure() plt.plot(slm.get_time_array(), slm.get(y=1)[:,zetaidx]) ax = plt.gca() ax.set_xlabel(xlabel_str, fontsize=14) ax.set_ylabel(r'$k_p Y_{b,\mathrm{tail}}$', fontsize=14) if magn_check(slm.get(y=1)[:,zetaidx]): ax.yaxis.set_major_formatter(FormatStrFormatter('%.1e')) plt.gcf().subplots_adjust(left=0.18) else: plt.gcf().subplots_adjust(left=0.15) saveas_eps_pdf(figYbtail, savepath, Yb_savename, h5plot=h5plot) plt.close(figYbtail)
def plot_save_proj_rms(slm, savepath, axdir='x', h5plot=True, verbose=True, t_is_z=True, versus_gamma=False): tot_charge = np.sum(slm.get_charge(), axis=1) if axdir == 'x': xsq = slm.get_proj(x=2) psq = slm.get_proj(px=2) xp = slm.get_proj(x=1,px=1) emittance_all_slices = np.sqrt( np.multiply(slm.get(x=2), slm.get(px=2)) - np.power(slm.get(x=1,px=1),2) ) sigma_xyr_lab = r'$k_p\sigma_x$' sigma_xyr_savename = 'sigma_x_proj' sigma_pxyr_lab = r'$\sigma_{p_x}/m_e c$' sigma_pxyr_savename = 'sigma_px_proj' corr_xy_lab = r'$ k_p \left \langle x\,p_x \right\rangle/m_e c$' corr_xy_savename = 'xpx_proj' epsilon_xy_proj_lab = r'$k_p \epsilon_x$' epsilon_xy_proj_savename = 'emittance_x_proj' epsilon_xy_sl_lab = r'$k_p \epsilon_{x,\mathrm{sliced}}$' epsilon_xy_sl_savename = 'emittance_x_sliced' elif axdir == 'y': xsq = slm.get_proj(y=2) psq = slm.get_proj(py=2) xp = slm.get_proj(y=1,py=1) emittance_all_slices = np.sqrt( np.multiply(slm.get(y=2), slm.get(py=2)) - np.power(slm.get(y=1,py=1),2) ) sigma_xyr_lab = r'$k_p\sigma_y$' sigma_xyr_savename = 'sigma_y_proj' sigma_pxyr_lab = r'$\sigma_{p_y}/m_e c$' sigma_pxyr_savename = 'sigma_py_proj' corr_xy_lab = r'$ k_p \left \langle y\,p_y \right\rangle/m_e c$' corr_xy_savename = 'ypy_proj' epsilon_xy_proj_lab = r'$k_p \epsilon_y$' epsilon_xy_proj_savename = 'emittance_y_proj' epsilon_xy_sl_lab = r'$k_p \epsilon_{y,\mathrm{sliced}}$' epsilon_xy_sl_savename = 'emittance_y_sliced' elif axdir == 'r': xavg = slm.project(slm.get(x=1) + slm.get(y=1)) pavg = slm.project(slm.get(px=1) + slm.get(py=1)) xsq_noncentral = slm.project(slm.get(x=2) + slm.get(y=2) \ + np.power(slm.get(x=1),2) + np.power(slm.get(y=1),2)) psq_noncentral = slm.project(slm.get(px=2) \ + slm.get(py=2) + np.power(slm.get(px=1),2) + np.power(slm.get(py=1),2)) xp_noncentral = slm.project(slm.get(x=1,px=1) + slm.get(y=1,py=1) \ + np.multiply(slm.get(x=1),slm.get(px=1)) + np.multiply(slm.get(y=1),slm.get(py=1))) xsq = xsq_noncentral - np.pow(xavg,2) psq = psq_noncentral - np.pow(pavg,2) xp = xp_noncentral = np.multiply(xavg,pavg) emittance_all_slices = np.sqrt( np.multiply(slm.get(x=2), slm.get(px=2)) - np.power(slm.get(x=1,px=1),2) + \ np.multiply(slm.get(y=2), slm.get(py=2)) - np.power(slm.get(y=1,py=1),2) ) sigma_xyr_lab = r'$k_p\sigma_r$' sigma_xyr_savename = 'sigma_r_proj' sigma_pxyr_lab = r'$(\sigma_{p_x}^2 + \sigma_{p_y}^2)^{1/2}/m_e c$' sigma_pxyr_savename = 'sigma_pr_proj' corr_xy_lab = r'$ k_p (\left \langle x\,p_x \right\rangle + \left \langle y\,p_y \right\rangle)/m_e c$' corr_xy_savename = 'xpx_ypy_proj' epsilon_xy_proj_lab = r'$k_p \epsilon_r$' epsilon_xy_proj_savename = 'emittance_r_proj' epsilon_xy_sl_lab = r'$k_p \epsilon_{r,\mathrm{sliced}}$' epsilon_xy_sl_savename = 'emittance_r_sliced' # Calculate emittance emittance_sliced = slm.project(emittance_all_slices) emittance_proj = np.sqrt(np.multiply(xsq,psq)-np.power(xp,2)) if versus_gamma: vs_gamma_append_str = '_vs_gamma' x_axis_array = slm.get_proj(pz=1) xlabel_str = r'$\overline{\gamma}$' sigma_xyr_savename += vs_gamma_append_str sigma_pxyr_savename += vs_gamma_append_str corr_xy_savename += vs_gamma_append_str epsilon_xy_proj_savename += vs_gamma_append_str epsilon_xy_sl_savename += vs_gamma_append_str adiabatic_matching_plots(savepath, slm.get_time_array(), emittance_proj, slm.get_proj(pz=1)) else: x_axis_array = slm.get_time_array() if t_is_z: xlabel_str = r'$k_p z$' else: xlabel_str = r'$\omega_p t$' fig_sx = plt.figure() plt.plot(x_axis_array, np.sqrt(xsq)) ax = plt.gca() ax.set_xlabel(xlabel_str, fontsize=14) ax.set_ylabel(sigma_xyr_lab, fontsize=14) if magn_check(np.sqrt(xsq)): ax.yaxis.set_major_formatter(FormatStrFormatter('%.1e')) plt.gcf().subplots_adjust(left=0.18) else: plt.gcf().subplots_adjust(left=0.15) saveas_eps_pdf(fig_sx, savepath, sigma_xyr_savename, h5plot=h5plot) plt.close(fig_sx) fig_sp = plt.figure() plt.plot(x_axis_array, np.sqrt(psq)) ax = plt.gca() ax.set_xlabel(xlabel_str, fontsize=14) ax.set_ylabel(sigma_pxyr_lab, fontsize=14) if magn_check(np.sqrt(psq)): ax.yaxis.set_major_formatter(FormatStrFormatter('%.1e')) plt.gcf().subplots_adjust(left=0.18) else: plt.gcf().subplots_adjust(left=0.15) saveas_eps_pdf(fig_sp, savepath, sigma_pxyr_savename, h5plot=h5plot) plt.close(fig_sp) fig_xp = plt.figure() plt.plot(x_axis_array, xp) ax = plt.gca() ax.set_xlabel(xlabel_str, fontsize=14) ax.set_ylabel(corr_xy_lab, fontsize=14) if magn_check(xp): ax.yaxis.set_major_formatter(FormatStrFormatter('%.1e')) plt.gcf().subplots_adjust(left=0.18) else: plt.gcf().subplots_adjust(left=0.15) saveas_eps_pdf(fig_xp, savepath, corr_xy_savename, h5plot=h5plot) plt.close(fig_xp) fig_e = plt.figure() plt.plot(x_axis_array, emittance_proj, label='projected emittance') ax = plt.gca() ax.set_xlabel(xlabel_str, fontsize=14) ax.set_ylabel(epsilon_xy_proj_lab, fontsize=14) if magn_check(emittance_proj): ax.yaxis.set_major_formatter(FormatStrFormatter('%.1e')) plt.gcf().subplots_adjust(left=0.18) else: plt.gcf().subplots_adjust(left=0.15) saveas_eps_pdf(fig_e, savepath, epsilon_xy_proj_savename, h5plot=h5plot) plt.close(fig_e) fig_esl = plt.figure() plt.plot(x_axis_array, emittance_sliced, label='sliced emittance') ax = plt.gca() ax.set_xlabel(xlabel_str, fontsize=14) ax.set_ylabel(epsilon_xy_sl_lab, fontsize=14) if magn_check(emittance_sliced): ax.yaxis.set_major_formatter(FormatStrFormatter('%.1e')) plt.gcf().subplots_adjust(left=0.18) else: plt.gcf().subplots_adjust(left=0.15) saveas_eps_pdf(fig_esl, savepath, epsilon_xy_sl_savename, h5plot=h5plot) plt.close(fig_esl) fig_e_slpr = plt.figure() epp = plt.plot(x_axis_array, emittance_proj, label='projected') esp = plt.plot(x_axis_array, emittance_sliced, color = epp[0].get_color(), linestyle='--', label='sliced') ax = plt.gca() ax.set_xlabel(xlabel_str, fontsize=14) ax.set_ylabel(epsilon_xy_proj_lab, fontsize=14) ax.legend() if magn_check(emittance_proj): ax.yaxis.set_major_formatter(FormatStrFormatter('%.1e')) plt.gcf().subplots_adjust(left=0.18) else: plt.gcf().subplots_adjust(left=0.15) saveas_eps_pdf(fig_e_slpr, savepath, epsilon_xy_proj_savename + '_sl', h5plot=h5plot) plt.close(fig_e_slpr)
def main(): parser = h5plot_parser() args = parser.parse_args() if args.maketitle: ''' Calculating plasma frequency for length ''' ELECTRON_CHARGE_IN_COUL = 1.60217657e-19 ELECTRON_MASS_IN_KG = 9.10938291e-31 VAC_PERMIT_FARAD_PER_M = 8.854187817e-12 SPEED_OF_LIGHT_IN_M_PER_S = 299792458.0 omega_p = np.sqrt( args.n0 * (ELECTRON_CHARGE_IN_COUL**2)/ (VAC_PERMIT_FARAD_PER_M * ELECTRON_MASS_IN_KG)); skin_depth = SPEED_OF_LIGHT_IN_M_PER_S/omega_p if len(sys.argv)==1: parser.print_help() sys.exit(1) saveformat = args.file_format h5lp = [] linestyles = ['-', '--', '-.', ':'] save_append_str = '' type_str = 'comp' if args.latexon: plt.rc('text', usetex=True) plt.rc('font', family='serif') j = 0 i = 0 for path in args.paths: fig = plt.figure(figsize=(7,5)) if not os.path.isfile(path): print('ERROR: File %s does not exist!' % path) sys.exit(1) h5lp.append(H5Plot()) h5lp[i].read(path) timestamp = path.split("_")[-1].split(".h5")[0] if args.data2: for files in args.data2: if timestamp in files or args.manual: print('Reading second data set...') h5secondlp = H5Plot() h5secondlp.read(files) if args.cno == None: argcolor = plt.cm.tab20(6) else: argcolor = plt.cm.tab20(args.cno[0]) j = 0 for (x, y, label, linestyle, color) in h5lp[i].get_line_plots(): if args.labels != None: save_append_str += '_' + args.labels[i] label = args.labels[i] if label == '_line0': label = path if args.latexon: if label[0] != '$' or label[-1] != '$': label = r'$\textrm{' + label + '}$' if args.abs: y = np.abs(y) if args.lstyle == None: linestyle_code = linestyles[0] else: linestyle_code = linestyles[args.lstyle[0]] if args.absylog: plt.semilogy( x, y, label=label, linestyle=linestyle_code, color=argcolor, zorder=20) else: #plt.plot(x, y, label=label, linestyle=linestyle, color=color) plt.plot(x, y, label=label, linestyle=linestyle_code, color=argcolor, zorder=20) ax = plt.gca() ax.tick_params('y', colors=argcolor) plt.gcf().subplots_adjust(left=0.15, bottom=0.15, right=1-0.15) if not (-3.0 < math.log(np.max(abs(y)),10) < 3.0): ax.yaxis.set_major_formatter(FormatStrFormatter('%.1e')) plt.gcf().subplots_adjust(left=0.18) ax2 = ax.twinx() ''' PLOTTING THE SECOND PLOT ''' if args.cno == None: argcolor = plt.cm.tab20(1) else: argcolor = plt.cm.tab20(args.cno[1]) ax2.tick_params('y', colors=argcolor) j = 0 for (x2, y2, label, linestyle, color) in h5secondlp.get_line_plots(): if args.labels != None: save_append_str += '_' + args.labels[i] label = args.labels[i] if label == '_line0': label = path if args.latexon: if label[0] != '$' or label[-1] != '$': label = r'$\textrm{' + label + '}$' if args.abs: y2 = np.abs(y2) if args.lstyle == None: linestyle_code = linestyles[0] else: linestyle_code = linestyles[args.lstyle[1]] if args.absy2log: plt.semilogy( x2, y2, label=label, linestyle=linestyle_code, color=argcolor, zorder=0) else: #plt.plot(x, y, label=label, linestyle=linestyle, color=color) plt.plot(x2, y2, label=label, linestyle=linestyle_code, color=argcolor, zorder=0) if args.xlim != None: ax.set_xlim(args.xlim[0],args.xlim[1]) if args.ylim != None: ax.set_ylim(args.ylim[0],args.ylim[1]) if args.y2lim != None: ax2.set_ylim(args.y2lim[0],args.y2lim[1]) if args.xlab != None: xlab = args.xlab else: xlab = h5lp[0].get_xlab() if args.ylab != None: ylab = args.ylab else: ylab = [] ylab.append(h5lp[0].get_ylab()) ylab.append(h5secondlp.get_ylab()) if args.latexon: if xlab[0] != '$' or xlab[-1] != '$': xlab = r'$' + xlab + '$' if args.latexon: if ylab[0,0] != '$' or ylab[0,-1] != '$': ylab[0] = r'$' + ylab[0] + '$' ylab[1] = r'$' + ylab[1] + '$' ax.set_xlabel(xlab, fontsize=14) ax.set_ylabel(ylab[0], fontsize=14) ax2.set_ylabel(ylab[1], fontsize=14) ax.set_zorder(ax2.get_zorder()+1) # put ax in front of ax2 ax.patch.set_visible(False) # hide the 'canvas' # handles, labels = ax.get_legend_handles_labels() # #plt.legend(flip(handles, 2), flip(labels, 2), ncol=2) # plt.legend(frameon=False) #plt.gcf().subplots_adjust(left=0.15, bottom=0.15, right=0.15) if not (-3.0 < math.log(np.max(abs(y2)),10) < 3.0): ax2.yaxis.set_major_formatter(FormatStrFormatter('%.1e')) plt.gcf().subplots_adjust(right=1-0.18) #plt.show() ''' make the title ''' if args.maketitle: time = np.around(float(timestamp)) distance = time*skin_depth*1e2 title = 'Time: ' + str(time) + r'$\,\omega_p^{-1}$ ' + ' Distance: ' + str(np.around(distance,2)) + r'$\,$cm' plt.title(title) if args.show: plt.show() spath, fname = os.path.split(args.paths[0]) if args.savepath != None: spath = args.savepath if not args.manual: save_name = type_str + '_' + fname + save_append_str + '_' + timestamp else: save_name = type_str + '_' + fname + save_append_str if saveformat==args.file_format: saveas_png(fig, spath, save_name, args.dpi) else: saveas_eps_pdf(fig, savepath=spath, savename=save_name) plt.close(fig)