From aef61c0ac69ec9ca47d067c813350e2423c0b0ac Mon Sep 17 00:00:00 2001 From: Nicolas Tessore Date: Sat, 25 May 2024 14:24:15 +0100 Subject: [PATCH] support for Gaussian random fields --- LICENSE | 2 +- docs/grf.rst | 34 +++++ docs/index.rst | 10 +- docs/{spectra.rst => twopoint.rst} | 32 +++-- pyproject.toml | 1 + src/angst/__init__.py | 22 ++- src/angst/_spectra.py | 43 ------ src/angst/_twopoint.py | 147 ++++++++++++++++++++ src/angst/grf.py | 214 +++++++++++++++++++++++++++++ 9 files changed, 442 insertions(+), 63 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 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/grf.rst b/docs/grf.rst new file mode 100644 index 0000000..34037fc --- /dev/null +++ b/docs/grf.rst @@ -0,0 +1,34 @@ +:mod:`angst.grf` --- Gaussian random fields +=========================================== + +.. currentmodule:: angst.grf +.. module:: angst.grf + + +Gaussian angular power spectra +------------------------------ + +.. autofunction:: solve + +.. class:: Transformation(Protocol) + + .. automethod:: __call__ + .. automethod:: inv + .. automethod:: der + + +Transformations +--------------- + +.. 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..978f9e3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,6 +18,7 @@ classifiers = [ requires-python = ">=3.8" 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..5817180 100644 --- a/src/angst/__init__.py +++ b/src/angst/__init__.py @@ -1,26 +1,38 @@ """angst -- angular statistics""" __all__ = [ + "cl2corr", + "cl2var", + "corr2cl", "displace", "displacement", - "enumerate_spectra", + "enumerate2", + "grf", "inv_triangle_number", - "spectra_indices", + "indices2", "__version__", "__version_tuple__", ] +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..dbf9848 --- /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(angst.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(angst.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..1684d4b --- /dev/null +++ b/src/angst/grf.py @@ -0,0 +1,214 @@ +""" +Transformations of Gaussian random fields. +""" + +from __future__ import annotations + +__all__ = [ + "Lognormal", + "LognormalXNormal", + "SquaredNormal", + "Transformation", + "solve", +] + +from dataclasses import dataclass + +import numpy as np + +# typing +from typing import Any, 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, + *, + initial: NDArray[Any] | None = None, + n: int | 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 : (m,) array + tfm : :class:`Transformation` + + Returns + ------- + gl : (m,) array + Gaussian angular power spectrum solution. + cl : (n,) 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 + + m = len(cl) + if n is not None: + if not isinstance(n, int): + raise TypeError("n must be integer") + elif n < m: + raise ValueError("n must be larger than the length of cl") + k = n - m + else: + k = 2 * len(cl) + + if initial is None: + gl = corr2cl(tfm.inv(cl2corr(cl), cl2var(cl))) + else: + gl = np.empty(m) + gl[: len(initial)] = initial[:m] + + if monopole is not None: + gl[0] = monopole + + gt = cl2corr(np.pad(gl, (0, k))) + var = cl2var(gl) + rl = corr2cl(tfm(gt, var)) + fl = rl[:m] - 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, k))) + dt = tfm.der(gt, var) + xl = -corr2cl(ft / dt)[:m] + if monopole is not None: + xl[0] = 0 + + while True: + gl_ = gl + xl + gt_ = cl2corr(np.pad(gl_, (0, k))) + var_ = cl2var(gl_) + rl_ = corr2cl(tfm(gt_, var_)) + fl_ = rl_[:m] - 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]: + return self.lamda1 * self.lamda2 * np.expm1(x) # type: ignore + + def inv(self, x: NDArray[Any], var: float) -> NDArray[Any]: + return np.log1p(x / (self.lamda1 * self.lamda2)) # type: ignore + + def der(self, x: NDArray[Any], var: float) -> NDArray[Any]: + return self.lamda1 * self.lamda2 * np.exp(x) # type: ignore + + +@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 = np.sqrt(1 - var) if self.a1 is None else self.a1 + a2 = np.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]: + aa, ll = self._pars(var) + return np.sqrt(x / (2 * ll) + aa**2) - aa # type: ignore + + def der(self, x: NDArray[Any], var: float) -> NDArray[Any]: + aa, ll = self._pars(var) + return 4 * ll * (x + aa)