From 53c48f4c98438cb38dc06467f6e65d575dabb3e0 Mon Sep 17 00:00:00 2001 From: Peter Hill Date: Fri, 20 Sep 2024 14:00:54 +0100 Subject: [PATCH] Switch to freeqdsk for reading/writing geqdsk files Closes #179 --- hypnotoad/cases/tokamak.py | 2 +- hypnotoad/cases/torpex.py | 4 +- hypnotoad/geqdsk/__init__.py | 0 hypnotoad/geqdsk/_fileutils.py | 126 ---------- hypnotoad/geqdsk/_geqdsk.py | 307 ------------------------- hypnotoad/test_suite/test_fileutils.py | 23 -- hypnotoad/test_suite/test_geqdsk.py | 49 ---- hypnotoad/test_suite/test_tokamak.py | 4 +- hypnotoad/test_suite/test_torpex.py | 2 +- pyproject.toml | 1 + 10 files changed, 7 insertions(+), 511 deletions(-) delete mode 100644 hypnotoad/geqdsk/__init__.py delete mode 100644 hypnotoad/geqdsk/_fileutils.py delete mode 100644 hypnotoad/geqdsk/_geqdsk.py delete mode 100644 hypnotoad/test_suite/test_fileutils.py delete mode 100644 hypnotoad/test_suite/test_geqdsk.py diff --git a/hypnotoad/cases/tokamak.py b/hypnotoad/cases/tokamak.py index ea5c98b9..c358f759 100644 --- a/hypnotoad/cases/tokamak.py +++ b/hypnotoad/cases/tokamak.py @@ -1719,7 +1719,7 @@ def read_geqdsk( if settings is None: settings = {} - from ..geqdsk._geqdsk import read as geq_read + from freeqdsk.geqdsk import read as geq_read data = geq_read(filehandle) diff --git a/hypnotoad/cases/torpex.py b/hypnotoad/cases/torpex.py index d03f8a7b..193dedd0 100644 --- a/hypnotoad/cases/torpex.py +++ b/hypnotoad/cases/torpex.py @@ -20,6 +20,7 @@ from collections import OrderedDict import warnings +from freeqdsk.geqdsk import read as geq_read import numpy from optionsfactory import WithMeta from optionsfactory.checks import is_positive @@ -31,7 +32,6 @@ EquilibriumRegion, SolutionError, ) -from ..geqdsk._geqdsk import read as geq_read from ..utils.utils import with_default # type for manipulating information about magnetic field coils @@ -688,7 +688,7 @@ def createMesh(filename): def createEqdsk(equilib, *, nR, Rmin, Rmax, nZ, Zmin, Zmax, filename="torpex_test.g"): - from ..geqdsk._geqdsk import write as geq_write + from freeqdsk.geqdsk import write as geq_write R = numpy.linspace(Rmin, Rmax, nR)[numpy.newaxis, :] Z = numpy.linspace(Zmin, Zmax, nZ)[:, numpy.newaxis] diff --git a/hypnotoad/geqdsk/__init__.py b/hypnotoad/geqdsk/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/hypnotoad/geqdsk/_fileutils.py b/hypnotoad/geqdsk/_fileutils.py deleted file mode 100644 index 5006b5f4..00000000 --- a/hypnotoad/geqdsk/_fileutils.py +++ /dev/null @@ -1,126 +0,0 @@ -""" -Utilities for writing and reading files compatible with Fortran - -""" - -import re - - -def f2s(f): - """ - Format a string containing a float - """ - s = "" - if f >= 0.0: - s += " " - return s + "%1.9E" % f - - -class ChunkOutput: - """ - This outputs values in lines, inserting - newlines when needed. - """ - - def __init__(self, filehandle, chunksize=5, extraspaces=0): - """ - filehandle output to write to - chunksize number of values on a line - extraspaces number of extra spaces between outputs - """ - self.fh = filehandle - self.counter = 0 - self.chunk = chunksize - self.extraspaces = extraspaces - - def write(self, value): - """ - Write a value to the output, adding a newline if needed - - Distinguishes between: - - list : Iterates over the list and writes each element - - int : Converts using str - - float : Converts using f2s to Fortran-formatted string - - """ - if isinstance(value, list): - for elt in value: - self.write(elt) - return - - self.fh.write(" " * self.extraspaces) - - if isinstance(value, int): - self.fh.write(" " + str(value)) - else: - self.fh.write(f2s(value)) - - self.counter += 1 - if self.counter == self.chunk: - self.fh.write("\n") - self.counter = 0 - - def newline(self): - """ - Ensure that the file is at the start of a new line - """ - if self.counter != 0: - self.fh.write("\n") - self.counter = 0 - - def endblock(self): - """ - Make sure next block of data is on new line - """ - self.fh.write("\n") - self.counter = 0 - - def __enter__(self): - return self - - def __exit__(self, type, value, traceback): - """Ensure that the chunk finishes with a new line""" - if self.counter != 0: - self.counter = 0 - self.fh.write("\n") - - -def write_1d(val, out): - """ - Writes a 1D variable val to the file handle out - """ - for i in range(len(val)): - out.write(val[i]) - out.newline() - - -def write_2d(val, out): - """ - Writes a 2D array. Note that this transposes - the array, looping over the first index fastest - """ - nx, ny = val.shape - for y in range(ny): - for x in range(nx): - out.write(val[x, y]) - out.newline() - - -def next_value(fh): - """ - A generator which yields values from a file handle - - Checks if the value is a float or int, returning - the correct type depending on if '.' is in the string - - """ - pattern = re.compile(r"[ +\-]?\d+(?:\.\d+[Ee][\+\-]\d\d)?") - - # Go through each line, extract values, then yield them one by one - for line in fh: - matches = pattern.findall(line) - for match in matches: - if "." in match: - yield float(match) - else: - yield int(match) diff --git a/hypnotoad/geqdsk/_geqdsk.py b/hypnotoad/geqdsk/_geqdsk.py deleted file mode 100644 index ed643216..00000000 --- a/hypnotoad/geqdsk/_geqdsk.py +++ /dev/null @@ -1,307 +0,0 @@ -""" -Low level routines for reading and writing G-EQDSK files - -Copyright 2016 Ben Dudson, University of York. Email: benjamin.dudson@york.ac.uk - -This file is part of FreeGS. - -FreeGS is free software: you can redistribute it and/or modify -it under the terms of the GNU Lesser General Public License as published by -the Free Software Foundation, either version 3 of the License, or -(at your option) any later version. - -FreeGS is distributed in the hope that it will be useful, -but WITHOUT ANY WARRANTY; without even the implied warranty of -MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -GNU Lesser General Public License for more details. - -You should have received a copy of the GNU Lesser General Public License -along with FreeGS. If not, see . - -""" - -from datetime import date -from numpy import zeros, pi - -from ._fileutils import f2s, ChunkOutput, write_1d, write_2d, next_value - - -def write(data, fh, label=None, shot=None, time=None): - """ - Write a GEQDSK equilibrium file, given a dictionary of data - - data - dictionary - nx, ny Number of points in R (x), Z (y) - rdim, zdim Sizes of the R,Z dimensions - rcentr Reference value of R - bcentr Vacuum toroidal magnetic field at rcentr - rleft R at left (inner) boundary - zmid Z at middle of domain - rmagx, zmagx R,Z at magnetic axis (O-point) - simagx Poloidal flux psi at magnetic axis - sibdry Poloidal flux psi at plasma boundary - cpasma Plasma current [Amps] - - fpol 1D array of f(psi)=R*Bt [meter-Tesla] - pres 1D array of p(psi) [Pascals] - qpsi 1D array of q(psi) - - psi 2D array (nx,ny) of poloidal flux - - fh - file handle - - label - Text label to put in the file - """ - - nx = data["nx"] - ny = data["ny"] - - if not label: - label = "FREEGS" - if len(label) > 11: - label = label[0:12] - print("WARNING: label too long, it will be shortened to {}".format(label)) - - creation_date = date.today().strftime("%d/%m/%Y") - - if not shot: - shot = 0 - - if isinstance(shot, int): - shot = "# {:d}".format(shot) - - if not time: - time = 0 - - if isinstance(time, int): - time = " {:d}ms".format(time) - - # I have no idea what idum is, here it is set to 3 - idum = 3 - header = "{0:11s}{1:10s} {2:>8s}{3:16s}{4:4d}{5:4d}{6:4d}\n".format( - label, creation_date, shot, time, idum, nx, ny - ) - - # First line: Identification string, followed by resolution - fh.write(header) - - # Second line - fh.write( - f2s(data["rdim"]) - + f2s(data["zdim"]) - + f2s(data["rcentr"]) - + f2s(data["rleft"]) - + f2s(data["zmid"]) - + "\n" - ) - - # Third line - fh.write( - f2s(data["rmagx"]) - + f2s(data["zmagx"]) - + f2s(data["simagx"]) - + f2s(data["sibdry"]) - + f2s(data["bcentr"]) - + "\n" - ) - - # 4th line - fh.write( - f2s(data["cpasma"]) - + f2s(data["simagx"]) - + f2s(0.0) - + f2s(data["rmagx"]) - + f2s(0.0) - + "\n" - ) - - # 5th line - fh.write( - f2s(data["zmagx"]) + f2s(0.0) + f2s(data["sibdry"]) + f2s(0.0) + f2s(0.0) + "\n" - ) - - # SCENE has actual ff' and p' data so can use that - # fill arrays - # Lukas Kripner (16/10/2018): uncommenting this, since you left there - # check for data existence bellow. This seems to as safer variant. - workk = zeros([nx]) - - # Write arrays - co = ChunkOutput(fh) - - write_1d(data["fpol"], co) - write_1d(data["pres"], co) - if "ffprime" in data: - write_1d(data["ffprime"], co) - else: - write_1d(workk, co) - if "pprime" in data: - write_1d(data["pprime"], co) - else: - write_1d(workk, co) - - write_2d(data["psi"], co) - write_1d(data["qpsi"], co) - - # Boundary / limiters - nbdry = 0 - nlim = 0 - if "rbdry" in data: - nbdry = len(data["rbdry"]) - if "rlim" in data: - nlim = len(data["rlim"]) - - co.newline() - fh.write("{0:5d}{1:5d}\n".format(nbdry, nlim)) - - if nbdry > 0: - for r, z in zip(data["rbdry"], data["zbdry"]): - co.write(r) - co.write(z) - co.newline() - - if nlim > 0: - for r, z in zip(data["rlim"], data["zlim"]): - co.write(r) - co.write(z) - co.newline() - - -def read(fh, cocos=1): - """ - Read a G-EQDSK formatted equilibrium file - - Format is specified here: - https://fusion.gat.com/theory/Efitgeqdsk - - cocos - COordinate COnventions. Not fully handled yet, - only whether psi is divided by 2pi or not. - if < 10 then psi is divided by 2pi, otherwise not. - - Returns - ------- - - A dictionary containing: - nx, ny Number of points in R (x), Z (y) - rdim, zdim Sizes of the R,Z dimensions - rcentr Reference value of R - bcentr Vacuum toroidal magnetic field at rcentr - rleft R at left (inner) boundary - zmid Z at middle of domain - rmagx, zmagx R,Z at magnetic axis (O-point) - simagx Poloidal flux psi at magnetic axis - sibdry Poloidal flux psi at plasma boundary - cpasma Plasma current [Amps] - - fpol 1D array of f(psi)=R*Bt [meter-Tesla] - pres 1D array of p(psi) [Pascals] - qpsi 1D array of q(psi) - - psi 2D array (nx,ny) of poloidal flux - - """ - - # Read the first line - header = fh.readline() - words = header.split() # Split on whitespace - if len(words) < 3: - raise ValueError("Expecting at least 3 numbers on first line") - - idum = int(words[-3]) # noqa: F841 - nx = int(words[-2]) - ny = int(words[-1]) - - print(" nx = {0}, ny = {1}".format(nx, ny)) - - # Dictionary to hold result - data = {"nx": nx, "ny": ny} - - # List of fields to read. None means discard value - fields = [ - "rdim", - "zdim", - "rcentr", - "rleft", - "zmid", - "rmagx", - "zmagx", - "simagx", - "sibdry", - "bcentr", - "cpasma", - "simagx", - None, - "rmagx", - None, - "zmagx", - None, - "sibdry", - None, - None, - ] - - values = next_value(fh) - - for f in fields: - val = next(values) - if f: - data[f] = val - - # Read arrays - - def read_1d(n): - """ - Read a 1D array of length n - """ - val = zeros(n) - for i in range(n): - val[i] = next(values) - return val - - def read_2d(n, m): - """ - Read a 2D (n,m) array in Fortran order - """ - val = zeros([n, m]) - for y in range(m): - for x in range(n): - val[x, y] = next(values) - return val - - data["fpol"] = read_1d(nx) - data["pres"] = read_1d(nx) - data["ffprime"] = read_1d(nx) - data["pprime"] = read_1d(nx) - - data["psi"] = read_2d(nx, ny) - - data["qpsi"] = read_1d(nx) - - # Ensure that psi is divided by 2pi - if cocos > 10: - for var in ["psi", "simagx", "sibdry"]: - data[var] /= 2 * pi - - nbdry = next(values) - nlim = next(values) - - # print(nbdry, nlim) - - if nbdry > 0: - # Read (R,Z) pairs - print(nbdry) - data["rbdry"] = zeros(nbdry) - data["zbdry"] = zeros(nbdry) - for i in range(nbdry): - data["rbdry"][i] = next(values) - data["zbdry"][i] = next(values) - - if nlim > 0: - # Read (R,Z) pairs - data["rlim"] = zeros(nlim) - data["zlim"] = zeros(nlim) - for i in range(nlim): - data["rlim"][i] = next(values) - data["zlim"][i] = next(values) - - return data diff --git a/hypnotoad/test_suite/test_fileutils.py b/hypnotoad/test_suite/test_fileutils.py deleted file mode 100644 index bba1ad89..00000000 --- a/hypnotoad/test_suite/test_fileutils.py +++ /dev/null @@ -1,23 +0,0 @@ -from hypnotoad.geqdsk import _fileutils -from io import StringIO - - -def test_f2s(): - assert _fileutils.f2s(0.0) == " 0.000000000E+00" - assert _fileutils.f2s(1234) == " 1.234000000E+03" - assert _fileutils.f2s(-1.65281e12) == "-1.652810000E+12" - assert _fileutils.f2s(-1.65281e-2) == "-1.652810000E-02" - - -def test_ChunkOutput(): - output = StringIO() - co = _fileutils.ChunkOutput(output) - - for val in [1.0, -3.2, 6.2e5, 8.7654e-12, 42.0, -76]: - co.write(val) - - assert ( - output.getvalue() - == """ 1.000000000E+00-3.200000000E+00 6.200000000E+05 8.765400000E-12 4.200000000E+01 - -76""" # noqa: E501 - ) diff --git a/hypnotoad/test_suite/test_geqdsk.py b/hypnotoad/test_suite/test_geqdsk.py deleted file mode 100644 index ea843cf4..00000000 --- a/hypnotoad/test_suite/test_geqdsk.py +++ /dev/null @@ -1,49 +0,0 @@ -import numpy - -from io import StringIO - -from hypnotoad.geqdsk import _geqdsk - - -def test_writeread(): - """ - Test that data can be written then read back - """ - nx = 65 - ny = 65 - - # Create a dummy dataset - data = { - "nx": nx, - "ny": ny, - "rdim": 2.0, - "zdim": 1.5, - "rcentr": 1.2, - "bcentr": 2.42, - "rleft": 0.5, - "zmid": 0.1, - "rmagx": 1.1, - "zmagx": 0.2, - "simagx": -2.3, - "sibdry": 0.21, - "cpasma": 1234521, - "fpol": numpy.random.rand(nx), - "pres": numpy.random.rand(nx), - "qpsi": numpy.random.rand(nx), - "psi": numpy.random.rand(nx, ny), - } - - output = StringIO() - - # Write to string - _geqdsk.write(data, output) - - # Move to the beginning of the buffer - output.seek(0) - - # Read from string - data2 = _geqdsk.read(output) - - # Check that data and data2 are the same - for key in data: - numpy.testing.assert_allclose(data2[key], data[key]) diff --git a/hypnotoad/test_suite/test_tokamak.py b/hypnotoad/test_suite/test_tokamak.py index f7fdd634..af655a82 100644 --- a/hypnotoad/test_suite/test_tokamak.py +++ b/hypnotoad/test_suite/test_tokamak.py @@ -1,9 +1,9 @@ +from freeqdsk.geqdsk import write as geq_write import numpy as np from io import StringIO import pytest from hypnotoad.cases import tokamak -from hypnotoad.geqdsk import _geqdsk @pytest.mark.parametrize("psi_interpolation_method", ["spline", "dct"]) @@ -146,7 +146,7 @@ def fpolprime_func(psi): # Write to string output = StringIO() - _geqdsk.write(data, output) + geq_write(data, output) # Move to the beginning of the buffer output.seek(0) diff --git a/hypnotoad/test_suite/test_torpex.py b/hypnotoad/test_suite/test_torpex.py index ddd69f36..acbb5124 100644 --- a/hypnotoad/test_suite/test_torpex.py +++ b/hypnotoad/test_suite/test_torpex.py @@ -53,7 +53,7 @@ def test_Eqdsk(self, equilib): R = numpy.linspace(Rmin, Rmax, nR)[numpy.newaxis, :] Z = numpy.linspace(Zmin, Zmax, nZ)[:, numpy.newaxis] - assert equilib.psi(R, Z) == pytest.approx(grid_equilib.psi(R, Z), rel=1.0e-9) + assert equilib.psi(R, Z) == pytest.approx(grid_equilib.psi(R, Z), rel=1.0e-8) import os diff --git a/pyproject.toml b/pyproject.toml index 374e5fbf..b4f1dd01 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,6 +36,7 @@ dependencies = [ "PyYAML>=5.1", "scipy~=1.6", "Qt.py~=1.2", + "freeqdsk>=0.4.0", ] [project.optional-dependencies]