diff --git a/.commitlint.rules.js b/.commitlint.rules.js index 34dc91a..084d2b5 100644 --- a/.commitlint.rules.js +++ b/.commitlint.rules.js @@ -9,6 +9,7 @@ module.exports = { "catalog", "cli", "core", + "ducc", "fields", "healpy", "io", diff --git a/heracles/cli.py b/heracles/cli.py index 85c3828..bc4d47e 100644 --- a/heracles/cli.py +++ b/heracles/cli.py @@ -165,11 +165,11 @@ 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 @@ -177,6 +177,13 @@ def mapper_from_config(config, section): 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) diff --git a/heracles/ducc.py b/heracles/ducc.py new file mode 100644 index 0000000..b69e827 --- /dev/null +++ b/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 . +""" +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 diff --git a/heracles/io.py b/heracles/io.py index 0e4b960..6468dc7 100644 --- a/heracles/io.py +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 24a0346..bf7f4b3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,6 +43,7 @@ optional-dependencies = {all = [ "sphinx", "sphinxcontrib-katex", ], test = [ + "ducc0", "pytest", "pytest-rerunfailures", ]} diff --git a/tests/test_cli.py b/tests/test_cli.py index 5cf6896..dc8f3e3 100644 --- a/tests/test_cli.py +++ b/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"): diff --git a/tests/test_ducc.py b/tests/test_ducc.py new file mode 100644 index 0000000..c2c346a --- /dev/null +++ b/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)