From 864cf23b7871a4319352267c1737ff659de05da6 Mon Sep 17 00:00:00 2001 From: Michele Guerini Rocco Date: Sat, 5 Oct 2024 15:03:11 +0200 Subject: [PATCH] scripts/gray_visual.py: rewrite and extend scope --- scripts/gray.py | 115 +++++ scripts/gray_visual.py | 1025 ++++++++++++++++++++-------------------- 2 files changed, 638 insertions(+), 502 deletions(-) create mode 100644 scripts/gray.py diff --git a/scripts/gray.py b/scripts/gray.py new file mode 100644 index 0000000..91ddaa5 --- /dev/null +++ b/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 diff --git a/scripts/gray_visual.py b/scripts/gray_visual.py index da9a81a..8604271 100644 --- a/scripts/gray_visual.py +++ b/scripts/gray_visual.py @@ -1,544 +1,565 @@ -import numpy as np -import matplotlib as mpl +''' +Script to quickly visualise both the inputs and output files of GRAY +Usage: python -m scripts.gray_visual -h +''' + +from pathlib import Path +from matplotlib.collections import LineCollection +from matplotlib.contour import ContourSet + +import argparse import matplotlib.pyplot as plt - -def white_to_alpha(r, g, b): - """Add transparency preserving the color when blended with white. - - RGB color components are mapped to HLS. A new color with lower luminance - (L_new < 0.5) and higher saturation (S_new = 1) and a transparency value - A_new are then computed, in such a way that the new color produces the - original color if blended with a white background. - - Parameters - ---------- - r, g, b : float - RGB components of the original color, in the [0, 1] range. - - Returns - ------- - r, g, b, a : float - RGBA components of the new color, with white component made transparent. - """ - import colorsys - - h, l, s = colorsys.rgb_to_hls(r, g, b) - if l > 0.5: - l_new = s/(1+s) - else: - l_new = 1 - (1-l)/(1 + l*(s-1)) - a = (1-l)/(1-l_new) - r_new, g_new, b_new = colorsys.hls_to_rgb(h, l_new, 1) - return r_new, g_new, b_new, a +import numpy as np +import scripts.gray as gray -def cmap_lin_alpha(name, min_loc=0, reverse_alpha=False, register=True): - """ - Set linear alpha (transparency) value in a colormap. - - A new colormap is returned, with unchanged RGB values, and transparency - changing linearly from A=0 (fully transparent) at min_alpha_loc to A=1 - (fully opaque) at the range ends. - - Parameters - ---------- - name : str - Name of builtin (or registered) colormap. - min_loc : float, default=0 - Location in the map of maximum transparency. - If min_loc=0, transparency decreases linearly from the lower end to - the upper end of the map. - If min_loc=1, transparency increases linearly from the lower end to - the upper end of the map. - If 0 argparse.Namespace: + ''' + Command line arguments + ''' + parser = argparse.ArgumentParser( + prog='gray_visual', + description='visualises the results of a GRAY simulation') + parser.add_argument('inputs', type=Path, + help='filepath of the inputs directory ' + '(containing gray.ini)') + parser.add_argument('outputs', type=Path, + help='filepath of the outputs directory') + parser.add_argument('--kind', '-k', default='raytracing', nargs='+', + choices=['raytracing', 'beam', 'profiles', 'inputs'], + help='which kind of information to visualise') + parser.add_argument('--outfile', '-o', + metavar='FILE', type=Path, default=None, + help='save the plot figure to FILE') + parser.add_argument('--interactive', '-i', action='store_true', + help='show the plot in an interactive window') + parser.add_argument('--legend', + action=argparse.BooleanOptionalAction, default=True, + help='whether to show plot legend') + return parser.parse_args() -def cmap_transp(name, register=True): - """ - Make transparent the white component in a colormap. - - A new colormap is returned, with new RGBA components that produce the - original RGB colors (with A=1) if blended with a white background. - - Parameters - ---------- - name : str - Name of builtin (or registered) colormap. - register : bool, default=True - Whether to register the new colormap. +def align_yaxis(*axes: [plt.Axes]): + ''' + Aligns the origins of two axes + ''' + extrema = np.array([ax.get_ylim() for ax in axes]) + tops = extrema[:, 1] / (extrema[:, 1] - extrema[:, 0]) - Returns - ------- - cmap : matplotlib.colors.ListedColormap - The new colormap. - """ - # Fetch the cmap callable - cmap = plt.get_cmap(name) - # Get the colors out from the colormap LUT - rgba = cmap(np.arange(cmap.N)) # N-by-4 - # Replace the old alpha channel - rgba = [white_to_alpha(*color[:3]) for color in rgba] - # Create a new Colormap object - cmap_alpha = mpl.colors.ListedColormap(rgba, name=name + '_transp') - if register: - mpl.cm.register_cmap(name=name + '_transp', cmap=cmap_alpha) - return cmap_alpha + # Ensure that plots (intervals) are ordered bottom to top: + if tops[0] > tops[1]: + axes, extrema, tops = [a[::-1] for a in (axes, extrema, tops)] + + # How much would the plot overflow if we kept current zoom levels? + tot_span = tops[1] + 1 - tops[0] + + extrema[0, 1] = extrema[0, 0] + tot_span * (extrema[0, 1] - extrema[0, 0]) + extrema[1, 0] = extrema[1, 1] + tot_span * (extrema[1, 0] - extrema[1, 1]) + [axes[i].set_ylim(*extrema[i]) for i in range(2)] -def cmap_partial(cmap, vmin=0, vmax=1): - """ - Make a new colormap from a section of an existing one. - - A new colormap is returned, with the same number of colors as the original - one, obtained by linear interpolation in a sub-range of the existing - colormap LUT. - - Parameters - ---------- - name : matplotlib.Colormap - The existing colormap a sub-section of which will become the new - colormap. - vmin : the lower end of the range in the LUT that will be mapped to the - minimum value (0) in the new colormap - vmax : the upper end of the range in the LUT that will be mapped to the - maximum value (1) in the new colormap +def add_alt_axis(ax: plt.Axes, x1: np.array, x2: np.array, + label: str) -> plt.Axes: + ''' + Adds an alternative x axis to the top of a plot. + x1: existing variable + x2: new variable + ''' + from matplotlib.ticker import AutoMinorLocator + from scipy.interpolate import interp1d - Returns - ------- - partial : matplotlib.colors.LinearSegmentedColormap - The new colormap. - """ - partial = mpl.colors.LinearSegmentedColormap.from_list( - cmap.name + '_partial', - cmap(np.linspace(vmin, vmax, cmap.N))) - return partial + ax.xaxis.set_minor_locator(AutoMinorLocator()) + ax.grid() + f = interp1d(x1, x2, fill_value='extrapolate') + g = interp1d(x2, x1, fill_value='extrapolate') + alt = ax.secondary_xaxis('top', functions=(f, g)) + alt.set_xlabel(label, labelpad=8) + alt.xaxis.set_minor_locator(AutoMinorLocator()) + return alt -def read_grayout(dirname, **kwargs): - """ - Make a new colormap from a section of an existing one. - - A new colormap is returned, with the same number of colors as the original - one, obtained by linear interpolation in a sub-range of the existing - colormap LUT. - - Parameters - ---------- - dirname : str - The path to the directory containing the output files of GRAY. - **kwargs : optional keyword arguments - Not yet implemented. Filenames will be passed here to override the - default (fort.*) values. - - Returns - ------- - run : dict - All the output data found in dirname, organized as a dictionary. The - field names in the structured arrays are derived from the column - headers in the output files. - - 'name' : str - The base name of the directory - - 'bres' : numpy.ndarray - Structured array containing the contours of the EC resonances. - - 'flux' : numpy.ndarray - Structured array conaining the contours of the poloidal flux. - - 'prof' : numpy.ndarray - Structured array containing the ECCD profiles. - - 'ray' : numpy.ndarray - Structured array describing the central ray of the beam. - - 'rays' : numpy.ndarray - Structured array describing the envelope of the beam. - - 'totals' : numpy.ndarray - Structured array conaining the integrated, 0D quantities. - """ - import os - import subprocess - import io - - HEADER_GRAY = 21 - UNITS_GRAY = { - 'ray': 'central_ray.4.txt', - 'totals': 'summary.7.txt', - 'rays': 'outer-rays.33.txt', - 'prof': 'ec-profiles.48.txt', - 'bres': 'ec-resonance.70.txt', - 'flux': 'flux-surfaces.71.txt', - } - - if not os.path.isdir(dirname): - print(f"{dirname} not a directory. Skipping...") - return None - if dirname.endswith("/"): dirname = dirname[:-1] - - name = os.path.basename(dirname) - run = dict(name=name) - for label, unit in UNITS_GRAY.items(): - file = dirname + f"/{unit}" - if label == 'bres' or label == 'flux': - # Replace blank lines with as many NaNs as the number of columns. - # In case of success, file is a string with the file's content. - # If the replacement fails, file remains the file's path. - try: - with open(file,'r') as f: - processed = subprocess.run(["awk", "BEGIN {ncol=0}; " - "($1 ~ /^[^#]/ && NF > ncol) {ncol=NF}; " - "{if (NF>0) {print $0} else " - r'{for (i=0; i plt.Circle: + ''' + Plots a circle at the origin + ''' + c = plt.Circle((0, 0), radius, linewidth=1.5, + fill=False, **args) + ax.add_patch(c) + ax.autoscale_view() + return c -def plot_ECprofiles(axs, profiles, label=None): - r""" - Plot EC absorption and ECCD profiles - - Parameters - ---------- - axs : list of matplotlib Axes - axs[0][0] : Power density - axs[0][1] : Power - axs[1][0] : Current density - axs[1][1] : Current - profiles : ndarray - Structured numpy array containing fields with names: - - 'rhot' : $\rho_t$, sqrt of normalized tor. flux, used as x variable - - 'dPdV' : $dP/dV$, power density - - 'Pins' : $P_{ins}(x) = \Int_0^x dP/dV dV/d\rho d\rho$, cumulated - absorbed power - - 'Jcd' : $J_{CD} = \ldots$, EC driven current density - - 'Icdins' : $I_{CD,ins}(x) = \Int_0^x J_\phi dA/d\rho d\rho$, cumulated - driven current +def plot_line_color(ax: plt.Axes, x: np.array, y: np.array, + c: np.array, **args) -> LineCollection: + ''' + Plots a line with variable color + ''' + points = np.array([x, y]).T.reshape(-1, 1, 2) + segments = np.concatenate([points[:-1], points[1:]], axis=1) - Returns - ------- - line_P, line_dP, line_I, line_J' : tuple of [lines2D] for the four plotted - elements. - """ - line_dP = axs[0][0].plot(profiles['rhot'], 1e3*profiles['dPdV'], label=label) - line_P = axs[0][1].plot(profiles['rhot'], profiles['Pins'], linestyle='--') - line_J = axs[1][0].plot(profiles['rhot'], 1e3*profiles['Jcdb']) - line_I = axs[1][1].plot(profiles['rhot'], 1e3*profiles['Icdins'], linestyle='--') - - return line_dP, line_P, line_J, line_I + lines = LineCollection(segments, **args) + lines.set_array(c) + ax.add_collection(lines) + ax.autoscale_view() + return lines -def plot_polview(ax, rays, flux, bres, - dpdsmax=None, dotsize=0.5, blobsize=20, marker='o', stride=1, label=None): - r""" - Plot rays path in 2D in a poloidal cross-section - - Expects to receive a structured numpy array with field names recognisable - as spatial coordinates ('x', 'y', 'z') or ('R', 'z'). If absorption - data are found ('alpha', 'tau'), the rays are coloured accordingly. - - Parameters - ---------- - ax : matplotlib Axes - Axes where the beam path will be plotted. - rays : ndarray - Structured numpy array with at least ('x', 'y', 'z') or ('R', - 'z') fields. If field names 'alpha' and 'tau' are found too, they will - be used to shade the beam path. - flux : ndarray - Structured numpy array with the (R, z) contours of the flux surfaces. - bres : ndarray - Structured numpy array with the (R, z) contours of the EC resonances. - dpdsmax : float, default=None - Normalization value for the size of markers along the beam. The - absorption markers will have size=blobsize where - $\alpha \exp(-\tau)$ = dpdsmax. - dotsize : float, optional, default=0.5 - Size of the dots used to represent the beam path. - blobsize : float, optional, default=20 - Size of the markers used to represent the absorption region along the - beam path. - marker : optional, default='o' - Shape of the absorption markers along the beam. - stride : int, optional, default=1 - label : str, optional, default=None - Unused. May be used in future for labelling. +def plot_poloidal(inputs: Path, outputs: Path, ax: plt.Axes): + ''' + Plots raytracing projected in the poloidal plane + ''' + conf = gray.read_conf(inputs / 'gray.ini') + central_ray = gray.read_table(outputs / 'central-ray.4.txt') - Returns - ------- - dots, blobs, flux, bres : matplotlib artists representing the beam path - with dots, the EC absorption region, the flux contours, and the EC - resonances. - """ - # Use colormaps with added transparency - # Warm (yellow to red) cmap for absorption coefficient - # Monochrome (black to white) cmap for residual power - cmap_alpha_name = 'autumn_r' - cmap_power_name = 'Blues' - dP_ds_threshold = 1e-8 + ax.set_title('poloidal view', loc='right') + ax.set_xlabel('$R$ / m') + ax.set_ylabel('$z$ / m') + + # load flux surfaces + surfaces = gray.read_table(outputs / 'flux-surfaces.71.txt') + surfaces = surfaces.reshape(-1, int(surfaces['i'].max())) + + # plot plasma boundary + bound = np.argmin(1 - surfaces['ψ_n'][:, 0]) + ax.plot(surfaces[bound]['R'], surfaces[bound]['z'], + c='xkcd:slate grey', label='plasma boundary') + + # manual contourplot of ψ(R,z) + for i, surf in enumerate(surfaces[:bound]): + ax.plot(surf['R'], surf['z'], c='xkcd:grey', + label='flux surfaces' if i == 0 else '') + + # label the rational surfaces + labels = ['q=2', 'q=3/2', 'q=1'] + for label, surf in zip(labels, surfaces[1+bound:][::-1]): + line = np.vstack([surf['R'], surf['z']]).T + cont = ContourSet(ax, [1], [[line]], filled=False, + colors='xkcd:ocean blue', + linestyles='--') + ax.clabel(cont, [1], inline=True, + fontsize=10, fmt={1: label}) + ax.plot(np.nan, np.nan, '--', color='xkcd:ocean blue', + label='rational surfaces') + + # load limiter + limiter = gray.get_limiter(conf, inputs) + ax.plot(*limiter, c='xkcd:black', label='limiter') + + # plot cold resonance curve try: - map_alpha = mpl.colormaps[cmap_alpha_name + "_lin_alpha"] - except KeyError: - map_alpha = cmap_lin_alpha(cmap_alpha_name, register=True) + resonance = gray.read_table(outputs / 'ec-resonance.70.txt') + for n in np.unique(resonance['n']): + harm = resonance['n'] == n + ax.plot(resonance['R'][harm], resonance['z'][harm], + c='xkcd:orange', alpha=n/resonance['n'].max(), + label=f"resonance $n={int(n)}$") + except FileNotFoundError: + pass + + # central ray + plot_line_color(ax, central_ray['R'], central_ray['z'], + central_ray['α'] * np.exp(-central_ray['τ']), + cmap='plasma', label='rays') + + # plot outer rays try: - map_power = mpl.colormaps[cmap_power_name + "_transp"] - except KeyError: - map_power = cmap_transp(cmap_power_name, register=True) - - if 'tau' in rays.dtype.names: - power = np.exp(-rays['tau']) - else: - power = np.ones(rays['z'].shape) - if 'alpha' in rays.dtype.names: - dP_ds = rays['alpha']*power - else: - dP_ds = np.zeros(rays['z'].shape) - - if 'R' in rays.dtype.names: - R = rays['R'] - elif np.all((label in rays.dtype.names for label in ('x', 'y'))): - R = np.hypot(rays['x'],rays['y']) - else: - raise ValueError("At least R, or x and y, fields must be " - "present for the poloidal cross-section plot.") - - if dpdsmax is None: dpdsmax=dP_ds.max() - dpdsmax = max(dpdsmax, dP_ds_threshold) - - fluxlines = ax.plot(flux['R'], flux['z'], color='gray', linestyle='-', alpha=0.5) - reslines = ax.plot(bres['R'], bres['z'], color='black', linestyle='-') - dots = ax.scatter(R[::stride], rays['z'][::stride], - c=power[::stride], s=dotsize, - marker='.', - cmap=map_power, norm=mpl.colors.Normalize(0, 1)) - blobs = ax.scatter(R, rays['z'], - c=dP_ds, s=blobsize*(dP_ds/dpdsmax)**2, edgecolors='none', - marker=marker, - cmap=map_alpha, norm=mpl.colors.Normalize(0, dpdsmax)) - - return dots, blobs, fluxlines, reslines + outer_rays = gray.read_table(outputs / 'outer-rays.33.txt') + for k in np.unique(outer_rays['k']): + ray = outer_rays[outer_rays['k'] == k] + plot_line_color(ax, ray['R'], ray['z'], ray['α']*np.exp(-ray['τ']), + cmap='plasma', alpha=0.3) + except FileNotFoundError: + pass -def plot_rays3d(ax, rays, dpdsmax=None, dotsize=0.5, blobsize=20, marker='o', - stride=1, label=None): - r""" - Plot rays path in 3D - - Expects to receive a structured numpy array with field names recognisable - as spatial coordinates ('x', 'y', 'z') or ('R', 'phi', 'z'). If absorption - data are found ('alpha', 'tau'), the rays are coloured accordingly. - - Parameters - ---------- - ax : matplotlib Axes - Axes where the beam path will be plotted. - rays : ndarray - Structured numpy array with at least ('x', 'y', 'z') or ('R', 'phi', - 'z') fields. If field names 'alpha' and 'tau' are found too, they will - be used to shade the beam path. - dpdsmax : float, default=None - Normalization value for the size of markers along the beam. The - absorption markers will have size=blobsize where - $\alpha \exp(-\tau)$ = dpdsmax. - dotsize : float, optional, default=0.5 - Size of the dots used to represent the beam path. - blobsize : float, optional, default=20 - Size of the markers used to represent the absorption region along the - beam path. - marker : optional, default='o' - Shape of the absorption markers along the beam. - stride : int, optional, default=1 - label : str, optional, default=None - Unused. May be used in future for labelling. +def plot_toroidal(inputs: Path, outputs: Path, ax: plt.Axes): + ''' + Plots raytracing projected in the poloidal plane + ''' + conf = gray.read_conf(inputs / 'gray.ini') + central_ray = gray.read_table(outputs / 'central-ray.4.txt') - Returns - ------- - dots : matplotlib artist representing the beam path with dots, shaded - according to the residual power fraction $P/P_0 = \exp(-\tau)$. - blobs : matplotlib artist representing the part of the beam where EC - absorption occurs, sized and shaded according to the local rate of - power absorption $dP/ds/P_0 = \alpha \exp(-\tau)$. - """ - # Use colormaps with added transparency - # Warm (yellow to red) cmap for absorption coefficient - # Monochrome (black to white) cmap for residual power - cmap_alpha_name = 'autumn_r' - cmap_power_name = 'Blues' - dP_ds_threshold = 1e-8 + ax.set_title('toroidal view', loc='right') + ax.set_xlabel('$x$ / m') + ax.set_ylabel('$y$ / m') + + limiter = gray.get_limiter(conf, inputs) + + # plot plasma boundary + surfaces = gray.read_table(outputs / 'flux-surfaces.71.txt') + boundary = surfaces[np.isclose(surfaces['ψ_n'], 1, 1e-3)] + + # plot plasma boundary + plot_circle(ax, radius=boundary['R'].min(), color='xkcd:slate gray') + plot_circle(ax, radius=boundary['R'].max(), color='xkcd:slate gray') + + # plot limiter + if limiter[0].size > 0: + Rmax = limiter[0].max() * 1.1 + plot_circle(ax, radius=limiter[0].min(), + color='xkcd:black') + plot_circle(ax, radius=limiter[0].max(), + color='xkcd:black') + # set bounds + ax.set_xlim(-Rmax, Rmax) + ax.set_ylim(-Rmax, Rmax) + + # plot cold resonance curve try: - map_alpha = mpl.colormaps[cmap_alpha_name + "_lin_alpha"] - except KeyError: - map_alpha = cmap_lin_alpha(cmap_alpha_name, register=True) + resonance = gray.read_table(outputs / 'ec-resonance.70.txt') + for n in np.unique(resonance['n']): + harm = resonance['n'] == n + plot_circle(ax, radius=resonance['R'][harm].max(), + color='xkcd:orange', alpha=n/resonance['n'].max()) + except FileNotFoundError: + pass + + # central ray + for rt in np.unique(central_ray['index_rt']): + ray = central_ray[central_ray['index_rt'] == rt] + x = ray['R'] * np.cos(np.radians(ray['φ'])) + y = ray['R'] * np.sin(np.radians(ray['φ'])) + dPds = ray['α'] * np.exp(-ray['τ']) + plot_line_color(ax, x, y, dPds, cmap='plasma') + + # outer rays try: - map_power = mpl.colormaps[cmap_power_name + "_transp"] - except KeyError: - map_power = cmap_transp(cmap_power_name, register=True) + outer_rays = gray.read_table(outputs / 'outer-rays.33.txt') + for k in np.unique(outer_rays['k']): + for rt in np.unique(outer_rays['index_rt']): + ray = outer_rays[ (outer_rays['k'] == k) + & (outer_rays['index_rt'] == rt)] + dPds = ray['α'] * np.exp(-ray['τ']) + plot_line_color(ax, ray['x'], ray['y'], dPds, + cmap='plasma', alpha=0.3) + except FileNotFoundError: + pass - if 'tau' in rays.dtype.names: - power = np.exp(-rays['tau']) + +def plot_beam(outputs: Path, axes: [plt.Axes]): + ''' + Plots the outermost rays in the local beam frame + ''' + rays = gray.read_table(outputs / 'beam-shape.8.txt') + nrays = int(max(rays['k'])) + + cm = 0.01 # 1cm in m + cmap = plt.cm.viridis + + # side view of rays, y(z) + axes['A'].set_title('rays, side view', loc='right') + axes['A'].set_ylabel('$y$ / cm') + + # top view of rays, x(z) + axes['B'].set_title('rays, top view', loc='right') + axes['B'].set_ylabel('$x$ / cm') + axes['B'].set_xlabel('$z$ / m') + + # 3D view + axes['C'].set_title('beam surface', loc='right') + axes['C'].set_xlabel('$z$ / m', labelpad=11) + axes['C'].set_ylabel('$x$ / cm', labelpad=-2) + axes['C'].set_zlabel('$y$ / cm', labelpad=-2) + axes['C'].set_box_aspect(aspect=(2.5, 1, 1), zoom=1.2) + axes['C'].view_init(azim=-44) + + # plot the beam envelope + try: + size = gray.read_table(outputs / 'beam-size.12.txt') + for i in ['A', 'B']: + style = dict(c='xkcd:grey', ls=':', lw=1) + axes[i].plot(size['s']*cm, +size['r_max'], **style) + axes[i].plot(size['s']*cm, -size['r_max'], **style) + axes[i].plot(size['s']*cm, +size['r_min'], **style) + axes[i].plot(size['s']*cm, -size['r_min'], **style) + except FileNotFoundError: + pass + + for k in np.arange(1, nrays+1): + ray = rays[rays['k'] == k] + color = cmap(k / max(rays['k'])) + z = (ray['s'] + ray['z']) * cm + # top projections + axes['B'].plot(z, ray['x'], alpha=0.6, c=color, lw=1) + axes['C'].plot(z, ray['x'], zs=rays['y'].min(), zdir='z', + alpha=0.6, c=color, lw=1) + # side projections + axes['A'].plot(z, ray['y'], alpha=0.8, c=color, lw=1) + axes['C'].plot(z, ray['y'], zs=rays['x'].max(), zdir='y', + alpha=0.6, c=color, lw=1) + + # adjust ticks spacing + plt.setp(axes['C'].get_yticklabels(), rotation=-25, ha='left') + axes['C'].tick_params(axis='y', pad=-5) + axes['C'].tick_params(axis='z', pad=0) + + # wrap arrays into 2D meshes for plot_surface + k = rays['k'].reshape(-1, nrays) + x = rays['x'].reshape(-1, nrays) + y = rays['y'].reshape(-1, nrays) + z = (rays['s'] + rays['z']).reshape(-1, nrays) * cm + + # make the xy slices closed + x = np.column_stack([x, x[:, 0]]) + y = np.column_stack([y, y[:, 0]]) + z = np.column_stack([z, z[:, 0]]) + + # plot beam surface, colored by ray indices + axes['C'].plot_surface(z, x, y, facecolors=cmap(k/nrays), alpha=0.6) + + +def plot_profiles(outputs: Path, axes: [plt.Axes]): + ''' + Plots the power and current profiles + ''' + profiles = gray.read_table(outputs / 'ec-profiles.48.txt') + + first = profiles['index_rt'] == profiles['index_rt'].min() + ρ_p, ρ_t = profiles[first]['ρ_p'], profiles[first]['ρ_t'] + + for _, ax in axes.items(): + # set ρ_t to the x axis + ax.set_xlabel('$ρ_t$') + ax.set_xlim(-0.05, 1.05) + ax.ticklabel_format(useMathText=True) + + # add secondary x axis for ρ_p + add_alt_axis(ax, ρ_t, ρ_p, label='$ρ_p$') + + cur_dens = axes['A'] + cur_dens.set_title('current density', loc='right') + cur_dens.set_ylabel(r'$J_\text{CD} \;/\;$ kA/m³') + + pow_dens = axes['B'] + pow_dens.set_title('power density', loc='right') + pow_dens.set_ylabel(r'$dP/dV \;/\;$ MW/m³') + + cur_ins = axes['C'] + cur_ins.set_title('current within $ρ$', loc='right') + cur_ins.set_ylabel(r'$I_\text{CD,inside} \;/\;$ kA') + + pow_ins = axes['D'] + pow_ins.set_title('power within $ρ$', loc='right') + pow_ins.set_ylabel(r'$P_\text{inside} \;/\;$ MW') + + styles = [':', '-', '--', '-.'] + colors = dict(O='xkcd:orange', X='xkcd:ocean blue') + + for i in np.unique(profiles['index_rt']): + mode, passes = gray.decode_index_rt(int(i)) + style = dict(c=colors[mode], ls=styles[passes % 4]) + + slice = profiles[ (profiles['index_rt'] == i) + & (profiles['dPdV'] > 1e-15)] + cur_dens.plot(slice['ρ_p'], slice['J_cdb'], **style) + pow_dens.plot(slice['ρ_p'], slice['dPdV'], **style) + + slice = profiles[ (profiles['index_rt'] == i) + & (profiles['P_inside'] > 1e-15)] + cur_ins.plot(slice['ρ_p'], slice['I_cd_inside'], **style) + pow_ins.plot(slice['ρ_p'], slice['P_inside'], **style) + + # legend + plt.plot(np.nan, np.nan, label='O mode', c='xkcd:orange') + plt.plot(np.nan, np.nan, label='X mode', c='xkcd:ocean blue') + plt.plot(np.nan, np.nan, label='1° pass', c='xkcd:black', ls=styles[1]) + plt.plot(np.nan, np.nan, label='2° pass', c='xkcd:black', ls=styles[2]) + plt.plot(np.nan, np.nan, label='3° pass', c='xkcd:black', ls=styles[3]) + plt.plot(np.nan, np.nan, label='4° pass', c='xkcd:black', ls=styles[0]) + + +def plot_inputs(inputs: Path, outputs: Path, axes: [plt.Axes]): + ''' + Plot the input plasma profiles and MHD equilibrium + ''' + from matplotlib.colors import TwoSlopeNorm + + maps = gray.read_table(outputs / 'inputs-maps.72.txt') + + # filter valid points + maps = maps[maps['ψ_n'] > 0] + flux = maps['R'], maps['z'], maps['ψ_n'] + + # contour levels + inner_levels = np.linspace(0, 0.99, 8) + outer_levels = np.linspace(0.9991, maps['ψ_n'].max()*0.99, 6) + all_levels = np.concatenate([inner_levels, outer_levels]) + + # contour style + norm = TwoSlopeNorm(vmin=0, vcenter=1, vmax=maps['ψ_n'].max()) + cmap = plt.cm.RdYlGn_r + borders = dict(colors='xkcd:grey', linestyles='-', linewidths=0.8) + + # interpolated equilibrium + interp = axes['B'] + interp.set_title('poloidal flux', loc='right') + interp.set_xlabel('$R$ / m') + interp.set_ylabel('$z$ / m') + interp.tricontourf(*flux, levels=all_levels, norm=norm, cmap=cmap) + interp.tricontour(*flux, levels=all_levels, **borders) + interp.tricontour(*flux, levels=[1], colors='xkcd:slate gray') + interp.plot(np.nan, np.nan, c='xkcd:slate gray', label='plasma boundary') + + # add limiter + conf = gray.read_conf(inputs / 'gray.ini') + limiter = gray.get_limiter(conf, inputs) + interp.plot(*limiter, c='xkcd:black', label='limiter') + + # original MHD equilibrium + orig = axes['A'] + if conf['equilibrium'].get('iequil') != 'EQ_ANALYTICAL': + file = conf['equilibrium']['filenm'].strip('"') + eqdsk = gray.read_eqdsk(inputs / file) + orig.set_title('input G-EQDSK flux', loc='right') + orig.set_xlabel('$R$ / m') + orig.set_ylabel('$z$ / m') + orig.contourf(*eqdsk.flux, levels=inner_levels, norm=norm, cmap=cmap) + orig.contour(*eqdsk.flux, levels=inner_levels, **borders) + orig.plot(*eqdsk.limiter, c='xkcd:black') + orig.plot(*eqdsk.boundary, c='xkcd:slate gray') else: - power = np.ones(rays['z'].shape) - if 'alpha' in rays.dtype.names: - dP_ds = rays['alpha']*power - else: - dP_ds = np.zeros(rays['z'].shape) + orig.axis('off') - if np.all((label in rays.dtype.names for label in ('x', 'y'))): - x = rays['x'] - y = rays['y'] - elif np.all((label in rays.dtype.names for label in ('R', 'phi'))): - x = rays['R']*np.cos(rays['phi']) - y = rays['R']*np.sin(rays['phi']) - else: - raise ValueError("At least (x, y, z), or (R, phi, z), fields must be " - "present for the 3D plot.") + # colorbar + bar = plt.colorbar(plt.cm.ScalarMappable(norm, cmap), cax=axes['C']) + bar.set_label(label='normalised poloidal flux', labelpad=10) - if dpdsmax is None: dpdsmax=dP_ds.max() - dpdsmax = max(dpdsmax, dP_ds_threshold) + # Plasma radial profiles + profiles = gray.read_table(outputs / 'kinetic-profiles.55.txt') - dots = ax.scatter(x[::stride], y[::stride], rays['z'][::stride], - c=power[::stride], s=dotsize, - marker='.', - cmap=map_power, norm=mpl.colors.Normalize(0, 1)) - blobs = ax.scatter(x, y, rays['z'], - c=dP_ds, s=blobsize*(dP_ds/dpdsmax)**2, edgecolors='none', - marker=marker, - cmap=map_alpha, norm=mpl.colors.Normalize(0, dpdsmax)) + axes['D'].set_title('plasma profiles', loc='right') + add_alt_axis(axes['D'], profiles['ρ_t'], profiles['ψ_n']**0.5, + label='$ρ_p$') - return dots, blobs + # density + dens = axes['D'] + dens.set_xlabel('$ρ_t$') + dens.set_ylabel('$n_e$ / 10²⁰m⁻³', color='xkcd:ocean blue') + dens.plot(profiles['ρ_t'], profiles['n_e'], c='xkcd:ocean blue') + + # temperature + temp = axes['D'].twinx() + temp.set_ylabel('$T_e$ / keV', color='xkcd:orange') + temp.plot(profiles['ρ_t'], profiles['T_e'], c='xkcd:orange') + + # safety factor + inside = profiles['ρ_t'] < 1 + saft = axes['D'].twinx() + saft.set_ylabel('$q$', color='xkcd:moss') + saft.plot(profiles['ρ_t'][inside], profiles['q'][inside], c='xkcd:moss') + saft.spines["right"].set_position(("axes", 1.11)) + + align_yaxis(dens, temp) + align_yaxis(temp, saft) + + # original data points + if conf['profiles']['iprof'] == 'PROF_NUMERIC': + # load input profiles + orig = np.loadtxt(inputs / conf['profiles']['filenm'].strip('"'), + skiprows=1) + + # covert 1st column to ρ_t + var = conf['profiles']['irho'] + if var == 'RHO_TOR': + prof_ρ = orig[:, 0] + elif var == 'RHO_POL': + prof_ρ = np.interp(orig[:, 0], profiles['ψ_n']**0.5, + profiles['ρ_t']) + elif var == 'RHO_PSI': + prof_ρ = np.interp(orig[:, 0], profiles['ψ_n'], profiles['ρ_t']) + + saft_ρ = np.interp(eqdsk.q[0], profiles['ψ_n']**0.5, profiles['ρ_t']) + + # add to existing plot + style = dict(marker='o', s=2) + temp.scatter(prof_ρ, orig[:, 1], c='xkcd:orange', **style) + dens.scatter(prof_ρ, orig[:, 2], c='xkcd:ocean blue', **style) + saft.scatter(saft_ρ, eqdsk.q[1], c='xkcd:moss', **style) + + # legend + axes['D'].plot(np.nan, np.nan, label='spline', c='k') + axes['D'].scatter(np.nan, np.nan, label='data', c='k', **style) + axes['D'].legend(loc='center left') -if __name__ == "__main__": - import sys +def main(kind: str, args: argparse.Namespace) -> (plt.Figure, [float]): + ''' + Draws a single figure of a given kind + ''' + cm = 0.3937 # 1cm in inches + rect = [0, 0, 1, 1] # default layout rect - runs = [] - for dirname in sys.argv[1:]: - run = read_grayout(dirname) - if run: - runs.append(run) - print(f"{run['name']} loaded.") - else: - print(f"No valid data in {dirname}.") - if not runs: exit() + if kind == 'raytracing': + fig, axes = plt.subplots( + 1, 2, figsize=(19.8*cm, 11*cm), + tight_layout=True, + subplot_kw=dict(aspect='equal')) + plot_poloidal(args.inputs, args.outputs, axes[0]) + plot_toroidal(args.inputs, args.outputs, axes[1]) - # Figure for profiles vs rho - # The list axsrho has shape (2,2) - # The axes axsrho[0][:] and axsrho[1][:] are vertically stacked and share - # the same x axis - # The axes axsrho[i][0] and axsrho[i][1] are overlapped with two - # independent y axis - figrho, axsrho = plt.subplots(2,1, sharex=True) - axsrho[1].set_xlabel(r"$\rho_t$") - axsrho = [[axrho, axrho.twinx()] for axrho in axsrho] - axsrho[0][0].set_ylabel(r"$dP/dV/P_0$ [kW/m$^{-3}$/MW]") - axsrho[0][1].set_ylabel(r"$P_{ins}/P_0$") - axsrho[1][0].set_ylabel(r"$J_{CD}/P_0$ [kA/m$^2$/MW]") - axsrho[1][1].set_ylabel(r"$I_{ins}/P_0$ [kA/MW]") + if args.legend: + rect = [0.2, 0, 1, 1] + fig.tight_layout(rect=rect) + fig.legend(loc='upper left') - plotsrho = [] - for i, run in enumerate(runs): - plotrho = plot_ECprofiles(axsrho, run['prof'], label=run['name']) - plotrho = { - 'name': run['name'], - 'dPdV': plotrho[0], - 'P': plotrho[1], - 'J': plotrho[2], - 'I': plotrho[3]} - plotsrho.append(plotrho) - figrho.legend(loc='upper center', ncols=3) + elif kind == 'beam': + fig, axes = plt.subplot_mosaic( + 'ACD;BCD', figsize=(24*cm, 10*cm), + per_subplot_kw={'C': dict(projection='3d')}, + width_ratios=[1, 1, 0.21]) - # Figure for beam path in poloidal cross-section - figpol, axpol = plt.subplots() - axpol.set_xlabel(r"$R$ [m]") - axpol.set_ylabel(r"$z$ [m]") - axpol.set_aspect('equal') + # manually adjust spacing (tight_layout is failing) + axes['D'].axis('off') + plt.subplots_adjust(hspace=0.5) + plot_beam(args.outputs, axes) - plotspol = [] - for i, run in enumerate(runs): - plotpol = plot_polview(axpol, run['rays'], run['flux'], run['bres'], - stride=5) - plotpol = { - 'name': run['name'], - 'pow': plotpol[0], - 'abs': plotpol[1], - 'flux': plotpol[2], - 'bres': plotpol[3], - } - plotspol.append(plotpol) + elif kind == 'profiles': + fig, axes = plt.subplot_mosaic( + 'AB;CD', figsize=(23.11*cm, 13*cm), + tight_layout=True) + plot_profiles(args.outputs, axes) - # Figure for 3D ray path - fig3d, ax3d = plt.subplots(subplot_kw=dict(projection='3d')) - xlabel3d = ax3d.set_xlabel('x [m]') - ylabel3d = ax3d.set_ylabel('y [m]') - zlabel3d = ax3d.set_zlabel('z [m]') - ax3d.set_box_aspect((1, 1, 1)) # must be set BEFORE calling .set_aspect() - limits = dict() - for label in ('x','y','z'): - limits[label] = ( - min([run['rays'][label].min() for run in runs]), - max([run['rays'][label].max() for run in runs])) -# ax.set( -# xlim=limits['x'], -# ylim=limits['y'], -# zlim=limits['z']) + if args.legend: + rect = [0.12, 0, 1, 1] + fig.tight_layout(rect=rect) + fig.legend(loc='upper left') - dP_ds_max = max([(run['rays']['alpha']*np.exp(-run['rays']['tau'])).max() - for run in runs]) + elif kind == 'inputs': + fig, axes = plt.subplot_mosaic( + 'ABC;DDD', figsize=(19.8*cm, 20*cm), + tight_layout=True, + height_ratios=[1.5, 1], + width_ratios=[1, 1, 0.05], + per_subplot_kw={'AB': dict(aspect='equal')}) + plot_inputs(args.inputs, args.outputs, axes) - scatters3d = [] - for run in runs: - scatter3d = plot_rays3d( - ax3d, run['rays'], dP_ds_max, blobsize=2, marker='o', stride=5) - scatter3d = { - 'name': run['name'], - 'pow': scatter3d[0], - 'abs': scatter3d[1]} - scatters3d.append(scatter3d) - ax3d.set_aspect('equal') + return fig, rect - colorbar_alpha = fig3d.colorbar(scatters3d[-1]['abs'], ax=ax3d, - label=r"$(dP/ds)/P_0$ $[\mathrm{m}^{-1}]$", extend='max', location='left') - colorbar_power = fig3d.colorbar(scatters3d[-1]['pow'], ax=ax3d, - label=r"$P/P_0$", extend='neither') - plt.show() +if __name__ == '__main__': + args = cli_args() + + def on_resize(event): + ''' + Update layout on window resizes + ''' + fig.tight_layout(rect=rect) + fig.canvas.draw() + + if args.interactive: + plt.ion() + + for kind in args.kind: + fig, rect = main(kind, args) + if args.interactive: + fig.canvas.mpl_connect('resize_event', on_resize) + if args.outfile is not None: + if len(args.kind) > 1: + # append plot kind to the output name + fname = args.outfile.with_stem(args.outfile.stem + '-' + kind) + else: + fname = args.outfile + fig.savefig(fname, bbox_inches='tight') + + if args.interactive: + input('press enter to quit')