Coaxial waveguide

"""Coaxial waveguide: an exact electromagnetic mode between two conducting cylinders.

Between two concentric perfectly conducting cylinders, Maxwell's equations have exact eigenmodes.
For azimuthal number m = 3 and no axial variation, the fields are a Bessel-function profile in the
radius times cos(m*theta - t): a pattern that rotates rigidly at frequency omega = 1. Struphy's
finite element Maxwell solver on the annulus is compared with this solution.

Adapted from Struphy's tutorial (tutorials/tutorial_maxwell.ipynb) and its verification test.

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 coaxial-waveguide.py

Figures are saved in the current directory; add --show to display them before saving.
"""

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 scipy.optimize import curve_fit
from scipy.special import jv, yn

from struphy import DerhamOptions, EnvironmentOptions, Simulation, Time, domains, equils, grids, perturbations
from struphy.models import Maxwell
from plasma_plots import save_figure

# The inner and outer radii are two zeros of the derivative of the radial profile
# J_m(r) - 0.28 Y_m(r), so that the tangential electric field vanishes on both conducting walls.
inner_radius = 2.326744
outer_radius = 3.686839
length = 2.0
mode_number = 3
bessel_ratio = 0.28


def create_simulation() -> Simulation:
    model = Maxwell()
    model.propagators.maxwell.options = model.propagators.maxwell.Options(algo="implicit")
    model.em_fields.e_field.add_perturbation(
        perturbations.CoaxialWaveguideElectric_r(m=mode_number, a1=inner_radius, a2=outer_radius)
    )
    model.em_fields.e_field.add_perturbation(
        perturbations.CoaxialWaveguideElectric_theta(m=mode_number, a1=inner_radius, a2=outer_radius)
    )
    model.em_fields.b_field.add_perturbation(
        perturbations.CoaxialWaveguideMagnetic(m=mode_number, a1=inner_radius, a2=outer_radius)
    )

    # The annulus in the (x, y) plane, with conducting walls (Dirichlet conditions) in the radial direction.
    domain = domains.HollowCylinder(a1=inner_radius, a2=outer_radius, Lz=length)
    grid = grids.TensorProductGrid(num_elements=(32, 64, 1))
    derham_opts = DerhamOptions(degree=(3, 3, 1), bcs=(("dirichlet", "dirichlet"), None, None))

    env = EnvironmentOptions(
        out_folders="struphy_gallery_runs",
        sim_folder="coaxial_waveguide",
    )
    sim = Simulation(
        model=model,
        name="Coaxial waveguide",
        description=(
            "An exact electromagnetic mode between two concentric conducting cylinders "
            "rotates around the annulus. Struphy's structure-preserving Maxwell solver "
            "reproduces its shape, its frequency and the conservation of its energy."
            r" With :math:`F(r)=J_3(r)-0.28Y_3(r)`, the initial mode is $$E_r=-\frac{3F(r)}{r}\cos(3\theta),\qquad E_\theta=F\prime(r)\sin(3\theta),\qquad B_z=F(r)\cos(3\theta),$$"
            r" between :math:`r=2.326744` and :math:`r=3.686839`; :math:`J_3` and :math:`Y_3` are Bessel functions."
        ),
        env=env,
        time_opts=Time(dt=0.05, Tend=20.0),
        domain=domain,
        equil=equils.HomogenSlab(),
        grid=grid,
        derham_opts=derham_opts,
    )
    return sim


