def test_fiesta_gfile_vs_database():
tmp_dir = tempfile.TemporaryDirectory()
file_name = '{}/g0000.0000'.format(tmp_dir.name)
resource_package = "pleque.resources"
gfile_file = pkg_resources.resource_filename(resource_package, "test00.gfile")
nc_file = pkg_resources.resource_filename(resource_package, "test00.nc")
basedata = xr.load_dataset(nc_file).load()
eq_nc = pleque.Equilibrium(basedata, cocos=13)
eq_gfile = read_geqdsk(gfile_file)
# with open(gfile_file, "r") as f:
# eq_gfile = _geqdsk.read(f)
eq_nc.to_geqdsk(file_name, q_positive=True)
eq2 = read_geqdsk(file_name)
def compare_two(eq_1: pleque.Equilibrium, eq_2: pleque.Equilibrium):
assert np.isclose(np.abs(eq_1.I_plasma), np.abs(eq_2.I_plasma), atol=1e4, rtol=1e-3)
assert np.isclose(eq_1.pressure(psi_n=0.2)[0], eq_2.pressure(psi_n=0.2)[0], atol=3e3, rtol=1e-2)
assert np.isclose(np.abs(eq_1.BvacR), np.abs(eq_2.BvacR))
assert np.isclose(np.abs(eq_1.q(psi_n=0.5)[0]), np.abs(eq_2.q(psi_n=0.5)[0]), atol=1e-3, rtol=1e-3)
compare_two(eq_nc, eq_gfile)
compare_two(eq_gfile, eq2)
compare_two(eq_nc, eq2)
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_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 test_directions(geqdsk_file, cocos):
print('g-file: {}'.format(geqdsk_file))
print('COCOS: {}'.format(cocos))
direction = -1
eq = read_geqdsk(geqdsk_file, cocos=cocos)
mgax = eq.magnetic_axis
maxx = eq.Z_max
R0 = mgax.R + 0.5 * (maxx - mgax.R)
Z0 = mgax.Z - 0.06
p0 = eq.coordinates(R=R0, Z=Z0, phi=0)
tr = eq.trace_field_line(R=R0, Z=Z0, direction=direction)[0]
br = eq.B_R(R=R0, Z=Z0)
bz = eq.B_Z(R=R0, Z=Z0)
btor = eq.B_tor(R=R0, Z=Z0)
coef = 0.01
p1 = eq.coordinates(R=R0 + coef * br, Z=Z0 + coef * bz, phi=coef * btor)
dphi = tr.phi[1] - tr.phi[0]
dtheta = tr.theta[1] - tr.theta[0]
assert np.sign(dphi) == np.sign(p1.phi - p0.phi) * direction
assert np.sign(dtheta) == np.sign(p1.theta - p0.theta) * direction
assert np.isclose(tr.theta[0], tr.theta[-1], atol=0.1, rtol=0.1)
dphidtheta = np.sign(dphi * dtheta)
my_dir = eq._cocosdic['sigma_pol'] * np.sign(eq.I_plasma) * np.sign(eq.F0)
# my_dir = eq._psi_sign * eq._cocosdic['sigma_pol'] * eq._cocosdic['sigma_cyl']* eq._cocosdic['sigma_pol'] * np.sign(eq.I_plasma) * np.sign(eq.F0)
print('sigma_cyl: {}\nsigma_pol: {}\nIp: {}\nF0: {}'.format(
eq._cocosdic['sigma_cyl'], eq._cocosdic['sigma_pol'], eq.I_plasma,
eq.F0))
print('sign psi: {}'.format(eq._psi_sign))
print('dphi = {}\ndtheta = {}'.format(dphi, dtheta))
assert my_dir == dphidtheta
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()
def test_coordinates_transforms(geqdsk_file, cocos):
"""
Octants: https://en.wikipedia.org/wiki/Octant_(solid_geometry)#/media/File:Octant_numbers.svg
:param geqdsk_file:
:return:
"""
eq = read_geqdsk(geqdsk_file, cocos=cocos)
mgax = eq.magnetic_axis
# 3d octants:
o1 = eq.coordinates(X=1, Y=1, Z=1)
o2 = eq.coordinates(X=-1, Y=1, Z=1)
o3 = eq.coordinates(X=-1, Y=-1, Z=1)
o4 = eq.coordinates(X=1, Y=-1, Z=1)
o5 = eq.coordinates(X=1, Y=1, Z=-1)
o6 = eq.coordinates(X=-1, Y=1, Z=-1)
o7 = eq.coordinates(X=-1, Y=-1, Z=-1)
o8 = eq.coordinates(X=1, Y=-1, Z=-1)
# 2d quadrants:
q1 = eq.coordinates(R=mgax.R + 0.5, Z=mgax.Z + 0.5)
q2 = eq.coordinates(R=mgax.R - 0.5, Z=mgax.Z + 0.5)
q3 = eq.coordinates(R=mgax.R - 0.5, Z=mgax.Z - 0.5)
q4 = eq.coordinates(R=mgax.R + 0.5, Z=mgax.Z - 0.5)
# toroidal angle:
if cocos in [1, 3, 5, 7]:
# Cnt-clockwise
assert 0 < np.mod(o1.phi, 2 * np.pi) < np.pi / 2
assert np.pi / 2 < np.mod(o2.phi, 2 * np.pi) < np.pi
assert np.pi < np.mod(o3.phi, 2 * np.pi) < 3 * np.pi / 2
assert 3 * np.pi / 2 < np.mod(o4.phi, 2 * np.pi) < 2 * np.pi
assert 0 < np.mod(o5.phi, 2 * np.pi) < np.pi / 2
assert np.pi / 2 < np.mod(o6.phi, 2 * np.pi) < np.pi
assert np.pi < np.mod(o7.phi, 2 * np.pi) < 3 * np.pi / 2
assert 3 * np.pi / 2 < np.mod(o8.phi, 2 * np.pi) < 2 * np.pi
else:
# Clockwise
assert 0 < np.mod(o4.phi, 2 * np.pi) < np.pi / 2
assert np.pi / 2 < np.mod(o3.phi, 2 * np.pi) < np.pi
assert np.pi < np.mod(o2.phi, 2 * np.pi) < 3 * np.pi / 2
assert 3 * np.pi / 2 < np.mod(o1.phi, 2 * np.pi) < 2 * np.pi
assert 0 < np.mod(o8.phi, 2 * np.pi) < np.pi / 2
assert np.pi / 2 < np.mod(o7.phi, 2 * np.pi) < np.pi
assert np.pi < np.mod(o6.phi, 2 * np.pi) < 3 * np.pi / 2
assert 3 * np.pi / 2 < np.mod(o5.phi, 2 * np.pi) < 2 * np.pi
# Poloidal angle:
if cocos in [2, 3, 5, 8]:
# Cnt-clockwise
assert 0 * np.pi < np.mod(q1.theta[0], 2 * np.pi) < np.pi / 2
assert np.pi / 2 < np.mod(q2.theta[0], 2 * np.pi) < np.pi
assert np.pi < np.mod(q3.theta[0], 2 * np.pi) < 3 * np.pi / 2
assert 3 * np.pi / 2 < np.mod(q4.theta[0], 2 * np.pi) < 2 * np.pi
else:
assert 0 * np.pi < np.mod(q4.theta[0], 2 * np.pi) < np.pi / 2
assert np.pi / 2 < np.mod(q3.theta[0], 2 * np.pi) < np.pi
assert np.pi < np.mod(q2.theta[0], 2 * np.pi) < 3 * np.pi / 2
assert 3 * np.pi / 2 < np.mod(q1.theta[0], 2 * np.pi) < 2 * np.pi
# Toroidal quadrants::
o1 = eq.coordinates(R=mgax.R, Z=mgax.Z, phi=np.pi / 4)
o2 = eq.coordinates(R=mgax.R, Z=mgax.Z, phi=3 * np.pi / 4)
o3 = eq.coordinates(R=mgax.R, Z=mgax.Z, phi=5 * np.pi / 4)
o4 = eq.coordinates(R=mgax.R, Z=mgax.Z, phi=7 * np.pi / 4)
# Poloidal quadrants:
q1 = eq.coordinates(r=0.1, theta=1 * np.pi / 4)
q2 = eq.coordinates(r=0.1, theta=3 * np.pi / 4)
q3 = eq.coordinates(r=0.1, theta=5 * np.pi / 4)
q4 = eq.coordinates(r=0.1, theta=7 * np.pi / 4)
# toroidal angle:
if cocos in [1, 3, 5, 7]:
# Cnt-clockwise
assert 0 < o1.X
assert 0 < o1.Y
assert 0 > o2.X
assert 0 < o2.Y
assert 0 > o3.X
assert 0 > o3.Y
assert 0 < o4.X
assert 0 > o4.Y
else:
# Clockwise
assert 0 < o1.X
assert 0 > o1.Y
assert 0 > o2.X
assert 0 > o2.Y
assert 0 > o3.X
assert 0 < o3.Y
assert 0 < o4.X
assert 0 < o4.Y
# Poloidal angle:
if cocos in [2, 3, 5, 8]:
# Cnt-clockwise
assert mgax.R < q1.R
assert mgax.Z < q1.Z
assert mgax.R > q2.R
assert mgax.Z < q2.Z
assert mgax.R > q3.R
assert mgax.Z > q3.Z
assert mgax.R < q4.R
assert mgax.Z > q4.Z
else:
# Clockwise
assert mgax.R < q1.R
assert mgax.Z > q1.Z
assert mgax.R > q2.R
assert mgax.Z > q2.Z
assert mgax.R > q3.R
assert mgax.Z < q3.Z
assert mgax.R < q4.R
assert mgax.Z < q4.Z
def load_testing_equilibrium(case=0, cocos=None):
"""
Return testing equilibrium file
:param case -
:return:
"""
# cases are numbered from one... for now :-)
res_file = get_test_equilibria_filenames()[case]
res_limiterfile = get_test_divertor()[0]
# equil = read_fiesta_equilibrium(res_file)
if 'eqdsk' in res_file or 'gfile' in res_file or '/g' in res_file or '\g' in res_file:
# load as NetCDF
if cocos is None:
equil = readers.read_geqdsk(res_file)
else:
equil = readers.read_geqdsk(res_file, cocos=cocos)
elif '.nc' in res_file:
# load as gfile
with xr.open_dataset(res_file) as ds:
basedata = ds.load()
fw = np.array([basedata.first_wall_R, basedata.first_wall_Z]).T
if cocos is None:
equil = pleque.Equilibrium(basedata=basedata, first_wall=fw)
else:
equil = pleque.Equilibrium(basedata=basedata,
first_wall=fw,
cocos=cocos)
else:
# note recognized: