sciencehub-team/tokamak-orbitsPublic

# /// script
# requires-python = ">=3.12"
# dependencies = ["marimo>=0.25.1", "numpy", "matplotlib", "contourpy"]
# ///

import marimo

__generated_with = "0.25.1"
app = marimo.App(width="medium")


@app.cell(hide_code=True)
def _(mo):
    mo.md(
        r"""
        # Will a fusion alpha stay in the plasma?

        Deuterium-tritium fusion releases a 3.5 MeV helium nucleus (an alpha
        particle) with every reaction. A burning plasma heats itself only if
        those alphas stay confined long enough to hand their energy to the fuel.

        In a tokamak a particle follows a magnetic field line around the torus
        while its guiding centre drifts. Three quantities are conserved on that
        path: the energy $E$, the magnetic moment $\mu = m v_\perp^2 / 2B$, and
        the canonical toroidal momentum

        $$P_\varphi = m R\, v_\parallel - Z e\, \psi(r),$$

        where $\psi$ is the poloidal magnetic flux. Every orbit is a curve on
        which all three stay fixed, so it can be drawn without integrating the
        motion. The machine below has the size and field of ITER.
        """
    )
    return


@app.cell
def _():
    import marimo as mo
    import numpy as np
    import matplotlib.pyplot as plt
    import contourpy
    return contourpy, mo, np, plt


@app.cell
def _(np):
    E_CHARGE = 1.602e-19
    AMU = 1.6605e-27
    R0, A, B0 = 6.2, 2.0, 5.3  # major radius (m), minor radius (m), field on axis (T)
    Q0, QA = 1.0, 3.0  # safety factor on axis and at the edge

    PARTICLES = {
        "Alpha (helium-4)": (4.0015 * AMU, 2),
        "Deuteron": (2.0136 * AMU, 1),
        "Triton": (3.0155 * AMU, 1),
    }

    def safety_factor(r):
        return Q0 + (QA - Q0) * (r / A) ** 2

    def poloidal_flux(r):
        """psi(r) = integral of B0 r / q(r) dr for the parabolic q profile."""
        k = (QA - Q0) / A**2
        return B0 / (2 * k) * np.log(1 + k * r**2 / Q0)

    def field(R):
        return B0 * R0 / R

    return A, B0, E_CHARGE, PARTICLES, R0, field, poloidal_flux, safety_factor


@app.cell
def _(A, PARTICLES, mo):
    particle = mo.ui.dropdown(list(PARTICLES), value="Alpha (helium-4)", label="Particle")
    energy = mo.ui.slider(0.05, 3.5, step=0.05, value=3.5, label="Energy (MeV)")
    start = mo.ui.slider(0.1, 0.9, step=0.01, value=0.5, label="Start radius r / a")
    pitch = mo.ui.slider(-1.0, 1.0, step=0.01, value=0.3, label="Pitch v∥ / v at start")
    mo.vstack([mo.hstack([particle, energy], justify="start"), start, pitch])
    return energy, particle, pitch, start


@app.cell
def _(A, E_CHARGE, R0, contourpy, field, np, poloidal_flux):
    GRID_R = np.linspace(R0 - A, R0 + A, 320)
    GRID_Z = np.linspace(-A, A, 320)
    MESH_R, MESH_Z = np.meshgrid(GRID_R, GRID_Z)
    MESH_r = np.hypot(MESH_R - R0, MESH_Z)

    def orbit(mass, charge, energy_mev, r_start, lam):
        """Guiding-centre orbit of a particle launched on the outboard midplane.

        Returns the orbit as an (n, 2) array of (R, Z) and its class:
        "passing", "trapped", or "lost" when it reaches the wall.
        """
        v = np.sqrt(2 * energy_mev * 1e6 * E_CHARGE / mass)
        R_s = R0 + r_start
        lam = np.clip(lam, -0.999, 0.999)
        if abs(lam) < 1e-3:
            lam = 1e-3
        mu = mass * v**2 * (1 - lam**2) / (2 * field(R_s))
        ze = charge * E_CHARGE
        p_phi = mass * R_s * lam * v - ze * poloidal_flux(r_start)
        # v_par from P_phi must agree with v_par from E and mu: the zero set of H.
        v_par = (p_phi + ze * poloidal_flux(MESH_r)) / (mass * MESH_R)
        H = (v_par**2 + 2 * mu * field(MESH_R) / mass - v**2) / v**2
        lines = contourpy.contour_generator(GRID_R, GRID_Z, H).lines(0.0)
        if not lines:
            return np.array([[R_s, 0.0]]), "lost"
        path = min(lines, key=lambda seg: np.min(np.hypot(seg[:, 0] - R_s, seg[:, 1])))
        r_path = np.hypot(path[:, 0] - R0, path[:, 1])
        closed = np.allclose(path[0], path[-1], atol=1e-6)
        if r_path.max() > 0.985 * A or not closed:
            return path, "lost"
        theta = np.unwrap(np.arctan2(path[:, 1], path[:, 0] - R0))
        return path, "passing" if np.ptp(theta) > 1.9 * np.pi else "trapped"

    return (orbit,)


