"""NetCDF reader and writer for NumericalEquilibrium files."""
import time
import warnings
import numpy as np
from pathlib import Path
from h5netcdf import File
from torx.geometry import Polygon2D
from torx.performance import get_git_hash, get_git_email
from torx.decorators import autodoc_function
[docs]
@autodoc_function
def read_netcdf_equilibrium(equi_file: Path) -> dict:
"""
Read a PARALLAX equilibrium NetCDF file into a standardized dictionary.
Parameters
----------
equi_file : Path
Path to the equilibrium NetCDF file.
Returns
-------
dict
Standardized dictionary with keys:
- axis_r: float, magnetic axis R in normalized units
- axis_z: float, magnetic axis Z in normalized units
- axis_Btor: float, toroidal field on axis in Tesla
- psi_axis: float, poloidal flux on axis in Weber
- psi_separatrix: float, poloidal flux on separatrix in Weber
- rho_min: float, minimum flux surface label
- rho_max: float, maximum flux surface label
- spline_basis_r: np.ndarray, R grid in normalized units
- spline_basis_z: np.ndarray, Z grid in normalized units
- psi_data: np.ndarray, poloidal flux on grid in Weber
- x_point_r: float or np.ndarray or None
- x_point_z: float or np.ndarray or None
- poloidal_field_factor: float
- divertor_polygon: Polygon2D
- exclusion_polygon: Polygon2D or None
- flux_limit_polygons: dict of Polygon2D
- flux_limit_rho_min: dict of float or None
- flux_limit_rho_max: dict of float or None
- axis_r_units: str
- x_point_units: str
- axis_Btor_units: str
- psi_units: str
"""
assert equi_file.exists(), \
f"Equilibrium file {equi_file} does not exist"
assert equi_file.suffix == ".nc", \
f"Expected .nc file, got {equi_file.suffix}"
with File(str(equi_file), "r") as f:
mg = f["Magnetic_geometry"]
pl = f["Psi_limits"]
mg_attrs = dict(mg.attrs)
pl_attrs = dict(pl.attrs)
axis_r = float(mg_attrs["magnetic_axis_R"])
axis_z = float(mg_attrs["magnetic_axis_Z"])
axis_r_units = _normalize_unit_string(
_decode_bs(mg_attrs.get("magnetic_axis_units", "normalized"))
)
axis_Btor = float(mg_attrs["axis_Btor"])
axis_Btor_units = _normalize_unit_string(
_decode_bs(mg_attrs.get("btor_units") or mg_attrs.get("axis_Btor_units", "T"))
)
poloidal_field_factor = float(mg_attrs.get("poloidal_field_factor", 1.0))
try:
x_point_r = mg_attrs["x_point_R"]
x_point_z = mg_attrs["x_point_Z"]
x_point_units = _normalize_unit_string(
_decode_bs(mg_attrs.get("x_point_units", "normalized"))
)
x_point_r = np.atleast_1d(np.array(x_point_r, dtype=float))
x_point_z = np.atleast_1d(np.array(x_point_z, dtype=float))
if x_point_r.size == 0:
x_point_r = None
x_point_z = None
elif x_point_r.size == 1:
x_point_r = float(x_point_r[0])
x_point_z = float(x_point_z[0])
except KeyError:
warnings.warn(
"No x-point data found in NetCDF. "
"Consider updating the equilibrium file."
)
x_point_r = None
x_point_z = None
x_point_units = "normalized"
spline_basis_r = np.array(mg.variables["R"], dtype=float)
spline_basis_z = np.array(mg.variables["Z"], dtype=float)
psi_data = np.array(mg.variables["psi"], dtype=float)
psi_units = _normalize_unit_string(
_decode_bs(
mg_attrs.get("psi_units")
or mg.variables["psi"].attrs.get("units", "Wb")
)
)
psi_axis = float(pl_attrs["psi_axis"])
rho_min = float(pl_attrs["rho_min"])
rho_max = float(pl_attrs["rho_max"])
if "psi_separatrix" in pl_attrs:
psi_separatrix = float(pl_attrs["psi_separatrix"])
elif "psi_seperatrix" in pl_attrs:
psi_separatrix = float(pl_attrs["psi_seperatrix"])
else:
raise KeyError(
"Neither psi_separatrix nor psi_seperatrix found in Psi_limits"
)
flux_limit_polygons = {}
flux_limit_rho_min = {}
flux_limit_rho_max = {}
for name in pl.groups:
grp = pl[name]
poly = _read_polygon_group(grp)
flux_limit_polygons[name] = poly
use_min = int(grp.attrs.get("use_local_min", 0))
use_max = int(grp.attrs.get("use_local_max", 0))
flux_limit_rho_min[name] = (
float(grp.attrs["local_rho_min"]) if use_min else None
)
flux_limit_rho_max[name] = (
float(grp.attrs["local_rho_max"]) if use_max else None
)
divertor_polygon = _read_polygon_group(f["divertor_polygon"])
exclusion_polygon = (
_read_polygon_group(f["exclusion_polygon"])
if "exclusion_polygon" in f.groups
else None
)
return {
"axis_r": axis_r,
"axis_z": axis_z,
"axis_Btor": axis_Btor,
"psi_axis": psi_axis,
"psi_separatrix": psi_separatrix,
"rho_min": rho_min,
"rho_max": rho_max,
"spline_basis_r": spline_basis_r,
"spline_basis_z": spline_basis_z,
"psi_data": psi_data,
"x_point_r": x_point_r,
"x_point_z": x_point_z,
"poloidal_field_factor": poloidal_field_factor,
"divertor_polygon": divertor_polygon,
"exclusion_polygon": exclusion_polygon,
"flux_limit_polygons": flux_limit_polygons,
"flux_limit_rho_min": flux_limit_rho_min,
"flux_limit_rho_max": flux_limit_rho_max,
"axis_r_units": axis_r_units,
"x_point_units": x_point_units,
"axis_Btor_units": axis_Btor_units,
"psi_units": psi_units,
}
[docs]
@autodoc_function
def write_netcdf_equilibrium(
file_path: Path,
equi,
description: str,
comment: str = "",
allow_overwrite: bool = False,
ordering: str = "Fortran",
):
"""
Write a NumericalEquilibrium or RawNumericalEquilibrium to a NetCDF file.
The file is written in NETCDF4_CLASSIC format for Fortran compatibility.
All integer attributes are written as np.int32 and all float attributes
as plain Python floats to satisfy the Fortran NetCDF library requirements.
All spatial quantities are written in normalized units.
Parameters
----------
file_path : Path
Output path, must have .nc suffix.
equi : NumericalEquilibrium or RawNumericalEquilibrium
The equilibrium to write. Must be normalized.
description : str
Human-readable description written to the file header.
comment : str, optional
Additional comment written to the file header.
allow_overwrite : bool, optional
If False, raises FileExistsError if file_path exists.
ordering : str, optional
Array ordering, "Fortran" or "C". Default is "Fortran".
"""
assert file_path.suffix == ".nc", \
f"Expected .nc suffix, got {file_path.suffix}"
assert ordering in ("Fortran", "C"), \
f"ordering must be Fortran or C"
if file_path.exists():
if not allow_overwrite:
raise FileExistsError(
f"{file_path} already exists and allow_overwrite is False"
)
file_path.unlink()
with File(str(file_path), "w", format="NETCDF4_CLASSIC") as f:
_write_header(f, description, comment)
_write_magnetic_geometry(f, equi, ordering)
_write_psi_limits(f, equi)
_write_polygon_group(f, "divertor_polygon", equi.divertor_polygon)
if equi.exclusion_polygon is not None:
_write_polygon_group(f, "exclusion_polygon", equi.exclusion_polygon)
assert file_path.exists(), \
f"Failed to create NetCDF file {file_path}"
def _normalize_unit_string(unit: str) -> str:
"""
Normalize legacy unit strings to canonical pint-compatible ones.
Maps the various spellings used in old files to the canonical form:
"normalised" / "normalized" -> "" (pint dimensionless)
anything else -> as-is
"""
if unit.lower() in ("normalised", "normalized"):
return "dimensionless"
return unit
def _decode_bs(val, default=""):
"""Decode bytestring from NetCDF to a string."""
if val is None:
return default
if isinstance(val, bytes):
return val.decode("utf-8")
return str(val)
def _read_polygon_group(grp) -> Polygon2D:
"""Read a polygon from an h5netcdf group."""
r_points = np.array(grp.variables["R_points"], dtype=float)
z_points = np.array(grp.variables["Z_points"], dtype=float)
invert = bool(int(grp.attrs.get("invert", 0)))
# writer appends closing point — strip it since Polygon2D handles closure
if np.isclose(r_points[0], r_points[-1]) and np.isclose(z_points[0], z_points[-1]):
r_points = r_points[:-1]
z_points = z_points[:-1]
return Polygon2D(r_points, z_points, invert_polygon=invert)
def _write_header(f: File, description: str, comment: str):
"""Write global attributes to the NetCDF file."""
f.attrs["description"] = description
f.attrs["build_url"] = "gitlab.mpcdf.mpg.de/phoenix/parallax_equilibrium"
f.attrs["build_hash"] = get_git_hash()
f.attrs["version"] = np.int32(2)
f.attrs["history"] = "Created " + time.ctime(time.time())
f.attrs["author"] = get_git_email()
if comment:
f.attrs["comment"] = comment
def _write_magnetic_geometry(f: File, equi, ordering: str):
"""
Write the Magnetic_geometry group.
All spatial quantities are written in normalized units.
The physical R0 in meters is stored in R0_m for reconstruction.
"""
group = f.create_group("Magnetic_geometry")
group.attrs["magnetic_axis_R"] = float(np.asarray(equi.axis_r).item())
group.attrs["magnetic_axis_Z"] = float(np.asarray(equi.axis_z).item())
assert equi.axis_r.norm.m == 1
group.attrs["magnetic_axis_units"] = str(equi.axis_r.norm.u)
x_point_r = equi.x_point_r
x_point_z = equi.x_point_z
if x_point_r is not None and x_point_z is not None:
group.attrs["x_point_R"] = np.atleast_1d(
np.asarray(x_point_r, dtype=float)
).tolist()
group.attrs["x_point_Z"] = np.atleast_1d(
np.asarray(x_point_z, dtype=float)
).tolist()
assert equi.x_point_r.norm.m == 1
assert equi.x_point_r.norm == equi.x_point_z.norm
group.attrs["x_point_units"] = str(equi.x_point_r.norm.u)
else:
group.attrs["x_point_R"] = []
group.attrs["x_point_Z"] = []
group.attrs["x_point_units"] = []
axis_Btor = float(np.asarray(equi.axis_Btor).item())
assert equi.axis_Btor.norm.m == 1
group.attrs["axis_Btor_units"] = str(equi.axis_Btor.norm.u)
if axis_Btor < 0:
group.attrs["axis_Btor"] = float(abs(axis_Btor))
group.attrs["toroidal_field_invert"] = np.int32(-1)
group.attrs["toroidal_field_direction"] = "favourable"
group.attrs["poloidal_field_factor"] = float(-equi.poloidal_field_factor)
else:
group.attrs["axis_Btor"] = float(axis_Btor)
group.attrs["toroidal_field_invert"] = np.int32(1)
group.attrs["toroidal_field_direction"] = "unfavourable"
group.attrs["poloidal_field_factor"] = float(equi.poloidal_field_factor)
r = np.asarray(equi.spline_basis_r, dtype=float)
z = np.asarray(equi.spline_basis_z, dtype=float)
psi = np.asarray(equi.psi_data, dtype=float)
group.dimensions = {"R": r.size, "Z": z.size}
dimensions_2d = ("Z", "R") if ordering == "Fortran" else ("R", "Z")
R_var = group.create_variable("R", dimensions=("R",), dtype=np.float64)
Z_var = group.create_variable("Z", dimensions=("Z",), dtype=np.float64)
psi_var = group.create_variable(
"psi", dimensions=dimensions_2d, dtype=np.float64
)
R_var[:] = r
Z_var[:] = z
psi_var[:, :] = psi
R_var.attrs["units"] = "normalized"
Z_var.attrs["units"] = "normalized"
psi_var.attrs["units"] = "Wb"
group.attrs["psi_units"] = "Wb"
def _write_psi_limits(f: File, equi):
"""Write the Psi_limits group including flux limit polygons."""
group = f.create_group("Psi_limits")
assert equi.rho_min is not None, "equi.rho_min must be set before writing"
assert equi.rho_max is not None, "equi.rho_max must be set before writing"
group.attrs["rho_min"] = float(equi.rho_min)
group.attrs["rho_max"] = float(equi.rho_max)
group.attrs["rho_units"] = "normalized"
group.attrs["psi_separatrix"] = float(
np.asarray(equi.psi_separatrix).item()
)
group.attrs["psi_axis"] = float(
np.asarray(equi.psi_axis).item()
)
group.attrs["psi_units"] = "Wb"
for name, polygon in equi.flux_limit_polygons.items():
rho_min = equi.flux_limit_rho_min.get(name)
rho_max = equi.flux_limit_rho_max.get(name)
_write_polygon_group(group, name, polygon, rho_min=rho_min, rho_max=rho_max)
def _write_polygon_group(
parent,
name: str,
polygon: Polygon2D,
rho_min: float = None,
rho_max: float = None,
):
"""Write a polygon into a subgroup of parent."""
poly_x, poly_y, n_pts = _unique_polygon_points(
polygon.x_points, polygon.y_points
)
poly_x = np.append(poly_x, poly_x[0])
poly_y = np.append(poly_y, poly_y[0])
n_pts += 1
check_poly = Polygon2D(poly_x, poly_y)
assert check_poly.signed_area() != 0, f"Polygon '{name}' has zero area"
if check_poly.signed_area() < 0:
poly_x = np.flip(poly_x)
poly_y = np.flip(poly_y)
group = parent.create_group(name)
group.dimensions = {"N_points": n_pts}
R_var = group.create_variable(
"R_points", dimensions=("N_points",), dtype=np.float64
)
Z_var = group.create_variable(
"Z_points", dimensions=("N_points",), dtype=np.float64
)
R_var[:] = poly_x
Z_var[:] = poly_y
group.attrs["invert"] = np.int32(int(polygon.invert_polygon))
group.attrs["extent_units"] = "normalized"
if rho_min is not None or rho_max is not None:
_write_flux_limit_attrs(group, rho_min, rho_max)
def _write_flux_limit_attrs(group, rho_min: float, rho_max: float):
"""Write flux limit rho attributes into a polygon group."""
assert not (rho_min is None and rho_max is None), \
"Must supply at least one flux limit"
assert rho_min is None or rho_max is None, \
"Cannot supply both rho_min and rho_max"
if rho_min is not None:
limit, blank = "min", "max"
limit_val = float(rho_min)
else:
limit, blank = "max", "min"
limit_val = float(rho_max)
group.attrs[f"use_local_{limit}"] = np.int32(1)
group.attrs[f"local_rho_{limit}"] = limit_val
group.attrs[f"use_local_{blank}"] = np.int32(0)
group.attrs[f"local_rho_{blank}"] = np.int32(0)
group.attrs["rho_units"] = "normalized"
def _unique_polygon_points(
polygon_x: np.ndarray, polygon_y: np.ndarray
) -> tuple:
"""Return unique polygon vertices preserving original order."""
polygon_x = np.array(polygon_x)
polygon_y = np.array(polygon_y)
pairs = np.column_stack((polygon_x, polygon_y))
_, index_sort = np.unique(pairs, axis=0, return_index=True)
unique_x = pairs[:, 0][np.sort(index_sort)]
unique_y = pairs[:, 1][np.sort(index_sort)]
n_unique = unique_x.size
assert unique_y.size == n_unique
return unique_x, unique_y, n_unique
