"""
Read and write flux-surface files in the Kisslinger format.
Turns a PoincarePolygonSet (the ordered flux-surface polygons plus its
mean-area and rho labels) into the ``{prefix}_isurf_NNN.txt`` surface
files and the ``{prefix}_flux_labels.txt`` label file. This is the
SOL-capable output of the field-line-traced flux surfing workflow: it
uses the per-surface area label already stored on the polygon set (see
PoincarePolygonSet.mean_area) rather than interpolating rho from the
raw Poincare data. The companion ``read_kisslinger_flux_surfaces``
parses those files back into coordinate arrays and labels.
"""
from pathlib import Path
import numpy as np
from torx.geometry import PoincarePolygonSet
from torx.decorators import autodoc_function
from torx.units import Quantity
[docs]
@autodoc_function
def write_kisslinger_flux_surfaces(out_directory: str, prefix: str,
polygon_set: PoincarePolygonSet,
r0_cm: float | Quantity,
z0_cm: float | Quantity=0.0,
mtor: int=1,
skip_first: bool=True) -> list[str]:
"""
Write Kisslinger-format flux-surface files from a PoincarePolygonSet.
For each exported surface a ``{prefix}_isurf_NNN.txt`` file is written
with the poloidal polygon at every toroidal plane, and a single
``{prefix}_flux_labels.txt`` file lists the unnormalized mean area and
the normalized rho label of each surface. Polygon coordinates are
assumed normalized to the axis radius and are scaled by ``r0_cm``
(with ``z0_cm`` added to Z) to obtain centimeters.
Parameters
----------
out_directory : str
Directory the files are written to. Created if it does not exist.
prefix : str
Output file-name prefix.
polygon_set : PoincarePolygonSet
The processed flux-surface polygons. Its ``mean_area`` and
``rho_data`` labels are written directly to the file.
r0_cm : float or Quantity
Axis major radius used to scale the normalized polygons. A
bare float is assumed to already be in cm; a Quantity is
converted to cm.
z0_cm : float or Quantity, optional
Vertical offset added to every Z coordinate. A bare float is
assumed to already be in cm; a Quantity is converted to cm.
mtor : int, optional
Number of toroidal field periods, written into each surface
header.
skip_first : bool, optional
Skip the first surface (the magnetic axis, a degenerate point),
as the reference tool does.
Returns
-------
list of str
The paths of the surface files written, in order.
"""
if isinstance(r0_cm, Quantity):
r0_cm = r0_cm.to("cm").magnitude
if isinstance(z0_cm, Quantity):
z0_cm = z0_cm.to("cm").magnitude
out_path = Path(out_directory)
out_path.mkdir(parents=True, exist_ok=True)
mean_area = np.asarray(polygon_set.mean_area.values, float)
rho = np.asarray(polygon_set.rho_data.values, float)
phi_deg = np.degrees(np.asarray(polygon_set.phi_array.values, float))
n_surfaces = polygon_set.n_surfaces
n_planes = polygon_set.n_planes
first = 1 if skip_first else 0
written = []
for s in range(first, n_surfaces):
fpath = out_path / f"{prefix}_isurf_{s:03d}.txt"
_write_surface_file(fpath, polygon_set, s, n_planes, phi_deg,
mean_area[s], r0_cm, z0_cm, mtor)
written.append(str(fpath))
label_path = out_path / f"{prefix}_flux_labels.txt"
with label_path.open("w") as label_file:
label_file.write(f"{n_surfaces - first}\n")
for s in range(first, n_surfaces):
label_file.write(
f"{r0_cm**2 * mean_area[s]:.5e} {rho[s]:.5e}\n")
return written
[docs]
@autodoc_function
def read_kisslinger_flux_surfaces(folder: str, prefix: str="") -> dict:
"""
Read Kisslinger-format flux-surface files back into arrays.
Parses every ``{prefix}*isurf_*.txt`` file in ``folder`` (and the companion
``{prefix}_flux_labels.txt`` label file if present) into coordinate
arrays and per-surface labels. Coordinates are returned in the units
they were written in (centimeters).
Parameters
----------
folder : str
Directory containing the isurf text files.
prefix : str, optional
Only read files whose name starts with this prefix.
Returns
-------
dict
Keys: ``paths`` (list of str, in order); ``phi_deg`` (n_planes
array of toroidal angles in degrees); ``mtor`` (int, field
periods); ``R`` and ``Z`` (each a list of (n_planes, n_points)
arrays, one per surface, in cm); and ``mean_area`` / ``rho``
(n_surfaces arrays from the label file, or None when it is
absent).
Raises
------
FileNotFoundError
If ``folder`` does not exist, or no matching surface files are
found within it.
"""
folder_path = Path(folder)
if not folder_path.is_dir():
raise FileNotFoundError(f"No such folder: {folder}")
files = sorted(folder_path.glob(f"{prefix}*isurf_*.txt"))
if not files:
raise FileNotFoundError(
f"No files matching {prefix}*isurf_*.txt in {folder}")
phi_deg = None
mtor = None
r_surfaces = []
z_surfaces = []
for path in files:
phi_deg, r, z, mtor = _read_surface_file(path)
r_surfaces.append(r)
z_surfaces.append(z)
result = {
"paths": [str(p) for p in files],
"phi_deg": phi_deg,
"mtor": mtor,
"R": r_surfaces,
"Z": z_surfaces,
"mean_area": None,
"rho": None,
}
label_path = folder_path / f"{prefix}_flux_labels.txt"
if label_path.exists():
lines = label_path.read_text().splitlines()
n_surfaces = int(lines[0])
labels = np.array([[float(v) for v in line.split()]
for line in lines[1:1 + n_surfaces]])
result["mean_area"] = labels[:, 0]
result["rho"] = labels[:, 1]
return result
def _read_surface_file(path: Path) -> tuple:
"""Read one Kisslinger surface file into phi_deg, R, Z and mtor."""
lines = path.read_text().splitlines()
header = lines[1].split()
n_planes, n_points, mtor = int(header[0]), int(header[1]), int(header[2])
phi_deg = np.empty(n_planes)
r = np.empty((n_planes, n_points))
z = np.empty((n_planes, n_points))
idx = 2
for ip in range(n_planes):
phi_deg[ip] = float(lines[idx])
idx += 1
coords = np.array(
[line.split() for line in lines[idx:idx + n_points]], float)
r[ip] = coords[:, 0]
z[ip] = coords[:, 1]
idx += n_points
return phi_deg, r, z, mtor
def _write_surface_file(fpath: Path, polygon_set: PoincarePolygonSet,
s: int, n_planes: int, phi_deg: np.ndarray,
mean_area_s: float, r0_cm: float, z0_cm: float,
mtor: int) -> None:
"""Write one surface's polygons across all planes to fpath."""
first_polygon = polygon_set.polygon_data.isel(
dim_surf=s, dim_plane=0).item()
n_points = len(first_polygon.x_points)
descr = (f"Flux surface is={s:03d}, mean_flux_area ={mean_area_s:.5e}")
with fpath.open("w") as surface_file:
surface_file.write(descr + "\n")
surface_file.write(f"{n_planes} {n_points} {mtor} "
f"{0.0:.5e} {0.0:.5e}\n")
for ip in range(n_planes):
surface_file.write(f"{phi_deg[ip]:.5e} \n")
polygon = polygon_set.polygon_data.isel(
dim_surf=s, dim_plane=ip).item()
r = r0_cm * np.asarray(polygon.x_points, float)
z = r0_cm * np.asarray(polygon.y_points, float) + z0_cm
for k in range(n_points):
surface_file.write(f"{r[k]:.5e} {z[k]:.5e} \n")
