changed _internal_coeffs() to not squeeze wavelen axis

Author: vnmanoharan

pymie/mie.py

@@ -612,7 +612,9 @@ def _internal_coeffs(m, x, n_max, eps1 = DEFAULT_EPS1, eps2 = DEFAULT_EPS2):
     dl = (m[..., np.newaxis] * ratio * (D3x - D1x)
           / (m[..., np.newaxis] * D3x - D1mx))
     # start from l = 1
-    return np.array([cl[..., 1:], dl[..., 1:]]).squeeze()
+    cldl = np.array([cl[..., 1:], dl[..., 1:]])
+    # remove unneeded layer axis
+    return cldl[..., 0, :]
 
 def _trans_coeffs(m, x, n_max, eps1 = DEFAULT_EPS1, eps2 = DEFAULT_EPS2):
     '''

pymie/tests/test_mie_vectorized.py

@@ -468,21 +468,16 @@ def test_vectorized_internal_coeffs(self, num_wavelen, num_layer):
             coeffs = mie._internal_coeffs(m, x, nstop)
 
             # make sure shape is correct
-            if num_wavelen == 1:
-                expected_shape = (2, nstop)
-                assert coeffs.shape == expected_shape
-                # no further test since no loop required in this case
-            else:
-                expected_shape = (2, num_wavelen, nstop)
-                assert coeffs.shape == expected_shape
+            expected_shape = (2, num_wavelen, nstop)
+            assert coeffs.shape == expected_shape
 
-                # we should get same values from loop
-                coeffs_loop = np.zeros(expected_shape, dtype=complex)
-                for i in range(m.shape[0]):
-                    c = mie._internal_coeffs(m[i], x[i], nstop)
-                    coeffs_loop[:, i] = c
+            # we should get same values from loop
+            coeffs_loop = []
+            for i in range(m.shape[0]):
+                coeffs_loop.append(mie._internal_coeffs(m[i], x[i], nstop))
+            coeffs_loop = np.concatenate(coeffs_loop, axis=1)
 
-                assert_equal(coeffs, coeffs_loop)
+            assert_equal(coeffs, coeffs_loop)
 
     @pytest.mark.parametrize("n_matrix",
                              [1.33, pytest.param(1.33+0.001j),