Add factor_squarefree to all poly/mpoly types

Author: oscarbenjamin

src/flint/test/test_all.py

@@ -3320,16 +3320,52 @@ def _all_polys_mpolys():
 def test_factor_poly_mpoly():
     """Test that factor() is consistent across different poly/mpoly types."""
 
-    def factor(p):
-        coeff, factors = p.factor()
+    def check(p, coeff, factors):
+        # Check all the types
         lc = p.leading_coefficient()
         assert type(coeff) is type(lc)
         assert isinstance(factors, list)
         assert all(isinstance(f, tuple) for f in factors)
         for fac, m in factors:
             assert type(fac) is type(p)
             assert type(m) is int
-        return coeff, sorted(factors, key=lambda p: (p[1], str(p[0])))
+
+        # Check the actual factorisation!
+        res = coeff
+        for fac, m in factors:
+            res *= fac ** m
+        assert res == p
+
+    def sort(factors):
+        def sort_key(p):
+            fac, m = p
+            return (m, sorted(str(i) for i in fac.coeffs()))
+        return sorted(factors, key=sort_key)
+
+    def factor(p):
+        coeff, factors = p.factor()
+        check(p, coeff, factors)
+        return coeff, sort(factors)
+
+    def factor_sqf(p):
+        coeff, factors = p.factor_squarefree()
+        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():
 
@@ -3351,9 +3387,40 @@ def factor(p):
         assert factor(x) == (S(1), [(x, 1)])
         assert factor(-x) == (S(-1), [(x, 1)])
         assert factor(x**2) == (S(1), [(x, 2)])
-        assert factor(x*(x+1)) == (S(1), [(x, 1), (x+1, 1)])
         assert factor(2*(x+1)) == (S(2), [(x+1, 1)])
 
+        assert factor_sqf(0*x) == (S(0), [])
+        assert factor_sqf(0*x + 1) == (S(1), [])
+        assert factor_sqf(0*x + 3) == (S(3), [])
+        assert factor_sqf(-x) == (S(-1), [(x, 1)])
+        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)])
+        if y is None:
+            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.
+            assert factor_sqf(x*(x+1)) == (S(1), [(x, 1), (x+1, 1)])
+
+        assert factor_sqf(x**2*(x+1)) == (S(1), [(x+1, 1), (x, 2)])
+
+        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)])
+
+        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:
+            assert factor(p) == (S(3), sort([(x-1, 2), (x+1, 2), (x**2+1, 3)]))
+
         if characteristic == 0:
             # primitive factors over Z for Z and Q.
             assert factor(2*x+1) == (S(1), [(2*x+1, 1)])
@@ -3372,16 +3439,20 @@ def factor(p):
             assert factor(x*y+1) == (S(1), [(x*y+1, 1)])
             assert factor(x*y) == (S(1), [(x, 1), (y, 1)])
 
+            assert factor_sqf((x*y+1)**2*(x*y-1)) == (S(1), [(x*y-1, 1), (x*y+1, 2)])
+
+            p = 2*x + y
             if characteristic == 0:
-                assert factor(2*x + y) == (S(1), [(2*x + y, 1)])
+                assert factor(p) == factor_sqf(p) == (S(1), [(p, 1)])
             else:
-                assert factor(2*x + y) == (S(2), [(x + y/2, 1)])
+                assert factor(p) == factor_sqf(p) == (S(2), [(p/2, 1)])
 
             if is_field:
+                p = (2*x + y)/7
                 if characteristic == 0:
-                    assert factor((2*x+y)/7) == (S(1)/7, [(2*x+y, 1)])
+                    assert factor(p) == factor_sqf(p) == (S(1)/7, [(7*p, 1)])
                 else:
-                    assert factor((2*x+y)/7) == (S(2)/7, [(x+y/2, 1)])
+                    assert factor(p) == factor_sqf(p) == (S(2)/7, [(7*p/2, 1)])
 
 
 def _all_matrices():

src/flint/types/fmpq_mpoly.pyx

@@ -985,7 +985,7 @@ cdef class fmpq_mpoly(flint_mpoly):
             c = fmpz.__new__(fmpz)
             fmpz_init_set((<fmpz>c).val, &fac.exp[i])
 
