tokamak

import numpy as np
import matplotlib.pyplot as plt

"""
    Explicit Euler integrator for guiding-center poloidal motion model

    Model equations integrated (Euler):
      dx/dt = -omega * z
      dz/dt =  omega * x + vdrift

    Input:
      R0       : reference major radius
      x0       : initial x(0) = R(0) - R0
      z0       : initial z(0)
      dt       : time step for explicit Euler
      t_end    : final simulation time

    omega_mode : selects angular frequency model:
                -- "const": omega = omega0 (constant angular frequency)
                -- "var"  : omega = vpar / (R0 + x) = vpar / R
    omega0     : constant angular frequency used if omega_mode == "const"
    vpar0      : initial parallel speed v_par(0)
    vdrift     : constant drift velocity added to dz/dt equation

    mirror_mode: Evolves vpar using a magnetic-mirror force from 2.20 Poli's notes
                -- "True" : dvpar/dt = - mu/m * B0 * R0 * z / (x + R0) ^ 2
                -- "False": vpar     = const
    mu_m       : mu/m = magnetic moment divided by mass
    B0         : magnetic field scale at R0 [T]

    Output:
      t[n]     : time
      z[n]     : poloidal-plane vertical coordinate Z
      R[n]     : major radius, R = R0 + x
      vpar[n]  : B-parallel guiding-center speed

"""

# ----------------------- Explicit Euler integrator  --------------------------------------- #
def euler_solution(R0, x0, z0, dt, t_end, omega_mode = "const", omega0=1.0, vpar0=1.0, vdrift=0.0, mirror = False, mu_m = 1.0, B0 = 1.0):

    N = int(t_end / dt) + 1                                     # number of steps
    t = np.linspace(0.0, (N - 1) * dt, N)                       # time array
    x = np.empty(N); z = np.empty(N); vpar = np.empty(N);       # solution arrays for x and z
    x[0] = x0; z[0] = z0; vpar[0] = vpar0;                      # initial conditions

    r0 = np.hypot(x0, z0)
    alpha = r0/(R0+x0)

    eps = 1e-14

    # system of equations to solve:
    for n in range(N - 1):

        # mirror force: update vpar
        if mirror:
            R = R0 + x[n]
            dvpar_dt = - mu_m * (B0 * R0) * z[n] / (R**2)
            vpar[n+1] = vpar[n] + dt * dvpar_dt
        else:
            vpar[n+1] = vpar[n]

        # drift velocity influence
        if omega_mode == "const":
            omega = omega0
        else:  # omega(R)=vpar/R with R=R0+x
            # omega = vpar[n] / (R0 + x[n])
            r = np.hypot(x[n], z[n])
            omega = (alpha * vpar[n]) / max(r, eps)

        x[n+1] = x[n] - dt * ( omega * z[n])
        z[n+1] = z[n] + dt * ( omega * x[n] + vdrift )

    # return to original variables for better plotting
    R = R0 + x
    return t, R, z, vpar


# ----------------------- Analytical solution (omega,vd are constant) ----------------------- #

def analytic_solution(R0, omega, vd, x0, z0, t):
    # Shift to remove constant vd term as from 1.3
    x_sh0 = x0 + vd / omega

    # constant amplitude/phase for the solution
    A = np.sqrt(x_sh0**2 + z0**2)
    phi = np.arctan2(z0, x_sh0)

    # solution
    x = A * np.cos(omega * t + phi) - vd / omega
    z = A * np.sin(omega * t + phi)

    # return to original variables
    R = R0 + x
    return R, z


# ----------------------- Main function: conditions, plotting -------------------------------- #


# ----------------------- A) No drift - compare Euler + Analytical solution ------------------ #
R0 = 1.0
vpar = 1.0
x0, z0 = 1.0, 0.0
vd0 = 0.0

omega0 = vpar/R0
T0 = 2*np.pi/omega0
t_end = 10*T0
dt_list = [T0/100, T0/500, T0/2000]

# Plot in (x,z)
plt.figure()
plt.title("Orbit in (x,z) for different dt")
plt.xlabel("x = R - R0")
plt.ylabel("z")
plt.axis("equal")

# analytical solution
t_an = np.linspace(0.0, t_end, 5000)
R_an, z_an = analytic_solution(R0, omega0, vd0, x0, z0, t_an)
plt.plot(R_an - R0, z_an, label="analytic")

# Euler solutions for each timestep dt

for dt in dt_list:
    t, R_eu, z_eu, vpar_eu = euler_solution(R0, x0, z0, dt, t_end, "const", omega0, vpar, vd0, mirror = False)
    plt.plot(R_eu - R0, z_eu, label=f"Euler dt={dt:.3g}")

plt.legend()
plt.show()

# ----------------------- B) Drift scan: different v_drift --------------------------------- #
x0, z0 = 0.1, 0.0
dt = T0/20000
vdrift_list = [0.01, 0.05]

plt.figure()
plt.title(f"Drift scan (dt={dt:.3g}): omega=vpar/(R0+x)")
plt.xlabel("x = R - R0")
plt.ylabel("z")
plt.axis("equal")

for vdrift in vdrift_list:
    t, R_eu, z_eu, vpar_eu = euler_solution(R0, x0, z0, dt, t_end, omega_mode="var", vpar0=vpar, vdrift=vdrift, mirror = False)
    plt.plot(R_eu - R0, z_eu, label=f"v_drift={vdrift}")

plt.legend()
plt.show()


# ----------------------- C) Mirror force  ------------------------------------------------- #
x0, z0 = 0.3, 0.0
vpar0 = 1.0
vdrift = 0.1
dt = T0/20000
t_end = 10*T0

mu_m = 2.0
B0 = 1.0

t, R_eu, z_eu, vpar_eu= euler_solution(R0, x0, z0, dt, t_end, omega_mode="var", vpar0=vpar0, vdrift=vdrift, mirror=True, mu_m=mu_m, B0=B0)

plt.figure()
plt.title("Mirror force induced trapped-like motion")
plt.xlabel("x = R - R0")
plt.ylabel("z")
plt.axis("equal")
plt.plot(R_eu - R0, z_eu)
plt.show()

## let's plot vpar changing with time to see the trapping itself
plt.figure()
plt.plot(t, vpar_eu)
plt.title("v_par(t)")
plt.xlabel("t"); plt.ylabel("v_par")
plt.show()