import unittest 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(result, degree, eq_idx, row_idx, value, monoms=None): if monoms is None: monoms = tuple() assert degree in result.terms, f'could not find {degree} in {result.terms}' degree_terms = result.terms[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' def test_param_matrix_param_d0_vector_degree_d0(self): """ param = [a11 a21 a31 a41 v11 v21] """ 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 = n_var, ) result = eq.create() offset_dict = eq.offset_dict # a11 v11 self.assert_term_in_eq( 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 = (offset_dict[(mat, 0)]+2, offset_dict[(vec, 0)]+1), value = 1, ) # a21 v11 self.assert_term_in_eq( result = result, degree = 2, eq_idx = 1, row_idx = (offset_dict[(mat, 0)]+1, offset_dict[(vec, 0)]), value = 1, ) # 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] """ 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 = n_var, ) result = eq.create() offset_dict = eq.offset_dict # a11 v011 self.assert_term_in_eq( result = result, degree = 2, eq_idx = 0, row_idx = (offset_dict[(mat, 0)], offset_dict[(vec, 0)]), value = 1, ) # a11 v011 self.assert_term_in_eq( result = result, degree = 2, eq_idx = 0, monoms=(0,), 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, ) # 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] """ 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 = n_var, ) result = eq.create() offset_dict = eq.offset_dict # a11 self.assert_term_in_eq( result = result, degree = 1, eq_idx = 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, ) def test_const_matrix_const_vector_degree_0(self): """ param = [a11 a21 a31 a41] """ 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 = n_var, ) result = eq.create() offset_dict = eq.offset_dict self.assert_term_in_eq( result = result, degree = 0, eq_idx = 0, row_idx = 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( degrees=(0,), re_index=skew_symmetric, 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 self.assert_term_in_eq( result = result, degree = 1, eq_idx = 0, row_idx = (offset_dict[(mat, 0)] + 2,), value = 1, ) self.assert_term_in_eq( result = result, degree = 1, eq_idx = 1, row_idx = (offset_dict[(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( 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( degrees=(1,), re_index=gradient, shape=(n_var, 1), ) eq = init_equation( terms = [(mat, vec)], n_var = n_var, ) result = eq.create() offset_dict = eq.offset_dict self.assert_term_in_eq( result = result, degree = 1, monoms=(0,), eq_idx = 0, row_idx = (offset_dict[(vec, 1)],), value = 2, ) self.assert_term_in_eq( result = result, degree = 1, monoms=(1,), eq_idx = 0, row_idx = (offset_dict[(vec, 1)]+2,), value = 1, ) self.assert_term_in_eq( result = result, degree = 1, monoms=(0,), eq_idx = 1, row_idx = (offset_dict[(vec, 1)]+2,), value = 1, )