From 81e0c80d0145795ec22df8889db0f9fbd4e2cf84 Mon Sep 17 00:00:00 2001 From: Nicolas Tessore Date: Fri, 24 Jan 2025 16:59:08 +0000 Subject: [PATCH] gh-3: support for Gaussian random fields (#4) Adds the functionality for Gaussian random fields from `transformcl` and `gaussiancl`. Closes: #3 --- .github/workflows/test.yml | 2 +- .pre-commit-config.yaml | 9 +- LICENSE | 2 +- docs/conf.py | 2 +- docs/grf.rst | 39 +++ docs/index.rst | 10 +- docs/{spectra.rst => twopoint.rst} | 32 +- pyproject.toml | 3 +- src/angst/__init__.py | 26 +- src/angst/_spectra.py | 43 --- src/angst/_twopoint.py | 147 +++++++++ src/angst/grf.py | 280 ++++++++++++++++++ .../{test_spectra.py => test_twopoint.py} | 14 +- 13 files changed, 528 insertions(+), 81 deletions(-) create mode 100644 docs/grf.rst rename docs/{spectra.rst => twopoint.rst} (51%) delete mode 100644 src/angst/_spectra.py create mode 100644 src/angst/_twopoint.py create mode 100644 src/angst/grf.py rename src/angst/test/{test_spectra.py => test_twopoint.py} (69%) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index af4d8df..0e5a773 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -12,7 +12,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"] + python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"] fail-fast: false steps: - uses: actions/checkout@v4 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index dc35ae8..a465521 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,12 +1,9 @@ repos: - - repo: https://github.com/psf/black - rev: 23.7.0 - hooks: - - id: black - - repo: https://github.com/charliermarsh/ruff-pre-commit - rev: v0.0.286 + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.7.3 hooks: - id: ruff + - id: ruff-format - repo: https://github.com/pre-commit/mirrors-prettier rev: v3.0.3 hooks: diff --git a/LICENSE b/LICENSE index 58f0a34..74900ca 100644 --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,6 @@ MIT License -Copyright (c) 2023 Nicolas Tessore +Copyright (c) 2023-2024 Nicolas Tessore Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/docs/conf.py b/docs/conf.py index e3354f1..9715e8e 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -9,7 +9,7 @@ from angst import __version__ as angst_version project = "angst" -copyright = "2023, Nicolas Tessore" +copyright = "2023-2024, Nicolas Tessore" author = "Nicolas Tessore" version = angst_version.partition("+")[0] release = version diff --git a/docs/grf.rst b/docs/grf.rst new file mode 100644 index 0000000..58d43d5 --- /dev/null +++ b/docs/grf.rst @@ -0,0 +1,39 @@ +:mod:`angst.grf` --- Gaussian random fields +=========================================== + +.. currentmodule:: angst.grf +.. module:: angst.grf + + +Gaussian angular power spectra +------------------------------ + +.. autofunction:: solve + +.. autofunction:: compute_generic + + +Transformations +--------------- + +.. class:: Transformation(Protocol) + + .. automethod:: __call__ + .. automethod:: inv + .. automethod:: der + + +.. class:: Lognormal + + Implements the :class:`Transformation` for lognormal fields. + + +.. class:: LognormalXNormal + + Implements the :class:`Transformation` for the cross-correlation between + :class:`Lognormal` and Gaussian fields. + + +.. class:: SquaredNormal + + Implements the :class:`Transformation` for squared normal fields. diff --git a/docs/index.rst b/docs/index.rst index 86f1b97..b93bc68 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -7,13 +7,19 @@ .. toctree:: :maxdepth: 2 - :caption: Contents: + :caption: Contents points - spectra + twopoint core glossary +.. toctree:: + :maxdepth: 2 + :caption: Modules + + grf + Indices and tables ================== diff --git a/docs/spectra.rst b/docs/twopoint.rst similarity index 51% rename from docs/spectra.rst rename to docs/twopoint.rst index 963630d..7d2ea6f 100644 --- a/docs/spectra.rst +++ b/docs/twopoint.rst @@ -1,22 +1,31 @@ -Angular power spectra -===================== +Two-point functions +=================== .. currentmodule:: angst -Sets of angular power spectra ------------------------------ +Spectra and correlation functions +--------------------------------- -.. autofunction:: spectra_indices -.. autofunction:: enumerate_spectra +.. autofunction:: cl2corr +.. autofunction:: corr2cl +.. autofunction:: cl2var -.. _spectra_order: + +Sets of two-point functions +--------------------------- + +.. autofunction:: indices2 +.. autofunction:: enumerate2 + + +.. _twopoint_order: Standard order -------------- -All functions that process sets of angular power spectra expect them as a +All functions that process sets of two-point functions expect them as a sequence using the following "Christmas tree" ordering: .. raw:: html @@ -32,9 +41,8 @@ In other words, the sequence begins as such: * index 5 describes the cross-correlation of field 2 and field 0, * etc. -In particular, cross-correlations for the first :math:`n` fields are contained +In particular, two-point functions for the first :math:`n` fields are contained in the first :math:`T_n = n \, (n + 1) / 2` entries of the sequence. -To easily generate or iterate over sequences of angular power spectra in -standard order, see the :func:`enumerate_spectra` and :func:`spectra_indices` -functions. +To easily generate or iterate over sequences of two-point functions in standard +order, see the :func:`enumerate2` and :func:`indices2` functions. diff --git a/pyproject.toml b/pyproject.toml index 9e9a157..eda8ded 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,9 +15,10 @@ classifiers = [ "License :: OSI Approved :: MIT License", "Operating System :: OS Independent", ] -requires-python = ">=3.8" +requires-python = ">=3.9" dependencies = [ "numpy>=1.20.0", + "flt>=2022.7.27", ] dynamic = ["version"] diff --git a/src/angst/__init__.py b/src/angst/__init__.py index 4cfbe39..a1425a4 100644 --- a/src/angst/__init__.py +++ b/src/angst/__init__.py @@ -1,26 +1,38 @@ """angst -- angular statistics""" __all__ = [ + "__version__", + "__version_tuple__", + "cl2corr", + "cl2var", + "corr2cl", "displace", "displacement", - "enumerate_spectra", + "enumerate2", + "grf", "inv_triangle_number", - "spectra_indices", - "__version__", - "__version_tuple__", + "indices2", ] +from . import grf + from ._core import ( inv_triangle_number, ) + from ._points import ( displace, displacement, ) -from ._spectra import ( - enumerate_spectra, - spectra_indices, + +from ._twopoint import ( + cl2corr, + cl2var, + corr2cl, + enumerate2, + indices2, ) + from ._version import ( __version__, __version_tuple__, diff --git a/src/angst/_spectra.py b/src/angst/_spectra.py deleted file mode 100644 index 42d41bc..0000000 --- a/src/angst/_spectra.py +++ /dev/null @@ -1,43 +0,0 @@ -"""Operations on angular power spectra.""" - -from __future__ import annotations - -# typing -from typing import Iterable, Iterator -from numpy.typing import ArrayLike - - -def enumerate_spectra( - spectra: Iterable[ArrayLike | None], -) -> Iterator[tuple[int, int, ArrayLike | None]]: - """ - Iterate over a set of angular power spectra in :ref:`standard order - `, returning a tuple of indices and their associated - spectrum from the input. - - >>> spectra = [[1, 2, 3], [4, 5, 6], [7, 8, 9]] - >>> list(enumerate_spectra(spectra)) - [(0, 0, [1, 2, 3]), (1, 1, [4, 5, 6]), (1, 0, [7, 8, 9])] - - """ - - for k, cl in enumerate(spectra): - i = int((2 * k + 0.25) ** 0.5 - 0.5) - j = i * (i + 3) // 2 - k - yield i, j, cl - - -def spectra_indices(n: int) -> Iterator[tuple[int, int]]: - """ - Return an iterator over indices in :ref:`standard order - ` for a set of spectra for *n* functions. Each item - is a tuple of indices *i*, *j*. - - >>> list(spectra_indices(3)) - [(0, 0), (1, 1), (1, 0), (2, 2), (2, 1), (2, 0)] - - """ - - for i in range(n): - for j in range(i, -1, -1): - yield i, j diff --git a/src/angst/_twopoint.py b/src/angst/_twopoint.py new file mode 100644 index 0000000..c742985 --- /dev/null +++ b/src/angst/_twopoint.py @@ -0,0 +1,147 @@ +""" +Module for two-point functions. +""" + +from __future__ import annotations + +import numpy as np + +# typing +from typing import Any, Iterable, Iterator +from numpy.typing import ArrayLike, NDArray + + +def enumerate2( + entries: Iterable[ArrayLike | None], +) -> Iterator[tuple[int, int, ArrayLike | None]]: + """ + Iterate over a set of two-point functions in :ref:`standard order + `, returning a tuple of indices and their associated entry + from the input. + + >>> spectra = [[1, 2, 3], [4, 5, 6], [7, 8, 9]] + >>> list(enumerate2(spectra)) + [(0, 0, [1, 2, 3]), (1, 1, [4, 5, 6]), (1, 0, [7, 8, 9])] + + """ + + for k, cl in enumerate(entries): + i = int((2 * k + 0.25) ** 0.5 - 0.5) + j = i * (i + 3) // 2 - k + yield i, j, cl + + +def indices2(n: int) -> Iterator[tuple[int, int]]: + """ + Return an iterator over indices in :ref:`standard order ` + for a set of two-point functions for *n* fields. Each item is a tuple of + indices *i*, *j*. + + >>> list(indices2(3)) + [(0, 0), (1, 1), (1, 0), (2, 2), (2, 1), (2, 0)] + + """ + + for i in range(n): + for j in range(i, -1, -1): + yield i, j + + +def cl2corr(cl: NDArray[Any], closed: bool = False) -> NDArray[Any]: + r"""transform angular power spectrum to correlation function + + Takes an angular power spectrum with :math:`\mathtt{n} = \mathtt{lmax}+1` + coefficients and returns the corresponding angular correlation function in + :math:`\mathtt{n}` points. + + The correlation function values can be computed either over the closed + interval :math:`[0, \pi]`, in which case :math:`\theta_0 = 0` and + :math:`\theta_{n-1} = \pi`, or over the open interval :math:`(0, \pi)`. + + Parameters + ---------- + cl : (n,) array_like + Angular power spectrum from :math:`0` to :math:`\mathtt{lmax}`. + closed : bool + Compute correlation function over open (``closed=False``) or closed + (``closed=True``) interval. + + Returns + ------- + corr : (n,) array_like + Angular correlation function. + + """ + + from flt import idlt # type: ignore [import-not-found] + + # length n of the transform + if cl.ndim != 1: + raise TypeError("cl must be 1d array") + n = cl.shape[-1] + + # DLT coefficients = (2l+1)/(4pi) * Cl + c = np.arange(1, 2 * n + 1, 2, dtype=float) + c /= 4 * np.pi + c *= cl + + # perform the inverse DLT + corr: NDArray[Any] = idlt(c, closed=closed) + + # done + return corr + + +def corr2cl(corr: NDArray[Any], closed: bool = False) -> NDArray[Any]: + r"""transform angular correlation function to power spectrum + + Takes an angular function in :math:`\mathtt{n}` points and returns the + corresponding angular power spectrum from :math:`0` to :math:`\mathtt{lmax} + = \mathtt{n}-1`. + + The correlation function must be given at the angles returned by + :func:`transformcl.theta`. These can be distributed either over the closed + interval :math:`[0, \pi]`, in which case :math:`\theta_0 = 0` and + :math:`\theta_{n-1} = \pi`, or over the open interval :math:`(0, \pi)`. + + Parameters + ---------- + corr : (n,) array_like + Angular correlation function. + closed : bool + Compute correlation function over open (``closed=False``) or closed + (``closed=True``) interval. + + Returns + ------- + cl : (n,) array_like + Angular power spectrum from :math:`0` to :math:`\mathtt{lmax}`. + + """ + + from flt import dlt + + # length n of the transform + if corr.ndim != 1: + raise TypeError("corr must be 1d array") + n = corr.shape[-1] + + # compute the DLT coefficients + cl: NDArray[Any] = dlt(corr, closed=closed) + + # DLT coefficients = (2l+1)/(4pi) * Cl + cl /= np.arange(1, 2 * n + 1, 2, dtype=float) + cl *= 4 * np.pi + + # done + return cl + + +def cl2var(cl: NDArray[Any]) -> float: + """ + Compute the variance of the spherical random field in a point from the + given angular power spectrum. The input can be multidimensional, with + the last axis representing the modes. + """ + ell = np.arange(np.shape(cl)[-1]) + return np.sum((2 * ell + 1) / (4 * np.pi) * cl) # type: ignore diff --git a/src/angst/grf.py b/src/angst/grf.py new file mode 100644 index 0000000..f3a0cac --- /dev/null +++ b/src/angst/grf.py @@ -0,0 +1,280 @@ +""" +Transformations of Gaussian random fields. +""" + +from __future__ import annotations + +__all__ = [ + "Lognormal", + "LognormalXNormal", + "SquaredNormal", + "Transformation", + "compute_generic", + "solve", +] + +import math +from dataclasses import dataclass + +import numpy as np + +# typing +from typing import Any, Callable, Protocol +from numpy.typing import NDArray + + +class Transformation(Protocol): + """ + Protocol for transformations of Gaussian random fields. + """ + + def __call__(self, x: NDArray[Any], var: float, /) -> NDArray[Any]: + """ + Transform a Gaussian correlation function. + """ + + def inv(self, x: NDArray[Any], var: float, /) -> NDArray[Any]: + """ + Inverse transform to a Gaussian correlation function. + """ + + def der(self, x: NDArray[Any], var: float, /) -> NDArray[Any]: + """ + Derivative of the transform. + """ + + +def _relerr(dx: NDArray[Any], x: NDArray[Any]) -> float: + """compute the relative error max(|dx/x|)""" + q = np.divide(dx, x, where=(dx != 0), out=np.zeros_like(dx)) + return np.fabs(q).max() # type: ignore + + +def solve( + cl: NDArray[Any], + tfm: Transformation, + pad: int = 0, + *, + initial: NDArray[Any] | None = None, + cltol: float = 1e-5, + gltol: float = 1e-5, + maxiter: int = 20, + monopole: float | None = None, +) -> tuple[NDArray[Any], NDArray[Any], int]: + """ + Solve for a Gaussian angular power spectrum. + + Parameters + ---------- + cl : (n,) array + tfm : :class:`Transformation` + pad : int + + Returns + ------- + gl : (n,) array + Gaussian angular power spectrum solution. + cl : (n + pad,) array + Realised transformed angular power spectrum. + info : {0, 1, 2, 3} + Indicates success of failure of the solution. Possible values are + + * ``0``, solution did not converge in *maxiter* iterations; + * ``1``, solution converged in *cl* relative error; + * ``2``, solution converged in *gl* relative error; + * ``3``, solution converged in both *cl* and *gl* relative error. + + """ + + from ._twopoint import corr2cl, cl2corr, cl2var + + n = len(cl) + if not isinstance(pad, int) or pad < 0: + raise TypeError("pad must be a positive integer") + + if initial is None: + gl = corr2cl(tfm.inv(cl2corr(cl), cl2var(cl))) + else: + gl = np.empty(n) + gl[: len(initial)] = initial[:n] + + if monopole is not None: + gl[0] = monopole + + gt = cl2corr(np.pad(gl, (0, pad))) + var = cl2var(gl) + rl = corr2cl(tfm(gt, var)) + fl = rl[:n] - cl + if monopole is not None: + fl[0] = 0 + clerr = _relerr(fl, cl) + + info = 0 + for i in range(maxiter): + if clerr <= cltol: + info |= 1 + if info > 0: + break + + ft = cl2corr(np.pad(fl, (0, pad))) + dt = tfm.der(gt, var) + xl = -corr2cl(ft / dt)[:n] + if monopole is not None: + xl[0] = 0 + + while True: + gl_ = gl + xl + gt_ = cl2corr(np.pad(gl_, (0, pad))) + var_ = cl2var(gl_) + rl_ = corr2cl(tfm(gt_, var_)) + fl_ = rl_[:n] - cl + if monopole is not None: + fl_[0] = 0 + clerr_ = _relerr(fl_, cl) + if clerr_ <= clerr: + break + xl /= 2 + + if _relerr(xl, gl) <= gltol: + info |= 2 + + gl, gt, var, rl, fl, clerr = gl_, gt_, var_, rl_, fl_, clerr_ + + return gl, rl, info + + +@dataclass +class Lognormal: + """ + Transformation for lognormal fields. + """ + + lamda1: float = 1.0 + lamda2: float = 1.0 + + def __call__(self, x: NDArray[Any], var: float) -> NDArray[Any]: + xp = x.__array_namespace__() + return self.lamda1 * self.lamda2 * xp.expm1(x) # type: ignore[no-any-return] + + def inv(self, x: NDArray[Any], var: float) -> NDArray[Any]: + xp = x.__array_namespace__() + return xp.log1p(x / (self.lamda1 * self.lamda2)) # type: ignore[no-any-return] + + def der(self, x: NDArray[Any], var: float) -> NDArray[Any]: + xp = x.__array_namespace__() + return self.lamda1 * self.lamda2 * xp.exp(x) # type: ignore[no-any-return] + + +@dataclass +class LognormalXNormal: + """ + Transformation for cross-correlation between lognormal and Gaussian + fields. + """ + + lamda: float = 1.0 + + def __call__(self, x: NDArray[Any], var: float) -> NDArray[Any]: + return self.lamda * x + + def inv(self, x: NDArray[Any], var: float) -> NDArray[Any]: + return x / self.lamda + + def der(self, x: NDArray[Any], var: float) -> NDArray[Any]: + return self.lamda + (0.0 * x) + + +@dataclass +class SquaredNormal: + """ + Squared normal field. The parameters *a1*, *a2* can be set to + ``None``, in which case they are inferred from the variance of + the field. + """ + + a1: float | None = None + a2: float | None = None + lamda1: float = 1.0 + lamda2: float = 1.0 + + def _pars(self, var: float) -> tuple[float, float]: + a1 = math.sqrt(1 - var) if self.a1 is None else self.a1 + a2 = math.sqrt(1 - var) if self.a2 is None else self.a2 + return a1 * a2, self.lamda1 * self.lamda2 + + def __call__(self, x: NDArray[Any], var: float) -> NDArray[Any]: + aa, ll = self._pars(var) + return 2 * ll * x * (x + 2 * aa) + + def inv(self, x: NDArray[Any], var: float) -> NDArray[Any]: + xp = x.__array_namespace__() + aa, ll = self._pars(var) + return xp.sqrt(x / (2 * ll) + aa**2) - aa # type: ignore[no-any-return] + + def der(self, x: NDArray[Any], var: float) -> NDArray[Any]: + aa, ll = self._pars(var) + return 4 * ll * (x + aa) + + +Alm = NDArray[np.complexfloating[Any, Any]] + + +def compute_generic( + cl: NDArray[Any], + tfm: Transformation, + sht: Callable[[NDArray[Any], int], Alm], + isht: Callable[[Alm], NDArray[Any]], +) -> NDArray[Any]: + """ + Compute a Gaussian angular power spectrum for the target spectrum + *cl* and the transformation *tfm*. Uses the spherical harmonic + transform pair *sht* and *isht* to compute a simple band-limited + Gaussian power spectrum, without the full machinery of + :func:`solve`. + + Examples + -------- + Compute a Gaussian angular power spectrum for a lognormal + transformation using the ``healpy`` spherical harmonic transform:: + + import healpy as hp + + gl = angst.grf.compute_generic( + cl, + tfm, + lambda m, lmax: hp.map2alm(m, lmax=lmax, use_pixel_weights=True), + lambda alm: hp.alm2map(alm, 2048), + ) + + """ + + xp = cl.__array_namespace__() + + # get lmax from cl + lmax = cl.size - 1 + + # store prefactor, compute variance + fl = (2 * xp.arange(lmax + 1) + 1) / (4 * xp.pi) + var = (fl * cl).sum() + fl = xp.sqrt(fl) + + # compute alms for m=0, rest zero + alm = xp.concat( + [ + fl * cl, + xp.zeros(lmax * (lmax + 1) // 2, dtype=complex), + ], + ) + + # convert to field f(\theta, \phi) = C(\theta) exp(im\phi) + m = isht(alm) + + # apply lognormal transform to C(theta) + # (for complex fields, modify absolute value here) + m = tfm.inv(m, var) + + # get transformed alms + alm = sht(m, lmax) + + # read Gaussian spectrum from m=0 + return alm.real[: lmax + 1] / fl # type: ignore[no-any-return] diff --git a/src/angst/test/test_spectra.py b/src/angst/test/test_twopoint.py similarity index 69% rename from src/angst/test/test_spectra.py rename to src/angst/test/test_twopoint.py index 233836b..44b75f5 100644 --- a/src/angst/test/test_spectra.py +++ b/src/angst/test/test_twopoint.py @@ -2,7 +2,7 @@ import pytest -def test_enumerate_spectra(): +def test_enumerate2(): import angst n = 100 @@ -15,7 +15,7 @@ def test_enumerate_spectra(): indices = [(i, j) for i in range(n) for j in range(i, -1, -1)] # iterator that will enumerate the spectra for checking - it = angst.enumerate_spectra(spectra) + it = angst.enumerate2(spectra) # go through expected indices and values and compare for k, (i, j) in enumerate(indices): @@ -26,13 +26,13 @@ def test_enumerate_spectra(): next(it) -def test_spectra_indices(): +def test_indices2(): import angst - assert list(angst.spectra_indices(0)) == [] - assert list(angst.spectra_indices(1)) == [(0, 0)] - assert list(angst.spectra_indices(2)) == [(0, 0), (1, 1), (1, 0)] - assert list(angst.spectra_indices(3)) == [ + assert list(angst.indices2(0)) == [] + assert list(angst.indices2(1)) == [(0, 0)] + assert list(angst.indices2(2)) == [(0, 0), (1, 1), (1, 0)] + assert list(angst.indices2(3)) == [ (0, 0), (1, 1), (1, 0),