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,
fac=1.,
linfac=1.,
phi=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,
symrange=False,
cb_label=None,
minval=None,
maxval=None,
nimrod=False,
make_cb=True):
if isinstance(field, str):
# 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,
nimrod=nimrod)
data = np.nan_to_num(fac * np.real(field.data))
if range is None:
if symrange:
vmin = -abs(data).max()
vmax = abs(data).max()
else:
vmin = data.min()
vmax = data.max()
else:
vmin = range[0]
vmax = range[1]
if minval is not None:
vmin = minval
if maxval is not None:
vmax = maxval
if palette is None:
if vmin >= 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 * 9, 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 make_cb:
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 cb_label is not None:
cb.ax.get_yaxis().labelpad = fs * 36
cb.ax.set_ylabel(cb_label, rotation=270, fontsize=fs * 24)
else:
cb = None
if lcfs is not None:
if not isinstance(filename, str):
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.)
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 isinstance(bound, tuple):
R0, a, delta, Z0, b = bound
theta = np.linspace(-np.pi / 2., 3 * np.pi / 2, 1000)
R = R0 + a * np.cos(theta + delta * np.sin(theta))
Z = Z0 + b * np.sin(theta)
ax.plot(R, Z, '-', color=col_lcfs, linewidth=fs * 3)
elif 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)
if f is not None:
f.tight_layout()
else:
plt.tight_layout()
return (f, ax, im, cb)
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)
for field in list:
print ' ', fio_py.get_field_name(field)
cs = fio_py.get_int_parameter(isrc, FIO_GEOMETRY)
if cs == fio_py.FIO_CYLINDRICAL:
print 'Using CYLINDRICAL coordinate system'
else:
print 'Using CARTESIAN coordinate system'
ntime = fio_py.get_int_parameter(isrc, FIO_NUM_TIMESLICES)
print 'Number of time slices: ', ntime
period = fio_py.get_real_parameter(isrc, FIO_PERIOD)
print 'Toroidal period = ', period
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.)
print 'Psi at magnetic axis: ', psi_axis
print 'Psi at lcfs: ', psi_lcfs
fio_py.close_series(ipsi_axis)
fio_py.close_series(ipsi_lcfs)
x = (10., 0., 0.)
(ar, aphi, az) = fio_py.eval_vector_field(ia, x)
(br, bphi, bz) = fio_py.eval_vector_field(imag, x)
(jr, jphi, jz) = fio_py.eval_vector_field(ij, x)
p = fio_py.eval_scalar_field(ipres, x)
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
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,
labels=None):
if ax is None:
sns.set_style('white')
f, ax = plt.subplots(figsize=[fs * 8, fs * 12])
else:
f = None
if isinstance(folder, str):
folder = [folder]
Nf = len(folder)
if Nf <= 6:
cols = sns.color_palette(n_colors=6, 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])
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)
if labels is None:
h.set_label(folder[i])
else:
h.set_label(labels[i])
if bound:
(Wi, Wo) = get_wall()
ax.plot(Wi[:, 0], Wi[:, 1], 'k-', linewidth=2)
ax.plot(Wo[:, 0], Wo[:, 1], 'k-', linewidth=2)
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=20 * fs, loc='upper right')
if f is not None:
plt.tight_layout()
return