"""Contains functionality to read and write AUG shotfile equilibrium data."""
import numpy as np
import pickle
from pathlib import Path
from typing import Union
from torx.decorators import autodoc_function
from ._aug_module_import_m import import_aug
AUG_KEY_MAP = {
"axis_r": "R0",
"axis_z": "Z0",
"axis_Btor": "Btor_sign",
"psi_axis": "psiO",
"psi_separatrix": "psiX",
"spline_basis_r": "R",
"spline_basis_z": "Z",
"psi_data": "psi",
"x_point_r": "R_xpt_lower",
"x_point_z": "Z_xpt_lower",
}
[docs]
@autodoc_function
def read_aug_shotfile(
shot_number: int,
time_point: float,
diag: str = "EQI",
) -> dict:
"""
Read equilibrium data from the AUG database.
Requires aug_sfutils to be installed and access to the IPP network.
Returns a dictionary that can be passed to
RawNumericalEquilibrium.from_dict() after key mapping.
Parameters
----------
shot_number : int
AUG shot number.
time_point : float
Time point in seconds.
diag : str, optional
Diagnostic to use, by default "EQI".
Returns
-------
dict
Dictionary containing the equilibrium data.
"""
sf = import_aug()
print(f"Reading equilibrium from AUG-database via {diag}-diagnostic:")
EQU = sf.EQU(nshot=shot_number, diag=diag)
EQU.read_pfm()
EQU.read_profiles()
EQU.read_scalars()
EQU.B_mesh()
t = EQU.time
t_idx = np.argmin(np.abs(t - time_point))
data = dict(
shot_number=shot_number,
diag=diag,
time=EQU.time[t_idx],
time_stamp=t_idx,
R=EQU.Rmesh,
Z=EQU.Zmesh,
R0=EQU.Rmag[t_idx],
Z0=EQU.Zmag[t_idx],
B_r=EQU.Br[t_idx, :, :].T,
B_z=EQU.Bz[t_idx, :, :].T,
B_t=EQU.Bt[t_idx, :, :].T,
Btor_sign=EQU.psi_sign,
R_xpt_lower=EQU.Rxpu[t_idx],
Z_xpt_lower=EQU.Zxpu[t_idx],
R_xpt_upper=EQU.Rxpo[t_idx],
Z_xpt_upper=EQU.Zxpo[t_idx],
psi=EQU.pfm[:, :, t_idx].T,
psi_chord=EQU.pfl[t_idx, :],
psiX=EQU.psix[t_idx],
psiO=EQU.psi0[t_idx],
psi_separatrix=EQU.psix[t_idx],
psi_axis=EQU.psi0[t_idx],
rho_chord=np.sqrt(
(EQU.pfl[t_idx, :] - EQU.psi0[t_idx])
/ (EQU.psix[t_idx] - EQU.psi0[t_idx])
),
q_chord=EQU.q[t_idx, :],
)
print(
f"AUG equilibrium from {diag} at t={data['time']:.3f} "
f"corresponding to time-stamp {t_idx}."
)
return data
[docs]
@autodoc_function
def save_aug_shotfile(
data: dict,
filepath: Union[str, Path],
):
"""
Save an AUG shotfile dictionary as a pickle file.
Useful for storing AUG data for use outside the IPP network.
Parameters
----------
data : dict
Dictionary as returned by read_aug_shotfile.
filepath : Union[str, Path]
Path to save the pickle file.
"""
filepath = Path(filepath)
with open(filepath, "wb") as f:
pickle.dump(data, f)
print(f"Saved AUG shotfile to {filepath}")
[docs]
@autodoc_function
def load_aug_shotfile(filepath: Union[str, Path]) -> dict:
"""
Load an AUG shotfile dictionary from a pickle file.
Parameters
----------
filepath : Union[str, Path]
Path to the pickle file.
Returns
-------
dict
Dictionary containing the equilibrium data.
"""
filepath = Path(filepath)
assert filepath.exists(), f"{filepath.absolute()} does not exist!"
with open(filepath, "rb") as f:
return pickle.load(f)
