changed _scatcoeffs() and all related funcs to not squeeze wavelen axis

Author: vnmanoharan

pymie/mie.py

@@ -437,7 +437,8 @@ def _scatcoeffs(m, x, nstop, eps1 = DEFAULT_EPS1, eps2 = DEFAULT_EPS2):
     # index ratio should be specified as a 2D array with shape
     # [num_values, num_layers] to calculate over a set of different values,
     # such as wavelengths. If specified as a 1D array, shape is [num_layers].
-    if np.atleast_2d(m).shape[-1] > 1:
+    num_layers = np.atleast_2d(m).shape[-1]
+    if num_layers > 1:
         return _scatcoeffs_multi(m, x)
 
     # Scattering coefficients for single-layer particles.
@@ -446,12 +447,17 @@ def _scatcoeffs(m, x, nstop, eps1 = DEFAULT_EPS1, eps2 = DEFAULT_EPS2):
     # nmx = np.array([nstop, np.round(np.absolute(m*x))]).max() + 20
     # Dnmx = mie_specfuncs.log_der_1(m*x, nmx, nstop)
     # above replaced with Lentz algorithm
-    z = np.atleast_1d(m * x).squeeze()
+    z = np.atleast_1d(m * x)
+    # remove unneded layer axis
+    z = z[..., 0]
     Dnmx = mie_specfuncs.dn_1_down(z, nstop + 1, nstop,
                                    mie_specfuncs.lentz_dn1(z, nstop + 1,
                                                            eps1, eps2))
+
     n = np.arange(nstop+1)
-    x = np.atleast_2d(x)
+    if np.ndim(x) <= 1:
+        x = np.atleast_1d(x)[..., np.newaxis]
+
     psi, xi = mie_specfuncs.riccati_psi_xi(x, nstop)
 
     # insert zeroes at the beginning of second axis (order axis)
@@ -460,9 +466,8 @@ def _scatcoeffs(m, x, nstop, eps1 = DEFAULT_EPS1, eps2 = DEFAULT_EPS2):
     an = ( (Dnmx/m + n/x)*psi - psishift ) / ( (Dnmx/m + n/x)*xi - xishift )
     bn = ( (Dnmx*m + n/x)*psi - psishift ) / ( (Dnmx*m + n/x)*xi - xishift )
 
-    # coefficient array has shape [2, num_values, nstop] or [2, nstop] if
-    # only one value (only one wavelength, for example)
-    return np.array([an[..., 1:nstop+1], bn[..., 1:nstop+1]]).squeeze()
+    # coefficient array has shape [2, num_values, nstop]
+    return np.array([an[..., 1:nstop+1], bn[..., 1:nstop+1]])
 
 def _scatcoeffs_multi(marray, xarray, nstop=None, eps1 = 1e-3, eps2 = 1e-16):
     '''Calculate scattered field expansion coefficients (in the Mie formalism)
@@ -578,7 +583,7 @@ def _scatcoeffs_multi(marray, xarray, nstop=None, eps1 = 1e-3, eps2 = 1e-16):
           / ((hbns*mlast + n/xlast)*xi - xishift))
 
     # output begins at n=1
-    return np.array([an[..., 1:nstop+1], bn[..., 1:nstop+1]]).squeeze()
+    return np.array([an[..., 1:nstop+1], bn[..., 1:nstop+1]])
 
 def _internal_coeffs(m, x, n_max, eps1 = DEFAULT_EPS1, eps2 = DEFAULT_EPS2):
     '''

pymie/tests/test_mie.py

@@ -88,8 +88,9 @@ def test_form_factor():
                            93.5508557840006])
 
     iparperp = mie.calc_ang_scat(m, x, angles)
-    assert_array_almost_equal(iparperp[0], ipar_bhmie)
-    assert_array_almost_equal(iparperp[1], iperp_bhmie)
+    # squeeze to remove singlet wavelength dimension
+    assert_array_almost_equal(iparperp[0].squeeze(), ipar_bhmie)
+    assert_array_almost_equal(iparperp[1].squeeze(), iperp_bhmie)
 
 def test_efficiencies():
     x = np.array([0.01, 0.01778279, 0.03162278, 0.05623413, 0.1, 0.17782794,
@@ -170,8 +171,9 @@ def test_absorbing_materials():
                            8.26505988320951, 47.4736966179677])
 
     iparperp = mie.calc_ang_scat(m, x, angles)
-    assert_array_almost_equal(iparperp[0], ipar_bhmie)
-    assert_array_almost_equal(iparperp[1], iperp_bhmie)
+    # squeeze to remove singlet wavelen axis before comparison
+    assert_array_almost_equal(iparperp[0].squeeze(), ipar_bhmie)
+    assert_array_almost_equal(iparperp[1].squeeze(), iperp_bhmie)
 
 def test_multilayer_spheres():
     # test that form factors and cross sections are the same for a
