diff --git a/heracles/maps.py b/heracles/maps.py index b3549b5..47439d2 100644 --- a/heracles/maps.py +++ b/heracles/maps.py @@ -22,7 +22,7 @@ import typing as t import warnings from abc import ABCMeta, abstractmethod -from collections.abc import Generator +from collections.abc import Generator, Mapping from functools import partial, wraps import healpy as hp @@ -829,6 +829,7 @@ def map_catalogs( def transform_maps( maps: t.Mapping[tuple[t.Any, t.Any], MapData], *, + lmax: t.Union[int, t.Mapping[t.Any, int], None] = None, out: t.MutableMapping[t.Any, t.Any] = None, progress: bool = False, **kwargs, @@ -849,6 +850,10 @@ def transform_maps( # convert maps to alms, taking care of complex and spin-weighted maps for (k, i), m in maps.items(): nside = hp.get_nside(m) + if isinstance(lmax, Mapping): + _lmax = lmax.get((k, i)) or lmax.get(k) + else: + _lmax = lmax md = m.dtype.metadata or {} spin = md.get("spin", 0) @@ -864,7 +869,7 @@ def transform_maps( msg = f"spin-{spin} maps not yet supported" raise NotImplementedError(msg) - alms = hp.map2alm(m, pol=pol, **kwargs) + alms = hp.map2alm(m, lmax=_lmax, pol=pol, **kwargs) if spin == 0: alms = {(k, i): alms} diff --git a/tests/test_maps.py b/tests/test_maps.py index 5b3b148..60495d5 100644 --- a/tests/test_maps.py +++ b/tests/test_maps.py @@ -351,6 +351,17 @@ def test_transform_maps(rng): assert alms["P_E", 1].dtype.metadata["nside"] == nside assert alms["P_B", 1].dtype.metadata["nside"] == nside + # explicit lmax per map + maps = {("T", 0): t, ("P", 1): p} + lmax = {"T": 10, "P": 20} + alms = transform_maps(maps, lmax=lmax) + + assert len(alms) == 3 + assert alms.keys() == {("T", 0), ("P_E", 1), ("P_B", 1)} + assert alms["T", 0].size == (lmax["T"] + 1) * (lmax["T"] + 2) // 2 + assert alms["P_E", 1].size == (lmax["P"] + 1) * (lmax["P"] + 2) // 2 + assert alms["P_B", 1].size == (lmax["P"] + 1) * (lmax["P"] + 2) // 2 + def test_update_metadata(): from heracles.maps import update_metadata