566 lines
20 KiB
Python
566 lines
20 KiB
Python
'''
|
||
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
|
||
import numpy as np
|
||
import scripts.gray as gray
|
||
|
||
|
||
def cli_args() -> 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 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])
|
||
|
||
# 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 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
|
||
|
||
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 plot_circle(ax: plt.Axes, radius: float, **args) -> 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_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)
|
||
|
||
lines = LineCollection(segments, **args)
|
||
lines.set_array(c)
|
||
ax.add_collection(lines)
|
||
ax.autoscale_view()
|
||
return lines
|
||
|
||
|
||
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')
|
||
|
||
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:
|
||
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:
|
||
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_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')
|
||
|
||
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:
|
||
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:
|
||
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
|
||
|
||
|
||
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_cd'], **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:
|
||
orig.axis('off')
|
||
|
||
# colorbar
|
||
bar = plt.colorbar(plt.cm.ScalarMappable(norm, cmap), cax=axes['C'])
|
||
bar.set_label(label='normalised poloidal flux', labelpad=10)
|
||
|
||
# Plasma radial profiles
|
||
profiles = gray.read_table(outputs / 'kinetic-profiles.55.txt')
|
||
|
||
axes['D'].set_title('plasma profiles', loc='right')
|
||
add_alt_axis(axes['D'], profiles['ρ_t'], profiles['ψ_n']**0.5,
|
||
label='$ρ_p$')
|
||
|
||
# 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')
|
||
|
||
|
||
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
|
||
|
||
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])
|
||
|
||
if args.legend:
|
||
rect = [0.2, 0, 1, 1]
|
||
fig.tight_layout(rect=rect)
|
||
fig.legend(loc='upper left')
|
||
|
||
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])
|
||
|
||
# manually adjust spacing (tight_layout is failing)
|
||
axes['D'].axis('off')
|
||
plt.subplots_adjust(hspace=0.5)
|
||
plot_beam(args.outputs, axes)
|
||
|
||
elif kind == 'profiles':
|
||
fig, axes = plt.subplot_mosaic(
|
||
'AB;CD', figsize=(23.11*cm, 13*cm),
|
||
tight_layout=True)
|
||
plot_profiles(args.outputs, axes)
|
||
|
||
if args.legend:
|
||
rect = [0.12, 0, 1, 1]
|
||
fig.tight_layout(rect=rect)
|
||
fig.legend(loc='upper left')
|
||
|
||
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)
|
||
|
||
return fig, rect
|
||
|
||
|
||
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')
|