add max_degree, max and filter operator

19 files changed, 818 insertions, 336 deletions

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)