Initial commit

1 files changed, 176 insertions, 0 deletions

diff --git a/test_polymatrix/test_polymatrix.py b/test_polymatrix/test_polymatrix.py
new file mode 100644
index 0000000..1f0bb73
--- /dev/null
+++ b/test_polymatrix/test_polymatrix.py
@@ -0,0 +1,176 @@
+import unittest
+
+from polymatrix.polystruct import init_equation, init_poly_matrix, init_poly_vector
+from polymatrix.utils import variable_to_index
+
+class TestPolyMatrix(unittest.TestCase):
+    @staticmethod
+    def assert_term_in_eq(terms, degree, eq_idx, row_idx, value, monoms=None):
+        if monoms is None:
+            monoms = tuple()
+
+        assert degree in terms, f'could not find {degree} in {terms}'
+        degree_terms = terms[degree]
+
+        key = eq_idx, monoms
+        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 eq_terms[row_idx] == value, f'value {eq_terms[row_idx]} and {value} do not match'
+
+    def test_param_matrix_param_vector_degree_0(self):
+        """
+        param = [a11 a21 a31 a41 v11 v21]
+        """
+
+        mat = init_poly_matrix(degrees=(0,))
+        vec = init_poly_vector(degrees=(0,))
+        n_param = 6
+
+        eq = init_equation(
+            terms = [(mat, vec)],
+            n_var = 2,
+        )
+
+        terms, offset_dict = list(eq.create())
+
+        # a11 v11
+        self.assert_term_in_eq(
+            terms = terms,
+            degree = 2,
+            eq_idx = 0,
+            row_idx = variable_to_index(n_param, (offset_dict[(mat, 0)], offset_dict[(vec, 0)])),
+            value = 1,
+        )
+
+        # a21 v11
+        self.assert_term_in_eq(
+            terms = terms,
+            degree = 2,
+            eq_idx = 1,
+            row_idx = variable_to_index(n_param, (offset_dict[(mat, 0)]+2, offset_dict[(vec, 0)])),
+            value = 1,
+        )
+
+    def test_param_matrix_param_vector_degree_01(self):
+        """
+        param = [a11 a21 a31 a41 v011 v021 v111 v112 v121 v122]
+        """
+
+        mat = init_poly_matrix(degrees=(0,))
+        vec = init_poly_vector(degrees=(0,1))
+        n_param = 10
+
+        eq = init_equation(
+            terms = [(mat, vec)],
+            n_var = 2,
+        )
+
+        terms, offset_dict = list(eq.create())
+
+        # a11 v011
+        self.assert_term_in_eq(
+            terms = terms,
+            degree = 2,
+            eq_idx = 0,
+            row_idx = variable_to_index(n_param, (offset_dict[(mat, 0)], offset_dict[(vec, 0)])),
+            value = 1,
+        )
+
+        # a11 v011
+        self.assert_term_in_eq(
+            terms = terms,
+            degree = 2,
+            eq_idx = 0,
+            monoms=(0,),
+            row_idx = variable_to_index(n_param, (offset_dict[(mat, 0)], offset_dict[(vec, 1)])),
+            value = 1,
+        )
+
+    def test_param_matrix_const_vector_degree_0(self):
+        """
+        param = [a11 a21 a31 a41]
+        """
+
+        mat = init_poly_matrix(degrees=(0,))
+        vec = init_poly_vector(subs={0: {(0, 0): 1, (1, 0): 1}})
+
+        eq = init_equation(
+            terms = [(mat, vec)],
+            n_var = 2,
+        )
+
+        terms, offset_dict = list(eq.create())
+
+        # a11
+        self.assert_term_in_eq(
+            terms = terms,
+            degree = 1,
+            eq_idx = 0,
+            row_idx = offset_dict[(mat, 0)],
+            value = 1,
+        )
+
+    def test_const_matrix_const_vector_degree_0(self):
+        """
+        param = [a11 a21 a31 a41]
+        """
+
+        mat = init_poly_matrix(subs={0: {(0, 0): 1, (1, 0): 1, (0, 1): 1, (1, 1): 1}})
+        vec = init_poly_vector(subs={0: {(0, 0): 1, (1, 0): 1}})
+
+        eq = init_equation(
+            terms = [(mat, vec)],
+            n_var = 2,
+        )
+
+        terms, offset_dict = list(eq.create())
+
+        self.assert_term_in_eq(
+            terms = terms,
+            degree = 0,
+            eq_idx = 0,
+            row_idx = 0,
+            value = 2,
+        )
+    
+    def test_param_matrix_const_vector_skew_symmetric(self):
+        """
+        param = [a11 a21 a31 a41]
+        """
+
+        def skew_symmetric(idx1, idx2):
+            if idx1 == idx2:
+                return idx1, idx2, 0
+            elif idx2 < idx1:
+                return idx2, idx1, -1
+
+        mat = init_poly_matrix(
+            degrees=(0,), 
+            re_index_func=skew_symmetric,
+        )
+        vec = init_poly_vector(subs={0: {(0, 0): 1, (1, 0): 1}})
+
+        eq = init_equation(
+            terms = [(mat, vec)],
+            n_var = 2,
+        )
+
+        terms, offset_dict = list(eq.create())
+
+        self.assert_term_in_eq(
+            terms = terms,
+            degree = 1,
+            eq_idx = 0,
+            row_idx = offset_dict[(mat, 0)] + 1,
+            value = 1,
+        )
+
+        self.assert_term_in_eq(
+            terms = terms,
+            degree = 1,
+            eq_idx = 1,
+            row_idx = offset_dict[(mat, 0)] + 1,
+            value = -1,
+        )
\ No newline at end of file