From 33ecb7ef184c0b2503117e51175e575e7bc537c1 Mon Sep 17 00:00:00 2001 From: damonge Date: Wed, 27 Dec 2023 15:50:16 +0000 Subject: [PATCH] streamlined NmtBin --- pymaster/bins.py | 146 ++++++++++++++++++----------------- pymaster/tests/test_bins.py | 44 ++++------- pymaster/tests/test_covar.py | 8 +- 3 files changed, 93 insertions(+), 105 deletions(-) diff --git a/pymaster/bins.py b/pymaster/bins.py index dc84bcee..b12bbd4e 100644 --- a/pymaster/bins.py +++ b/pymaster/bins.py @@ -2,6 +2,18 @@ import numpy as np +def _get_bpw_arrays_linear(lmax, nlb): + ells = np.arange(lmax+1, dtype=np.int32) + bpws = ((ells-2) // nlb).astype(np.int32) + bpws[:2] = -1 + # Remove last bandpower if smaller + if np.sum(bpws == bpws[-1]) != nlb: + good = bpws != bpws[-1] + ells = ells[good] + bpws = bpws[good] + return ells, bpws + + class NmtBin(object): """ An NmtBin object defines the set of bandpowers used in the \ @@ -12,80 +24,48 @@ class NmtBin(object): parameters (see :meth:`NmtBin.from_nside_linear`, :meth:`NmtBin.from_lmax_linear` and :meth:`Nmt.from_edges`). - :param int nside: HEALPix nside resolution parameter of the \ - maps you intend to correlate. The maximum multipole \ - considered for bandpowers will be 3*nside-1, unless \ - `lmax` is set. :param array-like ells: array of integers corresponding to \ different multipoles :param array-like bpws: array of integers that assign the \ multipoles in ells to different bandpowers. All negative \ values will be ignored. + :param int lmax: integer value corresponding to the maximum \ + multipole considered by these bandpowers (and by any \ + calculation that uses them - e.g. mode-coupling matrices, \ + input unbinned power spectra, etc.). If `None`, the maximum \ + of `ells` will be used. :param array-like weights: array of floats corresponding to \ the weights associated to each multipole in ells. The sum \ - of weights within each bandpower is normalized to 1. - :param int nlb: integer value corresponding to a constant \ - bandpower width. I.e. the bandpowers will be defined as \ - consecutive sets of nlb multipoles from l=2 to l=lmax (see \ - below) with equal weights. If this argument is provided, \ - the values of ells, bpws and weights are ignored. - :param int lmax: integer value corresponding to the maximum \ - multipole used by these bandpowers. If None, it will be set \ - to 3*nside-1. In any case the actual maximum multipole will \ - be chosen as the minimum of lmax, 3*nside-1 and the maximum \ - element of ells (e.g. if you are using CAR maps and don't \ - care about nside, you can pass whatever lmax you want and \ - e.g. nside=lmax). - :param boolean is_Dell: if True, the output of all pseudo-Cl \ - computations carried out using this bandpower scheme (e.g. \ - from :py:meth:`pymaster.workspaces.NmtWorkspace.decouple_cell`) \ - will be multiplied by `ell * (ell + 1) / 2 * PI`, where `ell` \ - is the multipole order (no prefactor otherwise). + of weights within each bandpower is normalized to 1. If \ + `None`, uniform weights are assumed. :param array-like f_ell: if present, this is array represents an \ `ell-dependent` function that will be multiplied by all \ pseudo-Cl computations carried out using this bandpower scheme. \ If not `None`, the value of `is_Dell` is ignored. """ - def __init__(self, nside=None, bpws=None, ells=None, weights=None, - nlb=None, lmax=None, is_Dell=False, f_ell=None): + def __init__(self, bpws, ells, lmax=None, weights=None, + f_ell=None): self.bin = None - if (bpws is None) and (ells is None) and (weights is None) \ - and (nlb is None): - raise KeyError("Must supply bandpower arrays or constant " - "bandpower width") + if (bpws is None) or (ells is None): + raise KeyError("Must supply bandpower arrays") if lmax is None: - if nside is None: - if ells is None: - raise ValueError("Must provide either `lmax`, `nside` " - "or `ells`.") - else: - lmax_in = np.amax(ells) - else: - lmax_in = 3 * nside - 1 - else: - lmax_in = lmax - - if nlb is None: - if (bpws is None) or (ells is None) or (weights is None): - raise KeyError("Must provide bpws, ells and weights") - if f_ell is None: - if is_Dell: - f_ell = ells * (ells + 1.) / (2 * np.pi) - else: - f_ell = np.ones(len(ells)) - self.bin = lib.bins_create_py(bpws.astype(np.int32), - ells.astype(np.int32), - weights, f_ell, int(lmax_in)) - else: - self.bin = lib.bins_constant(nlb, lmax_in, int(is_Dell)) - self.lmax = lmax_in + lmax = np.amax(ells) + self.lmax = int(lmax) + if weights is None: + weights = np.ones(len(ells)) + if f_ell is None: + f_ell = np.ones(len(ells)) + self.bin = lib.bins_create_py(bpws.astype(np.int32), + ells.astype(np.int32), + weights, f_ell, self.lmax) @classmethod - def from_nside_linear(NmtBin, nside, nlb, is_Dell=False): + def from_nside_linear(cls, nside, nlb, is_Dell=False, f_ell=None): """ - Convenience constructor for HEALPix maps with linear binning. + Convenience constructor for HEALPix maps with linear binning, + starting at `ell=2`. :param int nside: HEALPix nside resolution parameter of the \ maps you intend to correlate. The maximum multipole \ @@ -99,13 +79,24 @@ def from_nside_linear(NmtBin, nside, nlb, is_Dell=False): from :py:meth:`pymaster.workspaces.NmtWorkspace.decouple_cell`) \ will be multiplied by `ell * (ell + 1) / 2 * PI`, where `ell` \ is the multipole order (no prefactor otherwise). - """ - return NmtBin(nside=nside, nlb=nlb, is_Dell=is_Dell) + :param array-like f_ell: if present, this is array represents an \ + `ell-dependent` function that will be multiplied by all \ + pseudo-Cl computations carried out using this bandpower scheme. \ + If not `None`, the value of `is_Dell` is ignored. If provided, \ + it must be sampled at all ells up to (and including) 3*nside-1. + """ + ells, bpws = _get_bpw_arrays_linear(3*nside-1, nlb) + weights = np.ones(len(ells)) + if is_Dell and (f_ell is None): + f_ell = ells * (ells+1) / (2*np.pi) + return cls(lmax=3*nside-1, bpws=bpws, ells=ells, weights=weights, + f_ell=f_ell) @classmethod - def from_lmax_linear(NmtBin, lmax, nlb, is_Dell=False): + def from_lmax_linear(cls, lmax, nlb, is_Dell=False, f_ell=None): """ - Convenience constructor for generic linear binning. + Convenience constructor for generic linear binning, starting at + `ell=2`. :param int lmax: integer value corresponding to the maximum \ multipole used by these bandpowers. @@ -118,11 +109,21 @@ def from_lmax_linear(NmtBin, lmax, nlb, is_Dell=False): from :py:meth:`pymaster.workspaces.NmtWorkspace.decouple_cell`) \ will be multiplied by `ell * (ell + 1) / 2 * PI`, where `ell` \ is the multipole order (no prefactor otherwise). - """ - return NmtBin(lmax=lmax, nlb=nlb, is_Dell=is_Dell) + :param array-like f_ell: if present, this is array represents an \ + `ell-dependent` function that will be multiplied by all \ + pseudo-Cl computations carried out using this bandpower scheme. \ + If not `None`, the value of `is_Dell` is ignored. If provided, \ + it must be sampled at all ells up to (and including) lmax. + """ + ells, bpws = _get_bpw_arrays_linear(lmax, nlb) + weights = np.ones(len(ells)) + if is_Dell and (f_ell is None): + f_ell = ells * (ells+1) / (2*np.pi) + return cls(lmax=lmax, bpws=bpws, ells=ells, weights=weights, + f_ell=f_ell) @classmethod - def from_edges(NmtBin, ell_ini, ell_end, is_Dell=False): + def from_edges(cls, ell_ini, ell_end, is_Dell=False, f_ell=None): """ Convenience constructor for general equal-weight bands. All ells in the interval [ell_ini, ell_end) will be @@ -137,8 +138,13 @@ def from_edges(NmtBin, ell_ini, ell_end, is_Dell=False): from :py:meth:`pymaster.workspaces.NmtWorkspace.decouple_cell`) \ will be multiplied by `ell * (ell + 1) / 2 * PI`, where `ell` \ is the multipole order (no prefactor otherwise). + :param array-like f_ell: if present, this is array represents an \ + `ell-dependent` function that will be multiplied by all \ + pseudo-Cl computations carried out using this bandpower scheme. \ + If not `None`, the value of `is_Dell` is ignored. If provided, \ + it must be sampled at (and only at) the ells covered by `ell_ini` + and `ell_end`. """ - nls = int(np.amax(ell_end)) ells, bpws, weights = [], [], [] for ib, (li, le) in enumerate(zip(ell_ini, ell_end)): nlb = int(le - li) @@ -148,14 +154,13 @@ def from_edges(NmtBin, ell_ini, ell_end, is_Dell=False): ells = np.array(ells) bpws = np.array(bpws) weights = np.array(weights) - return NmtBin(bpws=bpws, - ells=ells, - weights=weights, - lmax=nls-1, - is_Dell=is_Dell) + if is_Dell and (f_ell is None): + f_ell = ells * (ells+1) / (2*np.pi) + return cls(lmax=np.amax(ells), bpws=bpws, ells=ells, weights=weights, + f_ell=f_ell) def __del__(self): - if self.bin is not None: + if getattr(self, 'bin', None) is not None: if lib.bins_free is not None: lib.bins_free(self.bin) self.bin = None @@ -260,7 +265,8 @@ def unbin_cell(self, cls_in): cls_in = np.array([cls_in]) if (cls_in.ndim > 2) or (len(cls_in[0]) != self.bin.n_bands): raise ValueError("Input Cl has wrong size") - cl1d = lib.unbin_cl(self.bin, cls_in, len(cls_in) * (self.lmax + 1)) + cl1d = lib.unbin_cl(self.bin, cls_in, + int(len(cls_in) * (self.lmax + 1))) clout = np.reshape(cl1d, [len(cls_in), self.lmax + 1]) if oned: clout = clout[0] diff --git a/pymaster/tests/test_bins.py b/pymaster/tests/test_bins.py index 3420c474..2c480d9e 100644 --- a/pymaster/tests/test_bins.py +++ b/pymaster/tests/test_bins.py @@ -9,28 +9,17 @@ def __init__(self): self.nside = 1024 self.lmax = 2000 self.nlb = 4 - self.bc = nmt.NmtBin(self.nside, nlb=4, lmax=self.lmax) + self.bc = nmt.NmtBin.from_lmax_linear(lmax=self.lmax, nlb=4) ells = np.arange(self.lmax - 4, dtype=int)+2 bpws = (ells - 2)//4 weights = 0.25*np.ones(self.lmax - 4) fell = ells*(ells+1.)/(2*np.pi) - self.bv = nmt.NmtBin(nside=self.nside, - bpws=bpws, ells=ells, - weights=weights, - lmax=self.lmax) - self.bcf = nmt.NmtBin(nside=self.nside, - nlb=4, lmax=self.lmax, - is_Dell=True) - self.bvf1 = nmt.NmtBin(nside=self.nside, - bpws=bpws, ells=ells, - weights=weights, - lmax=self.lmax, - is_Dell=True) - self.bvf2 = nmt.NmtBin(nside=self.nside, - bpws=bpws, ells=ells, - weights=weights, - lmax=self.lmax, - f_ell=fell) + self.bv = nmt.NmtBin(lmax=self.lmax, bpws=bpws, ells=ells, + weights=weights) + self.bcf = nmt.NmtBin.from_lmax_linear(lmax=self.lmax, nlb=4, + is_Dell=True) + self.bvf = nmt.NmtBin(lmax=self.lmax, bpws=bpws, ells=ells, + weights=weights, f_ell=fell) self.l_edges = np.arange(2, self.lmax+2, 4, dtype=int) self.be = nmt.NmtBin.from_edges(self.l_edges[:-1], self.l_edges[1:]) @@ -45,21 +34,16 @@ def test_bins_errors(): weights = 0.25*np.ones(BT.lmax - 4) weights[16:20] = 0 with pytest.raises(RuntimeError): - nmt.NmtBin(nside=BT.nside, + nmt.NmtBin(lmax=BT.lmax, bpws=bpws, ells=ells, - weights=weights, - lmax=BT.lmax) + weights=weights) with pytest.raises(ValueError): BT.bv.bin_cell(np.random.randn(3, 3, 3)) with pytest.raises(ValueError): BT.bv.unbin_cell(np.random.randn(3, 3, 3)) - with pytest.raises(KeyError): + with pytest.raises(TypeError): nmt.NmtBin() - with pytest.raises(ValueError): - nmt.NmtBin(nlb=10) - with pytest.raises(KeyError): - nmt.NmtBin(nside=16, weights=1) def test_bins_nell_list(): @@ -89,7 +73,7 @@ def test_bins_constant(): # Tests constant bandpower initialization assert (BT.bc.get_n_bands() == (BT.lmax - 2)//BT.nlb) assert (BT.bc.get_ell_list(5)[2] == 2+BT.nlb*5+2) - b = nmt.NmtBin(nside=1024, nlb=4, lmax=2000) + b = nmt.NmtBin.from_lmax_linear(lmax=2000, nlb=4) assert (b.bin.ell_max == 2000) @@ -137,11 +121,9 @@ def test_bins_binning_f_ell(): cls = np.arange(BT.lmax+1, dtype=float) fell = cls * (cls + 1.) / 2 / np.pi cl_b = BT.bcf.bin_cell(cls) - assert normdiff(cl_b, BT.bvf1.bin_cell(cls)) < 1E-5 - assert normdiff(cl_b, BT.bvf2.bin_cell(cls)) < 1E-5 + assert normdiff(cl_b, BT.bvf.bin_cell(cls)) < 1E-5 cl_u = BT.bcf.unbin_cell(cl_b) - assert normdiff(cl_u, BT.bvf1.unbin_cell(cl_b)) < 1E-5 - assert normdiff(cl_u, BT.bvf2.unbin_cell(cl_b)) < 1E-5 + assert normdiff(cl_u, BT.bvf.unbin_cell(cl_b)) < 1E-5 iend = 2+BT.nlb*((BT.lmax - 2)//BT.nlb) cl_b_p = np.mean((fell*cls)[2:iend].reshape([-1, BT.nlb]), axis=1) assert normdiff(cl_b_p, cl_b) < 1E-5 diff --git a/pymaster/tests/test_covar.py b/pymaster/tests/test_covar.py index bfb91133..72a73429 100644 --- a/pymaster/tests/test_covar.py +++ b/pymaster/tests/test_covar.py @@ -16,10 +16,10 @@ def __init__(self): self.nside = 64 self.nlb = 16 self.npix = hp.nside2npix(self.nside) - msk = hp.read_map("test/benchmarks/msk.fits", verbose=False, + msk = hp.read_map("test/benchmarks/msk.fits", dtype=float) mps = np.array(hp.read_map("test/benchmarks/mps.fits", - verbose=False, field=[0, 1, 2], + field=[0, 1, 2], dtype=float)) self.b = nmt.NmtBin.from_nside_linear(self.nside, self.nlb) self.f0 = nmt.NmtField(msk, [mps[0]]) @@ -225,8 +225,8 @@ def test_covar_rectangular(): npix = hp.nside2npix(nside) msk = np.ones(npix) f = nmt.NmtField(msk, None, spin=0) - b1 = nmt.NmtBin(nside, nlb=4) - b2 = nmt.NmtBin(nside, nlb=6) + b1 = nmt.NmtBin.from_nside_linear(nside, nlb=4) + b2 = nmt.NmtBin.from_nside_linear(nside, nlb=6) w1 = nmt.NmtWorkspace() w1.compute_coupling_matrix(f, f, b1) w2 = nmt.NmtWorkspace()