Add PolyIndex.as_multi_index for Newton polytope and minor cleanup

1 files changed, 61 insertions, 27 deletions

diff --git a/mdpoly/index.py b/mdpoly/index.py
index b1aa8c1..c085783 100644
--- a/mdpoly/index.py
+++ b/mdpoly/index.py
@@ -4,7 +4,7 @@ from .errors import InvalidShape
 from .util import partition, isclose
 
 from itertools import filterfalse
-from typing import Self, NamedTuple, Optional, TYPE_CHECKING
+from typing import Self, NamedTuple, Optional, Sequence, TYPE_CHECKING
 
 if TYPE_CHECKING:
     from .state import State
@@ -83,11 +83,8 @@ class PolyVarIndex(NamedTuple):
     """
     var_idx: int # Index in State.variables, not variable!
     power: Number
-
-    @classmethod
-    def from_var(cls, variable: Var, state: State, power: Number =1) -> Self:
-        """ Make an index from a variable object. """
-        return cls(var_idx=state.index(variable), power=power)
+    
+    # --- Magic methods ---
 
     def __eq__(self, other):
         if type(other) is not PolyVarIndex:
@@ -107,6 +104,24 @@ class PolyVarIndex(NamedTuple):
 
         return self.var_idx < other.var_idx
 
+    # -- Helper methods
+
+    @property
+    def exponent(self):
+        """ Alias for ``self.power``. """
+        return self.power
+
+    def is_constant(self) -> bool:
+        """ Check if is index of a constant term. """
+        return self.var_idx == -1
+    
+    # -- Helper methods for construction ---
+
+    @classmethod
+    def from_var(cls, variable: Var, state: State, power: Number =1) -> Self:
+        """ Make an index from a variable object. """
+        return cls(var_idx=state.index(variable), power=power)
+
     @classmethod
     def constant(cls) -> Self:
         """ Special index for constants.
@@ -116,11 +131,6 @@ class PolyVarIndex(NamedTuple):
         """
         return cls(var_idx=-1, power=0)
 
-    @staticmethod
-    def is_constant(index: PolyVarIndex) -> bool:
-        """ Check if is index of a constant term. """
-        return index.var_idx == -1
-
 
 class PolyIndex(tuple[PolyVarIndex]):
     """
@@ -132,17 +142,19 @@ class PolyIndex(tuple[PolyVarIndex]):
 
     .. code:: py
 
-        PolyIndex(PolyVarIndex(var_idx=0, power=1), PolyVarIndex(var_idx=1, power=2)
+       PolyIndex(PolyVarIndex(var_idx=0, power=1), PolyVarIndex(var_idx=1, power=2)
 
     Then given the polynomial
 
     .. math::
-        p(x, y) = 3x^2 + 5xy^2
+       p(x, y) = 3x^2 + 5xy^2
 
     The with the ``PolyIndex`` above we can retrieve the coefficient 5 in front
     of :math:`xy^2` from the ``Repr`` object (see also :py:meth:`mdpoly.abc.Repr.at`).
     """
 
+    # -- Magic methods ---
+
     def __repr__(self) -> str:
         name = self.__class__.__qualname__
         indices = ", ".join(map(repr, self))
@@ -151,11 +163,41 @@ class PolyIndex(tuple[PolyVarIndex]):
     def __contains__(self, var_idx):
         return any(map(lambda e: e.var_idx == var_idx, self))
 
+    # -- Helper methods ---
+
+    def as_multi_index(self) -> Sequence[Number]:
+        """ Get the multi-index of the monomial term.
+
+        For example for an index representing :math:`x^2 y` the multi-index is
+        a tuple :math:`(2, 1)`. Note that :math:`yx^2` has also the same
+        multi-index :math:`(2, 1)`. The order in which the exponents (powers)
+        are given in the tuple is ascending acording to the index of the
+        variables in the state (see total order of
+        :py:class:`mdpoly.index.PolyVarIndex`). In the given example then
+        :math:`x` would have index 0 and :math:`y` index 1.
+
+        The multi-index can be used to compute the Newton polytope of a
+        polynomial (which is the convex hull of the multi-indices).
+        """
+        def take_power(idx: PolyVarIndex) -> Number:
+            return idx.power
+
+        return tuple(map(take_power, sorted(self)))
+
+    def is_constant(self) -> bool:
+        """ Check if is index of a constant term. """
+        if len(self) != 1:
+            return False
+
+        return PolyVarIndex.is_constant(self[0])
+
+    # --- Helper methods for construction ---
+
     @classmethod
     def from_dict(cls, d) -> Self:
-        """ Construct an index froma dictionary, where the keys are the
-        variable index from the state object and the values are the exponent
-        (power). """
+        """ Construct an index from a dictionary, where the keys are the
+        variable indices from the state object and the values are the exponents
+        (powers). """
         return cls(PolyVarIndex(k, v) for k, v in d.items())
 
     @classmethod
@@ -163,14 +205,6 @@ class PolyIndex(tuple[PolyVarIndex]):
         """ Index of the constant term. """
         return cls((PolyVarIndex.constant(),))
 
-    @staticmethod
-    def is_constant(index: PolyIndex) -> bool:
-        """ Check if is index of a constant term. """
-        if len(index) != 1:
-            return False
-
-        return PolyVarIndex.is_constant(index[0])
-
     @classmethod
     def sort(cls, index: tuple | Self) -> Self:
         """ Sort a tuple of indices. """
@@ -183,10 +217,10 @@ class PolyIndex(tuple[PolyVarIndex]):
         For example if left is the index of :math:`xy` and right is the index
         of :math:`y^2` this functions returns the index of :math:`xy^3`.
         """
-        if cls.is_constant(left):
+        if left.is_constant():
             return right
 
-        if cls.is_constant(right):
+        if right.is_constant():
             return left
 
         result: dict[int, Number] = dict(filterfalse(PolyVarIndex.is_constant, left))
@@ -209,7 +243,7 @@ class PolyIndex(tuple[PolyVarIndex]):
         is a constant ``None`` is returned.
         """
 
-        if cls.is_constant(index):
+        if index.is_constant():
             return None
 
         if wrt not in index: