Repository: PlasmaPy/PlasmaPy

Issue:
Add legends to Langmuir plot examples
The plots [here](http://docs.plasmapy.org/en/stable/auto_examples/plot_langmuir_analysis.html#sphx-glr-auto-examples-plot-langmuir-analysis-py) could use some nice legends.

Discussion:
pep8speaks: Hello @StanczakDominik! Thanks for submitting your pull request.

* In the file [`plasmapy/diagnostics/langmuir.py`](https://github.com/PlasmaPy/PlasmaPy/blob/cb8db4d47784148deccc8e055ca912216faaae60/plasmapy/diagnostics/langmuir.py):

> [Line 370:100](https://github.com/PlasmaPy/PlasmaPy/blob/cb8db4d47784148deccc8e055ca912216faaae60/plasmapy/diagnostics/langmuir.py#L370): [E501](https://duckduckgo.com/?q=pep8%20E501) line too long (102 > 99 characters)
> [Line 379:100](https://github.com/PlasmaPy/PlasmaPy/blob/cb8db4d47784148deccc8e055ca912216faaae60/plasmapy/diagnostics/langmuir.py#L379): [E501](https://duckduckgo.com/?q=pep8%20E501) line too long (103 > 99 characters)



PlasmaPy follows the [PEP8 style guide](https://www.python.org/dev/peps/pep-0008/), which may be checked locally using [pycodestyle](http://pycodestyle.pycqa.org/en/latest/intro.html)
codecov[bot]: # [Codecov](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/623?src=pr&el=h1) Report
> Merging [#623](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/623?src=pr&el=desc) into [master](https://codecov.io/gh/PlasmaPy/PlasmaPy/commit/efd7bff5bf34c64e70259886c9b4111b046f4d79?src=pr&el=desc) will **not change** coverage.
> The diff coverage is `n/a`.

[![Impacted file tree graph](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/623/graphs/tree.svg?width=650&token=LzZ5mnXVad&height=150&src=pr)](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/623?src=pr&el=tree)

```diff
@@           Coverage Diff           @@
##           master     #623   +/-   ##
=======================================
  Coverage   96.49%   96.49%           
=======================================
  Files          51       51           
  Lines        4540     4540           
=======================================
  Hits         4381     4381           
  Misses        159      159
```


| [Impacted Files](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/623?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [plasmapy/diagnostics/langmuir.py](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/623/diff?src=pr&el=tree#diff-cGxhc21hcHkvZGlhZ25vc3RpY3MvbGFuZ211aXIucHk=) | `100% <ø> (ø)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/623?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/623?src=pr&el=footer). Last update [efd7bff...a77eec2](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/623?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/plasmapy/diagnostics/langmuir.py b/plasmapy/diagnostics/langmuir.py
index 5d8fff7b..e02209f1 100644
--- a/plasmapy/diagnostics/langmuir.py
+++ b/plasmapy/diagnostics/langmuir.py
@@ -350,13 +350,15 @@ def swept_probe_analysis(probe_characteristic, probe_area, gas_argument,
             fig, (ax1, ax2) = plt.subplots(2, 1)
             ax1.plot(probe_characteristic.bias,
                      probe_characteristic.current,
-                     marker='.', color='k', linestyle='')
+                     marker='.', color='k', linestyle='',
+                     label="Probe current")
             ax1.set_title("Probe characteristic")
             ax2.set_ylim(probe_characteristic.get_padded_limit(0.1))
 
             ax2.plot(probe_characteristic.bias,
                      np.abs(probe_characteristic.current),
-                     marker='.', color='k', linestyle='')
+                     marker='.', color='k', linestyle='',
+                     label="Probe current")
             ax2.set_title("Logarithmic")
             ax2.set_ylim(probe_characteristic.get_padded_limit(0.1, log=True))
 
@@ -364,8 +366,9 @@ def swept_probe_analysis(probe_characteristic, probe_area, gas_argument,
             ax1.axhline(y=I_es.value, color='grey', linestyle='--')
             ax1.axvline(x=V_F.value, color='k', linestyle='--')
             ax1.axhline(y=I_is.value, color='r', linestyle='--')
-            ax1.plot(ion_current.bias, ion_current.current, c='y')
-            ax1.plot(electron_current.bias, electron_current.current, c='c')
+            ax1.plot(ion_current.bias, ion_current.current, c='y', label="Ion current")
+            ax1.plot(electron_current.bias, electron_current.current,
+                     c='c', label="Electron current")
             tot_current = ion_current + electron_current
             ax1.plot(tot_current.bias, tot_current.current, c='g')
 
@@ -373,11 +376,14 @@ def swept_probe_analysis(probe_characteristic, probe_area, gas_argument,
             ax2.axhline(y=I_es.value, color='grey', linestyle='--')
             ax2.axvline(x=V_F.value, color='k', linestyle='--')
             ax2.axhline(y=np.abs(I_is.value), color='r', linestyle='--')
-            ax2.plot(ion_current.bias, np.abs(ion_current.current), c='y')
+            ax2.plot(ion_current.bias, np.abs(ion_current.current),
+                     label="Ion current", c='y')
             ax2.plot(electron_current.bias, np.abs(electron_current.current),
-                     c='c')
+                     label="Electron current", c='c')
             ax2.plot(tot_current.bias, np.abs(tot_current.current), c='g')
             ax2.set_yscale("log", nonposy='clip')
+            ax1.legend(loc='best')
+            ax2.legend(loc='best')
 
             fig.tight_layout()
 
@@ -869,7 +875,8 @@ def get_electron_temperature(exponential_section, bimaxwellian=False,
             plt.scatter(exponential_section.bias.to(u.V),
                         np.log(exponential_section.current.to(u.A).value),
                         color='k',
-                        marker='.')
+                        marker='.',
+                        label="Exponential section")
 
             if bimaxwellian:
                 plt.scatter(x0, y0, marker='o', c='g')
@@ -878,15 +885,19 @@ def get_electron_temperature(exponential_section, bimaxwellian=False,
                              exponential_section.bias.to(u.V).value,
                              fit[0], fit[1], fit[2] + fit[3]),
                          c='g',
-                         linestyle='--')
+                         linestyle='--',
+                         label="Bimaxwellian exponential section fit",
+                         )
 
             plt.plot(exponential_section.bias.to(u.V),
                      fit_func(exponential_section.bias.to(u.V).value,
                      *fit),
+                     label="Exponential fit",
                      c='g')
 
             plt.ylabel("Logarithmic current")
             plt.title("Exponential fit")
+            plt.legend(loc='best')
             plt.tight_layout()
 
     k = [T_e]
@@ -955,9 +966,12 @@ def extrapolate_electron_current(probe_characteristic, fit,
             plt.figure()
             plt.scatter(probe_characteristic.bias,
                         probe_characteristic.current.to(u.mA), marker='.',
+                        label="Probe characteristic",
                         c='k')
             plt.plot(electron_characteristic.bias,
-                     electron_characteristic.current.to(u.mA))
+                     electron_characteristic.current.to(u.mA),
+                     label="Estimated electron characteristic")
+            plt.legend()
 
     return electron_characteristic
 
@@ -1221,5 +1235,6 @@ def get_EEDF(probe_characteristic, visualize=False):
             plt.title("Electron Energy Distribution Function")
             plt.xlabel("Energy (eV)")
             plt.ylabel("Probability")
+            plt.grid()
 
     return energy, probability
diff --git a/plasmapy/examples/plot_langmuir_analysis.py b/plasmapy/examples/plot_langmuir_analysis.py
index 310e71bf..c509e9ef 100644
--- a/plasmapy/examples/plot_langmuir_analysis.py
+++ b/plasmapy/examples/plot_langmuir_analysis.py
@@ -9,9 +9,11 @@
 """
 
 from plasmapy.diagnostics.langmuir import Characteristic, swept_probe_analysis
+import matplotlib.pyplot as plt
 import astropy.units as u
 import numpy as np
 import os
+from pprint import pprint
 
 ######################################################
 # The first characteristic we analyze is a simple single-probe measurement in
@@ -25,8 +27,8 @@
 bias, current = np.load(path)
 
 # Create the Characteristic object, taking into account the correct units
-characteristic = Characteristic(np.array(bias) * u.V,
-                                np.array(current)*1e3 * u.mA)
+characteristic = Characteristic(u.Quantity(bias,  u.V),
+                                u.Quantity(current, u.A))
 
 # Calculate the cylindrical probe surface area
 probe_length = 1.145 * u.mm
@@ -38,7 +40,7 @@
 # Now we can actually perform the analysis. Since the plasma is in Helium an
 # ion mass number of 4 is entered. The results are visualized and the obtained
 # EEDF is also shown.
-print(swept_probe_analysis(characteristic,
+pprint(swept_probe_analysis(characteristic,
                            probe_area, 'He-4+',
                            visualize=True,
                            plot_EEDF=True))
@@ -55,8 +57,8 @@
 # Load the data from a file and create the Characteristic object
 path = os.path.join("langmuir_samples", "Pace2015.npy")
 bias, current = np.load(path)
-characteristic = Characteristic(np.array(bias) * u.V,
-                                np.array(current) * 1e3 * u.mA)
+characteristic = Characteristic(u.Quantity(bias, u.V),
+                                u.Quantity(current, u.A))
 
 ######################################################
 # Initially the electrons are assumed to be Maxwellian. To check this the fit
@@ -71,7 +73,7 @@
 # It can be seen that this plasma is slightly bi-Maxwellian, as there are two
 # distinct slopes in the exponential section. The analysis is now performed
 # with bimaxwellian set to True, which yields improved results.
-print(swept_probe_analysis(characteristic,
+pprint(swept_probe_analysis(characteristic,
                            0.738 * u.cm**2,
                            'Ar-40 1+',
                            bimaxwellian=True,
@@ -93,8 +95,8 @@
 # Import probe data and calculate probe surface area.
 path = os.path.join("langmuir_samples", "Beckers2017b.npy")
 bias, current = np.load(path)
-characteristic = Characteristic(np.array(bias) * u.V,
-                                np.array(current) * 1e3 * u.mA)
+characteristic = Characteristic(u.Quantity(bias, u.V),
+                                u.Quantity(current, u.A))
 probe_length = 1.145 * u.mm
 probe_diameter = 1.57 * u.mm
 probe_area = (probe_length * np.pi * probe_diameter +
@@ -103,7 +105,7 @@
 ######################################################
 # `plot_electron_fit` is set to True to check the bi-Maxwellian properties.
 # The fit converges nicely to the two slopes of the electron growth region.
-print(swept_probe_analysis(characteristic,
+pprint(swept_probe_analysis(characteristic,
                            probe_area,
                            'He-4+',
                            bimaxwellian=True,
diff --git a/plasmapy/physics/drifts.py b/plasmapy/physics/drifts.py
new file mode 100644
index 00000000..68595955
--- /dev/null
+++ b/plasmapy/physics/drifts.py
@@ -0,0 +1,106 @@
+import astropy.units as u
+from astropy.units.utils import np
+from plasmapy import utils
+
+
+@utils.check_quantity()
+def ExB_drift(E: u.V/u.m, B: u.T) -> u.m/u.s:
+    """
+    Calculate the "electric cross magnetic" particle drift.
+
+    Parameters
+    ----------
+    E : ~astropy.units.Quantity
+        Electric field vector
+    B : ~astropy.units.Quantity
+        Magnetic field vector
+
+    Returns
+    -------
+    v: ~astropy.units.Quantity
+        Drift velocity, in m/s
+
+    Examples
+    --------
+    >>> import astropy.units as u
+    >>> from astropy.constants.si import g0, e, m_e
+    >>> ex = np.array([1, 0, 0])
+    >>> ez = np.array([0, 0, 1])
+    >>> force_drift(-ez*g0*m_e, ex*0.01*u.T, e)
+    <Quantity [ 0.0000000e+00, -5.5756984e-09,  0.0000000e+00] m / s>
+    >>> force_drift(-ez*g0*m_e, ez*0.01*u.T, e)
+    <Quantity [ 0., -0.,  0.] m / s>
+    >>> force_drift(-ez*g0*m_e, ex*u.T, e)
+    <Quantity [ 0.0000000e+00, -5.5756984e-11,  0.0000000e+00] m / s>
+
+    Notes
+    -----
+    The E cross B drift is given by
+
+    .. math::
+
+        \vec{v} = \frac{\vec{E} \times \vec{B}}{|B|^2}
+
+    and is independent of particle charge.
+
+    References
+    ----------
+    - PM Bellan, Fundamentals of Plasma Physics, 3.57
+
+    """
+
+    # np.cross drops units right now, thus this hack: see
+    # https://github.com/PlasmaPy/PlasmaPy/issues/59 
+    cross = np.cross(E.si.value, B.si.value) * E.unit * B.unit
+    return (cross / (B*B).sum(-1)).to(u.m/u.s)
+
+
+@utils.check_quantity()
+def force_drift(F: u.N, B: u.T, q: u.C) -> u.m / u.s:
+    """
+    Calculate the general force drift for a particle in a magnetic field.
+
+    Parameters
+    ----------
+    F : ~astropy.units.Quantity
+        Force acting on particle
+    B : ~astropy.units.Quantity
+        Magnetic field
+    q : ~astropy.units.Quantity
+        Particle charge
+
+    Examples
+    --------
+    >>> import astropy.units as u
+    >>> ex = np.array([1, 0, 0])
+    >>> ey = np.array([0, 1, 0])
+    >>> ExB_drift(ex * u.V/u.m, ey * u.T)
+    <Quantity [0., 0., 1.] m / s>
+    >>> ExB_drift(ex * u.V/u.m, ex * u.T)
+    <Quantity [0., 0., 0.] m / s>
+    >>> ExB_drift(ex * u.V/u.m, 100 * ey * u.T)
+    <Quantity [0.  , 0.  , 0.01] m / s>
+
+    Returns
+    -------
+    v: ~astropy.units.Quantity
+        Drift velocity, in m/s
+
+    Notes
+    -----
+    The particle drift in a magnetic field and with a general
+    force (e.g. gravity) applied to it is given by
+
+    .. math::
+
+        \vec{v} = \frac{\vec{F} \times \vec{B}}{q |B|^2}
+
+    Note the charge dependency.
+
+    References
+    ----------
+    - PM Bellan, Fundamentals of Plasma Physics, 3.58
+
+    """
+    cross = np.cross(F.si.value, B.si.value) * F.unit * B.unit
+    return (cross / (q * (B*B).sum(-1))).to(u.m/u.s)


Test patch:
diff --git a/plasmapy/physics/tests/test_drifts.py b/plasmapy/physics/tests/test_drifts.py
new file mode 100644
index 00000000..af79d677
--- /dev/null
+++ b/plasmapy/physics/tests/test_drifts.py
@@ -0,0 +1,91 @@
+import pytest
+
+from plasmapy.physics import drifts
+import astropy.units as u
+from astropy.tests.helper import assert_quantity_allclose
+
+
+class Test_ExB_drift:
+    def test_E_x_B_1d_arrays(self):
+        E = u.Quantity([1, 0, 0], unit=u.V/u.m)
+        B = u.Quantity([0, 1, 0], unit=u.T)
+        result = drifts.ExB_drift(2*E, 3*B)
+        assert_quantity_allclose(result, (2/3)*u.Quantity([0, 0, 1], u.m/u.s))
+
+    def test_ExB_2d_array(self):
+        E = u.Quantity([[1, 0, 0],
+                        [1, 0, 0],
+                        [1, 0, 0]], unit=u.V/u.m)
+        B = u.Quantity([[0, 1, 0],
+                        [0, 1, 0],
+                        [0, 1, 0]], unit=u.T)
+
+        result = drifts.ExB_drift(2*E, 3*B)
+        assert_quantity_allclose(result, (2/3)*u.Quantity([[0, 0, 1],
+                                                           [0, 0, 1],
+                                                           [0, 0, 1],
+                                                           ], unit=u.m/u.s))
+
+    def test_nonsensical_units(self):
+        E = u.Quantity([[1, 0, 0],
+                        [1, 0, 0],
+                        [1, 0, 0]], unit=u.mm)
+        B = u.Quantity([[0, 1, 0],
+                        [0, 1, 0],
+                        [0, 1, 0]], unit=u.kg)
+
+        with pytest.raises(u.UnitConversionError):
+            drifts.ExB_drift(E, B)
+
+    def test_ExB_3d_array(self):
+        E = u.Quantity([[[1, 0, 0]]], unit=u.V/u.m)
+        B = u.Quantity([[[0, 1, 0]]], unit=u.T)
+
+        result = drifts.ExB_drift(2*E, 3*B)
+        assert_quantity_allclose(result, (2/3)*u.Quantity([[[0, 0, 1]]],
+                                                          unit=u.m/u.s))
+
+
+class Test_force_drift:
+    def test_force_x_B_1d_arrays(self):
+        F = u.Quantity([1, 0, 0], unit=u.N)
+        B = u.Quantity([0, 1, 0], unit=u.T)
+        q = 1*u.C
+        result = drifts.force_drift(2*F, 3*B, q)
+        assert_quantity_allclose(result, (2/3)*u.Quantity([0, 0, 1], u.m/u.s))
+
+    def test_force_B_2d_array(self):
+        F = u.Quantity([[1, 0, 0],
+                        [1, 0, 0],
+                        [1, 0, 0]], unit=u.N)
+        B = u.Quantity([[0, 1, 0],
+                        [0, 1, 0],
+                        [0, 1, 0]], unit=u.T)
+        q = 1*u.C
+
+        result = drifts.force_drift(2*F, 3*B, q)
+        assert_quantity_allclose(result, (2/3)*u.Quantity([[0, 0, 1],
+                                                           [0, 0, 1],
+                                                           [0, 0, 1],
+                                                           ], unit=u.m/u.s))
+
+    def test_force_B_3d_array(self):
+        F = u.Quantity([[[1, 0, 0]]], unit=u.N)
+        B = u.Quantity([[[0, 1, 0]]], unit=u.T)
+        q = 1*u.C
+
+        result = drifts.force_drift(2*F, 3*B, q)
+        assert_quantity_allclose(result, (2/3)*u.Quantity([[[0, 0, 1]]],
+                                                          unit=u.m/u.s))
+
+    def test_nonsensical_units(self):
+        F = u.Quantity([[1, 0, 0],
+                        [1, 0, 0],
+                        [1, 0, 0]], unit=u.mm)
+        B = u.Quantity([[0, 1, 0],
+                        [0, 1, 0],
+                        [0, 1, 0]], unit=u.kg)
+        q = 1*u.C
+
+        with pytest.raises(u.UnitConversionError):
+            drifts.force_drift(F, B, q)