Merge pull request #7 from budjensen/lee/fields

Add fields.py and new unit tests

Author: lkadz

src/test_particle_sim_1d/fields.py

@@ -0,0 +1,115 @@
+"""
+fields.py
+
+Defines analytic electric and magnetic field configurations for a
+1D-in-space, 3D-in-velocity (1D3V) plasma simulation.
+
+Each field function takes an array of particle positions `z` and returns
+a NumPy array of shape (N, 3), representing the (Ex, Ey, Ez) or (Bx, By, Bz)
+components at each position.
+
+Functions
+---------
+E_uniform : Return a uniform electric field of specified magnitude and direction.
+B_uniform : Return a uniform magnetic field of specified magnitude and direction.
+B_mirror  : Return a magnetic mirror field with strength increasing ∝ z².
+
+Examples
+--------
+>>> import numpy as np
+>>> from test_particle_sim_1d import fields
+>>> z = np.linspace(-0.1, 0.1, 5)
+>>> fields.E_uniform(z, E0=1.0, direction="x")
+array([[1., 0., 0.],
+       [1., 0., 0.],
+       [1., 0., 0.],
+       [1., 0., 0.],
+       [1., 0., 0.]])
+>>> fields.B_mirror(z, B0=1.0, z_mirror=0.05)[:3]
+array([[0., 0., 5.],
+       [0., 0., 2.],
+       [0., 0., 1.]])
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+# ----------------------------------------------------------------------
+# Electric field definitions
+# ----------------------------------------------------------------------
+
+
+def E_uniform(z: np.ndarray, E0: float = 0.0, direction: str = "z") -> np.ndarray:
+    """
+    Uniform electric field of magnitude E0 along a specified axis.
+
+    Parameters
+    ----------
+    z : np.ndarray
+        Particle positions (ignored for uniform field)
+    E0 : float, optional
+        Field strength [V/m] (default 0)
+    direction : {'x', 'y', 'z'}, optional
+        Axis of field direction (default 'z')
+
+    Returns
+    -------
+    np.ndarray
+        Electric field array of shape (len(z), 3)
+    """
+    E = np.zeros((len(z), 3))
+    idx = {"x": 0, "y": 1, "z": 2}[direction]
+    E[:, idx] = E0
+    return E
+
+
+# ----------------------------------------------------------------------
+# Magnetic field definitions
+# ----------------------------------------------------------------------
+
+
+def B_uniform(z: np.ndarray, B0: float = 0.0, direction: str = "z") -> np.ndarray:
+    """
+    Uniform magnetic field of magnitude B0 along a specified axis.
+
+    Parameters
+    ----------
+    z : np.ndarray
+        Particle positions (ignored)
+    B0 : float
+        Magnetic field strength [T]
+    direction : {'x', 'y', 'z'}, optional
+        Direction of the field vector
+
+    Returns
+    -------
+    np.ndarray
+        Magnetic field array of shape (len(z), 3)
+    """
+    B = np.zeros((len(z), 3))
+    idx = {"x": 0, "y": 1, "z": 2}[direction]
+    B[:, idx] = B0
+    return B
+
+
+def B_mirror(z: np.ndarray, B0: float = 1.0, z_mirror: float = 0.05) -> np.ndarray:
+    """
+    Simple magnetic mirror field that increases with z^2 near the ends.
+
+    Parameters
+    ----------
+    z : np.ndarray
+        Positions along z [m]
+    B0 : float
+        Base magnetic field [T]
+    z_mirror : float
+        Mirror length scale [m]
+
+    Returns
+    -------
+    np.ndarray
+        Magnetic field [T] along z
+    """
+    Bz = B0 * (1 + (z / z_mirror) ** 2)
+    return np.column_stack((np.zeros_like(Bz), np.zeros_like(Bz), Bz))

src/test_particle_sim_1d/integrators.py

@@ -18,10 +18,12 @@
 - 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?"
+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
@@ -38,26 +40,70 @@ def boris_1d3v_z(
     B_func,
     dt: float,
     n_steps: int,
+    record_history: bool = False,
 ):
     """
