diff --git a/src/powerbox/tools.py b/src/powerbox/tools.py index d03d500..e349f92 100644 --- a/src/powerbox/tools.py +++ b/src/powerbox/tools.py @@ -14,21 +14,47 @@ from . import dft -def _getbins(bins, coords, log): - if not np.iterable(bins): +def _getbins( + bins: np.ndarray | int, + coord_mags: np.ndarray, + log: bool, + bins_upto_boxlen: bool | None = None, +): + if np.iterable(bins): + return bins + + if bins_upto_boxlen is None: + warnings.warn( + ( + "In the future, bins will be generated by default up to the smallest " + "length over any dimension, instead of the largest magnitude for the box." + "Set bins_upto_boxlen to silence this warning." + ), + stacklevel=2, + category=FutureWarning, + ) + bins_upto_boxlen = False + + if bins_upto_boxlen: try: # Fails if coords is not a cube / inhomogeneous. - max_radius = np.min([np.max(coords, axis=i) for i in range(coords.ndim)]) + max_radius = np.min( + [np.max(coord_mags, axis=i) for i in range(coord_mags.ndim)] + ) except ValueError: - maxs = [np.max(coords, axis=i) for i in range(coords.ndim)] + maxs = [np.max(coord_mags, axis=i) for i in range(coord_mags.ndim)] maxs_flat = [] [maxs_flat.extend(m.ravel()) for m in maxs] max_radius = np.min(maxs_flat) - if not log: - bins = np.linspace(coords.min(), max_radius, bins + 1) - else: - mn = coords[coords > 0].min() - bins = np.logspace(np.log10(mn), np.log10(max_radius), bins + 1) + else: + max_radius = coord_mags.max() + + if not log: + bins = np.linspace(coord_mags.min(), max_radius, bins + 1) + else: + mn = coord_mags[coord_mags > 0].min() + bins = np.logspace(np.log10(mn), np.log10(max_radius), bins + 1) + return bins @@ -44,6 +70,7 @@ def angular_average( interpolation_method=None, interp_points_generator=None, return_sumweights=False, + bins_upto_boxlen: bool | None = None, ): r""" Average a given field within radial bins. @@ -107,6 +134,12 @@ def angular_average( (which can be adjusted by supplying a different interp_points_generator function with a different angular resolution). + bins_upto_boxlen : bool, optional + If set to True and the bins are determined automatically, calculate bins only + up to the maximum k along any dimension. Otherwise, calculate bins up to the + maximum magnitude of k (i.e. a factor of sqrt(ndim) higher). Default is False + for backwards compatibility. + Returns ------- field_1d : 1D-array @@ -149,7 +182,7 @@ def angular_average( raise ValueError("Only linear interpolation is supported.") if len(coords) == len(field.shape): # coords are a segmented list of dimensional co-ordinates - coords_grid = _magnitude_grid(coords) + coord_mags = _magnitude_grid(coords) elif interpolation_method is not None: raise ValueError( "Must supply a list of len(field.shape) of 1D coordinate arrays for coords when interpolating!" @@ -157,10 +190,17 @@ def angular_average( else: # coords are the magnitude of the co-ordinates # since we are not interpolating, then we can just use the magnitude of the co-ordinates - coords_grid = coords + coord_mags = coords + if interpolation_method is None: indx, bins, sumweights = _get_binweights( - coords_grid, weights, bins, average, bin_ave=bin_ave, log_bins=log_bins + coord_mags, + weights, + bins, + average, + bin_ave=bin_ave, + log_bins=log_bins, + bins_upto_boxlen=bins_upto_boxlen, ) if np.any(sumweights == 0): @@ -169,7 +209,8 @@ def angular_average( ) res = _field_average(indx, field, weights, sumweights) else: - bins = _getbins(bins, coords_grid, log_bins) + bins = _getbins(bins, coord_mags, log_bins, bins_upto_boxlen) + if bin_ave: if log_bins: bins = np.exp((np.log(bins[1:]) + np.log(bins[:-1])) / 2) @@ -207,18 +248,26 @@ def _magnitude_grid(x, dim=None): return np.sqrt(np.sum(np.meshgrid(*([X**2 for X in x]), indexing="ij"), axis=0)) -def _get_binweights(coords, weights, bins, average=True, bin_ave=True, log_bins=False): +def _get_binweights( + coord_mags, + weights, + bins, + average=True, + bin_ave=True, + log_bins=False, + bins_upto_boxlen: bool | None = None, +): # Get a vector of bin edges - bins = _getbins(bins, coords, log_bins) + bins = _getbins(bins, coord_mags, log_bins, bins_upto_boxlen=bins_upto_boxlen) - indx = np.digitize(coords.flatten(), bins) + indx = np.digitize(coord_mags.flatten(), bins) if average or bin_ave: if not np.isscalar(weights): - if coords.shape != weights.shape: + if coord_mags.shape != weights.shape: raise ValueError( "coords and weights must have the same shape!", - coords.shape, + coord_mags.shape, weights.shape, ) sumweights = np.bincount( @@ -236,7 +285,9 @@ def _get_binweights(coords, weights, bins, average=True, bin_ave=True, log_bins= if bin_ave: bins = ( np.bincount( - indx, weights=(weights * coords).flatten(), minlength=len(bins) + 1 + indx, + weights=(weights * coord_mags).flatten(), + minlength=len(bins) + 1, )[1:-1] / binweight ) @@ -557,6 +608,7 @@ def angular_average_nd( # noqa: C901 interpolation_method=None, interp_points_generator=None, return_sumweights=False, + bins_upto_boxlen: bool | None = None, ): """ Average the first n dimensions of a given field within radial bins. @@ -621,6 +673,12 @@ def angular_average_nd( # noqa: C901 (which can be adjusted by supplying a different interp_points_generator function with a different angular resolution). + bins_upto_boxlen : bool, optional + If set to True and the bins are determined automatically, calculate bins only + up to the maximum k along any dimension. Otherwise, calculate bins up to the + maximum magnitude of k (i.e. a factor of sqrt(ndim) higher). Default is False + for backwards compatibility. + Returns ------- field : (m-n+1)-array @@ -683,11 +741,12 @@ def angular_average_nd( # noqa: C901 interpolation_method=interpolation_method, interp_points_generator=interp_points_generator, return_sumweights=return_sumweights, + bins_upto_boxlen=bins_upto_boxlen, ) if len(coords) == len(field.shape): # coords are a segmented list of dimensional co-ordinates - coords_grid = _magnitude_grid([c for i, c in enumerate(coords) if i < n]) + coord_mags = _magnitude_grid([c for i, c in enumerate(coords) if i < n]) elif interpolation_method is not None: raise ValueError( "Must supply a list of len(field.shape) of 1D coordinate arrays for coords when interpolating!" @@ -695,18 +754,24 @@ def angular_average_nd( # noqa: C901 else: # coords are the magnitude of the co-ordinates # since we are not interpolating, then we can just use the magnitude of the co-ordinates - coords_grid = coords + coord_mags = coords - coords_grid = _magnitude_grid([c for i, c in enumerate(coords) if i < n]) + coord_mags = _magnitude_grid([c for i, c in enumerate(coords) if i < n]) n1 = np.prod(field.shape[:n]) n2 = np.prod(field.shape[n:]) if interpolation_method is None: indx, bins, sumweights = _get_binweights( - coords_grid, weights, bins, average, bin_ave=bin_ave, log_bins=log_bins + coord_mags, + weights, + bins, + average, + bin_ave=bin_ave, + log_bins=log_bins, + bins_upto_boxlen=bins_upto_boxlen, ) res = np.zeros((len(sumweights), n2), dtype=field.dtype) if interpolation_method is not None: - bins = _getbins(bins, coords_grid, log_bins) + bins = _getbins(bins, coord_mags, log_bins, bins_upto_boxlen) if bin_ave: if log_bins: bins = np.exp((np.log(bins[1:]) + np.log(bins[:-1])) / 2) @@ -977,6 +1042,7 @@ def get_power( interpolation_method=None, interp_points_generator=None, return_sumweights=False, + bins_upto_boxlen: bool | None = None, ): r""" Calculate isotropic power spectrum of a field, or cross-power of two similar fields. @@ -1065,6 +1131,11 @@ def get_power( Note that for the linear interpolation case, this corresponds to the number of samples averaged over (which can be adjusted with the angular_resolution parameter). + bins_upto_boxlen : bool, optional + If set to True and the bins are determined automatically, calculate bins only + up to the maximum k along any dimension. Otherwise, calculate bins up to the + maximum magnitude of k (i.e. a factor of sqrt(ndim) higher). Default is False + for backwards compatibility. Returns ------- @@ -1183,6 +1254,7 @@ def get_power( interpolation_method=interpolation_method, interp_points_generator=interp_points_generator, return_sumweights=return_sumweights, + bins_upto_boxlen=bins_upto_boxlen, ) res = list(res) # Remove shot-noise diff --git a/tests/test_discrete.py b/tests/test_discrete.py index 32cc4d8..9d99dab 100644 --- a/tests/test_discrete.py +++ b/tests/test_discrete.py @@ -1,9 +1,12 @@ import pytest import numpy as np +from functools import partial from powerbox import LogNormalPowerBox, PowerBox, get_power +get_power = partial(get_power, bins_upto_boxlen=True) + def test_discrete_power_gaussian(): pb = PowerBox( diff --git a/tests/test_lognormal.py b/tests/test_lognormal.py index 708b757..939e5b9 100644 --- a/tests/test_lognormal.py +++ b/tests/test_lognormal.py @@ -1,7 +1,10 @@ import numpy as np +from functools import partial from powerbox import LogNormalPowerBox, PowerBox, get_power +get_power = partial(get_power, bins_upto_boxlen=True) + def test_ln_vs_straight(): # Set up two boxes with exactly the same parameters diff --git a/tests/test_power.py b/tests/test_power.py index 931b307..0fcaba4 100644 --- a/tests/test_power.py +++ b/tests/test_power.py @@ -1,8 +1,11 @@ import numpy as np import warnings +from functools import partial from powerbox import PowerBox, get_power, ignore_zero_absk, ignore_zero_ki, power2delta +get_power = partial(get_power, bins_upto_boxlen=True) + def test_power1d(): p = [0] * 40 diff --git a/tests/test_tools.py b/tests/test_tools.py index 47c5ee8..0272881 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -2,9 +2,12 @@ import numpy as np import warnings +from functools import partial from powerbox.powerbox import PowerBox from powerbox.tools import ( + _getbins, + _magnitude_grid, above_mu_min_angular_generator, angular_average, angular_average_nd, @@ -12,6 +15,10 @@ regular_angular_generator, ) +get_power = partial(get_power, bins_upto_boxlen=True) +angular_average = partial(angular_average, bins_upto_boxlen=True) +angular_average_nd = partial(angular_average_nd, bins_upto_boxlen=True) + def test_warn_interp_weights(): x = np.linspace(-3, 3, 400) @@ -32,6 +39,34 @@ def test_warn_interp_weights(): ) +def test_bins_upto_boxlen_warning(): + with pytest.warns( + FutureWarning, + match="In the future, bins will be generated by default up to the smallest", + ): + _getbins(bins=10, coord_mags=np.ones((10, 10)), log=False) + + +@pytest.mark.parametrize("xmax", [1, 10, np.pi]) +@pytest.mark.parametrize("ndim", [1, 2, 3]) +def test_bins_upto_boxlen(xmax, ndim): + x = np.linspace(-xmax, xmax, 21) + mag = _magnitude_grid([x] * ndim) + + bins = _getbins(bins=12, coord_mags=mag, log=False, bins_upto_boxlen=True) + assert bins.max() == xmax + + +@pytest.mark.parametrize("xmax", [1, 10, np.pi]) +@pytest.mark.parametrize("ndim", [1, 2, 3]) +def test_bins_upto_maxmag(xmax, ndim): + x = np.linspace(-xmax, xmax, 21) + mag = _magnitude_grid([x] * ndim) + + bins = _getbins(bins=12, coord_mags=mag, log=False, bins_upto_boxlen=False) + assert np.isclose(bins.max(), xmax * np.sqrt(ndim)) + + @pytest.mark.parametrize("interpolation_method", [None, "linear"]) def test_angular_avg_nd_3(interpolation_method): x = np.linspace(-3, 3, 400) @@ -47,6 +82,7 @@ def test_angular_avg_nd_3(interpolation_method): n=2, interpolation_method=interpolation_method, return_sumweights=True, + bins_upto_boxlen=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 @@ -68,10 +104,7 @@ def test_weights_shape(): with pytest.raises(ValueError): p_k_lin, k_av_bins_lin = angular_average( - P, - freq, - bins=10, - weights=weights, + P, freq, bins=10, weights=weights, bins_upto_boxlen=True ) @@ -108,6 +141,7 @@ def test_interp_w_weights(n): weights=weights, interp_points_generator=regular_angular_generator(), log_bins=True, + bins_upto_boxlen=True, ) assert np.all(p_k_lin == 1.0) @@ -166,6 +200,7 @@ def test_interp_w_mu(n): interpolation_method="linear", weights=1.0, interp_points_generator=above_mu_min_angular_generator(), + bins_upto_boxlen=True, ) # Start from the 4th bin due to the average being a bit < 1 at low radii assert np.all(p_k_lin[3:] == 1.0) @@ -188,6 +223,7 @@ def test_error_coords_and_mask(): interpolation_method="linear", weights=1.0, interp_points_generator=above_mu_min_angular_generator(mu=0.97), + bins_upto_boxlen=True, ) @@ -196,36 +232,40 @@ def test_interp_method(): P = np.ones((40, 40, 40)) freq = [x, x, x] with pytest.raises(ValueError): - ave, coord, var = angular_average_nd( - P, freq, bins=20, get_variance=True, interpolation_method="abc" + angular_average_nd( + P, + freq, + bins=20, + get_variance=True, + interpolation_method="abc", + bins_upto_boxlen=True, ) with pytest.raises(ValueError): - ave, coord, var = angular_average( - P, freq, bins=20, get_variance=True, interpolation_method="abc" + angular_average( + P, + freq, + bins=20, + get_variance=True, + interpolation_method="abc", + bins_upto_boxlen=True, ) def test_error_w_kmag_coords(): + x = np.linspace(-3, 3, 40) + P = np.ones((40, 40, 40)) + X, Y = np.meshgrid(x, x) + with pytest.raises(ValueError): - x = np.linspace(-3, 3, 40) - P = np.ones((40, 40, 40)) - X, Y = np.meshgrid(x, x) - ( - ave, - coord, - ) = angular_average_nd( - P, X**2 + Y**2, bins=20, interpolation_method="linear" - ) + angular_average_nd(P, X**2 + Y**2, bins=20, interpolation_method="linear") + + x = np.linspace(-3, 3, 40) + P = np.ones((40, 40, 40)) + X, Y = np.meshgrid(x, x) with pytest.raises(ValueError): - x = np.linspace(-3, 3, 40) - P = np.ones((40, 40, 40)) - X, Y = np.meshgrid(x, x) - ( - ave, - coord, - ) = angular_average(P, X**2 + Y**2, bins=20, interpolation_method="linear") + angular_average(P, X**2 + Y**2, bins=20, interpolation_method="linear") def test_kmag_coords_nointerp(): @@ -233,19 +273,9 @@ def test_kmag_coords_nointerp(): P = np.ones((40, 40, 40)) X, Y = np.meshgrid(x, x) with pytest.raises(ValueError): - ( - ave, - coord, - ) = angular_average_nd( - P, np.sqrt(X**2 + Y**2), bins=20, interpolation_method=None - ) + angular_average_nd(P, np.sqrt(X**2 + Y**2), bins=20, interpolation_method=None) with pytest.raises(ValueError): - ( - ave, - coord, - ) = angular_average( - P, np.sqrt(X**2 + Y**2), bins=20, interpolation_method=None - ) + angular_average(P, np.sqrt(X**2 + Y**2), bins=20, interpolation_method=None) @pytest.mark.parametrize("n", range(1, 3)) @@ -260,7 +290,7 @@ def test_angular_avg_nd(n): freq = [x, x, x, np.linspace(-2, 2, 10)] p_k_lin, k_av_bins_lin = angular_average_nd( - P, freq, bins=10, n=n, interpolation_method="linear" + P, freq, bins=10, n=n, interpolation_method="linear", bins_upto_boxlen=True ) if n == 1: @@ -304,7 +334,7 @@ def test_angular_avg_nd_complex_interp(): P = np.repeat(P, 100).reshape(400, 400, 100) freq = [x, x, np.linspace(-2, 2, 100)] p_k_lin, k_av_bins_lin = angular_average_nd( - P, freq, bins=50, n=2, interpolation_method="linear" + P, freq, bins=50, n=2, interpolation_method="linear", bins_upto_boxlen=True ) real = np.real(p_k_lin) imag = np.imag(p_k_lin) @@ -327,7 +357,12 @@ def test_angular_avg_nd_4_2(interpolation_method): freq = [x, x, np.linspace(-2, 2, 10), np.linspace(-2, 2, 10)] p_k, k_av_bins = angular_average_nd(P, freq, bins=50, n=2) p_k_lin, k_av_bins_lin = angular_average_nd( - P, freq, bins=50, n=2, interpolation_method=interpolation_method + P, + freq, + bins=50, + n=2, + interpolation_method=interpolation_method, + bins_upto_boxlen=True, ) # The radially-averaged power is not very accurate # due to the low number of bins at small values of k_av_bins, so we start @@ -364,7 +399,13 @@ def test_angular_avg_nd_2_1_varnull(): coords = [x, np.linspace(-2, 2, 10)] p_k, k_av_bins, var, sw = angular_average_nd( - P, coords, bins=20, n=1, get_variance=True, return_sumweights=True + P, + coords, + bins=20, + n=1, + get_variance=True, + return_sumweights=True, + bins_upto_boxlen=True, ) assert np.all(var == 0) @@ -376,7 +417,11 @@ def test_null_variance_2d(): r2 = X**2 + Y**2 P = np.ones_like(r2) ave, coord, var = angular_average( - P, np.sqrt(r2), bins=np.linspace(0, x.max(), 20), get_variance=True + P, + np.sqrt(r2), + bins=np.linspace(0, x.max(), 20), + get_variance=True, + bins_upto_boxlen=True, ) assert np.all(var == 0) @@ -388,7 +433,11 @@ def test_variance_2d(): P = np.ones_like(r2) P += np.random.normal(scale=1, size=(len(x), len(x))) ave, coord, var = angular_average( - P, np.sqrt(r2), bins=np.linspace(0, x.max(), 20), get_variance=True + P, + np.sqrt(r2), + bins=np.linspace(0, x.max(), 20), + get_variance=True, + bins_upto_boxlen=True, ) assert np.all(np.diff(var) <= 0) @@ -400,7 +449,11 @@ def test_complex_variance(): P = np.ones_like(r2) + np.ones_like(r2) * 1j with pytest.raises(NotImplementedError): ave, coord, var = angular_average( - P, np.sqrt(r2), bins=np.linspace(0, x.max(), 20), get_variance=True + P, + np.sqrt(r2), + bins=np.linspace(0, x.max(), 20), + get_variance=True, + bins_upto_boxlen=True, ) @@ -410,7 +463,9 @@ def test_bin_edges(): r2 = X**2 + Y**2 P = r2**-1.0 bins = np.linspace(0, x.max(), 20) - ave, coord = angular_average(P, np.sqrt(r2), bins=bins, bin_ave=False) + ave, coord = angular_average( + P, np.sqrt(r2), bins=bins, bin_ave=False, bins_upto_boxlen=True + ) assert np.all(coord == bins) @@ -419,10 +474,14 @@ def test_sum(): X, Y = np.meshgrid(x, x) r2 = X**2 + Y**2 P = r2**-1.0 - ave, coord = angular_average(P, np.sqrt(r2), bins=20, bin_ave=False, average=False) + ave, coord = angular_average( + P, np.sqrt(r2), bins=20, bin_ave=False, average=False, bins_upto_boxlen=True + ) 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) + ave, coord = angular_average( + P, np.sqrt(r2), bins=20, bin_ave=True, average=False, bins_upto_boxlen=True + ) assert np.sum(P[r2 < 9.0]) == np.sum(ave) @@ -438,6 +497,7 @@ def test_var_trivial_weights(): bins=np.linspace(0, x.max(), 20), get_variance=True, weights=np.ones_like(r2), + bins_upto_boxlen=True, ) print(np.diff(var)) assert np.all(np.diff(var) <= 1e-6) @@ -448,7 +508,9 @@ def test_logbins(): X, Y = np.meshgrid(x, x) r2 = X**2 + Y**2 P = np.ones_like(r2) - ave, coord = angular_average(P, np.sqrt(r2), bins=10, bin_ave=False, log_bins=True) + ave, coord = angular_average( + P, np.sqrt(r2), bins=10, bin_ave=False, log_bins=True, bins_upto_boxlen=True + ) assert np.all(np.isclose(np.diff(coord[1:] / coord[:-1]), 0))