Repository: sympy/sympy

Issue:
Exact square root for ground domains of polynomials
Following discussions in https://github.com/sympy/sympy/issues/25087, exact square root for a ground domain is needed for square root of polynomials over the same domain, and thus a generic exact square root method for the abstract [sympy.polys.domains.domain.Domain](https://docs.sympy.org/latest/modules/polys/domainsref.html?highlight=domain#sympy.polys.domains.domain.Domain) for generic square root of polynomials over any domain.

Currently there is a `sqrt` method, but it's only implemented for the domain of integers as `floor(sqrt(x))`, which is not "a square root of a ring element" in the abstract sense, and cannot be generalized to other domains. Thus we need a new square root method that has a well-defined meaning applicable to all domains.

Proposed solution:
- Add `is_square` and `exsqrt` (exact square root, or `None` if it doesn't exist) methods to the abstract `Domain` class.
- Implement `is_square` and `exsqrt` for Z, Q, R, C, Fp. They are the most commonly used domains, and already have `is_square` and `exsqrt` functions in other modules of SymPy or in dependencies like GMPY.
- Add a deprecation note to the documentation of the `sqrt` method.

Discussion:

Patch:
diff --git a/sympy/external/gmpy.py b/sympy/external/gmpy.py
index 6cfe19bf74..9fe1c52901 100644
--- a/sympy/external/gmpy.py
+++ b/sympy/external/gmpy.py
@@ -100,6 +100,8 @@
 
     factorial = gmpy.fac
     sqrt = gmpy.isqrt
+    is_square = gmpy.is_square
+    sqrtrem = gmpy.isqrt_rem
     gcd = gmpy.gcd
     lcm = gmpy.lcm
     invert = gmpy.invert
@@ -116,6 +118,8 @@
 
     factorial = lambda x: int(mlib.ifac(x))
     sqrt = lambda x: int(mlib.isqrt(x))
+    is_square = lambda x: x >= 0 and mlib.sqrtrem(x)[1] == 0
+    sqrtrem = lambda x: tuple(int(r) for r in mlib.sqrtrem(x))
     if sys.version_info[:2] >= (3, 9):
         gcd = math.gcd
         lcm = math.lcm
diff --git a/sympy/polys/domains/complexfield.py b/sympy/polys/domains/complexfield.py
index b2d40f72f1..a5fe5cca35 100644
--- a/sympy/polys/domains/complexfield.py
+++ b/sympy/polys/domains/complexfield.py
@@ -147,5 +147,19 @@ def almosteq(self, a, b, tolerance=None):
         """Check if ``a`` and ``b`` are almost equal. """
         return self._context.almosteq(a, b, tolerance)
 
+    def is_square(self, a):
+        """Returns ``True``. Every complex number has a complex square root."""
+        return True
+
+    def exsqrt(self, a):
+        r"""Returns the principal complex square root of ``a``.
+
+        Explanation
+        ===========
+        The argument of the principal square root is always within
+        $(-\frac{\pi}{2}, \frac{\pi}{2}]$. The square root may be
+        slightly inaccurate due to floating point rounding error.
+        """
+        return a ** 0.5
 
 CC = ComplexField()
diff --git a/sympy/polys/domains/domain.py b/sympy/polys/domains/domain.py
index 36ac450f7d..afcbef6cb6 100644
--- a/sympy/polys/domains/domain.py
+++ b/sympy/polys/domains/domain.py
@@ -1277,7 +1277,51 @@ def log(self, a, b):
         raise NotImplementedError
 
     def sqrt(self, a):
-        """Returns square root of ``a``. """
+        """Returns a (possibly inexact) square root of ``a``.
+
+        Explanation
+        ===========
+        There is no universal definition of "inexact square root" for all
+        domains. It is not recommended to implement this method for domains
+        other then :ref:`ZZ`.
+
+        See also
+        ========
+        exsqrt
+        """
+        raise NotImplementedError
+
+    def is_square(self, a):
+        """Returns whether ``a`` is a square in the domain.
+
+        Explanation
+        ===========
+        Returns ``True`` if there is an element ``b`` in the domain such that
+        ``b * b == a``, otherwise returns ``False``. For inexact domains like
+        :ref:`RR` and :ref:`CC`, a tiny difference in this equality can be
+        tolerated.
+
+        See also
+        ========
+        exsqrt
+        """
+        raise NotImplementedError
+
+    def exsqrt(self, a):
+        """Principal square root of a within the domain if ``a`` is square.
+
+        Explanation
+        ===========
+        The implementation of this method should return an element ``b`` in the
+        domain such that ``b * b == a``, or ``None`` if there is no such ``b``.
+        For inexact domains like :ref:`RR` and :ref:`CC`, a tiny difference in
+        this equality can be tolerated. The choice of a "principal" square root
+        should follow a consistent rule whenever possible.
+
+        See also
+        ========
+        sqrt, is_square
+        """
         raise NotImplementedError
 
     def evalf(self, a, prec=None, **options):
diff --git a/sympy/polys/domains/finitefield.py b/sympy/polys/domains/finitefield.py
index cd018978ea..ea4ec88143 100644
--- a/sympy/polys/domains/finitefield.py
+++ b/sympy/polys/domains/finitefield.py
@@ -5,6 +5,7 @@
 
 from sympy.polys.domains.modularinteger import ModularIntegerFactory
 from sympy.polys.domains.simpledomain import SimpleDomain
+from sympy.polys.galoistools import gf_zassenhaus, gf_irred_p_rabin
 from sympy.polys.polyerrors import CoercionFailed
 from sympy.utilities import public
 from sympy.polys.domains.groundtypes import SymPyInteger
@@ -202,5 +203,28 @@ def from_RealField(K1, a, K0):
         if q == 1:
             return K1.dtype(K1.dom.dtype(p))
 
+    def is_square(self, a):
+        """Returns True if ``a`` is a quadratic residue modulo p. """
+        # a is not a square <=> x**2-a is irreducible
+        poly = [x.val for x in [self.one, self.zero, -a]]
+        return not gf_irred_p_rabin(poly, self.mod, self.dom)
+
+    def exsqrt(self, a):
+        """Square root modulo p of ``a`` if it is a quadratic residue.
+
+        Explanation
+        ===========
+        Always returns the square root that is no larger than ``p // 2``.
+        """
+        # x**2-a is not square-free if a=0 or the field is characteristic 2
+        if self.mod == 2 or a == 0:
+            return a
+        # Otherwise, use square-free factorization routine to factorize x**2-a
+        poly = [x.val for x in [self.one, self.zero, -a]]
+        for factor in gf_zassenhaus(poly, self.mod, self.dom):
+            if len(factor) == 2 and factor[1] <= self.mod // 2:
+                return self.dtype(factor[1])
+        return None
+
 
 FF = GF = FiniteField
diff --git a/sympy/polys/domains/groundtypes.py b/sympy/polys/domains/groundtypes.py
index e9dea487de..a694e48273 100644
--- a/sympy/polys/domains/groundtypes.py
+++ b/sympy/polys/domains/groundtypes.py
@@ -1,7 +1,7 @@
 """Ground types for various mathematical domains in SymPy. """
 
 import builtins
-from sympy.external.gmpy import HAS_GMPY, factorial, sqrt
+from sympy.external.gmpy import HAS_GMPY, factorial, sqrt, is_square, sqrtrem
 
 PythonInteger = builtins.int
 PythonReal = builtins.float
@@ -67,7 +67,7 @@ def __init__(self, obj):
     'gmpy_denom', 'gmpy_gcdex', 'gmpy_gcd', 'gmpy_lcm',
     'gmpy_qdiv',
 
-    'factorial', 'sqrt',
+    'factorial', 'sqrt', 'is_square', 'sqrtrem',
 
     'GMPYInteger', 'GMPYRational',
 ]
