ENH: discrete mapper with ducc

.commitlint.rules.js

@@ -9,6 +9,7 @@ module.exports = {
         "catalog",
         "cli",
         "core",
+        "ducc",
         "fields",
         "healpy",
         "io",

heracles/cli.py

@@ -165,18 +165,25 @@ def mapper_from_config(config, section):
     choices = {
         "none": "none",
         "healpix": "healpix",
+        "discrete": "discrete",
     }
 
     mapper = config.getchoice(section, "mapper", choices)
-    if mapper == "none":
-        return None
+
     if mapper == "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 HealpixMapper(nside, lmax, deconvolve=deconvolve)
+
+    if mapper == "discrete":
+        from .ducc import DiscreteMapper
+
+        lmax = config.getint(section, "lmax", fallback=None)
+        return DiscreteMapper(lmax)
+
     return None
 
 
@@ -223,6 +230,12 @@ def catalog_from_config(config, section, label=None, *, out=None):
     # check if visibility is per catalogue or per selection
     visibility: str | Mapping[str, str]
     visibility = config.get(section, "visibility", fallback=None)
+    visibility_transform = config.getboolean(
+        section,
+        "visibility-transform",
+        fallback=False,
+    )
+    visibility_lmax = config.getint(section, "visibility-lmax", fallback=None)
     # check if visibility is a mapping
     if visibility and "\n" in visibility:
         visibility = config.getdict(section, "visibility")
@@ -233,7 +246,11 @@ def catalog_from_config(config, section, label=None, *, out=None):
     # set base catalogue's visibility if just one was given
     if isinstance(visibility, str):
         try:
-            vmap = read_vmap(getpath(visibility))
+            vmap = read_vmap(
+                getpath(visibility),
+                transform=visibility_transform,
+                lmax=visibility_lmax,
+            )
         except (TypeError, ValueError, OSError) as exc:
             msg = f"Cannot load visibility: {exc!s}"
             raise ValueError(msg)
@@ -269,7 +286,11 @@ def catalog_from_config(config, section, label=None, *, out=None):
                 msg = f"Invalid value: unknown selection '{num}'"
                 raise ValueError(msg)
             try:
-                vmap = read_vmap(getpath(value))
+                vmap = read_vmap(
+                    getpath(value),
+                    transform=visibility_transform,
+                    lmax=visibility_lmax,
+                )
             except (TypeError, ValueError, OSError) as exc:
                 msg = f"Cannot load visibility: {exc!s}"
                 raise ValueError(msg)

heracles/ducc.py

@@ -0,0 +1,163 @@
+# Heracles: Euclid code for harmonic-space statistics on the sphere
+#
+# Copyright (C) 2023-2024 Euclid Science Ground Segment
+#
+# This file is part of Heracles.
+#
+# Heracles is free software: you can redistribute it and/or modify it
+# under the terms of the GNU Lesser General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# Heracles is distributed in the hope that it will be useful, but
+# WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
+# Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public
+# License along with Heracles. If not, see <https://www.gnu.org/licenses/>.
+"""
+Module for discrete spherical harmonic transforms with ducc.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import ducc0
+import numpy as np
+
+from heracles.core import update_metadata
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+    from typing import Any
+
+    from numpy.typing import ArrayLike, DTypeLike, NDArray
+
+
+class DiscreteMapper:
+    """
+    Mapper that creates alms directly.
+    """
+
+    def __init__(
+        self,
+        lmax: int,
+        *,
+        dtype: DTypeLike = np.complex128,
+        nthreads: int = 0,
+    ) -> None:
+        """
+        Mapper for alms.
+        """
+        self.__lmax = lmax
+        self.__dtype = np.dtype(dtype)
+        self.__nthreads = nthreads
+
+    @property
+    def lmax(self) -> int:
+        """
+        The maximum angular mode number.
+        """
+        return self.__lmax
+
+    @property
+    def area(self) -> float:
+        """
+        The effective area for this mapper.
+        """
+        return 1.0
+
+    def create(
+        self,
+        *dims: int,
+        spin: int = 0,
+    ) -> NDArray[Any]:
+        """
+        Create zero alms.
+        """
+        lmax = self.__lmax
+        m = np.zeros((*dims, (lmax + 1) * (lmax + 2) // 2), dtype=self.__dtype)
+        update_metadata(
+            m,
+            geometry="discrete",
+            kernel="none",
+            lmax=lmax,
+            spin=spin,
+        )
+        return m
+
+    def map_values(
+        self,
+        lon: NDArray[Any],
+        lat: NDArray[Any],
+        data: NDArray[Any],
+        values: NDArray[Any],
+    ) -> None:
+        """
+        Add values to alms.
+        """
+
+        md: Mapping[str, Any] = data.dtype.metadata or {}
+
+        flatten = values.ndim == 1
+        if flatten:
+            values = values.reshape(1, -1)
+
+        epsilon: float
+        if values.dtype == np.float64:
+            epsilon = 1e-12
+        elif values.dtype == np.float32:
+            epsilon = 1e-5
+        else:
+            values = values.astype(np.float64)
+            epsilon = 1e-12
+
+        spin = md.get("spin", 0)
+
+        loc = np.empty((lon.size, 2), dtype=np.float64)
+        loc[:, 0] = np.radians(90.0 - lat)
+        loc[:, 1] = np.radians(lon % 360.0)
+
+        alms = ducc0.sht.adjoint_synthesis_general(
+            map=values,
+            spin=spin,
+            lmax=self.__lmax,
+            loc=loc,
+            epsilon=epsilon,
+            nthreads=self.__nthreads,
+        )
+
+        if flatten:
+            alms = alms[0]
+
+        data += alms
+
+    def transform(
+        self,
+        data: ArrayLike,
+    ) -> ArrayLike | tuple[ArrayLike, ArrayLike]:
+        """
+        Does nothing, since inputs are alms already.
+        """
+        return data
+
+    def resample(self, data: NDArray[Any]) -> NDArray[Any]:
+        """
+        Change LMAX of alm.
+        """
+        *dims, n = data.shape
+        lmax_in = (int((8 * n + 1) ** 0.5 + 0.01) - 3) // 2
+        lmax_out = self.__lmax
+        lmax = min(lmax_in, lmax_out)
+        out = np.zeros(
+            (*dims, (lmax_out + 1) * (lmax_out + 2) // 2),
+            dtype=self.__dtype,
+        )
+        i = j = 0
+        for m in range(lmax + 1):
+            out[..., j : j + lmax - m + 1] = data[..., i : i + lmax - m + 1]
+            i += lmax_in - m + 1
+            j += lmax_out - m + 1
+        return out

heracles/io.py

@@ -300,7 +300,7 @@ def _read_twopoint(hdu):
     return arr
 
 
-def read_vmap(filename, nside=None, field=0):
+def read_vmap(filename, nside=None, field=0, *, transform=False, lmax=None):
     """read visibility map from a HEALPix map file"""
 
     import healpy as hp
@@ -315,6 +315,9 @@ def read_vmap(filename, nside=None, field=0):
         warn(f"{filename}: changing NSIDE to {nside}")
         vmap = hp.ud_grade(vmap, nside)
 
+    if transform:
+        vmap = hp.map2alm(vmap, lmax=lmax, use_pixel_weights=True)
+
     return vmap
 
 

pyproject.toml

@@ -43,6 +43,7 @@ optional-dependencies = {all = [
     "sphinx",
     "sphinxcontrib-katex",
 ], test = [
+    "ducc0",
     "pytest",
     "pytest-rerunfailures",
 ]}

tests/test_cli.py

@@ -191,6 +191,8 @@ def test_catalog_from_config(mock):
                     0 = vmap.0.fits
                     2 = vmap.2.fits
                 """,
+                "visibility-transform": "true",
+                "visibility-lmax": "100",
             },
         },
     )
