Merge pull request #190 from Jake-Moss/mpoly_update

Refactor fmpz_mpoly and fmpq_mpoly arithmatic functions

Author: oscarbenjamin

src/flint/types/fmpq_mpoly.pyx

@@ -94,6 +94,22 @@ cdef class fmpq_mpoly_ctx(flint_mpoly_context):
         fmpq_mpoly_ctx_init(self.val, nvars, ordering_py_to_c(ordering))
         super().__init__(nvars, names)
 
+    def any_as_scalar(self, other):
+        if isinstance(other, int):
+            return any_as_fmpq(other)
+        elif typecheck(other, fmpz):
+            return any_as_fmpq(other)
+        elif typecheck(other, fmpq):
+            res = fmpq.__new__(fmpq)
+            fmpq_set((<fmpq>res).val, (<fmpq>other).val)
+            return res
+        else:
+            return NotImplemented
+
+    def scalar_as_mpoly(self, other: fmpq):
+        # non-fmpq scalars should first be converted via self.any_as_scalar
+        return self.constant(<fmpq>other)
+
     def nvars(self):
         """
         Return the number of variables in the context
@@ -336,205 +352,77 @@ cdef class fmpq_mpoly(flint_mpoly):
         fmpq_mpoly_neg(res.val, (<fmpq_mpoly>self).val, res.ctx.val)
         return res
 
-    def __add__(self, other):
-        cdef fmpq_mpoly res
-        if typecheck(other, fmpq_mpoly):
-            if (<fmpq_mpoly>self).ctx is not (<fmpq_mpoly>other).ctx:
-                raise IncompatibleContextError(f"{(<fmpq_mpoly>self).ctx} is not {(<fmpq_mpoly>other).ctx}")
-            res = create_fmpq_mpoly(self.ctx)
-            fmpq_mpoly_add(res.val, (<fmpq_mpoly>self).val, (<fmpq_mpoly>other).val, res.ctx.val)
-            return res
-        else:
-            other = any_as_fmpq(other)
-            if other is not NotImplemented:
-                res = create_fmpq_mpoly(self.ctx)
-                fmpq_mpoly_add_fmpq(res.val, (<fmpq_mpoly>self).val, (<fmpq>other).val, res.ctx.val)
-                return res
-        return NotImplemented
-
-    def __radd__(self, other):
-        cdef fmpq_mpoly res
-        other = any_as_fmpq(other)
-        if other is not NotImplemented:
-            res = create_fmpq_mpoly(self.ctx)
-            fmpq_mpoly_add_fmpq(res.val, (<fmpq_mpoly>self).val, (<fmpq>other).val, res.ctx.val)
-            return res
-        return NotImplemented
 
-    def __sub__(self, other):
-        cdef fmpq_mpoly res
-        if typecheck(other, fmpq_mpoly):
-            if (<fmpq_mpoly>self).ctx is not (<fmpq_mpoly>other).ctx:
-                raise IncompatibleContextError(f"{(<fmpq_mpoly>self).ctx} is not {(<fmpq_mpoly>other).ctx}")
-            res = create_fmpq_mpoly(self.ctx)
-            fmpq_mpoly_sub(res.val, (<fmpq_mpoly>self).val, (<fmpq_mpoly>other).val, res.ctx.val)
-            return res
-        else:
-            other = any_as_fmpq(other)
-            if other is not NotImplemented:
-                res = create_fmpq_mpoly(self.ctx)
-                fmpq_mpoly_sub_fmpq(res.val, (<fmpq_mpoly>self).val, (<fmpq>other).val, res.ctx.val)
-                return res
-        return NotImplemented
-
-    def __rsub__(self, other):
-        cdef fmpq_mpoly res
-        other = any_as_fmpq(other)
-        if other is not NotImplemented:
-            res = create_fmpq_mpoly(self.ctx)
-            fmpq_mpoly_sub_fmpq(res.val, (<fmpq_mpoly>self).val, (<fmpq>other).val, res.ctx.val)
-            return -res
-        return NotImplemented
-
-    def __mul__(self, other):
-        cdef fmpq_mpoly res
-        if typecheck(other, fmpq_mpoly):
-            if (<fmpq_mpoly>self).ctx is not (<fmpq_mpoly>other).ctx:
-                raise IncompatibleContextError(f"{(<fmpq_mpoly>self).ctx} is not {(<fmpq_mpoly>other).ctx}")
-            res = create_fmpq_mpoly(self.ctx)
-            fmpq_mpoly_mul(res.val, (<fmpq_mpoly>self).val, (<fmpq_mpoly>other).val, res.ctx.val)
-            return res
-        else:
-            other = any_as_fmpq(other)
-            if other is not NotImplemented:
-                res = create_fmpq_mpoly(self.ctx)
-                fmpq_mpoly_scalar_mul_fmpq(res.val, (<fmpq_mpoly>self).val, (<fmpq>other).val, res.ctx.val)
-                return res
-        return NotImplemented
-
-    def __rmul__(self, other):
+    def _add_scalar_(self, other: fmpq):
         cdef fmpq_mpoly res
-        other = any_as_fmpq(other)
-        if other is not NotImplemented:
-            res = create_fmpq_mpoly(self.ctx)
-            fmpq_mpoly_scalar_mul_fmpq(res.val, (<fmpq_mpoly>self).val, (<fmpq>other).val, res.ctx.val)
-            return res
-        return NotImplemented
+        res = create_fmpq_mpoly(self.ctx)
+        fmpq_mpoly_add_fmpq(res.val, self.val, other.val, self.ctx.val)
+        return res
 
-    def __pow__(self, other, modulus):
+    def _add_mpoly_(self, other: fmpq_mpoly):
         cdef fmpq_mpoly res
-        if modulus is not None:
-            raise NotImplementedError
-        other = any_as_fmpz(other)
-        if other is NotImplemented:
-            return other
-        if other < 0:
-            raise ValueError("cannot raise fmpq_mpoly to negative power")
         res = create_fmpq_mpoly(self.ctx)
-        if fmpq_mpoly_pow_fmpz(res.val, (<fmpq_mpoly>self).val, (<fmpz>other).val, res.ctx.val) == 0:
-            raise ValueError("unreasonably large polynomial")  # pragma: no cover
+        fmpq_mpoly_add(res.val, self.val, other.val, res.ctx.val)
         return res
 
-    def __divmod__(self, other):
-        cdef fmpq_mpoly res, res2
-        if typecheck(other, fmpq_mpoly):
-            if not other:
-                raise ZeroDivisionError("fmpq_mpoly division by zero")
-            elif (<fmpq_mpoly>self).ctx is not (<fmpq_mpoly>other).ctx:
-                raise IncompatibleContextError(f"{(<fmpq_mpoly>self).ctx} is not {(<fmpq_mpoly>other).ctx}")
-            res = create_fmpq_mpoly(self.ctx)
-            res2 = create_fmpq_mpoly(self.ctx)
-            fmpq_mpoly_divrem(res.val, res2.val, (<fmpq_mpoly>self).val, (<fmpq_mpoly>other).val, res.ctx.val)
-            return (res, res2)
-        else:
-            other = any_as_fmpq(other)
-            if other is not NotImplemented:
-                other = fmpq_mpoly(other, self.ctx)
-                if not other:
-                    raise ZeroDivisionError("fmpq_mpoly division by zero")
-                res = create_fmpq_mpoly(self.ctx)
-                res2 = create_fmpq_mpoly(self.ctx)
-                fmpq_mpoly_divrem(res.val, res2.val, (<fmpq_mpoly>self).val, (<fmpq_mpoly>other).val, res.ctx.val)
-                return (res, res2)
-        return NotImplemented
-
-    def __rdivmod__(self, other):
-        cdef fmpq_mpoly res, res2
-        if not self:
-            raise ZeroDivisionError("fmpq_mpoly division by zero")
-        other = any_as_fmpq(other)
-        if other is not NotImplemented:
-            other = fmpq_mpoly(other, self.ctx)
-            res = create_fmpq_mpoly(self.ctx)
-            res2 = create_fmpq_mpoly(self.ctx)
-            fmpq_mpoly_divrem(res.val, res2.val, (<fmpq_mpoly>other).val, (<fmpq_mpoly>self).val, res.ctx.val)
-            return (res, res2)
-        return NotImplemented
-
-    def __floordiv__(self, other):
+    def _sub_scalar_(self, other: fmpq):
         cdef fmpq_mpoly res
-        if typecheck(other, fmpq_mpoly):
-            if not other:
-                raise ZeroDivisionError("fmpq_mpoly division by zero")
-            elif (<fmpq_mpoly>self).ctx is not (<fmpq_mpoly>other).ctx:
-                raise IncompatibleContextError(f"{(<fmpq_mpoly>self).ctx} is not {(<fmpq_mpoly>other).ctx}")
-            res = create_fmpq_mpoly(self.ctx)
-            fmpq_mpoly_div(res.val, (<fmpq_mpoly>self).val, (<fmpq_mpoly>other).val, res.ctx.val)
-            return res
-        else:
-            other = any_as_fmpq(other)
-            if other is not NotImplemented:
-                if not other:
-                    raise ZeroDivisionError("fmpq_mpoly division by zero")
-                other = fmpq_mpoly(other, self.ctx)
-                res = create_fmpq_mpoly(self.ctx)
-                fmpq_mpoly_div(res.val, (<fmpq_mpoly>self).val, (<fmpq_mpoly>other).val, res.ctx.val)
-                return res
-        return NotImplemented
-
-    def __rfloordiv__(self, other):
-        cdef fmpq_mpoly res
-        if not self:
-            raise ZeroDivisionError("fmpq_mpoly division by zero")
-        other = any_as_fmpq(other)
-        if other is not NotImplemented:
-            other = fmpq_mpoly(other, self.ctx)
-            res = create_fmpq_mpoly(self.ctx)
-            fmpq_mpoly_div(res.val, (<fmpq_mpoly>other).val, self.val, res.ctx.val)
-            return res
-        return NotImplemented
+        res = create_fmpq_mpoly(self.ctx)
+        fmpq_mpoly_sub_fmpq(res.val, self.val, other.val, self.ctx.val)
+        return res
 
-    def __truediv__(self, other):
-        cdef:
-            fmpq_mpoly res
+    def _sub_mpoly_(self, other: fmpq_mpoly):
+        cdef fmpq_mpoly res
+        res = create_fmpq_mpoly(self.ctx)
+        fmpq_mpoly_sub(res.val, self.val, other.val, res.ctx.val)
+        return res
 
-        if typecheck(other, fmpq_mpoly):
-            if not other:
-                raise ZeroDivisionError("fmpq_mpoly division by zero")
-            elif self.ctx is not (<fmpq_mpoly>other).ctx:
-                raise IncompatibleContextError(f"{self.ctx} is not {(<fmpq_mpoly>other).ctx}")
+    def _mul_scalar_(self, other: fmpq):
+        cdef fmpq_mpoly res
+        res = create_fmpq_mpoly(self.ctx)
+        fmpq_mpoly_scalar_mul_fmpq(res.val, self.val, other.val, self.ctx.val)
+        return res
 
-            res = create_fmpq_mpoly(self.ctx)
-            if fmpq_mpoly_divides(res.val, self.val, (<fmpq_mpoly>other).val, self.ctx.val):
-                return res
-            else:
-                raise DomainError("fmpq_mpoly division is not exact")
-        else:
-            o = any_as_fmpq(other)
-            if o is NotImplemented:
-                return NotImplemented
-            elif not o:
-                raise ZeroDivisionError("fmpq_mpoly division by zero")
-            res = create_fmpq_mpoly(self.ctx)
-            fmpq_mpoly_scalar_div_fmpq(res.val, self.val, (<fmpq>o).val, self.ctx.val)
-            return res
+    def _mul_mpoly_(self, other: fmpq_mpoly):
+        cdef fmpq_mpoly res
+        res = create_fmpq_mpoly(self.ctx)
+        fmpq_mpoly_mul(res.val, self.val, other.val, res.ctx.val)
+        return res
 
-    def __rtruediv__(self, other):
+    def _pow_(self, other: fmpz):
         cdef fmpq_mpoly res
-        if not self:
-            raise ZeroDivisionError("fmpq_mpoly division by zero")
-        o = any_as_fmpq(other)
-        if o is NotImplemented:
-            return NotImplemented
         res = create_fmpq_mpoly(self.ctx)
-        fmpq_mpoly_set_fmpq(res.val, (<fmpq>o).val, self.ctx.val)
-        return res / self
+        if fmpq_mpoly_pow_fmpz(res.val, self.val, other.val, res.ctx.val) == 0:
+            raise ValueError("unreasonably large polynomial")  # pragma: no cover
+        return res
 
-    def __mod__(self, other):
-        return divmod(self, other)[1]
+    def _divmod_mpoly_(self, other: fmpq_mpoly):
+        cdef fmpq_mpoly quotient, remainder
+        quotient = create_fmpq_mpoly(self.ctx)
+        remainder = create_fmpq_mpoly(self.ctx)
+        fmpq_mpoly_divrem(quotient.val, remainder.val, self.val, other.val, self.ctx.val)
+        return (quotient, remainder)
+
+    def _floordiv_mpoly_(self, other: fmpq_mpoly):
+        cdef fmpq_mpoly quotient
+        quotient = create_fmpq_mpoly(self.ctx)
+        fmpq_mpoly_div(quotient.val, self.val, other.val, self.ctx.val)
+        return quotient
+
+    def _truediv_mpoly_(self, other: fmpq_mpoly):
+        cdef fmpq_mpoly quotient
+        quotient = create_fmpq_mpoly(self.ctx)
+        if fmpq_mpoly_divides(quotient.val, self.val, other.val, self.ctx.val):
+            return quotient
+        else:
+            raise DomainError("fmpq_mpoly division is not exact")
 
-    def __rmod__(self, other):
-        return divmod(other, self)[1]
+    def _mod_mpoly_(self, other: fmpq_mpoly):
+        cdef fmpq_mpoly quotient, remainder
+        quotient = create_fmpq_mpoly(self.ctx)
+        remainder = create_fmpq_mpoly(self.ctx)
+        fmpq_mpoly_divrem(quotient.val, remainder.val, self.val, other.val, self.ctx.val)
+        return remainder
 
     def __call__(self, *args) -> fmpq:
         cdef:
@@ -551,83 +439,23 @@ cdef class fmpq_mpoly(flint_mpoly):
             raise ValueError("unreasonably large polynomial")  # pragma: no cover
         return vres
 
-    def iadd(self, other):
-        """
-        In-place addition, mutates self.
+    def _iadd_scalar_(self, other: fmpq):
+        fmpq_mpoly_add_fmpq(self.val, self.val, other.val, self.ctx.val)
 
