Merge pull request #322 from remyoudompheng/mullow

Extend mul_low and pow_trunc methods to fmpz_poly, fmpq_poly, nmod_poly

Author: oscarbenjamin

README.md

@@ -170,6 +170,9 @@ Contributors (0.9.0):
 
 Changes (0.9.0):
 
+- [gh-322](https://github.com/flintlib/python-flint/pull/322),
+  Add `mul_low` and `pow_trunc` methods to `fmpz_poly`, `fmpq_poly` and
+  `nmod_poly`. (RO)
 - [gh-318](https://github.com/flintlib/python-flint/pull/318),
   Add emscripten build in CI. Polynomial factors and roots are
   now sorted into a consistent order for `nmod_poly` and

src/flint/test/test_all.py

@@ -2964,6 +2964,20 @@ def setbad(obj, i, val):
     assert raises(lambda: P([1, 1]) ** -1, DomainError)
     assert raises(lambda: P([1, 1]) ** None, TypeError) # type: ignore
 
+    # Truncated operations
+    assert P([1, 2, 3]).mul_low(P([4, 5, 6]), 3) == P([4, 13, 28])
+    assert raises(lambda: P([1, 2, 3]).mul_low(None, 3), TypeError) # type: ignore
+    assert raises(lambda: P([1, 2, 3]).mul_low(P([4, 5, 6]), None), TypeError) # type: ignore
+
+    p = P([1, 2, 3])
+    assert p.pow_trunc(1234, 3) == P([1, 2468, 3046746])
+    assert raises(lambda: p.pow_trunc(None, 3), TypeError) # type: ignore
+    assert raises(lambda: p.pow_trunc(3, "A"), TypeError) # type: ignore
+    assert raises(lambda: p.pow_trunc(P([4, 5, 6]), 3), TypeError) # type: ignore
+    # Large exponents are allowed
+    assert p.pow_trunc(2**100, 2) == P([1, 2**101])
+    assert p.pow_trunc(6**60, 3) == p.pow_trunc(2**60, 3).pow_trunc(3**60, 3)
+
     # XXX: Not sure what this should do in general:
     p = P([1, 1])
     mod = P([1, 1])

src/flint/types/fmpq_poly.pyi

@@ -71,6 +71,8 @@ class fmpq_poly(flint_poly[fmpq]):
     def left_shift(self, n: int, /) -> fmpq_poly: ...
     def right_shift(self, n: int, /) -> fmpq_poly: ...
     def truncate(self, n: int, /) -> fmpq_poly: ...
+    def mul_low(self, other: fmpq_poly, n: int) -> fmpq_poly: ...
+    def pow_trunc(self, e: int, n: int) -> fmpq_poly: ...
 
     def gcd(self, other: ifmpq_poly, /) -> fmpq_poly: ...
     def discriminant(self) -> fmpq: ...

src/flint/types/fmpq_poly.pyx

@@ -502,6 +502,66 @@ cdef class fmpq_poly(flint_poly):
         fmpq_poly_pow(res.val, self.val, <ulong>exp)
         return res
 
+    def mul_low(self, other, slong n):
+        r"""
+        Returns the lowest ``n`` coefficients of the multiplication of ``self`` with ``other``
+
+        Equivalent to computing `f(x) \cdot g(x) \mod x^n`
+
+            >>> f = fmpq_poly([2,3,5,7,11])
+            >>> g = fmpq_poly([1,2,4,8,16])
+            >>> f.mul_low(g, 5)
+            101*x^4 + 45*x^3 + 19*x^2 + 7*x + 2
+            >>> f.mul_low(g, 3)
+            19*x^2 + 7*x + 2
+            >>> f.mul_low(g, 1)
+            2
+        """
+        # Only allow multiplication with other fmpq_poly
+        if not typecheck(other, fmpq_poly):
+            raise TypeError("other polynomial must be of type fmpq_poly")
+
+        cdef fmpq_poly res
+        res = fmpq_poly.__new__(fmpq_poly)
+        fmpq_poly_mullow(res.val, self.val, (<fmpq_poly>other).val, n)
+        return res
+
+    def pow_trunc(self, e, slong n):
+        r"""
+        Returns ``self`` raised to the power ``e`` modulo `x^n`:
+        :math:`f^e \mod x^n`/
+
+            >>> f = fmpq_poly([1, 2, 3])
+            >>> x = fmpq_poly([0, 1])
+            >>> f.pow_trunc(2**20, 4)
+            1537230871828889600*x^3 + 2199024304128*x^2 + 2097152*x + 1
+            >>> f.pow_trunc(5**25, 3)
+            177635683940025046765804290771484375*x^2 + 596046447753906250*x + 1
+        """
+        if e < 0:
+            raise ValueError("Exponent must be non-negative")
+
+        cdef slong e_c
+        cdef fmpq_poly res, tmp
+
+        try:
+            e_c = e
+        except OverflowError:
+            # Exponent does not fit slong
+            res = fmpq_poly.__new__(fmpq_poly)
+            tmp = fmpq_poly.__new__(fmpq_poly)
+            ebytes = e.to_bytes((e.bit_length() + 15) // 16 * 2, "big")
+            fmpq_poly_pow_trunc(res.val, self.val, ebytes[0] * 256 + ebytes[1], n)
+            for i in range(2, len(ebytes), 2):
+                fmpq_poly_pow_trunc(res.val, res.val, 1 << 16, n)
+                fmpq_poly_pow_trunc(tmp.val, self.val, ebytes[i] * 256 + ebytes[i+1], n)
+                fmpq_poly_mullow(res.val, res.val, tmp.val, n)
+            return res
+
+        res = fmpq_poly.__new__(fmpq_poly)
+        fmpq_poly_pow_trunc(res.val, self.val, e_c, n)
+        return res
+
     def gcd(self, other):
         """
         Returns the greatest common divisor of *self* and *other*.

src/flint/types/fmpz_mod_poly.pyx

@@ -1668,28 +1668,43 @@ cdef class fmpz_mod_poly(flint_poly):
         )
         return res
 
-    def pow_trunc(self, slong e, slong n):
+    def pow_trunc(self, e, slong n):
         r"""
         Returns ``self`` raised to the power ``e`` modulo `x^n`:
         :math:`f^e \mod x^n`/
 
-        Note: For exponents larger that 2^31 (which do not fit inside a ulong) use the
-        method :meth:`~.pow_mod` with the explicit modulus `x^n`.
-
             >>> R = fmpz_mod_poly_ctx(163)
             >>> x = R.gen()
             >>> f = 30*x**6 + 104*x**5 + 76*x**4 + 33*x**3 + 70*x**2 + 44*x + 65
             >>> f.pow_trunc(2**20, 30) == pow(f, 2**20, x**30)
             True
             >>> f.pow_trunc(2**20, 5)
             132*x^4 + 113*x^3 + 36*x^2 + 48*x + 6
+            >>> f.pow_trunc(5**25, 5)
+            147*x^4 + 98*x^3 + 95*x^2 + 33*x + 126
         """
         if e < 0:
             raise ValueError("Exponent must be non-negative")
 
-        cdef fmpz_mod_poly res
+        cdef fmpz_mod_poly res, tmp
+        cdef slong e_c
+
+        try:
+            e_c = e
+        except OverflowError:
+            # Exponent does not fit slong
+            res = self.ctx.new_ctype_poly()
+            tmp = self.ctx.new_ctype_poly()
+            ebytes = e.to_bytes((e.bit_length() + 15) // 16 * 2, "big")
+            fmpz_mod_poly_pow_trunc(res.val, self.val, ebytes[0] * 256 + ebytes[1], n, res.ctx.mod.val)
+            for i in range(2, len(ebytes), 2):
+                fmpz_mod_poly_pow_trunc(res.val, res.val, 1 << 16, n, res.ctx.mod.val)
+                fmpz_mod_poly_pow_trunc(tmp.val, self.val, ebytes[i] * 256 + ebytes[i+1], n, res.ctx.mod.val)
+                fmpz_mod_poly_mullow(res.val, res.val, tmp.val, n, res.ctx.mod.val)
+            return res
+
         res = self.ctx.new_ctype_poly()
-        fmpz_mod_poly_pow_trunc(res.val, self.val, e, n, res.ctx.mod.val)
+        fmpz_mod_poly_pow_trunc(res.val, self.val, e_c, n, res.ctx.mod.val)
         return res
 
     def inflate(self, ulong n):

src/flint/types/fmpz_poly.pyi

@@ -59,6 +59,8 @@ class fmpz_poly(flint_poly[fmpz]):
     def left_shift(self, n: int, /) -> fmpz_poly: ...
     def right_shift(self, n: int, /) -> fmpz_poly: ...
     def truncate(self, n: int, /) -> fmpz_poly: ...
+    def mul_low(self, other: fmpz_poly, n: int) -> fmpz_poly: ...
+    def pow_trunc(self, e: int, n: int) -> fmpz_poly: ...
 
     def gcd(self, other: ifmpz_poly, /) -> fmpz_poly: ...
     def content(self) -> fmpz: ...

src/flint/types/fmpz_poly.pyx

@@ -483,6 +483,66 @@ cdef class fmpz_poly(flint_poly):
         fmpz_poly_pow(res.val, self.val, <ulong>exp)
         return res
 
+    def mul_low(self, other, slong n):
+        r"""
+        Returns the lowest ``n`` coefficients of the multiplication of ``self`` with ``other``
+
+        Equivalent to computing `f(x) \cdot g(x) \mod x^n`
+
+            >>> f = fmpz_poly([2,3,5,7,11])
+            >>> g = fmpz_poly([1,2,4,8,16])
+            >>> f.mul_low(g, 5)
+            101*x^4 + 45*x^3 + 19*x^2 + 7*x + 2
+            >>> f.mul_low(g, 3)
+            19*x^2 + 7*x + 2
+            >>> f.mul_low(g, 1)
+            2
+        """
+        # Only allow multiplication with other fmpz_poly
+        if not typecheck(other, fmpz_poly):
+            raise TypeError("other polynomial must be of type fmpz_poly")
+
+        cdef fmpz_poly res
+        res = fmpz_poly.__new__(fmpz_poly)
+        fmpz_poly_mullow(res.val, self.val, (<fmpz_poly>other).val, n)
+        return res
+
+    def pow_trunc(self, e, slong n):
+        r"""
+        Returns ``self`` raised to the power ``e`` modulo `x^n`:
+        :math:`f^e \mod x^n`/
+
+            >>> f = fmpz_poly([1, 2, 3])
+            >>> x = fmpz_poly([0, 1])
+            >>> f.pow_trunc(2**20, 4)
+            1537230871828889600*x^3 + 2199024304128*x^2 + 2097152*x + 1
+            >>> f.pow_trunc(5**25, 3)
+            177635683940025046765804290771484375*x^2 + 596046447753906250*x + 1
+        """
+        if e < 0:
+            raise ValueError("Exponent must be non-negative")
+
+        cdef slong e_c
+        cdef fmpz_poly res, tmp
+
+        try:
+            e_c = e
+        except OverflowError:
+            # Exponent does not fit slong
+            res = fmpz_poly.__new__(fmpz_poly)
+            tmp = fmpz_poly.__new__(fmpz_poly)
+            ebytes = e.to_bytes((e.bit_length() + 15) // 16 * 2, "big")
+            fmpz_poly_pow_trunc(res.val, self.val, ebytes[0] * 256 + ebytes[1], n)
+            for i in range(2, len(ebytes), 2):
+                fmpz_poly_pow_trunc(res.val, res.val, 1 << 16, n)
+                fmpz_poly_pow_trunc(tmp.val, self.val, ebytes[i] * 256 + ebytes[i+1], n)
+                fmpz_poly_mullow(res.val, res.val, tmp.val, n)
+            return res
+
+        res = fmpz_poly.__new__(fmpz_poly)
+        fmpz_poly_pow_trunc(res.val, self.val, e_c, n)
+        return res
+
     def gcd(self, other):
         """
         Returns the greatest common divisor of self and other.

src/flint/types/fq_default_poly.pyx

@@ -1190,28 +1190,43 @@ cdef class fq_default_poly(flint_poly):
         )
         return res
 
-    def pow_trunc(self, slong e, slong n):
+    def pow_trunc(self, e, slong n):
         r"""
         Returns ``self`` raised to the power ``e`` modulo `x^n`:
         :math:`f^e \mod x^n`/
 
-        Note: For exponents larger that 2^31 (which do not fit inside a ulong) use the
-        method :meth:`~.pow_mod` with the explicit modulus `x^n`.
-
             >>> R = fq_default_poly_ctx(163)
             >>> x = R.gen()
             >>> f = 30*x**6 + 104*x**5 + 76*x**4 + 33*x**3 + 70*x**2 + 44*x + 65
             >>> f.pow_trunc(2**20, 30) == pow(f, 2**20, x**30)
             True
             >>> f.pow_trunc(2**20, 5)
             132*x^4 + 113*x^3 + 36*x^2 + 48*x + 6
+            >>> f.pow_trunc(5**25, 5)
+            147*x^4 + 98*x^3 + 95*x^2 + 33*x + 126
         """
         if e < 0:
             raise ValueError("Exponent must be non-negative")
 
-        cdef fq_default_poly res
+        cdef slong e_c
+        cdef fq_default_poly res, tmp
+
+        try:
+            e_c = e
+        except OverflowError:
+            # Exponent does not fit slong
+            res = self.ctx.new_ctype_poly()
+            tmp = self.ctx.new_ctype_poly()
+            ebytes = e.to_bytes((e.bit_length() + 15) // 16 * 2, "big")
+            fq_default_poly_pow_trunc(res.val, self.val, ebytes[0] * 256 + ebytes[1], n, res.ctx.field.val)
+            for i in range(2, len(ebytes), 2):
+                fq_default_poly_pow_trunc(res.val, res.val, 1 << 16, n, res.ctx.field.val)
+                fq_default_poly_pow_trunc(tmp.val, self.val, ebytes[i] * 256 + ebytes[i+1], n, res.ctx.field.val)
+                fq_default_poly_mullow(res.val, res.val, tmp.val, n, res.ctx.field.val)
+            return res
+
         res = self.ctx.new_ctype_poly()
-        fq_default_poly_pow_trunc(res.val, self.val, e, n, res.ctx.field.val)
+        fq_default_poly_pow_trunc(res.val, self.val, e_c, n, res.ctx.field.val)
         return res
 
     def sqrt_trunc(self, slong n):

src/flint/types/nmod_poly.pyi

@@ -56,6 +56,8 @@ class nmod_poly(flint_poly[nmod]):
     def __rdivmod__(self, other: inmod_poly) -> tuple[nmod_poly, nmod_poly]: ...
     def left_shift(self, n: int) -> nmod_poly: ...
     def right_shift(self, n: int) -> nmod_poly: ...
+    def mul_low(self, other: nmod_poly, n: int) -> nmod_poly: ...
+    def pow_trunc(self, e: int, n: int) -> nmod_poly: ...
     def __pow__(self, other: int, mod: inmod_poly | None = None) -> nmod_poly: ...
     def pow_mod(
         self, e: int, modulus: inmod_poly, mod_rev_inv: inmod_poly | None = None

src/flint/types/nmod_poly.pyx

@@ -724,6 +724,75 @@ cdef class nmod_poly(flint_poly):
             )
         return res
 
+    def mul_low(self, other, slong n):
+        r"""
+        Returns the lowest ``n`` coefficients of the multiplication of ``self`` with ``other``
+
+        Equivalent to computing `f(x) \cdot g(x) \mod x^n`
+
+            >>> f = nmod_poly([2,3,5,7,11], 163)
+            >>> g = nmod_poly([1,2,4,8,16], 163)
+            >>> f.mul_low(g, 5)
+            101*x^4 + 45*x^3 + 19*x^2 + 7*x + 2
+            >>> f.mul_low(g, 3)
+            19*x^2 + 7*x + 2
+            >>> f.mul_low(g, 1)
+            2
+        """
+        # Only allow multiplication with other nmod_poly
+        if not typecheck(other, nmod_poly):
+            raise TypeError("other polynomial must be of type nmod_poly")
+
+        if (<nmod_poly>self).val.mod.n != (<nmod_poly>other).val.mod.n:
+            raise ValueError("cannot multiply nmod_polys with different moduli")
+
+        cdef nmod_poly res = nmod_poly.__new__(nmod_poly)
+        res = nmod_poly.__new__(nmod_poly)
+        nmod_poly_init_preinv(res.val, self.val.mod.n, self.val.mod.ninv)
+        nmod_poly_mullow(res.val, self.val, (<nmod_poly>other).val, n)
+        return res
+
+    def pow_trunc(self, e, slong n):
+        r"""
+        Returns ``self`` raised to the power ``e`` modulo `x^n`:
+        :math:`f^e \mod x^n`/
+
+            >>> f = nmod_poly([65, 44, 70, 33, 76, 104, 30], 163)
+            >>> x = nmod_poly([0, 1], 163)
+            >>> f.pow_trunc(2**20, 30) == pow(f, 2**20, x**30)
+            True
+            >>> f.pow_trunc(2**20, 5)
+            132*x^4 + 113*x^3 + 36*x^2 + 48*x + 6
+            >>> f.pow_trunc(5**25, 5)
+            147*x^4 + 98*x^3 + 95*x^2 + 33*x + 126
+        """
+        if e < 0:
+            raise ValueError("Exponent must be non-negative")
+
+        cdef nmod_poly res, tmp
+        cdef slong e_c
+
+        try:
+            e_c = e
+        except OverflowError:
+            # Exponent does not fit slong
+            res = nmod_poly.__new__(nmod_poly)
+            tmp = nmod_poly.__new__(nmod_poly)
+            nmod_poly_init_preinv(res.val, self.val.mod.n, self.val.mod.ninv)
+            nmod_poly_init_preinv(tmp.val, self.val.mod.n, self.val.mod.ninv)
+            ebytes = e.to_bytes((e.bit_length() + 15) // 16 * 2, "big")
+            nmod_poly_pow_trunc(res.val, self.val, ebytes[0] * 256 + ebytes[1], n)
+            for i in range(2, len(ebytes), 2):
+                nmod_poly_pow_trunc(res.val, res.val, 1 << 16, n)
+                nmod_poly_pow_trunc(tmp.val, self.val, ebytes[i] * 256 + ebytes[i+1], n)
+                nmod_poly_mullow(res.val, res.val, tmp.val, n)
+            return res
+
+        res = nmod_poly.__new__(nmod_poly)
+        nmod_poly_init_preinv(res.val, self.val.mod.n, self.val.mod.ninv)
+        nmod_poly_pow_trunc(res.val, self.val, e_c, n)
+        return res
+
     def gcd(self, other):
         """
         Returns the monic greatest common divisor of self and other.