-            res[i] = (u, c)
+            res[i] = (u, int(c))
 
         c = fmpq.__new__(fmpq)
         fmpq_set((<fmpq>c).val, fac.constant)

src/flint/types/fmpq_poly.pyx

@@ -448,6 +448,27 @@ cdef class fmpq_poly(flint_poly):
 
         return c / self.denom(), fac
 
+    def factor_squarefree(self):
+        """
+        Factors *self* into square-free polynomials. Returns (*c*, *factors*)
+        where *c* is the leading coefficient and *factors* is a list of
+        (*poly*, *exp*).
+
+            >>> x = fmpq_poly([0, 1])
+            >>> p = x**2 * (x/2 - 1)**2 * (x + 1)**3
+            >>> p
+            1/4*x^7 + (-1/4)*x^6 + (-5/4)*x^5 + 1/4*x^4 + 2*x^3 + x^2
+            >>> p.factor_squarefree()
+            (1/4, [(x**2 - 2*x, 2), (x + 1, 3)])
+            >>> p.factor()
+            (1/4, [(x, 2), (x + (-2), 2), (x + 1, 3)])
+
+        """
+        c, fac = self.numer().factor_squarefree()
+        c = fmpq(c) / self.denom()
+        fac = [(fmpq_poly(f), m) for f, m in fac]
+        return c, fac
+
     def sqrt(self):
         """
         Return the exact square root of this polynomial or ``None``.

src/flint/types/fmpz.pyx

@@ -900,6 +900,7 @@ cdef class fmpz(flint_scalar):
         return u, v
 
     # warning: m should be prime!
+    # also if self is zero this crashes...
     def sqrtmod(self, m):
         cdef fmpz v
         v = fmpz()

src/flint/types/fmpz_mod_mpoly.pyx

@@ -896,7 +896,7 @@ cdef class fmpz_mod_mpoly(flint_mpoly):
             c = fmpz.__new__(fmpz)
             fmpz_set((<fmpz>c).val, &fac.exp[i])
 
-            res[i] = (u, c)
+            res[i] = (u, int(c))
 
         c = fmpz.__new__(fmpz)
         fmpz_set((<fmpz>c).val, fac.constant)

src/flint/types/fmpz_mpoly.pyx

@@ -934,7 +934,7 @@ cdef class fmpz_mpoly(flint_mpoly):
 
     def factor_squarefree(self):
         """
-        Factors self into irreducible factors, returning a tuple
+        Factors self into square-free factors, returning a tuple
         (c, factors) where c is the content of the coefficients and
         factors is a list of (poly, exp) pairs.
 
@@ -965,7 +965,7 @@ cdef class fmpz_mpoly(flint_mpoly):
             c = fmpz.__new__(fmpz)
             fmpz_set((<fmpz>c).val, &fac.exp[i])
 
-            res[i] = (u, c)
+            res[i] = (u, int(c))
 
         c = fmpz.__new__(fmpz)
         fmpz_set((<fmpz>c).val, fac.constant)

src/flint/types/fmpz_poly.pyx

