gh-509: inv_triangle_number -> nfields_from_nspectra (plus make it public) (#527)

Co-authored-by: Nicolas Tessore <[email protected]>

Author: Saransh-cpp

examples/2-advanced/legacy-mode.ipynb

@@ -64,7 +64,7 @@
    "source": [
     "def plot_spectra(*spectra, labels=[], log=True):\n",
     "    nspec = max(len(cls) for cls in spectra)\n",
-    "    nfields = glass.fields.inv_triangle_number(nspec)\n",
+    "    nfields = glass.nfields_from_nspectra(nspec)\n",
     "    fig, ax = plt.subplots(\n",
     "        nfields, nfields, figsize=(1.1 * nfields, 1.0 * nfields), sharex=True, sharey=True, layout=\"constrained\",\n",
     "    )\n",

glass/__init__.py

@@ -46,6 +46,7 @@
     "multalm",
     "multi_plane_matrix",
     "multi_plane_weights",
+    "nfields_from_nspectra",
     "partition",
     "position_weights",
     "positions_from_delta",
@@ -97,6 +98,7 @@
     lognormal_gls,
     lognormal_shift_hilbert2011,
     multalm,
+    nfields_from_nspectra,
     regularized_spectra,
     solve_gaussian_spectra,
     spectra_indices,

glass/fields.py

@@ -26,15 +26,16 @@
     Cls = Sequence[NDArray[Any]]
 
 
-def inv_triangle_number(triangle_number: int) -> int:
+def nfields_from_nspectra(nspectra: int) -> int:
     r"""
-    The :math:`n`-th triangle number is :math:`T_n = n \, (n+1)/2`.  If
-    the argument is :math:`T_n`, then :math:`n` is returned.  Otherwise,
-    a :class:`ValueError` is raised.
+    Returns the number of fields for a number of spectra.
+
+    Given the number of spectra *nspectra*, returns the number of
+    fields *n* such that ``n * (n + 1) // 2 == nspectra``.
     """
-    n = math.floor(math.sqrt(2 * triangle_number))
-    if n * (n + 1) // 2 != triangle_number:
-        msg = f"not a triangle number: {triangle_number}"
+    n = math.floor(math.sqrt(2 * nspectra))
+    if n * (n + 1) // 2 != nspectra:
+        msg = f"invalid number of spectra: {nspectra}"
         raise ValueError(msg)
     return n
 
@@ -274,11 +275,7 @@ def discretized_cls(
 
     """
     if ncorr is not None:
-        try:
-            n = inv_triangle_number(len(cls))
-        except ValueError:
-            msg = "length of cls array is not a triangle number"
-            raise ValueError(msg) from None
+        n = nfields_from_nspectra(len(cls))
         cls = [
             cls[i * (i + 1) // 2 + j] if j <= ncorr else np.asarray([])
             for i in range(n)
@@ -371,7 +368,7 @@ def generate_gaussian(
 
     # number of gls and number of fields
     ngls = len(gls)
-    ngrf = inv_triangle_number(ngls)
+    ngrf = nfields_from_nspectra(ngls)
 
     # number of correlated fields if not specified
     if ncorr is None:
@@ -591,11 +588,7 @@ def effective_cls(
 
     """
     # this is the number of fields
-    try:
-        n = inv_triangle_number(len(cls))
-    except ValueError:
-        msg = "length of cls is not a triangle number"
-        raise ValueError(msg) from None
+    n = nfields_from_nspectra(len(cls))
 
     # find lmax if not given
     if lmax is None:
@@ -844,7 +837,7 @@ def glass_to_healpix_spectra(spectra: Cls) -> Cls:
         Sequence of spectra in HEALPix order.
 
     """
-    n = inv_triangle_number(len(spectra))
+    n = nfields_from_nspectra(len(spectra))
 
     comb = [(i, j) for i, j in spectra_indices(n)]
     return [spectra[comb.index((i + k, i))] for k in range(n) for i in range(n - k)]
@@ -867,7 +860,7 @@ def healpix_to_glass_spectra(spectra: Cls) -> Cls:
         Sequence of spectra in GLASS order.
 
     """
-    n = inv_triangle_number(len(spectra))
+    n = nfields_from_nspectra(len(spectra))
 
     comb = [(i + k, i) for k in range(n) for i in range(n - k)]
     return [spectra[comb.index((i, j))] for i, j in spectra_indices(n)]
@@ -914,11 +907,7 @@ def cov_from_spectra(spectra: Cls, *, lmax: int | None = None) -> NDArray[Any]:
 
     """
     # recover the number of fields from the number of spectra
-    try:
-        n = inv_triangle_number(len(spectra))
-    except ValueError:
-        msg = "invalid number of spectra"
-        raise ValueError(msg) from None
+    n = nfields_from_nspectra(len(spectra))
 
     if lmax is None:  # noqa: SIM108
         # maximum length in input spectra

tests/test_fields.py

@@ -270,9 +270,7 @@ def test_discretized_cls() -> None:
 
     # check ValueError for triangle number
 
-    with pytest.raises(
-        ValueError, match="length of cls array is not a triangle number"
-    ):
+    with pytest.raises(ValueError, match="invalid number of spectra:"):
         glass.discretized_cls([np.arange(10), np.arange(10)], ncorr=1)
 
     # ncorr not None
@@ -307,7 +305,7 @@ def test_effective_cls() -> None:
 
     # check ValueError for triangle number
 
-    with pytest.raises(ValueError, match="length of cls is not a triangle number"):
+    with pytest.raises(ValueError, match="invalid number of spectra:"):
         glass.effective_cls([np.arange(10), np.arange(10)], np.ones((2, 1)))
 
     # check ValueError for triangle number
@@ -418,15 +416,15 @@ def test_getcl() -> None:
             np.testing.assert_array_equal(result[2:], expected)
 
 
-def test_inv_triangle_number():
+def test_nfields_from_nspectra():
     for n in range(10_000):
-        assert glass.fields.inv_triangle_number(n * (n + 1) // 2) == n
+        assert glass.nfields_from_nspectra(n * (n + 1) // 2) == n
 
     not_triangle_numbers = [2, 4, 5, 7, 8, 9, 11, 12, 13, 14, 16, 17, 18, 19, 20]
 
     for t in not_triangle_numbers:
-        with pytest.raises(ValueError, match="not a triangle number"):
-            glass.fields.inv_triangle_number(t)
+        with pytest.raises(ValueError, match=f"invalid number of spectra: {t}"):
+            glass.nfields_from_nspectra(t)
 
 
 def test_enumerate_spectra():