def plot_field( field, filename="C1.h5", points=200, slice=0, range=None, xrange=None, yrange=None, rrange=None, zrange=None, palette=None, lcfs=None, bound=None, linfac=1.0, phi=0.0, iabs=None, iphase=None, isum=None, iavg=None, idiff=None, ilinear=None, iequil=None, icomplex=None, ntor=None, title=None, fs=1.0, ax=None, ): if isinstance(field, basestring): # Read this field if title is None: title = field field = read_field( field, slice=slice, filename=filename, phi=phi, points=points, rrange=rrange, zrange=zrange, linfac=linfac, iabs=iabs, iphase=iphase, isum=isum, iavg=iavg, idiff=idiff, ilinear=ilinear, iequil=iequil, icomplex=icomplex, ntor=ntor, ) data = np.real(field.data) if range is None: vmin = data.min() vmax = data.max() else: vmin = range[0] vmax = range[1] if palette is None: if vmin >= 0.0: palette = "inferno" col_lcfs = "w" else: palette = "RdBu_r" col_lcfs = "k" else: col_lcfs = "w" cmap = mpl.colors.ListedColormap(sns.color_palette(palette, 256)) extent = [field.R.data[0], field.R.data[-1], field.Z.data[0], field.Z.data[-1]] if ax is None: f, ax = plt.subplots(figsize=[fs * 8, fs * 12]) else: f = None im = ax.imshow(data.T, origin="lower", vmin=vmin, vmax=vmax, extent=extent, cmap=cmap) if xrange is None: xrange = [field.R.data[0], field.R.data[-1]] if yrange is None: yrange = [field.Z.data[0], field.Z.data[-1]] ax.set_xlim(xrange) ax.set_xlabel(r"$R$ (m)", fontsize=fs * 28) ax.set_ylim(yrange) ax.set_ylabel(r"$Z$ (m)", fontsize=fs * 28) if title is not None: ax.set_title(title, fontsize=fs * 32) ax.tick_params(labelsize=fs * 24) if f is not None: cb = f.colorbar(im, format="%1.3g") cb.ax.tick_params(labelsize=fs * 24) else: div = make_axes_locatable(ax) cax = div.append_axes("right", size="10%", pad=0.02) cb = plt.colorbar(im, cax=cax, format="%1.3g") cb.ax.tick_params(labelsize=fs * 24) if lcfs is not None: if not isinstance(filename, basestring): filename = filename[0] isrc = fio_py.open_source(fio_py.FIO_M3DC1_SOURCE, filename) fio_py.get_options(isrc) psi = read_field("psi", slice=-1, filename=filename, points=points, rrange=rrange, zrange=zrange, iequil=1) ipsi_lcfs = fio_py.get_series(isrc, fio_py.FIO_LCFS_PSI) psi_lcfs = fio_py.eval_series(ipsi_lcfs, 0.0) fio_py.close_series(ipsi_lcfs) ax.contour(psi.data.T, [psi_lcfs], hold="on", origin="lower", extent=extent, colors=col_lcfs, linewidths=1) if bound is not None: (Wi, Wo) = get_wall() ax.plot(Wi[:, 0], Wi[:, 1], "--", color=col_lcfs, linewidth=1) ax.plot(Wo[:, 0], Wo[:, 1], "--", color=col_lcfs, linewidth=1) return (f, ax, im, cb)
def plot_shape(folder='./', rrange=None, zrange=None, bound=False, ax=None, legend=True, fs=1.0, Nlvl_in=10, Nlvl_out=1, linewidth=3, title=None): if ax is None: sns.set_style('white') f, ax = plt.subplots(figsize=[fs*8,fs*12]) else: f = None if isinstance(folder,basestring): folder = [folder] Nf = len(folder) if Nf <=6: cols = sns.color_palette(n_colors=Nf,palette='colorblind') else: cols = sns.color_palette(n_colors=Nf,palette='inferno') cols[-1] = (0.,0.,0.) for i in range(Nf): fn = folder[i]+'/C1.h5' isrc = fio_py.open_source(fio_py.FIO_M3DC1_SOURCE,fn) fio_py.get_options(isrc) ipsi_axis = fio_py.get_series(isrc, fio_py.FIO_MAGAXIS_PSI) ipsi_lcfs = fio_py.get_series(isrc, fio_py.FIO_LCFS_PSI) psi_axis = fio_py.eval_series(ipsi_axis, 0.) psi_lcfs = fio_py.eval_series(ipsi_lcfs, 0.) fio_py.close_series(ipsi_axis) fio_py.close_series(ipsi_lcfs) # print([psi_axis,psi_lcfs]) levels = (np.arange(Nlvl_in+Nlvl_out+1.)/Nlvl_in) levels = levels*(psi_lcfs-psi_axis)+psi_axis if psi_lcfs < psi_axis: levels = np.flipud(levels) psi = read_field('psi',slice=-1,filename=fn,rrange=rrange, zrange=zrange) col = mpl.colors.rgb2hex(cols[i-1]) if i==0: hold='off' else: hold='on' ax.contour(psi.R,psi.Z,psi.data.T,levels,hold=hold,colors=col, linewidths=linewidth,linestyles='solid') if legend: h, = ax.plot([np.inf,np.inf],[np.inf,np.inf],'-',color=col, linewidth=linewidth) h.set_label(folder[i]) if bound: (Wi,Wo) = get_wall() ax.plot(Wi[:,0],Wi[:,1],'r-',linewidth=1) ax.plot(Wo[:,0],Wo[:,1],'r-',linewidth=1) if rrange is None: rrange = [min(psi.R),max(psi.R)] if zrange is None: zrange = [min(psi.Z),max(psi.Z)] ax.set_xlim(rrange) ax.set_ylim(zrange) ax.set_xlabel(r'$R$ (m)',fontsize=28*fs) ax.set_ylabel(r'$Z$ (m)',fontsize=28*fs) ax.tick_params(labelsize=24*fs) if title is not None: ax.set_title(title,fontsize=32*fs) if legend: ax.legend(fontsize=24*fs) if f is not None: plt.tight_layout() return