Merge pull request #252 from astro-informatics/mmg/vectorized-signal-generators

Vectorize signal generator functions

Author: matt-graham

s2fft/transforms/wigner.py

@@ -156,6 +156,9 @@ def inverse_numpy(
         )
     fban = np.zeros(samples.f_shape(L, N, sampling, nside), dtype=np.complex128)
 
+    # Copy flmn argument to avoid in-place updates being propagated back to caller
+    flmn = flmn.copy()
+
     flmn[:, L_lower:] = np.einsum(
         "...nlm,...l->...nlm",
         flmn[:, L_lower:],

s2fft/utils/signal_generator.py

@@ -1,18 +1,80 @@
+from __future__ import annotations
+
 import numpy as np
 import torch
 
 from s2fft.sampling import s2_samples as samples
 from s2fft.sampling import so3_samples as wigner_samples
 
 
+def complex_normal(
+    rng: np.random.Generator,
+    size: int | tuple[int],
+    var: float,
+) -> np.ndarray:
+    """
+    Generate array of samples from zero-mean complex normal distribution.
+
+    For `z ~ ComplexNormal(0, var)` we have that `imag(z) ~ Normal(0, var/2)` and
+    `real(z) ~ Normal(0, var/2)` where `Normal(μ, σ²)` is the (real-valued) normal
+    distribution with mean parameter `μ` and variance parameter `σ²`.
+
+    Args:
+        rng: Numpy random generator object to generate samples using.
+        size: Output shape of array to generate.
+        var: Variance of complex normal distribution to generate samples from.
+
+    Returns:
+        Complex-valued array of shape `size` contained generated samples.
+
+    """
+    return (rng.standard_normal(size) + 1j * rng.standard_normal(size)) * (
+        var / 2
+    ) ** 0.5
+
+
+def complex_el_and_m_indices(L: int, min_el: int) -> tuple[np.ndarray, np.ndarray]:
+    """
+    Generate pairs of el, m indices for accessing complex harmonic coefficients.
+
+    Equivalent to nested list-comprehension based implementation
+
+    ```
+    el_indices, m_indices = np.array(
+        [(el, m) for el in range(min_el, L) for m in range(1, el + 1)]
+    ).T
+    ```
+
+    For `L, min_el = 1024, 0`, this implementation is around 80x quicker in
+    benchmarks compared to list-comprehension implementation.
+
+    Args:
+        L: Harmonic band-limit.
+        min_el: Inclusive lower-bound for el indices.
+
+    Returns:
+        Tuple `(el_indices, m_indices)` with both entries 1D integer-valued NumPy arrays
+        of same size, with values of corresponding entries corresponding to pairs of
+        el and m indices.
+
+    """
+    el_indices, m_indices = np.tril_indices(m=L, k=-1, n=L)
+    m_indices += 1
+    if min_el > 0:
+        in_range_el = el_indices >= min_el
+        el_indices = el_indices[in_range_el]
+        m_indices = m_indices[in_range_el]
+    return el_indices, m_indices
+
+
 def generate_flm(
     rng: np.random.Generator,
     L: int,
     L_lower: int = 0,
     spin: int = 0,
     reality: bool = False,
     using_torch: bool = False,
-) -> np.ndarray:
+) -> np.ndarray | torch.Tensor:
     r"""
     Generate a 2D set of random harmonic coefficients.
 
@@ -37,20 +99,24 @@ def generate_flm(
 
     """
     flm = np.zeros(samples.flm_shape(L), dtype=np.complex128)
