Fix radial functions in monofrequent case

Author: Till Rettberg

micarray/modal/radial.py

@@ -34,9 +34,7 @@ def spherical_pw(N, k, r, setup):
     n = np.arange(N+1)
 
     bn = weights(N, kr, setup)
-    for i, x in enumerate(kr):
-        bn[i, :] = bn[i, :] * 4*np.pi * (1j)**n
-    return bn
+    return 4*np.pi * (1j)**n * bn
 
 
 def spherical_ps(N, k, r, rs, setup):
@@ -72,11 +70,14 @@ def spherical_ps(N, k, r, rs, setup):
     n = np.arange(N+1)
 
     bn = weights(N, k*r, setup)
+    if len(k) == 1:
+        bn = bn[np.newaxis, :]
+
     for i, x in enumerate(krs):
         hn = special.spherical_jn(n, x) - 1j * special.spherical_yn(n, x)
         bn[i, :] = bn[i, :] * 4*np.pi * (-1j) * hn * k[i]
 
-    return bn
+    return np.squeeze(bn)
 
 
 def weights(N, kr, setup):
@@ -106,6 +107,7 @@ def weights(N, kr, setup):
         Radial weights for all orders up to N and the given wavenumbers.
 
     """
+    kr = util.asarray_1d(kr)
     n = np.arange(N+1)
     bns = np.zeros((len(kr), N+1), dtype=complex)
     for i, x in enumerate(kr):