def read_fiesta_equilibrium(filepath, first_wall=None):
"""
Current versions of the equilibria are stored in
`/compass/Shared/Common/COMPASS-UPGRADE/RP1 Design/Equilibria/v3.1`
:param filepath: Path to fiesta g-file equilibria
:param first_wall: Path to datafile with limiter line. (Fiesta doesn't store limiter contour into g-file).
If `None` IBA limiter v 3.1 is taken.
:return: Equilibrium: Instance of `Equilibrium`
"""
resource_package = 'pleque'
with open(filepath, 'r') as f:
data = read(f)
ds = data_as_ds(data)
# If there are some limiter data. Use them as and limiter.
if 'r_lim' in ds and 'z_lim' in ds and ds.r_lim.size > 3 and first_wall is None:
first_wall = np.stack((ds.r_lim.values, ds.z_lim.values)).T
if first_wall is None:
print('--- No limiter specified. The IBA v3.1 limiter will be used.')
first_wall = 'resources/limiter_v3_1_iba_v2.dat'
first_wall = pkg_resources.resource_filename(resource_package,
first_wall)
if isinstance(first_wall, str):
first_wall = np.loadtxt(first_wall)
eq = Equilibrium(ds, first_wall=first_wall)
eq._Ip = ds.attrs['cpasma']
return eq
def test_cocos_consistency(geqdsk_file, cocos):
print('COCOS: {}'.format(cocos))
equilibrium = read_geqdsk(geqdsk_file, cocos=cocos)
with open(geqdsk_file, 'r') as f:
eq_dict = read_geqdsk_as_dict(f)
eq_xr = data_as_ds(eq_dict)
eq_xr.psi.values = eq_xr.psi.values * (2 * np.pi)
eq_xr.pprime.values /= (2 * np.pi)
eq_xr.FFprime.values /= (2 * np.pi)
fw = np.stack(
(eq_xr['r_lim'].values, eq_xr['z_lim'].values)).T # first wall
equilibrium2 = Equilibrium(eq_xr, fw, cocos=cocos + 10)
sigma_Ip = np.sign(equilibrium.I_plasma)
sigma_B0 = np.sign(equilibrium.F0)
cocos_dict = equilibrium._cocosdic
Rax = equilibrium.magnetic_axis.R[0]
Zax = equilibrium.magnetic_axis.Z[0]
assert sigma_B0 * sigma_Ip == np.sign(
equilibrium.q(psi_n=0.5)) * cocos_dict['sigma_pol']
assert sigma_Ip == cocos_dict['sigma_Bp'] * equilibrium._psi_sign
assert np.sign(equilibrium.F(R=Rax, Z=Zax)) == sigma_B0
assert np.sign(equilibrium.B_tor(R=Rax, Z=Zax)) == sigma_B0
assert np.sign(equilibrium.j_tor(r=0.1, theta=0)) == sigma_Ip
assert np.sign(equilibrium.psi(psi_n=1) - equilibrium.psi(
psi_n=0)) == sigma_Ip * cocos_dict['sigma_Bp']
assert np.sign(equilibrium.tor_flux(r=0.1, theta=0)) == sigma_B0
assert np.sign(
equilibrium.pprime(psi_n=0.5)) == -sigma_Ip * cocos_dict['sigma_Bp']
assert np.sign(equilibrium.q(
psi_n=0.5)) == sigma_Ip * sigma_B0 * cocos_dict['sigma_pol']
assert np.isclose(
equilibrium.j_tor(R=equilibrium.magnetic_axis.R + 0.5,
Z=equilibrium.magnetic_axis.Z),
equilibrium2.j_tor(R=equilibrium2.magnetic_axis.R + 0.5,
Z=equilibrium2.magnetic_axis.Z))
assert np.isclose(equilibrium.q(psi_n=0.5), equilibrium2.q(psi_n=0.5))
assert np.isclose(equilibrium.I_plasma, equilibrium2.I_plasma, atol=100)
def plot_overview(eq: Equilibrium):
import matplotlib.pyplot as plt
r = np.linspace(eq.R_min, eq.R_max, 100)
z = np.linspace(eq.Z_min, eq.Z_max, 120)
psi = eq.psi(R=r, Z=z)
plt.figure(figsize=(8, 4))
plt.subplot(131)
plt.contour(r, z, psi, 20)
# plt.plot(eq._lcfs[:, 0], eq._lcfs[:, 1], label='lcfs')
plt.plot(eq.lcfs.R, eq.lcfs.Z, label='lcfs')
if eq._first_wall is not None:
plt.plot(eq._first_wall[:, 0], eq._first_wall[:, 1], 'k')
plt.plot(eq._mg_axis[0], eq._mg_axis[1], 'o', color='b', markersize=10)
plt.plot(eq._x_point[0], eq._x_point[1], 'x', color='r', markersize=10)
plt.plot(eq._x_point2[0], eq._x_point2[1], 'x', color='r', markersize=10)
return_axis = plt.gca()
plt.title(r'$\psi$')
plt.gca().set_aspect('equal')
plt.subplot(132)
cs = plt.contour(r, z, eq.B_pol(R=r, Z=z), 20)
plt.clabel(cs, inline=1)
plt.plot(eq._lcfs[:, 0], eq._lcfs[:, 1], label='lcfs')
if eq._first_wall is not None:
plt.plot(eq._first_wall[:, 0], eq._first_wall[:, 1], 'k')
plt.title(r'$B_\mathrm{pol}$')
plt.gca().set_aspect('equal')
plt.subplot(133)
cs = plt.contour(r, z, eq.B_tor(R=r, Z=z), 20)
plt.clabel(cs, inline=1)
plt.plot(eq._lcfs[:, 0], eq._lcfs[:, 1], label='lcfs')
if eq._first_wall is not None:
plt.plot(eq._first_wall[:, 0], eq._first_wall[:, 1], 'k')
plt.title(r'$B_\mathrm{tor}$')
plt.gca().set_aspect('equal')
# number of points
n = 100
# module automaticaly identify the type of the input:
midplane = eq.coordinates(r=np.linspace(-0.3, 0.3, n), theta=np.zeros(n))
fig, axs = plt.subplots(3, 1, sharex=True)
ax = axs[0]
# Profile of toroidal field:
ax.plot(midplane.r, eq.B_tor(midplane))
# Profile of poloidal field:
ax.plot(midplane.r, eq.B_pol(midplane))
return return_axis
def sal_jet(pulse, timex=47.0, time_unit="s"):
"""
Main loading routine, based on simple access layer, loads ppf data, calculates derivatives
:param pulse: JET pulse number
:param timex: time of slice
:param time_unit: (str) "s" or "ms"
:return: equilibrium
"""
if time_unit.lower() == "s":
time_factor = 1.
elif time_unit.lower() == "ms":
time_factor = 1000.
else:
raise ValueError("Unknown `time_unit`.")
data_path = '/pulse/{}/ppf/signal/jetppf/efit/{}:{}'
# default sequence
sequence = 0
# obtain psi data (reshape, transpose) and time axis
packed_psi = sal.get(data_path.format(pulse, 'psi', sequence))
psi = packed_psi
psi.data = packed_psi.data[:, :].reshape(len(packed_psi.dimensions[0]), 33,
33)
psi.data = np.swapaxes(psi.data, 1, 2)
time = packed_psi.dimensions[0].data
# psi grid axis
r = sal.get(data_path.format(pulse, 'psir', sequence)).data
z = sal.get(data_path.format(pulse, 'psiz', sequence)).data
# pressure profile
pressure = sal.get(data_path.format(pulse, 'p', sequence))
psi_n = pressure.dimensions[1].data
# f-profile
f = sal.get(data_path.format(pulse, 'f', sequence))
# q-profile
q = sal.get(data_path.format(pulse, 'q', sequence))
# calculate pprime and FFprime
deltapsi = deltapsi_calc(pulse)
pprime = pprime_calc(pressure, deltapsi, len(psi_n))
FFprime = FFprime_calc(f, deltapsi, len(psi_n))
#create dataset
dst = xr.Dataset(
{
'psi': (['time', 'R', 'Z'], psi.data),
'pressure': (['time', 'psi_n'], pressure.data),
'pprime': (['time', 'psi_n'], pprime),
'F': (['time', 'psi_n'], f.data),
'FFprime': (['time', 'psi_n'], FFprime),
'q': (['time', 'psi_n'], q.data),
'R': (['R'], r),
'Z': (['Z'], z),
},
coords={
'time': time,
'psi_n': psi_n,
})
# select desired time
ds = dst.sel(time=timex / time_factor, method='nearest')
# try to load limiter from ppfs
try:
limiter_r = sal.get(data_path.format(pulse, 'rlim', sequence)).data.T
limiter_z = sal.get(data_path.format(pulse, 'zlim', sequence)).data.T
except NodeNotFound:
limiter_r = sal.get(data_path.format(94508, 'rlim', sequence)).data.T
limiter_z = sal.get(data_path.format(94508, 'zlim', sequence)).data.T
print("Limiter points not present in #{}, loaded from #94508".format(
pulse))
limiter = np.column_stack([limiter_r, limiter_z])
# create pleque equilibrium
eq = Equilibrium(ds, limiter)
return eq
def read(ods: omas.ODS, time=None, time_unit="s"):
"""
:param ods: OMAS object with description of `wall` and `equilibrium`
:param time: (int, float or None) Time of required equilibria. The nearest time slice is taken. If `None`
the first time slice is taken.
:param time_unit: (str) "s" or "ms"
:return: Instance of `Equilibrioum` from the `OMAS` object keeping the IMAS data structure.
"""
if time_unit.lower() == "s":
time_factor = 1.
elif time_unit.lower() == "ms":
time_factor = 1000.
else:
raise ValueError("Unknown `time_unit`.")
try:
shot = ods["info"]["shot"]
except:
shot = ods["dataset_description"]["data_entry"]["pulse"]
if "wall" not in ods:
try:
# todo: (Need the latest OMAS)
from omas.omas_physics import add_wall
omas.add_wall(omas)
except ImportError:
print("The newest OMAS is required to add wall to IDS.")
except KeyError:
pass
# Reading wall:
try:
r_wall = ods["wall"]["description_2d"][0]["limiter"]["unit"][0][
"outline"]["r"]
z_wall = ods["wall"]["description_2d"][0]["limiter"]["unit"][0][
"outline"]["z"]
limiter = np.stack((r_wall, z_wall)).T
except ValueError:
limiter = None
# Times
ods_times = ods["equilibrium"]["time"]
if time is None:
time_idx = 0
else:
time_idx = np.argmin(np.abs(np.array(ods_times) - time / time_factor))
# Plasma boundary (LCFS)
rbnd = ods["equilibrium"]["time_slice"][time_idx]["boundary"]["outline"][
"r"]
zbnd = ods["equilibrium"]["time_slice"][time_idx]["boundary"]["outline"][
"z"]
psibnd = ods["equilibrium"]["time_slice"][time_idx]["global_quantities"][
"psi_boundary"]
# Magnetic axis
rmgax = ods["equilibrium"]["time_slice"][time_idx]["global_quantities"][
"magnetic_axis"]["r"]
zmgax = ods["equilibrium"]["time_slice"][time_idx]["global_quantities"][
"magnetic_axis"]["z"]
psimgax = ods["equilibrium"]["time_slice"][time_idx]["global_quantities"][
"psi_axis"]
# 1D profiles
psi_1d = ods["equilibrium"]["time_slice"][time_idx]["profiles_1d"]["psi"]
F = ods['equilibrium.time_slice'][time_idx]['profiles_1d.f']
pressure = ods['equilibrium.time_slice'][time_idx]['profiles_1d.pressure']
FFp = ods['equilibrium.time_slice'][time_idx]['profiles_1d.f_df_dpsi']
pprime = ods['equilibrium.time_slice'][time_idx][
'profiles_1d.dpressure_dpsi']
q = ods['equilibrium.time_slice'][time_idx]['profiles_1d.q']
psi_n = (psi_1d - psi_1d[0]) / (psibnd - psimgax)
# 2D profiles:
# todo: resolve this (!)
grid_idx = 0
gridtype = ods["equilibrium"]["time_slice"][time_idx]["profiles_2d"][0][
"grid_type"]["index"]
if gridtype != 1:
raise ValueError("Only rectangular is supported at the moment!")
# if 1 not in gridtype:
# raise ValueError("Only rectangular is supported at the moment!")
# else:
# grid_idx = gridtype.index(1)
r_grid = ods["equilibrium"]["time_slice"][time_idx]["profiles_2d"][
grid_idx]["grid"]["dim1"]
z_grid = ods["equilibrium"]["time_slice"][time_idx]["profiles_2d"][
grid_idx]["grid"]["dim2"]
psi = ods["equilibrium"]["time_slice"][time_idx]["profiles_2d"][grid_idx][
"psi"]
if psi.shape != (len(r_grid), len(z_grid)):
psi = psi.T
# x-array Dataset with time data:
ds = xr.Dataset(
{
'psi': (['R', 'Z'], psi),
'pressure': (['psi_n'], pressure),
'pprime': (['psi_n'], pprime),
'F': (['psi_n'], F),
'FFprime': (['psi_n'], FFp),
'q': (['psi_n'], q),
},
coords={
'time': ods_times[time_idx],
'R': r_grid,
'Z': z_grid,
'psi_n': psi_n,
},
attrs={
'shot': shot,
'time_unit': time_unit,
})
eq = Equilibrium(ds, limiter)
return eq
def write(equilibrium: Equilibrium,
grid_1d=None,
grid_2d=None,
gridtype=1,
ods=None,
cocosio=3):
"""
Function saving contents of equilibrium into the omas data structure.
:param equilibrium: Equilibrium object
:param grid_1d: Coordinate object with 1D grid (linearly spaced array over psi_n). It is used to save 1D equilibrium
characteristics into the omas data object. If is None, linearly spaced vector over psi_n inrange [0,1] with 200
points is used.
:param grid_2d: Coordinate object with 2D grid (can be the whole reconstruction space). It is used to save 2D
equilibrium profiles. If None, default Equilibrium.grid() is used to generate the parameter.
:param gridtype: Grid type specification for omas data structure. 1...rectangular.
(Any other grid type is not supported at them moment!)
:param ods: ods object to save the equilibrium into. If None, new ods object is created.
:return:
"""
# todo: Well for now it is just what ASCOT eats, we will extend this when we pile up more usage....
if ods is None:
ods = omas.ODS(cocosio=cocosio)
if grid_1d is None:
grid_1d = equilibrium.coordinates(psi_n=np.linspace(0, 1, 200))
if grid_2d is None:
grid_2d = equilibrium.grid(resolution=(1e-3, 1e-3), dim="step")
shot_time = equilibrium.time
if equilibrium.time_unit == "ms":
shot_time *= 1e3
elif equilibrium.time_unit != "s":
print(
"WARNING: unknown time unit ({}) is used. Seconds will be used insted for saving."
.format(equilibrium.time_unit))
# ods["info"]["shot"] = equilibrium.shot
ods["dataset_description"]["data_entry"]["pulse"] = equilibrium.shot
# fill the wall part
ods["wall"]["ids_properties"]["homogeneous_time"] = 1
# to MT: is this change correct?
ods["wall"]["time"] = np.array(shot_time, ndmin=1)
ods["wall"]["description_2d"][0]["limiter"]["unit"][0]["outline"][
"r"] = equilibrium.first_wall.R
ods["wall"]["description_2d"][0]["limiter"]["unit"][0]["outline"][
"z"] = equilibrium.first_wall.Z
############################
# The Equilibrium Part
############################
# time slices
ods["equilibrium"]["ids_properties"]["homogeneous_time"] = 1
# to MT: is this change correct?
ods["equilibrium"]["time"] = np.array(shot_time, ndmin=1)
# Vacuum
# todo: add vacuum Btor, not in equilibrium
# As R0 use magnetic axis. Isn't better other definition of R0?
# R0 = np.array((np.max(equilibrium.first_wall.R) - np.min(equilibrium.first_wall.R))/2, ndim=1)
R0 = equilibrium.magnetic_axis.R[0]
F0 = equilibrium.BvacR
B0 = F0 / R0 * np.ones_like(ods["equilibrium"]["time"])
ods["equilibrium"]["vacuum_toroidal_field"][
"b0"] = B0 # vacuum B tor at Rmaj
ods["equilibrium"]["vacuum_toroidal_field"][
"r0"] = R0 # vacuum B tor at Rmaj
# time slice time
ods["equilibrium"]["time_slice"][0]["time"] = shot_time
# plasma boundary (lcfs)
ods["equilibrium"]["time_slice"][0]["boundary"]["outline"][
"r"] = equilibrium.lcfs.R
ods["equilibrium"]["time_slice"][0]["boundary"]["outline"][
"z"] = equilibrium.lcfs.Z
ods["equilibrium"]["time_slice"][0]["global_quantities"][
"psi_boundary"] = equilibrium._psi_lcfs
# Magnetic axis
ods["equilibrium"]["time_slice"][0]["global_quantities"]["magnetic_axis"][
"r"] = equilibrium.magnetic_axis.R[0]
ods["equilibrium"]["time_slice"][0]["global_quantities"]["magnetic_axis"][
"z"] = equilibrium.magnetic_axis.Z[0]
ods["equilibrium"]["time_slice"][0]["global_quantities"][
"psi_axis"] = equilibrium._psi_axis
# define the 1 and 2d grids
# 1d profiles
ods["equilibrium"]["time_slice"][0]["profiles_1d"][
"psi"] = equilibrium.psi(grid_1d)
ods["equilibrium"]["time_slice"][0]["profiles_1d"][
"rho_tor"] = equilibrium.tor_flux(grid_1d)
ods['equilibrium.time_slice'][0]['profiles_1d.f'] = equilibrium.F(grid_1d)
ods['equilibrium.time_slice'][0][
'profiles_1d.pressure'] = equilibrium.pressure(grid_1d)
ods['equilibrium.time_slice'][0][
'profiles_1d.f_df_dpsi'] = equilibrium.FFprime(grid_1d)
ods['equilibrium.time_slice'][0][
'profiles_1d.dpressure_dpsi'] = equilibrium.pprime(grid_1d)
ods['equilibrium.time_slice'][0]['profiles_1d.q'] = equilibrium.q(grid_1d)
# get surface volumes and areas
surface_volume = np.zeros_like(grid_1d.psi)
surface_area = np.zeros_like(grid_1d.psi)
for i in range(grid_1d.psi.shape[0]):
psin_tmp = grid_1d.psi_n[i]
surface = 0
if not psin_tmp == 1:
coord_tmp = equilibrium.coordinates(psi_n=psin_tmp)
surface = equilibrium._flux_surface(coord_tmp)
elif psin_tmp == 1:
surface = [equilibrium._as_fluxsurface(equilibrium.lcfs)]
# todo: really 0 if open?
# todo: fix this 'surface[0].closed' ... maybe calculate it from boundary
if len(surface) > 0 and surface[0].closed:
surface_volume[i] = surface[0].volume
surface_area[i] = surface[0].area
else:
surface_volume[i] = 0
surface_area[i] = 0
ods["equilibrium"]["time_slice"][0]["profiles_1d"][
"volume"] = surface_volume
ods["equilibrium"]["time_slice"][0]["profiles_1d"]["area"] = surface_area
# 2D profiles
ods["equilibrium"]["time_slice"][0]["profiles_2d"][0]["grid_type"][
"index"] = gridtype
ods["equilibrium"]["time_slice"][0]["profiles_2d"][0]["grid"][
"dim1"] = grid_2d.R
ods["equilibrium"]["time_slice"][0]["profiles_2d"][0]["grid"][
"dim2"] = grid_2d.Z
ods["equilibrium"]["time_slice"][0]["profiles_2d"][0][
"psi"] = equilibrium.psi(grid_2d).T
ods["equilibrium"]["time_slice"][0]["profiles_2d"][0][
"b_field_tor"] = equilibrium.B_tor(grid_2d).T
ods['equilibrium.time_slice'][0][
'global_quantities.ip'] = equilibrium.I_plasma
return ods
def read_gfile(g_file: str, limiter: str = None):
from pleque.io import _geqdsk
import numpy as np
import xarray as xr
from pleque.core import Equilibrium
from scipy.interpolate import UnivariateSpline
with open(g_file, 'r') as f:
eq_gfile = _geqdsk.read(f)
psi = eq_gfile['psi']
r = eq_gfile['r'][:, 0]
z = eq_gfile['z'][0, :]
pressure = eq_gfile['pressure']
F = eq_gfile['F']
q = eq_gfile['q']
psi_n = np.linspace(0, 1, len(F))
eq_ds = xr.Dataset(
{
'psi': (['R', 'Z'], psi),
'pressure': ('psi_n', pressure),
'F': ('psi_n', F)
},
coords={
'R': r,
'Z': z,
'psi_n': psi_n
})
if limiter is not None:
lim = np.loadtxt(limiter)
print(lim.shape)
else:
lim = None
eq = Equilibrium(eq_ds, first_wall=lim)
eq._geqdsk = eq_gfile
eq._q_spl = UnivariateSpline(psi_n, q, s=0, k=3)
eq._dq_dpsin_spl = eq._q_spl.derivative()
eq._q_anideriv_spl = eq._q_spl.antiderivative()
def q(self,
*coordinates,
R=None,
Z=None,
psi_n=None,
coord_type=None,
grid=True,
**coords):
if R is not None and Z is not None:
psi_n = self.psi_n(R=R, Z=Z, grid=grid)
return self._q_spl(psi_n)
def diff_q(self: eq,
*coordinates,
R=None,
Z=None,
psi_n=None,
coord_type=None,
grid=True,
**coords):
"""
:param self:
:param coordinates:
:param R:
:param Z:
:param psi_n:
:param coord_type:
:param grid:
:param coords:
:return: Derivative of q with respect to psi.
"""
if R is not None and Z is not None:
psi_n = self.psi_n(R=R, Z=Z, grid=grid)
return self._dq_dpsin_spl(psi_n) * self._diff_psiN
def tor_flux(self: eq,
*coordinates,
R=None,
Z=None,
psi_n=None,
coord_type=None,
grid=True,
**coords):
if R is not None and Z is not None:
psi_n = self.psi_n(R=R, Z=Z, grid=grid)
return eq._q_anideriv_spl(psi_n) * (1 / self._diff_psi_n)
# eq.q = q
# eq.diff_q = diff_q
# eq.tor_flux = tor_flux
Equilibrium.q = q
Equilibrium.diff_q = diff_q
Equilibrium.tor_flux = tor_flux
return eq
def show_qprofiles(g_file: str, eq: Equilibrium):
# from tokamak.formats import geqdsk
from pleque.io._geqdsk import read
import matplotlib.pyplot as plt
with open(g_file, 'r') as f:
eq_gfile = read(f)
q = eq_gfile['q']
psi_n = np.linspace(0, 1, len(q))
print(eq_gfile.keys())
psin_axis = np.linspace(0, 1, 100)
r = np.linspace(eq.R_min, eq.R_max, 100)
z = np.linspace(eq.Z_min, eq.Z_max, 120)
psi = eq.psi(R=r, Z=z)
plt.figure()
plt.subplot(121)
ax = plt.gca()
cs = ax.contour(r, z, psi, 30)
ax.plot(eq._lcfs[:, 0], eq._lcfs[:, 1], label='lcfs')
if eq._first_wall is not None:
plt.plot(eq._first_wall[:, 0], eq._first_wall[:, 1], 'k')
ax.plot(eq._mg_axis[0], eq._mg_axis[1], 'o', color='b')
ax.plot(eq._x_point[0], eq._x_point[1], 'x', color='r')
ax.plot(eq._x_point2[0], eq._x_point2[1], 'x', color='r')
ax.set_xlabel('R [m]')
ax.set_ylabel('Z [m]')
ax.set_aspect('equal')
ax.set_title(r'$\psi$')
plt.colorbar(cs, ax=ax)
psi_mod = eq.psi(psi_n=psin_axis)
tor_flux = eq.tor_flux(psi_n=psin_axis)
q_as_grad = np.gradient(tor_flux, psi_mod)
plt.subplot(322)
ax = plt.gca()
ax.plot(psi_n, np.abs(q), 'x', label='g-file (abs)')
ax.plot(psin_axis,
q_as_grad,
'-',
label=r'$\mathrm{d} \Phi/\mathrm{d} \psi$')
ax.plot(psin_axis, q_as_grad, '--', label='Pleque')
ax.legend()
ax.set_xlabel(r'$\psi_\mathrm{N}$')
ax.set_ylabel(r'$q$')
plt.subplot(324)
ax = plt.gca()
ax.plot(tor_flux, psi_mod, label='Toroidal flux')
ax.set_ylabel(r'$\psi$')
ax.set_xlabel(r'$\Phi$')
plt.subplot(326)
ax = plt.gca()
ax.plot(psi_n, eq.pprime(psi_n=psi_n) / 1e3)
ax.set_xlabel(r'$\psi_\mathrm{N}$')
ax.set_ylabel(r"$p' (\times 10^3)$")
ax2 = ax.twinx()
ax2.plot(psi_n, eq.FFprime(psi_n=psi_n), 'C1')
ax2.set_ylabel(r"$ff'$")
def plot_overview(eq: Equilibrium):
import matplotlib.pyplot as plt
r = np.linspace(eq.R_min, eq.R_max, 100)
z = np.linspace(eq.Z_min, eq.Z_max, 120)
psi = eq.psi(R=r, Z=z)
plt.figure(figsize=(8, 4))
plt.subplot(131)
plt.contour(r, z, psi, 20)
# plt.plot(eq._lcfs[:, 0], eq._lcfs[:, 1], label='lcfs')
plt.plot(eq.lcfs.R, eq.lcfs.Z, label='lcfs')
if eq._first_wall is not None:
plt.plot(eq._first_wall[:, 0], eq._first_wall[:, 1], 'k')
plt.plot(eq._mg_axis[0], eq._mg_axis[1], 'o', color='b', markersize=10)
plt.plot(eq._x_point[0], eq._x_point[1], 'x', color='r', markersize=10)
plt.plot(eq._x_point2[0], eq._x_point2[1], 'x', color='r', markersize=10)
return_axis = plt.gca()
plt.title(r'$\psi$')
plt.gca().set_aspect('equal')
plt.subplot(132)
cs = plt.contour(r, z, eq.B_pol(R=r, Z=z), 20)
plt.clabel(cs, inline=1)
plt.plot(eq._lcfs[:, 0], eq._lcfs[:, 1], label='lcfs')
if eq._first_wall is not None:
plt.plot(eq._first_wall[:, 0], eq._first_wall[:, 1], 'k')
plt.title(r'$B_\mathrm{pol}$')
plt.gca().set_aspect('equal')
plt.subplot(133)
cs = plt.contour(r, z, eq.B_tor(R=r, Z=z), 20)
plt.clabel(cs, inline=1)
plt.plot(eq._lcfs[:, 0], eq._lcfs[:, 1], label='lcfs')
if eq._first_wall is not None:
plt.plot(eq._first_wall[:, 0], eq._first_wall[:, 1], 'k')
plt.title(r'$B_\mathrm{tor}$')
plt.gca().set_aspect('equal')
# number of points
n = 100
# module automaticaly identify the type of the input:
midplane = eq.coordinates(r=np.linspace(-0.3, 0.3, n), theta=np.zeros(n))
fig, axs = plt.subplots(3, 1, sharex=True)
ax = axs[0]
# Profile of toroidal field:
ax.plot(midplane.r, eq.B_tor(midplane))
# Profile of poloidal field:
ax.plot(midplane.r, eq.B_pol(midplane))
# # ----
# plt.figure(figsize=(8, 4))
# plt.subplot(131)
# plt.pcolormesh(r, z, eq.B_R(R=r, Z=z).T)
# if eq._first_wall is not None:
# plt.plot(eq._first_wall[:, 0], eq._first_wall[:, 1], 'k')
# plt.title(r'$B_\mathrm{R}$')
# plt.gca().set_aspect('equal')
#
# plt.subplot(132)
# plt.pcolormesh(r, z, eq.B_Z(R=r, Z=z).T)
# if eq._first_wall is not None:
# plt.plot(eq._first_wall[:, 0], eq._first_wall[:, 1], 'k')
# plt.title(r'$B_\mathrm{Z}$')
# plt.gca().set_aspect('equal')
#
# plt.subplot(133)
# plt.pcolormesh(r, z, np.sqrt(eq.B_R(R=r, Z=z) ** 2 + eq.B_Z(R=r, Z=z) ** 2).T)
# if eq._first_wall is not None:
# plt.plot(eq._first_wall[:, 0], eq._first_wall[:, 1], 'k')
# plt.title(r'$B_\mathrm{pol}$')
# plt.gca().set_aspect('equal')
#
# psi_n = np.linspace(0, 1, 100)
#
# plt.figure()
# plt.subplot(211)
# ax = plt.gca()
# ax.plot(psi_n, eq.pressure(psi_n=psi_n), 'C1')
# ax.set_xlabel(r'$\psi_\mathrm{N}$')
# ax.set_ylabel(r'$p [Pa]$', color='C1')
#
# ax2 = ax.twinx()
# ax2.plot(psi_n, eq.pprime(psi_n=psi_n), color='C2')
# ax2.set_ylabel(r"$p'$", color='C2')
#
# plt.subplot(212)
# ax = plt.gca()
# ax.plot(psi_n, eq.F(psi_n=psi_n), 'C1')
# ax.set_xlabel(r'$\psi_\mathrm{N}$')
# ax.set_ylabel(r'$f$ ', color='C1')
#
# ax2 = ax.twinx()
# ax2.plot(psi_n, eq.FFprime(psi_n=psi_n), 'C2')
# ax2.set_ylabel(r"$ff'$ ", color='C2')
return return_axis
def plot_psi_derivatives(eq: Equilibrium):
import matplotlib.pyplot as plt
r = np.linspace(eq.R_min, eq.R_max, 100)
z = np.linspace(eq.Z_min, eq.Z_max, 120)
psi = eq.psi(R=r, Z=z)
fig, axs = plt.subplots(1, 4, sharex=True, sharey=True, figsize=(8, 4))
ax = axs[0]
ax.contour(r, z, psi.T, 20)
ax.plot(eq._lcfs[:, 0], eq._lcfs[:, 1], label='lcfs')
if eq._first_wall is not None:
ax.plot(eq._first_wall[:, 0], eq._first_wall[:, 1], 'k')
ax.plot(eq._mg_axis[0], eq._mg_axis[1], 'o', color='b', markersize=10)
ax.plot(eq._x_point[0], eq._x_point[1], 'x', color='r', markersize=10)
ax.plot(eq._x_point2[0], eq._x_point2[1], 'x', color='r', markersize=10)
psi_axis = ax
plot_extremes(eq, psi_axis)
def psi_xysq_func(r, z):
return eq._spl_psi(r, z, dx=1, dy=0, grid=True) ** 2 \
+ eq._spl_psi(r, z, dx=0, dy=1, grid=True) ** 2
psi_xysq = psi_xysq_func(r, z)
psi_xyopt = (eq._spl_psi(r, z, dx=1, dy=1, grid=True))**2
plt.title(r'$\psi$')
ax.set_aspect('equal')
ax = axs[1]
cl = ax.contour(r, z, psi_xysq.T, np.linspace(0, 0.1, 50))
# plt.colorbar(cl)
ax.plot(eq._lcfs[:, 0], eq._lcfs[:, 1], label='lcfs')
if eq._first_wall is not None:
ax.plot(eq._first_wall[:, 0], eq._first_wall[:, 1], 'k')
ax.set_title(r'$\partial_x \psi^2 + \partial_y \psi^2$')
ax.set_aspect('equal')
ax = axs[2]
cl = ax.contour(r, z, psi_xyopt.T, np.linspace(0, 1, 50))
# plt.colorbar(cl)
ax.plot(eq._lcfs[:, 0], eq._lcfs[:, 1], label='lcfs')
if eq._first_wall is not None:
ax.plot(eq._first_wall[:, 0], eq._first_wall[:, 1], 'k')
ax.set_title(r'$\partial_{xy} \psi^2$')
ax.set_aspect('equal')
psi_xy = (eq._spl_psi(r, z, dx=1, dy=1, grid=True))**2
psi_xx = (eq._spl_psi(r, z, dx=2, dy=0, grid=True))
psi_yy = (eq._spl_psi(r, z, dx=0, dy=2, grid=True))
D = psi_xx * psi_yy - psi_xy
ax = axs[3]
cl = ax.contour(r, z, D.T, np.linspace(-50, 50, 50), cmap='PiYG')
plt.colorbar(cl)
ax.plot(eq._lcfs[:, 0], eq._lcfs[:, 1], label='lcfs')
if eq._first_wall is not None:
ax.plot(eq._first_wall[:, 0], eq._first_wall[:, 1], 'k')
ax.set_title(r'$D$')
ax.set_aspect('equal')