Repository: PlasmaPy/PlasmaPy
Pull request #485 (2018-06-05 03:23:44)

Numpy array operations in collision_frequency 's
Similar to before, the goal this time is to make it so that collision_frequency and friends work with len>1 array Quantities, such as `T_e = np.array([1, 2, 3]) * u.eV`.

Working on issue #469, hopefully this is the last PR needed that will touch the functions themselves.  Muahaha, can calculate so many numbers now! 

Still will probably need another PR to close the issue, just to add tests for array-capability to all the various functions that were already passing.  Aaaand probably Braginskii should be figured out to handle arrays, but... let's make that a separate issue. I'm not sure if it will be vectorize-able in all cases,  I think it will be.


Patch:
diff --git a/plasmapy/physics/transport/collisions.py b/plasmapy/physics/transport/collisions.py
index 5b9a03be..f782b19a 100644
--- a/plasmapy/physics/transport/collisions.py
+++ b/plasmapy/physics/transport/collisions.py
@@ -309,24 +309,36 @@ def _boilerPlate(T, particles, V):
     # obtaining reduced mass of 2 particle collision system
     reduced_mass = atomic.reduced_mass(*particles)
 
+    # getting thermal velocity of system if no velocity is given
+    V = _replaceNanVwithThermalV(V, T, reduced_mass)
+
+    _check_relativistic(V, 'V')
+
+    return T, masses, charges, reduced_mass, V
+
+
+def _replaceNanVwithThermalV(V, T, m):
+    """
+    Get thermal velocity of system if no velocity is given, for a given mass.
+    Handles vector checks for V, you must already know that T and m are okay.
+    """
     if np.any(V == 0):
         raise utils.exceptions.PhysicsError("You cannot have a collision for zero velocity!")
     # getting thermal velocity of system if no velocity is given
     if V is None:
-        V = parameters.thermal_speed(T, mass=reduced_mass)
+        V = parameters.thermal_speed(T, mass=m)
     elif np.any(np.isnan(V)):
         if np.isscalar(V.value) and np.isscalar(T.value):
-            V = parameters.thermal_speed(T, mass=reduced_mass)
+            V = parameters.thermal_speed(T, mass=m)
         elif np.isscalar(V.value):
-            V = parameters.thermal_speed(T, mass=reduced_mass)
+            V = parameters.thermal_speed(T, mass=m)
         elif np.isscalar(T.value):
             V = V.copy()
-            V[np.isnan(V)] = parameters.thermal_speed(T, mass=reduced_mass)
+            V[np.isnan(V)] = parameters.thermal_speed(T, mass=m)
         else:
             V = V.copy()
