add solver for roots of a set of polynomial equations

1 files changed, 168 insertions, 75 deletions

diff --git a/test_polymatrix/test_polymatrix.py b/test_polymatrix/test_polymatrix.py
index 3b239bf..f693df0 100644
--- a/test_polymatrix/test_polymatrix.py
+++ b/test_polymatrix/test_polymatrix.py
@@ -1,114 +1,195 @@
 import unittest
 
-from polymatrix.polystruct import init_equation, init_poly_matrix, init_poly_vector
+from polymatrix.polystruct import init_equation, init_poly_matrix
 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):
+    def assert_term_in_eq(result, 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]
+        assert degree in result.terms, f'could not find {degree} in {result.terms}'
+        degree_terms = result.terms[degree]
 
-        key = eq_idx, monoms
+        key = eq_idx, monoms, row_idx
         assert key in degree_terms, f'could not find {key} in {degree_terms}'
-        eq_terms = degree_terms[key]
+        # 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'
 
-        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):
+    def test_param_matrix_param_d0_vector_degree_d0(self):
         """
         param = [a11 a21 a31 a41 v11 v21]
         """
 
-        mat = init_poly_matrix(degrees=(0,))
-        vec = init_poly_vector(degrees=(0,))
-        n_param = 6
+        n_var = 2
+
+        mat = init_poly_matrix(degrees=(0,), shape=(n_var, n_var))
+        vec = init_poly_matrix(degrees=(0,), shape=(n_var, 1))
 
         eq = init_equation(
             terms = [(mat, vec)],
-            n_var = 2,
+            n_var = n_var,
         )
 
-        terms, offset_dict = list(eq.create())
+        result = eq.create()
+        offset_dict = eq.offset_dict
 
         # a11 v11
         self.assert_term_in_eq(
-            terms = terms,
+            result = result,
+            degree = 2,
+            eq_idx = 0,
+            row_idx = (offset_dict[(mat, 0)], offset_dict[(vec, 0)]),
+            value = 1,
+        )
+
+        # a12 v21
+        self.assert_term_in_eq(
+            result = result,
             degree = 2,
             eq_idx = 0,
-            row_idx = variable_to_index(n_param, (offset_dict[(mat, 0)], offset_dict[(vec, 0)])),
+            row_idx = (offset_dict[(mat, 0)]+2, offset_dict[(vec, 0)]+1),
             value = 1,
         )
 
         # a21 v11
         self.assert_term_in_eq(
-            terms = terms,
+            result = result,
             degree = 2,
             eq_idx = 1,
-            row_idx = variable_to_index(n_param, (offset_dict[(mat, 0)]+2, offset_dict[(vec, 0)])),
+            row_idx = (offset_dict[(mat, 0)]+1, offset_dict[(vec, 0)]),
             value = 1,
         )
 