diff --git a/sympy/polys/domains/integerring.py b/sympy/polys/domains/integerring.py
index 8145f16d9d..77345f18f6 100644
--- a/sympy/polys/domains/integerring.py
+++ b/sympy/polys/domains/integerring.py
@@ -5,7 +5,7 @@
 from sympy.polys.domains.groundtypes import (
     SymPyInteger,
     factorial,
-    gcdex, gcd, lcm, sqrt,
+    gcdex, gcd, lcm, sqrt, is_square, sqrtrem,
 )
 
 from sympy.polys.domains.characteristiczero import CharacteristicZero
@@ -223,6 +223,30 @@ def sqrt(self, a):
         """Compute square root of ``a``. """
         return sqrt(a)
 
+    def is_square(self, a):
+        """Return ``True`` if ``a`` is a square.
+
+        Explanation
+        ===========
+        An integer is a square if and only if there exists an integer
+        ``b`` such that ``b * b == a``.
+        """
+        return is_square(a)
+
+    def exsqrt(self, a):
+        """Non-negative square root of ``a`` if ``a`` is a square.
+
+        See also
+        ========
+        is_square
+        """
+        if a < 0:
+            return None
+        root, rem = sqrtrem(a)
+        if rem != 0:
+            return None
+        return root
+
     def factorial(self, a):
         """Compute factorial of ``a``. """
         return factorial(a)
