Toroidal shear-Alfvén waves

"""A small toroidal LinearMHD run with the m=10,11 shear-Alfvén perturbation.

Install dependencies and compile the kernels (Python 3.10 or newer):

    pip install "struphy[pproc]>=3.4.0"
    pip install "plasma-plots[plotly]>=0.1.1"
    struphy compile

PNG exports require Chrome or Chromium. If Chrome is not installed, run:

    kaleido_get_chrome

Save this file and run it from the directory where you want the output:

    python toroidal-shear-alfven.py

Figures are saved in the current directory; add --show to display them before saving.
The default is an exploratory local run, not a converged ITPA TAE benchmark.
Increase NUM_ELEMENTS to (24, 96, 16), DEGREE to (3, 3, 3), and END_TIME
to 500.0 to recover the supplied spatial/time resolution. The equilibrium,
sector, mode numbers, Gaussian profiles and amplitudes are retained.
"""

import argparse

import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from scipy.ndimage import map_coordinates

from struphy import (
    BaseUnits,
    DerhamOptions,
    EnvironmentOptions,
    Simulation,
    Time,
    domains,
    equils,
    grids,
    perturbations,
)
from struphy.models import LinearMHD
from plasma_plots import save_figure

STEM = "toroidal-shear-alfven"
NUM_ELEMENTS = (8, 48, 4)
DEGREE = (3, 3, 2)
END_TIME = 40.0
DT = 0.1
SAVE_STEP = 2

# These are logical H(div) / 2-form components, not physical unit-vector
# components. TorusModes uses phase 2*pi*(m*eta2 + n*eta3); n=-1 is one
# oscillation across the 1/6-torus sector, not a full-torus n=-1 mode.
modes = (10, 11)
amplitude = 1e-3


def physical_radial_component(output, field_xyz):
    """Project Cartesian components onto the local minor-radial unit vector."""
    params = output.domain.params
    theta = 2 * np.pi * field_xyz.eta2
    phi = -2 * np.pi * field_xyz.eta3 / params["tor_period"]
    fx, fy, fz = (field_xyz.isel(component=c, drop=True) for c in range(3))
    return (fx * np.cos(phi) + fy * np.sin(phi)) * np.cos(theta) + fz * np.sin(theta)


def radial_mode_amplitudes(output, field_xyz, poloidal_modes=(9, 10, 11, 12), sector_mode=-1):
    """Normalized radial profiles of physical (m, n_sector) Fourier amplitudes.

    ``field_xyz`` is one saved Cartesian vector-field snapshot over the whole
    toroidal sector. Normalize all requested harmonics together, preserving their
    relative strengths. This is a spatial FFT, not a temporal band selection.
    """
    params = output.domain.params
    radial = physical_radial_component(output, field_xyz)
    # Convert to the rotating radial basis BEFORE transforming in toroidal angle.
    # Cartesian components themselves are not periodic across the sector seam.
    # mode_spectrum drops the duplicate periodic endpoints and labels integer (m, n_sector).
    coefficients = radial.plasma.analysis.mode_spectrum(dims=("eta2", "eta3"), names=("m", "n"))
    if not set(poloidal_modes) <= set(coefficients.m.values) or sector_mode not in coefficients.n.values:
        raise ValueError("Angular sampling does not resolve the requested Fourier modes.")
    selected = coefficients.sel(m=list(poloidal_modes), n=sector_mode)
    amplitude = 2 * abs(selected)  # Conjugate-pair amplitude of a real spatial harmonic.
    if not np.isfinite(amplitude.values).all():
        raise RuntimeError("Non-finite radial Fourier amplitude.")
    peak = float(amplitude.max())
    normalized = amplitude / (peak if peak > 0 else 1.0)
    normalized = normalized.assign_coords(
        radius=params["a1"] + (params["a2"] - params["a1"]) * normalized.eta1,
    ).rename("normalized_fft_amplitude")
    normalized.attrs.update(normalization_amplitude=peak, sector_mode=sector_mode,
                            full_torus_mode=abs(sector_mode * params["tor_period"]),
                            snapshot_time=float(field_xyz.t))
    return normalized


