Add tests for sqrt/gcd/xgcd

Author: oscarbenjamin

src/flint/test/test_all.py

@@ -454,7 +454,7 @@ def test_fmpz_poly():
     assert Z([1,2,-4]).height_bits() == 3
     assert Z([1,2,-4]).height_bits(signed=True) == -3
     assert Z([1,2,1]).sqrt() == Z([1,1])
-    assert raises(lambda: Z([1,2,2]).sqrt(), ValueError)
+    assert raises(lambda: Z([1,2,2]).sqrt(), DomainError)
     assert Z([1,0,2,0,3]).deflation() == (Z([1,2,3]), 2)
     assert Z([]).deflation() == (Z([]), 1)
     assert Z([1,1]).deflation() == (Z([1,1]), 1)
@@ -2204,7 +2204,7 @@ def test_fmpz_mod_poly():
 
         # sqrt
         f1 = R_test.random_element(irreducible=True)
-        assert raises(lambda: f1.sqrt(), ValueError)
+        assert raises(lambda: f1.sqrt(), DomainError)
         assert (f1*f1).sqrt() in [f1, -f1]
 
         # deflation
@@ -2777,7 +2777,7 @@ def setbad(obj, i, val):
 
         if not composite_characteristic:
             assert P([1, 2, 1]).sqrt() == P([1, 1])
-            assert raises(lambda: P([1, 2, 2]).sqrt(), ValueError), f"{P}, {P([1, 2, 2]).sqrt()}"
+            assert raises(lambda: P([1, 2, 2]).sqrt(), DomainError)
         elif nmod_poly_will_crash:
             pass
         else:
@@ -3231,7 +3231,7 @@ def quick_poly():
             assert (f * f).sqrt() == f
             if P is flint.fmpz_mpoly:
                 assert (f * f).sqrt(assume_perfect_square=True) == f
-            assert raises(lambda: quick_poly().sqrt(), ValueError)
+            assert raises(lambda: quick_poly().sqrt(), DomainError)
 
         p = quick_poly()
         assert p.derivative(0) == p.derivative("x0") == mpoly({(0, 0): 3, (1, 2): 8})
@@ -3364,6 +3364,35 @@ def factor_sqf(p):
             # All tests below would raise
             continue
 
+        assert S(0).sqrt() == S(0)
+        assert S(1).sqrt() == S(1)
+        assert S(4).sqrt()**2 == S(4)
+        for i in range(100):
+            try:
+                assert S(i).sqrt() ** 2 == S(i)
+            except DomainError:
+                pass
+
+        if characteristic == 0:
+            assert raises(lambda: S(-1).sqrt(), DomainError)
+        else:
+            try:
+                assert S(-1).sqrt() ** 2 == S(-1)
+            except DomainError:
+                pass
+
+        assert (0*x).sqrt() == 0*x
+        assert (1*x/x).sqrt() == 0*x + 1
+        assert (4*x/x).sqrt()**2 == 0*x + 4
+        for i in range(100):
+            try:
+                assert (i*x).sqrt() ** 2 == i*x
+            except DomainError:
+                pass
+        assert (x**2).sqrt() == x
+        assert (S(4)*x**2).sqrt()**2 == S(4)*x**2
+        assert raises(lambda: (x**2 + 1).sqrt(), DomainError)
+
         assert factor(0*x) == (S(0), [])
         assert factor(0*x + 1) == (S(1), [])
         assert factor(0*x + 3) == (S(3), [])
@@ -3379,7 +3408,22 @@ def factor_sqf(p):
         assert factor_sqf(x**2) == (S(1), [(x, 2)])
         assert factor_sqf(2*(x+1)) == (S(2), [(x+1, 1)])
 
-        assert factor(x*(x+1)) == (S(1), [(x, 1), (x+1, 1)])
+        assert (2*x).gcd(x) == x
+        assert (2*x).gcd(x**2) == x
+        assert (2*x).gcd(x**2 + 1) == S(1)
+
+        if not is_field:
+            # primitive gcd over Z
+            assert (2*x).gcd(4*x**2) == 2*x
+        else:
+            # monic gcd over Q, Z/pZ and GF(p^d)
+            assert (2*x).gcd(4*x**2) == x
+
+        if is_field and y is None:
+            # xgcd is defined and consistent for all univariate polynomials
+            # over a field (Q, Z/pZ, GF(p^d)).
+            assert (2*x).xgcd(4*x) == (x, P(0), P(1)/4)
+            assert (2*x).xgcd(4*x**2+1) == (P(1), -2*x, P(1))
 
         # mpoly types have a slightly different squarefree factorisation
         # because they handle trivial factors differently. It looks like a
@@ -3388,6 +3432,7 @@ def factor_sqf(p):
         #
         # Maybe it is worth making them consistent by absorbing the power
         # of x into a factor of equal multiplicity.
+        assert factor(x*(x+1)) == (S(1), [(x, 1), (x+1, 1)])
         if y is None:
             # *_poly types
             assert factor_sqf(x*(x+1)) == (S(1), [(x**2+x, 1)])
