add 'degree' function that returns the degrees of each polynomial in the matrix

1 files changed, 45 insertions, 3 deletions

diff --git a/polymatrix/__init__.py b/polymatrix/__init__.py
index d0673e4..8f5d28e 100644
--- a/polymatrix/__init__.py
+++ b/polymatrix/__init__.py
@@ -21,6 +21,7 @@ from polymatrix.statemonad.init.initstatemonad import init_state_monad
 from polymatrix.statemonad.mixins.statemonadmixin import StateMonadMixin
 from polymatrix.expression.utils.monomialtoindex import monomial_to_index
 from polymatrix.expressionstate.init.initexpressionstate import init_expression_state as original_init_expression_state
+from polymatrix.statemonad.statemonad import StateMonad
 
 
 def init_expression_state():
@@ -33,10 +34,12 @@ def from_sympy(
         init_from_sympy_expr(data)
     )
 
-def from_(
-    data: tuple[tuple[float]],
+def from_state_monad(
+    data: StateMonad,
 ):
-    return from_sympy(data)
+    return init_expression(
+        data.flat_map(lambda inner_data: init_from_sympy_expr(inner_data)),
+    )
 
 def from_polymatrix(
     polymatrix: PolyMatrix,
@@ -45,6 +48,11 @@ def from_polymatrix(
         init_from_terms_expr(polymatrix)
     )
 
+def from_(
+    data: tuple[tuple[float]],
+):
+    return from_sympy(data)
+
 def v_stack(
     expressions: tuple[Expression],
 ):
@@ -588,6 +596,40 @@ def to_constant_repr(
 
     return init_state_monad(func)
 
+
+def degrees(
+    expr: Expression,
+    variables: Expression,
+) -> StateMonadMixin[ExpressionState, np.ndarray]:
+
+    def func(state: ExpressionState):
+        state, underlying = expr.apply(state)
+        state, variable_indices = get_variable_indices(state, variables)
+
+        def gen_rows():
+            for row in range(underlying.shape[0]):
+                def gen_cols():
+                    for col in range(underlying.shape[1]):
+                        
+                        def gen_degrees():
+                            polynomial = underlying.get_poly(row, col)
+
+                            if polynomial is None:
+                                yield 0
+
+                            else:
+                                for monomial, _ in polynomial.items():
+                                    yield sum(count for var, count in monomial if var in variable_indices)
+
+                        yield tuple(set(gen_degrees()))
+
+                yield tuple(gen_cols())
+
+        return state, tuple(gen_rows())
+
+    return init_state_monad(func)
+
+
 def to_sympy_repr(
     expr: Expression,
 ) -> StateMonadMixin[ExpressionState, sympy.Expr]: