Repository: PlasmaPy/PlasmaPy
Pull request #299 (2018-03-08 22:14:23)

Add type hints and update tests/docstrings in mathematics


Patch:
diff --git a/.pep8speaks.yml b/.pep8speaks.yml
index dd42a642..1daed89f 100644
--- a/.pep8speaks.yml
+++ b/.pep8speaks.yml
@@ -12,3 +12,6 @@ scanner:
 
 only_mention_files_with_errors: True
 descending_issues_order: False
+
+pycodestyle:
+    max-line-length: 99
\ No newline at end of file
diff --git a/plasmapy/mathematics/mathematics.py b/plasmapy/mathematics/mathematics.py
index dc45c9e1..df9b6a31 100644
--- a/plasmapy/mathematics/mathematics.py
+++ b/plasmapy/mathematics/mathematics.py
@@ -1,17 +1,16 @@
 """Functions related to the plasma dispersion function"""
 
 import numpy as np
-from scipy import special
 from astropy import units as u
 from mpmath import polylog
 from scipy.special import wofz as Faddeeva_function
+from typing import Union
 
-# TODO: Add type hint annotations to this module.
 
-
-def plasma_dispersion_func(zeta):
+def plasma_dispersion_func(zeta: Union[complex, int, float, np.ndarray, u.Quantity]
+                           ) -> Union[complex, float, np.ndarray, u.Quantity]:
     r"""
-    Calculate the plasma dispersion function
+    Calculate the plasma dispersion function.
 
     Parameters
     ----------
@@ -92,7 +91,8 @@ def plasma_dispersion_func(zeta):
     return Z
 
 
-def plasma_dispersion_func_deriv(zeta):
+def plasma_dispersion_func_deriv(zeta: Union[complex, int, float, np.ndarray, u.Quantity]
+                                 ) -> Union[complex, float, np.ndarray, u.Quantity]:
     r"""
     Calculate the derivative of the plasma dispersion function.
 
@@ -103,7 +103,7 @@ def plasma_dispersion_func_deriv(zeta):
 
     Returns
     -------
-    Zprime : complex, int, float, or ~numpy.ndarray
+    Zprime : complex, float, or ~numpy.ndarray
         First derivative of plasma dispersion function.
 
     Raises
@@ -127,7 +127,7 @@ def plasma_dispersion_func_deriv(zeta):
     The derivative of the plasma dispersion function is defined as:
 
     .. math::
-        Z'(\zeta) = \pi^{-0.5} \int_{-\infty}^{+\infty}
+        Z'(\zeta) = \pi^{-1/2} \int_{-\infty}^{+\infty}
         \frac{e^{-x^2}}{(x-\zeta)^2} dx
 
     where the argument is a complex number [fried.conte-1961]_.
@@ -152,34 +152,34 @@ def plasma_dispersion_func_deriv(zeta):
         if zeta.unit == u.dimensionless_unscaled:
             zeta = zeta.value
         else:
-            raise u.UnitsError("The argument to "
-                               "plasma_dispersion_function_deriv "
-                               "must be dimensionless if it is a Quantity")
+            raise u.UnitsError("The argument to plasma_dispersion_function_deriv "
+                               "must be dimensionless if it is a Quantity.")
 
     if not np.all(np.isfinite(zeta)):
-        raise ValueError("The argument to plasma_dispersion_function_deriv is "
-                         "not finite.")
+        raise ValueError("The argument to plasma_dispersion_function_deriv is not finite.")
 
     Zprime = -2 * (1 + zeta * plasma_dispersion_func(zeta))
 
     return Zprime
 
 
