Repository: PlasmaPy/PlasmaPy

Issue:
Add a particle stepper for given external fields
As discussed on HipChat recently, I'd like to give adapting my particle stepper a try. The way I see this working would be:

- [x] Add a `ParticleSpecies` class containing their basic information such as position and velocity (as numpy `ndarray`s), charge and mass. This seems like a minimal necessary set of information. I would also like to save historical data on their positions and velocities every few timesteps into a larger ndarray to avoid frequent file IO (the iteration step between saving could be given by user, according to the resolution required).
- [x] Give them a `push` method that takes electric and magnetic fields (`ndarray`s) at particle locations (so they're the same shape) and updates their positions and velocities. I have two Boris pushers, one nonrelativistic and one relativistic - I can port both. Of course, we want velocities at half timesteps to help with energy conservation - adding another method `init_push` (without updating positions) to start that is trivial.
- [x] Add functionality for interpolating fields from - initially - Cartesian grids to particle positions. I've done this in 1D. Given that we want 3D, I'll try to find an algorithm that implements that - I have a few papers with those lying around.
- [ ] I have to say that I have absolutely no idea on how to do this in non-cartesian geometries, so that's something I'd like to postpone until later. As far as I can tell our current idea for the `Plasma` class is to stay Cartesian for a bit, then branch out a little, so that seems reasonable to me.

I should be able to start working on this by about June 7th or so. Any help brainstorming before I get cracking on this would be much appreciated!


Discussion:
This sounds wonderful!  With regards to interpolation, the method I'm most familiar with is trilinear interpolation (e.g., [Haynes & Parnell, 2007](http://aip.scitation.org/doi/10.1063/1.2756751)).  I suspect that some of the interpolation routines in SciPy would do pretty much what we need them to do (e.g., [RegularGridInterpolator](https://docs.scipy.org/doc/scipy-0.19.0/reference/generated/scipy.interpolate.RegularGridInterpolator.html)).

Another thought related to interpolation is that we should allow for flexibility in terms of the interpolation method which we'll eventually need.  If we want to push test particles in a spectral or spectral element simulation, the most accurate way to interpolate would be to use the basis functions of the solution directly.  However, we do not need to do this yet.

A simple test would be to match against analytical results for a particle orbit in a uniform magnetic field.  Another possibility for a test would be to match particle drifts in simple configurations.  

Thank you!
... so this was implemented in scipy all along?

Well.

I really do need to read the full docs of that, scipy keeps positively surprising me.

Thank you very much for the heads up, this will definitely come in handy!

I'll try to make the interpolation easy to swap out.

Uniform mag fields are definitely one of the tests I'm considering; another would be a harmonic potential. The drifts are a neat idea. Thanks! 
Okay, a minor point of confusion. I started doing some minor work on porting my code, and I'm noticing that we don't actually have an electric field defined in the plasma class. I know, quasineutrality, everything runs on magnetic in MHD scenarios, but in this case this means I'm kind of at a loss on how to proceed for now. I can't really interpolate a non-existant electric field.

I had the idea of interpolating not the fields, but the actual momentum - scale that to the particle mass and we get a velocity for the particle. But that seems like a hack, like something you would implement in a video game where you don't actually care about the physics involved.
Hm...some options would be to have the electric field either inputted directly, or calculated from an Ohm's law.  @SolarDrew, do you have thoughts on how to address this?
Hey, just wanted to give a heads up that I still intend to get this done, but other commitments are kind of getting in the way right now. I hope I can get back to this soon.

Resolving patch:
diff --git a/plasmapy/classes/__init__.py b/plasmapy/classes/__init__.py
--- a/plasmapy/classes/__init__.py
+++ b/plasmapy/classes/__init__.py
@@ -1 +1,2 @@
 from .plasma import Plasma
+from .species import Species
diff --git a/plasmapy/classes/plasma.py b/plasmapy/classes/plasma.py
--- a/plasmapy/classes/plasma.py
+++ b/plasmapy/classes/plasma.py
@@ -53,8 +53,8 @@ def __init__(self, domain_x, domain_y, domain_z):
         self.y = domain_y
         self.z = domain_z
 
-        x, y, z = self.x.si.value, self.y.si.value, self.z.si.value
-        self.grid = np.meshgrid(x, y, z, indexing='ij') * u.m
+        self.grid = np.array(np.meshgrid(self.x, self.y, self.z,
+                                indexing='ij'))
         self.domain_shape = (len(self.x), len(self.y), len(self.z))
 
         # Initiate core plasma variables
@@ -62,6 +62,7 @@ def __init__(self, domain_x, domain_y, domain_z):
         self.momentum = np.zeros((3, *self.domain_shape)) * u.kg / (u.m**2 * u.s)
         self.pressure = np.zeros(self.domain_shape) * u.Pa
         self.magnetic_field = np.zeros((3, *self.domain_shape)) * u.T
+        self.electric_field = np.zeros((3, *self.domain_shape)) * u.V / u.m
 
     @property
     def velocity(self):
@@ -72,8 +73,13 @@ def magnetic_field_strength(self):
         B = self.magnetic_field
         return np.sqrt(np.sum(B * B, axis=0))
 
+    @property
+    def electric_field_strength(self):
+        E = self.electric_field
+        return np.sqrt(np.sum(E * E, axis=0))
+
     @property
     def alfven_speed(self):
         B = self.magnetic_field
         rho = self.density
-        return np.sqrt(np.sum(B * B, axis=0) / (mu0*rho))
+        return np.sqrt(np.sum(B * B, axis=0) / (mu0*rho))
\ No newline at end of file
diff --git a/plasmapy/classes/species.py b/plasmapy/classes/species.py
new file mode 100644
--- /dev/null
+++ b/plasmapy/classes/species.py
@@ -0,0 +1,257 @@
+"""
+Class representing a group of particles"""
+# coding=utf-8
+import numpy as np
+from scipy.interpolate import RegularGridInterpolator
+from plasmapy.constants import atomic
+from astropy import constants
+from astropy import units as u
+
+
+class Species:
+    """
+    Object representing a species of particles: ions, electrons, or simply
+    a group of particles with a particular initial velocity distribution.
+
+    Parameters
+    ----------
+    plasma : `Plasma`
+        plasma from which fields can be pulled
+    type : str
+        particle type. See `plasmapy.constants.atomic` for suitable arguments.
+        The default is a proton.
+    n_particles : int
+        number of macroparticles. The default is a single particle.
+    scaling : float
+        number of particles represented by each macroparticle.
+        The default is 1, which means a 1:1 correspondence between particles
+        and macroparticles.
+    dt : `astropy.units.Quantity`
+        length of timestep
+    nt : int
+        number of timesteps
+
+    Attributes
+    ----------
+    x : `astropy.units.Quantity`
+    v : `astropy.units.Quantity`
+        Current position and velocity, respectively. Shape (n, 3).
+    position_history : `astropy.units.Quantity`
+    velocity_history : `astropy.units.Quantity`
+        History of position and velocity. Shape (nt, n, 3).
+    q : `astropy.units.Quantity`
+    m : `astropy.units.Quantity`
+        Charge and mass of particle.
+    eff_q : `astropy.units.Quantity`
+    eff_m : `astropy.units.Quantity`
+        Total charge and mass of macroparticle.
+    kinetic_energy
+        calculated from `v`, as in, current velocity.
+    kinetic_energy_history
+        calculated from `velocity_history`.
+
+    Examples
+    ----------
+    See `plasmapy/examples/particle-stepper.ipynb.`
+
+    """
+    @u.quantity_input(dt=u.s)
+    def __init__(self, plasma, particle_type='p', n_particles=1, scaling=1,
+                 dt=np.inf * u.s, nt=np.inf):
+
+        if np.isinf(dt) and np.isinf(nt):
+            raise ValueError("Both dt and nt are infinite.")
+
+        self.q = atomic.charge_state(particle_type) * constants.e.si
+        self.m = atomic.ion_mass(particle_type)
+        self.N = int(n_particles)
+        self.scaling = scaling
+        self.eff_q = self.q * scaling
+        self.eff_m = self.m * scaling
+
+        self.plasma = plasma
+
+        self.dt = dt
+        self.NT = int(nt)
+        self.t = np.arange(nt) * dt
+
+        self.x = np.zeros((n_particles, 3), dtype=float) * u.m
+        self.v = np.zeros((n_particles, 3), dtype=float) * (u.m / u.s)
+        self.name = particle_type
+
+        self.position_history = np.zeros((self.NT, *self.x.shape),
+                                         dtype=float) * u.m
+        self.velocity_history = np.zeros((self.NT, *self.v.shape),
+                                         dtype=float) * (u.m / u.s)
+        self._B_interpolator = RegularGridInterpolator(
+            (self.plasma.x.si.value,
+             self.plasma.y.si.value,
+             self.plasma.z.si.value),
+            self.plasma.magnetic_field.T.si.value,
+            method="linear",
+            bounds_error=True)
+
+        self._E_interpolator = RegularGridInterpolator(
+            (self.plasma.x.si.value,
+             self.plasma.y.si.value,
+             self.plasma.z.si.value),
+            self.plasma.electric_field.T.si.value,
+            method="linear",
+            bounds_error=True)
+
+    def _interpolate_fields(self):
+        interpolated_b = self._B_interpolator(self.x.si.value) * u.T
+        interpolated_e = self._E_interpolator(self.x.si.value) * u.V / u.m
+        return interpolated_b, interpolated_e
+
+    @property
+    def kinetic_energy(self):
+        """
+        Calculates the kinetic energy for each macroparticle.
+
+        Returns
+        --------
+        `astropy.units.Quantity`
+            Array of kinetic energies, shape (n,).
+        """
+        return (self.v ** 2).sum(axis=-1) * self.eff_m / 2
+
+    @property
+    def kinetic_energy_history(self):
+        """
+        Calculates the kinetic energy history for each particle.
+
+        Returns
+        --------
+        `astropy.units.Quantity`
+            Array of kinetic energies, shape (nt, n).
+        """
+        return (self.velocity_history ** 2).sum(axis=-1) * self.eff_m / 2
+
+    def boris_push(self, init=False):
+        """
+        Implements the Boris algorithm for moving particles and updating their
+        velocities.
+
+        Arguments
+        ----------
+        init: bool (optional)
+            If `True`, does not change the particle positions and sets dt
+            to -dt/2.
+
+        Notes
+        ----------
+        The Boris algorithm is the standard energy conserving algorithm for
+        particle movement in plasma physics.
+
+        Conceptually, the algorithm has three phases:
+
+        1. Add half the impulse from electric field
+        2. Rotate the particle velocity about the direction of the magnetic
+        field
+        3. Add the second half of the impulse from the electric field
+
+        This ends up causing the magnetic field action to be properly
+        "centered" in time, and the algorithm conserves energy.
+
+        References
+        ----------
+        .. [1] C. K. Birdsall, A. B. Langdon, "Plasma Physics via Computer
+        Simulation", 2004, p. 58-63
+        """
+        dt = -self.dt / 2 if init else self.dt
+        b, e = self._interpolate_fields()
+
+        # add first half of electric impulse
+        vminus = self.v + self.eff_q * e / self.eff_m * dt * 0.5
+
+        # rotate to add magnetic field
+        t = -b * self.eff_q / self.eff_m * dt * 0.5
+        s = 2 * t / (1 + (t * t).sum(axis=1, keepdims=True))
+        vprime = vminus + np.cross(vminus.si.value, t) * u.m / u.s
+        vplus = vminus + np.cross(vprime.si.value, s) * u.m / u.s
+
+        # add second half of electric impulse
+        v_new = vplus + self.eff_q * e / self.eff_m * dt * 0.5
+
+        self.v = v_new
+        if not init:
+            self.x += self.v * dt
+
+    def run(self):
+        """
+        Runs a simulation instance.
+        """
+        self.boris_push(init=True)
+        self.position_history[0] = self.x
+        self.velocity_history[0] = self.v
+        for i in range(1, self.NT):
+            self.boris_push()
+            self.position_history[i] = self.x
+            self.velocity_history[i] = self.v
+
+    def __repr__(self, *args, **kwargs):
+        return f"Species(q={self.q:.4e},m={self.m:.4e},N={self.N}," \
+               f"name=\"{self.name}\",NT={self.NT})"
+
+    def __str__(self): # coveralls: ignore
+        return f"{self.N} {self.scaling:.2e}-{self.name} with " \
+               f"q = {self.q:.2e}, m = {self.m:.2e}, " \
+               f"{self.saved_iterations} saved history " \
+               f"steps over {self.NT} iterations"
+
+    def plot_trajectories(self):  # coveralls: ignore
+        """ Draws trajectory history."""
+        from astropy.visualization import quantity_support
+        import matplotlib.pyplot as plt
+        from mpl_toolkits.mplot3d import Axes3D
+
+        quantity_support()
+        fig = plt.figure()
+        ax = fig.add_subplot(111, projection='3d')
+        for p_index in range(self.N):
+            r = self.position_history[:, p_index]
+            x, y, z = r.T
+            ax.plot(x, y, z)
+        ax.set_title(self.name)
+        ax.set_xlabel("$x$ position")
+        ax.set_ylabel("$y$ position")
+        ax.set_zlabel("$z$ position")
+        plt.show()
+
+    def plot_time_trajectories(self, plot="xyz"):  # coveralls: ignore
+        """
+        Draws position history versus time.
+
+        Arguments
+        ----------
+        plot: str (optional)
+            Enable plotting of position component x, y, z for each of these
+            letters included in `plot`.
+        """
+        from astropy.visualization import quantity_support
+        import matplotlib.pyplot as plt
+        from mpl_toolkits.mplot3d import Axes3D
+
+        quantity_support()
+        fig, ax = plt.subplots()
+        for p_index in range(self.N):
+            r = self.position_history[:, p_index]
+            x, y, z = r.T
+            if "x" in plot:
+                ax.plot(self.t, x, label=f"x_{p_index}")
+            if "y" in plot:
+                ax.plot(self.t, y, label=f"y_{p_index}")
+            if "z" in plot:
+                ax.plot(self.t, z, label=f"z_{p_index}")
+        ax.set_title(self.name)
+        ax.legend(loc='best')
+        ax.grid()
+        plt.show()
+
+    def test_kinetic_energy(self):
+        """ Test conservation of kinetic energy. """
+        assert np.allclose(self.kinetic_energy_history,
+                           self.kinetic_energy_history.mean(),
+                           atol=3 * self.kinetic_energy_history.std()), \
+            "Kinetic energy is not conserved!"


Test patch:
diff --git a/plasmapy/classes/tests/test_Plasma.py b/plasmapy/classes/tests/test_Plasma.py
--- a/plasmapy/classes/tests/test_Plasma.py
+++ b/plasmapy/classes/tests/test_Plasma.py
@@ -54,6 +54,8 @@ def test_Plasma_setup(grid_dimensions, expected_size):
     assert test_plasma.magnetic_field.size == 3 * expected_size
     assert test_plasma.magnetic_field.si.unit == u.T
 
+    assert test_plasma.electric_field.size == 3 * expected_size
+    assert test_plasma.electric_field.si.unit == u.V / u.m
 
 # @pytest.mark.parametrize([()])
 def test_Plasma_derived_vars():
@@ -67,10 +69,11 @@ def test_Plasma_derived_vars():
                                 domain_z=np.linspace(0, 1, 1)*u.m)
 
     # Set an arbitrary uniform values throughout the plasma
