Add drop_gens, coerce_to_context, and infer_context_mapping functions

WIp

wip

Author: Jake-Moss

src/flint/flint_base/flint_base.pxd

@@ -1,4 +1,5 @@
 from flint.flintlib.types.mpoly cimport ordering_t
+from flint.flintlib.types.flint cimport slong
 
 cdef class flint_elem:
     pass
@@ -50,6 +51,7 @@ cdef class flint_mpoly(flint_elem):
     cdef _isub_mpoly_(self, other)
     cdef _imul_mpoly_(self, other)
 
+    cdef _compose_gens_(self, ctx, slong *mapping)
 
 cdef class flint_mat(flint_elem):
     pass

src/flint/flint_base/flint_base.pyx

@@ -2,6 +2,7 @@ from flint.flintlib.types.flint cimport (
     FLINT_BITS as _FLINT_BITS,
     FLINT_VERSION as _FLINT_VERSION,
     __FLINT_RELEASE as _FLINT_RELEASE,
+    slong,
 )
 from flint.utils.flint_exceptions import DomainError
 from flint.flintlib.types.mpoly cimport ordering_t
@@ -344,13 +345,20 @@ cdef class flint_mpoly_context(flint_elem):
         return tuple(self.gen(i) for i in range(self.nvars()))
 
     def variable_to_index(self, var: Union[int, str]) -> int:
-        """Convert a variable name string or possible index to its index in the context."""
+        """
+        Convert a variable name string or possible index to its index in the context.
+
+        If ``var`` is negative, return the index of the ``self.nvars() + var``
+        """
         if isinstance(var, str):
             try:
                 i = self.names().index(var)
             except ValueError:
                 raise ValueError("variable not in context")
         elif isinstance(var, int):
+            if var < 0:
+                var = self.nvars() + var
+
             if not 0 <= var < self.nvars():
                 raise IndexError("generator index out of range")
             i = var
@@ -379,7 +387,7 @@ cdef class flint_mpoly_context(flint_elem):
             names = (names,)
 
         for name in names:
-            if isinstance(name, str):
+            if isinstance(name, (str, bytes)):
                 res.append(name)
             else:
                 base, num = name
@@ -415,10 +423,14 @@ cdef class flint_mpoly_context(flint_elem):
         return ctx
 
     @classmethod
-    def from_context(cls, ctx: flint_mpoly_context):
+    def from_context(cls, ctx: flint_mpoly_context, names=None, ordering=None):
+        """
+        Get a new context from an existing one. Optionally override ``names`` or
+        ``ordering``.
+        """
         return cls.get(
-            ordering=ctx.ordering(),
-            names=ctx.names(),
+            names=ctx.names() if names is None else names,
+            ordering=ctx.ordering() if ordering is None else names,
         )
 
     def _any_as_scalar(self, other):
@@ -451,6 +463,47 @@ cdef class flint_mpoly_context(flint_elem):
             exp_vec = (0,) * self.nvars()
         return self.from_dict({tuple(exp_vec): coeff})
 
+    def drop_gens(self, *gens: str | int):
+        """
+        Get a context with the specified generators removed.
+
+            >>> from flint import fmpz_mpoly_ctx
+            >>> ctx = fmpz_mpoly_ctx.get(('x', 'y', 'z', 'a', 'b'))
+            >>> ctx.drop_gens('x', -2)
+            fmpz_mpoly_ctx(3, '<Ordering.lex: 'lex'>', ('y', 'z', 'b'))
+        """
+        nvars = self.nvars()
+        gen_idxs = set(self.variable_to_index(i) for i in gens)
+
+        if len(gens) > nvars:
+            raise ValueError(f"expected at most {nvars} unique generators, got {len(gens)}")
+
+        remaining_gens = []
+        for i in range(nvars):
+            if i not in gen_idxs:
+                remaining_gens.append(self.py_names[i])
+
+        return self.from_context(self, names=remaining_gens)
+
+    def infer_generator_mapping(self, ctx: flint_mpoly_context):
+        """
+        Infer a mapping of generator indexes from this contexts generators, to the
+        provided contexts generators. Inference is done based upon generator names.
+
+            >>> from flint import fmpz_mpoly_ctx
+            >>> ctx = fmpz_mpoly_ctx.get(('x', 'y', 'z', 'a', 'b'))
+            >>> ctx2 = fmpz_mpoly_ctx.get(('b', 'a'))
+            >>> ctx.infer_generator_mapping(ctx2)
+            {3: 1, 4: 0}
+        """
+        gens_to_idxs = {x: i for i, x in enumerate(self.names())}
+        other_gens_to_idxs = {x: i for i, x in enumerate(ctx.names())}
+        return {
+            gens_to_idxs[k]: other_gens_to_idxs[k]
+            for k in (gens_to_idxs.keys() & other_gens_to_idxs.keys())
+        }
+
+
 cdef class flint_mod_mpoly_context(flint_mpoly_context):
     @classmethod
     def _new_(_, flint_mod_mpoly_context self, names, prime_modulus):