-    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)
+    Advance particles in 1D (z) with 3D velocity using the Boris pusher.
+
+    This algorithm integrates the Lorentz force equation in a time-centered,
+    energy-conserving way. It applies a half-step electric acceleration,
+    a magnetic rotation, and a second half electric acceleration.
+
+    Parameters
+    ----------
+    z : np.ndarray
+        Particle positions along z-axis [m], shape (N,).
+    v : np.ndarray
+        Particle velocities [m/s], shape (N, 3).
+    q : np.ndarray
+        Particle charges [C], shape (N,).
+    m : np.ndarray
+        Particle masses [kg], shape (N,).
+    E_func : callable
+        Function that returns the electric field array for given z positions.
+        Must have signature ``E_func(z) -> np.ndarray`` of shape (N, 3).
+    B_func : callable
+        Function that returns the magnetic field array for given z positions.
+        Must have signature ``B_func(z) -> np.ndarray`` of shape (N, 3).
+    dt : float
+        Time step [s].
+    n_steps : int
+        Number of time steps to advance.
+    record_history : bool, optional
+        If True, store and return full velocity history as a 3D array of
+        shape (n_steps, N, 3). Default is False.
+
+    Returns
+    -------
+    tuple of np.ndarray
+        - z : np.ndarray
+            Final particle positions along z, shape (N,).
+        - v : np.ndarray
+            Final particle velocities, shape (N, 3).
+        - v_hist : np.ndarray, optional
+            Velocity history (n_steps, N, 3), only returned if
+            ``record_history=True``.
+
+    Raises
+    ------
+    ValueError
+        If z or v arrays have inconsistent shapes.
+    TypeError
+        If E_func or B_func do not return arrays of shape (N, 3).
+
+    Notes
+    -----
+    The Boris scheme is widely used in plasma and particle-in-cell simulations
+    due to its excellent long-term energy conservation and numerical stability.
+    It is exact for constant magnetic fields when E=0.
+
     """
     q_over_m = (q / m)[:, None]
 
-    for _ in range(n_steps):
+    if record_history:
+        v_hist = np.zeros((n_steps, len(v), 3), dtype=v.dtype)
+
+    for i in range(n_steps):
         # 1. Fields at current positions
         E = E_func(z)
         B = B_func(z)
@@ -79,4 +125,9 @@ def boris_1d3v_z(
         # 5. Position update using vz
         z += v[:, 2] * dt
 
+        if record_history:
+            v_hist[i] = v
+
+    if record_history:
+        return z, v, v_hist
     return z, v

tests/test_exb_drift.py

@@ -0,0 +1,65 @@
+"""
+tests/test_exb_drift.py
+
+Checks that a charged particle starting from rest in crossed E and B fields
+develops the correct ExB drift velocity over time and that gyration energy
+is conserved in the drift frame.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from test_particle_sim_1d import fields
+from test_particle_sim_1d.integrators import boris_1d3v_z
+from test_particle_sim_1d.particles import Species
+
+
+def test_exb_drift_from_rest():
+    q, m = 1.0, 1.0
+    E0, B0 = 0.5, 1.0
+    v_d_expected = -E0 / B0  # E (+x) x B (+z) = -ŷ
+
+    # Create species (start at rest)
+    sp = Species(q=q, m=m, capacity=1)
+    sp.add_particles(np.array([0.0]), np.array([[0.0, 0.0, 0.0]]))
+
+    # Define fields
+    def E_func(z):
+        return fields.E_uniform(z, E0=E0, direction="x")
+
+    def B_func(z):
+        return fields.B_uniform(z, B0=B0, direction="z")
+
+    # Integration setup
+    omega_c = q * B0 / m
+    T = 2 * np.pi / omega_c
+    dt = T / 200
+    n_steps = int(20 * 2 * np.pi / (omega_c * dt))  # ~20 gyroperiods
+
+    z, v = sp.z.copy(), np.column_stack((sp.vx, sp.vy, sp.vz))
+    _, _, v_hist = boris_1d3v_z(
+        z,
+        v,
+        np.array([q]),
+        np.array([m]),
+        E_func,
+        B_func,
+        dt,
+        n_steps,
+        record_history=True,
+    )
+
+    # Ignore transients (first few gyroperiods)
+    n_transient = n_steps // 2
+    v_hist_steady = v_hist[n_transient:, 0, :]
+
+    # 1. Average drift velocity (in y)
+    avg_vy = np.mean(v_hist_steady[:, 1])
+    np.testing.assert_allclose(avg_vy, v_d_expected, rtol=1e-2)
+
+    # 2. Check conservation of kinetic energy *in the drift frame*
+    v_d = np.array([0.0, v_d_expected, 0.0])
+    v_rel = v_hist_steady - v_d
+    v_mag_rel = np.linalg.norm(v_rel, axis=1)
+    np.testing.assert_allclose(v_mag_rel / v_mag_rel.mean(), 1.0, rtol=1e-3)