changed calc_g() to not squeeze wavelen axis

Author: vnmanoharan

pymie/mie.py

@@ -206,7 +206,7 @@ def calc_g(m, x, nstop=None):
     coeffs = _scatcoeffs(m, x, nstop)
 
     # for multilayer particle, need to scale by the x of the outermost layer
-    outer_x = np.array(x).max(axis=-1).squeeze()
+    outer_x = np.array(x).max(axis=-1)
     cscat = _cross_sections(coeffs[0], coeffs[1])[0] * 2./outer_x**2
     g = ((4./(outer_x**2 * cscat))
          * _asymmetry_parameter(coeffs[0], coeffs[1]))

pymie/tests/test_mie_vectorized.py

@@ -765,21 +765,16 @@ def test_vectorized_asymmetry_parameter(self, num_wavelen, num_layer):
         m, x = mx(num_wavelen, num_layer, **self.mxargs)
         # make sure shape is [num_wavelen]
         g = mie.calc_g(m,x)
-        if num_wavelen == 1:
-            expected_shape = ()
-            assert g.shape == expected_shape
-            # no further test needed since no loop is required in this case
-        else:
-            expected_shape = (num_wavelen,)
-            assert g.shape == expected_shape
-
-            # we should get same values from loop. Need to set nstop to the
-            # same value as used in the vectorized calculation.
-            g_loop = np.zeros(expected_shape, dtype=float)
-            nstop = mie._nstop(x.max())
-            for i in range(num_wavelen):
-                g_loop[i] = mie.calc_g(m[i], x[i], nstop=nstop)
-            assert_equal(g, g_loop)
+        expected_shape = (num_wavelen,)
+        assert g.shape == expected_shape
+
+        # we should get same values from loop. Need to set nstop to the
+        # same value as used in the vectorized calculation.
+        g_loop = np.zeros(expected_shape, dtype=float)
+        nstop = mie._nstop(x.max())
+        for i in range(num_wavelen):
+            g_loop[i] = mie.calc_g(m[i], x[i], nstop=nstop)
+        assert_equal(g, g_loop)
 
     @pytest.mark.parametrize("num_wavelen, num_layer",
                              [(10, 1), (1, 5), (10, 5)])