Running Whistler-Electron Test-Particle Simulations

import numpy as np
import WPIT.Environment_mod as env

# --- Location ---
L    = 5.0                          # L-shell
lamda = 0.0                         # magnetic latitude (rad)

# --- Geomagnetic field (dipole model) ---
B0 = env.Bmag_dipole(L, lamda)      # returns field magnitude in T
print(f"B0 = {B0:.3e} T")

# --- Electron density (Sheeley et al. 2001 empirical model) ---
# Returns mean, min, max in cm^-3
ne_mean, ne_min, ne_max = env.density_equ_sheeley(L)
Ne = ne_mean * 1e6                  # convert to m^-3

# --- Characteristic frequencies ---
wce = env.omega_cyclotron(B0, env.const.qe, env.const.me)   # electron gyrofrequency (rad/s)
wpe = env.omega_plasma(Ne, env.const.qe, env.const.me)       # electron plasma frequency (rad/s)
print(f"wce = {wce:.3e} rad/s,  wpe = {wpe:.3e} rad/s")

import WPIT.WaveProperties_mod as wave

# --- Wave frequency and normal angle ---
fwave   = 2000.0                          # wave frequency in Hz
wwave   = 2 * np.pi * fwave               # angular frequency (rad/s)
theta   = np.deg2rad(0.0)                 # wave normal angle (rad)

# Ion densities — set to 0 for a simple electron-only plasma
NH = NHe = NO = 0.0

# Stix cold-plasma parameters (S, D, P, R, L)
S, D, P, R_s, L_s = wave.stix_parameters(wwave, Ne, NH, NHe, NO, B0)

# Refractive index (Appleton–Hartree, parallel propagation)
eta_sq_plus, eta_sq_minus, refr_ind, kappa, kpar, kper = \
    wave.refr_index_appleton(wwave, wpe, wce, theta)

# Wave amplitude components (Bell 1984 polarisation relations)
Byw = 1e-10                              # wave magnetic amplitude (T)
Bxw, Byw, Bzw, Exw, Eyw, Ezw = \
    wave.wave_amplitudes_bell(refr_ind, P, D, S, Byw, theta)

Ekev  = 100.0                          # particle energy in keV
alpha = np.deg2rad(60.0)               # local pitch angle

upar0, uper0, ppar0, pper0, gamma0 = \
    env.initial_velocity(Ekev, alpha, env.const.me)

print(f"ppar0 = {ppar0:.3e} kg·m/s")
print(f"pper0 = {pper0:.3e} kg·m/s")
print(f"gamma0 = {gamma0:.6f}")

import WPIT.WPI_mod.whistler_electron_mod as wpi_mod

m_res = 1   # resonance order (fundamental cyclotron resonance)

gamma, w1, w2, wtau_sq, R1, R2, beta = wpi_mod.wpi_params(
    m_res, ppar0, pper0,
    Bxw, Byw, Exw, Eyw, Ezw,
    kpar, kper, wce
)

# Precompute spatial gradient of B (needed for drift and mirror terms)
dB = env.dB_ds(L, lamda)   # dB/ds along field line

# Example: evaluate all derivatives at t=0 for state
# state = [z, lamda, ppar, pper, Ek, gamma, alpha, aeq, eta]
z0    = 0.0
lamda0 = 0.0
aeq0  = alpha          # equatorial pitch angle (≈ local at equator)
eta0  = 0.0            # initial phase angle

dzdt_val     = wpi_mod.dzdt(ppar0, gamma0, env.const.me)
dlamdadt_val = wpi_mod.dlamdadt(ppar0, lamda0, gamma0, L)
dppardt_val  = wpi_mod.dppardt(pper0, eta0, wtau_sq, kpar, gamma0, wce, dB)
dpperdt_val  = wpi_mod.dpperdt(ppar0, pper0, eta0, w1, w2, beta,
                                gamma0, R1, R2, m_res, wce, dB)
detadt_val   = wpi_mod.detadt(ppar0, m_res, wce, wwave, gamma0, kpar)