diff --git a/sympy/polys/domains/rationalfield.py b/sympy/polys/domains/rationalfield.py
index 1179cd7e4e..a7762c510d 100644
--- a/sympy/polys/domains/rationalfield.py
+++ b/sympy/polys/domains/rationalfield.py
@@ -3,7 +3,7 @@
 
 from sympy.external.gmpy import MPQ
 
-from sympy.polys.domains.groundtypes import SymPyRational
+from sympy.polys.domains.groundtypes import SymPyRational, is_square, sqrtrem
 
 from sympy.polys.domains.characteristiczero import CharacteristicZero
 from sympy.polys.domains.field import Field
@@ -159,5 +159,31 @@ def denom(self, a):
         """Returns denominator of ``a``. """
         return a.denominator
 
+    def is_square(self, a):
+        """Return ``True`` if ``a`` is a square.
+
+        Explanation
+        ===========
+        A rational number is a square if and only if there exists
+        a rational number ``b`` such that ``b * b == a``.
+        """
+        return is_square(a.numerator) and is_square(a.denominator)
+
+    def exsqrt(self, a):
+        """Non-negative square root of ``a`` if ``a`` is a square.
+
+        See also
+        ========
+        is_square
+        """
+        if a.numerator < 0:  # denominator is always positive
+            return None
+        p_sqrt, p_rem = sqrtrem(a.numerator)
+        if p_rem != 0:
+            return None
+        q_sqrt, q_rem = sqrtrem(a.denominator)
+        if q_rem != 0:
+            return None
+        return MPQ(p_sqrt, q_sqrt)
 
 QQ = RationalField()
diff --git a/sympy/polys/domains/realfield.py b/sympy/polys/domains/realfield.py
index 4ae42359d2..c828648a23 100644
--- a/sympy/polys/domains/realfield.py
+++ b/sympy/polys/domains/realfield.py
@@ -128,5 +128,19 @@ def almosteq(self, a, b, tolerance=None):
         """Check if ``a`` and ``b`` are almost equal. """
         return self._context.almosteq(a, b, tolerance)
 
+    def is_square(self, a):
+        """Returns ``True`` if ``a >= 0`` and ``False`` otherwise. """
+        return a >= 0
+
+    def exsqrt(self, a):
+        """Non-negative square root for ``a >= 0`` and ``None`` otherwise.
+
+        Explanation
+        ===========
+        The square root may be slightly inaccurate due to floating point
+        rounding error.
+        """
+        return a ** 0.5 if a >= 0 else None
+
 
 RR = RealField()


