istp/gray
2
0
Fork 0
Code Issues Pull Requests Packages Projects Releases 1 Wiki Activity

Add first draft of Python plotting script

Browse Source
This commit is contained in:
Lorenzo Figini 2022-11-13 01:55:25 +01:00
parent 6181b6096e
commit 91fa6d84e0
1 changed files with 544 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

544
scripts/gray_visual.py Normal file
View File

@ -0,0 +1,544 @@
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': 4,
'totals': 7,
'rays': 33,
'prof': 48,
'bres': 70,
'flux': 71,
}
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"/fort.{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()