Repository: PlasmaPy/PlasmaPy

Issue:
Generalize `ion_sound_speed`
Currently `ion_sound_speed` under `parameters.py` only works in the limit `alpha = 1 / (k * lambda_De)` is large(?), such that the plasma is not dispersive. This should be generalized for all values of `alpha` and the docstring should be edited to reflect this generalization.

The general equation is on pg 2 of this [Google Doc](https://docs.google.com/document/d/1EcuxYoSzy-MpjjjqBE3E6zE8g6xYU2BH1MorU7_DQQA/edit#heading=h.88burrsgujyu). Please note that this equation is for `omega`, whereas what we want is `v = omega / k`, so just drop the `k` on the right hand side.


Discussion:
If it's not already being worked on, could I try resolving this issue?
@lemmatum (and/or @StanczakDominik): do we want to define `alpha` as a new parameter, or should `k` and `lambda_D` be passed into `ion_sound_speed` (or would you rather we called the `Debye_length` function inside `ion_sound_speed` instead)?
Go for it, and feel free to call debye length from within it! 
I've written a basic implementation of the updated function, but I've realised I'll need to update the examples and tests for `ion_sound_speed` since it's now got a couple of new input parameters. 

@StanczakDominik are there any tips/pointers (or rules) for choosing some example/test values for `n_e` (for `Debye_length`) and `k` (wavenumber)?
To be honest, I usually just put in some reasonable physical values - if I
have existing values I can compare test results to, I use those.


Resolving patch:
diff --git a/plasmapy/formulary/parameters.py b/plasmapy/formulary/parameters.py
--- a/plasmapy/formulary/parameters.py
+++ b/plasmapy/formulary/parameters.py
@@ -29,6 +29,8 @@
 from plasmapy import (atomic, utils)
 from astropy.constants.si import (m_p, m_e, c, mu0, k_B, e, eps0)
 from typing import Optional
+import warnings
+from plasmapy.utils.exceptions import PhysicsWarning
 
 
 def _grab_charge(ion, z_mean=None):
@@ -209,10 +211,14 @@ def Alfven_speed(B, density, ion="p+", z_mean=None):
 @utils.check_relativistic
 @utils.check_quantity(
     T_i={'units': u.K, 'can_be_negative': False},
-    T_e={'units': u.K, 'can_be_negative': False}
+    T_e={'units': u.K, 'can_be_negative': False},
+    n_e={'units': u.m ** -3, 'can_be_negative': False, 'none_shall_pass': True},
+    k={'units': u.m ** -1, 'can_be_negative': False, 'none_shall_pass': True}
     )
 def ion_sound_speed(T_e,
                     T_i,
+                    n_e=None,
+                    k=None,
                     gamma_e=1,
                     gamma_i=3,
                     ion='p+',
@@ -231,6 +237,18 @@ def ion_sound_speed(T_e,
         Ion temperature in units of temperature or energy per
         particle.  If this is not given, then the ion temperature is
         assumed to be zero.
+        
+    n_e : ~astropy.units.Quantity
+        Electron number density. If this is not given, then ion_sound_speed 
+        will be approximated in the non-dispersive limit 
+        (:math:`k^2 \lambda_{D}^2` will be assumed zero). If n_e is given, 
+        a value for k must also be given.
+        
+    k : ~astropy.units.Quantity
+        Wavenumber (in units of inverse length, e.g. per meter). If this 
+        is not given, then ion_sound_speed will be approximated in the 
+        non-dispersive limit (:math:`k^2 \lambda_{D}^2` will be assumed zero). 
+        If k is given, a value for n_e must also be given.
 
     gamma_e : float or int
         The adiabatic index for electrons, which defaults to 1.  This
@@ -275,7 +293,8 @@ def ion_sound_speed(T_e,
         If an adiabatic index is less than one.
 
     ~astropy.units.UnitConversionError
-        If the temperature is in incorrect units.
+        If the temperature, electron number density, or wavenumber 
+        is in incorrect units.
 
     Warns
     -----
@@ -284,42 +303,54 @@ def ion_sound_speed(T_e,
 
     ~astropy.units.UnitsWarning
         If units are not provided, SI units are assumed.
+        
+    PhysicsWarning
+        If only one of (k, n_e) is given, the non-dispersive limit 
+        is assumed.
 
     Notes
     -----
-    The ion sound speed :math:`V_S` is approximately given by
+    The ion sound speed :math:`V_S` is given by
 
     .. math::
 
-        V_S = \sqrt{\frac{\gamma_e Z k_B T_e + \gamma_i k_B T_i}{m_i}}
+        V_S = \sqrt{\frac{\gamma_e Z k_B T_e + \gamma_i k_B T_i}{m_i (1 + k^2 \lambda_{D}^2)}}
 
     where :math:`\gamma_e` and :math:`\gamma_i` are the electron and
     ion adiabatic indices, :math:`k_B` is the Boltzmann constant,
     :math:`T_e` and :math:`T_i` are the electron and ion temperatures,
-    :math:`Z` is the charge state of the ion, and :math:`m_i` is the
-    ion mass.
-
-    This function assumes that the product of the wavenumber and the
-    Debye length is small. In this limit, the ion sound speed is not
-    dispersive. In other words, it is frequency independent.
+    :math:`Z` is the charge state of the ion, :math:`m_i` is the
+    ion mass, :math:`\lambda_{D}` is the Debye length, and :math:`k` is the 
+    wavenumber.
+    
+    In the non-dispersive limit (:math:`k^2 \lambda_{D}^2` is small) the 
+    equation for :math:`V_S` is approximated (the denominator reduces 
+    to :math:`m_i`).
 
     When the electron temperature is much greater than the ion
     temperature, the ion sound velocity reduces to
-    :math:`\sqrt{\gamma_e k_B T_e / m_i}`.  Ion acoustic waves can
+    :math:`\sqrt{\gamma_e k_B T_e / m_i}`. Ion acoustic waves can
     therefore occur even when the ion temperature is zero.
 
     Example
     -------
     >>> from astropy import units as u
+    >>> n = 5e19*u.m**-3
+    >>> k_1 = 3e1*u.m**-1
+    >>> k_2 = 3e7*u.m**-1
     >>> ion_sound_speed(T_e=5e6*u.K, T_i=0*u.K, ion='p', gamma_e=1, gamma_i=3)
     <Quantity 203155.0764042 m / s>
-    >>> ion_sound_speed(T_e=5e6*u.K, T_i=0*u.K)
-    <Quantity 203155.0764042 m / s>
-    >>> ion_sound_speed(T_e=500*u.eV, T_i=200*u.eV, ion='D+')
-    <Quantity 229586.01860212 m / s>
+    >>> ion_sound_speed(T_e=5e6*u.K, T_i=0*u.K, n_e=n, k=k_1, ion='p', gamma_e=1, gamma_i=3)
+    <Quantity 203155.03286794 m / s>
+    >>> ion_sound_speed(T_e=5e6*u.K, T_i=0*u.K, n_e=n, k=k_2, ion='p', gamma_e=1, gamma_i=3)
+    <Quantity 310.31329069 m / s>
+    >>> ion_sound_speed(T_e=5e6*u.K, T_i=0*u.K, n_e=n, k=k_1)
+    <Quantity 203155.03286794 m / s>
+    >>> ion_sound_speed(T_e=500*u.eV, T_i=200*u.eV, n_e=n, k=k_1, ion='D+')
+    <Quantity 229585.96150738 m / s>
 
     """
-
+    
     m_i = atomic.particle_mass(ion)
     Z = _grab_charge(ion, z_mean)
 
@@ -333,9 +364,19 @@ def ion_sound_speed(T_e,
 
     T_i = T_i.to(u.K, equivalencies=u.temperature_energy())
     T_e = T_e.to(u.K, equivalencies=u.temperature_energy())
+    
+    # Assume non-dispersive limit if values for n_e (or k) are not specified
+    klD2 = 0.0
+    if (n_e is None) ^ (k is None):
+        warnings.warn("The non-dispersive limit has been assumed for "
+                      "this calculation. To prevent this, values must "
+                      "be specified for both n_e and k.", PhysicsWarning)
+    elif n_e is not None and k is not None:
+        lambda_D = Debye_length(T_e, n_e)
+        klD2 = (k * lambda_D) ** 2
 
     try:
-        V_S_squared = (gamma_e * Z * k_B * T_e + gamma_i * k_B * T_i) / m_i
+        V_S_squared = (gamma_e * Z * k_B * T_e + gamma_i * k_B * T_i) / (m_i * (1 + klD2))
         V_S = np.sqrt(V_S_squared).to(u.m / u.s)
     except Exception:
         raise ValueError("Unable to find ion sound speed.")
diff --git a/plasmapy/utils/pytest_helpers/pytest_helpers.py b/plasmapy/utils/pytest_helpers/pytest_helpers.py
--- a/plasmapy/utils/pytest_helpers/pytest_helpers.py
+++ b/plasmapy/utils/pytest_helpers/pytest_helpers.py
@@ -686,6 +686,9 @@ def _prepare_input(param_name, param_default, insert_some_nans, insert_all_nans,
         elif param_name == "characteristic_length":
             unit = u.m
             magnitude = 1.0
+        elif param_name == "k":
+            unit = u.m ** -1
+            magnitude = 1.0
 
         # else, last resort, if it has a default argument, go with that:
         elif not (param_default is inspect._empty):


Test patch:
diff --git a/plasmapy/formulary/tests/test_parameters.py b/plasmapy/formulary/tests/test_parameters.py
--- a/plasmapy/formulary/tests/test_parameters.py
+++ b/plasmapy/formulary/tests/test_parameters.py
@@ -6,7 +6,7 @@
 from astropy.tests.helper import assert_quantity_allclose
 
 from plasmapy.utils.exceptions import RelativityWarning, RelativityError
-from plasmapy.utils.exceptions import PhysicsError
+from plasmapy.utils.exceptions import PhysicsError, PhysicsWarning
 from plasmapy.atomic.exceptions import InvalidParticleError
 from plasmapy.utils.pytest_helpers import assert_can_handle_nparray
 from astropy.constants import c, m_p, m_e, e, mu0
@@ -37,6 +37,8 @@
 rho = n_i * m_i + n_e * m_e
 T_e = 1e6 * u.K
 T_i = 1e6 * u.K
+k_1 = 3e1 * u.m ** -1
+k_2 = 3e7 * u.m ** -1
 
 B_arr = np.array([0.001, 0.002]) * u.T
 B_nanarr = np.array([0.001, np.nan]) * u.T
@@ -188,23 +190,43 @@ def test_ion_sound_speed():
     assert np.isclose(ion_sound_speed(T_i=1.3232 * u.MK, T_e=1.831 * u.MK,
                                       ion='p', gamma_e=1, gamma_i=3).value,
                       218816.06086407552)
+    
+    assert np.isclose(ion_sound_speed(T_i=1.3232 * u.MK, T_e=1.831 * u.MK,
+                                      n_e=n_e, k=k_1, ion='p', gamma_e=1, 
+                                      gamma_i=3).value,
+                      218816.06086407552)
+    
+    assert np.isclose(ion_sound_speed(T_i=1.3232 * u.MK, T_e=1.831 * u.MK,
+                                      n_e=n_e, k=k_2, ion='p', gamma_e=1, 
+                                      gamma_i=3).value,
+                      552.3212936293337)
 
     assert np.isclose(ion_sound_speed(
-        T_i=0.88 * u.MK, T_e=1.28 * u.MK, ion='p', gamma_e=1.2,
-        gamma_i=3.4).value, 193328.52857788358)
+        T_i=0.88 * u.MK, T_e=1.28 * u.MK, n_e=n_e, k=0 * u.m ** -1,
+        ion='p', gamma_e=1.2, gamma_i=3.4).value, 193328.52857788358)
 
     # case when Z=1 is assumed
     # assert ion_sound_speed(T_i=T_i, T_e=T_e, ion='p+') == ion_sound_speed(T_i=T_i, T_e=T_e,
     # ion='H-1')
 
-    assert ion_sound_speed(T_i=T_i, T_e=0 * u.K, ion='p+').unit.is_equivalent(u.m / u.s)
+    assert ion_sound_speed(T_i=T_i, T_e=0 * u.K, n_e=n_e, 
+                           k=k_1, ion='p+').unit.is_equivalent(u.m / u.s)
 
     with pytest.raises(RelativityError):
-        ion_sound_speed(T_i=T_i, T_e=T_e, gamma_i=np.inf)
+        ion_sound_speed(T_i=T_i, T_e=T_e, n_e=n_e, 
+                        k=k_1, gamma_i=np.inf)
+        
+    with pytest.warns(PhysicsWarning):
+        ion_sound_speed(T_i=T_i, T_e=T_e, n_e=n_e)
+        
+    with pytest.warns(PhysicsWarning):
+        ion_sound_speed(T_i=T_i, T_e=T_e, k=k_1)
 
     with pytest.raises(ValueError):
         ion_sound_speed(T_i=np.array([5, 6, 5]) * u.K,
-                        T_e=np.array([3, 4]) * u.K)
+                        T_e=np.array([3, 4]) * u.K,
+                        n_e=np.array([5, 6, 5]) * u.m ** -3,
+                        k=np.array([3, 4]) * u.m ** -3)
 
     with pytest.raises(TypeError):  # Is this test right??????
         ion_sound_speed(5 * u.T)
@@ -228,7 +250,13 @@ def test_ion_sound_speed():
         ion_sound_speed(T_i=T_i, T_e=0 * u.K, ion='cupcakes')
 
     with pytest.raises(ValueError):
-        ion_sound_speed(T_i=-np.abs(T_i), T_e=0 * u.K, )
+        ion_sound_speed(T_i=-np.abs(T_i), T_e=0 * u.K)
+    
+    with pytest.raises(ValueError):
+        ion_sound_speed(T_i=T_i, T_e=0 * u.K, n_e=-np.abs(n_e), k=k_1)
+        
+    with pytest.raises(ValueError):
+        ion_sound_speed(T_i=T_i, T_e=0 * u.K, n_e=n_e, k=-np.abs(k_1))
 
     with pytest.warns(RelativityWarning):
         ion_sound_speed(T_i=5e11 * u.K, T_e=0 * u.K)
@@ -237,29 +265,30 @@ def test_ion_sound_speed():
         ion_sound_speed(T_i=5e19 * u.K, T_e=0 * u.K)
 
     with pytest.raises(u.UnitConversionError):
-        ion_sound_speed(T_i=5 * u.A, T_e=0 * u.K)
+        ion_sound_speed(T_i=5 * u.A, T_e=0 * u.K, n_e=n_e, k=k_1)
 
-    assert np.isnan(ion_sound_speed(T_i=T_nanarr, T_e=0 * u.K)[1])
+    assert np.isnan(ion_sound_speed(T_i=T_nanarr, T_e=0 * u.K, n_e=n_e, k=k_1)[1])
 
-    assert np.isnan(ion_sound_speed(T_e=T_nanarr, T_i=0 * u.K)[1])
+    assert np.isnan(ion_sound_speed(T_e=T_nanarr, T_i=0 * u.K, n_e=n_e, k=k_1)[1])
 
     with pytest.raises(ValueError):
-        ion_sound_speed(T_i=T_negarr, T_e=0 * u.K)
+        ion_sound_speed(T_i=T_negarr, T_e=0 * u.K, n_e=n_e, k=k_1)
 
     with pytest.raises(ValueError):
-        ion_sound_speed(T_e=T_negarr, T_i=0 * u.K)
+        ion_sound_speed(T_e=T_negarr, T_i=0 * u.K, n_e=n_e, k=k_1)
 
     with pytest.warns(u.UnitsWarning):
-        assert ion_sound_speed(T_e=1.2e6, T_i=0 * u.K) == ion_sound_speed(T_e=1.2e6 * u.K,
-                                                                          T_i=0 * u.K)
+        assert ion_sound_speed(T_e=1.2e6, T_i=0 * u.K, n_e=n_e, k=k_1) == \
+            ion_sound_speed(T_e=1.2e6 * u.K, T_i=0 * u.K, n_e=n_e, k=k_1)
 
     with pytest.warns(u.UnitsWarning):
-        assert ion_sound_speed(T_i=1.3e6, T_e=0 * u.K) == ion_sound_speed(T_i=1.3e6 * u.K,
-                                                                          T_e=0 * u.K)
+        assert ion_sound_speed(T_i=1.3e6, T_e=0 * u.K, n_e=n_e, k=k_1) == \
+            ion_sound_speed(T_i=1.3e6 * u.K, T_e=0 * u.K, n_e=n_e, k=k_1)
 
-    ion_sound_speed(T_e=1.2e6 * u.K, T_i=0 * u.K)
+    ion_sound_speed(T_e=1.2e6 * u.K, T_i=0 * u.K, n_e=n_e, k=k_1)
     # testing for user input z_mean
-    testMeth1 = ion_sound_speed(T_e=1.2e6 * u.K, T_i=0 * u.K, z_mean=0.8).si.value
+    testMeth1 = ion_sound_speed(T_e=1.2e6 * u.K, T_i=0 * u.K, n_e=n_e, 
+                                k=0 * u.m ** -1, z_mean=0.8).si.value
     testTrue1 = 89018.0944146141
     errStr = f"ion_sound_speed() gave {testMeth1}, should be {testTrue1}."
     assert np.isclose(testMeth1,