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 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 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