add max_degree, max and filter operator

47 files changed, 1214 insertions, 627 deletions

diff --git a/polymatrix/__init__.py b/polymatrix/__init__.py
index aa96d21..58729d2 100644
--- a/polymatrix/__init__.py
+++ b/polymatrix/__init__.py
@@ -298,9 +298,37 @@ def shape(
 
 
 @dataclasses.dataclass
+class MatrixBuffer:
+    data: dict[int, ...]
+    n_row: int
+    n_param: int
+
+    def add_buffer(self, index: int):
+        if index <= 1:
+            buffer = np.zeros((self.n_row, self.n_param**index), dtype=np.double)
+
+        else:
+            buffer = scipy.sparse.dok_array((self.n_row, self.n_param**index), dtype=np.double)
+            
+        self.data[index] = buffer
+
+    def add(self, row: int, col: int, index: int, value: float):
+        if index not in self.data:
+            self.add_buffer(index)
+
+        self.data[index][row, col] = value
+
+    def __getitem__(self, key):
+        if key not in self.data:
+            self.add_buffer(key)
+
+        return self.data[key]
+
+
+@dataclasses.dataclass
 class MatrixEquations:
-    matrix_equations: tuple[dict[int, np.ndarray], ...]
-    auxillary_matrix_equations: typing.Optional[dict[int, np.ndarray]]
+    matrix_equations: tuple[MatrixBuffer, ...]
+    auxillary_matrix_equations: typing.Optional[MatrixBuffer]
     variable_index: tuple[int, ...]
     state: ExpressionState
 
@@ -322,10 +350,6 @@ class MatrixEquations:
                 else:
                     yield index, scipy.sparse.vstack(tuple(gen_matrices(index)))
 
-        # matrix_1 = np.vstack(tuple(gen_matrices(0)))
-        # matrix_2 = np.vstack(tuple(gen_matrices(1)))
-        # matrix_3 = scipy.sparse.vstack(tuple(gen_matrices(2)))
-
         return dict(gen_matrices())
 
     def get_value(self, variable, value):
@@ -349,10 +373,52 @@ class MatrixEquations:
         vec[value_index] = value
         return vec
 
+    def get_func(self, eq_idx: int):
+        equations = self.matrix_equations[eq_idx].data
+        max_idx = max(equations.keys())
+
+        if 2 <= max_idx:
+            def func(x: np.ndarray) -> np.ndarray:
+                def acc_x_powers(acc, _):
+                    next = (acc @ x.T).reshape(-1, 1)
+                    return next
+
+                x_powers = tuple(itertools.accumulate(
+                    range(max_idx - 1),
+                    acc_x_powers,
+                    initial=x,
+                ))[1:]
+
+                def gen_value():
+                    for idx, equation in equations.items():
+                        if idx == 0:
+                            yield equation
+
+                        elif idx == 1:
+                            yield equation @ x
+
+                        else:
+                            yield equation @ x_powers[idx-2]
+
+                return sum(gen_value())
+
+        else:
+            def func(x: np.ndarray) -> np.ndarray:
+                def gen_value():
+                    for idx, equation in equations.items():
+                        if idx == 0:
+                            yield equation
+
+                        else:
+                            yield equation @ x
+
+                return sum(gen_value())
+
+        return func
 
 def to_matrix_equations(
     expressions: tuple[Expression],
-    variables: Expression = None,
+    variables: Expression,
 ) -> StateMonadMixin[ExpressionState, tuple[tuple[tuple[np.ndarray, ...], ...], tuple[int, ...]]]:
 
     if isinstance(expressions, Expression):
@@ -412,33 +478,6 @@ def to_matrix_equations(
 
         n_param = len(ordered_variable_index)
 
-        @dataclasses.dataclass
-        class MatrixBuffer:
-            data: dict
-            n_row: int
-            n_param: int
-
-            def add_buffer(self, index: int):
-                if index <= 1:
-                    buffer = np.zeros((self.n_row, self.n_param**index), dtype=np.double)
-
-                else:
-                    buffer = scipy.sparse.dok_array((self.n_row, self.n_param**index), dtype=np.double)
-                    
-                self.data[index] = buffer
-
-            def add(self, row: int, col: int, index: int, value: float):
-                if index not in self.data:
-                    self.add_buffer(index)
-
-                self.data[index][row, col] = value
-
-            def __getitem__(self, key):
-                if key not in self.data:
-                    self.add_buffer(key)
-
-                return self.data[key]
-
         def gen_underlying_matrices():
             for underlying in underlying_list:
 
@@ -461,9 +500,13 @@ def to_matrix_equations(
                             new_monomial = tuple(variable_index_map[var] for var, count in monomial for _ in range(count))
                         except KeyError:
                             raise KeyError(f'{monomial=} is incompatible with {variable_index_map=}')
-                        col = monomial_to_index(n_param, new_monomial)
 
-                        buffer.add(row, col, sum(count for _, count in monomial), value)
+                        cols = monomial_to_index(n_param, new_monomial)
+
+                        col_value = value / len(cols)
+
+                        for col in cols:
+                            buffer.add(row, col, sum(count for _, count in monomial), col_value)
                 
                 # yield A, B, C
                 yield buffer
@@ -493,11 +536,14 @@ def to_matrix_equations(
 
             for row, (key, monomial_terms) in enumerate(auxillary_equations):
                 for monomial, value in monomial_terms.items():
-                    # new_monomial = tuple(variable_index_map[var] for var, _ in monomial)
                     new_monomial = tuple(variable_index_map[var] for var, count in monomial for _ in range(count))
-                    col = monomial_to_index(n_param, new_monomial)
 
-                    buffer.add(row, col, sum(count for _, count in monomial), value)
+                    cols = monomial_to_index(n_param, new_monomial)
+
+                    col_value = value / len(cols)
+
+                    for col in cols:
+                        buffer.add(row, col, sum(count for _, count in monomial), col_value)
 
             auxillary_matrix_equations = buffer
 
diff --git a/polymatrix/expression/filterexpr.py b/polymatrix/expression/filterexpr.py
new file mode 100644
index 0000000..e7ca235
--- /dev/null
+++ b/polymatrix/expression/filterexpr.py
@@ -0,0 +1,4 @@
+from polymatrix.expression.mixins.filterexprmixin import FilterExprMixin 
+
+class FilterExpr(FilterExprMixin):
+	pass
diff --git a/polymatrix/expression/impl/parametrizetermsexprimpl.py b/polymatrix/expression/impl/filterexprimpl.py
index 32beac2..547ad9e 100644
--- a/polymatrix/expression/impl/parametrizetermsexprimpl.py
+++ b/polymatrix/expression/impl/filterexprimpl.py
@@ -1,10 +1,10 @@
 import dataclass_abc
-from polymatrix.expression.parametrizetermsexpr import ParametrizeTermsExpr 
+from polymatrix.expression.filterexpr import FilterExpr 
 
 from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin 
 
 @dataclass_abc.dataclass_abc(frozen=True)
-class ParametrizeTermsExprImpl(ParametrizeTermsExpr):
-	# name: str
+class FilterExprImpl(FilterExpr):
 	underlying: ExpressionBaseMixin
-	# variables: tuple
+	predicator: ExpressionBaseMixin
+	inverse: bool == None
diff --git a/polymatrix/expression/impl/maxdegreeexprimpl.py b/polymatrix/expression/impl/maxdegreeexprimpl.py
new file mode 100644
index 0000000..5327659
--- /dev/null
+++ b/polymatrix/expression/impl/maxdegreeexprimpl.py
@@ -0,0 +1,8 @@
+import dataclass_abc
+from polymatrix.expression.maxdegreeexpr import MaxDegreeExpr 
+
+from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin 
+
+@dataclass_abc.dataclass_abc(frozen=True)
+class MaxDegreeExprImpl(MaxDegreeExpr):
+	underlying: ExpressionBaseMixin
diff --git a/polymatrix/expression/impl/maxexprimpl.py b/polymatrix/expression/impl/maxexprimpl.py
new file mode 100644
index 0000000..ae2fe90
--- /dev/null
+++ b/polymatrix/expression/impl/maxexprimpl.py
@@ -0,0 +1,8 @@
+import dataclass_abc
+from polymatrix.expression.maxexpr import MaxExpr 
+
+from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin 
+
+@dataclass_abc.dataclass_abc(frozen=True)
+class MaxExprImpl(MaxExpr):
+	underlying: ExpressionBaseMixin
diff --git a/polymatrix/expression/impl/sosmonomialsexprimpl.py b/polymatrix/expression/impl/sosmonomialsexprimpl.py
index 9e43b23..400c2ae 100644
--- a/polymatrix/expression/impl/sosmonomialsexprimpl.py
+++ b/polymatrix/expression/impl/sosmonomialsexprimpl.py
@@ -1,9 +1,9 @@
 import dataclass_abc
-from polymatrix.expression.sosmonomialsexpr import SOSMonomialsExpr 
+from polymatrix.expression.quadraticmonomialsexpr import QuadraticMonomialsExpr 
 
 from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin 
 
 @dataclass_abc.dataclass_abc(frozen=True)
-class SOSMonomialsExprImpl(SOSMonomialsExpr):
+class SOSMonomialsExprImpl(QuadraticMonomialsExpr):
 	underlying: ExpressionBaseMixin
 	variables: tuple
diff --git a/polymatrix/expression/impl/truncateexprimpl.py b/polymatrix/expression/impl/truncateexprimpl.py
index 01c47d0..e4d10ae 100644
--- a/polymatrix/expression/impl/truncateexprimpl.py
+++ b/polymatrix/expression/impl/truncateexprimpl.py
@@ -7,4 +7,4 @@ from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin
 class TruncateExprImpl(TruncateExpr):
 	underlying: ExpressionBaseMixin
 	variables: ExpressionBaseMixin
-	degree: int
+	degrees: tuple[int]
diff --git a/polymatrix/expression/init/initfilterexpr.py b/polymatrix/expression/init/initfilterexpr.py
new file mode 100644
index 0000000..c47036f
--- /dev/null
+++ b/polymatrix/expression/init/initfilterexpr.py
@@ -0,0 +1,17 @@
+from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin
+from polymatrix.expression.impl.filterexprimpl import FilterExprImpl
+
+
+def init_filter_expr(
+		underlying: ExpressionBaseMixin,
+		predicator: ExpressionBaseMixin,
+		inverse: bool == None,
+):
+	if inverse is None:
+		inverse = False
+
+	return FilterExprImpl(
+		underlying=underlying,
+		predicator=predicator,
+		inverse=inverse,
+)
diff --git a/polymatrix/expression/init/initmaxdegreeexpr.py b/polymatrix/expression/init/initmaxdegreeexpr.py
new file mode 100644
index 0000000..60fc40a
--- /dev/null
+++ b/polymatrix/expression/init/initmaxdegreeexpr.py
@@ -0,0 +1,10 @@
+from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin
+from polymatrix.expression.impl.maxdegreeexprimpl import MaxDegreeExprImpl
+
+
+def init_max_degree_expr(
+		underlying: ExpressionBaseMixin,
+):
+	return MaxDegreeExprImpl(
+		underlying=underlying,
+)
diff --git a/polymatrix/expression/init/initmaxexpr.py b/polymatrix/expression/init/initmaxexpr.py
new file mode 100644
index 0000000..b0b1e4d
--- /dev/null
+++ b/polymatrix/expression/init/initmaxexpr.py
@@ -0,0 +1,10 @@
+from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin
+from polymatrix.expression.impl.maxexprimpl import MaxExprImpl
+
+
+def init_max_expr(
+		underlying: ExpressionBaseMixin,
+):
+	return MaxExprImpl(
+		underlying=underlying,
+)
diff --git a/polymatrix/expression/init/initparametrizetermsexpr.py b/polymatrix/expression/init/initparametrizetermsexpr.py
deleted file mode 100644
index 6e2a446..0000000
--- a/polymatrix/expression/init/initparametrizetermsexpr.py
+++ /dev/null
@@ -1,14 +0,0 @@
-from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin
-from polymatrix.expression.impl.parametrizetermsexprimpl import ParametrizeTermsExprImpl
-
-
-def init_parametrize_terms_expr(
-		# name: str,
-		underlying: ExpressionBaseMixin,
-		# variables: tuple,
-):
-	return ParametrizeTermsExprImpl(
-		# name=name,
-		underlying=underlying,
-		# variables=variables,
-)
diff --git a/polymatrix/expression/init/inittruncateexpr.py b/polymatrix/expression/init/inittruncateexpr.py
index 818981e..060ebaf 100644
--- a/polymatrix/expression/init/inittruncateexpr.py
+++ b/polymatrix/expression/init/inittruncateexpr.py
@@ -5,10 +5,13 @@ from polymatrix.expression.impl.truncateexprimpl import TruncateExprImpl
 def init_truncate_expr(
 		underlying: ExpressionBaseMixin,
 		variables: ExpressionBaseMixin,
-		degree: int,
+		degrees: tuple[int],
 ):
+	if isinstance(degrees, int):
+		degrees = (degrees,)
+
 	return TruncateExprImpl(
 		underlying=underlying,
 		variables=variables,
-		degree=degree,
+		degrees=degrees,
 )
diff --git a/polymatrix/expression/maxdegreeexpr.py b/polymatrix/expression/maxdegreeexpr.py
new file mode 100644
index 0000000..ec8ad47
--- /dev/null
+++ b/polymatrix/expression/maxdegreeexpr.py
@@ -0,0 +1,4 @@
+from polymatrix.expression.mixins.maxdegreeexprmixin import MaxDegreeExprMixin 
+
+class MaxDegreeExpr(MaxDegreeExprMixin):
+	pass
diff --git a/polymatrix/expression/maxexpr.py b/polymatrix/expression/maxexpr.py
new file mode 100644
index 0000000..dc0996c
--- /dev/null
+++ b/polymatrix/expression/maxexpr.py
@@ -0,0 +1,4 @@
+from polymatrix.expression.mixins.maxexprmixin import MaxExprMixin 
+
+class MaxExpr(MaxExprMixin):
+	pass
diff --git a/polymatrix/expression/mixins/evalexprmixin.py b/polymatrix/expression/mixins/evalexprmixin.py
index 7fa149e..0541e3d 100644
--- a/polymatrix/expression/mixins/evalexprmixin.py
+++ b/polymatrix/expression/mixins/evalexprmixin.py
@@ -32,7 +32,6 @@ class EvalExprMixin(ExpressionBaseMixin):
         state: ExpressionState,
     ) -> tuple[ExpressionState, PolyMatrix]:
         state, underlying = self.underlying.apply(state=state)
-
         state, variable_indices = get_variable_indices(state, self.variables)
 
         if len(self.values) == 1:
@@ -63,7 +62,7 @@ class EvalExprMixin(ExpressionBaseMixin):
 
                         if variable in variable_indices:
                             index = variable_indices.index(variable)
-                            new_value = value * values[index] * count
+                            new_value = value * values[index]**count
                             return new_monomial, new_value
 
                         else:
diff --git a/polymatrix/expression/mixins/expressionmixin.py b/polymatrix/expression/mixins/expressionmixin.py
index e75973e..ff2a146 100644
--- a/polymatrix/expression/mixins/expressionmixin.py
+++ b/polymatrix/expression/mixins/expressionmixin.py
@@ -13,6 +13,7 @@ from polymatrix.expression.init.initdivergenceexpr import init_divergence_expr
 from polymatrix.expression.init.initdivisionexpr import init_division_expr
 from polymatrix.expression.init.initelemmultexpr import init_elem_mult_expr
 from polymatrix.expression.init.initevalexpr import init_eval_expr
+from polymatrix.expression.init.initfilterexpr import init_filter_expr
 from polymatrix.expression.init.initlinearinexpr import init_linear_in_expr
 from polymatrix.expression.init.initlinearinmonomialsinexpr import init_linear_in_monomials_in
 from polymatrix.expression.init.initfromarrayexpr import init_from_array_expr
@@ -20,6 +21,8 @@ from polymatrix.expression.init.initgetitemexpr import init_get_item_expr
 from polymatrix.expression.init.initlinearmatrixinexpr import init_linear_matrix_in_expr
 from polymatrix.expression.init.initlinearmonomialsexpr import init_linear_monomials_expr
 from polymatrix.expression.init.initmatrixmultexpr import init_matrix_mult_expr
+from polymatrix.expression.init.initmaxdegreeexpr import init_max_degree_expr
+from polymatrix.expression.init.initmaxexpr import init_max_expr
 from polymatrix.expression.init.initparametrizeexpr import init_parametrize_expr
 from polymatrix.expression.init.initquadraticinexpr import init_quadratic_in_expr
 from polymatrix.expression.init.initrepmatexpr import init_rep_mat_expr
@@ -230,6 +233,20 @@ class ExpressionMixin(
             ),
         )
 
+    def filter(
+        self,
+        predicator: 'ExpressionMixin',
+        inverse: bool = None,
+    ) -> 'ExpressionMixin':
+        return dataclasses.replace(
+            self, 
+            underlying=init_filter_expr(
+                underlying=self.underlying,
+                predicator=predicator,
+                inverse=inverse,
+            ),
+        )
+
     def linear_matrix_in(self, variable: 'ExpressionMixin') -> 'ExpressionMixin':
         return dataclasses.replace(
             self, 
@@ -276,6 +293,22 @@ class ExpressionMixin(
             ),
         )
 
+    def max(self) -> 'ExpressionMixin':
+        return dataclasses.replace(
+            self, 
+            underlying=init_max_expr(
+                underlying=self.underlying,
+            ),
+        )
+
+    def max_degree(self) -> 'ExpressionMixin':
+        return dataclasses.replace(
+            self, 
+            underlying=init_max_degree_expr(
+                underlying=self.underlying,
+            ),
+        )
+
     def parametrize(self, name: str = None) -> 'ExpressionMixin':
         return dataclasses.replace(
             self, 
@@ -414,12 +447,12 @@ class ExpressionMixin(
             ),
         )
 
-    def truncate(self, variables: tuple, degree: int):
+    def truncate(self, variables: tuple, degrees: tuple[int]):
         return dataclasses.replace(
             self, 
             underlying=init_truncate_expr(
                 underlying=self.underlying,
                 variables=variables,
-                degree=degree,
+                degrees=degrees,
             ),
         )
diff --git a/polymatrix/expression/mixins/filterexprmixin.py b/polymatrix/expression/mixins/filterexprmixin.py
new file mode 100644
index 0000000..ad1253a
--- /dev/null
+++ b/polymatrix/expression/mixins/filterexprmixin.py
@@ -0,0 +1,65 @@
+
+import abc
+import collections
+import math
+import typing
+import dataclass_abc
+
+from polymatrix.expression.init.initpolymatrix import init_poly_matrix
+from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin
+from polymatrix.expression.polymatrix import PolyMatrix
+from polymatrix.expression.expressionstate import ExpressionState
+
+
+class FilterExprMixin(ExpressionBaseMixin):
+    @property
+    @abc.abstractmethod
+    def underlying(self) -> ExpressionBaseMixin:
+        ...
+
+    @property
+    @abc.abstractmethod
+    def predicator(self) -> ExpressionBaseMixin:
+        ...
+
+    @property
+    @abc.abstractmethod
+    def inverse(self) -> bool:
+        ...
+
+    # overwrites abstract method of `ExpressionBaseMixin`
+    def apply(
+        self, 
+        state: ExpressionState,
+    ) -> tuple[ExpressionState, PolyMatrix]:
+        state, underlying = self.underlying.apply(state=state)
+        state, predicator = self.predicator.apply(state=state)
+
+        assert underlying.shape[1] == 1
+        assert predicator.shape[1] == 1
+        assert underlying.shape[0] == predicator.shape[0]
+
+        terms = {}
+        row_index = 0
+
+        for row in range(underlying.shape[0]):
+
+            try:
+                underlying_terms = underlying.get_poly(row, 0)
+            except KeyError:
+                continue
+
+            predicator_value = predicator.get_poly(row, 0)[tuple()]
+
+            assert isinstance(predicator_value, int)
+
+            if (predicator_value != 0) is not self.inverse:
+                terms[row_index, 0] = underlying_terms
+                row_index += 1
+
+        poly_matrix = init_poly_matrix(
+            terms=terms,
+            shape=(row_index, 1),
+        )
+
+        return state, poly_matrix   
diff --git a/polymatrix/expression/mixins/maxdegreeexprmixin.py b/polymatrix/expression/mixins/maxdegreeexprmixin.py
new file mode 100644
index 0000000..a1c9b1e
--- /dev/null
+++ b/polymatrix/expression/mixins/maxdegreeexprmixin.py
@@ -0,0 +1,44 @@
+
+import abc
+
+from polymatrix.expression.init.initpolymatrix import init_poly_matrix
+from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin
+from polymatrix.expression.polymatrix import PolyMatrix
+from polymatrix.expression.expressionstate import ExpressionState
+
+
+class MaxDegreeExprMixin(ExpressionBaseMixin):
+    @property
+    @abc.abstractmethod
+    def underlying(self) -> ExpressionBaseMixin:
+        ...
+
+    # overwrites abstract method of `ExpressionBaseMixin`
+    def apply(
+        self, 
+        state: ExpressionState,
+    ) -> tuple[ExpressionState, PolyMatrix]:
+        state, underlying = self.underlying.apply(state=state)
+
+        terms = {}
+
+        for row in range(underlying.shape[0]):
+            for col in range(underlying.shape[1]):
+
+                try:
+                    underlying_terms = underlying.get_poly(row, col)
+                except KeyError:
+                    continue
+
+                def gen_degrees():
+                    for monomial, _ in underlying_terms.items():
+                        yield sum(count for _, count in monomial)
+                    
+                terms[row, col] = {tuple(): max(gen_degrees())}
+
+        poly_matrix = init_poly_matrix(
+            terms=terms,
+            shape=underlying.shape,
+        )
+
+        return state, poly_matrix   
diff --git a/polymatrix/expression/mixins/maxexprmixin.py b/polymatrix/expression/mixins/maxexprmixin.py
new file mode 100644
index 0000000..6070f73
--- /dev/null
+++ b/polymatrix/expression/mixins/maxexprmixin.py
@@ -0,0 +1,48 @@
+
+import abc
+
+from polymatrix.expression.init.initpolymatrix import init_poly_matrix
+from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin
+from polymatrix.expression.mixins.expressionstatemixin import ExpressionStateMixin
+from polymatrix.expression.mixins.polymatrixmixin import PolyMatrixMixin
+
+
+class MaxExprMixin(ExpressionBaseMixin):
+    @property
+    @abc.abstractclassmethod
+    def underlying(self) -> ExpressionBaseMixin:
+        ...
+
+    # overwrites abstract method of `ExpressionBaseMixin`
+    def apply(
+        self, 
+        state: ExpressionStateMixin,
+    ) -> tuple[ExpressionStateMixin, PolyMatrixMixin]:
+
+        state, underlying = self.underlying.apply(state)
+
+        terms = {}
+
+        for row in range(underlying.shape[0]):
+
+            def gen_values():
+                for col in range(underlying.shape[1]):
+
+                    try:
+                        underlying_terms = underlying.get_poly(row, col)
+                    except KeyError:
+                        continue
+
+                    yield underlying_terms[tuple()]
+
+            values = tuple(gen_values())
+
+            if 0 < len(values):
+                terms[row, 0] = {tuple(): max(values)}
+
+        poly_matrix = init_poly_matrix(
+            terms=terms,
+            shape=(underlying.shape[0], 1),
+        )
+
+        return state, poly_matrix
diff --git a/polymatrix/expression/mixins/parametrizetermsexprmixin.py b/polymatrix/expression/mixins/parametrizetermsexprmixin.py
deleted file mode 100644
index 51244d4..0000000
--- a/polymatrix/expression/mixins/parametrizetermsexprmixin.py
+++ /dev/null
@@ -1,101 +0,0 @@
-
-import abc
-import dataclasses
-
-from polymatrix.expression.init.initpolymatrix import init_poly_matrix
-from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin
-from polymatrix.expression.mixins.expressionstatemixin import ExpressionStateMixin
-from polymatrix.expression.mixins.polymatrixmixin import PolyMatrixMixin
-from polymatrix.expression.utils.getvariableindices import get_variable_indices
-from polymatrix.expression.utils.mergemonomialindices import merge_monomial_indices
-
-
-class ParametrizeTermsExprMixin(ExpressionBaseMixin):
-    # @property
-    # @abc.abstractclassmethod
-    # def name(self) -> str:
-    #     ...
-
-    @property
-    @abc.abstractclassmethod
-    def underlying(self) -> ExpressionBaseMixin:
-        ...
-
-    # @property
-    # @abc.abstractclassmethod
-    # def variables(self) -> tuple:
-    #     ...
-
-    # overwrites abstract method of `ExpressionBaseMixin`
-    def apply(
-        self, 
-        state: ExpressionStateMixin,
-    ) -> tuple[ExpressionStateMixin, PolyMatrixMixin]:
-
-        # cache polymatrix to not re-parametrize at every apply call
-        if self in state.cache:
-            return state, state.cache[self]
-
-        state, underlying = self.underlying.apply(state)
-
-        # state, variable_indices = get_variable_indices(state, self.variables)
-
-        start_index = state.n_param
-        terms = {}
-
-        for row in range(underlying.shape[0]):
-            for col in range(underlying.shape[1]):
-
-                try:
-                    underlying_terms = underlying.get_poly(row, col)
-                except KeyError:
-                    continue
-
-                # def gen_x_monomial_terms():
-                #     for monomial, value in underlying_terms.items():
-                #         x_monomial = tuple((var_idx, count) for var_idx, count in monomial if var_idx in variable_indices)
-                #         yield monomial, x_monomial, value
-
-                # x_monomial_terms = tuple(gen_x_monomial_terms())
-
-                # # use sort to make the parametrization predictable
-                # # x_monomials = (0, 1), then the parametrization is (2, 3)
-                # collected_terms = tuple(sorted(
-                #     set((x_monomial for _, x_monomial, _ in x_monomial_terms)), 
-                #     key=lambda m: (sum(count for _, count in m), m)),
-                # )
-                collected_terms = tuple(underlying_terms.keys())
-
-                # print(tuple(x_monomial for _, x_monomial, _ in x_monomial_terms))
-                # print(x_monomial_terms)
-
-                terms_row_col = {}
-
-                for monomial, value in underlying_terms.items():
-
-                    param_index = start_index + collected_terms.index(monomial)
-                        
-                    monomial_with_param = merge_monomial_indices(monomial, ((param_index, 1),))
-
-                    terms_row_col[monomial_with_param] = value
-
-                terms[row, col] = terms_row_col
-
-                start_index += len(collected_terms)
-
-        state = state.register(
-            key=self, 
-            n_param=start_index - state.n_param,
-        )
-
-        poly_matrix = init_poly_matrix(
-            terms=terms,
-            shape=underlying.shape,
-        )
-
-        state = dataclasses.replace(
-            state, 
-            cache=state.cache | {self: poly_matrix},
-        )
-
-        return state, poly_matrix
diff --git a/polymatrix/expression/mixins/polymatrixmixin.py b/polymatrix/expression/mixins/polymatrixmixin.py
index 89e05cb..bea80de 100644
--- a/polymatrix/expression/mixins/polymatrixmixin.py
+++ b/polymatrix/expression/mixins/polymatrixmixin.py
@@ -1,12 +1,6 @@
 import abc
-import dataclasses
-import numpy as np
-import scipy.sparse
 import typing
 
-from polymatrix.expression.mixins.expressionstatemixin import ExpressionStateMixin
-from polymatrix.utils import monomial_to_index
-
 
 class PolyMatrixMixin(abc.ABC):
     @property
diff --git a/polymatrix/expression/mixins/quadraticinexprmixin.py b/polymatrix/expression/mixins/quadraticinexprmixin.py
index dbfb76d..7e08d7c 100644
--- a/polymatrix/expression/mixins/quadraticinexprmixin.py
+++ b/polymatrix/expression/mixins/quadraticinexprmixin.py
@@ -49,8 +49,15 @@ class QuadraticInExprMixin(ExpressionBaseMixin):
 
             left, right = split_monomial_indices(x_monomial)
 
-            row = sos_monomials.index(left)
-            col = sos_monomials.index(right)
+            try:
+                col = sos_monomials.index(left)
+            except ValueError:
+                raise ValueError(f'{left=} not in {sos_monomials=}')
+
+            try:
+                row = sos_monomials.index(right)
+            except ValueError:
+                raise ValueError(f'{right=} not in {sos_monomials=}')
 
             monomial_terms = terms[row, col]
 
diff --git a/polymatrix/expression/mixins/truncateexprmixin.py b/polymatrix/expression/mixins/truncateexprmixin.py
index 7d0927a..badc01e 100644
--- a/polymatrix/expression/mixins/truncateexprmixin.py
+++ b/polymatrix/expression/mixins/truncateexprmixin.py
@@ -10,6 +10,7 @@ from polymatrix.expression.expressionstate import ExpressionState
 from polymatrix.expression.utils.getvariableindices import get_variable_indices
 
 
+# replace by filter operation?
 class TruncateExprMixin(ExpressionBaseMixin):
     @property
     @abc.abstractmethod
@@ -23,7 +24,7 @@ class TruncateExprMixin(ExpressionBaseMixin):
 
     @property
     @abc.abstractmethod
-    def degree(self) -> int:
+    def degrees(self) -> tuple[int]:
         ...
 
     # overwrites abstract method of `ExpressionBaseMixin`
@@ -32,7 +33,6 @@ class TruncateExprMixin(ExpressionBaseMixin):
         state: ExpressionState,
     ) -> tuple[ExpressionState, PolyMatrix]:
         state, underlying = self.underlying.apply(state=state)
-
         state, variable_indices = get_variable_indices(state, self.variables)
 
         terms = {}
@@ -51,7 +51,7 @@ class TruncateExprMixin(ExpressionBaseMixin):
 
                     degree = sum((count for var_idx, count in monomial if var_idx in variable_indices))
 
-                    if degree <= self.degree:
+                    if degree in self.degrees:
                         terms_row_col[monomial] = value
                 
                 terms[row, col] = terms_row_col
diff --git a/polymatrix/expression/parametrizetermsexpr.py b/polymatrix/expression/parametrizetermsexpr.py
deleted file mode 100644
index 18b4de2..0000000
--- a/polymatrix/expression/parametrizetermsexpr.py
+++ /dev/null
@@ -1,4 +0,0 @@
-from polymatrix.expression.mixins.parametrizetermsexprmixin import ParametrizeTermsExprMixin 
-
-class ParametrizeTermsExpr(ParametrizeTermsExprMixin):
-	pass
diff --git a/polymatrix/expression/sosmonomialsexpr.py b/polymatrix/expression/quadraticmonomialsexpr.py
index 63dfcef..77753a5 100644
--- a/polymatrix/expression/sosmonomialsexpr.py
+++ b/polymatrix/expression/quadraticmonomialsexpr.py
@@ -1,4 +1,4 @@
 from polymatrix.expression.mixins.quadraticmonomialsexprmixin import QuadraticMonomialsExprMixin 
 
-class SOSMonomialsExpr(QuadraticMonomialsExprMixin):
+class QuadraticMonomialsExpr(QuadraticMonomialsExprMixin):
 	pass
diff --git a/polymatrix/expression/utils/subtractmonomialindices.py b/polymatrix/expression/utils/subtractmonomialindices.py
index a8aef1c..766997a 100644
--- a/polymatrix/expression/utils/subtractmonomialindices.py
+++ b/polymatrix/expression/utils/subtractmonomialindices.py
@@ -10,9 +10,4 @@ def subtract_monomial_indices(m1, m2):
         if m1_dict[index] == 0:
             del m1_dict[index]
 
-    # return tuple(sorted(
-    #     m1_dict.items(), 
-    #     key=lambda m: m[0],
-    # ))
-
     return sort_monomial_indices(m1_dict.items())
diff --git a/polymatrix/sympyutils.py b/polymatrix/sympyutils.py
index 3e727ed..8fd74b2 100644
--- a/polymatrix/sympyutils.py
+++ b/polymatrix/sympyutils.py
@@ -4,8 +4,6 @@ import numpy as np
 import scipy.sparse
 import sympy
 
-from polymatrix.utils import variable_powers_to_index
-
 
 def poly_to_data_coord(poly_list, x, degree = None):
     """
@@ -87,81 +85,3 @@ def poly_to_matrix(poly_list, x, power = None):
                     yield degree, sparse_array.tocsr()
             
     return dict(gen_power_mat())
-
-
-# def poly_to_data_coord(poly_list, x, degree = None):
-#     """
-#     poly_list = [
-#         poly(x1*x3**2, x)
-#     ]
-#     power: up to which power 
-#     """
-
-#     sympy_poly_list = tuple(tuple(sympy.poly(p, x) for p in inner_poly_list) for inner_poly_list in poly_list)
-
-#     if degree is None:
-#         degree = max(degree for inner_poly_list in sympy_poly_list for poly in inner_poly_list for degree in poly.degree_list())
-
-#     # def all_equal(iterable):
-#     #     g = itertools.groupby(iterable)
-#     #     return next(g, True) and not next(g, False)
-
-#     # assert all_equal((p.gens for p in poly_list)), 'all polynomials need to have identical generators'
-
-#     def gen_power_mat():
-
-#         # for all powers generate a matrix
-#         for current_degree in range(degree + 1):
-        
-#             def gen_value_index():
-
-#                 # each polynomial defines a row in the matrix
-#                 for poly_row, inner_poly_list in enumerate(sympy_poly_list):
-
-#                     for poly_col, p in enumerate(inner_poly_list):
-
-#                         # a5 x1 x3**2 -> c=a5, m=(1, 0, 2)
-#                         for c, m in zip(p.coeffs(), p.monoms()):
-
-#                             if sum(m) == current_degree:
-
-#                                 index = variable_powers_to_index(m)
-#                                 yield (poly_row, poly_col, index), c
-                    
-#             data = dict(gen_value_index()) | collections.defaultdict(float)
-
-#             if len(data) > 0:
-#                 yield current_degree, data
-
-#     return dict(gen_power_mat())
-
-
-# def poly_to_matrix(poly_list, x, power = None):
-#     """
-    
-#     """
-
-#     data_coord_dict = poly_to_data_coord(poly_list, x, power)
-
-#     # n_free_symbols = len(poly_list[0][0].gens)
-#     n_free_symbols = len(x)
-
-#     def gen_power_mat():
-
-#         # for all powers generate a matrix
-#         for current_degree, data_coord in data_coord_dict:
-            
-#             # empty matrix
-#             shape = (len(poly_list), n_free_symbols**current_degree)
-#             # m = np.zeros((len(poly_list), n_free_symbols**current_power))
-
-#             # fill matrix
-#             if len(data_coord) == 0:
-#                 yield np.zeros((len(poly_list), n_free_symbols**current_degree))
-
-#             else:
-#                 rows, cols, data = list(zip(*data_coord))
-
-#                 yield scipy.sparse.coo_matrix((data, (rows, cols)), dtype=np.double, shape=shape)
-            
-#     return list(gen_power_mat())
diff --git a/polymatrix/utils.py b/polymatrix/utils.py
index 2595bf6..2234e4a 100644
--- a/polymatrix/utils.py
+++ b/polymatrix/utils.py
@@ -1,38 +1,43 @@
+import itertools
 import scipy.special
 
 
 def monomial_to_index(n_var, monomial):
-    return sum(idx*(n_var**level) for level, idx in enumerate(monomial))
+    # return sum(idx*(n_var**level) for level, idx in enumerate(monomial))
+
+    monomial_perm = itertools.permutations(monomial)
+    
+    return set(sum(idx*(n_var**level) for level, idx in enumerate(monomial)) for monomial in monomial_perm)
 
 
-def variable_to_index(n_var, combination):
-    """
-    example:
-        expr = x1 * x3**2
-        n_var = 3
-        combination = (2, 2, 0)
+# def variable_to_index(n_var, combination):
+#     """
+#     example:
+#         expr = x1 * x3**2
+#         n_var = 3
+#         combination = (2, 2, 0)
         
-        returns 5
-    """
+#         returns 5
+#     """
     
-    def gen_index_sum():
-        for k, x in enumerate(combination):
-            yield scipy.special.binom(n_var+k-1, k+1) - scipy.special.binom(n_var+k-1-x, k+1)
+#     def gen_index_sum():
+#         for k, x in enumerate(combination):
+#             yield scipy.special.binom(n_var+k-1, k+1) - scipy.special.binom(n_var+k-1-x, k+1)
             
-    return int(sum(gen_index_sum()))
+#     return int(sum(gen_index_sum()))
 
 
-def variable_powers_to_index(powers):
-    """
-    example:
-        expr = x1 * x3**2
-        powers = (1, 0, 2)
+# def variable_powers_to_index(powers):
+#     """
+#     example:
+#         expr = x1 * x3**2
+#         powers = (1, 0, 2)
 
-        returns 5
-    """
+#         returns 5
+#     """
 
-    n_var = len(powers)
+#     n_var = len(powers)
     
-    indices = sorted((idx for idx, p in enumerate(powers) for _ in range(p)), reverse=True)
+#     indices = sorted((idx for idx, p in enumerate(powers) for _ in range(p)), reverse=True)
     
-    return variable_to_index(n_var, indices)
\ No newline at end of file
+#     return variable_to_index(n_var, indices)
\ No newline at end of file
diff --git a/test_polymatrix/__init__.py b/test_polymatrix/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/test_polymatrix/__init__.py
diff --git a/test_polymatrix/test_expression/__init__.py b/test_polymatrix/test_expression/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/test_polymatrix/test_expression/__init__.py
diff --git a/test_polymatrix/test_expression/test_addition.py b/test_polymatrix/test_expression/test_addition.py
new file mode 100644
index 0000000..623b13f
--- /dev/null
+++ b/test_polymatrix/test_expression/test_addition.py
@@ -0,0 +1,64 @@
+import unittest
+
+from polymatrix.expression.init.initadditionexpr import init_addition_expr
+from polymatrix.expression.init.initexpressionstate import init_expression_state
+from polymatrix.expression.init.initfromtermsexpr import init_from_terms_expr
+
+
+class TestAddition(unittest.TestCase):
+
+    def test_1(self):
+        left_terms = {
+            (0, 0): {
+                tuple(): 1.0,
+                ((0, 1),): 1.0,
+            }, 
+            (1, 0): {
+                ((0, 2),): 1.0,
+            },
+        }
+
+        right_terms = {
+            (0, 0): {
+                tuple(): 3.0,
+                ((1, 1),): 2.0,
+            },
+            (1, 1): {
+                tuple(): 1.0,
+            },
+        }
+
+        left = init_from_terms_expr(
+            terms=left_terms,
+            shape=(2, 2),
+        )
+
+        right = init_from_terms_expr(
+            terms=right_terms,
+            shape=(2, 2),
+        )
+
+        expr = init_addition_expr(
+            left=left,
+            right=right,
+        )
+
+        state = init_expression_state(n_param=2)
+        state, val = expr.apply(state)
+
+        data = val.get_poly(0, 0)
+        self.assertDictContainsSubset({
+            tuple(): 4.0,
+            ((0, 1),): 1.0,
+            ((1, 1),): 2.0,
+        }, data)
+
+        data = val.get_poly(1, 0)
+        self.assertDictContainsSubset({
+            ((0, 2),): 1.0,
+        }, data)
+
+        data = val.get_poly(1, 1)
+        self.assertDictContainsSubset({
+            tuple(): 1.0,
+        }, data)
diff --git a/test_polymatrix/test_expression/test_blockdiag.py b/test_polymatrix/test_expression/test_blockdiag.py
new file mode 100644
index 0000000..69141a4
--- /dev/null
+++ b/test_polymatrix/test_expression/test_blockdiag.py
@@ -0,0 +1,58 @@
+import unittest
+
+from polymatrix.expression.init.initblockdiagexpr import init_block_diag_expr
+from polymatrix.expression.init.initexpressionstate import init_expression_state
+from polymatrix.expression.init.initfromtermsexpr import init_from_terms_expr
+
+
+class TestBlockDiag(unittest.TestCase):
+
+    def test_1(self):
+        terms1 = {
+            (0, 0): {
+                ((1, 1),): 1.0,
+            }, 
+            (1, 0): {
+                tuple(): 2.0,
+            },
+        }
+
+        terms2 = {
+            (0, 0): {
+                tuple(): 3.0,
+            },
+            (1, 1): {
+                tuple(): 4.0,
+            },
+        }
+
+        expr = init_block_diag_expr(
+            underlying=(
+                init_from_terms_expr(terms=terms1, shape=(2, 2),), 
+                init_from_terms_expr(terms=terms2, shape=(2, 2),),
+            ),
+        )
+
+        state = init_expression_state(n_param=2)
+        state, val = expr.apply(state)
+
+        data = val.get_poly(0, 0)
+        self.assertDictEqual({
+            ((1, 1),): 1.0,
+        }, data)
+
+        data = val.get_poly(1, 0)
+        self.assertDictEqual({
+            tuple(): 2.0,
+        }, data)
+
+        data = val.get_poly(2, 2)
+        self.assertDictEqual({
+            tuple(): 3.0,
+        }, data)
+
+        data = val.get_poly(3, 3)
+        self.assertDictEqual({
+            tuple(): 4.0,
+        }, data)
+
diff --git a/test_polymatrix/test_expression/test_derivative.py b/test_polymatrix/test_expression/test_derivative.py
new file mode 100644
index 0000000..a4fc6f6
--- /dev/null
+++ b/test_polymatrix/test_expression/test_derivative.py
@@ -0,0 +1,52 @@
+import unittest
+from polymatrix.expression.init.initderivativeexpr import init_derivative_expr
+from polymatrix.expression.init.initdivergenceexpr import init_divergence_expr
+
+from polymatrix.expression.init.initexpressionstate import init_expression_state
+from polymatrix.expression.init.initfromtermsexpr import init_from_terms_expr
+from polymatrix.expression.init.initlinearinexpr import init_linear_in_expr
+
+
+class TestDerivative(unittest.TestCase):
+
+    def test_1(self):
+        underlying_terms = {
+            (0, 0): {
+                ((0, 1),): 2.0,
+                ((1, 2),): 3.0,
+            },
+            (1, 0): {
+                tuple(): 5.0,
+                ((0, 1), (2, 3)): 4.0,
+            },
+        }
+
+        expr = init_derivative_expr(
+            underlying=init_from_terms_expr(terms=underlying_terms, shape=(2, 1)),
+            variables=(0, 1, 2),
+        )
+
+        state = init_expression_state(n_param=3)
+        state, val = expr.apply(state)
+
+        self.assertTupleEqual(val.shape, (2, 3))
+
+        data = val.get_poly(0, 0)
+        self.assertDictEqual({
+            tuple(): 2.0,
+        }, data)
+
+        data = val.get_poly(0, 1)
+        self.assertDictEqual({
+            ((1, 1),): 6.0,
+        }, data)
+
+        data = val.get_poly(1, 0)
+        self.assertDictEqual({
+            ((2, 3),): 4.0,
+        }, data)
+
+        data = val.get_poly(1, 2)
+        self.assertDictEqual({
+            ((0, 1), (2, 2)): 12.0,
+        }, data)
diff --git a/test_polymatrix/test_expression/test_divergence.py b/test_polymatrix/test_expression/test_divergence.py
new file mode 100644
index 0000000..412dcca
--- /dev/null
+++ b/test_polymatrix/test_expression/test_divergence.py
@@ -0,0 +1,39 @@
+import unittest
+from polymatrix.expression.init.initdivergenceexpr import init_divergence_expr
+
+from polymatrix.expression.init.initexpressionstate import init_expression_state
+from polymatrix.expression.init.initfromtermsexpr import init_from_terms_expr
+from polymatrix.expression.init.initlinearinexpr import init_linear_in_expr
+
+
+class TestDivergence(unittest.TestCase):
+
+    def test_1(self):
+        underlying_terms = {
+            (0, 0): {
+                ((0, 1),): 2.0,
+                ((1, 1),): 3.0,
+            },
+            (1, 0): {
+                tuple(): 5.0,
+                ((0, 1),): 3.0,
+            },
+            (2, 0): {
+                ((0, 1),): 2.0,
+                ((1, 1), (2, 3)): 3.0,
+            },
+        }
+
+        expr = init_divergence_expr(
+            underlying=init_from_terms_expr(terms=underlying_terms, shape=(3, 1)),
+            variables=(0, 1, 2),
+        )
+
+        state = init_expression_state(n_param=2)
+        state, val = expr.apply(state)
+
+        data = val.get_poly(0, 0)
+        self.assertDictEqual({
+            tuple(): 2.0,
+            ((1, 1), (2, 2)): 9.0,
+        }, data)
diff --git a/test_polymatrix/test_expression/test_eval.py b/test_polymatrix/test_expression/test_eval.py
new file mode 100644
index 0000000..662322a
--- /dev/null
+++ b/test_polymatrix/test_expression/test_eval.py
@@ -0,0 +1,41 @@
+import unittest
+
+from polymatrix.expression.init.initevalexpr import init_eval_expr
+from polymatrix.expression.init.initexpressionstate import init_expression_state
+from polymatrix.expression.init.initfromtermsexpr import init_from_terms_expr
+
+
+class TestEval(unittest.TestCase):
+
+    def test_1(self):
+        terms = {
+            (0, 0): {
+                ((0, 1), (2, 1)): 2.0,
+                ((0, 1), (1, 1), (3, 1)): 3.0,
+            }, (1, 0):{
+                tuple(): 1.0,
+                ((1, 2),): 1.0,
+                ((2, 1),): 1.0,
+            },
+        }
+
+        expr = init_eval_expr(
+            underlying=init_from_terms_expr(terms=terms, shape=(2, 1)),
+            variables=(0, 1),
+            values=(2.0, 3.0),
+        )
+
+        state = init_expression_state(n_param=2)
+        state, val = expr.apply(state)
+
+        data = val.get_poly(0, 0)
+        self.assertDictEqual({
+            ((2, 1),): 4.0,
+            ((3, 1),): 18.0,
+        }, data)
+
+        data = val.get_poly(1, 0)
+        self.assertDictEqual({
+            tuple(): 10.0,
+            ((2, 1),): 1,
+        }, data)
diff --git a/test_polymatrix/test_expression/test_linearin.py b/test_polymatrix/test_expression/test_linearin.py
new file mode 100644
index 0000000..67f2c3d
--- /dev/null
+++ b/test_polymatrix/test_expression/test_linearin.py
@@ -0,0 +1,50 @@
+import unittest
+
+from polymatrix.expression.init.initexpressionstate import init_expression_state
+from polymatrix.expression.init.initfromtermsexpr import init_from_terms_expr
+from polymatrix.expression.init.initlinearinexpr import init_linear_in_expr
+
+
+class TestLinearIn(unittest.TestCase):
+
+    def test_1(self):
+        underlying_terms = {
+            (0, 0): {
+                ((0, 1),): 2.0,
+                ((1, 1),): 3.0,
+            },
+        }
+
+        monomial_terms = {
+            (0, 0): {
+                ((0, 1),): 1.0,
+            },
+            (1, 0): {
+                ((2, 1),): 1.0,
+            },
+            (2, 0): {
+                ((1, 1),): 1.0,
+            },
+            (3, 0): {
+                ((3, 1),): 1.0,
+            },
+        }
+
+        expr = init_linear_in_expr(
+            underlying=init_from_terms_expr(terms=underlying_terms, shape=(2, 1)),
+            monomials=init_from_terms_expr(terms=monomial_terms, shape=(4, 1),),
+            variables=(0, 1),
+        )
+
+        state = init_expression_state(n_param=2)
+        state, val = expr.apply(state)
+
+        data = val.get_poly(0, 0)
+        self.assertDictEqual({
+            tuple(): 2.0,
+        }, data)
+
+        data = val.get_poly(0, 2)
+        self.assertDictEqual({
+            tuple(): 3.0,
+        }, data)
diff --git a/test_polymatrix/test_expression/test_matrixmult.py b/test_polymatrix/test_expression/test_matrixmult.py
new file mode 100644
index 0000000..6facb48
--- /dev/null
+++ b/test_polymatrix/test_expression/test_matrixmult.py
@@ -0,0 +1,64 @@
+import unittest
+from polymatrix.expression.init.initadditionexpr import init_addition_expr
+from polymatrix.expression.init.initexpressionstate import init_expression_state
+
+from polymatrix.expression.init.initfromtermsexpr import init_from_terms_expr
+from polymatrix.expression.init.initmatrixmultexpr import init_matrix_mult_expr
+
+
+class TestMatrixMult(unittest.TestCase):
+
+    def test_1(self):
+        left_terms = {
+            (0, 0): {
+                tuple(): 1.0,
+                ((0, 1),): 1.0,
+            },
+            (0, 1): {
+                ((0, 1),): 1.0,
+            }, 
+            (1, 1): {
+                ((0, 2),): 1.0,
+            },
+        }
+
+        right_terms = {
+            (0, 0): {
+                tuple(): 3.0,
+                ((1, 1),): 2.0,
+            },
+            (1, 0): {
+                tuple(): 1.0,
+            },
+        }
+
+        left = init_from_terms_expr(
+            terms=left_terms,
+            shape=(2, 2),
+        )
+
+        right = init_from_terms_expr(
+            terms=right_terms,
+            shape=(2, 1),
+        )
+
+        expr = init_matrix_mult_expr(
+            left=left,
+            right=right,
+        )
+
+        state = init_expression_state(n_param=2)
+        state, val = expr.apply(state)
+
+        data = val.get_poly(0, 0)
+        self.assertDictEqual({
+            tuple(): 3.0,
+            ((0, 1),): 4.0,
+            ((1, 1),): 2.0,
+            ((0, 1), (1, 1),): 2.0,
+        }, data)
+
+        data = val.get_poly(1, 0)
+        self.assertDictEqual({
+            ((0, 2),): 1.0,
+        }, data)
diff --git a/test_polymatrix/test_expression/test_quadraticin.py b/test_polymatrix/test_expression/test_quadraticin.py
new file mode 100644
index 0000000..c340d7f
--- /dev/null
+++ b/test_polymatrix/test_expression/test_quadraticin.py
@@ -0,0 +1,62 @@
+import unittest
+
+from polymatrix.expression.init.initexpressionstate import init_expression_state
+from polymatrix.expression.init.initfromtermsexpr import init_from_terms_expr
+from polymatrix.expression.init.initquadraticinexpr import init_quadratic_in_expr
+
+
+class TestQuadraticIn(unittest.TestCase):
+
+    def test_1(self):
+        underlying_terms = {
+            (0, 0): {
+                ((0, 1),): 1.0,  # x1
+                ((0, 1), (2, 1)): 2.0,  # x1 
+                ((0, 2), (3, 1)): 3.0,  # x1 x1
+                ((0, 2), (1, 2), (4, 1)): 4.0,  # x1 x1 x2 x2
+                ((0, 2), (1, 1), (5, 1)): 5.0,  # x1 x1 x2
+            }
+        }
+
+        monomial_terms = {
+            (0, 0): {
+                tuple(): 1.0,
+            },
+            (1, 0): {
+                ((0, 1),): 1.0,
+            },
+            (2, 0): {
+                ((0, 1), (1, 1)): 1.0,
+            },
+        }
+
+        expr = init_quadratic_in_expr(
+            underlying=init_from_terms_expr(terms=underlying_terms, shape=(1, 1)),
+            monomials=init_from_terms_expr(terms=monomial_terms, shape=(3, 1)),
+            variables=(0, 1),
+        )
+
+        state = init_expression_state(n_param=2)
+        state, val = expr.apply(state)
+
+        data = val.get_poly(0, 1)
+        self.assertDictContainsSubset({
+            tuple(): 1.0,
+            ((2, 1),): 2.0,
+        }, data)
+
+        data = val.get_poly(1, 1)
+        self.assertDictContainsSubset({
+            ((3, 1),): 3.0,
+        }, data)
+
+        data = val.get_poly(2, 2)
+        self.assertDictContainsSubset({
+            ((4, 1),): 4.0,
+        }, data)
+
+        data = val.get_poly(1, 2)
+        self.assertDictContainsSubset({
+            ((5, 1),): 5.0,
+        }, data)
+    
\ No newline at end of file
diff --git a/test_polymatrix/test_expression/test_substitude.py b/test_polymatrix/test_expression/test_substitude.py
new file mode 100644
index 0000000..9a6e875
--- /dev/null
+++ b/test_polymatrix/test_expression/test_substitude.py
@@ -0,0 +1,43 @@
+import unittest
+
+from polymatrix.expression.init.initexpressionstate import init_expression_state
+from polymatrix.expression.init.initfromtermsexpr import init_from_terms_expr
+from polymatrix.expression.init.initsubstituteexpr import init_substitute_expr
+
+
+class TestEval(unittest.TestCase):
+
+    def test_1(self):
+        terms = {
+            (0, 0): {
+                tuple(): 2.0,
+                ((0, 2),): 3.0,
+                ((1, 1),): 1.0,
+                ((2, 2),): 1.0,
+            },
+        }
+
+        substitution = {
+            (0, 0): {
+                ((1, 1),): 1.0,
+                ((2, 1),): 1.0,
+            },
+        }
+
+        expr = init_substitute_expr(
+            underlying=init_from_terms_expr(terms=terms, shape=(1, 1)),
+            variables=(0,),
+            substitutions=(init_from_terms_expr(terms=substitution, shape=(1, 1)),),
+        )
+
+        state = init_expression_state(n_param=2)
+        state, val = expr.apply(state)
+
+        data = val.get_poly(0, 0)
+        self.assertDictEqual({
+            tuple(): 2.0,
+            ((1, 1),): 1.0,
+            ((1, 2),): 3.0,
+            ((1, 1), (2, 1)): 6.0,
+            ((2, 2),): 4.0
+        }, data)
diff --git a/test_polymatrix/test_expression/test_subtractmonomials.py b/test_polymatrix/test_expression/test_subtractmonomials.py
new file mode 100644
index 0000000..f80f76a
--- /dev/null
+++ b/test_polymatrix/test_expression/test_subtractmonomials.py
@@ -0,0 +1,55 @@
+import unittest
+from polymatrix.expression.init.initderivativeexpr import init_derivative_expr
+from polymatrix.expression.init.initdivergenceexpr import init_divergence_expr
+
+from polymatrix.expression.init.initexpressionstate import init_expression_state
+from polymatrix.expression.init.initfromtermsexpr import init_from_terms_expr
+from polymatrix.expression.init.initlinearinexpr import init_linear_in_expr
+from polymatrix.expression.init.initsubtractmonomialsexpr import init_subtract_monomials_expr
+
+
+class TestDerivative(unittest.TestCase):
+
+    def test_1(self):
+        monomials1 = {
+            (0, 0): {
+                ((0, 1),): 1.0,
+            },
+            (1, 0): {
+                ((0, 1), (1, 2)): 1.0,
+            },
+        }
+
+        monomials2 = {
+            (0, 0): {
+                ((0, 1),): 1.0,
+            },
+            (1, 0): {
+                ((1, 1),): 1.0,
+            },
+        }
+
+        expr = init_subtract_monomials_expr(
+            underlying=init_from_terms_expr(terms=monomials1, shape=(2, 1)),
+            monomials=init_from_terms_expr(terms=monomials2, shape=(2, 1)),
+        )
+
+        state = init_expression_state(n_param=3)
+        state, val = expr.apply(state)
+
+        self.assertTupleEqual(val.shape, (3, 1))
+
+        data = val.get_poly(0, 0)
+        self.assertDictEqual({
+            tuple(): 1.0,
+        }, data)
+
+        data = val.get_poly(1, 0)
+        self.assertDictEqual({
+            ((1, 2),): 1.0,
+        }, data)
+
+        data = val.get_poly(2, 0)
+        self.assertDictEqual({
+            ((0, 1), (1, 1)): 1.0,
+        }, data)
diff --git a/test_polymatrix/test_expression/test_sum.py b/test_polymatrix/test_expression/test_sum.py
new file mode 100644
index 0000000..48ae073
--- /dev/null
+++ b/test_polymatrix/test_expression/test_sum.py
@@ -0,0 +1,35 @@
+import unittest
+
+from polymatrix.expression.init.initevalexpr import init_eval_expr
+from polymatrix.expression.init.initexpressionstate import init_expression_state
+from polymatrix.expression.init.initfromtermsexpr import init_from_terms_expr
+from polymatrix.expression.init.initsumexpr import init_sum_expr
+
+
+class TestSum(unittest.TestCase):
+
+    def test_1(self):
+        terms = {
+            (0, 0): {
+                tuple(): 2.0,
+                ((0, 1),): 3.0,
+            }, 
+            (0, 1):{
+                tuple(): 1.0,
+                ((0, 2),): 1.0,
+            },
+        }
+
+        expr = init_sum_expr(
+            underlying=init_from_terms_expr(terms=terms, shape=(1, 2)),
+        )
+
+        state = init_expression_state(n_param=2)
+        state, val = expr.apply(state)
+
+        data = val.get_poly(0, 0)
+        self.assertDictEqual({
+            tuple(): 3.0,
+            ((0, 1),): 3.0,
+            ((0, 2),): 1.0,
+        }, data)
diff --git a/test_polymatrix/test_expression/test_symmetric.py b/test_polymatrix/test_expression/test_symmetric.py
new file mode 100644
index 0000000..ac5eba6
--- /dev/null
+++ b/test_polymatrix/test_expression/test_symmetric.py
@@ -0,0 +1,53 @@
+import unittest
+
+from polymatrix.expression.init.initexpressionstate import init_expression_state
+from polymatrix.expression.init.initfromtermsexpr import init_from_terms_expr
+from polymatrix.expression.init.initquadraticinexpr import init_quadratic_in_expr
+from polymatrix.expression.init.initsymmetricexpr import init_symmetric_expr
+
+
+class TestQuadraticIn(unittest.TestCase):
+
+    def test_1(self):
+        terms = {
+            (0, 0): {
+                ((0, 1),): 1.0,
+            },
+            (1, 0): {
+                ((1, 1),): 1.0,
+            },
+            (0, 1): {
+                ((1, 1),): 1.0,
+                ((2, 1),): 1.0,
+            },
+        }
+
+        underlying = init_from_terms_expr(
+            terms=terms,
+            shape=(2, 2),
+        )
+
+        expr = init_symmetric_expr(
+            underlying=underlying,
+        )
+
+        state = init_expression_state(n_param=2)
+        state, val = expr.apply(state)
+
+        data = val.get_poly(0, 0)
+        self.assertDictContainsSubset({
+            ((0, 1),): 1.0,
+        }, data)
+
+        data = val.get_poly(0, 1)
+        self.assertDictContainsSubset({
+            ((1, 1),): 1.0,
+            ((2, 1),): 0.5,
+        }, data)
+
+        data = val.get_poly(1, 0)
+        self.assertDictContainsSubset({
+            ((1, 1),): 1.0,
+            ((2, 1),): 0.5,
+        }, data)
+    
\ No newline at end of file
diff --git a/test_polymatrix/test_expression/test_toconstant.py b/test_polymatrix/test_expression/test_toconstant.py
new file mode 100644
index 0000000..784aeec
--- /dev/null
+++ b/test_polymatrix/test_expression/test_toconstant.py
@@ -0,0 +1,46 @@
+import unittest
+
+from polymatrix.expression.init.initexpressionstate import init_expression_state
+from polymatrix.expression.init.initfromtermsexpr import init_from_terms_expr
+from polymatrix.expression.init.initquadraticinexpr import init_quadratic_in_expr
+from polymatrix.expression.init.initsymmetricexpr import init_symmetric_expr
+from polymatrix.expression.init.inittoconstantexpr import init_to_constant_expr
+from polymatrix.expression.init.inittruncateexpr import init_truncate_expr
+
+
+class TestToConstant(unittest.TestCase):
+
+    def test_1(self):
+        terms = {
+            (0, 0): {
+                tuple(): 2.0,
+                ((0, 1),): 1.0,
+            },
+            (1, 0): {
+                ((0, 2), (1, 1)): 1.0,
+                ((0, 3), (1, 1)): 1.0,
+            },
+            (0, 1): {
+                tuple(): 5.0,
+                ((0, 2), (2, 1),): 1.0,
+                ((3, 1),): 1.0,
+            },
+        }
+
+        expr = init_to_constant_expr(
+            underlying=init_from_terms_expr(terms=terms, shape=(2, 2)),
+        )
+
+        state = init_expression_state(n_param=2)
+        state, val = expr.apply(state)
+
+        data = val.get_poly(0, 0)
+        self.assertDictEqual({
+            tuple(): 2.0,
+        }, data)
+
+        data = val.get_poly(0, 1)
+        self.assertDictEqual({
+            tuple(): 5.0,
+        }, data)
+    
\ No newline at end of file
diff --git a/test_polymatrix/test_expression/test_truncate.py b/test_polymatrix/test_expression/test_truncate.py
new file mode 100644
index 0000000..e944229
--- /dev/null
+++ b/test_polymatrix/test_expression/test_truncate.py
@@ -0,0 +1,50 @@
+import unittest
+
+from polymatrix.expression.init.initexpressionstate import init_expression_state
+from polymatrix.expression.init.initfromtermsexpr import init_from_terms_expr
+from polymatrix.expression.init.initquadraticinexpr import init_quadratic_in_expr
+from polymatrix.expression.init.initsymmetricexpr import init_symmetric_expr
+from polymatrix.expression.init.inittruncateexpr import init_truncate_expr
+
+
+class TestTruncate(unittest.TestCase):
+
+    def test_1(self):
+        terms = {
+            (0, 0): {
+                ((0, 1),): 1.0,  # x1 x1
+            },
+            (1, 0): {
+                ((0, 2), (1, 1)): 1.0,  # x1 x1 x2
+                ((0, 3), (1, 1)): 1.0,  # x1 x1 x1 x2
+            },
+            (0, 1): {
+                ((0, 2), (2, 1),): 1.0,  # x1 x1
+                ((3, 1),): 1.0,
+            },
+        }
+
+        expr = init_truncate_expr(
+            underlying=init_from_terms_expr(terms=terms, shape=(2, 2)),
+            variables=(0, 1),
+            degrees=(1, 2),
+        )
+
+        state = init_expression_state(n_param=2)
+        state, val = expr.apply(state)
+
+        data = val.get_poly(0, 0)
+        self.assertDictEqual({
+            ((0, 1),): 1.0,
+        }, data)
+
+        data = val.get_poly(1, 0)
+        self.assertDictEqual(
+            {}, data
+        )
+
+        data = val.get_poly(0, 1)
+        self.assertDictEqual({
+            ((0, 2), (2, 1)): 1.0,  # x1 x1
+        }, data)
+    
\ No newline at end of file
diff --git a/test_polymatrix/test_expression/test_vstack.py b/test_polymatrix/test_expression/test_vstack.py
new file mode 100644
index 0000000..a50267f
--- /dev/null
+++ b/test_polymatrix/test_expression/test_vstack.py
@@ -0,0 +1,58 @@
+import unittest
+
+from polymatrix.expression.init.initexpressionstate import init_expression_state
+from polymatrix.expression.init.initfromtermsexpr import init_from_terms_expr
+from polymatrix.expression.init.initvstackexpr import init_v_stack_expr
+
+
+class TestVStack(unittest.TestCase):
+
+    def test_1(self):
+        terms1 = {
+            (0, 0): {
+                ((1, 1),): 1.0,
+            }, 
+            (1, 0): {
+                tuple(): 2.0,
+            },
+        }
+
+        terms2 = {
+            (0, 0): {
+                tuple(): 3.0,
+            },
+            (1, 1): {
+                tuple(): 4.0,
+            },
+        }
+
+        expr = init_v_stack_expr(
+            underlying=(
+                init_from_terms_expr(terms=terms1, shape=(2, 2),), 
+                init_from_terms_expr(terms=terms2, shape=(2, 2),),
+            ),
+        )
+
+        state = init_expression_state(n_param=2)
+        state, val = expr.apply(state)
+
+        data = val.get_poly(0, 0)
+        self.assertDictEqual({
+            ((1, 1),): 1.0,
+        }, data)
+
+        data = val.get_poly(1, 0)
+        self.assertDictEqual({
+            tuple(): 2.0,
+        }, data)
+
+        data = val.get_poly(2, 0)
+        self.assertDictEqual({
+            tuple(): 3.0,
+        }, data)
+
+        data = val.get_poly(3, 1)
+        self.assertDictEqual({
+            tuple(): 4.0,
+        }, data)
+
diff --git a/test_polymatrix/test_polymatrix.py b/test_polymatrix/test_polymatrix.py
deleted file mode 100644
index d8abeef..0000000
--- a/test_polymatrix/test_polymatrix.py
+++ /dev/null
@@ -1,336 +0,0 @@
-import unittest
-from polymatrix.init.initoptimization import init_optimization
-from polymatrix.init.initpolymatrix import init_poly_matrix
-from polymatrix.optimization import Optimization
-from polymatrix.polymatrix import PolyMatrix
-
-class TestPolyMatrix(unittest.TestCase):
-    # @staticmethod
-    # def assert_term_in_eq(result, degree, eq_idx, row_idx, value, monoms=None):
-    #     if monoms is None:
-    #         monoms = tuple()
-
-    #     assert degree in result.data, f'could not find {degree} in {result.data}'
-    #     degree_terms = result.data[degree]
-
-    #     key = eq_idx, monoms, row_idx
-    #     assert key in degree_terms, f'could not find {key} in {degree_terms}'
-    #     # eq_terms = degree_terms[key]
-    #     # assert row_idx in eq_terms, f'could not find {row_idx} in {eq_terms}'
-    #     assert degree_terms[key] == value, f'value {degree_terms[key]} and {value} do not match'
-
-    # @staticmethod
-    # def assert_term_in_eq(result, degree, row_idx, monoms, value):
-
-    #     assert degree in result.data, f'could not find {degree} in {result.data}'
-    #     degree_terms = result.data[degree]
-
-    #     key = row_idx, variable_to_index(result.n_param, monoms)
-    #     assert key in degree_terms, f'could not find {key} in {degree_terms}'
-
-    #     assert degree_terms[key] == value, f'value {degree_terms[key]} and {value} do not match'
-
-    @staticmethod
-    def assert_term_in_eq(
-        problem: Optimization,
-        poly_row: int, 
-        p_monomial: tuple[tuple[PolyMatrix, int, int], ...], 
-        value: float,
-        x_monomial: tuple[int, ...] = None,
-    ):
-        if x_monomial is None:
-            x_monomial = tuple()
-
-        offset_dict = problem.state.offset_dict
-        p_monomial = tuple(offset_dict[(polymat, degree)][0] + offset for polymat, degree, offset in p_monomial)
-
-        equality_constraint = problem.equality_constraints[0]
-
-        key = (poly_row, x_monomial, p_monomial)
-        assert key in equality_constraint, f'could not find {key} in {equality_constraint}'
-
-        assert equality_constraint[key] == value, f'value {equality_constraint[key]} and {value} do not match'
-
-    def test_param_matrix_param_d0_vector_degree_d0(self):
-        """
-        param = [a11 a21 a31 a41 v11 v21]
-        """
-
-        n_var = 2
-
-        mat = init_poly_matrix(name='mat', degrees=(0,), shape=(n_var, n_var))
-        vec = init_poly_matrix(name='vec', degrees=(0,), shape=(n_var, 1))
-
-        problem = init_optimization(
-            n_var=n_var,
-        ).add_equality_constraints(
-            expr=[(mat, vec)],
-        )
-
-        # a11 v11
-        self.assert_term_in_eq(
-            problem = problem,
-            poly_row = 0,
-            p_monomial = ((mat, 0, 0), (vec, 0, 0)),
-            value = 1,
-        )
-
-        # a12 v21
-        self.assert_term_in_eq(
-            problem = problem,
-            poly_row = 0,
-            p_monomial = ((mat, 0, 2), (vec, 0, 1)),
-            value = 1,
-        )
-
-        # a21 v11
-        self.assert_term_in_eq(
-            problem = problem,
-            poly_row = 1,
-            p_monomial = ((mat, 0, 1), (vec, 0, 0)),
-            value = 1,
-        )
-
-        # a22 v21
-        self.assert_term_in_eq(
-            problem = problem,
-            poly_row = 1,
-            p_monomial = ((mat, 0, 3), (vec, 0, 1)),
-            value = 1,
-        )
-
-    def test_param_matrix_d0_param_vector_d01(self):
-        """
-        param = [a11 a21 a31 a41 v011 v021 v111 v112 v121 v122]
-        """
-
-        n_var = 2
-
-        mat = init_poly_matrix(name='mat', degrees=(0,), shape=(n_var, n_var))
-        vec = init_poly_matrix(name='vec', degrees=(0, 1), shape=(n_var, 1))
-
-        problem = init_optimization(
-            n_var=n_var,
-        ).add_equality_constraints(
-            expr=[(mat, vec)],
-        )
-
-        # a11 v011
-        self.assert_term_in_eq(
-            problem = problem,
-            poly_row = 0,
-            p_monomial = ((mat, 0, 0), (vec, 0, 0)),
-            value = 1,
-        )
-
-        # a11 v011
-        self.assert_term_in_eq(
-            problem = problem,
-            poly_row = 0,
-            x_monomial=(0,),
-            p_monomial = ((mat, 0, 0), (vec, 1, 0)),
-            value = 1,
-        )
-
-    def test_param_matrix_d0_const_vector_d0(self):
-        """
-        param = [a11 a21 a31 a41]
-        """
-
-        n_var = 2
-
-        mat = init_poly_matrix(name='mat', degrees=(0,), shape=(n_var, n_var))
-        vec = init_poly_matrix(name='vec', subs={0: {(0, 0, 0): 1, (1, 0, 0): 1}}, shape=(n_var, 1))
-
-        problem = init_optimization(
-            n_var=n_var,
-        ).add_equality_constraints(
-            expr=[(mat, vec)],
-        )
-
-
-        # a11
-        self.assert_term_in_eq(
-            problem = problem,
-            poly_row = 0,
-            p_monomial = ((mat, 0, 0),),
-            value = 1,
-        )
-
-        # a21
-        self.assert_term_in_eq(
-            problem = problem,
-            poly_row = 1,
-            p_monomial = ((mat, 0, 1),),
-            value = 1,
-        )
-
-    def test_param_matrix_d0_const_vector_d1(self):
-        """
-        param = [a11 a21 a31 a41]
-        """
-
-        n_var = 2
-        
-        mat = init_poly_matrix(name='mat', degrees=(0,), shape=(n_var, n_var))
-        vec = init_poly_matrix(name='vec', subs={1: {(0, 0, 0): 1, (0, 0, 1): 0, (1, 0, 0): 0, (1, 0, 1): 1}}, shape=(n_var, 1))
-
-        problem = init_optimization(
-            n_var=n_var,
-        ).add_equality_constraints(
-            expr=[(mat, vec)],
-        )
-
-        # a11
-        self.assert_term_in_eq(
-            problem = problem,
-            poly_row = 0,
-            x_monomial=(0,),
-            p_monomial = ((mat, 0, 0),),
-            value = 1,
-        )
-
-        # a12
-        self.assert_term_in_eq(
-            problem = problem,
-            poly_row = 0,
-            x_monomial=(1,),
-            p_monomial = ((mat, 0, 2),),
-            value = 1,
-        )
-
-        # a21
-        self.assert_term_in_eq(
-            problem = problem,
-            poly_row = 1,
-            x_monomial=(0,),
-            p_monomial = ((mat, 0, 1),),
-            value = 1,
-        )
-
-    def test_const_matrix_const_vector_degree_0(self):
-        """
-        param = [a11 a21 a31 a41]
-        """
-
-        n_var = 2
-
-        mat = init_poly_matrix(name='mat', subs={0: {(0, 0, 0): 1, (0, 1, 0): 1, (1, 0, 0): 1, (1, 1, 0): 1}}, shape=(n_var, n_var))
-        vec = init_poly_matrix(name='vec', subs={0: {(0, 0, 0): 1, (1, 0, 0): 1}}, shape=(n_var, 1))
-
-        problem = init_optimization(
-            n_var=n_var,
-        ).add_equality_constraints(
-            expr=[(mat, vec)],
-        )
-
-        self.assert_term_in_eq(
-            problem = problem,
-            poly_row = 0,
-            p_monomial = tuple(),
-            value = 2,
-        )
-    
-    def test_skew_symmetric_param_matrix_const_vector(self):
-        """
-        param = [a11 a21 a31 a41]
-        """
-
-        def skew_symmetric(degree, poly_row, poly_col, monom):
-            if poly_row == poly_col:
-                return poly_row, poly_col, monom, 0
-            elif poly_col < poly_row:
-                return poly_col, poly_row, monom, -1
-
-        n_var = 2
-
-        mat = init_poly_matrix(
-            name='mat', 
-            degrees=(0,), 
-            re_index=skew_symmetric,
-            shape=(n_var, n_var),
-        )
-        vec = init_poly_matrix(name='vec', subs={0: {(0, 0, 0): 1, (1, 0, 0): 1}}, shape=(n_var, 1))
-
-        problem = init_optimization(
-            n_var=n_var,
-        ).add_equality_constraints(
-            expr=[(mat, vec)],
-        )
-
-        self.assert_term_in_eq(
-            problem = problem,
-            poly_row = 0,
-            p_monomial = ((mat, 0, 2),),
-            value = 1,
-        )
-
-        self.assert_term_in_eq(
-            problem = problem,
-            poly_row = 1,
-            p_monomial = ((mat, 0, 2),),
-            value = -1,
-        )
-
-    def test_const_matrix_d0_param_gradient_vector_d1(self):
-        """
-        param = [v11 v12 v21 v22]
-        """
-
-        def gradient(degree, p_row, p_col, monom):
-            if degree == 1:
-                factor = sum(p_row==e for e in monom) + 1
-
-                if monom[-1] < p_row:
-                    n_p_row = monom[-1]
-                    n_monom = sorted(monom + (p_row,), reverse=True)
-
-                if p_row <= monom[-1]:
-                    n_p_row = p_row
-                    n_monom = monom
-
-                return n_p_row, p_col, n_monom, factor
-
-        n_var = 2
-
-        mat = init_poly_matrix(
-            name='mat', 
-            subs={0: {(0, 0, 0): 1, (0, 1, 0): 1, (1, 0, 0): 1, (1, 1, 0): 1}}, 
-            shape=(n_var, n_var),
-        )
-        vec = init_poly_matrix(
-            name='vec', 
-            degrees=(1,),
-            re_index=gradient,
-            shape=(n_var, 1),
-        )
-
-        problem = init_optimization(
-            n_var=n_var,
-        ).add_equality_constraints(
-            expr=[(mat, vec)],
-        )
-
-        self.assert_term_in_eq(
-            problem = problem,
-            x_monomial=(0,),
-            poly_row = 0,
-            p_monomial = ((vec, 1, 0),),
-            value = 2,
-        )
-
-        self.assert_term_in_eq(
-            problem = problem,
-            x_monomial=(1,),
-            poly_row = 0,
-            p_monomial = ((vec, 1, 2),),
-            value = 1,
-        )
-
-        self.assert_term_in_eq(
-            problem = problem,
-            x_monomial=(0,),
-            poly_row = 1,
-            p_monomial = ((vec, 1, 2),),
-            value = 1,
-        )
diff --git a/test_polymatrix/test_tomatrixrepr.py b/test_polymatrix/test_tomatrixrepr.py
new file mode 100644
index 0000000..7486c8e
--- /dev/null
+++ b/test_polymatrix/test_tomatrixrepr.py
@@ -0,0 +1,48 @@
+import unittest
+import polymatrix
+
+from polymatrix.expression.init.initexpressionstate import init_expression_state
+from polymatrix.expression.init.initfromtermsexpr import init_from_terms_expr
+from polymatrix.expression.init.initlinearinexpr import init_linear_in_expr
+
+
+class TestLinearIn(unittest.TestCase):
+
+    def test_1(self):
+        underlying_terms = {
+            (0, 0): {
+                tuple(): 1.0,
+                ((1, 1),): 2.0,
+            },
+            (1, 0): {
+                ((0, 1),): 4.0,
+                ((0, 1), (1, 1)): 3.0,
+                ((1, 2),): 5.0,
+            },
+            (2, 0): {
+                ((0, 1), (1, 2)): 3.0,
+            },
+        }
+
+        expr = init_from_terms_expr(terms=underlying_terms, shape=(3, 1))
+
+        state = init_expression_state(n_param=2)
+        state, result = polymatrix.to_matrix_equations((expr,), (0, 1)).apply(state)
+
+        A0 = result.matrix_equations[0][0]
+        A1 = result.matrix_equations[0][1]
+        A2 = result.matrix_equations[0][2]
+        A3 = result.matrix_equations[0][3]
+
+        self.assertEquals(A0[0, 0], 1.0)
+
+        self.assertEquals(A1[0, 1], 2.0)
+        self.assertEquals(A1[1, 0], 4.0)
+
+        self.assertEquals(A2[1, 1], 1.5)
+        self.assertEquals(A2[1, 2], 1.5)
+        self.assertEquals(A2[1, 3], 5.0)
+
+        self.assertEquals(A3[2, 3], 1.0)
+        self.assertEquals(A3[2, 5], 1.0)
+        self.assertEquals(A3[2, 6], 1.0)