@@ -485,8 +487,6 @@ def test_differential_cross_section():
 
     # With far-field Mie solutions
     I_parperp_cad = mie.calc_ang_scat(m, x, theta)
-    # temp fix to re-insert wavelength dimension (which gets squeezed out)
-    I_parperp_cad = I_parperp_cad[:, np.newaxis, :]
 
     # With Mie solutions at surface of particle (but neglecting near-fields)
     kd = (k*distance).to("").magnitude
@@ -748,11 +748,7 @@ def test_diff_scat_intensity_complex_medium_cartesian():
     I_par_perp_mag = np.sqrt((I_parperp**2).sum(axis=0))
 
     # check that the magnitudes are equal
-    #
-    # Because of the way this test is done (using a 2D array of angles
-    # for theta), we don't end up with a wavelength axis for I_parperp.  Have
-    # to squeeze to compare.
-    assert_allclose(I_xy_mag.squeeze(), I_par_perp_mag, rtol=1e-15)
+    assert_allclose(I_xy_mag, I_par_perp_mag, rtol=1e-15)
 
 def test_integrate_intensity_complex_medium_cartesian():
     '''
@@ -901,7 +897,12 @@ def test_dwell_time_and_energy():
     distance_calc = dwell_time*c
     distance_calc = distance_calc.to('um')
 
-    assert_approx_equal(cscat_reported.magnitude, cscat_calc.magnitude, significant=2)
-    assert_approx_equal(W_star_reported, np.real(W_star_calc), significant=2)
-    assert_approx_equal(W_reported.magnitude, np.real(W_calc.magnitude), significant=2)
-    assert_approx_equal(distance_reported.magnitude, np.real(distance_calc.magnitude), significant=2)
+    # squeeze to remove singlet wavelength dimension on calculated values
+    assert_allclose(cscat_calc.squeeze().magnitude,
+                    cscat_reported.magnitude, rtol=1e-2)
+    assert_allclose(np.real(W_star_calc).squeeze(), W_star_reported, rtol=1e-1)
+    assert_allclose(np.real(W_calc.magnitude).squeeze(),
+                    W_reported.magnitude, rtol=1e-1)
+    assert_allclose(np.real(distance_calc.magnitude).squeeze(),
+                    distance_reported.magnitude, rtol=1e-1)
+

pymie/tests/test_mie_vectorized.py

@@ -430,27 +430,23 @@ def test_vectorized_scatcoeffs(self, num_wavelen, num_layer):
             assert_equal(coeffs, coeffs_direct)
 
         # 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
-
-            # we should get same value from loop
-            coeffs_loop = np.zeros(expected_shape, dtype=complex)
-            for i in range(m.shape[0]):
-                if num_layer == 1:
-                    c = mie._scatcoeffs(m[i], x[i], nstop)
-                else:
-                    # need to specify nstop here; otherwise we will get a
-                    # different number of scattering coefficients for each
-                    # wavelength, since _scatcoeffs_multi() picks the largest x
-                    # for each wavelength.
-                    c = mie._scatcoeffs_multi(m[i], x[i], nstop)
-                coeffs_loop[:, i] = c
-            assert_equal(coeffs, coeffs_loop)
+        expected_shape = (2, num_wavelen, nstop)
+        assert coeffs.shape == expected_shape
+
+        # we should get same value from loop
+        coeffs_loop = []
+        for i in range(m.shape[0]):
+            if num_layer == 1:
+                coeffs_loop.append(mie._scatcoeffs(m[i], x[i], nstop))
+            else:
+                # need to specify nstop here; otherwise we will get a
+                # different number of scattering coefficients for each
+                # wavelength, since _scatcoeffs_multi() picks the largest x
+                # for each wavelength.
+                coeffs_loop.append(mie._scatcoeffs_multi(m[i], x[i], nstop))
+        # concatenate along wavelength axis
+        coeffs_loop = np.concatenate(coeffs_loop, axis=1)
+        assert_equal(coeffs, coeffs_loop)
 
     @pytest.mark.parametrize("num_wavelen,num_layer",
                              [(1, 1), (10, 1), (1, 5), (10, 5)])
@@ -631,10 +627,7 @@ def test_vectorized_angular_functions(self, num_wavelen,
 
         # check that shapes of all the computed quantities are correct
         for element in vsa + mat:
-            if num_wavelen > 1:
-                assert element.shape == (num_wavelen, ) + thetas.shape
-            else:
-                assert element.shape == thetas.shape
+            assert element.shape == (num_wavelen, ) + thetas.shape
         assert i12.shape == (2, num_wavelen) + thetas.shape
 
         # check that vectorized calculations match looped calculations over
@@ -691,13 +684,13 @@ def test_vectorized_angular_functions(self, num_wavelen,
             dsigma_1[i] = integral_loop[3].squeeze()
             dsigma_2[i] = integral_loop[4].squeeze()
 
-        assert_equal(mat[0], S1.squeeze())
-        assert_equal(mat[1], S2.squeeze())
-        assert_equal(mat[2], S3.squeeze())
-        assert_equal(mat[3], S4.squeeze())
+        assert_equal(mat[0], S1)
+        assert_equal(mat[1], S2)
+        assert_equal(mat[2], S3)
+        assert_equal(mat[3], S4)
 
-        assert_equal(vsa[0], amp0.squeeze())
-        assert_equal(vsa[1], amp1.squeeze())
+        assert_equal(vsa[0], amp0)
+        assert_equal(vsa[1], amp1)
 
         assert_equal(i12[0], i1)
         assert_equal(i12[1], i2)
@@ -763,10 +756,11 @@ def test_vectorized_asymmetry_parameter(self, num_wavelen, num_layer):
 
         # 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)
+        g_loop = []
         nstop = mie._nstop(x.max())
         for i in range(num_wavelen):
-            g_loop[i] = mie.calc_g(m[i], x[i], nstop=nstop)
+            g_loop.append(mie.calc_g(m[i], x[i], nstop=nstop))
+        g_loop = np.concatenate(g_loop, axis=0)
         assert_equal(g, g_loop)
 
     @pytest.mark.parametrize("num_wavelen, num_layer",
@@ -782,10 +776,7 @@ def test_vectorized_cross_sections(self, num_wavelen, num_layer):
         cscat, cext, cback, cabs, asym = mie.calc_cross_sections(m, x)
 
         # test shape
-        if num_wavelen > 1:
-            expected_shape = (num_wavelen,)
-        else:
-            expected_shape = ()
+        expected_shape = (num_wavelen,)
         for cs in [cscat, cext, cback, cabs, asym]:
             assert cs.shape == expected_shape
 
@@ -798,7 +789,7 @@ def test_vectorized_cross_sections(self, num_wavelen, num_layer):
         for i in range(num_wavelen):
             cs = mie.calc_cross_sections(m[i], x[i])
             cscat_loop[i], cext_loop[i], cback_loop[i], \
-                cabs_loop[i], asym_loop[i] = (c for c in cs)
+                cabs_loop[i], asym_loop[i] = (c.squeeze() for c in cs)
         assert_equal(cscat, cscat_loop)
         assert_equal(cext, cext_loop)
         assert_equal(cback, cback_loop)
@@ -855,20 +846,16 @@ def test_vectorized_calc_ang_scat(self, num_wavelen, num_layer):
         """
         m, x = mx(num_wavelen, num_layer, **self.mxargs)
         form_factor = mie.calc_ang_scat(m, x, self.angles)
-        if num_wavelen == 1:
-            expected_shape = (2, self.num_angle,)
-            assert form_factor.shape == expected_shape
-            # no further test required since there is only one wavelength
-            return
-        else:
-            expected_shape = (2, num_wavelen, self.num_angle)
-            assert form_factor.shape == expected_shape
+        expected_shape = (2, num_wavelen, self.num_angle)
+        assert form_factor.shape == expected_shape
 
         # we should get same values from loop
-        iparperp_loop = np.zeros(expected_shape, dtype=float)
+        iparperp_loop = []
         for i in range(num_wavelen):
             iparperp = mie.calc_ang_scat(m[i], x[i], self.angles)
-            iparperp_loop[:, i] = iparperp
+            iparperp_loop.append(iparperp)
+        # concatenate along wavelength axis
+        iparperp_loop = np.concatenate(iparperp_loop, axis=1)
         assert_equal(form_factor, iparperp_loop)
 
         # check vectorization for Rayleigh-Gans approximation

pymie/tests/test_multilayer.py

@@ -92,6 +92,8 @@ def test_sooty_particles():
     def efficiencies_from_scat_units(m, x):
         asbs = mie._scatcoeffs_multi(m, x)
         qs = np.array(mie._cross_sections(*asbs)) * 2 / x_L**2
+        # squeeze out singleton wavelen dimension
+        qs = qs.squeeze()
         # there is a factor of 2 conventional difference between
         # "backscattering" and "radar backscattering" efficiencies.
         return np.array([qs[1], qs[0], qs[2]*4*np.pi/2.])