Add shift and truncate methods to nmod_poly

README.md

@@ -166,9 +166,14 @@ Contributors
 
 - Oscar Benjamin (OB)
 - Robert Dougherty-Bliss (RDB)
+- Rémy Oudompheng (RO)
 
 Changes
 
+- [gh-310](https://github.com/flintlib/python-flint/pull/310),
+  Add `truncate`, `left_shift` and `right_shift` methods to
+  `fmpz_poly`, `fmpq_poly`, `nmod_poly`, `acb_poly`, `arb_poly`
+  to match other univariate polynomial types. (RO)
 - [gh-300](https://github.com/flintlib/python-flint/pull/300), Fix `arb.repr`
   which now returns a Python representation that round trips. (OB)
 - [gh-289](https://github.com/flintlib/python-flint/pull/289),

src/flint/test/test_all.py

@@ -2952,6 +2952,12 @@ def setbad(obj, i, val):
         assert raises(lambda: 1 / P([1, 1]), DomainError)
         assert raises(lambda: P([1, 2, 1]) / P([1, 2]), DomainError)
 
+    # Shifts and truncation.
+    assert P([1, 2, 3]).left_shift(3) == P([0, 0, 0, 1, 2, 3])
+    assert P([1, 2, 3]).right_shift(1) == P([2, 3])
+    assert P([1, 2, 3]).truncate(4) == P([1, 2, 3])
+    assert P([1, 2, 3]).truncate(2) == P([1, 2])
+
     assert P([1, 1]) ** 0 == P([1])
     assert P([1, 1]) ** 1 == P([1, 1])
     assert P([1, 1]) ** 2 == P([1, 2, 1])

src/flint/types/acb_poly.pyx

@@ -314,6 +314,94 @@ cdef class acb_poly(flint_poly):
             return s
         return divmod(t, s)
 
+    def truncate(self, slong n):
+        r"""
+        Notionally truncate the polynomial to have length ``n``. If
+        ``n`` is larger than the length of the input, then a copy of ``self`` is
+        returned. If ``n`` is not positive, then the zero polynomial
+        is returned.
+
+        Effectively returns this polynomial :math:`\mod x^n`.
+
+            >>> f = acb_poly([1,2,3])
+            >>> f.truncate(3)
+            3.00000000000000*x^2 + 2.00000000000000*x + 1.00000000000000
+            >>> f.truncate(2)
+            2.00000000000000*x + 1.00000000000000
+            >>> f.truncate(1)
+            1.00000000000000
+            >>> f.truncate(0)
+            0
+            >>> f.truncate(-1)
+            0
+
+        """
+        cdef acb_poly res
+        res = acb_poly.__new__(acb_poly)
+
+        length = acb_poly_length(self.val)
+        if n <= 0:  # return zero
+            return res
+        elif n > length:  # do nothing
+            acb_poly_set(res.val, self.val)
+        else:
+            acb_poly_set_trunc(res.val, self.val, n)
+
+        return res
+
+    def left_shift(self, slong n):
+        """
+        Returns ``self`` shifted left by ``n`` coefficients by inserting
+        zero coefficients. This is equivalent to multiplying the polynomial
+        by x^n
+
+            >>> f = acb_poly([1,2,3])
+            >>> f.left_shift(0)
+            3.00000000000000*x^2 + 2.00000000000000*x + 1.00000000000000
+            >>> f.left_shift(1)
+            3.00000000000000*x^3 + 2.00000000000000*x^2 + 1.00000000000000*x
+            >>> f.left_shift(4)
+            3.00000000000000*x^6 + 2.00000000000000*x^5 + 1.00000000000000*x^4
+
+        """
+        cdef acb_poly res
+        res = acb_poly.__new__(acb_poly)
+
+        if n < 0:
+            raise ValueError("Value must be shifted by a non-negative integer")
+        if n > 0:
+            acb_poly_shift_left(res.val, self.val, n)
+        else:  # do nothing, just copy self
+            acb_poly_set(res.val, self.val)
+
+        return res
+
+    def right_shift(self, slong n):
+        """
+        Returns ``self`` shifted right by ``n`` coefficients.
+        This is equivalent to the floor division of the polynomial
+        by x^n
+
+            >>> f = acb_poly([1,2,3])
+            >>> f.right_shift(0)
+            3.00000000000000*x^2 + 2.00000000000000*x + 1.00000000000000
+            >>> f.right_shift(1)
+            3.00000000000000*x + 2.00000000000000
+            >>> f.right_shift(4)
+            0
+        """
+        cdef acb_poly res
+        res = acb_poly.__new__(acb_poly)
+
+        if n < 0:
+            raise ValueError("Value must be shifted by a non-negative integer")
+        if n > 0:
+            acb_poly_shift_right(res.val, self.val, n)
+        else:  # do nothing, just copy self
+            acb_poly_set(res.val, self.val)
+
+        return res
+
     def __pow__(acb_poly s, ulong exp, mod):
         if mod is not None:
             raise NotImplementedError("acb_poly modular exponentiation")

src/flint/types/arb_poly.pyx

@@ -312,6 +312,94 @@ cdef class arb_poly(flint_poly):
             return s
         return divmod(t, s)
 
+    def truncate(self, slong n):
+        r"""
+        Notionally truncate the polynomial to have length ``n``. If
+        ``n`` is larger than the length of the input, then a copy of ``self`` is
+        returned. If ``n`` is not positive, then the zero polynomial
+        is returned.
+
+        Effectively returns this polynomial :math:`\mod x^n`.
+
+            >>> f = arb_poly([1,2,3])
+            >>> f.truncate(3)
+            3.00000000000000*x^2 + 2.00000000000000*x + 1.00000000000000
+            >>> f.truncate(2)
+            2.00000000000000*x + 1.00000000000000
+            >>> f.truncate(1)
+            1.00000000000000
+            >>> f.truncate(0)
+            0
+            >>> f.truncate(-1)
+            0
+
+        """
+        cdef arb_poly res
+        res = arb_poly.__new__(arb_poly)
+
+        length = arb_poly_length(self.val)
+        if n <= 0:  # return zero
+            return res
+        elif n > length:  # do nothing
+            arb_poly_set(res.val, self.val)
+        else:
+            arb_poly_set_trunc(res.val, self.val, n)
+
+        return res
+
+    def left_shift(self, slong n):
+        """
+        Returns ``self`` shifted left by ``n`` coefficients by inserting
+        zero coefficients. This is equivalent to multiplying the polynomial
+        by x^n
+
+            >>> f = arb_poly([1,2,3])
+            >>> f.left_shift(0)
+            3.00000000000000*x^2 + 2.00000000000000*x + 1.00000000000000
+            >>> f.left_shift(1)
+            3.00000000000000*x^3 + 2.00000000000000*x^2 + 1.00000000000000*x
+            >>> f.left_shift(4)
+            3.00000000000000*x^6 + 2.00000000000000*x^5 + 1.00000000000000*x^4
+
+        """
+        cdef arb_poly res
+        res = arb_poly.__new__(arb_poly)
+
+        if n < 0:
+            raise ValueError("Value must be shifted by a non-negative integer")
+        if n > 0:
+            arb_poly_shift_left(res.val, self.val, n)
+        else:  # do nothing, just copy self
+            arb_poly_set(res.val, self.val)
+
+        return res
+
+    def right_shift(self, slong n):
+        """
+        Returns ``self`` shifted right by ``n`` coefficients.
+        This is equivalent to the floor division of the polynomial
+        by x^n
+
+            >>> f = arb_poly([1,2,3])
+            >>> f.right_shift(0)
+            3.00000000000000*x^2 + 2.00000000000000*x + 1.00000000000000
+            >>> f.right_shift(1)
+            3.00000000000000*x + 2.00000000000000
+            >>> f.right_shift(4)
+            0
+        """
+        cdef arb_poly res
+        res = arb_poly.__new__(arb_poly)
+
+        if n < 0:
+            raise ValueError("Value must be shifted by a non-negative integer")
+        if n > 0:
+            arb_poly_shift_right(res.val, self.val, n)
+        else:  # do nothing, just copy self
+            arb_poly_set(res.val, self.val)
+
+        return res
+
     def __pow__(arb_poly s, ulong exp, mod):
         if mod is not None:
             raise NotImplementedError("arb_poly modular exponentiation")

src/flint/types/fmpq_poly.pyi

@@ -68,6 +68,10 @@ class fmpq_poly(flint_poly[fmpq]):
     def __rdivmod__(self, other: ifmpq_poly, /) -> tuple[fmpq_poly, fmpq_poly]: ...
     def __pow__(self, exp: int | ifmpz, /) -> fmpq_poly: ...
 
+    def left_shift(self, n: int, /) -> fmpq_poly: ...
+    def right_shift(self, n: int, /) -> fmpq_poly: ...
+    def truncate(self, n: int, /) -> fmpq_poly: ...
+
     def gcd(self, other: ifmpq_poly, /) -> fmpq_poly: ...
     def resultant(self, other: ifmpq_poly, /) -> fmpq: ...
     def xgcd(self, other: ifmpq_poly, /) -> tuple[fmpq_poly, fmpq_poly, fmpq_poly]: ...

src/flint/types/fmpq_poly.pyx

@@ -209,6 +209,41 @@ cdef class fmpq_poly(flint_poly):
         """
         return <bint>fmpq_poly_is_gen(self.val)
 
+    def truncate(self, slong n):
+        r"""
+        Notionally truncate the polynomial to have length ``n``. If
+        ``n`` is larger than the length of the input, then a copy of ``self`` is
+        returned. If ``n`` is not positive, then the zero polynomial
+        is returned.
+
+        Effectively returns this polynomial :math:`\mod x^n`.
+
+            >>> f = fmpq_poly([1,2,3])
+            >>> f.truncate(3) == f
+            True
+            >>> f.truncate(2)
+            2*x + 1
+            >>> f.truncate(1)
+            1
+            >>> f.truncate(0)
+            0
+            >>> f.truncate(-1)
+            0
+
+        """
+        cdef fmpq_poly res
+        res = fmpq_poly.__new__(fmpq_poly)
+
+        length = fmpq_poly_length(self.val)
+        if n <= 0:  # return zero
+            return res
+        elif n > length:  # do nothing
+            fmpq_poly_set(res.val, self.val)
+        else:
+            fmpq_poly_set_trunc(res.val, self.val, n)
+
+        return res
+
     def leading_coefficient(self):
         """
         Returns the leading coefficient of the polynomial.
@@ -402,6 +437,59 @@ cdef class fmpq_poly(flint_poly):
             return t
         return t._divmod_(s)
 
+    def left_shift(self, slong n):
+        """
+        Returns ``self`` shifted left by ``n`` coefficients by inserting
+        zero coefficients. This is equivalent to multiplying the polynomial
+        by x^n
+
+            >>> f = fmpq_poly([1,2,3])
+            >>> f.left_shift(0)
+            3*x^2 + 2*x + 1
+            >>> f.left_shift(1)
+            3*x^3 + 2*x^2 + x
+            >>> f.left_shift(4)
+            3*x^6 + 2*x^5 + x^4
+
+        """
+        cdef fmpq_poly res
+        res = fmpq_poly.__new__(fmpq_poly)
+
+        if n < 0:
+            raise ValueError("Value must be shifted by a non-negative integer")
+        if n > 0:
+            fmpq_poly_shift_left(res.val, self.val, n)
+        else:  # do nothing, just copy self
+            fmpq_poly_set(res.val, self.val)
+
+        return res
+
+    def right_shift(self, slong n):
+        """
+        Returns ``self`` shifted right by ``n`` coefficients.
+        This is equivalent to the floor division of the polynomial
+        by x^n
+
+            >>> f = fmpq_poly([1,2,3])
+            >>> f.right_shift(0)
+            3*x^2 + 2*x + 1
+            >>> f.right_shift(1)
+            3*x + 2
+            >>> f.right_shift(4)
+            0
+        """
+        cdef fmpq_poly res
+        res = fmpq_poly.__new__(fmpq_poly)
+
+        if n < 0:
+            raise ValueError("Value must be shifted by a non-negative integer")
+        if n > 0:
+            fmpq_poly_shift_right(res.val, self.val, n)
+        else:  # do nothing, just copy self
+            fmpq_poly_set(res.val, self.val)
+
+        return res
+
     def __pow__(fmpq_poly self, exp, mod):
         cdef fmpq_poly res
         if mod is not None:

src/flint/types/fmpz_mod_poly.pyx

@@ -990,7 +990,7 @@ cdef class fmpz_mod_poly(flint_poly):
     def truncate(self, slong n):
         r"""
         Notionally truncate the polynomial to have length ``n``. If
-        ``n`` is larger than the length of the input, then ``self`` is
+        ``n`` is larger than the length of the input, then a copy of ``self`` is
         returned. If ``n`` is not positive, then the zero polynomial
         is returned.
 

src/flint/types/fmpz_poly.pyi

@@ -55,6 +55,11 @@ class fmpz_poly(flint_poly[fmpz]):
     def __divmod__(self, other: ifmpz_poly, /) -> tuple[fmpz_poly, fmpz_poly]: ...
     def __rdivmod__(self, other: ifmpz, /) -> tuple[fmpz_poly, fmpz_poly]: ...
     def __pow__(self, other: int, /) -> fmpz_poly: ...
+
+    def left_shift(self, n: int, /) -> fmpz_poly: ...
+    def right_shift(self, n: int, /) -> fmpz_poly: ...
+    def truncate(self, n: int, /) -> fmpz_poly: ...
+
     def gcd(self, other: ifmpz_poly, /) -> fmpz_poly: ...
     def content(self) -> fmpz: ...
     def resultant(self, other: ifmpz_poly, /) -> fmpz: ...