# modified from stellopt.pySTEL
import numpy as np
import xarray
import matplotlib.pyplot as plt
from .misc import trig2real
from .surface import FourSurf
__all__ = ["VMECout"]
[docs]class VMECout(object):
"""
Class to parse a VMEC wout file
:param filename: filename passed to xarray.open_dataset()
The entire dataset is stored in `self.wout` and you can access
to variables via `self.wout['iotaf'].values`.
The flux surfaces are all parsed as `FourSurf` classes and stored
in the list of `self.surface`.
Magnetic fields are Fourier transformed to real space
(only the stellarator symmetric part, `bmnc`) and store in `self.data['b]`.
`self.plot` has several options to plot the profiles, etc.
"""
def __init__(self, filename, **kwargs):
self.wout = xarray.open_dataset(filename)
self.data = {}
self.data["ns"] = int(self.wout["ns"].values)
self.data["nfp"] = int(self.wout["nfp"].values)
self.data["nu"] = int(self.wout["mpol"].values * 4)
self.data["nv"] = int(self.wout["ntor"].values * 4)
self.data["curpol"] = (
2.0 * np.pi / self.data["nfp"] * self.wout["rbtor"].values * 1e6
) # unit: Ampere
self.data["nflux"] = np.linspace(
0, 1, self.data["ns"]
) # np.ndarray((self.data['ns'],1))
self.data["theta"] = np.linspace(
0, 2 * np.pi, self.data["nu"]
) # np.ndarray((self.data['nu'],1))
self.data["zeta"] = np.linspace(
0, 2 * np.pi, self.data["nv"]
) # np.ndarray((self.data['nv'],1))
self.surface = []
self.data["b"] = []
for i in range(self.data["ns"]):
self.surface.append(
FourSurf(
xm=self.wout["xm"].values,
xn=self.wout["xn"].values,
rbc=self.wout["rmnc"][i].values,
rbs=np.zeros_like(self.wout["rmnc"][i].values),
zbs=self.wout["zmns"][i].values,
zbc=np.zeros_like(self.wout["zmns"][i].values),
)
)
self.data["b"].append(
trig2real(
self.data["theta"],
self.data["zeta"],
self.wout["xm_nyq"].values,
self.wout["xn_nyq"].values / self.data["nfp"],
self.wout["bmnc"][i].values,
)
)
return
[docs] def plot(self, plot_name="none", ax=None, **kwargs):
"""Plot various VMEC quantities
Args:
plot_name (str, optional): The quantity to be plotted, should be one of
iota, q, pressue, <Buco>, <Bvco>, <jcuru>, <jcurv>,
<j.B>, LPK, none. Defaults to 'none'.
ax (Matplotlib axis, optional): The Matplotlib axis to be plotted on. Defaults to None.
"""
if ax is None:
fig, ax = plt.subplots()
else:
fig = ax.get_figure()
if plot_name == "none":
print("You can plot: iota, q, pressue, <Buco>, <Bvco>, <jcuru>, <jcurv>, ")
print(" <j.B>, LPK")
elif plot_name == "iota":
ax.plot(self.data["nflux"], self.wout["iotaf"].values, **kwargs)
ax.set_xlabel("Normalized Flux")
ax.set_ylabel("iota")
ax.set_title("Rotational Transform")
# ax.set(xlabel='s',ylabel='iota',aspect='square')
elif plot_name == "q":
ax.plot(self.data["nflux"], 1.0 / self.wout["iotaf"].values, **kwargs)
ax.set_xlabel("Normalized Flux")
ax.set_ylabel("q")
ax.set_title("Safety Factor")
elif plot_name == "pressure":
ax.plot(self.data["nflux"], self.wout["presf"].values / 1000, **kwargs)
ax.set_xlabel("Normalized Flux")
ax.set_ylabel("Pressure [kPa]")
ax.set_title("Pressure Profile")
elif plot_name == "<Buco>":
ax.plot(self.data["nflux"], self.wout["buco"].values, **kwargs)
ax.set_xlabel("Normalized Flux")
ax.set_ylabel("<B^u> [T]")
ax.set_title("Flux surface Averaged B^u")
elif plot_name == "<Bvco>":
ax.plot(self.data["nflux"], self.wout["bvco"].values, **kwargs)
ax.set_xlabel("Normalized Flux")
ax.set_ylabel("<B^v> [T]")
ax.set_title("Flux surface Averaged B^v")
elif plot_name == "<jcuru>":
ax.plot(self.data["nflux"], self.wout["jcuru"].values / 1000, **kwargs)
ax.set_xlabel("Normalized Flux")
ax.set_ylabel("<j^u> [kA/m^2]")
ax.set_title("Flux surface Averaged j^u")
elif plot_name == "<jcurv>":
ax.plot(self.data["nflux"], self.wout["jcurv"].values / 1000, **kwargs)
ax.set_xlabel("Normalized Flux")
ax.set_ylabel("<j^v> [kA/m^2]")
ax.set_title("Flux surface Averaged j^v")
elif plot_name == "<j.B>":
ax.plot(self.data["nflux"], self.wout["jdotb"].values / 1000, **kwargs)
ax.set_xlabel("Normalized Flux")
ax.set_ylabel("<j.B> [T*kA/m^2]")
ax.set_title("Flux surface Averaged j.B")
elif plot_name == "LPK":
self.surface[-1].plot(zeta=0, color="red", label=r"$\phi=0$", **kwargs)
self.surface[-1].plot(
zeta=0.5 * np.pi / self.data["nfp"],
color="green",
label=r"$\phi=0.25$",
**kwargs
)
self.surface[-1].plot(
zeta=np.pi / self.data["nfp"],
color="blue",
label=r"$\phi=0.5$",
**kwargs
)
ax.set_title("LPK Plot")
elif plot_name[0] == "-":
print(plot_name)
else:
return ax