diff --git a/croissant/core/healpix.py b/croissant/core/healpix.py index f554189..aa887d0 100644 --- a/croissant/core/healpix.py +++ b/croissant/core/healpix.py @@ -453,7 +453,7 @@ def from_grid( def switch_coords(self, to_coord, loc=None, time=None): to_coord = coord_rep(to_coord) rot = Rotator(coord=[self.coord, to_coord], loc=loc, time=time) - rot.rotate_alm(self.alm, lmax=self.lmax, inplace=True) + self.alm = rot.rotate_alm(self.alm, lmax=self.lmax) self.coord = to_coord def getlm(self, i=None): diff --git a/croissant/core/rotations.py b/croissant/core/rotations.py index b579441..0120abf 100644 --- a/croissant/core/rotations.py +++ b/croissant/core/rotations.py @@ -116,7 +116,7 @@ def __init__( rot=rot, coord=coord, inv=inv, deg=deg, eulertype=eulertype ) - def rotate_alm(self, alm, lmax=None, inplace=False): + def rotate_alm(self, alm, lmax=None, polarized=False): """ Rotate an alm or a list of alms. @@ -126,23 +126,22 @@ def rotate_alm(self, alm, lmax=None, inplace=False): The alm or list of alms to rotate. lmax : int The maximum ell value to rotate. - inplace : bool - If True, the alm is rotated in place. Otherwise, a copy is - rotated and returned. + polarized : bool + If true, the alm is assumed to be a sequence of TEB alms + corresponding to I, Q, U maps, where I is spin-0 and Q, U are + spin-2. In this case, ``alm'' has two dimensions where the first + has size 3. Multiple frequency maps are not yet supported in this + case. Returns ------- rotated_alm : array_like - The rotated alm or list of alms. This is only returned if - inplace=False. + The rotated alm or list of alms. """ - if inplace: - rotated_alm = alm - else: - rotated_alm = np.array(alm, copy=True, dtype=np.complex128) + rotated_alm = np.array(alm, copy=True, dtype=np.complex128) - if rotated_alm.ndim == 1: + if rotated_alm.ndim == 1 or polarized: super().rotate_alm(rotated_alm, lmax=lmax, inplace=True) elif rotated_alm.ndim == 2: # iterate over the list of alms @@ -153,10 +152,9 @@ def rotate_alm(self, alm, lmax=None, inplace=False): f"alm must have 1 or 2 dimensions, not {alm.ndim}." ) - if not inplace: - return rotated_alm + return rotated_alm - def rotate_map_alms(self, m, lmax=None, inplace=False): + def rotate_map_alms(self, m, lmax=None, polarized=False): """ Rotate a map or a list of maps in spherical harmonics space. @@ -166,28 +164,26 @@ def rotate_map_alms(self, m, lmax=None, inplace=False): The map or list of maps to rotate. lmax : int The maximum ell value to rotate. - inplace : bool - If True, the map is rotated in place. Otherwise, a copy is - rotated and returned. + polarized : bool + If true, ``m'' is assumed to be a list of I, Q, U polarizations. + I is spin-0 and Q, U are spin-2. In this case, ``m'' has two + dimensions where the first has size 3. Multiple frequency maps + are not yet supported in this case. Returns ------- rotated_m : np.ndarray - The rotated map or list of maps. This is only returned if - inplace=False. + The rotated map or list of maps. """ npix = m.shape[-1] nside = hp.npix2nside(npix) - alm = map2alm(m, lmax=lmax) - self.rotate_alm(alm, lmax=lmax, inplace=True) - rotated_m = alm2map(alm, nside, lmax=lmax) - if inplace: - m = rotated_m - else: - return rotated_m + alm = map2alm(m, lmax=lmax, polarized=polarized) + alm = self.rotate_alm(alm, lmax=lmax, polarized=polarized) + rotated_m = alm2map(alm, nside, lmax=lmax, polarized=polarized) + return rotated_m - def rotate_map_pixel(self, m, inplace=False): + def rotate_map_pixel(self, m, polarized=False): """ Rotate a map or a list of maps in pixel space. @@ -195,18 +191,19 @@ def rotate_map_pixel(self, m, inplace=False): ----------- m : array-like The map or list of maps to rotate. - inplace : bool - If True, the map is rotated in place. Otherwise, a copy is - rotated and returned. + polarized : bool + If true, ``m'' is assumed to be a list of I, Q, U polarizations. + I is spin-0 and Q, U are spin-2. In this case, ``m'' has two + dimensions where the first has size 3. Multiple frequency maps + are not yet supported in this case. Returns ------- rotated_m : np.ndarray - The rotated map or list of maps. This is only returned if - inplace=False. + The rotated map or list of maps. """ - if m.ndim == 1: + if m.ndim == 1 or polarized: rotated_m = super().rotate_map_pixel(m) elif m.ndim == 2: rotated_m = np.empty_like(m) @@ -214,7 +211,4 @@ def rotate_map_pixel(self, m, inplace=False): rotated_m[i] = super().rotate_map_pixel(m[i]) else: raise ValueError(f"m must have 1 or 2 dimensions, not {m.ndim}.") - if inplace: - m = rotated_m - else: - return rotated_m + return rotated_m diff --git a/croissant/core/sphtransform.py b/croissant/core/sphtransform.py index 613e4b5..aebf20a 100644 --- a/croissant/core/sphtransform.py +++ b/croissant/core/sphtransform.py @@ -3,40 +3,34 @@ from ..constants import PIX_WEIGHTS_NSIDE -def alm2map(alm, nside, lmax=None): +def alm2map(alm, nside, lmax=None, polarized=False): """ - Compute the healpix map from the spherical harmonics coefficients. + Compute the healpix map from the spherical harmonics coefficients. - Parameters - ---------- - alm : array-like - The spherical harmonics coefficients in the healpy convention. Shape - ([nfreq], hp.Alm.getsize(lmax)). - nside : int - The nside of the output map(s). - lmax : int - The lmax of the spherical harmonics transform. Defaults to 3*nside-1. + Parameters + ---------- + alm : array-like + The spherical harmonics coefficients in the healpy convention. Shape + ([nfreq], hp.Alm.getsize(lmax)). + nside : int + The nside of the output map(s). + lmax : int + The lmax of the spherical harmonics transform. Defaults to 3*nside-1. + polarized : bool + If true, alm's are assumed to be TEB and output will be I, Q, U maps. + I is spin-0 and Q, U are spin-2. Multiple frequency maps are not yet + supported in this case. - Returns - ------- - map : np.ndarray - The healpix map. Shape ([nfreq], hp.nside2npix(nside)). + Returns + ------- + map : np.ndarray + The healpix map or sequence of maps. """ - alm = np.array(alm, copy=True) - if alm.ndim == 1: - return hp.alm2map(alm, nside, lmax=lmax) - else: - npix = hp.nside2npix(nside) - nfreqs = alm.shape[0] - hp_map = np.empty((nfreqs, npix)) - for i in range(nfreqs): - map_i = hp.alm2map(alm[i], nside, lmax=lmax) - hp_map[i] = map_i - return hp_map + return hp.alm2map(alm, nside, lmax=lmax, pol=polarized) -def map2alm(data, lmax=None): +def map2alm(data, lmax=None, polarized=False): """ Compute the spherical harmonics coefficents of a healpix map. @@ -46,6 +40,10 @@ def map2alm(data, lmax=None): The healpix map(s). Shape ([nfreq], hp.nside2npix(nside)). lmax : int The lmax of the spherical harmonics transform. Defaults to 3*nside-1. + polarized : bool + If true, map is assumed to be I, Q, U maps. I is spin-0 and Q, U are + spin-2. Output alm's will be TEB. + Multiple frequency maps are not yet supported in this case. Returns ------- @@ -59,19 +57,10 @@ def map2alm(data, lmax=None): nside = hp.npix2nside(npix) use_pix_weights = nside in PIX_WEIGHTS_NSIDE use_ring_weights = not use_pix_weights - kwargs = { - "lmax": lmax, - "pol": False, - "use_weights": use_ring_weights, - "use_pixel_weights": use_pix_weights, - } - if data.ndim == 1: - alm = hp.map2alm(data, **kwargs) - else: - # compute the alms of the first map to determine the size of the array - alm0 = hp.map2alm(data[0], **kwargs) - alm = np.empty((len(data), alm0.size), dtype=alm0.dtype) - alm[0] = alm0 - for i in range(1, len(data)): - alm[i] = hp.map2alm(data[i], **kwargs) - return alm + return hp.map2alm( + data, + lmax=lmax, + pol=polarized, + use_weights=use_ring_weights, + use_pixel_weights=use_pix_weights, + ) diff --git a/croissant/core/tests/test_rotations.py b/croissant/core/tests/test_rotations.py index 43976ab..7491309 100644 --- a/croissant/core/tests/test_rotations.py +++ b/croissant/core/tests/test_rotations.py @@ -6,6 +6,8 @@ from croissant import rotations +rng = np.random.default_rng(1234) + def test_rotator_init(): # invalid eulertype @@ -35,10 +37,10 @@ def test_rotator_init(): rotations.Rotator(coord=["T", "M"], loc=loc) -def test_rotate_alm(): - lmax = 10 +@pytest.mark.parametrize("lmax", [8, 16, 32, 64]) +def test_rotate_alm(lmax): size = hp.Alm.getsize(lmax) - alm = np.arange(size) + 5j * np.arange(size) ** 2 + alm = rng.standard_normal(size) + 1j * rng.standard_normal(size) # galactic -> equatorial rot = rotations.Rotator(coord=["G", "C"]) @@ -55,6 +57,17 @@ def test_rotate_alm(): hp_rot_alm = hp_rot.rotate_alm(alm) assert np.allclose(rot_alm, hp_rot_alm) + # test polarization maps + t = alm.copy() + e = rng.standard_normal(size) + 1j * rng.standard_normal(size) + b = rng.standard_normal(size) + 1j * rng.standard_normal(size) + alm_pol = np.array([t, e, b]) + rot = rotations.Rotator(coord=["G", "C"]) + rot_alm = rot.rotate_alm(alm_pol, lmax=lmax, polarized=True) + hp_rot = hp.Rotator(coord=["G", "C"]) + hp_rot_alm = hp_rot.rotate_alm(alm_pol, lmax=lmax) + assert np.allclose(rot_alm, hp_rot_alm) + # multiple sets of alms at once alm.shape = (1, size) f = np.linspace(1, 50, 50).reshape(-1, 1) @@ -88,6 +101,18 @@ def test_rotate_map(): hprm = hp_rot.rotate_map_pixel(m) assert np.allclose(rm, hprm) + # test polarization maps + i_map = m.copy() + q_map = rng.standard_normal(npix) + u_map = rng.standard_normal(npix) + map_pol = np.array([i_map, q_map, u_map]) + rm = rot.rotate_map_alms(map_pol, polarized=True) + hprm = hp_rot.rotate_map_alms(map_pol, use_pixel_weights=True) + assert np.allclose(rm, hprm) + rm = rot.rotate_map_pixel(map_pol, polarized=True) + hprm = hp_rot.rotate_map_pixel(map_pol) + assert np.allclose(rm, hprm) + # several maps at once f = np.linspace(1, 50, 50).reshape(-1, 1) maps = m.reshape(1, -1) * f**3.2 diff --git a/croissant/core/tests/test_sphtransform.py b/croissant/core/tests/test_sphtransform.py index 73b5fb2..7329e34 100644 --- a/croissant/core/tests/test_sphtransform.py +++ b/croissant/core/tests/test_sphtransform.py @@ -1,13 +1,15 @@ import healpy as hp import numpy as np - +import pytest from croissant.constants import Y00 from croissant.core.sphtransform import alm2map, map2alm +pytestmark = pytest.mark.parametrize("lmax", [16, 32, 64]) +rng = np.random.default_rng(seed=1234) + -def test_alm2map(): +def test_alm2map(lmax): # make constant map - lmax = 10 a00 = 3 size = hp.Alm.getsize(lmax) alm = np.zeros(size, dtype=np.complex128) @@ -30,20 +32,10 @@ def test_alm2map(): assert np.allclose(hp_map, expected_maps) # inverting map2alm - lmax = 10 size = hp.Alm.getsize(lmax) - alm = np.zeros(size, dtype=np.complex128) - lm_dict = { - (0, 0): 5.1, - (2, 0): 7.4, - (3, 2): 3 + 5j, - (4, 1): -4.3 + 1.3j, - (7, 7): 11, - } - for ell, emm in lm_dict: - ix = hp.Alm.getidx(lmax, ell, emm) - alm[ix] = lm_dict[(ell, emm)] - + alm = rng.standard_normal(size) + 1j * rng.standard_normal(size) + # m = 0 modes are real + alm[: lmax + 1] = alm[: lmax + 1].real nside = 64 m = alm2map(alm, nside, lmax=lmax) alm_ = map2alm(m, lmax=lmax) @@ -51,12 +43,22 @@ def test_alm2map(): m_ = alm2map(alm_, nside, lmax=lmax) assert np.allclose(m, m_) + # polarized case + t_alm = alm.copy() + e_alm = rng.standard_normal(size) + 1j * rng.standard_normal(size) + b_alm = rng.standard_normal(size) + 1j * rng.standard_normal(size) + alm = np.array([t_alm, e_alm, b_alm]) + # m = 0 modes are real + alm[:, : lmax + 1] = alm[:, : lmax + 1].real + m = alm2map(alm, nside, lmax=lmax, polarized=True) + hp_map = hp.alm2map(alm, nside, lmax=lmax, pol=True) + assert np.allclose(m, hp_map) -def test_map2alm(): + +def test_map2alm(lmax): nside = 32 npix = hp.nside2npix(nside) data = np.ones(npix) - lmax = 5 alm = map2alm(data, lmax=lmax) assert np.isclose(alm[0], Y00 * 4 * np.pi) # a00 = Y00 * 4pi assert np.allclose(alm[1:], 0) # rest of alms should be 0 @@ -75,3 +77,11 @@ def test_map2alm(): # d1 is d0 + constant value so they should only differ at a00 assert np.allclose(alm0[1:], alm1[1:]) assert np.isclose(alm0[0] + c * Y00 * 4 * np.pi, alm1[0]) + + # polarized case + nside = 64 + npix = hp.nside2npix(nside) + iqu_map = rng.standard_normal((3, npix)) + alm = map2alm(iqu_map, lmax=lmax, polarized=True) + hp_alm = hp.map2alm(iqu_map, lmax=lmax, pol=True, use_pixel_weights=True) + assert np.allclose(alm, hp_alm)