-    test_plasma.density = 2.0 * np.ones((8, 8)) * u.kg / u.m**3
-    test_plasma.momentum = 10.0 * np.ones((3, 8, 8)) * u.kg / (u.m**2 * u.s)
-    test_plasma.pressure = np.ones((8, 8)) * u.Pa
-    test_plasma.magnetic_field = 0.01 * np.ones((3, 8, 8)) * u.T
+    test_plasma.density[...] = 2.0 *  u.kg / u.m**3
+    test_plasma.momentum[...] = 10.0 * u.kg / (u.m**2 * u.s)
+    test_plasma.pressure[...] = 1 * u.Pa
+    test_plasma.magnetic_field[...] = 0.01 * u.T
+    test_plasma.electric_field[...] = 0.01 * u.V / u.m
 
     # Test derived variable units and shapes
     assert test_plasma.velocity.shape == test_plasma.momentum.shape
@@ -80,6 +83,9 @@ def test_Plasma_derived_vars():
     assert test_plasma.magnetic_field_strength.si.unit == u.T
     assert np.allclose(test_plasma.magnetic_field_strength.value, 0.017320508)
 
+    assert test_plasma.electric_field_strength.shape == test_plasma.electric_field.shape[1:]
+    assert test_plasma.electric_field_strength.si.unit == u.V/u.m
+
     assert test_plasma.alfven_speed.shape == test_plasma.density.shape
     assert test_plasma.alfven_speed.unit.si == u.m / u.s
     assert np.allclose(test_plasma.alfven_speed.value, 10.92548431)