def pproc(sim: Simulation, show: bool = False):
    from plasma_plots.analysis import evaluate_on

    output = sim.output
    output.pproc(physical=True)

    # The axial magnetic field on the (r, theta) evaluation grid, and the exact mode at the same points.
    b_z = output.evaluate("em_fields/b_field_xyz").isel(component=2, eta3=0)  # (t, e1, e2)
    times = b_z.t.values

    def exact_b_z(x, y, z, time):
        radius, angle = np.hypot(x, y), np.arctan2(y, x)
        profile = jv(mode_number, radius) - bessel_ratio * yn(mode_number, radius)
        return profile * np.cos(mode_number * angle - time)

    exact = evaluate_on(b_z, exact_b_z)
    error = b_z.plasma.analysis.error(exact, norm="pointwise").values
    amplitude = float(np.abs(exact.isel(t=0)).max())
    relative_error = b_z.plasma.analysis.error(exact, norm="max") / amplitude
    print(f"Largest relative error of B_z: {float(relative_error.max()):.2e}")

    # The mode's frequency, from a sinusoid fitted to B_z at a probe in the middle of the gap. The
    # exact mode is proportional to cos(m theta - t), so its frequency is 1.
    probe = {"eta1": b_z.sizes["eta1"] // 2, "eta2": 0}
    signal = b_z.isel(probe).assign_attrs(label="B_z at the probe")
    exact_signal = exact.isel(probe)
    fit, _ = curve_fit(
        lambda t, a, w, phase: a * np.cos(w * t + phase), times, signal.values, p0=[float(signal.max()), 1.0, 0.0]
    )
    measured_frequency = float(fit[1])
    frequency_error = abs(measured_frequency - 1.0)
    print(f"Measured frequency: {measured_frequency:.5f} (exact: 1, error {frequency_error:.1e})")

    energy = output.scalars["total_energy"]
    energy_drift = float(energy.plasma.analysis.relative_error().max())
    print(f"Largest relative change of the total energy: {energy_drift:.1e}")
    electric = output.scalars["electric_energy"]
    magnetic = output.scalars["magnetic_energy"]

    # The animation: both fields are interpolated from the (r, theta) grid onto a Cartesian grid of
    # pixels, which are left empty outside the annulus. Every frame is embedded in the page, so the
    # smooth error map gets fewer pixels than the field.
    def pixel_grid(pixels):
        axis = np.linspace(-outer_radius, outer_radius, pixels)
        pixel_x, pixel_y = np.meshgrid(axis, axis)
        pixel_radius = np.hypot(pixel_x, pixel_y)
        pixel_angle = np.mod(np.arctan2(pixel_y, pixel_x), 2.0 * np.pi)
        outside = (pixel_radius < inner_radius) | (pixel_radius > outer_radius)
        index_r = (pixel_radius - inner_radius) / (outer_radius - inner_radius) * (b_z.sizes["eta1"] - 1)
        index_theta = pixel_angle / (2.0 * np.pi) * (b_z.sizes["eta2"] - 1)
        return axis, outside, index_r, index_theta

    def on_pixels(values, grid):
        _, outside, index_r, index_theta = grid
        image = map_coordinates(values, [index_r, index_theta], order=1, mode="nearest").astype(np.float32)
        image[outside] = np.nan
        return image

    field_grid = pixel_grid(112)
    error_grid = pixel_grid(96)

    frame_indices = np.linspace(0, len(times) - 1, 60, dtype=int)
    error_limit = float(np.abs(error).max())

    def frame_traces(index):
        return [
            go.Heatmap(
                z=on_pixels(b_z.isel(t=index).values, field_grid),
                x=field_grid[0],
                y=field_grid[0],
                zmin=-amplitude,
                zmax=amplitude,
                colorscale="RdBu",
                colorbar={"title": "B_z", "x": 0.45, "len": 0.75},
                xaxis="x",
                yaxis="y",
            ),
            go.Heatmap(
                z=on_pixels(error[index], error_grid),
                x=error_grid[0],
                y=error_grid[0],
                zmin=-error_limit,
                zmax=error_limit,
                colorscale="PuOr",
                colorbar={"title": "error", "x": 1.0, "len": 0.75, "exponentformat": "e"},
                xaxis="x2",
                yaxis="y2",
            ),
        ]

    figure = make_subplots(
        rows=1, cols=2, subplot_titles=("B_z (Struphy)", "Error against the exact mode"), horizontal_spacing=0.14
    )
    first, second = frame_traces(0)
    figure.add_trace(first, row=1, col=1)
    figure.add_trace(second, row=1, col=2)
    frames = [go.Frame(name=f"{times[index]:.1f}", data=frame_traces(index), traces=[0, 1]) for index in frame_indices]
    figure.frames = frames
    limit = 1.02 * outer_radius
    figure.update_xaxes(range=[-limit, limit], constrain="domain", title_text="x [a.u.]")
    figure.update_yaxes(range=[-limit, limit], title_text="y [a.u.]", row=1, col=1)
    figure.update_yaxes(range=[-limit, limit], row=1, col=2)
    figure.update_yaxes(scaleanchor="x", scaleratio=1, row=1, col=1)
    figure.update_yaxes(scaleanchor="x2", scaleratio=1, row=1, col=2)
    figure.update_layout(
        title="Coaxial waveguide: a rotating electromagnetic mode",
        template="plotly_white",
        autosize=True,
        margin={"l": 70, "r": 30, "t": 100, "b": 130},
        updatemenus=[
            {
                "type": "buttons",
                "showactive": False,
                "x": 0.0,
                "xanchor": "left",
                "y": -0.2,
                "yanchor": "top",
                "buttons": [
                    {
                        "label": "Play",
                        "method": "animate",
                        "args": [None, {"frame": {"duration": 60, "redraw": True}, "fromcurrent": True}],
                    },
                ],
            },
        ],
        sliders=[
            {
                "steps": [
                    {
                        "args": [[frame.name], {"frame": {"duration": 0, "redraw": True}, "mode": "immediate"}],
                        "label": frame.name,
                        "method": "animate",
                    }
                    for frame in frames
                ],
                "active": 0,
                "x": 0.12,
                "len": 0.88,
                "y": -0.1,
                "currentvalue": {"prefix": "t = "},
            },
        ],
    )

    # The still image and thumbnail show the mode a little after the start of the run.
    still_position = len(frame_indices) // 4
    save_figure(figure, "coaxial-waveguide", height=650, frame=still_position, show=show)

    # The probe signal against the exact mode.
    probe_figure = signal.plasma.plot.timeseries(
        logy=False,
        reference={"exact mode": exact_signal},
        title=f"Coaxial waveguide: B_z at a probe, measured frequency {measured_frequency:.4f} (exact: 1)",
        backend="plotly",
    )

    # The energies, and the relative change of the total energy.
    energy_figure = make_subplots(
        rows=2,
        cols=1,
        shared_xaxes=True,
        vertical_spacing=0.1,
        subplot_titles=("Energies", "Relative change of the total energy"),
    )
    for name, series, color, dash in (
        ("electric", electric, "#168aad", "solid"),
        ("magnetic", magnetic, "#f77f00", "dash"),
        ("total", energy, "#111111", "solid"),
    ):
        energy_figure.add_scatter(
            x=times,
            y=series.values,
            mode="lines",
            name=name,
            line={"color": color, "width": 3, "dash": dash},
            row=1,
            col=1,
        )
    energy_figure.add_scatter(
        x=times,
        y=np.abs(energy.values / energy.values[0] - 1.0) + 1e-17,
        mode="lines",
        showlegend=False,
        line={"color": "#111111", "width": 2},
        row=2,
        col=1,
    )
    energy_figure.update_yaxes(title_text="energy [a.u.]", rangemode="tozero", row=1, col=1)
    energy_figure.update_yaxes(
        title_text="|ΔW / W₀|", type="log", dtick=1, exponentformat="power", range=[-16, -13], row=2, col=1
    )
    energy_figure.update_xaxes(title_text="t [a.u.]", row=2, col=1)
    energy_figure.update_layout(template="plotly_white", autosize=True, margin={"l": 80, "r": 30, "t": 80, "b": 60})

    save_figure(probe_figure, "coaxial-waveguide-frequency", show=show)
    save_figure(energy_figure, "coaxial-waveguide-energy", show=show)


if __name__ == "__main__":
    argparser = argparse.ArgumentParser(description="Run the coaxial waveguide 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)