"""Contains functionality to write eqdsk equilibrium files."""
import numpy as np
from pathlib import Path
from warnings import warn
from typing import Union
from fortranformat import FortranRecordWriter
from torx.autodoc_decorators_m import autodoc_function
_required_keys = ["nw", "nh", "rdim", "zdim", "rcentr", "rleft", "zmid",
"rmaxis", "zmaxis", "simag", "sibry", "bcentr", "current",
"fpol", "pres", "ffprim", "pprime", "psirz", "qpsi",
"nbbbs", "limitr", "rbbbs", "zbbbs", "rlim", "zlim"]
[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:
# Header
line = f2000.write(case + [idum, content["nw"], content["nh"]])
filew.write(line + "\n")
# Block with scalars
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")
# Block with arrays
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")
# NOTE: Psi needs to be in row major ordering
psi = content["psirz"].flatten(order="C")
line = f2020.write(psi)
filew.write(line + "\n")
line = f2020.write(content["qpsi"])
filew.write(line + "\n")
# Block with boundary and limiter
line = f2022.write([content["nbbbs"], content["limitr"]])
filew.write(line + "\n")
boundary = np.stack([content["rbbbs"], content["zbbbs"]]\
).flatten(order="C")
line = f2020.write(boundary)
filew.write(line + "\n")
limiter = np.stack([content["rlim"], content["zlim"]]\
).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.
"""
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."""
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.
"""
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.
"""
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}!"
