add size parameter to generate_flm()

Author: ntessore

s2fft/utils/signal_generator.py

@@ -74,6 +74,7 @@ def generate_flm(
     spin: int = 0,
     reality: bool = False,
     using_torch: bool = False,
+    size: tuple[int, ...] | int | None = None,
 ) -> np.ndarray | torch.Tensor:
     r"""
     Generate a 2D set of random harmonic coefficients.
@@ -94,29 +95,39 @@ def generate_flm(
 
         using_torch (bool, optional): Desired frontend functionality. Defaults to False.
 
+        size (tuple[int, ...] | int | None, optional): Shape of realisations.
+
     Returns:
         np.ndarray: Random set of spherical harmonic coefficients.
 
     """
-    flm = np.zeros(samples.flm_shape(L), dtype=np.complex128)
+    # always turn size into a tuple of int
+    if size is None:
+        size = ()
+    elif isinstance(size, int):
+        size = (size,)
+    elif not (isinstance(size, tuple) and all(isinstance(_, int) for _ in size)):
+        raise TypeError("size must be int or tuple of int")
+
+    flm = np.zeros((*size, *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)
+    flm[..., min_el:L, L - 1] = rng.standard_normal((*size, 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)
+    rand_size = (*size, len(m_indices))
     # Generate independent complex coefficients for positive m
-    flm[el_indices, L - 1 + m_indices] = complex_normal(rng, len_indices, var=2)
+    flm[..., el_indices, L - 1 + m_indices] = complex_normal(rng, rand_size, 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()
+        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)
+        flm[..., el_indices, L - 1 - m_indices] = complex_normal(rng, rand_size, var=2)
     return torch.from_numpy(flm) if using_torch else flm
 
 

tests/test_signal_generator.py

@@ -55,6 +55,14 @@ def check_flm_conjugate_symmetry(flm, L, min_el):
             assert flm[el, L - 1 - m] == (-1) ** m * flm[el, L - 1 + m].conj()
 
 
+def check_flm_unequal(flm1, flm2, L, min_el):
+    """assert that two passed flm are elementwise unequal"""
+    for el in range(L):
+        for m in range(L):
+            if not (el < min_el or m > el):
+                assert flm1[el, L - 1 + m] != flm2[el, L - 1 - m]
+
+
 @pytest.mark.parametrize("L", L_values_to_test)
 @pytest.mark.parametrize("L_lower", L_lower_to_test)
 @pytest.mark.parametrize("spin", spin_to_test)
@@ -76,6 +84,24 @@ def test_generate_flm(rng, L, L_lower, spin, reality):
         assert np.allclose(f_complex.real, f_real)
 
 
+@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)
+def test_generate_flm_size(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, size=2)
+    assert flm.shape == (2,) + smp.s2_samples.flm_shape(L)
+    check_flm_zeros(flm[0], L, max(L_lower, abs(spin)))
+    check_flm_zeros(flm[1], L, max(L_lower, abs(spin)))
+    check_flm_unequal(flm[0], flm[1], L, max(L_lower, abs(spin)))
+
+    flm = gen.generate_flm(rng, L, L_lower, spin, reality, size=(3, 4))
+    assert flm.shape == (3, 4) + smp.s2_samples.flm_shape(L)
+
+
 def check_flmn_zeros(flmn, L, N, L_lower):
     for n in range(-N + 1, N):
         min_el = max(L_lower, abs(n))