"""Contains functionality to read equilibrium files into RawNumericalEquilibrium."""
import numpy as np
from pathlib import Path
from typing import Union
from .eqdsk_io_m import read_eqdsk_file
from .xml_io_m import read_xml_file
from torx.equilibrium import RawNumericalEquilibrium
from torx.specializations.aug import AUG_KEY_MAP, load_aug_shotfile
from torx.decorators import autodoc_function
from torx.geometry import Polygon2D
_EQDSK_KEY_MAP = {
"axis_r": "rmaxis",
"axis_z": "zmaxis",
"axis_Btor": "bcentr",
"psi_axis": "simag",
"psi_separatrix": "sibry",
"spline_basis_r": "r",
"spline_basis_z": "z",
"psi_data": "psirz",
"x_point_r": None,
"x_point_z": None,
}
_XML_KEY_MAP = {
"axis_r": ("op", 1),
"axis_z": ("op", 2),
"axis_Btor": "B0",
"psi_axis": None,
"psi_separatrix": None,
"spline_basis_r": "r",
"spline_basis_z": "z",
"psi_data": "psi",
"x_point_r": ("xp", 1),
"x_point_z": ("xp", 2),
}
_SUPPORTED_FILETYPES = ["eqdsk", "xml", "pickle", "pkl"]
def _apply_key_map(raw: dict, key_map: dict) -> dict:
"""
Apply a key map to a raw dictionary.
Parameters
----------
raw : dict
Raw dictionary from a file reader.
key_map : dict
Mapping from standardized keys to raw keys. Values can be:
- str: direct key lookup
- tuple: (key, index) for array indexing
- None: key is not available in this format
Returns
-------
dict
Dictionary with standardized keys.
"""
result = {}
for std_key, raw_key in key_map.items():
if raw_key is None:
continue
if isinstance(raw_key, tuple):
key, idx = raw_key
result[std_key] = raw[key][idx]
else:
result[std_key] = raw[raw_key]
result["raw"] = raw
return result
def _psi_at_location(
psi: np.ndarray,
r: np.ndarray,
z: np.ndarray,
r_loc: float,
z_loc: float,
) -> float:
"""
Look up psi at the nearest grid point to a given location.
Parameters
----------
psi : np.ndarray
2D array of poloidal flux values with shape (nz, nr).
r : np.ndarray
R grid vector.
z : np.ndarray
Z grid vector.
r_loc : float
R location to look up.
z_loc : float
Z location to look up.
Returns
-------
float
Poloidal flux at the nearest grid point.
"""
r_idx = np.argmin(np.abs(r - r_loc))
z_idx = np.argmin(np.abs(z - z_loc))
return float(psi[z_idx, r_idx])
[docs]
@autodoc_function
def equilibrium_from_file(
filepath: Union[str, Path],
filetype: str = "auto",
) -> RawNumericalEquilibrium:
"""
Read an equilibrium file and return a RawNumericalEquilibrium.
The returned equilibrium is not normalized. Call normalize(B0) before
use as a NumericalEquilibrium.
Currently supported formats are:
- eqdsk (G-EQDSK format)
- xml (as generated by pyequil)
- pickle/pkl (AUG shotfile saved via save_aug_shotfile)
Parameters
----------
filepath : Union[str, Path]
Path to the equilibrium file.
filetype : str, optional
File format, by default "auto" which detects from extension.
Returns
-------
RawNumericalEquilibrium
An unnormalized equilibrium ready for preprocessing.
"""
filepath = Path(filepath)
assert filepath.exists(), f"{filepath.absolute()} does not exist!"
if filetype == "auto":
filetype = filepath.suffix[1:]
assert filetype in _SUPPORTED_FILETYPES, (
f"Unsupported filetype '{filetype}'. "
f"Currently supported are: {_SUPPORTED_FILETYPES}"
)
if filetype == "eqdsk":
raw = read_eqdsk_file(filepath)
d = _apply_key_map(raw, _EQDSK_KEY_MAP)
d["divertor_polygon"] = Polygon2D(raw["rlim"], raw["zlim"])
elif filetype == "xml":
raw = read_xml_file(filepath)
d = _apply_key_map(raw, _XML_KEY_MAP)
d["psi_axis"] = _psi_at_location(
raw["psi"], raw["r"], raw["z"],
d["axis_r"], d["axis_z"],
)
d["psi_separatrix"] = _psi_at_location(
raw["psi"], raw["r"], raw["z"],
d["x_point_r"], d["x_point_z"],
)
elif filetype in ["pickle", "pkl"]:
raw = load_aug_shotfile(filepath)
d = _apply_key_map(raw, AUG_KEY_MAP)
return RawNumericalEquilibrium.from_dict(d)