Test patch:
diff --git a/sympy/polys/domains/tests/test_domains.py b/sympy/polys/domains/tests/test_domains.py
index 9fc621cb51..0a0f3cb876 100644
--- a/sympy/polys/domains/tests/test_domains.py
+++ b/sympy/polys/domains/tests/test_domains.py
@@ -1273,3 +1273,65 @@ def test_AlgebraicField_alias():
     k = QQ.algebraic_field(alpha)
     beta = k.to_alg_num(k([1, 2, 3]))
     assert beta.alias is alpha.alias
+
+
+def test_exsqrt():
+    assert ZZ.is_square(ZZ(4)) is True
+    assert ZZ.exsqrt(ZZ(4)) == ZZ(2)
+    assert ZZ.is_square(ZZ(42)) is False
+    assert ZZ.exsqrt(ZZ(42)) is None
+    assert ZZ.is_square(ZZ(0)) is True
+    assert ZZ.exsqrt(ZZ(0)) == ZZ(0)
+    assert ZZ.is_square(ZZ(-1)) is False
+    assert ZZ.exsqrt(ZZ(-1)) is None
+
+    assert QQ.is_square(QQ(9, 4)) is True
+    assert QQ.exsqrt(QQ(9, 4)) == QQ(3, 2)
+    assert QQ.is_square(QQ(18, 8)) is True
+    assert QQ.exsqrt(QQ(18, 8)) == QQ(3, 2)
+    assert QQ.is_square(QQ(-9, -4)) is True
+    assert QQ.exsqrt(QQ(-9, -4)) == QQ(3, 2)
+    assert QQ.is_square(QQ(11, 4)) is False
+    assert QQ.exsqrt(QQ(11, 4)) is None
+    assert QQ.is_square(QQ(9, 5)) is False
+    assert QQ.exsqrt(QQ(9, 5)) is None
+    assert QQ.is_square(QQ(4)) is True
+    assert QQ.exsqrt(QQ(4)) == QQ(2)
+    assert QQ.is_square(QQ(0)) is True
+    assert QQ.exsqrt(QQ(0)) == QQ(0)
+    assert QQ.is_square(QQ(-16, 9)) is False
+    assert QQ.exsqrt(QQ(-16, 9)) is None
+
+    assert RR.is_square(RR(6.25)) is True
+    assert RR.exsqrt(RR(6.25)) == RR(2.5)
+    assert RR.is_square(RR(2)) is True
+    assert RR.almosteq(RR.exsqrt(RR(2)), RR(1.4142135623730951), tolerance=1e-15)
+    assert RR.is_square(RR(0)) is True
+    assert RR.exsqrt(RR(0)) == RR(0)
+    assert RR.is_square(RR(-1)) is False
+    assert RR.exsqrt(RR(-1)) is None
+
+    assert CC.is_square(CC(2)) is True
+    assert CC.almosteq(CC.exsqrt(CC(2)), CC(1.4142135623730951), tolerance=1e-15)
+    assert CC.is_square(CC(0)) is True
+    assert CC.exsqrt(CC(0)) == CC(0)
+    assert CC.is_square(CC(-1)) is True
+    assert CC.exsqrt(CC(-1)) == CC(0, 1)
+    assert CC.is_square(CC(0, 2)) is True
+    assert CC.exsqrt(CC(0, 2)) == CC(1, 1)
+    assert CC.is_square(CC(-3, -4)) is True
+    assert CC.exsqrt(CC(-3, -4)) == CC(1, -2)
+
+    F2 = FF(2)
+    assert F2.is_square(F2(1)) is True
+    assert F2.exsqrt(F2(1)) == F2(1)
+    assert F2.is_square(F2(0)) is True
+    assert F2.exsqrt(F2(0)) == F2(0)
+
+    F7 = FF(7)
+    assert F7.is_square(F7(2)) is True
+    assert F7.exsqrt(F7(2)) == F7(3)
+    assert F7.is_square(F7(3)) is False
+    assert F7.exsqrt(F7(3)) is None
+    assert F7.is_square(F7(0)) is True
+    assert F7.exsqrt(F7(0)) == F7(0)