-            >>> from flint import Ordering
-            >>> ctx = fmpq_mpoly_ctx.get_context(2, Ordering.lex, 'x')
-            >>> f = ctx.from_dict({(1, 0): 2, (0, 1): 3, (1, 1): 4})
-            >>> f
-            4*x0*x1 + 2*x0 + 3*x1
-            >>> f.iadd(5)
-            >>> f
-            4*x0*x1 + 2*x0 + 3*x1 + 5
+    def _iadd_mpoly_(self, other: fmpq_mpoly):
+        fmpq_mpoly_add(self.val, self.val, other.val, self.ctx.val)
 
-        """
-        if typecheck(other, fmpq_mpoly):
-            if (<fmpq_mpoly>self).ctx is not (<fmpq_mpoly>other).ctx:
-                raise IncompatibleContextError(f"{(<fmpq_mpoly>self).ctx} is not {(<fmpq_mpoly>other).ctx}")
-            fmpq_mpoly_add((<fmpq_mpoly>self).val, (<fmpq_mpoly>self).val, (<fmpq_mpoly>other).val, self.ctx.val)
-            return
-        else:
-            other = any_as_fmpq(other)
-            if other is not NotImplemented:
-                fmpq_mpoly_add_fmpq((<fmpq_mpoly>self).val, (<fmpq_mpoly>self).val, (<fmpq>other).val, self.ctx.val)
-                return
-        raise NotImplementedError(f"in-place addition not implemented between {type(self)} and {type(other)}")
-
-    def isub(self, other):
-        """
-        In-place subtraction, mutates self.
-
-            >>> from flint import Ordering
-            >>> ctx = fmpq_mpoly_ctx.get_context(2, Ordering.lex, 'x')
-            >>> f = ctx.from_dict({(1, 0): 2, (0, 1): 3, (1, 1): 4})
-            >>> f
-            4*x0*x1 + 2*x0 + 3*x1
-            >>> f.isub(5)
-            >>> f
-            4*x0*x1 + 2*x0 + 3*x1 - 5
-
-        """
-        if typecheck(other, fmpq_mpoly):
-            if (<fmpq_mpoly>self).ctx is not (<fmpq_mpoly>other).ctx:
-                raise IncompatibleContextError(f"{(<fmpq_mpoly>self).ctx} is not {(<fmpq_mpoly>other).ctx}")
-            fmpq_mpoly_sub((<fmpq_mpoly>self).val, (<fmpq_mpoly>self).val, (<fmpq_mpoly>other).val, self.ctx.val)
-            return
-        else:
-            other = any_as_fmpq(other)
-            if other is not NotImplemented:
-                fmpq_mpoly_sub_fmpq((<fmpq_mpoly>self).val, (<fmpq_mpoly>self).val, (<fmpq>other).val, self.ctx.val)
-                return
-        raise NotImplementedError(f"in-place subtraction not implemented between {type(self)} and {type(other)}")
+    def _isub_scalar_(self, other: fmpq):
+        fmpq_mpoly_sub_fmpq(self.val, self.val, other.val, self.ctx.val)
 
-    def imul(self, other):
-        """
-        In-place multiplication, mutates self.
+    def _isub_mpoly_(self, other: fmpq_mpoly):
+        fmpq_mpoly_sub(self.val, self.val, other.val, self.ctx.val)
 
