def test_calculated_profiles(geqdsk_file):
"""
Octants: https://en.wikipedia.org/wiki/Octant_(solid_geometry)#/media/File:Octant_numbers.svg
:param geqdsk_file:
:return:
"""
eq = read_geqdsk(geqdsk_file, cocos=3)
with open(geqdsk_file, 'r') as f:
eq_dict = _geqdsk.read(f)
pressure = eq_dict['pres']
pprime = eq_dict['pprime']
F = eq_dict['F']
FFprime = eq_dict['FFprime']
psi_n = np.linspace(0, 1, len(pressure))
assert np.allclose(pressure, eq.pressure(psi_n=psi_n), atol=1)
assert np.allclose(pprime, eq.pprime(psi_n=psi_n), atol=100)
assert np.allclose(F, eq.F(psi_n=psi_n), atol=1e-2)
assert np.allclose(FFprime, eq.FFprime(psi_n=psi_n), atol=1e-3)
def test_profiles_integration(geqdsk_file):
with open(geqdsk_file, 'r') as f:
eq_in = _geqdsk.read(f)
psi_ax = eq_in['simagx']
psi_bnd = eq_in['sibdry']
p = eq_in['pres']
if np.isclose(p[-1], 0):
psi_n = np.linspace(0, 1, len(p), endpoint=True)
else:
psi_n = np.linspace(0, 1, len(p), endpoint=False)
pprime = eq_in['pprime']
pprime = xa.DataArray(pprime, [psi_n], ['psi_n'])
ffprime = eq_in['FFprime']
f = eq_in['F']
f0 = f[-1]
p_calc = pprime2p(pprime, psi_ax, psi_bnd)
f_calc = ffprime2f(ffprime, psi_ax, psi_bnd, f0)
assert np.allclose(p, p_calc)
assert np.allclose(f, f_calc)
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 load_gfile(g_file):
from pleque.io._geqdsk import read
eq_gfile = read(g_file)
psi = eq_gfile['psi']
r = eq_gfile['r'][:, 0]
z = eq_gfile['z'][0, :]
pressure = eq_gfile['pressure']
F = eq_gfile['F']
psi_n = np.linspace(0, 1, len(F))
eq_ds = xa.Dataset(
{
'psi': (['Z', 'R'], psi),
'pressure': ('psi_n', pressure),
'F': ('psi_n', F)
},
coords={
'R': r,
'Z': z,
'psi_n': psi_n
})
print(eq_ds)
if False:
import matplotlib.pyplot as plt
plt.figure()
plt.contour(r, z, psi)
plt.gca().set_aspect('equal')
return eq_ds
def read_write_geqdsk(file_in, file_out):
'''
Read the equilibrium file and then save it again. In principle,
newly written file may be according to g-file standard.
:param file_in:
:param file_out:
:return:
'''
with open(file_in, 'r') as f:
eq_in = _geqdsk.read(f)
with open(file_out, 'w') as f:
_geqdsk.write(eq_in, f)
def read_write_geqsk_flip_direction(file_in, file_out, flip=True, plot=False):
"""
Read possibly corrupted g-eqdsk file and change direction of B and I_plasma.
:param file_in: file name
:param file_out: file name
:param flip: bool, if true flip direction of I_plasma and B_tor.
:param plot: bool, if true plot
:return:
"""
with open(file_in, 'r') as f:
eq_in = _geqdsk.read(f)
if flip:
# Change of current -> change of psi
eq_in['psi'] = -eq_in['psi']
eq_in['simagx'] = -eq_in['simagx']
eq_in['sibdry'] = -eq_in['sibdry']
eq_in['pprime'] = -eq_in['pprime']
eq_in['FFprime'] = -eq_in['FFprime']
eq_in['cpasma'] = -eq_in['cpasma']
# Change of toroidal magnetic field
eq_in['F'] = -eq_in['F']
eq_in['bcentr'] = -eq_in['bcentr']
with open(file_out, 'w') as f:
_geqdsk.write(eq_in, f)
if plot:
eq_new = readers.read_geqdsk(file_out)
eq_new.plot_geometry()
plt.show()
if __name__ == '__main__':
from pleque.io import _geqdsk, _readgeqdsk
from pleque.tests.utils import get_test_equilibria_filenames
from collections import *
eqdsk_file = get_test_equilibria_filenames()[0]
# Read Matheeesek:
eq_1 = _readgeqdsk._readeqdsk(eqdsk_file)
# Read Ben:
with open(eqdsk_file, 'r') as f:
eq_2 = _geqdsk.read(f)
print(eq_1.keys())
print(eq_2.keys())
def compare(eq_1, eq_2):
'''
Compare content of eq_1 according to content of eq_2
:param eq_1:
:param eq_2:
:return:
'''
for k, v in eq_1.items():
if k in eq_2:
if isinstance(v, Iterable):
if len(v) == len(eq_2[k]):
print("{} OK".format(k))
else:
print("{} in both, but {} != {}".format(k, len(v), len(eq_2[k])))
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'$")