diff --git a/plasmapy/classes/tests/test_Species.py b/plasmapy/classes/tests/test_Species.py
new file mode 100644
--- /dev/null
+++ b/plasmapy/classes/tests/test_Species.py
@@ -0,0 +1,191 @@
+import numpy as np
+import pytest
+from astropy import units as u
+from astropy.modeling import models, fitting
+from scipy.optimize import curve_fit
+
+from plasmapy.classes import Plasma, Species
+
+
+@pytest.fixture()
+def uniform_magnetic_field(N=3, max_x=1):
+    x = np.linspace(-max_x, max_x, N) * u.m
+    test_plasma = Plasma(x, x, x)
+    magfieldstr = 1 * u.T
+    test_plasma.magnetic_field[2] = magfieldstr
+    return test_plasma
+
+
+def test_basic_species_functionality():
+    plasma = uniform_magnetic_field()
+
+    s = Species(plasma=plasma, dt=1e-3 * u.s, nt=1)
+    assert np.isclose(s.kinetic_energy, 0 * u.J, atol=1e-4 * u.J)
+
+    # this should crash as neither `dt` nor `NT` are not provided
+    with pytest.raises(ValueError):
+        Species(plasma=plasma)
+
+
+def fit_sine_curve(position, t, expected_gyrofrequency, phase=0):
+    def sine(t, amplitude, omega, phase, mean):
+        return amplitude * np.sin(omega * t + phase) + mean
+
+    mean = position.mean().si.value
+    amplitude = 3 * position.std().si.value
+    omega = expected_gyrofrequency.si.value
+    params, covariances = curve_fit(sine, position.si.value, t.si.value,
+                                    p0=(amplitude, omega, phase, mean))
+    stds = np.sqrt(np.diag(covariances))
+    return params, stds
+
+
+def test_particle_uniform_magnetic():
+    """
+        Tests the particle stepper for a uniform magnetic field motion.
+    """
+    test_plasma = uniform_magnetic_field()
+
+    particle_type = 'N-14++'
+    s = Species(test_plasma, particle_type=particle_type, dt=1e-2 * u.s,
+                nt=int(1e4))
+
+    perp_speed = 0.01 * u.m / u.s
+    parallel_speed = 1e-5 * u.m / u.s
+    mean_B = test_plasma.magnetic_field_strength.mean()
+    expected_gyrofrequency = (s.q * mean_B / s.m).to(1 / u.s)
+    expected_gyroradius = perp_speed / expected_gyrofrequency
+    expected_gyroperiod = 2 * np.pi / expected_gyrofrequency
+
+    dt = expected_gyroperiod / 100
+
+    s = Species(test_plasma, particle_type=particle_type, dt=dt, nt=int(1e4))
+    s.v[:, 1] = perp_speed
+
+    s.v[:, 2] = parallel_speed
+    s.run()
+
+    x = s.position_history[:, 0, 0]
+    z = s.position_history[:, 0, 2]
+
+    try:
+        params, stds = fit_sine_curve(x, s.t, expected_gyrofrequency)
+    except RuntimeError as e:
+        print(s)
+        raise e
+    estimated_gyroradius = np.abs(params[0]) * u.m
+    estimated_gyroradius_std = np.abs(stds[0]) * u.m
+    estimated_gyrofrequency = np.abs(params[1]) / u.s
+    estimated_gyrofrequency_std = np.abs(stds[1]) / u.s
+
+    assert np.isclose(expected_gyroradius, estimated_gyroradius,
+                      atol=estimated_gyroradius_std), \
+        "Gyroradii don't match!"
+
+    assert np.isclose(expected_gyrofrequency, estimated_gyrofrequency,
+                      atol=estimated_gyrofrequency_std), \
+        "Gyrofrequencies don't match!"
+
+    p_init = models.Polynomial1D(degree=1)
+    fit_p = fitting.LinearLSQFitter()
+    p = fit_p(p_init, s.t, z)
+
+    assert np.allclose(z, p(s.t), atol=1e-4 * u.m), \
+        "z-velocity doesn't stay constant!"
+
+    s.test_kinetic_energy()
+    # s.plot_trajectories()
+
+
+def test_particle_exb_drift():
+    """
+        Tests the particle stepper for a field with magnetic field in the Z
+        direction, electric field in the y direction. This should produce a
+        drift in the negative X direction, with the drift velocity
+
+        v_e = ExB / B^2
+
+        which is independent of ion charge.
+    """
+    test_plasma = uniform_magnetic_field()
+    test_plasma.electric_field[1] = 1 * u.V / u.m
+    expected_drift_velocity = -(test_plasma.electric_field_strength /
+                                test_plasma.magnetic_field_strength).mean() \
+        .to(u.m / u.s)
+
+    s = Species(test_plasma, 'p', 5, dt=1e-10 * u.s, nt=int(1e4))
+    s.v[:, 2] += np.random.normal(size=s.N) * u.m / u.s
+
+    s.run()
+
+    p_init = models.Polynomial1D(degree=1)
+    for x in s.position_history[:, :, 0].T:
+        fit_p = fitting.LinearLSQFitter()
+        p = fit_p(p_init, s.t, x)
+        fit_velocity = p.parameters[1] * u.m / u.s
+
+        assert np.allclose(x, p(s.t), atol=1e-3 * u.m), \
+            "x position doesn't follow linear fit!"
+
+        assert np.isclose(expected_drift_velocity, fit_velocity,
+                          atol=1e-3 * u.m / u.s), \
+            "x velocity doesn't agree with expected drift velocity!"
+
+    s.test_kinetic_energy()
+
+    # s.plot_trajectories()
+
+
+""" TODO: figure out how to get this test to work
+def test_particle_exb_nonuniform_drift():
+        Tests the particle stepper for a field with magnetic field in the Z
+        direction, electric field in the y direction. This should produce a
+        drift in the negative X direction, with the drift velocity
+
+        v_e = ExB / B^2
+
+        which is independent of ion charge.
+    test_plasma = uniform_magnetic_field(10, 1e-3)
+    c1 = 10 * u.V / u.m**3
+    test_plasma.electric_field[1] = c1 * test_plasma.x**2 / 2
+    perp_speed = 1e3 * u.m / u.s
+    parallel_speed = 1e-5 * u.m / u.s
+
+
+    s = Species(test_plasma, particle_type='p', dt=1e-10 * u.s, nt=int(1e4))
+    s.v[:, 1] = parallel_speed
+    s.v[:, 2] = perp_speed
+
+    expected_gyrofrequency = (s.q * test_plasma.magnetic_field_strength.mean()
+                              / s.m).to(1 / u.s)
+    expected_gyroradius = perp_speed / expected_gyrofrequency
+    s.x[:, 0] = expected_gyroradius / 2
+
+    expected_drift_velocity = -(test_plasma.electric_field_strength /
+                                test_plasma.magnetic_field_strength).to(u.m/u.s)
+    expected_drift_velocity += 0.25 * expected_gyroradius**2 * c1 / u.T
+    expected_drift_velocity = expected_drift_velocity.mean().to(u.m / u.s)
+
+    s.run()
+
+    x = s.position_history[:, 0, 0]
+
+    # s.plot_time_trajectories("x")
+
+    # p_init = models.Polynomial1D(degree=1)
+    # fit_p = fitting.LinearLSQFitter()
+    # p = fit_p(p_init, s.t, x)
+    # fit_velocity = p.parameters[1] * u.m / u.s
+    #
+    # # s.plot_trajectories()
+
+    # assert np.allclose(x, p(s.t), atol=1e-3 * u.m), \
+    #     "x position doesn't follow linear fit!"
+    #
+
+    # assert np.isclose(expected_drift_velocity, fit_velocity,
+    #                   atol=1e-3 * u.m / u.s), \
+    #     "x velocity doesn't agree with expected drift velocity!"
+    #
+    # s.test_kinetic_energy()
+"""