Improve implementation of DegreeExprMixin

1 files changed, 20 insertions, 28 deletions

diff --git a/polymatrix/expression/mixins/degreeexprmixin.py b/polymatrix/expression/mixins/degreeexprmixin.py
index 5518d0a..522af7c 100644
--- a/polymatrix/expression/mixins/degreeexprmixin.py
+++ b/polymatrix/expression/mixins/degreeexprmixin.py
@@ -5,13 +5,18 @@ from polymatrix.polymatrix.init import init_poly_matrix
 from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin
 from polymatrix.polymatrix.abc import PolyMatrix
 from polymatrix.expressionstate.abc import ExpressionState
+from polymatrix.polymatrix.index import MatrixIndex, PolyMatrixDict, PolyDict, MonomialIndex
 from polymatrix.utils.getstacklines import FrameSummary
 
 
-# NP: IIRC this gets the highest degree.
-# NP: If I am correct consider renaming to MaxDegree
-# NP: (also what is the use case?)
 class DegreeExprMixin(ExpressionBaseMixin):
+    """
+    Elementwise maximum degree.
+
+    The result is a constant matrix with each entry containing the highest
+    degree of the polynomial in the entry of the argument.
+    """
+
     @property
     @abc.abstractmethod
     def underlying(self) -> ExpressionBaseMixin:
@@ -23,32 +28,19 @@ class DegreeExprMixin(ExpressionBaseMixin):
         ...
 
     # overwrites the abstract method of `ExpressionBaseMixin`
-    def apply(
-        self, 
-        state: ExpressionState,
-    ) -> tuple[ExpressionState, PolyMatrix]:
-        state, underlying = self.underlying.apply(state=state)
-
-        poly_matrix_data = {}
-
-        for row in range(underlying.shape[0]):
-            for col in range(underlying.shape[1]):
-
-                polynomial = underlying.get_poly(row, col)
+    def apply(self, state: ExpressionState) -> tuple[ExpressionState, PolyMatrix]:
 
-                if polynomial is None or len(polynomial) == 0:
-                    poly_matrix_data[row, col] =0
+        state, p = self.underlying.apply(state=state)
+        result = PolyMatrixDict.empty()
 
-                else:
-                    def gen_degrees():
-                        for monomial, _ in polynomial.items():
-                            yield sum(count for _, count in monomial)
-                        
-                    poly_matrix_data[row, col] = {tuple(): max(gen_degrees())}
+        for entry, poly in p.entries():
+            max_degree = 0
+            for monomial, coeff in poly.terms():
+                if monomial.degree > max_degree:
+                    max_degree = monomial.degree
 
-        poly_matrix = init_poly_matrix(
-            data=poly_matrix_data,
-            shape=underlying.shape,
-        )
+            result[*entry] = PolyDict({
+                MonomialIndex.constant(): max_degree
+            })
 
-        return state, poly_matrix   
+        return state, init_poly_matrix(data=result, shape=p.shape)