from .hdf5 import HDF5
from .misc import get_figure, map_matrix
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.ticker import FuncFormatter
[docs]class FOCUSHDF5(HDF5):
"""FOCUS output hdf5 file"""
# initialization, test = FOCUSHDF5('focus_test.h5')
def __init__(self, filename, periodic=False, **kwargs):
"""Initialization
Keyword arguments:
filenmae -- string, path and name to FOCUS output hdf5 file,
usually in the format of 'focus_*.h5'
periodic -- logical, map all 2D surface data automatically,
(default: True)
"""
# read data
super().__init__(filename)
# print version
try:
abc = ""
for i in self.version:
abc += i.decode("utf-8")
self.version = abc
print("FOCUS version: " + self.version)
except AttributeError:
print(filename + " is not a valid FOCUS output. Please check.")
raise
# add additional colume and row for plotting
if periodic:
self.xsurf = map_matrix(self.xsurf)
self.ysurf = map_matrix(self.ysurf)
self.zsurf = map_matrix(self.zsurf)
self.nx = map_matrix(self.nx)
self.ny = map_matrix(self.ny)
self.nz = map_matrix(self.nz)
self.nn = map_matrix(self.nn)
self.Bx = map_matrix(self.Bx)
self.By = map_matrix(self.By)
self.Bz = map_matrix(self.Bz)
self.Bn = map_matrix(self.Bn)
self.plas_Bn = map_matrix(self.plas_Bn)
return
# convergence plot
[docs] def convergence(self, term="bnorm", iteration=True, axes=None, **kwargs):
# get figure
fig, axes = get_figure(axes)
# set default plotting parameters
kwargs["linewidth"] = kwargs.get("linewidth", 2.5) # line width
# kwargs['marker'] = kwargs.get('marker', 'o') # extent
# get iteration data
if iteration:
abscissa = 1 + np.arange(self.iout)
_xlabel = "iteration"
else:
abscissa = self.evolution[0, :] # be careful; DF is not saving wall-time
_xlabel = "wall time [Second]"
# plot data
if term.lower() == "chi":
data = self.evolution[1, :]
kwargs["label"] = kwargs.get("label", r"$\chi^2$") # line label
elif term.lower() == "gradient":
data = self.evolution[2, :]
kwargs["label"] = kwargs.get(
"label", r"$|d \chi^2 / d {\bf X}|$"
) # line label
elif term.lower() == "bnorm":
data = self.evolution[3, :]
kwargs["label"] = kwargs.get("label", r"$f_{B_n}$") # line label
elif term.lower() == "bharm":
data = self.evolution[4, :]
kwargs["label"] = kwargs.get("label", r"$f_{B_{mn}}$") # line label
elif term.lower() == "tflux":
data = self.evolution[5, :]
kwargs["label"] = kwargs.get("label", r"$f_{\Psi}$") # line label
elif term.lower() == "ttlen":
data = self.evolution[6, :]
kwargs["label"] = kwargs.get("label", r"$f_L$") # line label
elif term.lower() == "cssep":
data = self.evolution[7, :]
kwargs["label"] = kwargs.get("label", r"$f_{CS}$") # line label
elif term.lower() == "curv":
data = self.evolution[8, :]
kwargs["label"] = kwargs.get("label", r"$f_{curv}$") # line label
elif term.lower() == "all":
lines = []
lines.append(
self.convergence(
term="chi",
iteration=iteration,
axes=axes,
linestyle="-",
color="k",
**kwargs
)
)
lines.append(
self.convergence(
term="gradient",
iteration=iteration,
axes=axes,
linestyle="--",
color="k",
**kwargs
)
)
lines.append(
self.convergence(
term="bnorm",
iteration=iteration,
axes=axes,
linestyle="-.",
color="r",
**kwargs
)
)
lines.append(
self.convergence(
term="tflux",
iteration=iteration,
axes=axes,
linestyle="-.",
color="g",
**kwargs
)
)
lines.append(
self.convergence(
term="bharm",
iteration=iteration,
axes=axes,
linestyle="-.",
color="b",
**kwargs
)
)
lines.append(
self.convergence(
term="ttlen",
iteration=iteration,
axes=axes,
linestyle="-.",
color="c",
**kwargs
)
)
lines.append(
self.convergence(
term="cssep",
iteration=iteration,
axes=axes,
linestyle="-.",
color="m",
**kwargs
)
)
lines.append(
self.convergence(
term="curv",
iteration=iteration,
axes=axes,
linestyle="-.",
color="y",
**kwargs
)
)
# fig.legend(loc='upper right', frameon=False, ncol=2, prop={'size':16})
plt.legend()
return lines
else:
raise ValueError("unsupported option for term")
line = axes.semilogy(abscissa, data, **kwargs)
axes.tick_params(axis="both", which="major", labelsize=15)
axes.set_xlabel(_xlabel, fontsize=15)
axes.set_ylabel("cost functions", fontsize=15)
# fig.legend(loc='upper right', frameon=False, prop={'size':24, 'weight':'bold'})
plt.legend()
return line
# poincare plot
[docs] def poincare_plot(self, color=None, prange="full", **kwargs):
"""Poincare plot from FOCUS output.
Args:
color (matplotlib color, or None): dot colors. Defaults to None (rainbow).
prange (str, optional): Plot range, one of ["upper", "lower", "all"]. Defaults to "full".
kwargs : matplotlib scatter keyword arguments
Returns:
None
"""
import matplotlib.pyplot as plt
from matplotlib import cm
# get figure and ax data
if plt.get_fignums():
fig = plt.gcf()
ax = plt.gca()
else:
fig, ax = plt.subplots()
# get colors
if color is None:
colors = cm.rainbow(np.linspace(1, 0, self.pp_ns))
else:
colors = [color] * self.pp_ns
kwargs["s"] = kwargs.get("s", 0.1) # dotsize
# scatter plot
for i in range(self.pp_ns):
# determine whether plot upper or lower
if prange == "upper":
cond = self.ppz[:, i] > 0
elif prange == "lower":
cond = self.ppz[:, i] < 0
else:
cond = np.ones_like(self.ppz[:, i], dtype=bool)
ax.scatter(
self.ppr[:, i][cond], self.ppz[:, i][cond], color=colors[i], **kwargs
)
plt.axis("equal")
plt.xlabel("R [m]", fontsize=20)
plt.ylabel("Z [m]", fontsize=20)
plt.xticks(fontsize=16)
plt.yticks(fontsize=16)
return
[docs] def iota_plot(self, axes=None, **kwargs):
"""Plot the rotational transform profile
Args:
axes (matplotlib.pyplot.axis, optional): axis to be plotted on. Defaults to None.
kwargs (optional): Keyword arguments passed to matplotlib.pyplot.plot, e.g. color='r'.
Returns:
matplotlib.pyplot.axis: The updated axis
"""
import matplotlib.pyplot as plt
fig, axes = get_figure(axes)
kwargs.setdefault("marker", "o")
axes.plot(self.ppr[0, :], self.iota, **kwargs)
axes.set_ylabel(r"$\,\iota\!\!$-", fontsize=18)
axes.set_xlabel("R [m]", fontsize=14)
return axes
# Bnorm plot
[docs] def Bnorm(self, plottype="2D", source="all", axes=None, flip=False, **kwargs):
"""Plot Bn distribution.
Keyword arguments:
plottype -- string, '2D' (default) or '3D', determine the plottype
source -- string, 'coil', 'plasma', 'sum' or 'all' (default), data source
axes -- matplotlib.pyplot or mayavi.mlab axis, axis to be
plotted on (default None)
flip -- logical, determine how to calculate Bn from coil,
True: coil_Bn = Bn - plas_Bn; False: coil_Bn = Bn + plas_Bn
Returns:
obj -- matplotlib.pyplot or mayavi.mlab plotting object
"""
obj = []
# check coil_Bn
if flip:
coil_Bn = self.Bn - self.plas_Bn
else:
coil_Bn = self.Bn + self.plas_Bn
# 2D plots
if plottype.lower() == "2d":
# prepare coordinates
ntheta = self.Nteta
nzeta = self.Nzeta
nt, nz = self.Bn.shape
if self.IsSymmetric == 0:
zeta_end = 2 * np.pi
elif self.IsSymmetric == 1:
zeta_end = 2 * np.pi / self.Nfp
elif self.IsSymmetric == 2:
zeta_end = np.pi / self.Nfp
else:
raise ValueError("Something wrong with IsSymmetric.")
def theta(x, pos):
return "{:3.2f}".format(np.pi / ntheta + x * (2 * np.pi) / ntheta)
def zeta(x, pos):
return "{:3.2f}".format(0.5 * zeta_end / nzeta + x * zeta_end / nzeta)
theta_format = FuncFormatter(theta)
zeta_format = FuncFormatter(zeta)
# planar plotting
if source.lower() == "all":
fig, axes = get_figure(axes, ncols=3, sharex=True)
axes = np.atleast_1d(axes)
# kwargs['aspect'] = kwargs.get('aspect', nt/(2.0*nz*self.Nfp))
kwargs["aspect"] = kwargs.get("aspect", "auto")
else:
fig, axes = get_figure(axes)
axes = np.atleast_1d(axes)
kwargs["aspect"] = kwargs.get("aspect", float(nt / nz))
# prepare axes
for ax in axes:
ax.xaxis.set_major_formatter(zeta_format)
ax.yaxis.set_major_formatter(theta_format)
# set default plotting parameters
kwargs["cmap"] = kwargs.get("cmap", "RdBu_r") # colormap
kwargs["origin"] = kwargs.get("origin", "lower") # number of contours
kwargs["extent"] = kwargs.get("extent", [0, nt, 0, nz]) # extent
plt.subplots_adjust(hspace=0.05)
# imshow
if source.lower() == "coil":
obj.append(axes[0].imshow(np.transpose(coil_Bn), **kwargs))
axes[0].set_title("Bn from coils", fontsize=15)
axes[0].set_ylabel(r"$\theta$", fontsize=14)
axes[0].set_xlabel(r"$\phi$", fontsize=14)
elif source.lower() == "plasma":
obj.append(axes[0].imshow(np.transpose(self.plas_Bn), **kwargs))
axes[0].set_title("Bn from plasma", fontsize=15)
axes[0].set_ylabel(r"$\theta$", fontsize=14)
axes[0].set_xlabel(r"$\phi$", fontsize=14)
elif source.lower() == "sum":
obj.append(axes[0].imshow(np.transpose(self.Bn), **kwargs))
axes[0].set_title("Residual Bn", fontsize=15)
axes[0].set_ylabel(r"$\theta$", fontsize=14)
axes[0].set_xlabel(r"$\phi$", fontsize=14)
elif source.lower() == "all":
vmin = np.min([self.plas_Bn, self.Bn, coil_Bn])
vmax = np.max([self.plas_Bn, self.Bn, coil_Bn])
obj.append(
axes[0].imshow(
np.transpose(self.plas_Bn), vmin=vmin, vmax=vmax, **kwargs
)
)
# plt.setp(axes[0].get_xticklabels(), visible=False)
axes[1].set_title(
"Bn from plasma (left), coil (mid), overall(right).", fontsize=14
)
obj.append(
axes[1].imshow(
np.transpose(coil_Bn), vmin=vmin, vmax=vmax, **kwargs
)
)
# axes[1].set_title('Bn from coils', fontsize=15)
plt.setp(axes[1].get_yticklabels(), visible=False)
obj.append(
axes[2].imshow(
np.transpose(self.Bn), vmin=vmin, vmax=vmax, **kwargs
)
)
plt.setp(axes[2].get_yticklabels(), visible=False)
# axes[2].set_title('Residual Bn', fontsize=15)
axes[0].set_ylabel(r"$\theta$", fontsize=14)
# axes[1].set_ylabel(r'$\theta$', fontsize=14)
# axes[2].set_ylabel(r'$\theta$', fontsize=14)
axes[0].set_xlabel(r"$\phi$", fontsize=14)
axes[1].set_xlabel(r"$\phi$", fontsize=14)
axes[2].set_xlabel(r"$\phi$", fontsize=14)
else:
raise ValueError("unsupported option for source")
fig.subplots_adjust(right=0.85)
cbar_ax = fig.add_axes([0.86, 0.10, 0.04, 0.8])
fig.colorbar(obj[0], cax=cbar_ax)
return obj
# Bmod plot
[docs] def plot(
self, engine="plotly", scalars="bn", fig=None, ax=None, show=True, **kwargs
):
"""Plot 3D field on the plasma from FOCUS ouput
Args:
engine (str, optional): 3D plot engine, one of ["mayavi", "pyplot", "plotly"]. Defaults to "plotly".
scalars (str, optional): Scalar function on the surface, one of ["Bn", "B", "plas_Bn"] or 2D arrays. Defaults to "bn".
fig (, optional): Figure to be plotted on. Defaults to None.
ax (, optional): Axis to be plotted on. Defaults to None.
show (bool, optional): If show the plotly figure immediately. Defaults to True.
kwargs: optional keyword arguments for plotting.
Raises:
ValueError: Plot engine should be one value of ["mayavi", "pyplot", "plotly"].
"""
full = True if self.IsSymmetric == 0 else False
xsurf = map_matrix(self.xsurf, first=full)
ysurf = map_matrix(self.ysurf, first=full)
zsurf = map_matrix(self.zsurf, first=full)
label = "plasma"
if scalars.lower() == "bn":
scalars = self.Bn
label = "Bn"
elif scalars.lower() == "plas_bn":
scalars = self.plas_Bn
label = "Bn_plasma"
elif scalars.lower() == "b":
scalars = np.sqrt(self.Bx ** 2 + self.By ** 2 + self.Bz ** 2)
label = "|B|"
scalars = map_matrix(scalars, first=full)
if engine == "pyplot":
import matplotlib.pyplot as plt
# plot in matplotlib.pyplot
if ax is None or ax.name != "3d":
fig = plt.figure()
ax = fig.add_subplot(111, projection="3d")
kwargs.setdefault("cmap", "coolwarm")
ax.plot_surface(xsurf, ysurf, zsurf, c=scalars, **kwargs)
ax.set_title(label)
fig.colorbar()
elif engine == "mayavi":
# plot 3D surface in mayavi.mlab
from mayavi import mlab # to overrid plt.mlab
kwargs.setdefault("colormap", "coolwarm")
mlab.mesh(xsurf, ysurf, zsurf, scalars=scalars, **kwargs)
mlab.colorbar(title=label)
elif engine == "plotly":
import plotly.graph_objects as go
kwargs.setdefault("colorbar", go.surface.ColorBar(title=label))
if fig is None:
fig = go.Figure()
fig.add_trace(
go.Surface(x=xsurf, y=ysurf, z=zsurf, surfacecolor=scalars, **kwargs)
)
fig.update_layout(scene_aspectmode="data")
if show:
fig.show()
else:
raise ValueError("Invalid engine option {pyplot, mayavi, noplot}")
# write vtk
[docs] def toVTK(self, name=None, full=False, **kwargs):
"""Save surface and magnetic field data into VTK file
Arguments:
name -- string, VTK file name. default: None, if None, using self.filename
full -- boolean, if save the entire torus
**kwargs -- external data will be saved in VTK file
Return:
VTK file name
"""
from pyevtk.hl import gridToVTK
# automatically get a file name, focus_*.h5 -> vtk_*
if name is None:
name = self.filename[:-3].replace("focus_", "vtk_")
xx = self.xsurf
yy = self.ysurf
zz = self.zsurf
Bx = self.Bx
By = self.By
Bz = self.Bz
Bn = self.Bn
plas_Bn = self.plas_Bn
toroidal = full
xx = np.atleast_3d(map_matrix(xx, first=toroidal))
yy = np.atleast_3d(map_matrix(yy, first=toroidal))
zz = np.atleast_3d(map_matrix(zz, first=toroidal))
Bn = np.atleast_3d(map_matrix(Bn, first=toroidal))
plas_Bn = np.atleast_3d(map_matrix(plas_Bn, first=toroidal))
B = (
np.atleast_3d(map_matrix(Bx, first=toroidal)),
np.atleast_3d(map_matrix(By, first=toroidal)),
np.atleast_3d(map_matrix(Bz, first=toroidal)),
)
if full:
pass
data = {"Bn": Bn, "plas_Bn": plas_Bn, "B": B}
data.update(kwargs)
return gridToVTK(name, xx, yy, zz, pointData=data)
[docs] def curvature(
self,
iteration=True,
vlines=True,
shift_ind=0,
icoil=1,
NS=128,
):
"""Curvature ploting for the FOCUS-spline paper [arXiv:2107.02123] by N. Lonigro
Args:
iteration (bool, optional): [description]. Defaults to True.
vlines (bool, optional): [description]. Defaults to True.
shift_ind (int, optional): [description]. Defaults to 0.
icoil (int, optional): [description]. Defaults to 1.
NS (int, optional): [description]. Defaults to 128.
Returns:
[type]: [description]
"""
import math
# get figure
max_ind = -1
min_ind = -1
# get iteration data
if iteration:
abscissa = np.arange(NS)
_xlabel = "Segment number"
else:
abscissa = self.evolution[0, :] # be careful; DF is not saving wall-time
_xlabel = "wall time [Second]"
# get data
data = object.__getattribute__(self, "curvature_{:}".format(icoil))
data_s = object.__getattribute__(self, "straight_{:}".format(icoil))
if (data_s[0] == 0) and (data_s[NS - 1] != 0) and (data_s[k + 1] == 0):
max_ind = 0
if (data_s[0] == 0) and (data_s[1] != 0) and (data_s[k - 1] == 0):
min_ind = 0
if (data_s[NS - 1] == 0) and (data_s[NS - 2] != 0) and (data_s[0] == 0):
max_ind = NS - 1
if (data_s[NS - 1] == 0) and (data_s[0] != 0) and (data_s[NS - 2] == 0):
min_ind = NS - 1
for k in range(1, len(data_s) - 2):
if (data_s[k] == 0) and (data_s[k - 1] != 0) and (data_s[k + 1] == 0):
max_ind = np.max([max_ind, k])
if (
(data_s[k] == 0)
and (data_s[k + 1] != 0)
and (data_s[k - 1] == 0)
and min_ind != -1
):
min_ind = np.min([min_ind, k])
if (
(data_s[k] == 0)
and (data_s[k + 1] != 0)
and (data_s[k - 1] == 0)
and min_ind == -1
):
min_ind = k
if max_ind > min_ind and shift_ind == 0:
shift_ind = math.floor(0.5 * (NS - max_ind - min_ind))
if max_ind < min_ind and shift_ind == 0:
shift_ind = math.floor(0.5 * NS + 0.5 * (NS - max_ind - min_ind))
return [shift_ind, max_ind - min_ind]