@@ -213,7 +215,7 @@ def test_catalog_from_config(mock):
     assert catalog[0].visibility is catalog[0].base.visibility
     assert catalog[1].visibility is catalog[0].base.visibility
     assert catalog[2].visibility is catalog[0].base.visibility
-    assert mock.call_args_list == [(("vmap.fits",),)]
+    assert mock.call_args_list == [(("vmap.fits",), dict(transform=False, lmax=None))]
 
     mock.reset_mock()
 
@@ -236,8 +238,8 @@ def test_catalog_from_config(mock):
     assert catalog[1].visibility is None
     assert catalog[2].visibility is mock.return_value
     assert mock.call_args_list == [
-        (("vmap.0.fits",),),
-        (("vmap.2.fits",),),
+        (("vmap.0.fits",), dict(transform=True, lmax=100)),
+        (("vmap.2.fits",), dict(transform=True, lmax=100)),
     ]
 
     with pytest.raises(ValueError, match="Duplicate selection"):

tests/test_ducc.py

@@ -0,0 +1,45 @@
+import importlib.util
+
+import numpy as np
+import numpy.testing as npt
+import pytest
+
+HAVE_DUCC = importlib.util.find_spec("ducc0") is not None
+
+skipif_no_ducc = pytest.mark.skipif(not HAVE_DUCC, reason="test requires ducc")
+
+
+@skipif_no_ducc
+def test_resample():
+    from heracles.ducc import DiscreteMapper
+
+    lmax = 200
+
+    alm = np.concatenate(
+        [np.arange(m, lmax + 1) for m in range(lmax + 1)],
+        dtype=complex,
+    )
+
+    out = DiscreteMapper(lmax).resample(alm)
+
+    npt.assert_array_equal(out, alm)
+
+    lmax_out = lmax // 2
+    out = DiscreteMapper(lmax_out).resample(alm)
+
+    assert out.shape == ((lmax_out + 1) * (lmax_out + 2) // 2,)
+    i = j = 0
+    for m in range(lmax_out + 1):
+        i, j = j, j + lmax_out - m + 1
+        npt.assert_array_equal(out[i:j], np.arange(m, lmax_out + 1))
+
+    lmax_out = lmax * 2
+    out = DiscreteMapper(lmax_out).resample(alm)
+
+    assert out.shape == ((lmax_out + 1) * (lmax_out + 2) // 2,)
+    i = j = 0
+    for m in range(lmax + 1):
+        i, j = j, j + lmax_out - m + 1
+        expected = np.pad(np.arange(m, lmax + 1), (0, lmax_out - lmax))
+        npt.assert_array_equal(out[i:j], expected)
+    npt.assert_array_equal(out[j:], 0.0)