"""Contains functionality to read and write eqdsk equilibrium files."""
import numpy as np
from pathlib import Path
from warnings import warn
from typing import Union
from fortranformat import FortranRecordReader
from fortranformat import FortranRecordWriter
from torx.decorators import autodoc_function
_eqdsk_key_descriptions = {
"case": "Identification character string (6 elements of 8 chars)",
"nw": "Number of horizontal R grid points",
"nh": "Number of vertical Z grid points",
"rdim": "Horizontal dimension in meter of computational box",
"zdim": "Vertical dimension in meter of computational box",
"rcentr": "R in meter of vacuum toroidal magnetic field BCENTR",
"rleft": "Minimum R in meter of rectangular computational box",
"zmid": "Z of center of computational box in meter",
"rmaxis": "R of magnetic axis in meter",
"zmaxis": "Z of magnetic axis in meter",
"simag": "Poloidal flux at magnetic axis in Weber/rad",
"sibry": "Poloidal flux at plasma boundary in Weber/rad",
"bcentr": "Vacuum toroidal magnetic field in Tesla at RCENTR",
"current": "Plasma current in Ampere",
"fpol": "Poloidal current function in m-T, F = RBT on flux grid",
"pres": "Plasma pressure in N/m^2 on uniform flux grid",
"ffprim": "FF'(psi) in (mT)^2 / (Weber/rad) on uniform flux grid",
"pprime": "P'(psi) in (N/m^2) / (Weber/rad) on uniform flux grid",
"psirz": "Poloidal flux in Weber/rad on rectangular grid points",
"qpsi": "Safety factor q on uniform flux grid from axis to boundary",
"nbbbs": "Number of boundary points",
"limitr": "Number of limiter points",
"rbbbs": "R of boundary points in meter",
"zbbbs": "Z of boundary points in meter",
"rlim": "R of surrounding limiter contour in meter",
"zlim": "Z of surrounding limiter contour in meter",
"r": "R grid vector in meter, derived from rleft, rdim and nw",
"z": "Z grid vector in meter, derived from zmid, zdim and nh",
}
_required_keys = [k for k in _eqdsk_key_descriptions if k not in ("case", "r", "z")]
[docs]
@autodoc_function
def eqdsk_info(key: str = None):
"""
Return description of one or all canonical EQDSK keys.
Parameters
----------
key : str, optional
Canonical EQDSK key name. If None, prints a table of all keys
and their descriptions.
Returns
-------
str
Description of the requested key, or None if no key is provided.
"""
if key is None:
col_width = max(len(k) for k in _eqdsk_key_descriptions)
for k, v in _eqdsk_key_descriptions.items():
print(f" {k:<{col_width}} {v}")
return None
assert key in _eqdsk_key_descriptions, f"Unknown key: {key}"
return _eqdsk_key_descriptions[key]
[docs]
@autodoc_function
def read_eqdsk_file(filepath: Path) -> dict:
"""
Read an eqdsk file and return a dictionary with canonical EQDSK keys.
Parameters
----------
filepath : Path
Path to the eqdsk file.
Returns
-------
dict
Dictionary with canonical EQDSK keys. See G-EQDSK format documentation
for variable descriptions. Also includes derived keys ``r`` and ``z``
for the grid vectors.
"""
lines = filepath.read_text().split("\n")
header = FortranRecordReader("(6a8,3i4)").read(lines[0])
nw = header[-2]
nh = header[-1]
case = header[:6]
rdim, zdim, rcentr, rleft, zmid = FortranRecordReader("(5e16.9)").read(
lines[1]
)
rmaxis, zmaxis, simag, sibry, bcentr = FortranRecordReader("(5e16.9)").read(
lines[2]
)
current, simag_dup, _, rmaxis_dup, _ = FortranRecordReader("(5e16.9)").read(
lines[3]
)
zmaxis_dup, _, sibry_dup, _, _ = FortranRecordReader("(5e16.9)").read(
lines[4]
)
assert np.isclose(simag, simag_dup)
assert np.isclose(rmaxis, rmaxis_dup)
assert np.isclose(zmaxis, zmaxis_dup)
assert np.isclose(sibry, sibry_dup)
value_reader = _read_single_values(lines, start_at_line=5)
fpol, _, _ = _read_1d_array(value_reader, number_of_values=nw)
pres, _, _ = _read_1d_array(value_reader, number_of_values=nw)
ffprim, _, _ = _read_1d_array(value_reader, number_of_values=nw)
pprime, _, _ = _read_1d_array(value_reader, number_of_values=nw)
psirz, _, _ = _read_2d_array(value_reader, rows=nh, columns=nw)
qpsi, n_line, _ = _read_1d_array(value_reader, number_of_values=nw)
nbbbs, limitr = FortranRecordReader("(2i5)").read(lines[n_line + 1])
value_reader = _read_single_values(lines, start_at_line=n_line + 2)
bbbs, _, _ = _read_1d_array(value_reader, number_of_values=2 * nbbbs)
rbbbs, zbbbs = bbbs[0::2], bbbs[1::2]
lim, _, _ = _read_1d_array(value_reader, number_of_values=2 * limitr)
rlim, zlim = lim[0::2], lim[1::2]
r = np.linspace(rleft, rleft + rdim, nw)
z = np.linspace(zmid - 0.5 * zdim, zmid + 0.5 * zdim, nh)
return dict(
case=case,
nw=nw,
nh=nh,
rdim=rdim,
zdim=zdim,
rcentr=rcentr,
rleft=rleft,
zmid=zmid,
rmaxis=rmaxis,
zmaxis=zmaxis,
simag=simag,
sibry=sibry,
bcentr=bcentr,
current=current,
fpol=fpol,
pres=pres,
ffprim=ffprim,
pprime=pprime,
psirz=psirz,
qpsi=qpsi,
nbbbs=nbbbs,
limitr=limitr,
rbbbs=rbbbs,
zbbbs=zbbbs,
rlim=rlim,
zlim=zlim,
r=r,
z=z,
)
def _read_single_values(lines: list, format_spec: str = "5e16.9",
start_at_line: int = 5):
"""
Generate individual float values from a list of Fortran-formatted lines.
Parameters
----------
lines : list
List of strings representing the lines of the eqdsk file.
format_spec : str, optional
Fortran format specifier, by default "5e16.9".
start_at_line : int, optional
Line number to start reading from, by default 5.
"""
for line_number, line in enumerate(lines):
if line_number < start_at_line:
continue
values = FortranRecordReader(f"({format_spec})").read(line)
for value_position, value in enumerate(values):
if value is not None:
yield value, line_number, value_position
def _read_1d_array(generator, number_of_values: int):
"""
Read a 1D array of given length from a value generator.
Parameters
----------
generator : generator
Generator yielding (value, line_number, value_position) tuples.
number_of_values : int
Number of values to read.
Returns
-------
tuple
(array, line_number, value_position)
"""
array = np.zeros(number_of_values)
line_number, value_position = 0, 0
for i in range(number_of_values):
value, line_number, value_position = next(generator)
array[i] = value
return array, line_number, value_position
def _read_2d_array(generator, rows: int, columns: int):
"""
Read a 2D array of given shape from a value generator.
Parameters
----------
generator : generator
Generator yielding (value, line_number, value_position) tuples.
rows : int
Number of rows.
columns : int
Number of columns.
Returns
-------
tuple
(array, line_number, value_position) where array has shape
(rows, columns).
"""
array, line_number, value_position = _read_1d_array(
generator, rows * columns
)
return array.reshape((rows, columns)), line_number, value_position
[docs]
@autodoc_function
def write_eqdsk_file(filepath: Path, content: dict):
"""
Create an eqdsk equilibrium file with path and name given by filepath.
The content will be extracted from the dictionary provided. The presence of
all required quantities is assumed (see eqdsk documentation), except the
"idum" and "xdum" variables (dummies) and "case". The "case" variable is
built automatically from keys "name", "shot", "time" and "comment". These
are assumed to consist of 8 chars each (24 for comment) and are truncated
if longer.
Example for "case":
name = "AUG"
shot = "30000"
time = "1.0"
comment = "generated for testing"
"""
_check_content(content)
_check_dimensions(content)
case = _create_case(content)
idum = 3
xdum = 0.0
f2000 = FortranRecordWriter("(6a8,3i4)")
f2020 = FortranRecordWriter("(5e16.9)")
f2022 = FortranRecordWriter("(2i5)")
with open(filepath, "w") as filew:
line = f2000.write(case + [idum, content["nw"], content["nh"]])
filew.write(line + "\n")
line = f2020.write([content["rdim"], content["zdim"],
content["rcentr"], content["rleft"],
content["zmid"]])
filew.write(line + "\n")
line = f2020.write([content["rmaxis"], content["zmaxis"],
content["simag"], content["sibry"],
content["bcentr"]])
filew.write(line + "\n")
line = f2020.write([content["current"], content["simag"], xdum,
content["rmaxis"], xdum])
filew.write(line + "\n")
line = f2020.write([content["zmaxis"], xdum, content["sibry"],
xdum, xdum])
filew.write(line + "\n")
line = f2020.write(content["fpol"])
filew.write(line + "\n")
line = f2020.write(content["pres"])
filew.write(line + "\n")
line = f2020.write(content["ffprim"])
filew.write(line + "\n")
line = f2020.write(content["pprime"])
filew.write(line + "\n")
psi = content["psirz"].flatten(order="C")
line = f2020.write(psi)
filew.write(line + "\n")
line = f2020.write(content["qpsi"])
filew.write(line + "\n")
line = f2022.write([content["nbbbs"], content["limitr"]])
filew.write(line + "\n")
boundary = np.stack([content["rbbbs"], content["zbbbs"]], axis=1
).flatten(order="C")
line = f2020.write(boundary)
filew.write(line + "\n")
limiter = np.stack([content["rlim"], content["zlim"]], axis=1
).flatten(order="C")
line = f2020.write(limiter)
filew.write(line + "\n")
def _create_case(content: dict):
"""
Create the variable case which consists of a 6 element string array.
The elements may contain more than the allowed 8 characters, since it is
assumed that the strings will be truncated when writing. A warning will
be given for each element that is longer than 8 characters.
The content of case will be determined based on the fields "name", "shot",
"time" and "comment" in the content dict. From comment, 23 characters will
be taken, a space will be prepended.
Parameters
----------
content : dict
Dictionary containing equilibrium data.
Returns
-------
list
List of 6 strings representing the case variable.
"""
case = [""] * 6
keys = content.keys()
if "name" in keys:
case[0] = content["name"]
if "shot" in keys:
case[1] = content["shot"]
if "time" in keys:
case[2] = content["time"]
if "comment" in keys:
case[3] = "" + content["comment"][0:7]
case[4] = content["comment"][7:15]
case[5] = content["comment"][15:]
for i, c in enumerate(case):
if len(c) > 8:
warn(f"Case entry {i} which is {c} contains more "
+ "characters than 8 and will be truncated!")
case[i] = c[0:8]
return case
def _check_content(content: dict):
"""
Check the content dictionary if all required keys are present.
Parameters
----------
content : dict
Dictionary containing equilibrium data.
"""
for key in _required_keys:
assert key in content.keys(), f"Required key {key} not found!"
def _check_dimensions(content: dict):
"""
Check for all required keys in the content dictionary if they are valid.
Validity is described in a separate check array function.
Parameters
----------
content : dict
Dictionary containing equilibrium data.
"""
nw_arrays = ["fpol", "pres", "ffprim", "pprime", "qpsi"]
for na in nw_arrays:
_check_array(content, na, "nw")
_check_array(content, "psirz", ("nw", "nh"))
_check_array(content, "rbbbs", "nbbbs")
_check_array(content, "zbbbs", "nbbbs")
_check_array(content, "rlim", "limitr")
_check_array(content, "zlim", "limitr")
skip = nw_arrays + ["psirz", "rbbbs", "zbbbs", "rlim", "zlim",
"name", "shot", "time", "comment"]
for key in content.keys():
if key in skip:
continue
_check_array(content, key, "1")
def _check_array(content: dict, array_name: str, dim_name: Union[str, tuple]):
"""
Check a single array contained in the content dictionary for validity.
This includes size and dimensionality checks against the dimension
specified in dim_name. For scalars dim_name should be "1". For 2D arrays,
a tuple of 2 strings is expected.
Parameters
----------
content : dict
Dictionary containing equilibrium data.
array_name : str
Name of the array to check.
dim_name : Union[str, tuple]
Name of the dimension to check against, or tuple of two names for 2D.
"""
array = content[array_name]
assert (isinstance(array, np.ndarray) or isinstance(array, (int, float))), \
f"Array {array_name} must be of float or array type, was {type(array)}!"
if isinstance(dim_name, tuple):
assert len(dim_name) == 2, \
f"For dim tuple the expected length is 2, was {len(dim_name)}!)"
dim = content[dim_name[0]] * content[dim_name[1]]
ndim = 2
shape = (content[dim_name[0]], content[dim_name[1]])
assert array.shape == shape, \
f"Shape of {array_name} does not match {shape}!"
elif dim_name == "1":
dim = 1
ndim = 0
else:
dim = content[dim_name]
ndim = 1
if not isinstance(array, np.ndarray):
array = np.array(array)
assert len(array.shape) >= 0, f"Quantity {array_name} must be >= 0D!"
assert len(array.shape) <= 2, f"Quantity {array_name} must be <= 2D!"
assert len(array.shape) == ndim, f"Quantity {array_name} must be {ndim}D!"
assert array.size == dim, f"Size of {array_name} does not match {dim_name}!"
