From 1f0fbf6161d401d1aba6ea58bf097af6f88d3069 Mon Sep 17 00:00:00 2001
From: Nao Pross <np@0hm.ch>
Date: Tue, 7 May 2024 18:26:55 +0200
Subject: Add helper methods to retrieve variables from ExpressionState
---
polymatrix/expressionstate/mixins.py | 21 +++++++++++++++++++++
1 file changed, 21 insertions(+)
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():
--
cgit v1.2.1