Add shifts and truncate methods for arb_poly and acb_poly

Author: remyoudompheng

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")