Finished

src/powerbox/tools.py

@@ -321,7 +321,7 @@ def generator():
     return generator()
 
 
-def regular_angular_generator(bins, dims2avg, angular_resolution=0.1):
+def regular_angular_generator(bins, dims2avg, angular_resolution=0.05):
     r"""
     Returns a set of spherical coordinates regularly sampled at a given angular resolution.
 
@@ -399,7 +399,6 @@ def _sample_coords_interpolate(coords, bins, weights, interp_points_generator):
     else:
         sample_coords = bins.reshape(1, -1)
         r_n = bins
-
     # Remove sample coords that are not even on the coords grid (e.g. due to phi)
     mask1 = np.all(
         sample_coords >= np.array([c.min() for c in coords])[..., np.newaxis], axis=0
@@ -424,10 +423,6 @@ def _field_average_interpolate(coords, field, bins, weights, sample_coords, r_n)
         weights = weights.reshape(field.shape)
     else:
         weights = np.ones_like(field) * weights
-    # if field.dtype.kind == "c":
-    #    field = np.array(field, dtype=np.complex64)
-    # else:
-    #    field = np.array(field, dtype=np.float32)
     # Set 0 weights to NaNs
     field = field * weights
     field[weights == 0] = np.nan
@@ -449,14 +444,22 @@ def _field_average_interpolate(coords, field, bins, weights, sample_coords, r_n)
         warnings.warn("Interpolator returned all NaNs.", stacklevel=2)
     # Average over the spherical shells for each radius / bin value
     avged_field = np.array([np.nanmean(interped_field[r_n == b]) for b in bins])
-    sumweights = np.unique(r_n[~np.isnan(interped_field)], return_counts=True)[1]
-    return avged_field, sumweights
+    unique_rn, sumweights = np.unique(
+        r_n[~np.isnan(interped_field)], return_counts=True
+    )
+    final_sumweights = []
+    for b in bins:
+        if b in unique_rn:
+            final_sumweights.append(sumweights[unique_rn == b][0])
+        else:
+            final_sumweights.append(0)
+    return avged_field, np.array(final_sumweights)
 
 
 def _field_average(indx, field, weights, sumweights):
     if not np.isscalar(weights) and field.shape != weights.shape:
         raise ValueError(
-            "the field and weights must have the same shape!",
+            "The field and weights must have the same shape!",
             field.shape,
             weights.shape,
         )
@@ -801,8 +804,7 @@ def ignore_zero_ki(freq: list, kmag: np.ndarray = None):
     res_ndim = len(kmag.shape)
 
     coords = np.array(np.meshgrid(*freq[:res_ndim], sparse=False))
-    k_weights = np.all(coords != 0, axis=0)
-
+    k_weights = ~np.any(coords == 0, axis=0)
     return k_weights
 
 

tests/test_tools.py

@@ -20,8 +20,13 @@ def test_angular_avg_nd_3(interpolation_method):
     P = r2**-1.0
     P = np.repeat(P, 100).reshape(400, 400, 100)
     freq = [x, x, np.linspace(-2, 2, 100)]
-    p_k, k_av_bins = angular_average_nd(
-        P, freq, bins=50, n=2, interpolation_method=interpolation_method
+    p_k, k_av_bins, sw = angular_average_nd(
+        P,
+        freq,
+        bins=50,
+        n=2,
+        interpolation_method=interpolation_method,
+        return_sumweights=True,
     )
     if interpolation_method == "linear":
         assert np.max(np.abs((p_k[:, 0] - k_av_bins**-2.0) / k_av_bins**-2.0)) < 0.05
@@ -35,6 +40,21 @@ def test_angular_avg_nd_3(interpolation_method):
         )
 
 
+def test_weights_shape():
+    x = np.linspace(-3, 3, 40)
+    P = np.ones(3 * [40])
+    weights = np.ones(3 * [20])
+    freq = [x for _ in range(3)]
+
+    with pytest.raises(ValueError):
+        p_k_lin, k_av_bins_lin = angular_average(
+            P,
+            freq,
+            bins=10,
+            weights=weights,
+        )
+
+
 @pytest.mark.parametrize("n", range(1, 5))
 def test_interp_w_weights(n):
     x = np.linspace(-3, 3, 40)
@@ -67,16 +87,45 @@ def test_interp_w_weights(n):
         interpolation_method="linear",
         weights=weights,
         interp_points_generator=regular_angular_generator,
+        log_bins=True,
     )
 
     assert np.all(p_k_lin == 1.0)
 
 
-def test_interp_w_mu():
+@pytest.mark.parametrize("n", range(1, 3))
+def test_zero_ki(n):
+    x = np.arange(-100, 100, 1)
+    from powerbox.tools import ignore_zero_ki
+
+    # needed only for shape
+    freq = n * [x]
+    coords = np.array(np.meshgrid(*freq))
+    kmag = np.sqrt(np.sum(coords**2, axis=0))
+    weights = ignore_zero_ki(freq, kmag)
+    L = x[-1] - x[0] + 1
+    masked_points = np.sum(weights == 0)
+    if n == 1:
+        assert masked_points == 1
+    elif n == 2:
+        assert masked_points == n * L - 1
+    elif n == 3:
+        assert masked_points == n * L**2 - n * L + 1
+    else:
+        assert masked_points == n * L**3 - n * L**2 + n * L - 1
+
+
+@pytest.mark.parametrize("n", range(2, 3))
+def test_interp_w_mu(n):
     x = np.linspace(0.0, 3, 40)
-    kpar_mesh, kperp_mesh = np.meshgrid(x, x)
-    theta = np.arctan(kperp_mesh / kpar_mesh)
-    mu_mesh = np.cos(theta)
+    if n == 2:
+        kpar_mesh, kperp_mesh = np.meshgrid(x, x)
+        theta = np.arctan(kperp_mesh / kpar_mesh)
+        mu_mesh = np.cos(theta)
+    else:
+        kx_mesh, ky_mesh, kz_mesh = np.meshgrid(x, x, x, indexing="ij")
+        theta = np.arccos(kz_mesh / np.sqrt(kx_mesh**2 + ky_mesh**2 + kz_mesh**2))
+        mu_mesh = np.cos(theta)
 
     # Need a little cushion so we test against data at mu = 0.95
     # If we test for mu that is higher (default is mu >= 0.97)
@@ -88,7 +137,7 @@ def test_interp_w_mu():
 
     p_k_lin, k_av_bins_lin = angular_average(
         P,
-        [x, x],
+        n * [x],
         bins=10,
         interpolation_method="linear",
         weights=1.0,
@@ -317,7 +366,6 @@ def test_variance_2d():
     ave, coord, var = angular_average(
         P, np.sqrt(r2), bins=np.linspace(0, x.max(), 20), get_variance=True
     )
-    print(np.diff(var))
     assert np.all(np.diff(var) <= 0)
 
 
@@ -350,6 +398,9 @@ def test_sum():
     ave, coord = angular_average(P, np.sqrt(r2), bins=20, bin_ave=False, average=False)
     assert np.sum(P[r2 < 9.0]) == np.sum(ave)
 
+    ave, coord = angular_average(P, np.sqrt(r2), bins=20, bin_ave=True, average=False)
+    assert np.sum(P[r2 < 9.0]) == np.sum(ave)
+
 
 def test_var_trivial_weights():
     x = np.linspace(-3, 3, 400)