-
-    for el in range(max(L_lower, abs(spin)), L):
-        if reality:
-            flm[el, 0 + L - 1] = rng.normal()
-        else:
-            flm[el, 0 + L - 1] = rng.normal() + 1j * rng.normal()
-
-        for m in range(1, el + 1):
-            flm[el, m + L - 1] = rng.normal() + 1j * rng.normal()
-            if reality:
-                flm[el, -m + L - 1] = (-1) ** m * np.conj(flm[el, m + L - 1])
-            else:
-                flm[el, -m + L - 1] = rng.normal() + 1j * rng.normal()
-
+    min_el = max(L_lower, abs(spin))
+    # m = 0 coefficients are always real
+    flm[min_el:L, L - 1] = rng.standard_normal(L - min_el)
+    # Construct arrays of m and el indices for entries in flm corresponding to complex-
+    # valued coefficients (m > 0)
+    el_indices, m_indices = complex_el_and_m_indices(L, min_el)
+    len_indices = len(m_indices)
+    # Generate independent complex coefficients for positive m
+    flm[el_indices, L - 1 + m_indices] = complex_normal(rng, len_indices, var=2)
+    if reality:
+        # Real-valued signal so set complex coefficients for negative m using conjugate
+        # symmetry such that flm[el, L - 1 - m] = (-1)**m * flm[el, L - 1 + m].conj
+        flm[el_indices, L - 1 - m_indices] = (-1) ** m_indices * (
+            flm[el_indices, L - 1 + m_indices].conj()
+        )
+    else:
+        # Non-real signal so generate independent complex coefficients for negative m
+        flm[el_indices, L - 1 - m_indices] = complex_normal(rng, len_indices, var=2)
     return torch.from_numpy(flm) if using_torch else flm
 
 
@@ -61,7 +127,7 @@ def generate_flmn(
     L_lower: int = 0,
     reality: bool = False,
     using_torch: bool = False,
-) -> np.ndarray:
+) -> np.ndarray | torch.Tensor:
     r"""
     Generate a 3D set of random Wigner coefficients.
 
@@ -87,26 +153,49 @@ def generate_flmn(
 
     """
     flmn = np.zeros(wigner_samples.flmn_shape(L, N), dtype=np.complex128)
-
     for n in range(-N + 1, N):
-        for el in range(max(L_lower, abs(n)), L):
-            if reality:
-                flmn[N - 1 + n, el, 0 + L - 1] = rng.normal()
-                flmn[N - 1 - n, el, 0 + L - 1] = (-1) ** n * flmn[
-                    N - 1 + n,
-                    el,
-                    0 + L - 1,
-                ]
-            else:
-                flmn[N - 1 + n, el, 0 + L - 1] = rng.normal() + 1j * rng.normal()
-
-            for m in range(1, el + 1):
-                flmn[N - 1 + n, el, m + L - 1] = rng.normal() + 1j * rng.normal()
-                if reality:
-                    flmn[N - 1 - n, el, -m + L - 1] = (-1) ** (m + n) * np.conj(
-                        flmn[N - 1 + n, el, m + L - 1]
-                    )
-                else:
-                    flmn[N - 1 + n, el, -m + L - 1] = rng.normal() + 1j * rng.normal()
-
+        min_el = max(L_lower, abs(n))
+        # Separately deal with m = 0 case
+        if reality:
+            if n == 0:
+                # For m = n = 0
+                # flmn[N - 1, el, L - 1] = flmn[N - 1, el, L - 1].conj (real-valued)
+                # Generate independent real coefficients for n = 0
+                flmn[N - 1, min_el:L, L - 1] = rng.standard_normal(L - min_el)
+            elif n > 0:
+                # Generate independent complex coefficients for positive n
+                flmn[N - 1 + n, min_el:L, L - 1] = complex_normal(
+                    rng, L - min_el, var=2
+                )
+                # For m = 0, n > 0
+                # flmn[N - 1 - n, el, L - 1] = (-1)**n * flmn[N - 1 + n, el, L - 1].conj
+                flmn[N - 1 - n, min_el:L, L - 1] = (-1) ** n * (
+                    flmn[N - 1 + n, min_el:L, L - 1].conj()
+                )
+        else:
+            flmn[N - 1 + n, min_el:L, L - 1] = complex_normal(rng, L - min_el, var=2)
+        # Construct arrays of m and el indices for entries in flmn slices for n
+        # corresponding to complex-valued coefficients (m > 0)
+        el_indices, m_indices = complex_el_and_m_indices(L, min_el)
+        len_indices = len(m_indices)
+        # Generate independent complex coefficients for positive m
+        flmn[N - 1 + n, el_indices, L - 1 + m_indices] = complex_normal(
+            rng, len_indices, var=2
+        )
+        if reality:
+            # Real-valued signal so set complex coefficients for negative m using
+            # conjugate symmetry relationship
+            #     flmn[N - 1 - n, el, L - 1 - m] =
+            #         (-1)**(m + n) * flmn[N - 1 + n, el, L - 1 + m].conj
+            # As (m_indices + n) can be negative use floating point value (-1.0) as
+            # base of exponentation operation to avoid Numpy
+            # 'ValueError: Integers to negative integer powers are not allowed' error
+            flmn[N - 1 - n, el_indices, L - 1 - m_indices] = (-1.0) ** (
+                m_indices + n
+            ) * flmn[N - 1 + n, el_indices, L - 1 + m_indices].conj()
+        else:
+            # Complex signal so generate independent complex coefficients for negative m
+            flmn[N - 1 + n, el_indices, L - 1 - m_indices] = complex_normal(
+                rng, len_indices, var=2
+            )
     return torch.from_numpy(flmn) if using_torch else flmn

