Add .sqrt() for all scalar types and add docstrings

Author: oscarbenjamin

src/flint/test/test_all.py

@@ -285,7 +285,7 @@ def test_fmpz_factor():
         (1296814508839693536173209832765271992846610925502473758289451540212712414540699659186801, 1)]
 
 def test_fmpz_functions():
-    T, F, VE, OE = True, False, ValueError, OverflowError
+    T, F, VE, OE, DE = True, False, ValueError, OverflowError, DomainError
     cases = [
         # (f, [f(-1), f(0), f(1), f(2), ... f(10)]),
         (lambda n: flint.fmpz(n).is_prime(),
@@ -331,11 +331,11 @@ def test_fmpz_functions():
         (lambda n: flint.fmpz(n).euler_phi(),
             [0, 0, 1, 1, 2, 2, 4, 2, 6, 4, 6, 4]),
         (lambda n: flint.fmpz(n).isqrt(),
-            [VE, 0, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3]),
+            [DE, 0, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3]),
         (lambda n: flint.fmpz(n).sqrtrem(),
-            [VE, (0, 0), (1, 0), (1, 1), (1, 2), (2, 0), (2, 1), (2, 2), (2, 3), (2, 4), (3, 0), (3, 1)]),
+            [DE, (0, 0), (1, 0), (1, 1), (1, 2), (2, 0), (2, 1), (2, 2), (2, 3), (2, 4), (3, 0), (3, 1)]),
         (lambda n: flint.fmpz(n).sqrtmod(11),
-            [VE, 0, 1, VE, 5, 2, 4, VE, VE, VE, 3, VE]),
+            [DE, 0, 1, DE, 5, 2, 4, DE, DE, DE, 3, DE]),
         (lambda n: flint.fmpz(n).root(3),
             [VE, 0, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2]),
         (lambda n: flint.fmpz(n).jacobi(3),
@@ -2781,8 +2781,6 @@ def setbad(obj, i, val):
         elif nmod_poly_will_crash:
             pass
         else:
-            # fmpz_mod_poly.sqrt() also crashes here:
-            # GR_MUST_SUCCEED failure: src/fmpz_mod_poly/sqrt_series.c
             assert raises(lambda: P([1, 2, 1]).sqrt(), DomainError)
 
         if P == flint.fmpq_poly:
@@ -3352,21 +3350,6 @@ def factor_sqf(p):
         check(p, coeff, factors)
         return coeff, sort(factors)
 
-    def sqrt(a):
-        if type(x) is flint.fq_default_poly:
-            try:
-                return S(a).sqrt()
-            except ValueError:
-                return None
-        elif characteristic != 0:
-            # XXX: fmpz(0).sqrtmod crashes
-            try:
-                return flint.fmpz(-1).sqrtmod(characteristic)
-            except ValueError:
-                return None
-        else:
-            return None
-
     for P, S, [x, y], is_field, characteristic in _all_polys_mpolys():
 
         if characteristic != 0 and not characteristic.is_prime():
@@ -3397,25 +3380,39 @@ def sqrt(a):
         assert factor_sqf(2*(x+1)) == (S(2), [(x+1, 1)])
 
         assert factor(x*(x+1)) == (S(1), [(x, 1), (x+1, 1)])
+
+        # mpoly types have a slightly different squarefree factorisation
+        # because they handle trivial factors differently. It looks like a
+        # monomial gcd is extracted but not recombined so the square-free
+        # factors might not have unique multiplicities.
+        #
+        # Maybe it is worth making them consistent by absorbing the power
+        # of x into a factor of equal multiplicity.
         if y is None:
+            # *_poly types
             assert factor_sqf(x*(x+1)) == (S(1), [(x**2+x, 1)])
         else:
-            # mpoly types have a different squarefree factorisation because
-            # they handle trivial factors differently...
-            #
-            # Maybe it is worth making them consistent by absorbing the power
-            # of x into a factor of equal multiplicity.
+            # *_mpoly types
             assert factor_sqf(x*(x+1)) == (S(1), [(x, 1), (x+1, 1)])
 
+        # This is the same for all types because the extracted monomial has
+        # 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
+        # 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)])
 
+        # Some finite fields have sqrt(-1) so we can factor x**2 + 1
+        try:
+            i = S(-1).sqrt()
+        except DomainError:
+            i = None
+
         p = 3*(x-1)**2*(x+1)**2*(x**2 + 1)**3
         assert factor_sqf(p) == (S(3), [(x**2 - 1, 2), (x**2 + 1, 3)])
 
-        i = sqrt(-1)
         if i is not None:
             assert factor(p) == (S(3), sort([(x+1, 2), (x-1, 2), (x+i, 3), (x-i, 3)]))
         else:
@@ -4163,7 +4160,7 @@ def test_fq_default():
             nqr = gf.random_element()
             if not nqr.is_square():
                 break
-        assert raises(lambda: nqr.sqrt(), ValueError)
+        assert raises(lambda: nqr.sqrt(), DomainError)
 
 
 def test_fq_default_poly():

src/flint/types/fmpq.pyx

@@ -486,3 +486,17 @@ cdef class fmpq(flint_scalar):
             return v
         else:
             raise OverflowError("fmpq_pow_fmpz(): exponent too large")
+
+    def sqrt(self):
+        """
+        Return exact rational square root of self or raise an error.
+
+            >>> fmpq(9, 4).sqrt()
+            3/2
+            >>> fmpq(8).sqrt()
+            Traceback (most recent call last):
+                ...
+            flint.utils.flint_exceptions.DomainError: not a square number
+
+        """
+        return fmpq(self.numer().sqrt(), self.denom().sqrt())

src/flint/types/fmpq_mpoly.pyx

@@ -859,6 +859,19 @@ cdef class fmpq_mpoly(flint_mpoly):
         return res
 
     def leading_coefficient(self):
+        """
+        Leading coefficient in the monomial ordering.
+
+            >>> from flint import Ordering
+            >>> ctx = fmpq_mpoly_ctx(2, Ordering.lex, ['x', 'y'])
+            >>> x, y = ctx.gens()
+            >>> p = 2*x*y + 3*x + 4*y**2 + 5
+            >>> p
+            2*x*y + 3*x + 4*y^2 + 5
+            >>> p.leading_coefficient()
+            2
+
+        """
         if fmpq_mpoly_is_zero(self.val, self.ctx.val):
             return fmpq(0)
         else:

src/flint/types/fmpz.pyx

@@ -883,30 +883,102 @@ cdef class fmpz(flint_scalar):
             return self.bit_length()
 
     def isqrt(self):
+        """
+        Return square root rounded down.
+
+            >>> fmpz(9).isqrt()
+            3
+            >>> fmpz(8).isqrt()
+            2
+
+        """
+        cdef fmpz v
+
+        if fmpz_sgn(self.val) < 0:
+            raise DomainError("integer square root of a negative number")
+
+        v = fmpz()
+        fmpz_sqrt(v.val, self.val)
+        return v
+
+    def sqrt(self):
+        """
+        Return exact integer square root of self or raise an error.
+
+            >>> fmpz(9).sqrt()
+            3
+            >>> fmpz(8).sqrt()
+            Traceback (most recent call last):
+                ...
+            flint.utils.flint_exceptions.DomainError: not a square number
+
+        """
         cdef fmpz v
+
         if fmpz_sgn(self.val) < 0:
-            raise ValueError("integer square root of a negative number")
+            raise DomainError("integer square root of a negative number")
+
         v = fmpz()
         fmpz_sqrt(v.val, self.val)
+
+        c = fmpz()
+        fmpz_mul(c.val, v.val, v.val)
+        if not fmpz_equal(c.val, self.val):
+            raise DomainError("not a square number")
+
         return v
 
     def sqrtrem(self):
+        """
+        Return the integer square root of self and remainder.
+
+            >>> fmpz(9).sqrtrem()
+            (3, 0)
+            >>> fmpz(8).sqrtrem()
+            (2, 4)
+            >>> c = fmpz(123456789012345678901234567890)
+            >>> u, v = c.sqrtrem()
+            >>> u ** 2 + v == c
+            True
+
+        """
         cdef fmpz u, v
+
         if fmpz_sgn(self.val) < 0:
-            raise ValueError("integer square root of a negative number")
+            raise DomainError("integer square root of a negative number")
+
         u = fmpz()
         v = fmpz()
         fmpz_sqrtrem(u.val, v.val, self.val)
+
         return u, v
 
     # warning: m should be prime!
-    # also if self is zero this crashes...
-    def sqrtmod(self, m):
+    def sqrtmod(self, p):
+        """
+        Return modular square root of self modulo *p* or raise an error.
+
+            >>> fmpz(10).sqrtmod(13)
+            6
+            >>> (6**2) % 13
+            10
+            >>> fmpz(11).sqrtmod(13)
+            Traceback (most recent call last):
+                ...
+            flint.utils.flint_exceptions.DomainError: modular square root does not exist
+
+        The modulus *p* must be a prime number.
+        """
         cdef fmpz v
+
         v = fmpz()
-        m = fmpz(m)
-        if not fmpz_sqrtmod(v.val, self.val, (<fmpz>m).val):
-            raise ValueError("unable to compute modular square root")
+        if fmpz_is_zero(self.val):
+            return v
+
+        p = fmpz(p)
+        if not fmpz_sqrtmod(v.val, self.val, (<fmpz>p).val):
+            raise DomainError("modular square root does not exist")
+
         return v
 
     def root(self, long n):

src/flint/types/fmpz_mod.pyx

@@ -10,6 +10,7 @@ from flint.flintlib.fmpz cimport(
     fmpz_is_probabprime,
     fmpz_mul,
     fmpz_invmod,
+    fmpz_sqrtmod,
     fmpz_divexact,
     fmpz_gcd,
     fmpz_is_one,
@@ -29,6 +30,9 @@ from flint.types.fmpz cimport(
 cimport cython
 cimport libc.stdlib
 
+from flint.utils.flint_exceptions import DomainError
+
+
 cdef class fmpz_mod_ctx:
     r"""
     Context object for creating :class:`~.fmpz_mod` initalised 
@@ -578,3 +582,34 @@ cdef class fmpz_mod(flint_scalar):
             )
 
         return res
+
+    def sqrt(self):
+        """
+        Return the square root of this ``fmpz_mod`` or raise an exception.
+
+            >>> ctx = fmpz_mod_ctx(13)
+            >>> s = ctx(10).sqrt()
+            >>> s
+            fmpz_mod(6, 13)
+            >>> s * s
+            fmpz_mod(10, 13)
+            >>> ctx(11).sqrt()
+            Traceback (most recent call last):
+                ...
+            flint.utils.flint_exceptions.DomainError: no square root exists for 11 mod 13
+
+        The modulus must be prime.
+
+        """
+        cdef fmpz_mod v
+
+        v = fmpz_mod.__new__(fmpz_mod)
+        v.ctx = self.ctx
+
+        if fmpz_is_zero(self.val):
+            return v
+
+        if not fmpz_sqrtmod(v.val, self.val, self.ctx.val.n):
+            raise DomainError("no square root exists for {} mod {}".format(self, self.ctx.modulus()))
+
+        return v

src/flint/types/fmpz_mod_mpoly.pyx

@@ -758,6 +758,19 @@ cdef class fmpz_mod_mpoly(flint_mpoly):
         return res
 
     def leading_coefficient(self):
+        """
+        Leading coefficient in the monomial ordering.
+
+            >>> from flint import Ordering
+            >>> ctx = fmpz_mod_mpoly_ctx(2, Ordering.lex, ['x', 'y'], 11)
+            >>> x, y = ctx.gens()
+            >>> p = 2*x*y + 3*x + 4*y**2 + 5
+            >>> p
+            2*x*y + 3*x + 4*y^2 + 5
+            >>> p.leading_coefficient()
+            2
+
+        """
         if fmpz_mod_mpoly_is_zero(self.val, self.ctx.val):
             return fmpz(0)
         else:

src/flint/types/fmpz_mpoly.pyx

@@ -837,6 +837,19 @@ cdef class fmpz_mpoly(flint_mpoly):
         return res
 
     def leading_coefficient(self):
+        """
+        Leading coefficient in the monomial ordering.
+
+            >>> from flint import Ordering
+            >>> ctx = fmpz_mpoly_ctx(2, Ordering.lex, ['x', 'y'])
+            >>> x, y = ctx.gens()
+            >>> p = 2*x*y + 3*x + 4*y**2 + 5
+            >>> p
+            2*x*y + 3*x + 4*y^2 + 5
+            >>> p.leading_coefficient()
+            2
+
+        """
         if fmpz_mpoly_is_zero(self.val, self.ctx.val):
             return fmpz(0)
         else:

src/flint/types/fq_default.pyx

@@ -7,6 +7,8 @@ from flint.types.fmpz_mod_poly cimport fmpz_mod_poly, fmpz_mod_poly_ctx
 from flint.types.nmod_poly cimport nmod_poly
 from flint.utils.typecheck cimport typecheck
 
+from flint.utils.flint_exceptions import DomainError
+
 # Allow the type to be denoted by strings or integers
 FQ_TYPES = {
     "FQ_ZECH" : 1,
@@ -750,7 +752,7 @@ cdef class fq_default(flint_scalar):
         check = fq_default_sqrt(res.val, self.val, self.ctx.val)
         if check:
             return res
-        raise ValueError("element is not a square")
+        raise DomainError("element is not a square")
 
     def is_square(self):
         """