-            >>> from flint import Ordering
-            >>> ctx = fmpq_mpoly_ctx.get_context(2, Ordering.lex, 'x')
-            >>> f = ctx.from_dict({(1, 0): 2, (0, 1): 3, (1, 1): 4})
-            >>> f
-            4*x0*x1 + 2*x0 + 3*x1
-            >>> f.imul(2)
-            >>> f
-            8*x0*x1 + 4*x0 + 6*x1
+    def _imul_scalar_(self, other: fmpq):
+        fmpq_mpoly_scalar_mul_fmpq(self.val, self.val, other.val, self.ctx.val)
 
-        """
-        if typecheck(other, fmpq_mpoly):
-            if (<fmpq_mpoly>self).ctx is not (<fmpq_mpoly>other).ctx:
-                raise IncompatibleContextError(f"{(<fmpq_mpoly>self).ctx} is not {(<fmpq_mpoly>other).ctx}")
-            fmpq_mpoly_mul((<fmpq_mpoly>self).val, (<fmpq_mpoly>self).val, (<fmpq_mpoly>other).val, self.ctx.val)
-            return
-        else:
-            other = any_as_fmpq(other)
-            if other is not NotImplemented:
-                fmpq_mpoly_scalar_mul_fmpq(self.val, (<fmpq_mpoly>self).val, (<fmpq>other).val, self.ctx.val)
-                return
-        raise NotImplementedError(f"in-place multiplication not implemented between {type(self)} and {type(other)}")
+    def _imul_mpoly_(self, other: fmpq_mpoly):
+        fmpq_mpoly_mul(self.val, self.val, other.val, self.ctx.val)
 
     def monoms(self):
         """

src/flint/types/fmpz_mod_mpoly.pyx

@@ -137,7 +137,9 @@ cdef class fmpz_mod_mpoly_ctx(flint_mpoly_context):
                 )
             return any_as_fmpz((<nmod>other).val)
         elif typecheck(other, fmpz):
-            return fmpz(other)
+            res = fmpz.__new__(fmpz)
+            fmpz_set((<fmpz>res).val, (<fmpz>other).val)
+            return res
         elif typecheck(other, fmpz_mod):
             if (<fmpz_mod>other).ctx.modulus() != self.modulus():
                 raise DomainError(
@@ -150,7 +152,7 @@ cdef class fmpz_mod_mpoly_ctx(flint_mpoly_context):
             return NotImplemented
 
     def scalar_as_mpoly(self, other: fmpz):
-        # non-fmpz scalars should first be converted via cls.any_as_scalar
+        # non-fmpz scalars should first be converted via self.any_as_scalar
         return self.constant(<fmpz>other)
 
     def nvars(self):