545 lines
19 KiB
Python
545 lines
19 KiB
Python
import numpy as np
|
|
import matplotlib as mpl
|
|
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
|
|
|
|
|
|
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<min_loc<1, colors are fully opaque at both ends and transparency
|
|
increases linearly towards min_loc.
|
|
reverse_alpha : bool, default=False
|
|
Whether alpha channel is reversed, so that minimum transparency
|
|
occurs at min_loc and maximum transparency at range ends.
|
|
register : bool, default=True
|
|
Whether to register the new colormap.
|
|
|
|
Returns
|
|
-------
|
|
cmap : matplotlib.colors.ListedColormap
|
|
Colormap with linear alpha set.
|
|
"""
|
|
# 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
|
|
# Map the max transparency location from (0, 1) to index in LUT
|
|
index_alpha_min = int(np.rint(np.clip(min_loc*cmap.N, 0, cmap.N)))
|
|
# Build the linear alpha channel with minimum at index_alpha_min
|
|
alpha = np.concatenate(
|
|
(np.linspace(1, 0, index_alpha_min, endpoint=False),
|
|
np.linspace(0, 1, cmap.N - index_alpha_min)))
|
|
# Replace the old alpha channel
|
|
rgba[:,-1] = alpha
|
|
# Create a new Colormap object
|
|
cmap_alpha = mpl.colors.ListedColormap(rgba, name=name + '_lin_alpha')
|
|
if register:
|
|
mpl.cm.register_cmap(name=name + '_lin_alpha', cmap=cmap_alpha)
|
|
return cmap_alpha
|
|
|
|
|
|
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.
|
|
|
|
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
|
|
|
|
|
|
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
|
|
|
|
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
|
|
|
|
|
|
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<ncol; i++) printf " -"; printf "\n"}}'],
|
|
stdout=subprocess.PIPE, text=True,
|
|
input="".join(f.readlines()))
|
|
file = io.StringIO(processed.stdout)
|
|
except OSError:
|
|
pass
|
|
skip=0
|
|
if label == 'flux': names=['i', 'psin', 'R', 'z']
|
|
else:
|
|
skip=HEADER_GRAY
|
|
names=True
|
|
try:
|
|
run[label] = np.genfromtxt(file, skip_header=skip,
|
|
names=names, missing_values=['-'], filling_values=np.nan)
|
|
except OSError:
|
|
print(f"Cannot open {file}. Skipping...")
|
|
except ValueError:
|
|
print(f"Bad format in {file}.")
|
|
raise
|
|
return run
|
|
|
|
|
|
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
|
|
|
|
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
|
|
|
|
|
|
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.
|
|
|
|
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
|
|
try:
|
|
map_alpha = mpl.colormaps[cmap_alpha_name + "_lin_alpha"]
|
|
except KeyError:
|
|
map_alpha = cmap_lin_alpha(cmap_alpha_name, register=True)
|
|
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
|
|
|
|
|
|
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.
|
|
|
|
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
|
|
try:
|
|
map_alpha = mpl.colormaps[cmap_alpha_name + "_lin_alpha"]
|
|
except KeyError:
|
|
map_alpha = cmap_lin_alpha(cmap_alpha_name, register=True)
|
|
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 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.")
|
|
|
|
if dpdsmax is None: dpdsmax=dP_ds.max()
|
|
dpdsmax = max(dpdsmax, dP_ds_threshold)
|
|
|
|
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))
|
|
|
|
return dots, blobs
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
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()
|
|
|
|
# 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]")
|
|
|
|
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)
|
|
|
|
# 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')
|
|
|
|
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)
|
|
|
|
# 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'])
|
|
|
|
dP_ds_max = max([(run['rays']['alpha']*np.exp(-run['rays']['tau'])).max()
|
|
for run in runs])
|
|
|
|
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')
|
|
|
|
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()
|