-    def test_param_matrix_param_vector_degree_01(self):
+        # a22 v21
+        self.assert_term_in_eq(
+            result = result,
+            degree = 2,
+            eq_idx = 1,
+            row_idx = (offset_dict[(mat, 0)]+3, offset_dict[(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]
         """
 
-        mat = init_poly_matrix(degrees=(0,))
-        vec = init_poly_vector(degrees=(0,1))
-        n_param = 10
+        n_var = 2
+
+        mat = init_poly_matrix(degrees=(0,), shape=(n_var, n_var))
+        vec = init_poly_matrix(degrees=(0,1), shape=(n_var, 1))
 
         eq = init_equation(
             terms = [(mat, vec)],
-            n_var = 2,
+            n_var = n_var,
         )
 
-        terms, offset_dict = list(eq.create())
+        result = eq.create()
+        offset_dict = eq.offset_dict
 
         # a11 v011
         self.assert_term_in_eq(
-            terms = terms,
+            result = result,
             degree = 2,
             eq_idx = 0,
-            row_idx = variable_to_index(n_param, (offset_dict[(mat, 0)], offset_dict[(vec, 0)])),
+            row_idx = (offset_dict[(mat, 0)], offset_dict[(vec, 0)]),
             value = 1,
         )
 
         # a11 v011
         self.assert_term_in_eq(
-            terms = terms,
+            result = result,
             degree = 2,
             eq_idx = 0,
             monoms=(0,),
-            row_idx = variable_to_index(n_param, (offset_dict[(mat, 0)], offset_dict[(vec, 1)])),
+            row_idx = (offset_dict[(mat, 0)], offset_dict[(vec, 1)]),
+            value = 1,
+        )
+
+    def test_param_matrix_d0_const_vector_d0(self):
+        """
+        param = [a11 a21 a31 a41]
+        """
+
+        n_var = 2
+
+        mat = init_poly_matrix(degrees=(0,), shape=(n_var, n_var))
+        vec = init_poly_matrix(subs={0: {(0, 0, 0): 1, (1, 0, 0): 1}}, shape=(n_var, 1))
+
+        eq = init_equation(
+            terms = [(mat, vec)],
+            n_var = n_var,
+        )
+
+        result = eq.create()
+        offset_dict = eq.offset_dict
+
+        # a11
+        self.assert_term_in_eq(
+            result = result,
+            degree = 1,
+            eq_idx = 0,
+            row_idx = (offset_dict[(mat, 0)],),
             value = 1,
         )
 
-    def test_param_matrix_const_vector_degree_0(self):
+        # a21
+        self.assert_term_in_eq(
+            result = result,
+            degree = 1,
+            eq_idx = 1,
+            row_idx = (offset_dict[(mat, 0)]+1,),
+            value = 1,
+        )
+
+    def test_param_matrix_d0_const_vector_d1(self):
         """
         param = [a11 a21 a31 a41]
         """
 
-        mat = init_poly_matrix(degrees=(0,))
-        vec = init_poly_vector(subs={0: {(0, 0): 1, (1, 0): 1}})
+        n_var = 2
+
+        mat = init_poly_matrix(degrees=(0,), shape=(n_var, n_var))
+        vec = init_poly_matrix(subs={1: {(0, 0, 0): 1, (0, 0, 1): 0, (1, 0, 0): 0, (1, 0, 1): 1}}, shape=(n_var, 1))
 
         eq = init_equation(
             terms = [(mat, vec)],
-            n_var = 2,
+            n_var = n_var,
         )
 
-        terms, offset_dict = list(eq.create())
+        result = eq.create()
+        offset_dict = eq.offset_dict
 
         # a11
         self.assert_term_in_eq(
-            terms = terms,
+            result = result,
             degree = 1,
             eq_idx = 0,
-            row_idx = offset_dict[(mat, 0)],
+            monoms=(0,),
+            row_idx = (offset_dict[(mat, 0)],),
+            value = 1,
+        )
+
+        # a12
+        self.assert_term_in_eq(
+            result = result,
+            degree = 1,
+            eq_idx = 0,
+            monoms=(1,),
+            row_idx = (offset_dict[(mat, 0)]+2,),
+            value = 1,
+        )
+
+        # a21
+        self.assert_term_in_eq(
+            result = result,
+            degree = 1,
+            eq_idx = 1,
+            monoms=(0,),
+            row_idx = (offset_dict[(mat, 0)]+1,),
             value = 1,
         )
 
@@ -117,21 +198,24 @@ class TestPolyMatrix(unittest.TestCase):
         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}})
+        n_var = 2
+
+        mat = init_poly_matrix(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(subs={0: {(0, 0, 0): 1, (1, 0, 0): 1}}, shape=(n_var, 1))
 
         eq = init_equation(
             terms = [(mat, vec)],
-            n_var = 2,
+            n_var = n_var,
         )
 
-        terms, offset_dict = list(eq.create())
+        result = eq.create()
+        offset_dict = eq.offset_dict
 
         self.assert_term_in_eq(
-            terms = terms,
+            result = result,
             degree = 0,
             eq_idx = 0,
-            row_idx = 0,
+            row_idx = tuple(),
             value = 2,
         )
     
@@ -140,98 +224,107 @@ class TestPolyMatrix(unittest.TestCase):
         param = [a11 a21 a31 a41]
         """
 
-        def skew_symmetric(degree, idx1, idx2):
-            if idx1 == idx2:
-                return idx1, idx2, 0
-            elif idx2 < idx1:
-                return idx2, idx1, -1
+        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(
             degrees=(0,), 
-            re_index_func=skew_symmetric,
+            re_index=skew_symmetric,
+            shape=(n_var, n_var),
         )
-        vec = init_poly_vector(subs={0: {(0, 0): 1, (1, 0): 1}})
+        vec = init_poly_matrix(subs={0: {(0, 0, 0): 1, (1, 0, 0): 1}}, shape=(n_var, 1))
 
         eq = init_equation(
             terms = [(mat, vec)],
-            n_var = 2,
+            n_var = n_var,
         )
 
-        terms, offset_dict = list(eq.create())
+        result = eq.create()
+        offset_dict = eq.offset_dict
 
         self.assert_term_in_eq(
-            terms = terms,
+            result = result,
             degree = 1,
             eq_idx = 0,
-            row_idx = offset_dict[(mat, 0)] + 1,
+            row_idx = (offset_dict[(mat, 0)] + 2,),
             value = 1,
         )
 
         self.assert_term_in_eq(
-            terms = terms,
+            result = result,
             degree = 1,
             eq_idx = 1,
-            row_idx = offset_dict[(mat, 0)] + 1,
+            row_idx = (offset_dict[(mat, 0)] + 2,),
             value = -1,
         )
 
-    def test_const_matrix_param_gradient_vector(self):
+    def test_const_matrix_d0_param_gradient_vector_d1(self):
         """
         param = [v11 v12 v21 v22]
         """
 
-        def gradient(degree, v_row, monom):
+        def gradient(degree, p_row, p_col, monom):
             if degree == 1:
-                factor = sum(v_row==e for e in monom) + 1
+                factor = sum(p_row==e for e in monom) + 1
 
-                if monom[-1] < v_row:
-                    n_v_row = monom[-1]
-                    n_monom = sorted(monom + (v_row,), reverse=True)
+                if monom[-1] < p_row:
+                    n_p_row = monom[-1]
+                    n_monom = sorted(monom + (p_row,), reverse=True)
 
-                if v_row <= monom[-1]:
-                    n_v_row = v_row
+                if p_row <= monom[-1]:
+                    n_p_row = p_row
                     n_monom = monom
 
-                return n_v_row, n_monom, factor
+                return n_p_row, p_col, n_monom, factor
+
+        n_var = 2
 
         mat = init_poly_matrix(
-            subs={0: {(0, 0): 1, (1, 0): 1, (0, 1): 1, (1, 1): 1}}, 
+            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_vector(
+        vec = init_poly_matrix(
             degrees=(1,),
-            re_index_func_2=gradient,
+            re_index=gradient,
+            shape=(n_var, 1),
         )
 
         eq = init_equation(
             terms = [(mat, vec)],
-            n_var = 2,
+            n_var = n_var,
         )
 
-        terms, offset_dict = list(eq.create())
+        result = eq.create()
+        offset_dict = eq.offset_dict
 
         self.assert_term_in_eq(
-            terms = terms,
+            result = result,
             degree = 1,
             monoms=(0,),
             eq_idx = 0,
-            row_idx = offset_dict[(vec, 1)],
+            row_idx = (offset_dict[(vec, 1)],),
             value = 2,
         )
 
         self.assert_term_in_eq(
-            terms = terms,
+            result = result,
             degree = 1,
             monoms=(1,),
             eq_idx = 0,
-            row_idx = offset_dict[(vec, 1)]+1,
+            row_idx = (offset_dict[(vec, 1)]+2,),
             value = 1,
         )
 
         self.assert_term_in_eq(
-            terms = terms,
+            result = result,
             degree = 1,
             monoms=(0,),
             eq_idx = 1,
-            row_idx = offset_dict[(vec, 1)]+1,
+            row_idx = (offset_dict[(vec, 1)]+2,),
             value = 1,
         )