Merge pull request #74 from christianhbye/rotate_pol

add support for I, Q, U polarized maps

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):

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,55 +164,51 @@ 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.
 
         Parameters
         -----------
         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)
             for i in range(len(m)):
                 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

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,
+    )

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

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,33 +32,33 @@ 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)
     assert np.allclose(alm, alm_)
     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)