@@ -472,11 +525,15 @@ cdef class flint_mod_mpoly_context(flint_mpoly_context):
         return *super().create_context_key(names, ordering), modulus
 
     @classmethod
-    def from_context(cls, ctx: flint_mod_mpoly_context):
+    def from_context(cls, ctx: flint_mod_mpoly_context, names=None, ordering=None, modulus=None):
+        """
+        Get a new context from an existing one. Optionally override ``names``,
+        ``modulus``, or ``ordering``.
+        """
         return cls.get(
-            names=ctx.names(),
-            modulus=ctx.modulus(),
-            ordering=ctx.ordering(),
+            names=ctx.names() if names is None else names,
+            modulus=ctx.modulus() if modulus is None else modulus,
+            ordering=ctx.ordering() if ordering is None else ordering,
         )
 
     def is_prime(self):
@@ -869,6 +926,88 @@ cdef class flint_mpoly(flint_elem):
         """
         return zip(self.monoms(), self.coeffs())
 
+    def drop_unused_gens(self):
+        """
+        Remove unused generators from this polynomial. Returns a potentially new
+        context, and potentially new polynomial.
+
+        A generator is unused if it's maximum degree is 0.
+
+            >>> from flint import fmpz_mpoly_ctx
+            >>> ctx = fmpz_mpoly_ctx.get(('x', 4))
+            >>> ctx2 = fmpz_mpoly_ctx.get(('x1', 'x2'))
+            >>> f = sum(ctx.gens()[1:3])
+            >>> f
+            x1 + x2
+            >>> new_ctx, new_f = f.drop_unused_gens()
+            >>> new_ctx
+            fmpz_mpoly_ctx(2, '<Ordering.lex: 'lex'>', ('x1', 'x2'))
+            >>> new_f
+            x1 + x2
+        """
+        new_ctx = self.context().drop_gens(
+            *(i for i, x in enumerate(self.degrees()) if not x)
+        )
+        return new_ctx, self.coerce_to_context(new_ctx)
+
+    def coerce_to_context(self, ctx, mapping: dict[str | int, str | int] = None):
+        """
+        Coerce this polynomial to a different context.
+
+        This is equivalent to composing this polynomial with the generators of another
+        context. By default the mapping between contexts is inferred based on the name
+        of the generators. Generators with names that are not found within the other
+        context are mapped to 0. The mapping can be explicitly provided.
+
+        The resulting polynomial is also normalised.
+
+            >>> from flint import fmpz_mpoly_ctx
+            >>> ctx = fmpz_mpoly_ctx.get(('x', 'y', 'a', 'b'))
+            >>> ctx2 = fmpz_mpoly_ctx.get(('a', 'b'))
+            >>> x, y, a, b = ctx.gens()
+            >>> f = x + 2*y + 3*a + 4*b
+            >>> f.coerce_to_context(ctx2)
+            3*a + 4*b
+            >>> f.coerce_to_context(ctx2, mapping={"x": "a", "b": 0})
+            5*a
+        """
+        cdef:
+            slong *c_mapping
+            slong i
+
+        if not typecheck(ctx, type(self.ctx)):
+            raise ValueError(
+                f"provided context is not a {self.ctx.__class__.__name__}"
+            )
+        elif self.ctx is ctx:
+            return self
+
+        if mapping is None:
+            mapping = self.ctx.infer_generator_mapping(ctx)
+        else:
+            mapping = {
+                self.ctx.variable_to_index(k): ctx.variable_to_index(v)
+                for k, v in mapping.items()
+            }
+
+        try:
+            c_mapping = <slong *> libc.stdlib.malloc(self.ctx.nvars() * sizeof(slong *))
+            if c_mapping is NULL:
+                raise MemoryError("malloc returned a null pointer")
+
+            for i in range(self.ctx.nvars()):
+                c_mapping[i] = <slong>-1
+
+            for k, v in mapping.items():
+                c_mapping[k] = <slong>v
+
+            return self._compose_gens_(ctx, c_mapping)
+        finally:
+            libc.stdlib.free(c_mapping)
+
+    cdef _compose_gens_(self, ctx, slong *mapping):
+        raise NotImplementedError("abstract method")
+
 
 cdef class flint_series(flint_elem):
     """

src/flint/flintlib/types/flint.pxd

