Merge pull request #4 from budjensen/lee/integrators

Add 1D3V Boris integrator and gyromotion test

Author: lkadz

src/test_particle_sim_1d/integrators.py

@@ -0,0 +1,82 @@
+"""
+integrators.py
+
+Boris integrator for 1D spatial / 3D velocity (1D3V) particle motion.
+
+Particles move along z but have full 3D velocity.
+The electric and magnetic fields depend only on z.
+
+This module advances particle positions and velocities according to the
+Lorentz force law:
+
+    m dv/dt = q (E + v x B)
+    dz/dt = v_z
+
+The Boris algorithm is a second-order, time-centered scheme that splits the
+electric and magnetic effects into separate half steps:
+- Half electric acceleration
+- Magnetic rotation
+- Second half electric acceleration
+
+References:
+    J. P. Boris, "Relativistic Plasma Simulation—Optimization of a Hybrid Code",
+    Proceedings of the Fourth Conference on Numerical Simulation of Plasmas, 1970.
+    Qin, H. et al. (2013). "Why is Boris algorithm so good?"
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+def boris_1d3v_z(
+    z: np.ndarray,
+    v: np.ndarray,
+    q: np.ndarray,
+    m: np.ndarray,
+    E_func,
+    B_func,
+    dt: float,
+    n_steps: int,
+):
+    """
+    Advance 1D-in-z, 3D-in-velocity particles using the Boris algorithm.
+
+    Args:
+        z (np.ndarray): Particle positions along z [m], shape (N,)
+        v (np.ndarray): Particle velocities [m/s], shape (N, 3)
+        q (np.ndarray): Charges [C], shape (N,)
+        m (np.ndarray): Masses [kg], shape (N,)
+        E_func (callable): E(z) -> (N, 3) electric field [V/m]
+        B_func (callable): B(z) -> (N, 3) magnetic field [T]
+        dt (float): Time step [s]
+        n_steps (int): Number of steps
+
+    Returns:
+        tuple[np.ndarray, np.ndarray]: Final (z, v)
+    """
+    q_over_m = (q / m)[:, None]
+
+    for _ in range(n_steps):
+        # 1. Fields at current positions
+        E = E_func(z)
+        B = B_func(z)
+
+        # 2. Half-step electric acceleration
+        v_minus = v + 0.5 * dt * q_over_m * E
+
+        # 3. Magnetic rotation
+        t = q_over_m * B * (0.5 * dt)
+        t_mag2 = np.sum(t * t, axis=1, keepdims=True)
+        s = 2 * t / (1 + t_mag2)
+
+        v_prime = v_minus + np.cross(v_minus, t)
+        v_plus = v_minus + np.cross(v_prime, s)
+
+        # 4. Second half-step electric acceleration
+        v = v_plus + 0.5 * dt * q_over_m * E
+
+        # 5. Position update using vz
+        z += v[:, 2] * dt
+
+    return z, v

tests/test_gyromotion.py

@@ -0,0 +1,72 @@
+"""
+test_gyromotion.py
+
+Test that the Boris 1D3V integrator produces correct gyromotion in a uniform B field.
+
+Physical expectation:
+---------------------
+In a uniform magnetic field (B = Bz * e_z) and zero electric field,
+a charged particle with initial velocity perpendicular to B
+undergoes circular motion with constant |v|.
+
+The cyclotron frequency is:
+    ω_c = |q| * B / m
+
+After one cyclotron period:
+    T = 2π / ω_c
+the velocity vector should return to its initial direction (within numerical tolerance).
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from test_particle_sim_1d.integrators import boris_1d3v_z
+
+
+def test_gyromotion_conservation():
+    # --- Physical constants ---
+    q = np.array([-1.602e-19])  # electron charge [C]
+    m = np.array([9.109e-31])  # electron mass [kg]
+    B0 = 1.0  # magnetic field [T]
+    E0 = 0.0  # no electric field
+
+    # --- Initial conditions ---
+    z = np.zeros(1)
+    v0 = np.array([[1e6, 0.0, 0.0]])  # velocity purely in x (perpendicular to Bz)
+
+    # --- Field definitions ---
+    def E_field(z):
+        return np.stack(
+            [np.full_like(z, E0), np.full_like(z, 0.0), np.full_like(z, 0.0)], axis=1
+        )
+
+    def B_field(z):
+        return np.stack(
+            [np.full_like(z, 0.0), np.full_like(z, 0.0), np.full_like(z, B0)], axis=1
+        )
+
+    # --- Cyclotron frequency and time step ---
+    omega_c = np.abs(q[0]) * B0 / m[0]  # cyclotron frequency (rad/s)
+    T_c = 2 * np.pi / omega_c  # cyclotron period (s)
+    n_steps = 1000  # number of integration steps to use for one full gyro orbit
+    dt = T_c / n_steps  # time step so that one period is divided into n_steps steps
+
+    # --- Run one full gyro orbit ---
+    _z_final, v_final = boris_1d3v_z(z, v0.copy(), q, m, E_field, B_field, dt, n_steps)
+
+    # --- Expected: velocity magnitude conserved ---
+    v_mag_initial = np.linalg.norm(v0)
+    v_mag_final = np.linalg.norm(v_final)
+    np.testing.assert_allclose(v_mag_final, v_mag_initial, rtol=1e-6)
+
+    # --- Expected: velocity returns to (vx, vy) ≈ (v0, 0) after one full period ---
+    np.testing.assert_allclose(
+        v_final[0, 0], v0[0, 0], rtol=1e-4
+    )  # vx close to initial
+    np.testing.assert_allclose(v_final[0, 1], 0.0, atol=1e-4 * v_mag_initial)  # vy ≈ 0
+    np.testing.assert_allclose(v_final[0, 2], 0.0, atol=1e-12)  # vz stays constant
+
+
+if __name__ == "__main__":
+    test_gyromotion_conservation()