Add helper methods to retrieve variables from ExpressionState

1 files changed, 21 insertions, 0 deletions

diff --git a/polymatrix/expressionstate/mixins.py b/polymatrix/expressionstate/mixins.py
index 36dc5e7..9b4817b 100644
--- a/polymatrix/expressionstate/mixins.py
+++ b/polymatrix/expressionstate/mixins.py
@@ -8,6 +8,7 @@ import dataclasses
 
 from polymatrix.variable.abc import Variable
 from polymatrix.utils.deprecation import deprecated
+from polymatrix.polymatrix.index import MonomialIndex, VariableIndex
 
 from polymatrix.statemonad.mixins import StateCacheMixin
 
@@ -67,12 +68,32 @@ class ExpressionStateMixin(
 
         When a variable is not a scalar multiple indices will be associated to
         the variables, one for each entry.
+
+        See also :py:meth:`get_variable_indices`, :py:meth:`get_monomial_indices`.
         """
         if var not in self.indices:
             raise IndexError(f"There is no variable {var} in this state object.")
 
         yield from range(*self.indices[var])
 
+    def get_indices_as_variable_index(self, var: Variable) -> Iterable[VariableIndex]:
+        """
+        Get all indices associated to a variable, wrapped in a `VariableIndex`.
+
+        See also :py:meth:`get_indices`,
+        :py:class:`polymatrix.polymatrix.index.VariableIndex`.
+        """
+        yield from map(lambda i: VariableIndex(index=i, power=1), self.get_indices(var))
+
+    def get_indices_as_monomial_index(self, var: Variable) -> Iterable[MonomialIndex]:
+        """
+        Get all indices associated to a variable, wrapped in a `MonomialIndex`.
+
+        See also :py:meth:`get_indices`,
+        :py:class:`polymatrix.polymatrix.index.MonomialIndex`.
+        """
+        yield from map(lambda v: MonomialIndex((v,)), self.get_indices_as_variable_index(var))
+
     def get_variable(self, index: int) -> Variable:
         """ Get the variable object from its index. """
         for variable, (start, end) in self.indices.items():