def fixed_theta_amplitudes(output, field_xyz, angles=(0.0, 45.0), sector_mode=-1):
    """Toroidal Fourier amplitudes at fixed poloidal angles (in degrees).

    Retain the coherent sum of all poloidal harmonics. One normalization over
    both angles and all radii preserves the difference between the two rays.
    """
    radial = physical_radial_component(output, field_xyz)
    # Interpolate the physical radial field only if an angle is off the display
    # grid. The default 0 and 45 degree rays lie exactly on that grid.
    rays = radial.interp(eta2=np.asarray(angles) / 360.0)
    coefficients = rays.plasma.analysis.mode_spectrum(dims="eta3", names="n")
    if sector_mode not in coefficients.n.values:
        raise ValueError("Toroidal sampling does not resolve the requested Fourier mode.")
    selected = coefficients.sel(n=sector_mode)
    amplitude = 2 * abs(selected)
    if not np.isfinite(amplitude.values).all():
        raise RuntimeError("Non-finite fixed-angle Fourier amplitude.")
    peak = float(amplitude.max())
    normalized = (amplitude / (peak if peak > 0 else 1.0)).rename({"eta2": "theta_degrees"})
    params = output.domain.params
    normalized = normalized.assign_coords(
        theta_degrees=list(angles),
        radius=params["a1"] + (params["a2"] - params["a1"]) * normalized.eta1,
    ).rename("normalized_fft_amplitude")
    normalized.attrs.update(normalization_amplitude=peak, snapshot_time=float(field_xyz.t),
                            full_torus_mode=abs(sector_mode * params["tor_period"]))
    return normalized


def save_fixed_theta_fft_figures(output, show=False):
    """Save (and show) the two angle comparisons; usable directly with saved output."""
    for field_name, label, key in (
        ("mhd/velocity_xyz", "u_r", "fixed-theta-fft-velocity"),
        ("em_fields/b_field_xyz", "δB_r", "fixed-theta-fft-magnetic"),
    ):
        profiles = fixed_theta_amplitudes(output, output.evaluate(field_name).isel(t=-1))
        time = profiles.attrs["snapshot_time"]
        toroidal_mode = profiles.attrs["full_torus_mode"]
        figure = go.Figure()
        for angle, color, dash in ((0.0, "#0072B2", "solid"), (45.0, "#D55E00", "dash")):
            figure.add_scatter(
                x=profiles.radius.values, y=profiles.sel(theta_degrees=angle).values,
                mode="lines+markers", name=f"θ={angle:g}°", marker={"size": 4},
                line={"color": color, "dash": dash, "width": 2.5},
                hovertemplate=f"θ={angle:g}°<br>r=%{{x:.3f}}<br>Normalized amplitude=%{{y:.4f}}<extra></extra>",
            )
        figure.update_layout(
            title=f"Fixed-angle radial FFT amplitude · {label} · |n|={toroidal_mode:g}, t={time:g}",
            xaxis_title="Minor radius r", yaxis_title="Normalized FFT amplitude",
            xaxis={"range": [float(profiles.radius[0]), float(profiles.radius[-1])]},
            yaxis={"range": [0, 1.05]}, template="plotly_white",
            margin={"l": 85, "r": 35, "t": 95, "b": 75},
            legend={"title": {"text": "Poloidal angle"}},
        )
        save_figure(figure, f"{STEM}-{key}", show=show)


