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_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_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 read_field(name, slice=0, filename='C1.h5', points=200, phi=0.,
rrange=None, zrange=None, linfac=1., iabs=None, iphase=None,
isum=None, iavg=None, idiff=None,ilinear=None, iequil=None,
icomplex=None, ntor=None):
# Redefine zero options to None for convenience
if iabs==0:
iabs = None
if iphase==0:
iphase = None
if isum==0:
isum = None
if iavg==0:
iavg = None
if idiff==0:
idiff = None
if ilinear==0:
ilinear = None
if iequil==0:
iequil = None
if icomplex==0:
icomplex = None
if (isum is not None) or (iavg is not None) or (idiff is not None):
if isinstance(filename,basestring):
filename = [filename]
if isinstance(slice,int):
slice = [slice]
N_files = len(filename)
N_slices = len(slice)
N = N_files*N_slices
if (N_files > 1) and (N_slices > 1):
print("Error: can't combine mutliple filenames and slices")
return None
elif N_files > 1:
slice = slice*N
elif N_slices > 1:
filename = filename*N
if not hasattr(linfac,"__len__"):
# Duplicate constant linear factor
linfac = np.ones(N)*linfac
elif len(linfac) != N:
print("Error: incorrect number of linfacs")
return None
if not hasattr(phi,"__len__"):
# Duplicate constant linear factor
phi = np.ones(N)*phi
elif len(phi) != N:
print("Error: incorrect number of phis")
return None
field = read_field(name, slice=slice[0], filename=filename[0],
points=points, phi=phi[0], rrange=rrange,
zrange=zrange, linfac=linfac[0],
ilinear=ilinear, iequil=iequil,
icomplex=icomplex, ntor=ntor)
if N==1:
return field
for i in [x+1 for x in range(N-1)]:
field2 = read_field(name, slice=slice[i], filename=filename[i],
points=points, phi=phi[i], rrange=rrange,
zrange=zrange, linfac=linfac[i],
ilinear=ilinear, iequil=iequil,
icomplex=icomplex, ntor=ntor)
if (isum is not None) or (iavg is not None):
field = field + field2
elif (idiff is not None):
field = field + field2*((-1)**i)
if iavg is not None:
return field/N
else:
return field
else:
if not isinstance(filename,basestring):
print('Warning: Only considering first filename')
filename = filename[0]
if hasattr(slice,"__len__"):
print('Warning: Only considering first slice')
slice = slice[0]
if hasattr(linfac,"__len__"):
print('Warning: Only considering first linfac')
linfac = linfac[0]
if hasattr(phi,"__len__"):
print('Warning: Only considering first phi')
phi = phi[0]
# initialize the M3D-C1 file
isrc = fio_py.open_source(fio_py.FIO_M3DC1_SOURCE,filename)
fio_py.get_options(isrc)
# Set up options given as imput
fio_py.set_int_option(fio_py.FIO_TIMESLICE,slice)
if ilinear is not None:
fio_py.set_int_option(fio_py.FIO_PART,fio_py.FIO_PERTURBED_ONLY)
elif iequil is not None:
fio_py.set_int_option(fio_py.FIO_PART,fio_py.FIO_EQUILIBRIUM_ONLY)
else:
fio_py.set_int_option(fio_py.FIO_PART,fio_py.FIO_TOTAL)
if rrange is None:
rrange = [0.75,2.5] # I want to read this from the file somehow
if zrange is None:
zrange = [-1.5,1.5] # I want to read this from the file somehow
R = np.linspace(rrange[0],rrange[1],points)
Z = np.linspace(zrange[0],zrange[1],points)
# Are we reading a real field?
if icomplex is None:
ntor = None
elif ntor is None:
print('Warning: ntor must be assigned for complex evaluation')
print('Reading real field')
icomplex = None
elif ilinear is None:
print('Warning: Trying to access complex values but ilinear=None')
print(' This has not been properly implemented yet')
print('Reading real field')
icomplex = None
ntor = None
if icomplex is None:
dtype = np.float64
else:
dtype = np.complex128
data = np.zeros((points,points),dtype=dtype)
# Lists for scalar fields read from scalars
total_pressure = ['total pressure','p']
ion_pressure = ['ion pressure','pi']
electron_pressure = ['electron pressure','pe']
ion_density = ['ion density', 'ni', 'den']
electron_density = ['electron density', 'ne']
# Dictionaries for scalar fields read from vectors
B_comp = {'bx':'R', 'by':'phi', 'bt':'phi', 'bphi':'phi', 'bz':'Z'}
J_comp = {'jx':'R', 'jy':'phi', 'jt':'phi', 'jphi':'phi', 'jz':'Z'}
E_comp = {'ex':'R', 'ey':'phi', 'et':'phi', 'ephi':'phi', 'ez':'Z'}
v_comp = {'vx':'R', 'vy':'phi', 'vt':'phi', 'vphi':'phi', 'vz':'Z'}
A_comp = {'ax':'R', 'ay':'phi', 'at':'phi', 'aphi':'phi', 'az':'Z'}
comps = ['R', 'phi', 'Z']
# Lists for vector fields
magnetic_field = ['magnetic field', 'b field', 'bfield', 'b_field']
current_density = ['current density', 'j']
electric_field = ['electric field', 'e', 'efield', 'e_field']
fluid_velocity = ['fluid velocity', 'v']
vector_potential = ['vector potential', 'a']
# Lists for composite fields
major_radius = ['major radius', 'r', 'x']
height = ['height', 'z']
ion_temperature = ['ion_temperature', 'ti']
electron_temperature = ['electron_temperature', 'te']
name_lc = name.lower()
# Primitive scalar fields
if name_lc in total_pressure:
field = xr.DataArray(data,[('R',R),('Z',Z)])
field.attrs['type'] = fio_py.FIO_TOTAL_PRESSURE
field.name = total_pressure[0]
elif name_lc in ion_pressure:
field = xr.DataArray(data,[('R',R),('Z',Z)])
field.attrs['species'] = fio_py.FIO_MAIN_ION
field.attrs['type'] = fio_py.FIO_PRESSURE
field.name = ion_pressure[0]
elif name_lc in ion_density:
field = xr.DataArray(data,[('R',R),('Z',Z)])
field.attrs['species'] = fio_py.FIO_MAIN_ION
field.attrs['type'] = fio_py.FIO_DENSITY
field.name = ion_density[0]
elif name_lc in electron_pressure:
field = xr.DataArray(data,[('R',R),('Z',Z)])
field.attrs['species'] = fio_py.FIO_ELECTRON
field.attrs['type'] = fio_py.FIO_PRESSURE
field.name = electron_pressure[0]
elif name_lc in electron_density:
field = xr.DataArray(data,[('R',R),('Z',Z)])
field.attrs['species'] = fio_py.FIO_ELECTRON
field.attrs['type'] = fio_py.FIO_DENSITY
field.name = electron_density[0]
# Primitive vector fields
elif name_lc in magnetic_field:
data = np.array([data,data,data])
field = xr.DataArray(data,[('component',comps),('R',R),('Z',Z)])
field.attrs['type'] = fio_py.FIO_MAGNETIC_FIELD
field.name = magnetic_field[0]
elif name_lc in current_density:
data = np.array([data,data,data])
field = xr.DataArray(data,[('component',comps),('R',R),('Z',Z)])
field.attrs['type'] = fio_py.FIO_CURRENT_DENSITY
field.name = current_density[0]
elif name_lc in electric_field:
data = np.array([data,data,data])
field = xr.DataArray(data,[('component',comps),('R',R),('Z',Z)])
field.attrs['type'] = fio_py.FIO_ELECTRIC_FIELD
field.name = electric_field[0]
elif name_lc in fluid_velocity:
data = np.array([data,data,data])
field = xr.DataArray(data,[('component',comps),('R',R),('Z',Z)])
field.attrs['type'] = fio_py.FIO_FLUID_VELOCITY
field.name = fluid_velocity[0]
elif name_lc in vector_potential:
data = np.array([data,data,data])
field = xr.DataArray(data,[('component',comps),('R',R),('Z',Z)])
field.attrs['type'] = fio_py.FIO_VECTOR_POTENTIAL
field.name = vector_potential[0]
# Vector components
elif name_lc in B_comp:
B = read_field('bfield', slice=slice, filename=filename, points=points,
phi=phi, rrange=rrange, zrange=zrange, ilinear=ilinear,
iequil=iequil, icomplex=icomplex, ntor=ntor)
field = B.sel(component=B_comp[name_lc])
field.attrs['type'] = 'component'
elif name_lc in J_comp:
J = read_field('j', slice=slice, filename=filename, points=points,
phi=phi, rrange=rrange, zrange=zrange, ilinear=ilinear,
iequil=iequil, icomplex=icomplex, ntor=ntor)
field = J.sel(component=J_comp[name_lc])
field.attrs['type'] = 'component'
elif name_lc in E_comp:
E = read_field('efield', slice=slice, filename=filename, points=points,
phi=phi, rrange=rrange, zrange=zrange, ilinear=ilinear,
iequil=iequil, icomplex=icomplex, ntor=ntor)
field = E.sel(component=E_comp[name_lc])
field.attrs['type'] = 'component'
elif name_lc in v_comp:
v = read_field('v', slice=slice, filename=filename, points=points,
phi=phi, rrange=rrange, zrange=zrange, ilinear=ilinear,
iequil=iequil, icomplex=icomplex, ntor=ntor)
field = v.sel(component=v_comp[name_lc])
field.attrs['type'] = 'component'
elif name_lc in A_comp:
A = read_field('a', slice=slice, filename=filename, points=points,
phi=phi, rrange=rrange, zrange=zrange, ilinear=ilinear,
iequil=iequil, icomplex=icomplex, ntor=ntor)
field = A.sel(component=A_comp[name_lc])
field.attrs['type'] = 'component'
# Composite fields
elif name_lc == 'zero':
field = xr.DataArray(data,[('R',R),('Z',Z)])
field.attrs['type'] = 'composite'
field.name = name_lc
elif name_lc in major_radius:
field = xr.DataArray(data,[('R',R),('Z',Z)])
for j in range(points):
field[dict(Z=j)]=field.coords['R']
field.attrs['type'] = 'composite'
field.name = major_radius[0]
elif name_lc in height:
for i in range(points):
field[dict(R=i)]=field.coords['Z']
field.attrs['type'] = 'composite'
field.name = height[0]
elif name_lc == 'beta':
if ilinear is not None:
print(name + ': Cannot calculate linear version')
print('Running for total instead')
P = read_field('p', slice=slice, filename=filename, points=points,
phi=phi, rrange=rrange, zrange=zrange,iequil=iequil)
B = read_field('b', slice=slice, filename=filename, points=points,
phi=phi, rrange=rrange, zrange=zrange,iequil=iequil)
field = 2.0*spc.mu_0*P/(B**2)
field.attrs['type'] = 'composite'
field.name = name_lc
elif name_lc in ion_temperature:
Pi = read_field('pi', slice=slice, filename=filename, points=points,
phi=phi, rrange=rrange, zrange=zrange, ilinear=ilinear,
iequil=iequil, icomplex=icomplex, ntor=ntor)
ni = read_field('ni', slice=slice, filename=filename, points=points,
phi=phi, rrange=rrange, zrange=zrange, ilinear=ilinear,
iequil=iequil, icomplex=icomplex, ntor=ntor)
if ilinear is None:
Ti = Pi/ni
else:
Pi0 = read_field('pi', slice=slice, filename=filename, points=points,
phi=phi, rrange=rrange, zrange=zrange, iequil=1)
ni0 = read_field('ni', slice=slice, filename=filename, points=points,
phi=phi, rrange=rrange, zrange=zrange, iequil=1)
Ti = Pi/ni0 - (Pi0*ni)/(ni0**2)
J_eV = spc.physical_constants['electron volt'][0]
field = Ti/J_eV
field.attrs['type'] = 'composite'
field.name = ion_temperature[0]
elif name_lc in electron_temperature:
Pe = read_field('pe', slice=slice, filename=filename, points=points,
phi=phi, rrange=rrange, zrange=zrange, ilinear=ilinear,
iequil=iequil, icomplex=icomplex, ntor=ntor)
ne = read_field('ne', slice=slice, filename=filename, points=points,
phi=phi, rrange=rrange, zrange=zrange, ilinear=ilinear,
iequil=iequil, icomplex=icomplex, ntor=ntor)
if ilinear is not None:
Pe0 = read_field('pe', slice=slice, filename=filename, points=points,
phi=phi, rrange=rrange, zrange=zrange, iequil=1)
ne0 = read_field('ne', slice=slice, filename=filename, points=points,
phi=phi, rrange=rrange, zrange=zrange, iequil=1)
Te = Pe/ne0 - (Pe0*ne)/(ne0**2)
else:
Te = Pe/ne
J_eV = spc.physical_constants['electron volt'][0]
field = Te/J_eV
field.attrs['type'] = 'composite'
field.name = electron_temperature[0]
elif name_lc == 'b':
B = read_field('bfield', slice=slice, filename=filename,
points=points, phi=phi, rrange=rrange, zrange=zrange,
ilinear=ilinear, iequil=iequil, icomplex=icomplex,
ntor=ntor)
Bx = B.sel(component='R')
By = B.sel(component='phi')
Bz = B.sel(component='Z')
if ilinear is not None:
B0 = read_field('bfield', slice=slice, filename=filename, phi=phi,
points=points, rrange=rrange, zrange=zrange,
iequil=1)
Bx0 = B0.sel(component='R')
By0 = B0.sel(component='phi')
Bz0 = B0.sel(component='Z')
b0 = (Bx0**2 + By0**2 + Bz0**2)**0.5
b = (Bx*Bx0 + By*By0 + Bz*Bz0)/b0
else:
b = (Bx**2 + By**2 + Bz**2)**0.5
field = b.drop('component')
field.attrs['type'] = 'composite'
field.name = name_lc
elif name_lc == 'b2':
B = read_field('bfield', slice=slice, filename=filename,
points=points, phi=phi, rrange=rrange, zrange=zrange,
ilinear=ilinear, iequil=iequil, icomplex=icomplex,
ntor=ntor)
Bx = B.sel(component='R')
By = B.sel(component='phi')
Bz = B.sel(component='Z')
if ilinear is not None:
B0 = read_field('bfield', slice=slice, filename=filename, phi=phi,
points=points, rrange=rrange, zrange=zrange,
iequil=1)
Bx0 = B0.sel(component='R')
By0 = B0.sel(component='phi')
Bz0 = B0.sel(component='Z')
b2 = 2.0*(Bx*Bx0 + By*By0 + Bz*Bz0)
else:
b2 = Bx**2 + By**2 + Bz**2
field = b2.drop('component')
field.attrs['type'] = 'composite'
field.name = name_lc
elif name_lc == 'omega':
vy = read_field('vy', slice=slice, filename=filename, phi=phi,
points=points, rrange=rrange, zrange=zrange,
ilinear=ilinear, iequil=iequil, icomplex=icomplex,
ntor=ntor)
R = read_field('r', slice=slice, filename=filename, phi=phi,
points=points, rrange=rrange, zrange=zrange,
ilinear=ilinear, iequil=iequil, icomplex=icomplex,
ntor=ntor)
field = vy/R
field.attrs['type'] = 'composite'
field.name = name_lc
elif name_lc == 'psi':
Aphi = read_field('ay', slice=slice, filename=filename, phi=phi,
points=points, rrange=rrange, zrange=zrange,
ilinear=ilinear, iequil=iequil, icomplex=icomplex,
ntor=ntor)
R = read_field('r', slice=slice, filename=filename, phi=phi,
points=points, rrange=rrange, zrange=zrange,
ilinear=ilinear, iequil=iequil, icomplex=icomplex,
ntor=ntor)
field = R*Aphi
field = field.drop('component')
field.attrs['type'] = 'composite'
field.name = name_lc
else:
print "Field '" + name_lc + "' is+- not defined"
return None
field.attrs['ntor'] = ntor
#**************************************************************************
if not isinstance(field.attrs['type'],basestring):
if 'species' in field.attrs:
fio_py.set_int_option(fio_py.FIO_SPECIES, field.attrs['species'])
handle = fio_py.get_field(isrc,field.attrs['type'])
if len(field.shape) == 2:
enum = np.ndenumerate(field)
eval_field = fio_py.eval_scalar_field
elif len(field.shape) == 3:
enum = np.ndenumerate(field.data[0,:,:])
eval_field = fio_py.eval_vector_field
for (i,j), dummy in enum:
# Real part
x = (R[i], phi*np.pi/180., Z[j])
val = np.asarray(eval_field(handle, x),dtype=dtype)
# Imaginary part
if icomplex is not None:
y = (x[0], x[1] + 3.0*np.pi/(2.0*ntor), x[2])
val += 1j*np.asarray(eval_field(handle, y),dtype=dtype)
field[dict(R=i,Z=j)] = val
return linfac*field
def read_field(name,
slice=0,
filename='C1.h5',
points=200,
phi=0.,
rrange=None,
zrange=None,
linfac=1.,
iabs=None,
iphase=None,
isum=None,
iavg=None,
idiff=None,
ilinear=None,
iequil=None,
icomplex=None,
ntor=None,
nimrod=False):
# Redefine zero options to None for convenience
if iabs == 0:
iabs = None
if iphase == 0:
iphase = None
if isum == 0:
isum = None
if iavg == 0:
iavg = None
if idiff == 0:
idiff = None
if ilinear == 0:
ilinear = None
if iequil == 0:
iequil = None
if icomplex == 0:
icomplex = None
if (isum is not None) or (iavg is not None) or (idiff is not None):
if isinstance(filename, str):
filename = [filename]
if isinstance(slice, int):
slice = [slice]
N_files = len(filename)
N_slices = len(slice)
N = N_files * N_slices
if (N_files > 1) and (N_slices > 1):
print("Error: can't combine mutliple filenames and slices")
return None
elif N_files > 1:
slice = slice * N
elif N_slices > 1:
filename = filename * N
if not hasattr(linfac, "__len__"):
# Duplicate constant linear factor
linfac = np.ones(N) * linfac
elif len(linfac) != N:
print("Error: incorrect number of linfacs")
return None
if not hasattr(phi, "__len__"):
# Duplicate constant linear factor
phi = np.ones(N) * phi
elif len(phi) != N:
print("Error: incorrect number of phis")
return None
field = read_field(name,
slice=slice[0],
filename=filename[0],
points=points,
phi=phi[0],
rrange=rrange,
zrange=zrange,
linfac=linfac[0],
ilinear=ilinear,
iequil=iequil,
icomplex=icomplex,
ntor=ntor)
if N == 1:
return field
for i in [x + 1 for x in range(N - 1)]:
field2 = read_field(name,
slice=slice[i],
filename=filename[i],
points=points,
phi=phi[i],
rrange=rrange,
zrange=zrange,
linfac=linfac[i],
ilinear=ilinear,
iequil=iequil,
icomplex=icomplex,
ntor=ntor)
if (isum is not None) or (iavg is not None):
field = field + field2
elif (idiff is not None):
field = field + field2 * ((-1)**i)
if iavg is not None:
return field / N
else:
return field
else:
if not isinstance(filename, str):
print('Warning: Only considering first filename')
filename = filename[0]
if hasattr(slice, "__len__"):
print('Warning: Only considering first slice')
slice = slice[0]
if hasattr(linfac, "__len__"):
print('Warning: Only considering first linfac')
linfac = linfac[0]
if hasattr(phi, "__len__"):
print('Warning: Only considering first phi')
phi = phi[0]
if nimrod:
#filename is actually a folder
R = np.loadtxt(filename + '/r.txt')
Z = np.loadtxt(filename + '/z.txt')
data = np.loadtxt(filename + '/' + name + ('%d' % slice).zfill(5) +
'.txt')
data = np.transpose(np.nan_to_num(data))
field = xr.DataArray(data, coords=[('Z', Z), ('R', R)])
return field
# initialize the M3D-C1 file
isrc = fio_py.open_source(fio_py.FIO_M3DC1_SOURCE, filename)
fio_py.get_options(isrc)
# Set up options given as imput
fio_py.set_int_option(fio_py.FIO_TIMESLICE, slice)
if ilinear is not None:
fio_py.set_int_option(fio_py.FIO_PART, fio_py.FIO_PERTURBED_ONLY)
elif iequil is not None:
fio_py.set_int_option(fio_py.FIO_PART, fio_py.FIO_EQUILIBRIUM_ONLY)
else:
fio_py.set_int_option(fio_py.FIO_PART, fio_py.FIO_TOTAL)
if rrange is None:
rrange = [0.75, 2.5] # I want to read this from the file somehow
if zrange is None:
zrange = [-1.5, 1.5] # I want to read this from the file somehow
R = np.linspace(rrange[0], rrange[1], points)
Z = np.linspace(zrange[0], zrange[1], points)
# Are we reading a real field?
if icomplex is None:
ntor = None
elif ntor is None:
print('Warning: ntor must be assigned for complex evaluation')
print('Reading real field')
icomplex = None
elif ilinear is None:
print('Warning: Trying to access complex values but ilinear=None')
print(' This has not been properly implemented yet')
print('Reading real field')
icomplex = None
ntor = None
if icomplex is None:
dtype = np.float64
else:
dtype = np.complex128
data = np.zeros((points, points), dtype=dtype)
# Lists for scalar fields read from scalars
total_pressure = ['total pressure', 'p']
ion_pressure = ['ion pressure', 'pi']
electron_pressure = ['electron pressure', 'pe']
ion_density = ['ion density', 'ni', 'den']
electron_density = ['electron density', 'ne']
# Dictionaries for scalar fields read from vectors
B_comp = {'bx': 'R', 'by': 'phi', 'bt': 'phi', 'bphi': 'phi', 'bz': 'Z'}
J_comp = {'jx': 'R', 'jy': 'phi', 'jt': 'phi', 'jphi': 'phi', 'jz': 'Z'}
E_comp = {'ex': 'R', 'ey': 'phi', 'et': 'phi', 'ephi': 'phi', 'ez': 'Z'}
v_comp = {'vx': 'R', 'vy': 'phi', 'vt': 'phi', 'vphi': 'phi', 'vz': 'Z'}
A_comp = {'ax': 'R', 'ay': 'phi', 'at': 'phi', 'aphi': 'phi', 'az': 'Z'}
comps = ['R', 'phi', 'Z']
# Lists for vector fields
magnetic_field = ['magnetic field', 'b field', 'bfield', 'b_field']
current_density = ['current density', 'j']
electric_field = ['electric field', 'e', 'efield', 'e_field']
fluid_velocity = ['fluid velocity', 'v']
vector_potential = ['vector potential', 'a']
# Lists for composite fields
major_radius = ['major radius', 'r', 'x']
height = ['height', 'z']
ion_temperature = ['ion_temperature', 'ti']
electron_temperature = ['electron_temperature', 'te']
Bpl_comp = {
'bx_plasma': 'R',
'by_plasma': 'phi',
'bt_plasma': 'phi',
'bphi_plasma': 'phi',
'bz_plasma': 'Z'
}
Jpl_comp = {
'jx_plasma': 'R',
'jy_plasma': 'phi',
'jt_plasma': 'phi',
'jphi_plasma': 'phi',
'jz_plasma': 'Z'
}
name_lc = name.lower()
# Primitive scalar fields
if name_lc in total_pressure:
field = xr.DataArray(data, coords=[('R', R), ('Z', Z)])
field.attrs['type'] = fio_py.FIO_TOTAL_PRESSURE
field.name = total_pressure[0]
elif name_lc in ion_pressure:
field = xr.DataArray(data, coords=[('R', R), ('Z', Z)])
field.attrs['species'] = fio_py.FIO_MAIN_ION
field.attrs['type'] = fio_py.FIO_PRESSURE
field.name = ion_pressure[0]
elif name_lc in ion_density:
field = xr.DataArray(data, coords=[('R', R), ('Z', Z)])
field.attrs['species'] = fio_py.FIO_MAIN_ION
field.attrs['type'] = fio_py.FIO_DENSITY
field.name = ion_density[0]
elif name_lc in electron_pressure:
field = xr.DataArray(data, coords=[('R', R), ('Z', Z)])
field.attrs['species'] = fio_py.FIO_ELECTRON
field.attrs['type'] = fio_py.FIO_PRESSURE
field.name = electron_pressure[0]
elif name_lc in electron_density:
field = xr.DataArray(data, coords=[('R', R), ('Z', Z)])
field.attrs['species'] = fio_py.FIO_ELECTRON
field.attrs['type'] = fio_py.FIO_DENSITY
field.name = electron_density[0]
# Primitive vector fields
elif name_lc in magnetic_field:
data = np.array([data, data, data])
field = xr.DataArray(data,
coords=[('component', comps), ('R', R), ('Z', Z)])
field.attrs['type'] = fio_py.FIO_MAGNETIC_FIELD
field.name = magnetic_field[0]
elif name_lc in current_density:
data = np.array([data, data, data])
field = xr.DataArray(data,
coords=[('component', comps), ('R', R), ('Z', Z)])
field.attrs['type'] = fio_py.FIO_CURRENT_DENSITY
field.name = current_density[0]
elif name_lc in electric_field:
data = np.array([data, data, data])
field = xr.DataArray(data,
coords=[('component', comps), ('R', R), ('Z', Z)])
field.attrs['type'] = fio_py.FIO_ELECTRIC_FIELD
field.name = electric_field[0]
elif name_lc in fluid_velocity:
data = np.array([data, data, data])
field = xr.DataArray(data,
coords=[('component', comps), ('R', R), ('Z', Z)])
field.attrs['type'] = fio_py.FIO_FLUID_VELOCITY
field.name = fluid_velocity[0]
elif name_lc in vector_potential:
data = np.array([data, data, data])
field = xr.DataArray(data,
coords=[('component', comps), ('R', R), ('Z', Z)])
field.attrs['type'] = fio_py.FIO_VECTOR_POTENTIAL
field.name = vector_potential[0]
# Vector components
elif name_lc in B_comp:
B = read_field('bfield',
slice=slice,
filename=filename,
points=points,
phi=phi,
rrange=rrange,
zrange=zrange,
ilinear=ilinear,
iequil=iequil,
icomplex=icomplex,
ntor=ntor)
field = B.sel(component=B_comp[name_lc])
field.attrs['type'] = 'component'
elif name_lc in J_comp:
J = read_field('j',
slice=slice,
filename=filename,
points=points,
phi=phi,
rrange=rrange,
zrange=zrange,
ilinear=ilinear,
iequil=iequil,
icomplex=icomplex,
ntor=ntor)
field = J.sel(component=J_comp[name_lc])
field.attrs['type'] = 'component'
elif name_lc in E_comp:
E = read_field('efield',
slice=slice,
filename=filename,
points=points,
phi=phi,
rrange=rrange,
zrange=zrange,
ilinear=ilinear,
iequil=iequil,
icomplex=icomplex,
ntor=ntor)
field = E.sel(component=E_comp[name_lc])
field.attrs['type'] = 'component'
elif name_lc in v_comp:
v = read_field('v',
slice=slice,
filename=filename,
points=points,
phi=phi,
rrange=rrange,
zrange=zrange,
ilinear=ilinear,
iequil=iequil,
icomplex=icomplex,
ntor=ntor)
field = v.sel(component=v_comp[name_lc])
field.attrs['type'] = 'component'
elif name_lc in A_comp:
A = read_field('a',
slice=slice,
filename=filename,
points=points,
phi=phi,
rrange=rrange,
zrange=zrange,
ilinear=ilinear,
iequil=iequil,
icomplex=icomplex,
ntor=ntor)
field = A.sel(component=A_comp[name_lc])
field.attrs['type'] = 'component'
# Composite fields
elif name_lc == 'zero':
field = xr.DataArray(data, coords=[('R', R), ('Z', Z)])
field.attrs['type'] = 'composite'
field.name = name_lc
elif name_lc in major_radius:
field = xr.DataArray(data, coords=[('R', R), ('Z', Z)])
for j in range(points):
field[dict(Z=j)] = field.coords['R']
field.attrs['type'] = 'composite'
field.name = major_radius[0]
elif name_lc in height:
for i in range(points):
field[dict(R=i)] = field.coords['Z']
field.attrs['type'] = 'composite'
field.name = height[0]
elif name_lc == 'beta':
if ilinear is not None:
print(name + ': Cannot calculate linear version')
print('Running for total instead')
P = read_field('p',
slice=slice,
filename=filename,
points=points,
phi=phi,
rrange=rrange,
zrange=zrange,
iequil=iequil)
B = read_field('b',
slice=slice,
filename=filename,
points=points,
phi=phi,
rrange=rrange,
zrange=zrange,
iequil=iequil)
field = 2.0 * spc.mu_0 * P / (B**2)
field.attrs['type'] = 'composite'
field.name = name_lc
elif name_lc in ion_temperature:
Pi = read_field('pi',
slice=slice,
filename=filename,
points=points,
phi=phi,
rrange=rrange,
zrange=zrange,
ilinear=ilinear,
iequil=iequil,
icomplex=icomplex,
ntor=ntor)
ni = read_field('ni',
slice=slice,
filename=filename,
points=points,
phi=phi,
rrange=rrange,
zrange=zrange,
ilinear=ilinear,
iequil=iequil,
icomplex=icomplex,
ntor=ntor)
if ilinear is None:
Ti = Pi / ni
else:
Pi0 = read_field('pi',
slice=slice,
filename=filename,
points=points,
phi=phi,
rrange=rrange,
zrange=zrange,
iequil=1)
ni0 = read_field('ni',
slice=slice,
filename=filename,
points=points,
phi=phi,
rrange=rrange,
zrange=zrange,
iequil=1)
Ti = Pi / ni0 - (Pi0 * ni) / (ni0**2)
J_eV = spc.physical_constants['electron volt'][0]
field = Ti / J_eV
field.attrs['type'] = 'composite'
field.name = ion_temperature[0]
elif name_lc in electron_temperature:
Pe = read_field('pe',
slice=slice,
filename=filename,
points=points,
phi=phi,
rrange=rrange,
zrange=zrange,
ilinear=ilinear,
iequil=iequil,
icomplex=icomplex,
ntor=ntor)
ne = read_field('ne',
slice=slice,
filename=filename,
points=points,
phi=phi,
rrange=rrange,
zrange=zrange,
ilinear=ilinear,
iequil=iequil,
icomplex=icomplex,
ntor=ntor)
if ilinear is not None:
Pe0 = read_field('pe',
slice=slice,
filename=filename,
points=points,
phi=phi,
rrange=rrange,
zrange=zrange,
iequil=1)
ne0 = read_field('ne',
slice=slice,
filename=filename,
points=points,
phi=phi,
rrange=rrange,
zrange=zrange,
iequil=1)
Te = Pe / ne0 - (Pe0 * ne) / (ne0**2)
else:
Te = Pe / ne
J_eV = spc.physical_constants['electron volt'][0]
field = Te / J_eV
field.attrs['type'] = 'composite'
field.name = electron_temperature[0]
elif name_lc == 'b':
B = read_field('bfield',
slice=slice,
filename=filename,
points=points,
phi=phi,
rrange=rrange,
zrange=zrange,
ilinear=ilinear,
iequil=iequil,
icomplex=icomplex,
ntor=ntor)
Bx = B.sel(component='R')
By = B.sel(component='phi')
Bz = B.sel(component='Z')
if ilinear is not None:
B0 = read_field('bfield',
slice=slice,
filename=filename,
phi=phi,
points=points,
rrange=rrange,
zrange=zrange,
iequil=1)
Bx0 = B0.sel(component='R')
By0 = B0.sel(component='phi')
Bz0 = B0.sel(component='Z')
b0 = (Bx0**2 + By0**2 + Bz0**2)**0.5
b = (Bx * Bx0 + By * By0 + Bz * Bz0) / b0
else:
b = (Bx**2 + By**2 + Bz**2)**0.5
field = b.drop('component')
field.attrs['type'] = 'composite'
field.name = name_lc
elif name_lc == 'b2':
B = read_field('bfield',
slice=slice,
filename=filename,
points=points,
phi=phi,
rrange=rrange,
zrange=zrange,
ilinear=ilinear,
iequil=iequil,
icomplex=icomplex,
ntor=ntor)
Bx = B.sel(component='R')
By = B.sel(component='phi')
Bz = B.sel(component='Z')
if ilinear is not None:
B0 = read_field('bfield',
slice=slice,
filename=filename,
phi=phi,
points=points,
rrange=rrange,
zrange=zrange,
iequil=1)
Bx0 = B0.sel(component='R')
By0 = B0.sel(component='phi')
Bz0 = B0.sel(component='Z')
b2 = 2.0 * (Bx * Bx0 + By * By0 + Bz * Bz0)
else:
b2 = Bx**2 + By**2 + Bz**2
field = b2.drop('component')
field.attrs['type'] = 'composite'
field.name = name_lc
elif name_lc in Bpl_comp:
B0 = read_field('bfield',
slice=0,
filename=filename,
points=points,
phi=phi,
rrange=rrange,
zrange=zrange,
ilinear=ilinear,
iequil=iequil,
icomplex=icomplex,
ntor=ntor)
B1 = read_field('bfield',
slice=slice,
filename=filename,
points=points,
phi=phi,
rrange=rrange,
zrange=zrange,
ilinear=ilinear,
iequil=iequil,
icomplex=icomplex,
ntor=ntor)
Bpl = B1 - B0
field = Bpl.sel(component=Bpl_comp[name_lc])
field.attrs['type'] = 'component'
elif name_lc in Jpl_comp:
J0 = read_field('j',
slice=0,
filename=filename,
points=points,
phi=phi,
rrange=rrange,
zrange=zrange,
ilinear=ilinear,
iequil=iequil,
icomplex=icomplex,
ntor=ntor)
J1 = read_field('j',
slice=slice,
filename=filename,
points=points,
phi=phi,
rrange=rrange,
zrange=zrange,
ilinear=ilinear,
iequil=iequil,
icomplex=icomplex,
ntor=ntor)
Jpl = J1 - J0
field = Jpl.sel(component=Jpl_comp[name_lc])
field.attrs['type'] = 'component'
elif name_lc == 'omega':
vy = read_field('vy',
slice=slice,
filename=filename,
phi=phi,
points=points,
rrange=rrange,
zrange=zrange,
ilinear=ilinear,
iequil=iequil,
icomplex=icomplex,
ntor=ntor)
R = read_field('r',
slice=slice,
filename=filename,
phi=phi,
points=points,
rrange=rrange,
zrange=zrange,
ilinear=ilinear,
iequil=iequil,
icomplex=icomplex,
ntor=ntor)
field = vy / R
field.attrs['type'] = 'composite'
field.name = name_lc
elif name_lc == 'psi':
Aphi = read_field('ay',
slice=slice,
filename=filename,
phi=phi,
points=points,
rrange=rrange,
zrange=zrange,
ilinear=ilinear,
iequil=iequil,
icomplex=icomplex,
ntor=ntor)
R = read_field('r',
slice=slice,
filename=filename,
phi=phi,
points=points,
rrange=rrange,
zrange=zrange,
ilinear=ilinear,
iequil=iequil,
icomplex=icomplex,
ntor=ntor)
field = R * Aphi
field = field.drop('component')
field.attrs['type'] = 'composite'
field.name = name_lc
else:
print("Field '" + name_lc + "' is not defined")
return None
field.attrs['ntor'] = ntor
#**************************************************************************
if not isinstance(field.attrs['type'], str):
if 'species' in field.attrs:
fio_py.set_int_option(fio_py.FIO_SPECIES, field.attrs['species'])
handle = fio_py.get_field(isrc, field.attrs['type'])
hint = fio_py.allocate_hint(isrc)
if len(field.shape) == 2:
enum = np.ndenumerate(field)
eval_field = fio_py.eval_scalar_field
elif len(field.shape) == 3:
enum = np.ndenumerate(field.data[0, :, :])
eval_field = fio_py.eval_vector_field
for (i, j), dummy in enum:
# Real part
x = (R[i], phi * np.pi / 180., Z[j])
val = np.asarray(eval_field(handle, x, hint), dtype=dtype)
# Imaginary part
if icomplex is not None:
y = (x[0], x[1] + 3.0 * np.pi / (2.0 * ntor), x[2])
val += 1j * np.asarray(eval_field(handle, y, hint),
dtype=dtype)
field[dict(R=i, Z=j)] = val
return linfac * field
import fio_py
filename = "C1.h5"
isrc = fio_py.open_source(fio_py.FIO_M3DC1_SOURCE, filename)
fio_py.get_options(isrc)
ia = fio_py.get_field(isrc, fio_py.FIO_VECTOR_POTENTIAL)
imag = fio_py.get_field(isrc, fio_py.FIO_MAGNETIC_FIELD)
ipres = fio_py.get_field(isrc, fio_py.FIO_TOTAL_PRESSURE)
ivel = fio_py.get_field(isrc, fio_py.FIO_FLUID_VELOCITY)
ij = fio_py.get_field(isrc, fio_py.FIO_CURRENT_DENSITY)
fio_py.set_int_option(fio_py.FIO_SPECIES, fio_py.FIO_ELECTRON)
idens = fio_py.get_field(isrc, fio_py.FIO_DENSITY)
fio_py.set_int_option(fio_py.FIO_SPECIES, fio_py.FIO_MAIN_ION)
ipi = fio_py.get_field(isrc, fio_py.FIO_PRESSURE)
list = fio_py.get_available_fields(isrc)
print 'Available fields:'
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'
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