-def Fermi_integral(x, j):
+def Fermi_integral(
+        x: Union[float, int, complex, np.ndarray],
+        j: Union[float, int, complex, np.ndarray]) -> Union[float, complex, np.ndarray]:
     r"""
     Calculate the complete Fermi-Dirac integral.
 
     Parameters
     ----------
-    x : ~astropy.units.Quantity
+    x : float, int, complex, or ~numpy.ndarray
         Argument of the Fermi-Dirac integral function.
 
-    j : ~astropy.units.Quantity
+    j : float, int, complex, or ~numpy.ndarray
         Order/index of the Fermi-Dirac integral function.
 
     Returns
     -------
-    integral : ~astropy.units.Quantity, complex
+    integral : float, complex, or ~numpy.ndarray
         Complete Fermi-Dirac integral for given argument and order.
 
     Raises
@@ -194,32 +194,30 @@ def Fermi_integral(x, j):
     ValueError
         If the argument is not entirely finite.
 
-    See also
+    See Also
     --------
     plasma_dispersion_func
 
     Notes
     -----
-    The complete Fermi-Dirac integral is defined as: [1]_
+    The `complete Fermi-Dirac integral
+    <https://en.wikipedia.org/wiki/Complete_Fermi-Dirac_integral>`_ is
+    defined as:
 
     .. math::
-        F_j (x) = \frac{1}{\Gamma (j+1)} \int_0^{\inf} \frac{t^j}{\exp{(t-x)} + 1} dt
+        F_j (x) = \frac{1}{\Gamma (j+1)} \int_0^{\infty} \frac{t^j}{\exp{(t-x)} + 1} dt
 
     for j > 0.
 
-    This is equivalent to the following polylogarithm [2]_ function:
+    This is equivalent to the following `polylogarithm
+    <https://en.wikipedia.org/wiki/Polylogarithm>`_ function:
 
     .. math::
         F_j (x) = -Li_{j+1}\left(-e^{x}\right)
 
     Warning: at present this function is limited to relatively small
-    arguments due to limitations in the mpmath package's implementation
-    of polylog.
-
-    References
-    ----------
-    .. [1] https://en.wikipedia.org/wiki/Complete_Fermi-Dirac_integral
-    .. [2] https://en.wikipedia.org/wiki/Polylogarithm
+    arguments due to limitations in the `~mpmath` package's
+    implementation of `~mpmath.polylog`.
 
     Examples
     --------
@@ -229,6 +227,7 @@ def Fermi_integral(x, j):
     (1.3132616875182228-0j)
     >>> Fermi_integral(1, 1)
     (1.8062860704447743-0j)
+
     """
     if isinstance(x, (int, float, complex)):
         arg = -np.exp(x)
@@ -240,4 +239,4 @@ def Fermi_integral(x, j):
             integral_arr[idx] = -1 * complex(polylog(j + 1, -np.exp(val)))
         return integral_arr
     else:
-        raise ValueError(f"Improper type {type(x)} given for argument x.")
+        raise TypeError(f"Improper type {type(x)} given for argument x.")


Test patch:
diff --git a/plasmapy/mathematics/tests/test_Fermi_integral.py b/plasmapy/mathematics/tests/test_Fermi_integral.py
index 8447fb90..7ef4165a 100644
--- a/plasmapy/mathematics/tests/test_Fermi_integral.py
+++ b/plasmapy/mathematics/tests/test_Fermi_integral.py
@@ -10,32 +10,31 @@
 
 class Test_Fermi_integral:
     def setup_method(self):
-        """initializing parameters for tests """
+        """Initialize parameters for tests."""
         self.arg1 = 3.889780
-        self.True1 = (6.272518847136373-8.673617379884035e-19j)
+        self.True1 = (6.272518847136373 - 8.673617379884035e-19j)
         self.argFail1 = 3.889781
         self.order1 = 0.5
         self.args = np.array([0.5, 1, 2])
-        self.Trues = np.array([(1.1173314873128224-0j),
-                               (1.5756407761513003-0j),
-                               (2.8237212774015843-2.6020852139652106e-18j)])
+        self.Trues = np.array([(1.1173314873128224 - 0j),
+                               (1.5756407761513003 - 0j),
+                               (2.8237212774015843 - 2.6020852139652106e-18j)])
 
     def test_known1(self):
         """
-        Tests Fermi_integral for expected value.
+        Test Fermi_integral for expected value.
         """
         methodVal = Fermi_integral(self.arg1, self.order1)
         testTrue = np.isclose(methodVal,
                               self.True1,
                               rtol=1e-16,
                               atol=0.0)
