API: reorganise maps module (#147)

Refactors the `maps` module, which is about "mapping" more than about
"maps", into a set of individual modules `heracles.mapper` for the
`Mapper` protocol and `heracles.healpy` for the healpy-based mapper. The
`map_catalogs()` and `transform()` functions are exported via the main
`heracles` module.

Closes: #145

.commitlint.rules.js

@@ -10,9 +10,10 @@ module.exports = {
         "cli",
         "core",
         "covariance",
-        "io",
         "fields",
-        "maps",
+        "healpy",
+        "io",
+        "mapper",
         "plot",
         "progress",
         "twopoint",

examples/example.ipynb

@@ -439,14 +439,14 @@
    "outputs": [],
    "source": [
     "from heracles.fields import Positions, Shears\n",
-    "from heracles.maps import Healpix"
+    "from heracles.healpy import HealpixMapper"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "To turn catalogues into spectra, *Heracles* requires a so-called mapper object, which knows how to turn positions and values into spherical functions. Here, we construct a `Healpix` mapper instance with our desired *NSIDE* and *LMAX* parameters. When computing spherical harmonic coefficients, the HEALPix mapper will also remove the pixel window function, unless `deconvolve=False` is passed."
+    "To turn catalogues into spectra, *Heracles* requires a so-called mapper object, which knows how to turn positions and values into spherical functions. Here, we construct a `HealpixMapper` instance with our desired *NSIDE* and *LMAX* parameters. When computing spherical harmonic coefficients, the HEALPix mapper will also remove the pixel window function, unless `deconvolve=False` is passed."
    ]
   },
   {
@@ -455,7 +455,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "mapper = Healpix(nside, lmax)"
+    "mapper = HealpixMapper(nside, lmax)"
    ]
   },
   {
@@ -499,7 +499,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from heracles.maps import map_catalogs"
+    "from heracles import map_catalogs"
    ]
   },
   {
@@ -672,7 +672,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from heracles.maps import transform"
+    "from heracles import transform"
    ]
   },
   {
@@ -873,7 +873,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "mapper_mm = Healpix(2*nside, 2*lmax)"
+    "mapper_mm = HealpixMapper(2*nside, 2*lmax)"
    ]
   },
   {

heracles/__init__.py

@@ -18,7 +18,16 @@
 # License along with Heracles. If not, see <https://www.gnu.org/licenses/>.
 """Main module of the *Heracles* package."""
 
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "map_catalogs",
+    "transform",
+]
+
 try:
-    from ._version import __version__, __version_tuple__  # noqa: F401
+    from ._version import __version__, __version_tuple__
 except ModuleNotFoundError:
     pass
+
+from ._mapping import map_catalogs, transform

heracles/cli.py

@@ -171,12 +171,12 @@ def mapper_from_config(config, section):
     if mapper == "none":
         return None
     if mapper == "healpix":
-        from .maps import Healpix
+        from .healpy import HealpixMapper
 
         nside = config.getint(section, "nside")
         lmax = config.getint(section, "lmax", fallback=None)
         deconvolve = config.getboolean(section, "deconvolve", fallback=None)
-        return Healpix(nside, lmax, deconvolve=deconvolve)
+        return HealpixMapper(nside, lmax, deconvolve=deconvolve)
     return None
 
 