def create_simulation() -> Simulation:
    model = LinearMHD(base_units=BaseUnits())
    model.propagators.shear_alf.options = model.propagators.shear_alf.Options()
    model.propagators.mag_sonic.options = model.propagators.mag_sonic.Options()
    for field in (model.em_fields.b_field, model.mhd.density, model.mhd.velocity, model.mhd.pressure):
        field.save_data = True

    domain = domains.HollowTorus(a1=0.1, a2=1.0, R0=10.0, sfl=False, pol_period=1, tor_period=6)
    equil = equils.AdhocTorus(
        a=1.0, R0=10.0, B0=3.0, q_kind=0, p_kind=1,
        q0=1.71, q1=1.87, p1=0.95, p2=0.05, beta=0.0018,
    )
    grid = grids.TensorProductGrid(num_elements=NUM_ELEMENTS)
    derham_opts = DerhamOptions(degree=DEGREE, bcs=(("dirichlet", "dirichlet"), None, None))
    model.mhd.velocity.add_perturbation(perturbations.TorusModesSin(
        ms=modes, ns=(-1, -1), amps=(amplitude, amplitude),
        pfuns=("exp", "exp"), pfun_params=([0.5, 0.1], [0.5, 0.1]),
        comp=0, given_in_basis="2",
    ))
    model.mhd.velocity.add_perturbation(perturbations.TorusModesCos(
        ms=modes, ns=(-1, -1), amps=tuple(amplitude / (2 * np.pi * m) for m in modes),
        pfuns=("d_exp", "d_exp"), pfun_params=([0.5, 0.1], [0.5, 0.1]),
        comp=1, given_in_basis="2",
    ))

    env = EnvironmentOptions(
        out_folders="struphy_gallery_runs", sim_folder="toroidal_shear_alfven",
        save_step=SAVE_STEP,
    )
    time_opts = Time(dt=DT, Tend=END_TIME)
    sim = Simulation(
        model=model,
        name="Toroidal shear-Alfvén waves",
        description=(
            "A local preview of coupled poloidal shear-Alfvén perturbations in a circular tokamak, "
            "evolved with LinearMHD, including its magnetosonic propagator. "
            r"The hollow torus has :math:`0.1\le r\le1`, major radius :math:`R_0=10`, "
            r"and a periodic one-sixth toroidal sector. The AdhocTorus equilibrium uses "
            r":math:`B_0=3`, :math:`q_0=1.71`, :math:`q_1=1.87` and :math:`\beta=0.0018`. "
            r"The initial logical velocity combines :math:`m=10,11`, :math:`n=-1` sector modes "
            r"with amplitude parameter :math:`10^{-3}` and Gaussian profiles centered at "
            r":math:`\eta_1=0.5` with width :math:`0.1`; the poloidal component uses their radial derivatives. "
            "The coarse grid and short duration are intended for exploration, not a converged TAE benchmark. "
            "Plots show physical minor-radial, poloidal and toroidal velocity on the φ=0 slice; "
            "each component keeps its own fixed color range throughout the animation."
        ),
        params_path=__file__, env=env, time_opts=time_opts,
        domain=domain, equil=equil, grid=grid, derham_opts=derham_opts,
    )
    return sim


