Ordinary waves and the plasma cutoff

"""Ordinary electromagnetic waves and the plasma-frequency cutoff.

For k perpendicular to B0 and E parallel to B0, the cold-plasma O mode obeys
omega**2 = omega_p**2 + c**2*k**2. Four standing modes are evolved together;
their frequencies are measured independently from electric-field zero crossings.
This periodic initial-value experiment measures propagating modes, not reflection
from an interface. The cutoff is the k -> 0 limit of the dispersion relation.

Reference: https://farside.ph.utexas.edu/teaching/315/Waveshtml/node75.html
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 ordinary-mode-dispersion.py

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

import argparse

import numpy as np

from struphy import DerhamOptions, EnvironmentOptions, Simulation, Time, domains, equils, grids, perturbations
from struphy.models import ColdPlasma
from struphy.linear_algebra.solver import SolverParameters
from plasma_plots.theory.waves import plasma_light_wave
from plasma_plots import save_figure

stem = "ordinary-mode-dispersion"
length = 8.0 * np.pi
mode_numbers = (1, 2, 4, 8)
amplitude = 0.02  # electric-field amplitude per mode
wavenumbers = 2.0 * np.pi * np.array(mode_numbers) / length
frequencies = plasma_light_wave(wavenumbers).real  # omega**2 = omega_p**2 + c**2 k**2, c = omega_p = 1


def create_simulation() -> Simulation:
    model = ColdPlasma(alpha=1.0, epsilon=1.0)
    for propagator in (model.propagators.maxwell, model.propagators.ohm, model.propagators.jxb):
        propagator.options = propagator.Options(solver_params=SolverParameters(tol=1e-12))
    model.em_fields.e_field.add_perturbation(
        perturbations.ModesCos(ls=mode_numbers, amps=(amplitude,) * len(mode_numbers),
                               comp=2, Lx=length, given_in_basis="physical"),
    )
    domain = domains.Cuboid(r1=length)
    grid = grids.TensorProductGrid(num_elements=(64, 1, 1))
    derham_opts = DerhamOptions(degree=(3, 1, 1))
    time_opts = Time(dt=0.02, Tend=30.0, split_algo="Strang")
    sim = Simulation(
        model=model, name="Ordinary waves and the plasma cutoff",
        description=(
            "Four ordinary electromagnetic waves oscillate across a uniform magnetic field. "
            "Their electric field is parallel to the background field, so the electrons feel no "
            "magnetic force in this polarization. Measured frequencies follow the cold-plasma dispersion relation"
            r" $$\omega^2=\omega_p^2+c^2k^2,\qquad \omega_p=c=1.$$"
            r" Initially :math:`E_z(x,0)=0.02\sum_{n\in\{1,2,4,8\}}\cos(nx/4)` and the magnetic perturbation and current vanish. "
            "The dispersion curve approaches the plasma-frequency cutoff as the wavelength grows; "
            "below it, a uniform cold plasma has no propagating ordinary mode."
        ),
        env=EnvironmentOptions(out_folders="struphy_gallery_runs", sim_folder="ordinary_mode_dispersion"),
        time_opts=time_opts, domain=domain, grid=grid, derham_opts=derham_opts,
        equil=equils.HomogenSlab(B0z=1.0, n0=1.0),
    )
    return sim


def pproc(sim: Simulation, show: bool = False):
    time_opts = sim.time_opts
    from plotly.subplots import make_subplots

    output = sim.output
    output.pproc(physical=True)
    field = output.evaluate("em_fields/e_field_xyz").isel(component=2, eta2=0, eta3=0)
    energy = output.scalars["total_energy"]
    times, x = field.t.values, field.eta1.values * length
    if not np.isfinite(field.values).all() or not np.isfinite(energy.values).all() or not np.isclose(times[-1], time_opts.Tend):
        raise RuntimeError("The ordinary-mode run is incomplete or non-finite")
    # cos(k x) is mode n along eta1; project_mode drops the repeated periodic endpoint.
    signals = np.column_stack([
        field.plasma.analysis.project_mode(dim="eta1", number=n, kind="cos").values for n in mode_numbers
    ])
    measured = []
    for signal in signals.T:
        crossing = np.flatnonzero(np.diff(np.signbit(signal)))
        if crossing.size < 4:
            raise RuntimeError("Too few zero crossings to measure an ordinary-mode frequency")
        roots = times[crossing] - signal[crossing] * np.diff(times)[crossing] / np.diff(signal)[crossing]
        measured.append(float(np.pi / np.mean(np.diff(roots))))
    measured = np.array(measured)
    reference = amplitude * np.cos(times[:, None] * frequencies[None, :])
    frequency_error = float(np.max(np.abs(measured / frequencies - 1.0)))
    mode_error = float(np.max(np.abs(signals - reference)) / amplitude)
    energy_drift = float(energy.plasma.analysis.relative_error().max())
    print(f"Measured frequencies: {np.round(measured, 4).tolist()} (exact: {np.round(frequencies, 4).tolist()})")
    if frequency_error > 0.01 or mode_error > 0.1 or energy_drift > 1e-6:
        raise RuntimeError(f"Ordinary-mode check failed: frequency={frequency_error:.3g}, field={mode_error:.3g}, energy={energy_drift:.3g}")

    figure = make_subplots(rows=1, cols=2, horizontal_spacing=0.13,
                           subplot_titles=("Dispersion and cutoff", "Four independently measured modes"))
    k = np.linspace(0, 2.2, 200)
    figure.add_scatter(x=k, y=plasma_light_wave(k).real, name="ω² = 1 + k²", line={"color": "#222"}, row=1, col=1)
    figure.add_scatter(x=wavenumbers, y=measured, mode="markers", name="Measured frequencies",
                       marker={"size": 10, "color": "#d62828"}, row=1, col=1)
    figure.add_hrect(y0=0, y1=1, fillcolor="#168aad", opacity=0.08, line_width=0, row=1, col=1)
    figure.add_hline(y=1, line_dash="dot", annotation_text="ω_p: cutoff", row=1, col=1)
    for j, color in enumerate(("#168aad", "#d62828", "#f77f00", "#6a4c93")):
        figure.add_scatter(x=times, y=signals[:, j] / amplitude + 2.5 * j, name=f"k = {wavenumbers[j]:g}",
                           line={"color": color}, row=1, col=2)
        figure.add_scatter(x=times[::25], y=reference[::25, j] / amplitude + 2.5 * j, mode="markers",
                           name="Exact oscillation", showlegend=j == 0,
                           marker={"color": "#222", "symbol": "circle-open", "size": 4}, row=1, col=2)
    figure.update_xaxes(title_text="k c / ω_p", row=1, col=1)
    figure.update_yaxes(title_text="ω / ω_p", range=[0, 2.5], row=1, col=1)
    figure.update_xaxes(title_text="t ω_p", row=1, col=2)
    figure.update_yaxes(title_text="E-mode / A + vertical offset", row=1, col=2)
    figure.update_layout(title="Ordinary electromagnetic waves in a cold plasma", template="plotly_white",
                         legend={"orientation": "h", "y": -0.2}, margin={"l": 70, "r": 30, "t": 100, "b": 140})
    save_figure(figure, stem, height=650, show=show)
    space_time = field.assign_coords(eta1=x).plasma.plot.slice(
        x="eta1", y="t", symmetric=True, cmap="RdBu_r", title="Superposed ordinary waves: E_z(x, t)",
        xlabel="x", ylabel="t [a.u.]", colorbar_label="E_z", backend="plotly",
    )
    save_figure(space_time, f"{stem}-space-time", show=show)


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