scripts/gray_visual.py: rewrite and extend scope

scripts/gray.py

@@ -0,0 +1,115 @@
+'''
+Python interface for GRAY
+'''
+
+# output files handling
+import numpy as np
+import collections
+import io
+from configparser import ConfigParser
+
+# standard GRAY simulation
+from pathlib import Path
+
+
+# Type of a parsed EQDSK file
+Eqdsk = collections.namedtuple('Eqdsk', ['flux', 'boundary', 'limiter', 'q'])
+
+
+def read_eqdsk(filepath: Path) -> Eqdsk:
+    '''
+    Loads a full G-EQDSK file and returns:
+    R, z, ψ(R, z), q(ρ)
+    '''
+    def read(file, n):
+        '''
+        Reads n records of a G-EQDSK file
+        '''
+        x = np.empty(n)
+        for i in range(n):
+            x[i] = float(file.read(16).strip())
+        return x
+
+    def skip(file, n):
+        '''
+        Skip n records of a G-EQDSK file
+        '''
+        file.seek(file.tell() + 16 * n)
+
+    with open(filepath) as file:
+        file = io.StringIO(''.join(file.read().splitlines()))
+
+    # comment string
+    file.read(48)
+
+    # number of (R, z) grid points
+    _, nr, nz = [int(i) for i in file.read(4 * 3).split()]
+
+    # build (R, z) grid
+    r_dim, z_dim, _, r_left, z_mid = read(file, 5)
+    r = np.linspace(r_left, r_left + r_dim, nr)
+    z = np.linspace(z_mid - z_dim/2, z_mid + z_dim/2, nz)
+
+    # ψ on the axis/boundary
+    psi_min, psi_max = read(file, 5)[2:4]
+
+    # skip other MHD quantities
+    skip(file, 5 * 2)
+    skip(file, nr * 4)
+
+    # read flux and safety factor
+    psi = read(file, nr * nz).reshape(nr, nz)
+    q = read(file, nr)
+
+    # normalise flux
+    psi = abs(psi - psi_min) / abs(psi_max - psi_min)
+
+    # plasma boundary and limiter contour
+    nbound, nlim = [int(i) for i in file.read(5 * 2).split()]
+    rbound, zbound = read(file, nbound * 2).reshape(nbound, 2).T
+    rlim, zlim = read(file, nlim * 2).reshape(nlim, 2).T
+
+    return Eqdsk(flux=(r, z, psi),
+                 boundary=(rbound, zbound),
+                 limiter=(rlim, zlim),
+                 q=(np.sqrt(np.linspace(0, 1, nr)), abs(q)))
+
+
+def read_conf(fname: Path):
+    '''
+    Parses gray.ini configuration file
+    '''
+    config = ConfigParser()
+    with open(fname) as ini:
+        config.read_file(ini)
+    return config
+
+
+def read_table(fname: Path):
+    '''
+    Loads a GRAY output unit file as a structured numpy array
+    '''
+    return np.genfromtxt(fname, skip_header=21, names=True)
+
+
+def get_limiter(conf: ConfigParser, inputs: Path) -> np.array:
+    '''
+    Returns the limiter contour
+    '''
+    file = conf['equilibrium']['filenm'].strip('"')
+
+    if conf['equilibrium'].get('iequil') == 'EQ_ANALYTICAL':
+        return np.loadtxt(inputs / file, skiprows=4).T
+    else:
+        eqdsk = read_eqdsk(inputs / file)
+        return eqdsk.limiter
+
+
+def decode_index_rt(i: int) -> (str, int):
+    '''
+    Given the index_rt of a beam, returns the polarisation
+    mode ('O' or 'X') and the number of passes
+    '''
+    mode = ['X', 'O'][i % 2]
+    passes = int(np.floor(np.log2(1 + i)))
+    return mode, passes