-        errStr = (f"Fermi integral value should be {self.True1} and not "
-                  f"{methodVal}.")
+        errStr = f"Fermi integral value should be {self.True1} and not {methodVal}."
         assert testTrue, errStr
 
     def test_fail1(self):
         """
-        Tests if test_known1() would fail if we slightly adjusted the
+        Test if test_known1() would fail if we slightly adjusted the
         value comparison by some quantity close to numerical error.
         """
         fail1 = self.True1 + 1e-15
@@ -50,27 +49,26 @@ def test_fail1(self):
 
     def test_polog_fail(self):
         """
-        Tests whether Fermi_integral() fails due to polylog from mpmath
+        Test whether Fermi_integral fails due to `~mpmath.polylog`
         not having an implementation for larger argument values.
         """
         with pytest.raises(NotImplementedError):
             Fermi_integral(self.argFail1, self.order1)
 
     def test_array(self):
-        """Testing Fermi_integral where argument is an array of inputs."""
+        """Test Fermi_integral where argument is an array of inputs."""
         methodVals = Fermi_integral(self.args, self.order1)
         testTrue = np.allclose(methodVals,
                                self.Trues,
                                rtol=1e-16,
                                atol=0.0)
-        errStr = (f"Fermi integral value should be {self.Trues} and not "
-                  f"{methodVals}.")
+        errStr = (f"Fermi integral value should be {self.Trues} and not {methodVals}.")
         assert testTrue, errStr
 
     def test_invalid_type(self):
         """
-        Testing whether ValueError is raised when an invalid argument
-        type is passed to Fermi integral.
+        Test whether `TypeError` is raised when an invalid argument
+        type is passed to `~plasmapy.mathematics.Fermi_integral`.
         """
-        with pytest.raises(ValueError):
+        with pytest.raises(TypeError):
             Fermi_integral([1, 2, 3], self.order1)
diff --git a/plasmapy/mathematics/tests/test_dispersion.py b/plasmapy/mathematics/tests/test_dispersion.py
index ae0eed14..25ba9e54 100644
--- a/plasmapy/mathematics/tests/test_dispersion.py
+++ b/plasmapy/mathematics/tests/test_dispersion.py
@@ -15,7 +15,7 @@
     (0, 1j * np.sqrt(π)),
     (1, -1.076_159_013_825_536_8 + 0.652_049_332_173_292_2j),
     (1j, 0.757_872_156_141_311_87j),
-    (1.2 + 4.4j, -0.054_246_157_069_223_27+0.207_960_584_359_855_62j),
+    (1.2 + 4.4j, -0.054_246_157_069_223_27 + 0.207_960_584_359_855_62j),
     (9.2j, plasma_dispersion_func(9.2j * u.dimensionless_unscaled)),
     (5.4 - 3.1j, -0.139_224_873_051_713_11 - 0.082_067_822_640_155_802j),
     (9.9 - 10j, 2.013_835_257_947_027_6 - 25.901_274_737_989_727j),
@@ -92,7 +92,7 @@ def test_plasma_dispersion_func_power_series_expansion():
     Z_power_series = np.zeros_like(w_array)
 
     for n in range(0, 200):
-        Z_power_series += 1j * np.sqrt(π) * (1j * w_array)**n / Γ(n / 2 + 1)
+        Z_power_series += 1j * np.sqrt(π) * (1j * w_array) ** n / Γ(n / 2 + 1)
 
     assert np.allclose(Z_of_w_array, Z_power_series,
                        atol=1e-15 * (1 + 1j), rtol=1e-15), \
@@ -154,7 +154,7 @@ def test_plasma_dispersion_func_deriv(w, expected):
 
     Z_deriv = plasma_dispersion_func_deriv(w)
 
-    assert np.isclose(Z_deriv, expected, atol=5e-5*(1+1j), rtol=5e-6), \
+    assert np.isclose(Z_deriv, expected, atol=5e-5 * (1 + 1j), rtol=5e-6), \
         (f"The derivative of the plasma dispersion function does not match "
          f"the expected value for w = {w}.  The value of "
          f"plasma_dispersion_func_deriv({w}) equals {Z_deriv} whereas the "