generate flm from spectra

Author: ntessore

pyproject.toml

@@ -72,6 +72,7 @@ tests = [
     "numpy<2",  # Required currently due to lack of Numpy v2 compatible pyssht release
     "pytest",
     "pytest-cov",
+    "pytest-rerunfailures",
     "so3",
     "pyssht",
 ]

s2fft/utils/signal_generator.py

@@ -6,6 +6,10 @@
 from s2fft.sampling import s2_samples as samples
 from s2fft.sampling import so3_samples as wigner_samples
 
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    import jax
+
 
 def complex_normal(
     rng: np.random.Generator,
@@ -210,3 +214,60 @@ def generate_flmn(
                 rng, len_indices, var=2
             )
     return torch.from_numpy(flmn) if using_torch else flmn
+
+
+def generate_flm_from_spectra(
+    rng: np.random.Generator,
+    spectra: np.ndarray | jax.Array,
+) -> np.ndarray | jax.Array:
+    r"""
+    Generate a stack of random harmonic coefficients from power spectra.
+
+    The input power spectra must be a stack of shape *(K, K, L)* where
+    *K* is the number of fields to be sampled, and *L* is the harmonic
+    band-limit.
+
+    Args:
+        rng (Generator): Random number generator.
+
+        spectra (np.ndarray | jax.Array): Stack of angular power spectra.
+
+    Returns:
+        np.ndarray | jax.Array: A stack of random spherical harmonic
+        coefficients with the given power spectra.
+
+    """
+    # get the Array API namespace from spectra
+    xp = spectra.__array_namespace__()
+
+    # check input
+    if spectra.ndim != 3 or spectra.shape[0] != spectra.shape[1]:
+        raise ValueError("shape of spectra must be (K, K, L)")
+
+    # K is the number of fields, L is the band limit
+    *_, K, L = spectra.shape
+
+    # permute shape (K, K, L) -> (L, K, K)
+    spectra = xp.permute_dims(spectra, (2, 0, 1))
+
+    # SVD for matrix square root
+    # not using cholesky() here because matrix may be semi-definite
+    # divides spectra by 2 for correct amplitude
+    u, s, vh = xp.linalg.svd(spectra / 2, full_matrices=False)
+
+    # compute the matrix square root for sampling
+    a = u @ (xp.sqrt(s[..., None]) * vh)
+
+    # permute shape (L, K, K) -> (K, K, L)
+    a = xp.permute_dims(a, (1, 2, 0))
+
+    # sample the random coefficients
+    # always use reality=True, this could be real fields or E/B modes
+    # shape of flm is (K, L, M)
+    flm = generate_flm(rng, L, reality=True, size=K)
+
+    # compute the matrix multiplication by hand, because we have a mix of
+    # contraction (dim=K) and broadcasting (dim=L)
+    flm = (a[..., None] * flm).sum(axis=-3)
+
+    return flm

tests/test_signal_generator.py

@@ -1,5 +1,7 @@
+import jax.numpy as jnp
 import numpy as np
 import pytest
+from jax.test_util import check_grads
 
 import s2fft
 import s2fft.sampling as smp
@@ -143,3 +145,93 @@ def test_generate_flmn(rng, L, N, L_lower, reality):
         assert np.allclose(f_complex.imag, 0)
         f_real = s2fft.wigner.inverse(flmn, L, N, reality=True, L_lower=L_lower)
         assert np.allclose(f_complex.real, f_real)
+
+
+def gaussian_covariance(spectra):
+    """Gaussian covariance for a stack of spectra.
+
+    If the shape of *spectra* is *(K, K, L)*, the shape of the
+    covariance is *(L, C, C)*, where ``C = K * (K + 1) // 2``
+    is the number of independent spectra.
+
+    """
+    _, K, L = spectra.shape
+    row, col = np.tril_indices(K)
+    cov = np.zeros((L, row.size, col.size))
+    ell = np.arange(L)
+    for i, j in np.ndindex(row.size, col.size):
+        cov[:, i, j] = (
+            spectra[row[i], row[j]] * spectra[col[i], col[j]]
+            + spectra[row[i], col[j]] * spectra[col[i], row[j]]
+        ) / (2 * ell + 1)
+    return cov
+
+
+@pytest.mark.flaky
+@pytest.mark.parametrize("L", L_values_to_test)
+@pytest.mark.parametrize("xp", [np, jnp])
+def test_generate_flm_from_spectra(rng, L, xp):
+    # number of fields to generate
+    K = 4
+
+    # correlation matrix for fields, applied to all ell
+    corr = xp.asarray(
+        [
+            [1.0, 0.1, -0.1, 0.1],
+            [0.1, 1.0, 0.1, -0.1],
+            [-0.1, 0.1, 1.0, 0.1],
+            [0.1, -0.1, 0.1, 1.0],
+        ],
+    )
+
+    ell = xp.arange(L)
+
+    # auto-spectra are power laws
+    powers = xp.arange(1, K + 1)
+    auto = 1 / (2 * ell + 1) ** powers[:, None]
+
+    # compute the spectra from auto and corr
+    spectra = xp.sqrt(auto[:, None, :] * auto[None, :, :]) * corr[:, :, None]
+    assert spectra.shape == (K, K, L)
+
+    # generate random flm from spectra
+    flm = s2fft.utils.signal_generator.generate_flm_from_spectra(rng, spectra)
+    assert flm.shape == (K, L, 2 * L - 1)
+
+    # compute the realised spectra
+    re, im = flm.real, flm.imag
+    result = (
+        re[None, :, :, :] * re[:, None, :, :] + im[None, :, :, :] * im[:, None, :, :]
+    )
+    result = result.sum(axis=-1) / (2 * ell + 1)
+
+    # compute covariance of sampled spectra
+    cov = gaussian_covariance(spectra)
+
+    # data vector, remove duplicate entries, and put L dim first
+    x = result - spectra
+    x = x[np.tril_indices(K)]
+    x = x.T
+
+    # compute chi2/n of realised spectra
+    y = xp.linalg.solve(cov, x[..., None])[..., 0]
+    n = x.size
+    chi2_n = (x * y).sum() / n
+
+    # make sure chi2/n is as expected
+    sigma = np.sqrt(2 / n)
+    assert np.fabs(chi2_n - 1.0) < 3 * sigma
+
+
+@pytest.mark.parametrize("L", L_values_to_test)
+def test_generate_flm_from_spectra_grads(L):
+    # fixed set of power spectra
+    ell = jnp.arange(L)
+    cl = 1 / (2 * ell + 1)
+    spectra = cl.reshape(1, 1, L)
+
+    def func(x):
+        rng = np.random.default_rng(42)
+        return s2fft.utils.signal_generator.generate_flm_from_spectra(rng, x)
+
+    check_grads(func, (spectra,), 1)