Whistler and ion-cyclotron waves in Hall MHD

"""Whistler and ion-cyclotron waves of Hall MHD, with Struphy's LinearExtendedMHDuniform model.

Ideal MHD has one transverse wave along the magnetic field, the shear Alfvén wave with omega = v_A k.
The Hall term splits it into two circularly polarized branches once the wavelength approaches the ion
inertial length d_i: the right-hand whistler, whose frequency grows like k^2, and the left-hand
ion-cyclotron wave, which saturates at the ion cyclotron frequency. Broadband noise in the velocity
excites both, and the sound wave along the field. The (k, omega) power spectra of the magnetic field and
the pressure are compared with the analytic Hall-MHD branches, and the phase velocities read off the
spectrum with the exact ones.

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 hall-mhd-waves.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 struphy import DerhamOptions, EnvironmentOptions, Simulation, Time, domains, equils, grids, perturbations
from struphy.models import LinearExtendedMHDuniform
from plasma_plots import save_figure

# Background field along z with n0 = B0 = 1, so the Alfvén speed is 1. With epsilon = 1 the ion cyclotron
# frequency and the ion inertial length are 1 as well, i.e. k is measured in units of 1/d_i.
B0z, n0, beta, gamma, epsilon = 1.0, 1.0, 0.5, 5.0 / 3.0, 1.0
p0 = beta * B0z**2 / 2.0
alfven_speed = B0z / np.sqrt(n0)
sound_speed = np.sqrt(gamma * p0 / n0)
cyclotron_frequency = B0z / epsilon


def hall_branches(k):
    """Parallel Hall-MHD branches: omega = v_A k (sqrt(1 + k^2 d_i^2 / 4) +- k d_i / 2) with d_i = v_A / Omega_i, and sound."""
    from plasma_plots.theory.waves import hall_mhd_parallel

    hall = hall_mhd_parallel(k, alfven_speed=alfven_speed, ion_inertial_length=alfven_speed / cyclotron_frequency)
    return {"whistler": hall["whistler"].real, "ion-cyclotron": hall["ion cyclotron"].real, "sound": sound_speed * k}


def create_simulation() -> Simulation:
    model = LinearExtendedMHDuniform(epsilon=epsilon)

    # Broadband noise in all three velocity components: the transverse ones drive the two Hall branches,
    # the parallel one the sound wave.
    for component in range(3):
        model.mhd.velocity.add_perturbation(perturbations.Noise(amp=0.1, comp=component, seed=123))

    domain = domains.Cuboid(r3=60.0)
    grid = grids.TensorProductGrid(num_elements=(1, 1, 64))
    derham_opts = DerhamOptions(degree=(1, 1, 3))
    equil = equils.HomogenSlab(B0x=0.0, B0y=0.0, B0z=B0z, beta=beta, n0=n0)
    time_opts = Time(dt=0.1, Tend=100.0)

    env = EnvironmentOptions(out_folders="struphy_gallery_runs", sim_folder="hall_mhd_waves")
    sim = Simulation(
        model=model,
        name="Whistler and ion-cyclotron waves in Hall MHD",
        description=(
            "Broadband noise in a uniform magnetized plasma excites the two circularly polarized branches into which "
            "the Hall term splits the shear Alfvén wave: the whistler, whose frequency grows like k², and the "
            "ion-cyclotron wave, which saturates at the ion cyclotron frequency. Their power spectra and phase "
            "velocities are compared with the analytic Hall-MHD dispersion relation."
            r" The initial background has $$\mathbf{B}_0=\mathbf{e}_z,\qquad n_0=1,\qquad p_0=0.25,\qquad \epsilon=1,$$"
            r" with coefficient-noise amplitude :math:`0.1` in all three velocity components and periodic length :math:`L_z=60`."
        ),
        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):
    from plotly.subplots import make_subplots

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

    from plasma_plots.analysis import power_spectrum

    length = output.domain.params["r3"] - output.domain.params["l3"]

    def spectrum_along_z(field):
        """omega >= 0, k >= 0 and the Fourier amplitude (the square root of the power) of a field along z."""
        line = field.isel(eta1=0, eta2=0)
        line = line.isel(component=0) if "component" in line.dims else line
        spectrum = power_spectrum(line.assign_coords(eta3=line.eta3 * length), dim="eta3")  # physical z
        quadrant = spectrum.sel(omega=spectrum.omega >= 0, k=spectrum.k >= 0)
        return quadrant.omega.values, quadrant.k.values, np.sqrt(quadrant.values)

    omega_b, k_b, spectrum_b = spectrum_along_z(output.fields.em_fields.b_field)
    omega_p, k_p, spectrum_p = spectrum_along_z(output.fields.mhd.pressure)

    k_top, omega_top = 3.0, 6.0
    k_fine = np.linspace(0.0, k_top, 300)
    exact = hall_branches(k_fine)

    # Read each branch off the spectrum: at every k, the frequency of the strongest power within a window
    # around the exact branch. The window is narrow enough to keep the whistler and the ion-cyclotron wave apart.
    def measured_branch(omega, k, spectrum, branch):
        ks, omegas = [], []
        for column, kk in enumerate(k):
            if not 0.0 < kk <= k_top:
                continue
            target = hall_branches(np.array([kk]))[branch][0]
            window = np.flatnonzero(np.abs(omega - target) < 0.2 * target + 2.0 * (omega[1] - omega[0]))
            if window.size == 0:
                continue
            ks.append(kk)
            omegas.append(omega[window[np.argmax(np.abs(spectrum[window, column]))]])
        return np.array(ks), np.array(omegas)

    measured = {
        "whistler": measured_branch(omega_b, k_b, spectrum_b, "whistler"),
        "ion-cyclotron": measured_branch(omega_b, k_b, spectrum_b, "ion-cyclotron"),
        "sound": measured_branch(omega_p, k_p, spectrum_p, "sound"),
    }
    errors = {}
    for branch, (ks, omegas) in measured.items():
        reference = hall_branches(ks)[branch]
        errors[branch] = float(np.median(np.abs(omegas / reference - 1.0)))
        print(f"{branch}: median relative frequency error {errors[branch]:.3f} over {ks.size} wavenumbers")
    if not all(np.isfinite(error) for error in errors.values()):
        raise RuntimeError("A wave branch could not be read off the spectrum")

    def log_power(spectrum):
        power = spectrum**2
        return np.log10(np.clip(power / power.max(), 1e-12, None))

    colors = {"whistler": "#48cae4", "ion-cyclotron": "#e9c46a", "sound": "#2a9d8f", "alfven": "#ffffff"}
    labels = {
        "whistler": "whistler (R)",
        "ion-cyclotron": "ion-cyclotron (L)",
        "sound": "sound",
        "alfven": "ideal-MHD Alfvén wave, ω = v_A k",
    }
    panels = (
        ("Magnetic field B_x", omega_b, k_b, spectrum_b, ("whistler", "ion-cyclotron")),
        ("Pressure", omega_p, k_p, spectrum_p, ("sound",)),
    )
    figure = make_subplots(rows=1, cols=2, subplot_titles=[panel[0] for panel in panels], horizontal_spacing=0.12)
    for column, (_, omega, k, spectrum, branches) in enumerate(panels, start=1):
        figure.add_trace(
            go.Heatmap(
                x=k, y=omega, z=log_power(spectrum), zmin=-8, zmax=0, colorscale="Plasma",
                showscale=column == 2, colorbar={"title": {"text": "log₁₀ P"}, "len": 0.9},
                hovertemplate="k=%{x:.3f}<br>ω=%{y:.3f}<br>log₁₀ P=%{z:.2f}<extra></extra>",
            ),
            row=1, col=column,
        )
        for branch in branches:
            figure.add_trace(
                go.Scatter(x=k_fine, y=exact[branch], mode="lines", name=labels[branch], legendgroup=branch,
                           line={"color": colors[branch], "width": 2.5, "dash": "dash"}),
                row=1, col=column,
            )
        figure.update_xaxes(title_text="k d_i", range=[0, k_top], row=1, col=column)
        figure.update_yaxes(title_text="ω / Ω_i", range=[0, omega_top], row=1, col=column)
    figure.add_trace(
        go.Scatter(x=k_fine, y=alfven_speed * k_fine, mode="lines", name=labels["alfven"],
                   line={"color": colors["alfven"], "width": 1.5, "dash": "dot"}),
        row=1, col=1,
    )
    figure.add_hline(y=cyclotron_frequency, line={"color": "rgba(255,255,255,0.6)", "width": 1, "dash": "dot"},
                     annotation_text="Ω_i", annotation_font_color="white", annotation_position="bottom right",
                     row=1, col=1)
    figure.update_layout(
        title="Hall MHD along B₀: the Alfvén wave splits into whistler and ion-cyclotron branches",
        template="plotly_white", legend={"orientation": "h", "y": -0.2},
        margin={"l": 70, "r": 40, "t": 90, "b": 110},
    )
    save_figure(figure, "hall-mhd-waves", width=1300, height=650, show=show)

    # Phase velocities: the ideal-MHD Alfvén wave is not dispersive (omega / k = v_A for every k); the
    # Hall branches are, one speeding up and one slowing down as k d_i grows.
    phase = go.Figure()
    positive = k_fine > 0
    for branch in ("whistler", "ion-cyclotron", "sound"):
        phase.add_scatter(x=k_fine[positive], y=exact[branch][positive] / k_fine[positive], mode="lines",
                          name=f"{labels[branch]}, exact", legendgroup=branch,
                          line={"color": colors[branch] if branch != "whistler" else "#0077b6", "width": 2.5})
        ks, omegas = measured[branch]
        phase.add_scatter(x=ks, y=omegas / ks, mode="markers", name=f"{labels[branch]}, from the spectrum",
                          legendgroup=branch,
                          marker={"color": colors[branch] if branch != "whistler" else "#0077b6", "size": 7,
                                  "line": {"color": "#264653", "width": 1}})
    phase.add_hline(y=alfven_speed, line={"color": "#6c757d", "dash": "dot"},
                    annotation_text="ideal MHD: v_A", annotation_position="top left")
    phase.update_layout(
        title="Phase velocity of the Hall-MHD waves", xaxis_title="k d_i", yaxis_title="ω / k  [v_A]",
        template="plotly_white", autosize=True, legend={"orientation": "h", "y": -0.2},
        margin={"l": 75, "r": 30, "t": 80, "b": 120},
    )
    phase.update_xaxes(range=[0, k_top])
    save_figure(phase, "hall-mhd-waves-phase-velocity", show=show)


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