@@ -504,16 +504,17 @@ def alms(
 
     """
 
+    from . import transform
+    from .healpy import HealpixMapper
     from .io import AlmFits
-    from .maps import Healpix, transform
 
     # load the config file, this contains alms setting and maps definition
     logger.info("reading configuration from %s", files)
     config = loader(files)
 
     # set the HEALPix datapath
     if healpix_datapath is not None:
-        Healpix.DATAPATH = healpix_datapath
+        HealpixMapper.DATAPATH = healpix_datapath
 
     # construct fields to get mappers for transform
     fields = fields_from_config(config)

heracles/maps/_healpix.py → heracles/healpy.py

@@ -63,7 +63,7 @@ def _map(ipix, maps, values):
         maps[..., i] += values[..., j]
 
 
-class Healpix:
+class HealpixMapper:
     """
     Mapper for HEALPix maps.
     """
@@ -178,7 +178,7 @@ def transform(self, data: NDArray[Any]) -> NDArray[Any]:
             msg = f"spin-{spin} maps not yet supported"
             raise NotImplementedError(msg)
 
-        alms = hp.map2alm(
+        alm = hp.map2alm(
             data,
             lmax=self.__lmax,
             pol=pol,
@@ -189,16 +189,16 @@ def transform(self, data: NDArray[Any]) -> NDArray[Any]:
         if pw is not None:
             fl = np.ones(self.__lmax + 1)
             fl[abs(spin) :] /= pw[abs(spin) :]
-            for alm in alms:
-                hp.almxfl(alm, fl, inplace=True)
+            for i in np.ndindex(*alm.shape[:-1]):
+                alm[i] = hp.almxfl(alm[i], fl)
             del fl
 
         if spin != 0:
-            alms = alms[1:].copy()
+            alm = alm[1:].copy()
 
-        update_metadata(alms, **md)
+        update_metadata(alm, **md)
 
-        return alms
+        return alm
 
     def resample(self, data: NDArray[Any]) -> NDArray[Any]:
         """

tests/test_fields.py

@@ -18,9 +18,9 @@ def sigma_e():
 # TODO: mock mapper
 @pytest.fixture
 def mapper(nside):
-    from heracles.maps import Healpix
+    from heracles.healpy import HealpixMapper
 
-    return Healpix(nside)
+    return HealpixMapper(nside)
 
 
 @pytest.fixture
@@ -132,7 +132,7 @@ def test_visibility(nside, vmap):
     from unittest.mock import Mock
 
     from heracles.fields import Visibility
-    from heracles.maps import Healpix
+    from heracles.healpy import HealpixMapper
 
     fsky = vmap.mean()
 
@@ -141,7 +141,7 @@ def test_visibility(nside, vmap):
         catalog.visibility = vmap
         catalog.metadata = {"catalog": catalog.label}
 
-        mapper_out = Healpix(nside_out)
+        mapper_out = HealpixMapper(nside_out)
 
         f = Visibility(mapper_out)
 

tests/test_maps.py → tests/test_healpy.py

@@ -4,18 +4,27 @@
 import numpy.testing as npt
 import pytest
 
-
-def test_healpix_maps(rng):
+try:
     import healpy as hp
+except ImportError:
+    HAVE_HEALPY = False
+else:
+    HAVE_HEALPY = True
+
+skipif_not_healpy = pytest.mark.skipif(not HAVE_HEALPY, reason="test requires healpy")
 
-    from heracles.maps import Healpix, Mapper
+
+@skipif_not_healpy
+def test_healpix_maps(rng):
+    from heracles.healpy import HealpixMapper
+    from heracles.mapper import Mapper
 
     nside = 1 << rng.integers(1, 10)
     npix = hp.nside2npix(nside)
     lmax = 1 << rng.integers(1, 10)
     deconv = rng.choice([True, False])
 
-    mapper = Healpix(nside, lmax, deconvolve=deconv)
+    mapper = HealpixMapper(nside, lmax, deconvolve=deconv)
 
     assert isinstance(mapper, Mapper)
 
@@ -69,15 +78,16 @@ def test_healpix_maps(rng):
     npt.assert_array_equal(m, expected)
 
 
+@skipif_not_healpy
 @unittest.mock.patch("healpy.map2alm")
 def test_healpix_transform(mock_map2alm, rng):
     from heracles.core import update_metadata
-    from heracles.maps import Healpix
+    from heracles.healpy import HealpixMapper
 
     nside = 32
     npix = 12 * nside**2
 
-    mapper = Healpix(nside)
+    mapper = HealpixMapper(nside)
 
     # single scalar map
     m = rng.standard_normal(npix)
@@ -106,63 +116,48 @@ def test_healpix_transform(mock_map2alm, rng):
     assert alms.dtype.metadata["nside"] == nside
 
 
-class MockCatalog:
-    size = 10
-    page_size = 1
-
-    def __iter__(self):
-        for i in range(0, self.size, self.page_size):
-            yield {}
-
-
-@pytest.mark.parametrize("parallel", [False, True])
-def test_map_catalogs(parallel):
-    from unittest.mock import AsyncMock
-
-    from heracles.maps import map_catalogs
-
-    fields = {"a": AsyncMock(), "b": AsyncMock(), "z": AsyncMock()}
-    catalogs = {"x": MockCatalog(), "y": MockCatalog()}
-
-    maps = map_catalogs(fields, catalogs, parallel=parallel)
-
-    for k in fields:
-        for i in catalogs:
-            fields[k].assert_any_call(catalogs[i], progress=None)
-            assert maps[k, i] is fields[k].return_value
-
-
-def test_map_catalogs_match():
-    from unittest.mock import AsyncMock
-
-    from heracles.maps import map_catalogs
+@skipif_not_healpy
+@unittest.mock.patch("healpy.map2alm")
+def test_healpix_deconvolve(mock_map2alm):
+    from heracles.core import update_metadata
+    from heracles.healpy import HealpixMapper
 
-    fields = {"a": AsyncMock(), "b": AsyncMock(), "c": AsyncMock()}
-    catalogs = {"x": MockCatalog(), "y": MockCatalog()}
+    nside = 32
+    npix = 12 * nside**2
 
-    maps = map_catalogs(fields, catalogs, include=[(..., "y")])
+    lmax = 48
+    nlm = (lmax + 1) * (lmax + 2) // 2
 
-    assert set(maps.keys()) == {("a", "y"), ("b", "y"), ("c", "y")}
+    pw0, pw2 = hp.pixwin(nside, lmax=lmax, pol=True)
+    pw2[:2] = 1.0
 
-    maps = map_catalogs(fields, catalogs, exclude=[("a", ...)])
+    mapper = HealpixMapper(nside, lmax, deconvolve=True)
 
-    assert set(maps.keys()) == {("b", "x"), ("b", "y"), ("c", "x"), ("c", "y")}
+    # single scalar map
+    data = np.zeros(npix)
+    update_metadata(data, spin=0)
 
+    mock_map2alm.return_value = np.ones(nlm, dtype=complex)
 
-def test_transform(rng):
-    from unittest.mock import Mock
+    alm = mapper.transform(data)
 
-    from heracles.maps import transform
+    assert alm.shape == (nlm,)
+    stop = 0
+    for m in range(lmax + 1):
+        start, stop = stop, stop + lmax - m + 1
+        npt.assert_array_equal(alm[start:stop], 1.0 / pw0[m:])
 
-    x = Mock()
-    y = Mock()
+    # polarisation map
+    data = np.zeros((2, npix))
+    update_metadata(data, spin=2)
 
-    fields = {"X": x, "Y": y}
-    maps = {("X", 0): Mock(), ("Y", 1): Mock()}
+    mock_map2alm.return_value = np.ones((3, nlm), dtype=complex)
 
-    alms = transform(fields, maps)
+    alm = mapper.transform(data)
 
-    assert len(alms) == 2
-    assert alms.keys() == {("X", 0), ("Y", 1)}
-    assert alms["X", 0] is x.mapper_or_error.transform.return_value
-    assert alms["Y", 1] is y.mapper_or_error.transform.return_value
+    assert alm.shape == (2, nlm)
+    stop = 0
+    for m in range(lmax + 1):
+        start, stop = stop, stop + lmax - m + 1
+        npt.assert_array_equal(alm[0, start:stop], 1.0 / pw2[m:])
+        npt.assert_array_equal(alm[1, start:stop], 1.0 / pw2[m:])