diff --git a/src/test_particle_sim_1d/particles.py b/src/test_particle_sim_1d/particles.py index e69de29..528702a 100644 --- a/src/test_particle_sim_1d/particles.py +++ b/src/test_particle_sim_1d/particles.py @@ -0,0 +1,219 @@ +""" +particles.py + +Defines the Species class: a Structure-of-Arrays (SoA) container that stores +properties of particles belonging to one species (e.g., electrons or ions). + +""" + +# import statements +from __future__ import annotations + +import numpy as np + + +# species class definition +class Species: + """ + A simple Structure-of-Arrays (SoA) representation of a particle species. + + Each physical quantity (position, velocity, etc.) is stored in its own + contiguous NumPy array for efficient vectorized operations. + + Attributes + ---------- + q : float + charge of one particle (Coulombs, C) + m : float + mass of one particle (kilograms, kg) + name : str + species name, "electron", "ion", etc. + dtype : np.dtype + floating-point precision for numeric arrays + capacity : int + allocated length of all arrays (number of available slots) + N : int + current number of active particles + x, vx : np.ndarray + arrays for position (m) and velocity (m/s) (1D simulation) + weight : np.ndarray + particle weights (used for scaling particle contribution to charge density) + alive : np.ndarray + boolean array marking whether each particle is active + """ + + # constructor + def __init__( + self, + q: float, + m: float, + name: str = "electron", + capacity: int = 0, + dtype=np.float64, + ) -> None: + """Initialize empty species container""" + # store physical constants and identifiers + self.q = float(q) # float + self.m = float(m) # float + self.name = name # string + self.dtype = dtype # floating point type + + self.capacity = int(capacity) + self.N = 0 # number of active particles + + # allocate SoA storage (1D arrays for 1D sim); preallocate if capacity > 0 + self.x = np.zeros(self.capacity, dtype=self.dtype) + self.vx = np.zeros(self.capacity, dtype=self.dtype) +self.vy = np.zeros(self.capacity, dtype=self.dtype) +self.vz = np.zeros(self.capacity, dtype=self.dtype) + # allocate auxiliary bookkeeping arrays + self.weight = np.ones(self.capacity, dtype=self.dtype) # weighting factors + self.alive = np.ones(self.capacity, dtype=bool) # true = active + + # public methods + def add_particles(self, z_init: np.ndarray, v_init: np.ndarray) -> None: + """ + Append new particles to this species + + Parameters + ---------- + x_init : np.ndarray + Initial positions of new particles + vx_init : np.ndarray + Initial velocities of new particles + + Notes + ----- + function assumes both arrays are the same length + arrays are appended to the end of the current data + """ + + # number of new particles being added + n_new = len(x_init) + + # check if enough space; if not, allocate more + self._ensure_capacity(self.N + n_new) + + # compute slice indices for where to put new data + start = self.N + end = self.N + n_new + + # copy incoming values into array slices + # np.asarray ensures consistent dtype + self.z[start:end] = np.asarray(z_init, dtype=self.dtype) + self.vx[start:end] = v_init[:, 0] + self.vy[start:end] = v_init[:, 1] + self.vz[start:end] = v_init[:, 2] + + # initialize other arrays for new particles + self.weight[start:end] = 1.0 # default weight + self.alive[start:end] = True # mark as active + + # update particle count + self.N = end + + def initialize_uniform( + self, + n: int, + z_min: float, + z_max: float, + v0: float = 0.0, + seed: int | None = None, + ) -> None: + """ + Create N particles uniformly distributed in space + + Parameters + ---------- + n : int + number of particles to create + x_min, x_max : float + spatial bounds for uniform distribution + v0 : float, optional + initial velocity assigned to all particles (default 0) + seed : int or None, optional + random seed for reproducibility + """ + # initialize random number generator (NumPy's modern API) + rng = np.random.default_rng(seed) + + # generate uniformly spaced random positions + z_new = rng.uniform(z_min, z_max, size=n).astype(self.dtype) + + # all velocities start with the same value v0 + if v0 is None: + v_new = np.zeros((n, 3), dtype=self.dtype) + else: + v_new = np.tile(v0, (n, 1)) + + # reuse add_particles to avoid repeating logic (DRY principle) + self.add_particles(z_new, v_new) + + def __len__(self) -> int: + return self.N + + # internal helpers + def _ensure_capacity(self, needed: int) -> None: + """Expand storage arrays if needed""" + + # guard pattern: early return if no expansion needed + if needed <= self.capacity: + return + + # compute new capacity + new_cap = max(needed, 1, int(self.capacity * 2)) + + # grow all SoA arrays to the same new size + self.z = self._grow(self.z, new_cap) + self.vx = self._grow(self.vx, new_cap) + self.vy = self._grow(self.vy, new_cap) + self.vz = self._grow(self.vz, new_cap) + self.weight = self._grow(self.weight, new_cap, fill=1.0) + self.alive = self._grow(self.alive, new_cap, fill=True, array_dtype=bool) + + # update capacity + self.capacity = new_cap + + @staticmethod + def _grow( + arr: np.ndarray, + new_cap: int, + fill=0, + array_dtype=None, + ) -> np.ndarray: + """ + Create a larger copy of existing array + + Parameters + ---------- + arr : np.ndarray + old array to copy from + new_cap : int + desired total length + fill : scalar + value to fill the new (empty) portion with + array_dtype : np.dtype or None + optionally override dtype (used for bool arrays) + + Returns + ------- + np.ndarray + new array containing the old data plus new filler elements + """ + # pick the right data type + dtype = array_dtype or arr.dtype + + # allocate a new array of desired size + out = np.empty(new_cap, dtype=dtype) + + # copy existing values into the front section + n_old = len(arr) + out[:n_old] = arr + + # fill the remainder with a sensible default value + if dtype is bool: + out[n_old:] = bool(fill) + else: + out[n_old:] = fill + + return out diff --git a/tests/test_species1.py b/tests/test_species1.py new file mode 100644 index 0000000..7d4a9ed --- /dev/null +++ b/tests/test_species1.py @@ -0,0 +1,35 @@ +# tests/test_species.py +from __future__ import annotations + +import numpy as np + +from test_particle_sim_1d.particles import Species + + +def test_species_initialization(): + """Test that a species object can be created and is initially empty""" + s = Species(q=-1.602e-19, m=9.109e-31, name="electron", capacity=10) + + # Basic sanity checks + assert isinstance(s, Species) + assert s.N == 0 + assert s.capacity == 10 + assert np.all(s.z == 0.0) + assert np.all(s.vx == 0.0) + assert np.all(s.vy == 0.0) + assert np.all(s.vz == 0.0) + + +def test_add_particles_increases_count(): + """Test adding new particles updates arrays correctly""" + s = Species(q=-1.602e-19, m=9.109e-31) + z_init = np.array([0.1, 0.2, 0.3]) + v_init = np.array([[1.0, 0.0, 0.0],[0.0, 1.0, 0.0],[0.0, 0.0, 1.0]]) + + s.add_particles(z_init, v_init) + + assert len(s) == 3 + np.testing.assert_allclose(s.x[:3], z_init) + np.testing.assert_allclose(s.vx[:3], v_init[:, 0]) + np.testing.assert_allclose(s.vy[:3], v_init[:, 1]) + np.testing.assert_allclose(s.vz[:3], v_init[:, 2])