Repository: glass-dev/glass

Issue:
Add `inv_triangle_number` back along with a helper function
> Maybe we should add the private function back in, then make two helper functions with clear purpose (and error msg) instead of using it directly?

_Originally posted by @ntessore in https://github.com/glass-dev/glass/pull/533#discussion_r1985333055_

Discussion:

Solution patch:
diff --git a/glass/fields.py b/glass/fields.py
index facf941..c39fd52 100644
--- a/glass/fields.py
+++ b/glass/fields.py
@@ -51,17 +51,32 @@ except ImportError:
         return decorator
 
 
+def _inv_triangle_number(triangle_number: 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.
+    """
+    n = math.floor(math.sqrt(2 * triangle_number))
+    if n * (n + 1) // 2 != triangle_number:
+        msg = f"not a triangle number: {triangle_number}"
+        raise ValueError(msg)
+    return n
+
+
 def nfields_from_nspectra(nspectra: int) -> int:
     r"""
     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``.
+    fields *n* such that ``n * (n + 1) // 2 == nspectra`` or raises
+    a :class:`ValueError` if the number of spectra is invalid.
     """
-    n = math.floor(math.sqrt(2 * nspectra))
-    if n * (n + 1) // 2 != nspectra:
+    try:
+        n = _inv_triangle_number(nspectra)
+    except ValueError:
         msg = f"invalid number of spectra: {nspectra}"
-        raise ValueError(msg)
+        raise ValueError(msg) from None
     return n
 
 


Test patch:
diff --git a/tests/test_fields.py b/tests/test_fields.py
index 0db89b6..ced908d 100644
--- a/tests/test_fields.py
+++ b/tests/test_fields.py
@@ -5,6 +5,11 @@ import pytest
 import glass
 
 
+@pytest.fixture(scope="session")
+def not_triangle_numbers() -> list[int]:
+    return [2, 4, 5, 7, 8, 9, 11, 12, 13, 14, 16, 17, 18, 19, 20]
+
+
 def test_iternorm() -> None:
     # check output shapes and types
 
@@ -404,11 +409,18 @@ def test_getcl() -> None:
             np.testing.assert_array_equal(result[2:], expected)
 
 
-def test_nfields_from_nspectra():
+def test_is_inv_triangle_number(not_triangle_numbers: list[int]):
     for n in range(10_000):
-        assert glass.nfields_from_nspectra(n * (n + 1) // 2) == n
+        assert glass.fields._inv_triangle_number(n * (n + 1) // 2) == n
+
+    for t in not_triangle_numbers:
+        with pytest.raises(ValueError, match=f"not a triangle number: {t}"):
+            glass.fields._inv_triangle_number(t)
+
 
-    not_triangle_numbers = [2, 4, 5, 7, 8, 9, 11, 12, 13, 14, 16, 17, 18, 19, 20]
+def test_nfields_from_nspectra(not_triangle_numbers: list[int]):
+    for n in range(10_000):
+        assert glass.nfields_from_nspectra(n * (n + 1) // 2) == n
 
     for t in not_triangle_numbers:
         with pytest.raises(ValueError, match=f"invalid number of spectra: {t}"):