@@ -3399,7 +3444,7 @@ def factor_sqf(p):
         # a unique multiplicity.
         assert factor_sqf(x**2*(x+1)) == (S(1), [(x+1, 1), (x, 2)])
 
-        # This is the same for all types because there is no tivial monomial
+        # This is the same for all types because there is no trivial monomial
         # factor to extract.
         assert factor((x-1)*(x+1)) == (S(1), sort([(x-1, 1), (x+1, 1)]))
         assert factor_sqf((x-1)*(x+1)) == (S(1), [(x**2-1, 1)])
@@ -3433,6 +3478,8 @@ def factor_sqf(p):
 
         if y is not None:
 
+            # *_mpoly types
+
             assert factor(x*y+1) == (S(1), [(x*y+1, 1)])
             assert factor(x*y) == (S(1), [(x, 1), (y, 1)])
 
@@ -3451,6 +3498,13 @@ def factor_sqf(p):
                 else:
                     assert factor(p) == factor_sqf(p) == (S(2)/7, [(7*p/2, 1)])
 
+            if not is_field:
+                # primitive gcd over Z
+                assert (2*(x+y)).gcd(4*(x+y)**2) == 2*(x+y)
+            else:
+                # monic gcd over Q, Z/pZ and GF(p^d)
+                assert (2*(x+y)).gcd(4*(x+y)**2) == x + y
+
 
 def _all_matrices():
     """Return a list of matrix types and scalar types."""

src/flint/types/fmpq_mpoly.pyx

@@ -925,7 +925,7 @@ cdef class fmpq_mpoly(flint_mpoly):
         if fmpq_mpoly_sqrt(res.val, self.val, self.ctx.val):
             return res
         else:
-            raise ValueError("polynomial is not a perfect square")
+            raise DomainError("polynomial is not a perfect square")
 
     def factor(self):
         """

src/flint/types/fmpz_mod_mpoly.pyx

@@ -827,7 +827,7 @@ cdef class fmpz_mod_mpoly(flint_mpoly):
         if fmpz_mod_mpoly_sqrt(res.val, self.val, self.ctx.val):
             return res
         else:
-            raise ValueError("polynomial is not a perfect square")
+            raise DomainError("polynomial is not a perfect square")
 
     def factor(self):
         """

src/flint/types/fmpz_mod_poly.pyx

@@ -1496,7 +1496,7 @@ cdef class fmpz_mod_poly(flint_poly):
             res.val, self.val, res.ctx.mod.val
         )
         if check != 1:
-            raise ValueError(
+            raise DomainError(
                 f"Cannot compute square-root {self}"
             )
         return res

src/flint/types/fmpz_mpoly.pyx

@@ -903,7 +903,7 @@ cdef class fmpz_mpoly(flint_mpoly):
         if fmpz_mpoly_sqrt_heap(res.val, self.val, self.ctx.val, not assume_perfect_square):
             return res
         else:
-            raise ValueError("polynomial is not a perfect square")
+            raise DomainError("polynomial is not a perfect square")
 
     def factor(self):
         """

src/flint/types/fmpz_poly.pyx

@@ -591,7 +591,7 @@ cdef class fmpz_poly(flint_poly):
         if fmpz_poly_sqrt(v.val, self.val):
             return v
         else:
-            raise ValueError(f"Cannot compute square root of {self}")
+            raise DomainError(f"Cannot compute square root of {self}")
 
     def deflation(self):
         cdef fmpz_poly v

src/flint/types/fq_default_poly.pyx

@@ -982,7 +982,7 @@ cdef class fq_default_poly(flint_poly):
             res.val, self.val, res.ctx.field.val
         )
         if check != 1:
-            raise ValueError(
+            raise DomainError(
                 f"Cannot compute square-root {self}"
             )
         return res

src/flint/types/nmod.pyx

@@ -261,11 +261,16 @@ cdef class nmod(flint_scalar):
 
         """
         cdef nmod r
+        cdef mp_limb_t val
         r = nmod.__new__(nmod)
         r.mod = self.mod
-        r.val = n_sqrtmod(self.val, self.mod.n)
 
-        if r.val == 0:
+        if self.val == 0:
+            return r
+
+        val = n_sqrtmod(self.val, self.mod.n)
+        if val == 0:
             raise DomainError("no square root exists for %s mod %s" % (self.val, self.mod.n))
 
+        r.val = val
         return r

src/flint/types/nmod_mpoly.pyx

@@ -803,7 +803,7 @@ cdef class nmod_mpoly(flint_mpoly):
         if nmod_mpoly_sqrt(res.val, self.val, self.ctx.val):
             return res
         else:
-            raise ValueError("polynomial is not a perfect square")
+            raise DomainError("polynomial is not a perfect square")
 
     def factor(self):
         """

src/flint/types/nmod_poly.pyx

@@ -700,7 +700,7 @@ cdef class nmod_poly(flint_poly):
         if nmod_poly_sqrt(res.val, self.val):
             return res
         else:
-            raise ValueError(f"Cannot compute square root of {self}")
+            raise DomainError(f"Cannot compute square root of {self}")
 
     def deflation(self):
         cdef nmod_poly v