@@ -384,11 +384,36 @@ cdef class fmpz_poly(flint_poly):
             (1, [(6*x^5 + 5*x^4 + 4*x^3 + 3*x^2 + 2*x + 1, 1)])
 
         """
+        return self._factor('irreducible')
+
+    def factor_squarefree(self):
+        """
+        Factors self into square-free factors, returning a tuple
+        (c, factors) where c is the content of the coefficients and
+        factors is a list of (poly, exp) pairs.
+
+            >>> x = fmpz_poly([0, 1])
+            >>> p = (-3 * x**2 * (x + 1)**2 * (x - 1)**3)
+            >>> p.factor_squarefree()
+            (-3, [(x**2 + x, 2), (x - 1, 3)])
+            >>> p.factor()
+            (-3, [(x, 2), (x + 1, 2), (x - 1, 3)])
+
+        """
+        return self._factor('squarefree')
+
+    def _factor(self, factor_type):
         cdef fmpz_poly_factor_t fac
         cdef int i
         fmpz_poly_factor_init(fac)
-        # should use fmpz_poly_factor, but not available in 2.5.2
-        fmpz_poly_factor(fac, self.val)
+
+        if factor_type == 'squarefree':
+            fmpz_poly_factor_squarefree(fac, self.val)
+        elif factor_type == 'irreducible':
+            fmpz_poly_factor(fac, self.val)
+        else:
+            assert False
+
         res = [0] * fac.num
         for 0 <= i < fac.num:
             u = fmpz_poly.__new__(fmpz_poly)
@@ -398,6 +423,7 @@ cdef class fmpz_poly(flint_poly):
         c = fmpz.__new__(fmpz)
         fmpz_set((<fmpz>c).val, &fac.c)
         fmpz_poly_factor_clear(fac)
+
         return c, res
 
     def complex_roots(self, bint verbose=False):

src/flint/types/nmod_mpoly.pyx

@@ -871,9 +871,9 @@ cdef class nmod_mpoly(flint_mpoly):
             c = fmpz.__new__(fmpz)
             fmpz_set((<fmpz>c).val, &fac.exp[i])
 
-            res[i] = (u, c)
+            res[i] = (u, int(c))
 
-        constant = fac.constant
+        constant = nmod(fac.constant, self.ctx.modulus())
         nmod_mpoly_factor_clear(fac, self.ctx.val)
         return constant, res
 

src/flint/types/nmod_poly.pyx

@@ -632,16 +632,49 @@ cdef class nmod_poly(flint_poly):
             (3, [(x + 4, 1), (x + 2, 1), (x^2 + 4*x + 1, 1)])
 
         """
+        if algorithm is None:
+            algorithm = 'irreducible'
+        elif algorithm not in ('berlekamp', 'cantor-zassenhaus'):
+            raise ValueError(f"unknown factorization algorithm: {algorithm}")
+        return self._factor(algorithm)
+
+    def factor_squarefree(self):
+        """
+        Factors *self* into square-free polynomials. Returns (*c*, *factors*)
+        where *c* is the leading coefficient and *factors* is a list of
+        (*poly*, *exp*).
+
+            >>> x = nmod_poly([0, 1], 7)
+            >>> p = x**2 * (x/2 - 1)**2 * (x + 1)**3
+            >>> p
+            2*x^7 + 5*x^6 + 4*x^5 + 2*x^4 + 2*x^3 + x^2
+            >>> p.factor_squarefree()
+            (2, [(x**2 - 2*x, 2), (x + 1, 3)])
+            >>> p.factor()
+            (1/4, [(x, 2), (x + 5, 2), (x + 1, 3)])
+
+        """
+        return self._factor('squarefree')
+
+    def _factor(self, factor_type):
         cdef nmod_poly_factor_t fac
         cdef mp_limb_t lead
         cdef int i
+
         nmod_poly_factor_init(fac)
-        if algorithm == 'berlekamp':
+
+        if factor_type == 'berlekamp':
             lead = nmod_poly_factor_with_berlekamp(fac, self.val)
-        elif algorithm == 'cantor-zassenhaus':
+        elif factor_type == 'cantor-zassenhaus':
             lead = nmod_poly_factor_with_cantor_zassenhaus(fac, self.val)
-        else:
+        elif factor_type == 'irreducible':
             lead = nmod_poly_factor(fac, self.val)
+        elif factor_type == 'squarefree':
+            nmod_poly_factor_squarefree(fac, self.val)
+            lead = (<nmod>self.leading_coefficient()).val
+        else:
+            assert False
+
         res = [None] * fac.num
         for 0 <= i < fac.num:
             u = nmod_poly.__new__(nmod_poly)
@@ -650,10 +683,13 @@ cdef class nmod_poly(flint_poly):
             nmod_poly_set((<nmod_poly>u).val, &fac.p[i])
             exp = fac.exp[i]
             res[i] = (u, exp)
+
         c = nmod.__new__(nmod)
         (<nmod>c).mod = self.val.mod
         (<nmod>c).val = lead
+
         nmod_poly_factor_clear(fac)
+
         return c, res
 
     def sqrt(nmod_poly self):