tests/test_signal_generator.py

@@ -0,0 +1,119 @@
+import numpy as np
+import pytest
+
+import s2fft
+import s2fft.sampling as smp
+import s2fft.utils.signal_generator as gen
+
+L_values_to_test = [6, 7, 16]
+L_lower_to_test = [0, 1]
+spin_to_test = [-2, 0, 1]
+reality_values_to_test = [False, True]
+
+
+@pytest.mark.parametrize("size", (10, 100, 1000))
+@pytest.mark.parametrize("var", (1, 2))
+def test_complex_normal(rng, size, var):
+    samples = gen.complex_normal(rng, size, var)
+    assert samples.dtype == np.complex128
+    assert samples.size == size
+    mean = samples.mean()
+    # Error in real + imag components of mean estimate ~ Normal(0, (var / 2) / size)
+    # Therefore difference between mean estimate and true zero value should be
+    # less than 3 * sqrt(var / (2 * size)) with probability 0.997
+    mean_error_tol = 3 * (var / (2 * size)) ** 0.5
+    assert abs(mean.imag) < mean_error_tol and abs(mean.real) < mean_error_tol
+    # If S is (unbiased) sample variance estimate then (size - 1) * S / var is a
+    # chi-squared distributed random variable with (size - 1) degrees of freedom
+    # For size >> 1, S ~approx Normal(var, 2 * var**2 / (size - 1)) so error in
+    # variance estimate should be less than 3 * sqrt(2 * var**2 / (size - 1))
+    # with high probability
+    assert abs(samples.var(ddof=1) - var) < 3 * (2 * var**2 / (size - 1)) ** 0.5
+
+
+@pytest.mark.parametrize("L", L_values_to_test)
+@pytest.mark.parametrize("min_el", [0, 1])
+def test_complex_el_and_m_indices(L, min_el):
+    expected_el_indices, expected_m_indices = np.array(
+        [(el, m) for el in range(min_el, L) for m in range(1, el + 1)]
+    ).T
+    el_indices, m_indices = gen.complex_el_and_m_indices(L, min_el)
+    assert (el_indices == expected_el_indices).all()
+    assert (m_indices == expected_m_indices).all()
+
+
+def check_flm_zeros(flm, L, min_el):
+    for el in range(L):
+        for m in range(L):
+            if el < min_el or m > el:
+                assert flm[el, L - 1 + m] == flm[el, L - 1 - m] == 0
+
+
+def check_flm_conjugate_symmetry(flm, L, min_el):
+    for el in range(min_el, L):
+        for m in range(el + 1):
+            assert flm[el, L - 1 - m] == (-1) ** m * flm[el, L - 1 + m].conj()
+
+
+@pytest.mark.parametrize("L", L_values_to_test)
+@pytest.mark.parametrize("L_lower", L_lower_to_test)
+@pytest.mark.parametrize("spin", spin_to_test)
+@pytest.mark.parametrize("reality", reality_values_to_test)
+@pytest.mark.filterwarnings("ignore::RuntimeWarning")
+def test_generate_flm(rng, L, L_lower, spin, reality):
+    if reality and spin != 0:
+        pytest.skip("Reality only valid for scalar fields (spin=0).")
+    flm = gen.generate_flm(rng, L, L_lower, spin, reality)
+    assert flm.shape == smp.s2_samples.flm_shape(L)
+    assert flm.dtype == np.complex128
+    assert np.isfinite(flm).all()
+    check_flm_zeros(flm, L, max(L_lower, abs(spin)))
+    if reality:
+        check_flm_conjugate_symmetry(flm, L, max(L_lower, abs(spin)))
+        f_complex = s2fft.inverse(flm, L, spin=spin, reality=False, L_lower=L_lower)
+        assert np.allclose(f_complex.imag, 0)
+        f_real = s2fft.inverse(flm, L, spin=spin, reality=True, L_lower=L_lower)
+        assert np.allclose(f_complex.real, f_real)
+
+
+def check_flmn_zeros(flmn, L, N, L_lower):
+    for n in range(-N + 1, N):
+        min_el = max(L_lower, abs(n))
+        for el in range(L):
+            for m in range(L):
+                if el < min_el or m > el:
+                    assert (
+                        flmn[N - 1 + n, el, L - 1 + m]
+                        == flmn[N - 1 + n, el, L - 1 - m]
+                        == 0
+                    )
+
+
+def check_flmn_conjugate_symmetry(flmn, L, N, L_lower):
+    for n in range(-N + 1, N):
+        min_el = max(L_lower, abs(n))
+        for el in range(min_el, L):
+            for m in range(el + 1):
+                assert (
+                    flmn[N - 1 - n, el, L - 1 - m]
+                    == (-1) ** (m + n) * flmn[N - 1 + n, el, L - 1 + m].conj()
+                )
+
+
+@pytest.mark.parametrize("L", L_values_to_test)
+@pytest.mark.parametrize("N", [1, 2, 3])
+@pytest.mark.parametrize("L_lower", L_lower_to_test)
+@pytest.mark.parametrize("reality", reality_values_to_test)
+@pytest.mark.filterwarnings("ignore::RuntimeWarning")
+def test_generate_flmn(rng, L, N, L_lower, reality):
+    flmn = gen.generate_flmn(rng, L, N, L_lower, reality)
+    assert flmn.shape == smp.so3_samples.flmn_shape(L, N)
+    assert flmn.dtype == np.complex128
+    assert np.isfinite(flmn).all()
+    check_flmn_zeros(flmn, L, N, L_lower)
+    if reality:
+        check_flmn_conjugate_symmetry(flmn, L, N, L_lower)
+        f_complex = s2fft.wigner.inverse(flmn, L, N, reality=False, L_lower=L_lower)
+        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)

tests/test_spherical_precompute.py

@@ -306,7 +306,7 @@ def test_transform_inverse_high_spin(
     kernel = c.spin_spherical_kernel(L, spin, reality, sampling, forward=False)
 
     f = inverse(flm, L, spin, kernel, sampling, reality, "numpy")
-    tol = 1e-8 if sampling.lower() in ["dh", "gl"] else 1e-12
+    tol = 1e-8 if sampling.lower() in ["dh", "gl"] else 1e-11
     np.testing.assert_allclose(f, f_check, atol=tol, rtol=tol)