from scipy.integrate import solve_ivp

def equations_of_motion(t, state):
    z_, lamda_, ppar_, pper_, Ek_, gam_, alpha_, aeq_, eta_ = state

    # Recompute environment at new position
    B_    = env.Bmag_dipole(L, lamda_)
    wce_  = env.omega_cyclotron(B_, env.const.qe, env.const.me)
    dB_   = env.dB_ds(L, lamda_)

    # Recompute wave parameters at current momentum
    _, w1_, w2_, wtau_sq_, R1_, R2_, beta_ = wpi_mod.wpi_params(
        m_res, ppar_, pper_, Bxw, Byw, Exw, Eyw, Ezw, kpar, kper, wce_
    )

    dz      = wpi_mod.dzdt(ppar_, gam_, env.const.me)
    dlam    = wpi_mod.dlamdadt(ppar_, lamda_, gam_, L)
    dppar   = wpi_mod.dppardt(pper_, eta_, wtau_sq_, kpar, gam_, wce_, dB_)
    dpper   = wpi_mod.dpperdt(ppar_, pper_, eta_, w1_, w2_, beta_,
                               gam_, R1_, R2_, m_res, wce_, dB_)
    dEk     = wpi_mod.dEkdt(pper_, ppar_, eta_, m_res, Exw, Eyw, Ezw, beta_, gam_)
    dgamma  = wpi_mod.dgammadt(pper_, ppar_, eta_, m_res, Exw, Eyw, Ezw, beta_, gam_)
    dalpha  = wpi_mod.dalphadt(pper_, ppar_, eta_, w1_, w2_, R1_, R2_,
                                wtau_sq_, kpar, beta_, m_res, gam_, wce_, dB_)
    daeq    = wpi_mod.daeqdt(ppar_, pper_, alpha_, aeq_, eta_, w1_, R1_,
                              w2_, R2_, gam_, beta_, wtau_sq_, kpar, m_res)
    deta    = wpi_mod.detadt(ppar_, m_res, wce_, wwave, gam_, kpar)

    return [dz, dlam, dppar, dpper, dEk, dgamma, dalpha, daeq, deta]

# Initial state vector
state0 = [z0, lamda0, ppar0, pper0,
          Ekev * 1e3 * env.const.qe,   # energy in Joules
          gamma0, alpha, aeq0, eta0]

# Integration time span
t_span = (0.0, 0.5)       # seconds
t_eval = np.linspace(0, 0.5, 5000)

sol = solve_ivp(equations_of_motion, t_span, state0,
                method='RK45', t_eval=t_eval,
                rtol=1e-8, atol=1e-10)

# Unpack results
z_sol, lamda_sol, ppar_sol, pper_sol = sol.y[0], sol.y[1], sol.y[2], sol.y[3]
Ek_sol, gamma_sol, alpha_sol         = sol.y[4], sol.y[5], sol.y[6]
aeq_sol, eta_sol                     = sol.y[7], sol.y[8]

import matplotlib.pyplot as plt

fig, axes = plt.subplots(2, 2, figsize=(10, 8))

axes[0, 0].plot(sol.t, np.rad2deg(alpha_sol))
axes[0, 0].set_xlabel("Time (s)")
axes[0, 0].set_ylabel("Pitch angle α (deg)")

axes[0, 1].plot(sol.t, Ek_sol / (1e3 * env.const.qe))
axes[0, 1].set_xlabel("Time (s)")
axes[0, 1].set_ylabel("Kinetic energy (keV)")

axes[1, 0].plot(sol.t, np.rad2deg(aeq_sol))
axes[1, 0].set_xlabel("Time (s)")
axes[1, 0].set_ylabel("Equatorial pitch angle α_eq (deg)")

axes[1, 1].plot(sol.t, eta_sol % (2 * np.pi))
axes[1, 1].set_xlabel("Time (s)")
axes[1, 1].set_ylabel("Phase angle η (rad)")

plt.tight_layout()
plt.show()