diff options
| -rw-r--r-- | polymatrix/expressionstate/mixins.py | 21 |
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(): |