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 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'$")