@@ -59,8 +59,10 @@ cdef extern from *:
 cdef extern from "flint/flint.h":
     """
     #define SIZEOF_ULONG sizeof(ulong)
+    #define SIZEOF_SLONG sizeof(slong)
     """
     int SIZEOF_ULONG
+    int SIZEOF_SLONG
     const char * FLINT_VERSION
     const int __FLINT_RELEASE
     const int FLINT_BITS

src/flint/types/fmpq_mpoly.pyx

@@ -20,7 +20,9 @@ from flint.flintlib.functions.fmpq_mpoly cimport (
     fmpq_mpoly_add,
     fmpq_mpoly_add_fmpq,
     fmpq_mpoly_clear,
+    fmpq_mpoly_combine_like_terms,
     fmpq_mpoly_compose_fmpq_mpoly,
+    fmpq_mpoly_compose_fmpq_mpoly_gen,
     fmpq_mpoly_ctx_init,
     fmpq_mpoly_degrees_fmpz,
     fmpq_mpoly_derivative,
@@ -40,6 +42,7 @@ from flint.flintlib.functions.fmpq_mpoly cimport (
     fmpq_mpoly_get_term_coeff_fmpq,
     fmpq_mpoly_get_term_exp_fmpz,
     fmpq_mpoly_integral,
+    fmpq_mpoly_is_canonical,
     fmpq_mpoly_is_one,
     fmpq_mpoly_is_zero,
     fmpq_mpoly_length,
@@ -1119,6 +1122,22 @@ cdef class fmpq_mpoly(flint_mpoly):
         fmpz_mpoly_deflation(shift.val, stride.val, self.val.zpoly, self.ctx.val.zctx)
         return list(stride), list(shift)
 
+    cdef _compose_gens_(self, ctx, slong *mapping):
+        cdef fmpq_mpoly res = create_fmpq_mpoly(ctx)
+        fmpq_mpoly_compose_fmpq_mpoly_gen(
+            res.val,
+            self.val,
+            mapping,
+            self.ctx.val,
+            (<fmpq_mpoly_ctx>ctx).val
+        )
+
+        if not fmpq_mpoly_is_canonical(res.val, res.ctx.val):
+            fmpq_mpoly_sort_terms(res.val, res.ctx.val)
+            fmpq_mpoly_combine_like_terms(res.val, res.ctx.val)
+
+        return res
+
 
 cdef class fmpq_mpoly_vec:
     """

src/flint/types/fmpz_mod_mpoly.pyx

@@ -18,7 +18,9 @@ from flint.flintlib.functions.fmpz_mod_mpoly cimport (
     fmpz_mod_mpoly_add,
     fmpz_mod_mpoly_add_fmpz,
     fmpz_mod_mpoly_clear,
+    fmpz_mod_mpoly_combine_like_terms,
     fmpz_mod_mpoly_compose_fmpz_mod_mpoly,
+    fmpz_mod_mpoly_compose_fmpz_mod_mpoly_gen,
     fmpz_mod_mpoly_ctx_get_modulus,
     fmpz_mod_mpoly_ctx_init,
     fmpz_mod_mpoly_deflate,
@@ -40,6 +42,7 @@ from flint.flintlib.functions.fmpz_mod_mpoly cimport (
     fmpz_mod_mpoly_get_term_coeff_fmpz,
     fmpz_mod_mpoly_get_term_exp_fmpz,
     fmpz_mod_mpoly_inflate,
+    fmpz_mod_mpoly_is_canonical,
     fmpz_mod_mpoly_is_one,
     fmpz_mod_mpoly_is_zero,
     fmpz_mod_mpoly_length,
@@ -1136,6 +1139,22 @@ cdef class fmpz_mod_mpoly(flint_mpoly):
         fmpz_mod_mpoly_deflation(shift.val, stride.val, self.val, self.ctx.val)
         return list(stride), list(shift)
 
+    cdef _compose_gens_(self, ctx, slong *mapping):
+        cdef fmpz_mod_mpoly res = create_fmpz_mod_mpoly(ctx)
+        fmpz_mod_mpoly_compose_fmpz_mod_mpoly_gen(
+            res.val,
+            self.val,
+            mapping,
+            self.ctx.val,
+            (<fmpz_mod_mpoly_ctx>ctx).val
+        )
+
+        if not fmpz_mod_mpoly_is_canonical(res.val, res.ctx.val):
+            fmpz_mod_mpoly_sort_terms(res.val, res.ctx.val)
+            fmpz_mod_mpoly_combine_like_terms(res.val, res.ctx.val)
+
+        return res
+
 
 cdef class fmpz_mod_mpoly_vec:
     """