-            V[np.isnan(V)] = parameters.thermal_speed(T[np.isnan(V)], mass=reduced_mass)
-    _check_relativistic(V, 'V')
-    return T, masses, charges, reduced_mass, V
+            V[np.isnan(V)] = parameters.thermal_speed(T[np.isnan(V)], mass=m)
+    return V
 
 
 @check_quantity({"T": {"units": u.K, "can_be_negative": False}
@@ -739,11 +751,10 @@ def collision_frequency(T,
     # reduced mass
     V_reduced = V_r
     if particles[0] in ('e','e-') and particles[1] in ('e','e-'):
+        # electron-electron collision
         # if a velocity was passed, we use that instead of the reduced
         # thermal velocity
-        if np.isnan(V):
-            V = V_reduced
-        # electron-electron collision
+        V = _replaceNanVwithThermalV(V, T, reduced_mass)
         # impact parameter for 90 degree collision
         bPerp = impact_parameter_perp(T=T,
                                       particles=particles,
@@ -759,10 +770,9 @@ def collision_frequency(T,
         # electron-ion collision
         # Need to manually pass electron thermal velocity to obtain
         # correct perpendicular collision radius
-        if np.isnan(V):
-            # we ignore the reduced velocity and use the electron thermal
-            # velocity instead
-            V = np.sqrt(2 * k_B * T / m_e)
+        # we ignore the reduced velocity and use the electron thermal
+        # velocity instead
+        V = _replaceNanVwithThermalV(V, T, m_e)
         # need to also correct mass in collision radius from reduced
         # mass to electron mass
         bPerp = impact_parameter_perp(T=T,
@@ -778,11 +788,10 @@ def collision_frequency(T,
                                     V=np.nan * u.m / u.s,
                                     method=method)
     else:
+        # ion-ion collision
         # if a velocity was passed, we use that instead of the reduced
         # thermal velocity
-        if np.isnan(V):
-            V = V_reduced
-        # ion-ion collision
+        V = _replaceNanVwithThermalV(V, T, reduced_mass)
         bPerp = impact_parameter_perp(T=T,
                                       particles=particles,
                                       V=V)
@@ -941,9 +950,9 @@ def fundamental_electron_collision_freq(T_e,
     fundamental_ion_collision_freq
     """
     T_e = T_e.to(u.K, equivalencies=u.temperature_energy())
-    if not V:
-        # electron thermal velocity (most probable)
-        V = np.sqrt(2 * k_B * T_e / m_e)
+
+    # specify to use electron thermal velocity (most probable), not based on reduced mass
+    V = _replaceNanVwithThermalV(V, T_e, m_e)
 
     particles = [ion_particle, 'e-']
     Z_i = atomic.integer_charge(ion_particle)
@@ -1079,10 +1088,12 @@ def fundamental_ion_collision_freq(T_i,
     T_i = T_i.to(u.K, equivalencies=u.temperature_energy())
     m_i = atomic.particle_mass(ion_particle)
     particles = [ion_particle, ion_particle]
-    if not V:
-        # ion thermal velocity (most probable)
-        V = np.sqrt(2 * k_B * T_i / m_i)
+
+    # specify to use ion thermal velocity (most probable), not based on reduced mass
+    V = _replaceNanVwithThermalV(V, T_i, m_i)
+
     Z_i = atomic.integer_charge(ion_particle)
+
     nu = collision_frequency(T_i,
                              n_i,
                              particles,
@@ -1108,6 +1119,7 @@ def fundamental_ion_collision_freq(T_i,
         nu_i = coeff * nu_mod
     else:
         nu_i = coeff * nu
+
     return nu_i.to(1 / u.s)
 
 


Test patch:
diff --git a/plasmapy/physics/transport/tests/test_collisions.py b/plasmapy/physics/transport/tests/test_collisions.py
index 2fa7fc9e..eb48ef4b 100644
--- a/plasmapy/physics/transport/tests/test_collisions.py
+++ b/plasmapy/physics/transport/tests/test_collisions.py
@@ -11,7 +11,9 @@
                                         mobility,
                                         Knudsen_number,
                                         coupling_parameter)
-from plasmapy.physics.transport.collisions import Spitzer_resistivity
+from plasmapy.physics.transport.collisions import (Spitzer_resistivity,
+                                                   fundamental_electron_collision_freq,
+                                                   fundamental_ion_collision_freq)
 from plasmapy.utils import exceptions
 from plasmapy.constants import m_p, m_e, c
 
@@ -41,6 +43,14 @@ def setup_class(self):
         self.gms6 = 3.635342040477818
         self.gms6_negative = 0.030720859361047514
 
+    def test_unknown_method(self):
+        """Test that function will raise ValueError on non-existent method"""
+        with pytest.raises(ValueError):
+            Coulomb_logarithm(self.T_arr[0],
+                              self.n_arr[0],
+                              self.particles,
+                              method="welcome our new microsoft overlords")
+
     def test_handle_invalid_V(self):
         """Test that V default, V = None, and V = np.nan all give the same result"""
         methodVal_0 = Coulomb_logarithm(self.T_arr[0],
@@ -674,6 +684,16 @@ def test_handle_numpy_array(self, method):
                                        method=method)
         assert_quantity_allclose((methodVal[0][0], methodVal[1][0]), methodVal_0)
 
+    def test_extend_scalar_bmin(self):
+        """
+        Test to verify that if T is scalar and n is vector, bmin will be extended
+        to the same length as bmax
+        """
+        (bmin, bmax) = impact_parameter(1 * u.eV,
+                                        self.n_e_arr,
+                                        self.particles)
+        assert(len(bmin) == len(bmax))
+
 
 class Test_collision_frequency:
     @classmethod
@@ -794,6 +814,50 @@ def test_zmean(self):
         assert testTrue, errStr
 
 
+class Test_fundamental_electron_collision_freq():
+    @classmethod
+    def setup_class(self):
+        """initializing parameters for tests """
+        self.T_arr = np.array([1, 2]) * u.eV
+        self.n_arr = np.array([1e20, 2e20]) * u.cm ** -3
+        self.ion_particle = 'p'
+        self.coulomb_log = 10
+
+    def test_handle_numpy_array(self):
+        """Tests to verify that can handle Quantities with numpy array as the value"""
+        methodVal = fundamental_electron_collision_freq(self.T_arr,
+                                                        self.n_arr,
+                                                        self.ion_particle,
+                                                        coulomb_log=self.coulomb_log)
+        methodVal_0 = fundamental_electron_collision_freq(self.T_arr[0],
+                                                          self.n_arr[0],
+                                                          self.ion_particle,
+                                                          coulomb_log=self.coulomb_log)
+        assert_quantity_allclose(methodVal[0], methodVal_0)
+
+
+class Test_fundamental_ion_collision_freq():
+    @classmethod
+    def setup_class(self):
+        """initializing parameters for tests """
+        self.T_arr = np.array([1, 2]) * u.eV
+        self.n_arr = np.array([1e20, 2e20]) * u.cm ** -3
+        self.ion_particle = 'p'
+        self.coulomb_log = 10
+
+    def test_handle_numpy_array(self):
+        """Tests to verify that can handle Quantities with numpy array as the value"""
+        methodVal = fundamental_ion_collision_freq(self.T_arr,
+                                                   self.n_arr,
+                                                   self.ion_particle,
+                                                   coulomb_log=self.coulomb_log)
+        methodVal_0 = fundamental_ion_collision_freq(self.T_arr[0],
+                                                     self.n_arr[0],
+                                                     self.ion_particle,
+                                                     coulomb_log=self.coulomb_log)
+        assert_quantity_allclose(methodVal[0], methodVal_0)
+
+
 class Test_mean_free_path:
     @classmethod
     def setup_class(self):