Guiding-center orbits in a tokamak

"""Guiding-center orbits in a tokamak: passing and trapped particles.

Eight guiding centers start on the outboard midplane of one flux surface with the same energy and
different pitch v_parallel / v. In the field of a circular tokamak the field strength is larger on
the inner side, so particles with a small parallel velocity are reflected there and bounce back and
forth (trapped, banana-shaped orbits in the poloidal plane), while faster particles circle the
magnetic axis (passing). The static magnetic field does no work, so the energy and the canonical
toroidal momentum of every particle are conserved, which checks the guiding-center pusher.

Adapted from Struphy's particle-tracing tutorial (tutorials/tutorial_particle_tracing.ipynb).

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 guiding-center-orbits.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 struphy import (
    BoundaryParameters,
    DerhamOptions,
    EnvironmentOptions,
    LoadingParameters,
    SavingParameters,
    Simulation,
    Time,
    WeightsParameters,
    domains,
    equils,
    grids,
    maxwellians,
)
from struphy.models import GuidingCenter
from plasma_plots import save_figure

# Eight markers on the outboard midplane of the flux surface eta1 = 0.5, all with speed 3 and
# different pitch. The state of a guiding center is (position, v_parallel, mu), with the magnetic
# moment mu = v_perp^2 / (2 B) chosen from the pitch at the starting point.
speed = 3.0
pitches = (-0.9, -0.75, -0.5, -0.25, 0.25, 0.5, 0.75, 0.9)
start_eta1 = 0.5


def create_simulation() -> Simulation:
    # A circular tokamak (minor radius 1, major radius 3) with a field strength of 2 on the axis.
    equil = equils.AdhocTorus(B0=2.0)
    domain = domains.Tokamak(equilibrium=equil, num_elements=(4, 16), degree=(2, 3))
    equil.domain = domain
    b_start = float(equil.absB0(start_eta1, 0.0, 0.0, squeeze_out=True))
    initial = tuple(
        (start_eta1, 0.0, 0.0, speed * pitch, speed**2 * (1.0 - pitch**2) / (2.0 * b_start)) for pitch in pitches
    )

    model = GuidingCenter()
    model.kinetic_ions.set_markers(
        loading_params=LoadingParameters(Np=len(initial), seed=1, specific_markers=initial),
        weights_params=WeightsParameters(),
        # A marker that leaves through the plasma edge is removed; the angles are periodic.
        boundary_params=BoundaryParameters(bc=("remove", "periodic", "periodic")),
        saving_params=SavingParameters(n_markers=1.0),
        bufsize=2.0,
    )
    model.propagators.push_bxe.options = model.propagators.push_bxe.Options(tol=1e-5)
    model.propagators.push_parallel.options = model.propagators.push_parallel.Options(tol=1e-5)
    model.kinetic_ions.var.add_background(maxwellians.GyroMaxwellian2D(n=(1.0, None), B0=b_start))

    # The equilibrium is projected onto splines, which the guiding-center pusher evaluates.
    grid = grids.TensorProductGrid(num_elements=(32, 64, 1))
    derham_opts = DerhamOptions(degree=(3, 3, 1), bcs=(("free", "free"), None, None))

    env = EnvironmentOptions(
        out_folders="struphy_gallery_runs",
        sim_folder="guiding_center_orbits",
    )
    sim = Simulation(
        model=model,
        name="Guiding-center orbits in a tokamak",
        description=(
            "Guiding centers with different pitch angles follow the field of a circular tokamak: "
            "particles with large parallel velocity circle the magnetic axis; others are reflected on the high-field "
            "side and bounce in banana-shaped orbits."
            r" At the logical launch point :math:`\boldsymbol{\eta}_0=(0.5,0,0)`, the markers have $$v_{\parallel,0}=3\xi,\qquad \mu=\frac{9(1-\xi^2)}{2B_{\mathrm{start}}},$$"
            r" with pitch :math:`\xi\in\{\pm0.25,\pm0.5,\pm0.75,\pm0.9\}` and :math:`B_{\mathrm{start}}=|\mathbf{B}(\boldsymbol{\eta}_0)|`."
        ),
        env=env,
        time_opts=Time(dt=0.05, Tend=100.0, split_algo="Strang"),
        domain=domain,
        equil=equil,
        grid=grid,
        derham_opts=derham_opts,
    )
    return sim


def pproc(sim: Simulation, show: bool = False):
    domain = sim.domain
    equil = sim.equil
    b_start = float(equil.absB0(start_eta1, 0.0, 0.0, squeeze_out=True))
    model = sim.model

    output = sim.output
    output.pproc()

    # (t, marker, quantity) with the quantities named x, y, z, v1, v2, ...: the position in Cartesian
    # coordinates, then the parallel velocity (v1) and the magnetic moment (v2) of a 5D marker. The
    # labels after v2 do not describe these columns and are not used here.
    orbits = output.orbits.kinetic_ions.to_dataarray("quantity").transpose("t", "marker", "quantity")
    times = orbits.t.values
    x, y, z = (orbits.sel(quantity=name).values for name in ("x", "y", "z"))
    v_parallel = orbits.sel(quantity="v_par").values
    moment = orbits.sel(quantity="mu").values
    major_radius = np.hypot(x, y)
    n_markers = x.shape[1]

    # Trapped particles are reflected: their parallel velocity changes sign along the orbit.
    reflections = (np.diff(np.sign(v_parallel), axis=0) != 0).sum(axis=0)
    trapped = reflections > 0
    labels = [f"v∥/v = {pitch:+.2f} ({'trapped' if trapped[i] else 'passing'})" for i, pitch in enumerate(pitches)]
    print("reflections per marker:", reflections.tolist())

    # The trapping boundary at the outboard midplane from the field strength on the starting surface:
    # a particle is trapped if v_parallel^2 < 2 mu (B_max - B_start).
    poloidal_angle = np.linspace(0.0, 1.0, 400, endpoint=False)
    surface_x, surface_y, surface_z = domain(start_eta1, poloidal_angle, 0.0, squeeze_out=True)
    surface_field = np.linalg.norm(
        np.array(equil.b_xyz(surface_x.ravel(), surface_y.ravel(), surface_z.ravel())), axis=0
    )
    boundary_pitch = float(np.sqrt(1.0 - b_start / surface_field.max()))
    print(f"Trapping boundary from the field on the starting surface: |v∥/v| = {boundary_pitch:.3f}")

    # Conserved quantities per marker: the energy v_par^2 / 2 + mu B, and the canonical toroidal momentum
    # R v_par b_phi - psi / epsilon (psi is set to zero on the magnetic axis). The magnetic moment is a
    # phase-space coordinate of the model and constant by construction.
    field = np.array(equil.b_xyz(x.ravel(), y.ravel(), z.ravel())).reshape(3, *x.shape)
    field_strength = np.linalg.norm(field, axis=0)
    toroidal_angle = np.arctan2(-y, x)
    toroidal_unit = np.array([-np.sin(toroidal_angle), -np.cos(toroidal_angle), np.zeros_like(x)])
    b_toroidal = (field * toroidal_unit).sum(axis=0) / field_strength
    epsilon = float(model.kinetic_ions.equation_params.epsilon)
    flux = equil.psi(major_radius, z) - float(equil.psi_range[0])
    energy = 0.5 * v_parallel**2 + moment * field_strength
    momentum = major_radius * v_parallel * b_toroidal - flux / epsilon
    energy_drift = np.abs(energy - energy[0]) / np.abs(energy[0])
    # P_phi can be nearly zero by cancellation. Normalize by the sum of the
    # magnitudes of its initial mechanical and flux contributions instead.
    momentum_scale = np.abs(major_radius[0] * v_parallel[0] * b_toroidal[0]) + np.abs(flux[0] / epsilon)
    momentum_drift = np.abs(momentum - momentum[0]) / momentum_scale
    print(f"Largest relative drift: energy {energy_drift.max():.2e}, canonical momentum {momentum_drift.max():.2e}")

    # The half period between two reflections is the time a trapped particle needs to go from one
    # mirror point to the other.
    half_periods = []
    for i in np.flatnonzero(reflections >= 2):
        crossings = times[1:][np.diff(np.sign(v_parallel[:, i])) != 0]
        half_periods.extend(np.diff(crossings))
    bounce_period = 2.0 * float(np.mean(half_periods))
    print(f"Bounce period of the trapped particles: {bounce_period:.1f}")

    trapped_colors = ["#d62828", "#f77f00", "#c9184a", "#9d4edd"]
    passing_colors = ["#168aad", "#1a759f", "#52b69a", "#34a0a4"]
    colors = []
    for i in range(n_markers):
        palette = trapped_colors if trapped[i] else passing_colors
        colors.append(palette[sum(1 for j in range(i) if trapped[j] == trapped[i]) % len(palette)])

    # The flux surfaces in the poloidal plane, for context.
    surfaces = np.array([0.25, 0.5, 0.75, 1.0])
    surface_x, surface_y, surface_z = domain(surfaces, np.linspace(0.0, 1.0, 200), 0.0, squeeze_out=True)
    surface_r = np.hypot(surface_x, surface_y)

    def add_flux_surfaces(figure, **position):
        for i in range(len(surfaces)):
            figure.add_trace(
                go.Scatter(
                    x=surface_r[i],
                    y=surface_z[i],
                    mode="lines",
                    line={"color": "rgba(110,120,130,0.45)", "width": 2 if surfaces[i] == 1.0 else 1},
                    hoverinfo="skip",
                    showlegend=False,
                ),
                **position,
            )

    # The animation: every guiding center is a bright head with a trail of its last 15 time units,
    # over its faint full path.
    tail = 300
    n_frames = 150
    picks = np.linspace(0, len(times) - 1, n_frames, dtype=int)
    figure = go.Figure()
    add_flux_surfaces(figure)
    for i in range(n_markers):
        figure.add_trace(
            go.Scatter(
                x=major_radius[:, i],
                y=z[:, i],
                mode="lines",
                line={"color": colors[i], "width": 1},
                opacity=0.2,
                hoverinfo="skip",
                showlegend=False,
                legendgroup=str(i),
            )
        )

    def moving_traces(index):
        traces = []
        for i in range(n_markers):
            start = max(0, index - tail)
            traces.append(
                go.Scatter(
                    x=major_radius[start : index + 1, i],
                    y=z[start : index + 1, i],
                    mode="lines",
                    line={"color": colors[i], "width": 3},
                    name=labels[i],
                    legendgroup=str(i),
                )
            )
            traces.append(
                go.Scatter(
                    x=[major_radius[index, i]],
                    y=[z[index, i]],
                    mode="markers",
                    marker={"color": colors[i], "size": 10, "line": {"color": "white", "width": 1}},
                    showlegend=False,
                    legendgroup=str(i),
                )
            )
        return traces

    first_moving = len(figure.data)
    for trace in moving_traces(0):
        figure.add_trace(trace)
    moving_indices = list(range(first_moving, len(figure.data)))
    frames = [go.Frame(name=f"{times[index]:.1f}", data=moving_traces(index), traces=moving_indices) for index in picks]
    figure.frames = frames
    figure.update_xaxes(range=[1.9, 4.1], title_text="R [a.u.]", constrain="domain")
    figure.update_yaxes(range=[-1.1, 1.1], title_text="Z [a.u.]", scaleanchor="x", scaleratio=1)
    figure.update_layout(
        title="Guiding-center orbits in a tokamak: poloidal plane",
        template="plotly_white",
        autosize=True,
        margin={"l": 70, "r": 30, "t": 80, "b": 130},
        legend={"x": 1.02, "y": 0.5, "yanchor": "middle"},
        updatemenus=[
            {
                "type": "buttons",
                "showactive": False,
                "x": 0.0,
                "xanchor": "left",
                "y": -0.2,
                "yanchor": "top",
                "buttons": [
                    {
                        "label": "Play",
                        "method": "animate",
                        "args": [None, {"frame": {"duration": 40, "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 image shows the last frame, the end of the run, with all trails.
    save_figure(figure, "guiding-center-orbits", frame=-1, height=750, show=show)

    # One panel per particle: the orbit in the poloidal plane.
    panels = make_subplots(rows=2, cols=4, subplot_titles=labels, horizontal_spacing=0.03, vertical_spacing=0.12)
    for i in range(n_markers):
        row, col = divmod(i, 4)
        add_flux_surfaces(panels, row=row + 1, col=col + 1)
        panels.add_trace(
            go.Scatter(
                x=major_radius[:, i],
                y=z[:, i],
                mode="lines",
                line={"color": colors[i], "width": 2.5},
                showlegend=False,
            ),
            row=row + 1,
            col=col + 1,
        )
        panels.add_trace(
            go.Scatter(
                x=[major_radius[0, i]],
                y=[z[0, i]],
                mode="markers",
                marker={"color": "#222", "size": 7, "symbol": "circle-open", "line": {"width": 2}},
                showlegend=False,
            ),
            row=row + 1,
            col=col + 1,
        )
        axis = i + 1
        panels.update_xaxes(range=[1.9, 4.1], showticklabels=row == 1, row=row + 1, col=col + 1)
        panels.update_yaxes(
            range=[-1.1, 1.1],
            showticklabels=col == 0,
            scaleanchor="x" if axis == 1 else f"x{axis}",
            scaleratio=1,
            row=row + 1,
            col=col + 1,
        )
    panels.update_xaxes(title_text="R [a.u.]", row=2)
    panels.update_yaxes(title_text="Z [a.u.]", col=1)
    panels.update_layout(
        title="Guiding-center orbits: one panel per particle (circle: start)",
        template="plotly_white",
        autosize=True,
        margin={"l": 70, "r": 30, "t": 100, "b": 60},
    )
    panels.update_annotations(font_size=12)

    # The parallel velocity over time: it changes sign for trapped particles.
    velocity = go.Figure()
    for i in range(n_markers):
        velocity.add_scatter(
            x=times,
            y=v_parallel[:, i] / speed,
            mode="lines",
            name=labels[i],
            line={"color": colors[i], "width": 2.5, "dash": "solid" if trapped[i] else "dot"},
        )
    velocity.add_hline(y=0.0, line={"color": "#888", "width": 1})
    velocity.update_layout(
        title="Parallel velocity of the guiding centers",
        xaxis_title="t [a.u.]",
        yaxis_title="v∥ / v",
        template="plotly_white",
        autosize=True,
        margin={"l": 70, "r": 30, "t": 80, "b": 60},
    )

    # The drift of the conserved quantities, on a logarithmic axis.
    floor = 1e-12
    conservation = make_subplots(rows=1, cols=2, subplot_titles=("Energy", "Canonical toroidal momentum"))
    for i in range(n_markers):
        for column, drift in enumerate((energy_drift, momentum_drift), start=1):
            conservation.add_scatter(
                x=times,
                y=np.maximum(drift[:, i], floor),
                mode="lines",
                name=labels[i],
                showlegend=column == 1,
                legendgroup=str(i),
                line={"color": colors[i], "width": 2, "dash": "solid" if trapped[i] else "dot"},
                row=1,
                col=column,
            )
    conservation.update_yaxes(type="log", title_text="relative change", col=1)
    conservation.update_yaxes(type="log", col=2)
    conservation.update_yaxes(exponentformat="power")
    conservation.update_xaxes(title_text="t [a.u.]")
    conservation.update_layout(
        title="Conserved quantities along the orbits",
        template="plotly_white",
        autosize=True,
        margin={"l": 70, "r": 30, "t": 100, "b": 60},
    )

    save_figure(panels, "guiding-center-orbits-panels", show=show)
    save_figure(velocity, "guiding-center-orbits-velocity", show=show)
    save_figure(conservation, "guiding-center-orbits-conservation", show=show)


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