@app.cell
def _(A, PARTICLES, R0, energy, np, orbit, particle, pitch, plt, start):
    mass, charge = PARTICLES[particle.value]
    path, kind = orbit(mass, charge, energy.value, start.value * A, pitch.value)
    width = np.ptp(np.hypot(path[:, 0] - R0, path[:, 1]))

    colour = {"passing": "#0969da", "trapped": "#9a6700", "lost": "#cf222e"}[kind]
    fig, ax = plt.subplots(figsize=(5.4, 6))
    circle = np.linspace(0, 2 * np.pi, 200)
    for rho in (0.25, 0.5, 0.75):
        ax.plot(R0 + rho * A * np.cos(circle), rho * A * np.sin(circle), color="#d0d7de", lw=0.8)
    ax.plot(R0 + A * np.cos(circle), A * np.sin(circle), color="#57606a", lw=1.5)
    ax.plot(*path.T, color=colour, lw=2)
    ax.plot(R0 + start.value * A, 0, "o", color=colour)
    ax.plot(R0, 0, "+", color="#57606a")
    ax.set(aspect="equal", xlabel="R (m)", ylabel="Z (m)")
    ax.set_title(f"{kind.capitalize()} orbit, radial width {width * 100:.0f} cm")
    fig.tight_layout()
    fig
    return charge, kind, mass


@app.cell(hide_code=True)
def _(mo):
    mo.md(
        r"""
        The grey circles are magnetic surfaces at a quarter, half, and three
        quarters of the minor radius; the dark circle is the plasma edge.
        The magnetic field is stronger on the inside of the torus ($B \propto
        1/R$). A particle moving mostly across the field is reflected by that
        stronger field, the way a mirror reflects it, and its drift traces the
        **banana** shape. A particle moving mostly along the field passes
        around the whole torus. A 3.5 MeV alpha has a wide banana. Lower its
        energy, or switch to a deuteron of the same energy, and the orbit
        shrinks towards its magnetic surface.
        """
    )
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md(
        r"""
        ## Which pitches are trapped

        In the limit of thin orbits a particle starting on the outboard midplane
        at inverse aspect ratio $\varepsilon = r/R_0$ is trapped when
        $|v_\parallel / v| < \sqrt{2\varepsilon / (1 + \varepsilon)}$.
        The scan below launches 81 particles from the start radius you chose and
        compares each orbit with that estimate.
        """
    )
    return


@app.cell
def _(A, R0, charge, energy, mass, np, orbit, plt, start):
    pitches = np.linspace(-0.99, 0.99, 81)
    scan = [orbit(mass, charge, energy.value, start.value * A, lam) for lam in pitches]
    widths = np.array([np.ptp(np.hypot(_p[:, 0] - R0, _p[:, 1])) for _p, _ in scan])
    kinds = np.array([_k for _, _k in scan])
    eps = start.value * A / R0
    lam_c = np.sqrt(2 * eps / (1 + eps))

    fig2, ax2 = plt.subplots(figsize=(7, 3.4))
    for _kind, _colour in (("passing", "#0969da"), ("trapped", "#9a6700"), ("lost", "#cf222e")):
        _chosen = kinds == _kind
        ax2.scatter(pitches[_chosen], widths[_chosen] * 100, s=14, color=_colour, label=_kind)
    for edge in (-lam_c, lam_c):
        ax2.axvline(edge, color="#57606a", ls="--", lw=0.8)
    ax2.set(xlabel="pitch v∥ / v at start", ylabel="orbit width (cm)")
    ax2.set_ylim(0, widths.max() * 125)
    ax2.legend(frameon=False, ncol=3, loc="upper center")
    fig2.tight_layout()
    fig2
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md(
        r"""
        ---
        **Model and limits.** Circular, concentric magnetic surfaces; toroidal
        field $B_0 R_0 / R$ with $R_0 = 6.2$ m, $a = 2.0$ m, $B_0 = 5.3$ T;
        safety factor $q(r) = 1 + 2(r/a)^2$; guiding-centre approximation with
        $|B| \approx B_\varphi$. No collisions, no electric field, no ripple
        from the discrete coils, so this shows orbit topology, not loss rates.
        An orbit is counted as lost when it reaches the plasma edge.
        **References.** J. Wesson, *Tokamaks*, 4th ed. (2011), ch. 3;
        ITER Physics Basis, Nucl. Fusion 39 (1999) 2137, ch. 5.
        Fork it to add a realistic plasma shape or a ripple field.
        """
    )
    return


if __name__ == "__main__":
    app.run()