Add more unit tests for integer roots

galois/_math.py

@@ -133,7 +133,7 @@ def iroot(n: int, k: int) -> int:
     n : int
         A non-negative integer.
     k : int
-        The root :math:`k`, must be at least 2.
+        The positive root :math:`k`.
 
     Returns
     -------
@@ -154,12 +154,14 @@ def iroot(n: int, k: int) -> int:
         raise TypeError(f"Argument `k` must be an integer, not {type(k)}.")
     if not n >= 0:
         raise ValueError(f"Argument `n` must be non-negative, not {n}.")
-    if not k >= 2:
-        raise ValueError(f"Argument `k` must be at least 2, not {k}.")
+    if not k >= 1:
+        raise ValueError(f"Argument `k` must be at least 1, not {k}.")
     n, k = int(n), int(k)
 
     if n == 0:
         return 0
+    if k == 1:
+        return n
 
     # https://stackoverflow.com/a/39191163/11694321
     u = n

scripts/generate_int_test_vectors.py

@@ -11,7 +11,7 @@
 
 import sage
 import numpy as np
-from sage.all import xgcd, lcm, prod, isqrt
+from sage.all import Integer, xgcd, lcm, prod, isqrt
 
 PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "tests")
 FOLDER = os.path.join(PATH, "data")
@@ -96,3 +96,15 @@ def save_pickle(d, folder, name):
     Z[i] = int(z)
 d = {"X": X, "Z": Z}
 save_pickle(d, FOLDER, "isqrt.pkl")
+
+set_seed(SEED + 106)
+X = [random.randint(0, 1000) for _ in range(20)] + [random.randint(1000, 1_000_000_000) for _ in range(20)]
+R = [random.randint(1, 6) for _ in range(40)]
+Z = [0,]*len(X)
+for i in range(len(X)):
+    x = X[i]
+    r = R[i]
+    z = Integer(x).nth_root(r, truncate_mode=True)[0]
+    Z[i] = int(z)
+d = {"X": X, "R": R, "Z": Z}
+save_pickle(d, FOLDER, "iroot.pkl")

tests/conftest.py

@@ -52,3 +52,11 @@ def isqrt():
         print(f"Loading {f}...")
         d = pickle.load(f)
     return d
+
+
+@pytest.fixture(scope="session")
+def iroot():
+    with open(os.path.join(FOLDER, "iroot.pkl"), "rb") as f:
+        print(f"Loading {f}...")
+        d = pickle.load(f)
+    return d

tests/test_math.py

@@ -107,3 +107,21 @@ def test_isqrt(isqrt):
     X, Z = isqrt["X"], isqrt["Z"]
     for i in range(len(X)):
         assert galois.isqrt(X[i]) == Z[i]
+
+
+def test_iroot_exceptions():
+    with pytest.raises(TypeError):
+        galois.iroot(9.0, 3)
+    with pytest.raises(TypeError):
+        galois.iroot(9, 3.0)
+    with pytest.raises(ValueError):
+        galois.iroot(-9, 3)
+
+
+def test_iroot(iroot):
+    X, R, Z = iroot["X"], iroot["R"], iroot["Z"]
+    for i in range(len(X)):
+        assert galois.iroot(X[i], R[i]) == Z[i]
+
+    galois.iroot(10, 1) == 10
+    assert galois.iroot(0, 2) == 0

tests/test_number_theory.py

@@ -209,37 +209,6 @@ def test_is_cyclic():
 # Integer arithmetic
 ###############################################################################
 
-def test_iroot_exceptions():
-    with pytest.raises(TypeError):
-        galois.iroot(9.0, 3)
-    with pytest.raises(TypeError):
-        galois.iroot(9, 3.0)
-    with pytest.raises(ValueError):
-        galois.iroot(-9, 3)
-    with pytest.raises(ValueError):
-        galois.iroot(9, 1)
-
-
-def test_iroot():
-    """
-    Sage:
-        N = 20
-        lut = []
-        for _ in range(N):
-            n = Integer(randint(0, 1_000_000))
-            k = randint(2, 6)
-            x = n.nth_root(k, truncate_mode=True)[0]
-            lut.append((n, k, x))
-        print(lut)
-    """
-    LUT = [(779174, 4, 29), (867742, 4, 30), (709111, 2, 842), (616365, 6, 9), (615576, 2, 784), (259784, 2, 509), (862570, 2, 928), (553097, 2, 743), (929919, 4, 31), (841722, 6, 9), (658636, 3, 87), (326492, 5, 12), (195217, 4, 21), (969412, 3, 98), (95064, 3, 45), (550943, 3, 81), (171374, 3, 55), (881656, 3, 95), (915960, 6, 9), (810062, 2, 900)]
-    for item in LUT:
-        n, k, x = item
-        assert galois.iroot(n, k) == x
-
-    assert galois.iroot(0, 2) == 0
-
-
 def test_ilog_exceptions():
     with pytest.raises(TypeError):
         galois.ilog(9.0, 2)