Make coefficients real for m=0 and factor out functions

Author: matt-graham

s2fft/utils/signal_generator.py

@@ -7,6 +7,66 @@
 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))]
+    )
+    ```
+
+    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,
@@ -39,24 +99,24 @@ def generate_flm(
 
     """
     flm = np.zeros(samples.flm_shape(L), dtype=np.complex128)
-
     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)
-    if not reality:
-        flm[min_el:L, L - 1] += 1j * rng.standard_normal(L - min_el)
-    m_indices, el_indices = np.triu_indices(n=L, k=1, m=L) + np.array([[1], [0]])
-    if min_el > 0:
-        in_range_el = el_indices >= min_el
-        m_indices = m_indices[in_range_el]
-        el_indices = el_indices[in_range_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)
-    flm[el_indices, L - 1 - m_indices] = rng.standard_normal(
-        len_indices
-    ) + 1j * rng.standard_normal(len_indices)
-    flm[el_indices, L - 1 + m_indices] = (-1) ** m_indices * np.conj(
-        flm[el_indices, L - 1 - 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