From ecc816b4a7196911c19259c50be4ee34e0253837 Mon Sep 17 00:00:00 2001 From: Peter Hill Date: Fri, 20 Sep 2024 15:55:57 +0100 Subject: [PATCH] Return a dataclass from `geqdsk.read` Closes #17 --- docs/source/geqdsk.rst | 1 + freeqdsk/geqdsk.py | 174 +++++++++++++++++++++++++++++++++++++++-- tests/test_geqdsk.py | 5 +- 3 files changed, 173 insertions(+), 7 deletions(-) diff --git a/docs/source/geqdsk.rst b/docs/source/geqdsk.rst index 5b11d7c..3192b3f 100644 --- a/docs/source/geqdsk.rst +++ b/docs/source/geqdsk.rst @@ -6,3 +6,4 @@ geqdsk .. automodule:: freeqdsk.geqdsk .. automethod:: freeqdsk.geqdsk.read .. automethod:: freeqdsk.geqdsk.write +.. autoclass:: freeqdsk.geqdsk.GEQDSKFile diff --git a/freeqdsk/geqdsk.py b/freeqdsk/geqdsk.py index da585f1..353831b 100644 --- a/freeqdsk/geqdsk.py +++ b/freeqdsk/geqdsk.py @@ -138,8 +138,9 @@ from __future__ import annotations # noqa import warnings +from dataclasses import asdict, dataclass from datetime import date -from typing import Dict, Optional, TextIO, Union, TypedDict +from typing import Optional, TextIO, TypedDict import numpy as np @@ -182,6 +183,163 @@ ) +def _synonym(canonical: str) -> property: + """Set a property which is just a synonym for another attribute""" + + return property( + lambda self: getattr(self, canonical), + lambda self, value: setattr(self, canonical, value), + ) + + +@dataclass() +class GEQDSKFile: + r"""G-EQDSK equilibrium. + + Includes some common synonyms for the "canonical" names, as well as some + more human-readable ones. In the *Attributes* section below, canonical names + are listed first. + + When creating an instance, if keyword arguments are used, they should be the + canonical names. + + Attributes + ---------- + comment: + Header comment + shot: + Header shot number + nx, nw, nr: + Number of radial points + ny, nh, nz: + Number of vertical points + rdim: + Width of computational domain in R direction, float [meter] + zdim: + Height of computational domain in Z direction, float [meter] + rcentr: + Reference value of R, float [meter] + rleft: + R at left (inner) boundary, float [meter] + zmid: + Z at middle of domain, float [meter] + rmagx, rmaxis: + R at magnetic axis (0-point), float [meter] + zmagx, zmaxis: + Z at magnetic axis (0-point), float [meter] + simagx, psi_axis: + Poloidal flux :math:`\psi` at magnetic axis, float [weber / radian] + sibdry, psi_boundary: + Poloidal flux :math:`\psi` at plasma boundary, float [weber / radian] + bcentr: + Vacuum toroidal magnetic field at rcentr, float [tesla] + cpasma, current: + Plasma current, float [ampere] + fpol: + Poloidal current function :math:`F(\psi)=RB_t`, 1D array [meter * tesla] + pres, pressure: + Plasma pressure :math:`p(\psi)`, 1D array [pascal] + ffprime: + :math:`FF'(\psi)`, 1D array [meter**2 * tesla**2 * radian / weber] + pprime: + :math:`p'(\psi)`, 1D array [pascal * radian / weber] + psi, f, psirz: + Poloidal flux :math:`\psi`, 2D array [weber / radian] + qpsi: + Safety factor :math:`q(\psi)`, 1D array [dimensionless] + nbdry, nbbbs: + Number of points in the boundary grid, int + nlim, limitr: + Number of points in the limiter grid, int + rbdry, rbbbs: + R of boundary points, 1D array [meter] + zbdry, zbbbs: + Z of boundary points, 1D array [meter] + rlim: + R of limiter points, 1D array [meter] + zlim: + Z of limiter points, 1D array [meter] + + Examples + -------- + + Variables can be accessed with either the attribute dot syntax or ``dict`` + item access. For example, the poloidal flux on the magnetic axis can be + accessed via any one of the following: + + >>> gfile = GEQDSKFile(...) + >>> print(gfile["simagx"]) + >>> print(gfile.simagx) + >>> print(gfile.psi_axis) + + """ + + comment: str + shot: int + nx: int + ny: int + rdim: float + zdim: float + rcentr: float + rleft: float + zmid: float + rmagx: float + zmagx: float + simagx: float + sibdry: float + bcentr: float + cpasma: float + fpol: FloatArray + pres: FloatArray + ffprime: FloatArray + pprime: FloatArray + psi: FloatArray + qpsi: FloatArray + nbdry: int + nlim: int + rbdry: FloatArray | None = None + zbdry: FloatArray | None = None + rlim: FloatArray | None = None + zlim: FloatArray | None = None + + def __post_init__(self): + r = np.zeros((self.nx, self.ny), float) + z = np.zeros((self.nx, self.ny), float) + + for i in range(0, self.nx): + for j in range(0, self.ny): + r[i, j] = self.rleft + self.rdim * i / (self.nx - 1) + z[i, j] = (self.zmid - 0.5 * self.zdim) + self.zdim * j / (self.ny - 1) + + self.r_grid = r + self.z_grid = z + + nw = _synonym("nx") + nr = _synonym("nx") + nh = _synonym("ny") + nz = _synonym("ny") + rmaxis = _synonym("rmagx") + zmaxis = _synonym("zmagx") + psi_axis = _synonym("simagx") + psi_boundary = _synonym("sibdry") + pressure = _synonym("pres") + current = _synonym("cpasma") + f = _synonym("psi") + psirz = _synonym("psi") + limitr = _synonym("nlim") + xlim = _synonym("rlim") + ylim = _synonym("zlim") + nbbbs = _synonym("nbdry") + rbbbs = _synonym("rbdry") + zbbbs = _synonym("zbdry") + + def __getitem__(self, name: str): + return getattr(self, name) + + def __setitem__(self, name: str, value): + return setattr(self, name, any) + + class GeqdskDataDict(TypedDict): """Names and expected types of keys in dict of G-EQDSK data""" @@ -212,7 +370,7 @@ class GeqdskDataDict(TypedDict): def write( - data: GeqdskDataDict, + data: GEQDSKFile | GeqdskDataDict, fh: TextIO, label: Optional[str] = None, shot: int = 0, @@ -270,6 +428,12 @@ def write( if bdry_lim_fmt is None: bdry_lim_fmt = _bdry_lim_fmt + if isinstance(data, GEQDSKFile): + # Convert to dict just so we handle the rest of the function the + # same. Although, if we do have a `GEQDSKFile`, then we should already + # know we have everything we need + data = asdict(data) + # Get dimensions and check data is correct nx = data.get("nx", np.shape(data["psi"])[0]) ny = data.get("ny", np.shape(data["psi"])[1]) @@ -353,7 +517,7 @@ def read( header_fmt: Optional[str] = None, data_fmt: Optional[str] = None, bdry_lim_fmt: Optional[str] = None, -) -> Dict[str, Union[int, float, np.ndarray]]: +) -> GEQDSKFile: r""" Read a G-EQDSK formatted equilibrium file. The format is specified `here `_. @@ -397,7 +561,7 @@ def read( comment, integer, nx, ny = read_line(fh, header_fmt) # Dictionary to hold result - data = {"comment": comment, "int": integer, "nx": nx, "ny": ny} + data = {"comment": comment, "shot": integer, "nx": nx, "ny": ny} # Read first four lines floats = read_array(20, fh, data_fmt) @@ -440,4 +604,4 @@ def read( lim = read_array(2 * nlim, fh, data_fmt) data["rlim"], data["zlim"] = lim[0::2], lim[1::2] - return data + return GEQDSKFile(**data) diff --git a/tests/test_geqdsk.py b/tests/test_geqdsk.py index 1f1a175..a331124 100644 --- a/tests/test_geqdsk.py +++ b/tests/test_geqdsk.py @@ -4,6 +4,7 @@ SPDX-License-Identifier: MIT """ +from dataclasses import asdict from io import StringIO from difflib import unified_diff from pathlib import Path @@ -170,7 +171,7 @@ def test_write(path, tmp_path): def test_write_unrecoverable_missing_data(tmp_path): # read in test data with open(_data_path / "test_1.geqdsk") as f: - data = geqdsk.read(f) + data = asdict(geqdsk.read(f)) # Delete necessary data data.pop("psi") @@ -186,7 +187,7 @@ def test_write_unrecoverable_missing_data(tmp_path): def test_write_recoverable_missing_data(tmp_path): # read in test data with open(_data_path / "test_1.geqdsk") as f: - data = geqdsk.read(f) + data = asdict(geqdsk.read(f)) # Delete superfluous data data.pop("nx")