def pproc(sim: Simulation, show: bool = False):
    """Save (and show) the figures from an existing simulation."""
    domain = sim.domain
    grid = sim.grid

    output = sim.output
    # More samples improve the display of the spline, not the simulation resolution.
    output.pproc(physical=True, celldivide=(3, 3, 1), create_vtk=False)
    velocity = output.evaluate("mhd/velocity_xyz").isel(eta3=0).transpose("t", "component", "eta1", "eta2")
    times = velocity.t.values
    if not np.isclose(times[-1], END_TIME):
        raise RuntimeError(f"Run stopped at t={times[-1]}, before requested t={END_TIME}.")
    if not np.isfinite(velocity.values).all():
        raise RuntimeError("Non-finite velocity: refusing to export the example.")

    # At phi=0: e_R=e_x, e_phi=e_y, e_Z=e_z. Rotate the Cartesian
    # push-forward into orthonormal minor-radial and poloidal directions.
    theta = 2 * np.pi * velocity.eta2.values
    ux, uy, uz = (velocity.isel(component=c).values for c in range(3))
    components = np.stack((
        ux * np.cos(theta) + uz * np.sin(theta),
        -ux * np.sin(theta) + uz * np.cos(theta),
        uy,
    ), axis=1)
    labels = ("u_r", "u_θ", "u_φ")
    titles = ("Minor-radial velocity", "Poloidal velocity", "Toroidal velocity")
    limits = np.maximum(np.abs(components).max(axis=(0, 2, 3)), 1e-16)

    # Resample the evaluated annulus onto pixels in the physical (R,Z) plane.
    # Use the actual evaluation coordinates; close the periodic angular seam.
    axis = np.linspace(-domain.params["a2"], domain.params["a2"], 112)
    pixel_r, pixel_z = np.meshgrid(axis, axis)
    radius = np.hypot(pixel_r, pixel_z)
    angle = np.mod(np.arctan2(pixel_z, pixel_r), 2 * np.pi) / (2 * np.pi)
    outside = (radius < domain.params["a1"]) | (radius > domain.params["a2"])
    eta1 = (radius - domain.params["a1"]) / (domain.params["a2"] - domain.params["a1"])
    radial_index = np.interp(eta1, velocity.eta1.values, np.arange(velocity.sizes["eta1"]))
    periodic = velocity.eta2.values < 1.0 - 1e-12
    angular_grid = np.append(velocity.eta2.values[periodic], 1.0)
    angular_index = np.interp(angle, angular_grid, np.arange(len(angular_grid)))

    def traces(index):
        result = []
        for component, label in enumerate(labels):
            values = components[index, component][:, periodic]
            values = np.concatenate((values, values[:, :1]), axis=1)
            pixels = map_coordinates(values, [radial_index, angular_index], order=1, mode="nearest")
            pixels[outside] = np.nan
            result.append(go.Heatmap(
                x=axis + domain.params["R0"], y=axis, z=pixels.astype(np.float32),
                zmin=-limits[component], zmax=limits[component], colorscale="RdBu",
                colorbar={
                    "title": label, "x": (0.28, 0.64, 1.0)[component], "len": 0.55,
                    "thickness": 10, "tickformat": ".1e", "tickfont": {"size": 10},
                },
                hovertemplate=f"R=%{{x:.3f}}<br>Z=%{{y:.3f}}<br>{label}=%{{z:.3e}}<extra></extra>",
                xaxis=f"x{component + 1}" if component else "x",
                yaxis=f"y{component + 1}" if component else "y",
            ))
        return result

    figure = make_subplots(rows=1, cols=3, subplot_titles=titles, horizontal_spacing=0.12)
    for column, trace in enumerate(traces(0), 1):
        figure.add_trace(trace, row=1, col=column)
    frame_indices = np.unique(np.linspace(0, len(times) - 1, min(26, len(times)), dtype=int))
    figure.frames = [go.Frame(name=f"{times[i]:g}", data=traces(i), traces=[0, 1, 2]) for i in frame_indices]
    figure.update_xaxes(
        title_text="R", range=[domain.params["R0"] - domain.params["a2"], domain.params["R0"] + domain.params["a2"]],
        constrain="domain", showgrid=False, zeroline=False,
    )
    for column in range(1, 4):
        figure.update_yaxes(
            title_text="Z" if column == 1 else None, range=[-domain.params["a2"], domain.params["a2"]],
            constrain="domain", showticklabels=column == 1, showgrid=False, zeroline=False,
            scaleanchor="x" if column == 1 else f"x{column}", scaleratio=1, row=1, col=column,
        )
    figure.update_layout(
        title="Toroidal shear-Alfvén waves · physical velocity at φ = 0",
        template="plotly_white", autosize=True, margin={"l": 55, "r": 95, "t": 95, "b": 130},
        sliders=[{
            "active": 0, "x": 0.12, "len": 0.86, "y": -0.12,
            "currentvalue": {"prefix": "t = "},
            "steps": [{"label": f.name, "method": "animate", "args": [[f.name], {
                "mode": "immediate", "frame": {"duration": 0, "redraw": True}, "transition": {"duration": 0},
            }]} for f in figure.frames],
        }],
        updatemenus=[{"type": "buttons", "direction": "left", "x": 0, "y": -0.12, "buttons": [
            {"label": "Play", "method": "animate", "args": [None, {
                "fromcurrent": True, "frame": {"duration": 100, "redraw": True}, "transition": {"duration": 0},
            }]},
            {"label": "Pause", "method": "animate", "args": [[None], {
                "mode": "immediate", "frame": {"duration": 0, "redraw": False},
            }]},
        ]}],
    )
    # A later snapshot also shows the generated toroidal component in the thumbnail.
    still = len(frame_indices) // 2
    save_figure(figure, STEM, frame=still, show=show)

    radial_probe = int(np.abs(velocity.eta1.values - 0.5).argmin())
    probe_radius = domain.params["a1"] + (domain.params["a2"] - domain.params["a1"]) * float(velocity.eta1.values[radial_probe])
    history = make_subplots(rows=1, cols=3, subplot_titles=labels, horizontal_spacing=0.12)
    for component in range(3):
        history.add_trace(go.Heatmap(
            x=theta / (2 * np.pi), y=times, z=components[:, component, radial_probe],
            zmin=-limits[component], zmax=limits[component], colorscale="RdBu",
            colorbar={
                "title": labels[component], "x": (0.28, 0.64, 1.0)[component], "len": 0.8,
                "thickness": 10, "tickformat": ".1e", "tickfont": {"size": 10},
            },
        ), row=1, col=component + 1)
        history.update_yaxes(showticklabels=component == 0, range=[times[0], times[-1]], row=1, col=component + 1)
    history.update_xaxes(title_text="θ / 2π")
    history.update_yaxes(title_text="t", row=1, col=1)
    history.update_layout(
        title=f"Velocity around the ring r = {probe_radius:.3f}, φ = 0",
        template="plotly_white", margin={"l": 60, "r": 95, "t": 90, "b": 70},
    )

    # Compare radial rays at theta=0 and theta=45 degrees on the phi=0 plane.
    radii = domain.params["a1"] + (domain.params["a2"] - domain.params["a1"]) * velocity.eta1.values
    angle_probes = (
        (int(np.abs(theta - 0.0).argmin()), 0.0, "solid"),
        (int(np.abs(theta - np.pi / 4).argmin()), 45.0, "dash"),
    )
    initial_radius = domain.params["a1"] + 0.5 * (domain.params["a2"] - domain.params["a1"])
    snapshot_indices = np.unique(np.linspace(0, len(times) - 1, min(5, len(times)), dtype=int))
    snapshot_colors = ("#0072B2", "#E69F00", "#009E73", "#CC79A7", "#D55E00")
    radial_profile_figures = []
    for angle_probe, angle_degrees, angle_dash in angle_probes:
        radial_profiles = make_subplots(rows=1, cols=3, subplot_titles=titles, horizontal_spacing=0.12)
        for component, label in enumerate(labels):
            for snapshot, color in zip(snapshot_indices, snapshot_colors):
                radial_profiles.add_scatter(
                    x=radii, y=components[snapshot, component, :, angle_probe],
                    mode="lines+markers", name=f"t = {times[snapshot]:g}",
                    legendgroup=str(snapshot), showlegend=component == 0,
                    line={"color": color, "width": 2, "dash": angle_dash},
                    marker={"size": 4},
                    hovertemplate=(
                        f"r=%{{x:.3f}}<br>{label}=%{{y:.3e}}"
                        f"<extra>t={times[snapshot]:g}</extra>"
                    ),
                    row=1, col=component + 1,
                )
            radial_profiles.update_yaxes(
                title_text=label, exponentformat="power", tickfont={"size": 10}, zeroline=True,
                row=1, col=component + 1,
            )
        radial_profiles.add_vline(x=initial_radius, line_dash="dot", line_color="#9ca3af", line_width=1)
        radial_profiles.update_xaxes(title_text="Minor radius r", range=[radii[0], radii[-1]])
        radial_profiles.update_layout(
            title=f"Radial velocity profiles · θ = {angle_degrees:g}°, φ = 0",
            template="plotly_white", margin={"l": 85, "r": 35, "t": 110, "b": 105},
            legend={"orientation": "h", "x": 0.5, "xanchor": "center", "y": -0.18},
        )
        radial_profile_figures.append((angle_degrees, radial_profiles))

    # Angular RMS avoids cancellation between opposite signs of a wave.
    # Samples are uniform in theta; exclude the duplicate periodic endpoint.
    # This is an angular average on the phi=0 plane, not a volume average.
    angular_rms = np.sqrt(np.mean(components[..., periodic] ** 2, axis=-1))
    radial_history = make_subplots(rows=1, cols=3, subplot_titles=labels, horizontal_spacing=0.12)
    for component, label in enumerate(labels):
        radial_history.add_trace(go.Heatmap(
            x=radii, y=times, z=angular_rms[:, component],
            zmin=0, zmax=max(float(angular_rms[:, component].max()), 1e-16), colorscale="Viridis",
            colorbar={
                "title": f"RMS {label}", "x": (0.28, 0.64, 1.0)[component], "len": 0.8,
                "thickness": 10, "tickformat": ".1e", "tickfont": {"size": 10},
            },
            hovertemplate=f"r=%{{x:.3f}}<br>t=%{{y:g}}<br>RMS {label}=%{{z:.3e}}<extra></extra>",
        ), row=1, col=component + 1)
        radial_history.update_yaxes(
            showticklabels=component == 0, range=[times[0], times[-1]], row=1, col=component + 1,
        )
    radial_history.add_vline(x=initial_radius, line_dash="dot", line_color="white", line_width=1)
    radial_history.update_xaxes(title_text="Minor radius r", range=[radii[0], radii[-1]])
    radial_history.update_yaxes(title_text="t", row=1, col=1)
    radial_history.update_layout(
        title="Radial evolution · poloidal RMS velocity at φ = 0",
        template="plotly_white", margin={"l": 60, "r": 100, "t": 90, "b": 70},
    )

    # Transform signed physical components, before taking RMS or squaring:
    # transforming a velocity magnitude/energy would change its frequencies.
    # Omit the duplicated poloidal endpoint from spatial sums.
    plane = velocity.copy(data=components).assign_coords(component=list(labels))
    plane = plane.plasma.analysis.drop_periodic_endpoint("eta2").rename("physical_poloidal_velocity")
    temporal = output.analysis.time_fft(plane)
    band = output.analysis.filter_time(plane, dims=("eta1", "eta2"), pad_bins=0)
    positive = temporal.power.isel(omega=slice(1, None))
    omega = positive.omega.values
    frequency_resolution = temporal.attrs["frequency_resolution"]
    mean_power = positive.mean(("eta1", "eta2"))
    radius_power = positive.mean("eta2")
    frequency_plot = make_subplots(rows=1, cols=3, subplot_titles=titles, horizontal_spacing=0.1)
    radius_spectrum = make_subplots(rows=1, cols=3, subplot_titles=labels, horizontal_spacing=0.12)
    filtered_probe = make_subplots(rows=1, cols=3, subplot_titles=labels, horizontal_spacing=0.12)
    for component, label in enumerate(labels):
        power = mean_power.isel(component=component).values
        frequency_plot.add_scatter(
            x=omega, y=power / max(float(power.max()), 1e-30), mode="lines+markers",
            name="Power", showlegend=False, line={"color": "#0072B2"}, row=1, col=component + 1,
        )
        selected = band.spectrum.isel(component=component)
        if bool(selected.has_peak):
            # Bin edges make a one-bin FWHM band visible; metadata retains bin centers.
            frequency_plot.add_vrect(
                x0=max(0, float(selected.omega_lo) - frequency_resolution / 2),
                x1=float(selected.omega_hi) + frequency_resolution / 2,
                fillcolor="#E69F00", opacity=0.2, line_width=0, row=1, col=component + 1,
            )
        local_power = radius_power.isel(component=component).transpose("omega", "eta1").values
        log_power = np.log10(np.maximum(local_power / max(float(local_power.max()), 1e-30), 1e-6))
        radius_spectrum.add_trace(go.Heatmap(
            x=radii, y=omega, z=log_power, zmin=-6, zmax=0, colorscale="Magma",
            colorbar={"title": "log₁₀ P/Pmax", "x": (0.28, 0.64, 1.0)[component], "len": 0.8,
                      "thickness": 10, "title_font": {"size": 10}},
            hovertemplate="r=%{x:.3f}<br>ω=%{y:.3f}<br>log₁₀ P/Pmax=%{z:.2f}<extra></extra>",
        ), row=1, col=component + 1)
        radius_spectrum.update_yaxes(showticklabels=component == 0, row=1, col=component + 1)
        for field, name, color, dash in ((plane, "Original", "#0072B2", "solid"),
                                         (band.filtered, "Dominant band", "#D55E00", "dash")):
            probe = field.isel(component=component, eta1=radial_probe, eta2=angle_probe)
            filtered_probe.add_scatter(
                x=times, y=probe.values, mode="lines+markers", name=name, legendgroup=name,
                showlegend=component == 0, line={"color": color, "dash": dash}, marker={"size": 4},
                row=1, col=component + 1,
            )
        filtered_probe.update_yaxes(title_text=label, exponentformat="power", tickfont={"size": 10},
                                    row=1, col=component + 1)
    frequency_plot.update_xaxes(title_text="ω (normalized)", range=[0, temporal.attrs["nyquist_frequency"]])
    frequency_plot.update_yaxes(range=[0, 1.05])
    frequency_plot.update_yaxes(title_text="Power / peak power", row=1, col=1)
    frequency_plot.update_layout(
        title=f"Temporal velocity spectra · Δω = {frequency_resolution:.3f}, N = {len(times)}",
        template="plotly_white", margin={"l": 65, "r": 30, "t": 100, "b": 80},
    )
    radius_spectrum.update_xaxes(title_text="Minor radius r", range=[radii[0], radii[-1]])
    radius_spectrum.update_yaxes(title_text="ω (normalized)", row=1, col=1)
    radius_spectrum.update_layout(
        title="Temporal power versus radius · poloidal average at φ = 0",
        template="plotly_white", margin={"l": 85, "r": 100, "t": 90, "b": 70},
    )
    filtered_probe.update_xaxes(title_text="t")
    filtered_probe.update_layout(
        title=f"Dominant-band reconstruction · r = {probe_radius:.3f}, θ = 45°, φ = 0",
        template="plotly_white", margin={"l": 85, "r": 35, "t": 100, "b": 100},
        legend={"orientation": "h", "x": 0.5, "xanchor": "center", "y": -0.18},
    )

    # Check the seeded poloidal modes in the logical radial component, before
    # physical-basis rotation introduces additional geometric harmonics.
    # On the post-processing grid of `velocity` above (celldivide 3 in eta1 and eta2), in the plane eta3 = 0.
    logical_initial = output.evaluate(
        "mhd/velocity",
        eta1=np.linspace(0.0, 1.0, 3 * grid.num_elements[0] + 1),
        eta2=np.linspace(0.0, 1.0, 3 * grid.num_elements[1] + 1),
        eta3=0.0,
        representation="2",
    ).isel(t=0, component=0)
    poloidal_fft = logical_initial.plasma.analysis.mode_spectrum(dims="eta2", names="m")
    mode_numbers = poloidal_fft.m.values
    modal_amplitude = np.sqrt((abs(poloidal_fft) ** 2).mean("eta1")).values
    positive_modes = (mode_numbers > 0) & (mode_numbers <= grid.num_elements[1] // 2)
    mode_plot = go.Figure(go.Scatter(
        x=mode_numbers[positive_modes], y=2 * modal_amplitude[positive_modes], mode="lines+markers",
        line={"color": "#0072B2"}, name="Initial radial mode amplitude",
    ))
    for m in modes:
        mode_plot.add_vline(x=m, line_dash="dot", line_color="#D55E00", annotation_text=f"m={m}",
                            annotation_position="top left" if m == modes[0] else "top right")
    mode_plot.update_layout(
        title="Initial poloidal Fourier modes · logical radial velocity",
        xaxis_title="Poloidal mode number m", yaxis_title="Mode amplitude (RMS over sampled radii)",
        template="plotly_white", margin={"l": 85, "r": 35, "t": 100, "b": 70},
    )

    # Radial mode structures in the style of Fig. 5 of arXiv:2510.04385:
    # spatial (m,n) amplitudes at a single time, for velocity and perturbed B.
    for field_name, label, key in (
        ("mhd/velocity_xyz", "u_r", "radial-fft-velocity"),
        ("em_fields/b_field_xyz", "δB_r", "radial-fft-magnetic"),
    ):
        snapshot = output.evaluate(field_name).isel(t=-1)
        profiles = radial_mode_amplitudes(output, snapshot)
        snapshot_time = profiles.attrs["snapshot_time"]
        toroidal_mode = profiles.attrs["full_torus_mode"]
        radial_mode_plot = go.Figure()
        for m, color in zip(profiles.m.values, ("#0072B2", "#D55E00", "#009E73", "#CC79A7")):
            radial_mode_plot.add_scatter(
                x=profiles.radius.values, y=profiles.sel(m=m).values,
                mode="lines+markers", name=f"m={m}",
                line={"color": color, "width": 2.5}, marker={"size": 4},
                hovertemplate=f"m={m}<br>r=%{{x:.3f}}<br>Normalized amplitude=%{{y:.4f}}<extra></extra>",
            )
        radial_mode_plot.update_layout(
            title=f"Radial Fourier mode structure · {label} · |n|={toroidal_mode:g}, t={snapshot_time:g}",
            xaxis_title="Minor radius r", yaxis_title="Normalized FFT amplitude",
            xaxis={"range": [radii[0], radii[-1]]}, yaxis={"range": [0, 1.05]},
            template="plotly_white", margin={"l": 85, "r": 35, "t": 95, "b": 75},
            legend={"title": {"text": "Poloidal harmonic"}},
        )
        save_figure(radial_mode_plot, f"{STEM}-{key}", show=show)

    energy = go.Figure()
    for key, label in (("en_U", "Kinetic"), ("en_B", "Magnetic"), ("en_thermal", "Compressional")):
        values = output.scalars[key]
        if not np.isfinite(values.values).all():
            raise RuntimeError(f"Non-finite energy diagnostic: {key}")
        energy.add_scatter(x=values.t.values, y=values.values, mode="lines", name=label)
    energy.update_layout(
        title="Perturbation energies", xaxis_title="t", yaxis_title="Energy [normalized units]",
        template="plotly_white", margin={"l": 80, "r": 30, "t": 80, "b": 65},
    )
    save_figure(history, f"{STEM}-velocity-history", show=show)
    for angle_degrees, radial_profiles in radial_profile_figures:
        save_figure(radial_profiles, f"{STEM}-radial-profiles-theta-{int(angle_degrees)}", show=show)
    save_figure(radial_history, f"{STEM}-radial-history", show=show)
    save_figure(mode_plot, f"{STEM}-poloidal-fft", show=show)
    save_fixed_theta_fft_figures(output, show=show)
    save_figure(frequency_plot, f"{STEM}-time-fft", show=show)
    save_figure(radius_spectrum, f"{STEM}-radial-time-fft", show=show)
    save_figure(filtered_probe, f"{STEM}-filtered-velocity", show=show)
    save_figure(energy, f"{STEM}-energy", show=show)
    dominant = ", ".join(f"{label} {value:.3f}" for label, value in zip(labels, band.spectrum.dominant_frequency.values))
    print(f"Dominant temporal frequencies: {dominant} (resolution {frequency_resolution:.3f})")


if __name__ == "__main__":
    argparser = argparse.ArgumentParser(description="Run the toroidal shear alfven example.")
    argparser.add_argument(
        "--pproc-only",
        action="store_true",
        help="Run post-processing on an existing simulation instead of running a new one.",
    )
    argparser.add_argument("--show", action="store_true", help="Show the figures before saving them.")
    args = argparser.parse_args()

    simulation = create_simulation()
    if not args.pproc_only:
        simulation.run()
    pproc(simulation, show=args.show)