changed integrate_intensity_complex_medium to not squeeze

Author: vnmanoharan

pymie/mie.py

@@ -1368,6 +1368,11 @@ def integrate_intensity_complex_medium(dscat, distance, thetas, k,
         if isinstance(dsigma_2, Quantity):
             sigma_2 = Quantity(sigma_2, dsigma_2.units)
 
+    # k has trailing axes for theta (and possibly phi) that are no longer
+    # needed after the integration.  We remove them here
+    k_shape = k.shape
+    k = k.reshape(num_values)
+
     # multiply by factor that accounts for attenuation in the incident light
     # (see Sudiarta and Chylek (2001), eq 10).
     # if the imaginary part of k is close to 0 (because the medium index is
@@ -1381,16 +1386,12 @@ def integrate_intensity_complex_medium(dscat, distance, thetas, k,
                           1 / (exponent / (2*distance*k.imag)
                                + (1 - exponent) / (2*distance*k.imag)**2))
 
-    # prepare for broadcasting (this will add trailing axes of size 1)
-    sigma_1 = sigma_1.reshape(factor.shape)
-    sigma_2 = sigma_2.reshape(factor.shape)
-
     # calculate the averaged sigma
     sigma = (sigma_1 + sigma_2)/2 * factor
 
-    return(sigma.squeeze(), (sigma_1*factor).squeeze(),
-           (sigma_2*factor).squeeze(), (dsigma_1*factor/2).squeeze(),
-           (dsigma_2*factor/2).squeeze())
+    return(sigma, (sigma_1*factor),
+           (sigma_2*factor), (dsigma_1*factor.reshape(k_shape)/2),
+           (dsigma_2*factor.reshape(k_shape)/2))
 
 def diff_abs_intensity_complex_medium(m, x, thetas, ktd):
     '''

pymie/tests/test_mie_vectorized.py

@@ -689,9 +689,9 @@ def test_vectorized_angular_functions(self, num_wavelen,
             # by trying to assign units to each element, so we have to take
             # the magnitudes here
             units = integral_loop[0].units
-            sigma[i] = integral_loop[0].magnitude
-            sigma_1[i] = integral_loop[1].magnitude
-            sigma_2[i] = integral_loop[2].magnitude
+            sigma[i] = integral_loop[0].magnitude.squeeze()
+            sigma_1[i] = integral_loop[1].magnitude.squeeze()
+            sigma_2[i] = integral_loop[2].magnitude.squeeze()
             dsigma_1[i] = integral_loop[3].magnitude
             dsigma_2[i] = integral_loop[4].magnitude
 
@@ -706,11 +706,11 @@ def test_vectorized_angular_functions(self, num_wavelen,
         assert_equal(i12[0], i1)
         assert_equal(i12[1], i2)
 
-        assert_equal(integral[0].magnitude, sigma.squeeze())
-        assert_equal(integral[1].magnitude, sigma_1.squeeze())
-        assert_equal(integral[2].magnitude, sigma_2.squeeze())
-        assert_equal(integral[3].magnitude, dsigma_1.squeeze())
-        assert_equal(integral[4].magnitude, dsigma_2.squeeze())
+        assert_equal(integral[0].magnitude, sigma)
+        assert_equal(integral[1].magnitude, sigma_1)
+        assert_equal(integral[2].magnitude, sigma_2)
+        assert_equal(integral[3].magnitude, dsigma_1)
+        assert_equal(integral[4].magnitude, dsigma_2)
 
 
 class TestVectorizedUserFunctions():
@@ -940,7 +940,8 @@ def test_vectorized_calc_reflectance(self, n_medium, num_wavelen,
         for i in range(num_wavelen):
             reflectance = mie.calc_reflectance(radius, n_medium, n_particle[i],
                                                wavelen[i]).magnitude
-            refl_loop[i] = reflectance
+            # squeeze to remove singleton wavelength dimension
+            refl_loop[i] = reflectance.squeeze()
 
         assert_allclose(refl.magnitude, refl_loop, rtol=1e-14)
         assert refl.units == 1/wavelen.units**2