diff options
| -rw-r--r-- | main.py | 124 | ||||
| -rw-r--r-- | polymatrix/polystruct.py | 199 | ||||
| -rw-r--r-- | test_polymatrix/test_polymatrix.py | 67 |
3 files changed, 259 insertions, 131 deletions
@@ -1,5 +1,6 @@ import sympy import scipy.special +import scipy.sparse import numpy as np from polymatrix.polystruct import init_equation, init_poly_matrix, init_poly_vector from polymatrix.sympyutils import poly_to_data_coord, poly_to_matrix @@ -7,6 +8,28 @@ from polymatrix.utils import variable_to_index binom = scipy.special.binom +def skew_symmetric(degree, p_row, p_col): + if p_col < p_row: + return p_col, p_row, -1 + +def zero_diagonal_for_uneven_degree(degree, p_row, p_col): + if degree % 2 == 1 and p_row == p_col: + return 0 + +def gradient(degree, v_row, monom): + if degree == 1: + factor = sum(v_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 v_row <= monom[-1]: + n_v_row = v_row + n_monom = monom + + return n_v_row, n_monom, factor + x = sympy.symbols('v_dc, i_d, i_q') vdc, id, iq = x @@ -40,17 +63,10 @@ nabla_h = init_poly_vector(subs=dh_terms) JR = init_poly_matrix(subs=jr_terms) mg = init_poly_matrix(subs=mg_terms) -def skew_symmetric(idx1, idx2): - if idx2 < idx1: - return idx2, idx1, -1 - -def zero_diagonal_for_uneven_degree(degree, idx1, idx2): - if degree % 2 == 1 and idx1 == idx2: - return 0 - - - -nabla_ha = init_poly_vector(degrees=(1,2)) +nabla_ha = init_poly_vector( + degrees=(1,), + re_index_func_2=gradient, +) JRa = init_poly_matrix( degrees=(0,1,2), re_index_func=skew_symmetric, @@ -58,6 +74,13 @@ JRa = init_poly_matrix( ) u = init_poly_vector(degrees=(1,2)) +eq = init_equation( + terms = [(JR, nabla_ha), (JRa, nabla_ha), (JRa, nabla_h), (mg, u)], + n_var = 2, + # terms = [(JR, nabla_ha)], + # n_var = 2, +) + # n_var = 2 # n_param_nabla_ha = sum(n_var*binom(n_var+p-1, p) for p in nabla_ha.degrees) # n_param_JRa = sum(n_var**2*binom(n_var+p-1, p) for p in JRa.non_zero_degrees) @@ -69,33 +92,68 @@ u = init_poly_vector(degrees=(1,2)) # print(f'{n_param_u=}') # print(f'{total=}') -# print(n_var*binom(n_var+1-1, 1)) -# print(n_var*binom(n_var+2-1, 2)) - # print(binom(total+2-1, 2)) +# mat = init_poly_matrix( +# degrees=(1,), +# re_index_func=skew_symmetric, +# subs_func=zero_diagonal_for_uneven_degree, +# ) +# vec = init_poly_vector( +# subs={0: {(0, 0): 1, (1, 0): 1}}, +# ) + +# # mat = init_poly_matrix( +# # subs={0: {(0, 0): 1, (1, 0): 1, (0, 1): 1, (1, 1): 1}}, +# # ) + +# # vec = init_poly_vector( +# # degrees=(1,), +# # re_index_func_2=gradient, +# # ) + # eq = init_equation( -# # terms = [(JR, nabla_ha), (JRa, nabla_ha), (JRa, nabla_h), (mg, u)], -# # n_var = 3, -# terms = [(JR, nabla_ha)], +# terms = [(mat, vec)], # n_var = 2, # ) -# print(eq.create()) +eq_tuples, offset_dict, n_param = eq.create() -mat = init_poly_matrix(degrees=(0,), re_index_func=skew_symmetric) -# vec = init_poly_vector(degrees=(0,1)) -vec = init_poly_vector(subs={0: {(0, 0): 1, (1, 0): 1}}) +print(f'{list(offset_dict.values())=}') +print(f'{n_param=}') -eq = init_equation( - terms = [(mat, vec)], - n_var = 2, -) -eq_tuples, offset_dict = eq.create() - -# print(eq_tuples) -# print(eq_tuples[1][(1, tuple())]) -# print(offset_dict[(mat, 0)]) -# print(offset_dict[(vec, 0)]) -print(offset_dict) -# # print(variable_to_index(10, (offset_dict[(mat, 0)], offset_dict[(vec, 1)]))) +# mapping from row index to entry in eq_tuples +rows_to_eq = list(set(key for eq_tuple_degree in eq_tuples.values() for key in eq_tuple_degree.keys())) +eq_to_rows = {eq: idx for idx, eq in enumerate(rows_to_eq)} + +print(f'n_eq = {len(rows_to_eq)}') + +# # mapping from col index to entry in eq_tuples +# cols_to_var = list(set(key for eq in eq_tuples[degree].values() for key in eq.keys())) +# var_to_cols = {var: idx for idx, var in enumerate(cols_to_var)} + +def gen_matrix_per_degree(): + for degree, eq_tuple_degree in eq_tuples.items(): + + if 0 < degree: + def gen_coords(eq_tuple_degree=eq_tuple_degree): + for eq, var_dict in eq_tuple_degree.items(): + zero_degree_val = eq_tuples[0][eq][0] + row = eq_to_rows[eq] + for var, value in var_dict.items(): + # col = variable_to_index(n_param, var) + # col = var_to_cols[var] + col = var + yield row, col, value / zero_degree_val + + row, col, data = list(zip(*gen_coords())) + + n_col = int(binom(n_param+degree-1, degree)) + n_row = int(max(row) + 1) + sparse_matrix = scipy.sparse.coo_matrix((data, (row, col)), dtype=np.float64, shape=(n_row, n_col)).toarray() + + yield degree, sparse_matrix + +result = dict(gen_matrix_per_degree()) + +print(result[1].shape) diff --git a/polymatrix/polystruct.py b/polymatrix/polystruct.py index 69a05d6..64c0f42 100644 --- a/polymatrix/polystruct.py +++ b/polymatrix/polystruct.py @@ -70,7 +70,9 @@ class EquationMixin(abc.ABC): subs: dict[PolyMatrixMixin, dict[DegreeType, dict[int, int, float]]] = None, ): if subs is None: - subs = {} + added_subs = {} + else: + added_subs = subs binom = scipy.special.binom @@ -81,106 +83,102 @@ class EquationMixin(abc.ABC): acc = 0 for struct in all_structs: - added_subs = subs.get(struct, {}) for degree in struct.degrees: - - if struct.subs is not None and degree in struct.subs: - struct_subs = struct.subs[degree] - - elif struct.subs_func is not None: - def gen_coords(): - for eq in range(self.n_var): - - if struct.is_vector: - col_indices = (0,) - else: - col_indices = range(self.n_var) - for col in col_indices: - coord = (eq, col) + yield (struct, degree), acc - subs_val = struct.subs_func(*coord) + if struct.is_vector: + number_of_terms = int(self.n_var * binom(self.n_var+degree-1, degree)) + else: + number_of_terms = int(self.n_var**2 * binom(self.n_var+degree-1, degree)) - if subs_val is not None: - yield coord, subs_val + acc += number_of_terms - struct_subs = dict(gen_coords()) - - else: - struct_subs = {} + yield None, acc #, None - if struct.re_index is not None and degree in struct.re_index: - struct_re_index = struct.re_index[degree] - - elif struct.re_index_func is not None: - def gen_coords(): - for eq in range(self.n_var): - - if struct.is_vector: - col_indices = (0,) - else: - col_indices = range(self.n_var) + param_list = list(gen_n_param_per_struct()) + n_param = param_list[-1][1] + offset_dict = dict(((e[0], e[1]) for e in param_list[:-1])) + + def gen_substitutions(): + for struct in all_structs: + added_struct_subs = added_subs.get(struct, None) - for col in col_indices: - coord = (eq, col) + if added_struct_subs is not None: + if struct.subs is not None: + def gen_merged_subs(): + for degree in struct.degrees: + all_subs = struct.subs.get(degree, {}) | added_struct_subs.get(degree, {}) + yield degree, all_subs - re_index_val = struct.re_index_func(*coord) + all_subs = dict(gen_merged_subs()) + + else: + all_subs = added_struct_subs - match re_index_val: - case (row, col, value): - yield coord, ((row, col), value) - struct_re_index = dict(gen_coords()) + else: + if struct.subs is not None: + all_subs = struct.subs else: - struct_re_index = {} + all_subs = None + + def subs_func(degree, idx_eq, idx_col, all_subs=all_subs, struct=struct): - all_subs = struct_subs | added_subs.get(degree, {}) + if struct.re_index_func is not None: + re_index = struct.re_index_func(degree, idx_eq, idx_col) + else: + re_index = None + + if re_index is None: + n_coord = (idx_eq, idx_col) + factor = 1 + else: + n_coord = (re_index[0], re_index[1]) + factor = re_index[2] - yield (struct, degree), acc, (all_subs, struct_re_index) + if all_subs is not None and degree in all_subs: + all_degree_subs = all_subs[degree] - # def gen_predefined_coord(): - # re_index = re_index_dict[struct].get(degree, {}) + if n_coord in all_degree_subs: + subs_val = all_degree_subs[n_coord] * factor + else: + subs_val = None - # for coord_dict in (all_subs, re_index): - # for coord in coord_dict.keys(): - # # in case the n_var is smaller than expected - # if max(coord) <= self.n_var: - # yield coord + elif struct.subs_func is not None: + sub_value = struct.subs_func(degree, n_coord[0], n_coord[1]) - # n_predefined = len(set(gen_predefined_coord())) + if sub_value is not None: + subs_val = sub_value * factor + else: + subs_val = None - if struct.is_vector: - number_of_terms = int(self.n_var * binom(self.n_var+degree-1, degree)) else: - number_of_terms = int(self.n_var**2 * binom(self.n_var+degree-1, degree)) + subs_val = None - acc += number_of_terms + return n_coord, factor, subs_val - yield None, acc, None + yield struct, subs_func - param_list = list(gen_n_param_per_struct()) - n_param = param_list[-1][1] - offset_dict = dict(((e[0], e[1]) for e in param_list[:-1])) - subs_dict = dict(((e[0], e[2][0]) for e in param_list[:-1])) - re_index_dict = dict(((e[0], e[2][1]) for e in param_list[:-1])) + subs_dict = dict(gen_substitutions()) terms = collections.defaultdict(lambda: collections.defaultdict(lambda: collections.defaultdict(int))) for left, right in self.terms: + + subs_1 = subs_dict[left] + subs_2 = subs_dict[right] + for d1 in left.degrees: n_param_1 = binom(self.n_var+d1-1, d1) offset_1 = offset_dict[(left, d1)] - subs_1 = subs_dict[(left, d1)] - re_index_1 = re_index_dict[(left, d1)] for d2 in right.degrees: n_param_2 = binom(self.n_var+d2-1, d2) offset_2 = offset_dict[(right, d2)] - subs_2 = subs_dict[(right, d2)] - re_index_2 = re_index_dict[(right, d2)] total_degree = d1 + d2 @@ -202,67 +200,73 @@ class EquationMixin(abc.ABC): if non_increasing(grp1) and non_increasing(grp2): left_idx = variable_to_index(self.n_var, grp1) - right_idx = variable_to_index(self.n_var, grp2) + d_right_idx = variable_to_index(self.n_var, grp2) # for each column of the poly matrix, and row of the poly vector for idx_col in range(self.n_var): - coord_2 = (idx_col, 0) - n_coord_2, factor_2 = re_index_2.get(coord_2, (coord_2, 1)) - - if factor_2 == 0: - continue + if right.re_index_func_2 is None: + re_index_2 = None + else: + re_index_2 = right.re_index_func_2(d2, idx_col, grp2) - if n_coord_2 in subs_2: - subs_val_2 = factor_2 * subs_2[n_coord_2] - if subs_val_2 == 0: - continue + if re_index_2 is None: + v_idx_row = idx_col + factor_22 = 1 + right_idx = d_right_idx else: - subs_val_2 = None + v_idx_row, n_grp2, factor_22 = re_index_2 + right_idx = variable_to_index(self.n_var, n_grp2) + + n_coord_2, factor_21, subs_val_2 = subs_2(d2, v_idx_row, 0) + + factor_2 = factor_21 * factor_22 + + if factor_2 == 0 or subs_val_2 == 0: + continue # for each polynomial equation for idx_eq in range(self.n_var): - coord_1 = (idx_eq, idx_col) - n_coord_1, factor_1 = re_index_1.get(coord_1, (coord_1, 1)) + n_coord_1, factor_1, subs_val_1 = subs_1(d1, idx_eq, idx_col) - if factor_1 == 0: + if factor_1 == 0 or subs_val_1 == 0: continue - - if n_coord_1 in subs_1: - subs_val_1 = factor_1 * subs_1[n_coord_1] - if subs_val_1 == 0: - continue - else: - subs_val_1 = None left_param_idx = int(offset_1 + left_idx + (self.n_var * n_coord_1[0] + n_coord_1[1]) * n_param_1) right_param_idx = int(offset_2 + right_idx + n_coord_2[0] * n_param_2) match (subs_val_1, subs_val_2): case (None, None): - row_idx = variable_to_index(n_param, (left_param_idx, right_param_idx)) + col_idx = variable_to_index(n_param, (left_param_idx, right_param_idx)) + # col_idx = (left_param_idx, right_param_idx) degree = 2 value = factor_1*factor_2 + + terms[1][(idx_eq, perm)][left_param_idx] += 2*value + terms[1][(idx_eq, perm)][right_param_idx] += 2*value + terms[0][(idx_eq, perm)][0] += value case (subs_val, None): - # row_idx = variable_to_index(n_param, (right_param_idx,)) - row_idx = right_param_idx + col_idx = right_param_idx + # col_idx = (right_param_idx,) degree = 1 value = subs_val_1*factor_2 + terms[degree][(idx_eq, perm)][0] += value + case (None, subs_val): - # row_idx = variable_to_index(n_param, (left_param_idx,)) - row_idx = left_param_idx + col_idx = left_param_idx + # col_idx = (left_param_idx,) degree = 1 value = subs_val_2*factor_1 case _: - degree, row_idx, value = 0, 0, subs_val_1*subs_val_2 + degree, col_idx, value = 0, tuple(), subs_val_1*subs_val_2*factor_1*factor_2 - terms[degree][(idx_eq, perm)][row_idx] += value + terms[degree][(idx_eq, perm)][col_idx] += value - return terms, offset_dict + return terms, offset_dict, n_param ######################################## # Classes @@ -286,6 +290,7 @@ class PolyMatrixImpl(PolyMatrix): subs_func: typing.Callable[[int, int, int], tuple[float]] re_index: dict[int, dict[tuple[int, int], tuple[int, int, float]]] re_index_func: typing.Callable[[int, int], tuple[int, int, float]] + re_index_func_2: typing.Callable[[int, int, tuple[int, ...]], tuple[int, tuple[int, ...], float]] is_vector: bool @@ -304,6 +309,7 @@ def init_poly_vector( subs_func: typing.Callable[[int, int, int], tuple[float]] = None, re_index: dict[int, dict[tuple[int, int], tuple[int, int, float]]] = None, re_index_func: typing.Callable[[int, int], tuple[int, int, float]] = None, + re_index_func_2: typing.Callable[[int, int], tuple[int, int, float]] = None, ): return init_poly_matrix( degrees=degrees, @@ -312,6 +318,7 @@ def init_poly_vector( re_index=re_index, re_index_func=re_index_func, is_vector=True, + re_index_func_2=re_index_func_2, ) @@ -321,6 +328,7 @@ def init_poly_matrix( subs_func: typing.Callable[[int, int, int], tuple[float]] = None, re_index: dict[int, dict[tuple[int, int], tuple[int, int, float]]] = None, re_index_func: typing.Callable[[int, int], tuple[int, int, float]] = None, + re_index_func_2: typing.Callable[[int, int], tuple[int, int, float]] = None, is_vector: bool = None, ): if degrees is None: @@ -342,6 +350,7 @@ def init_poly_matrix( subs_func=subs_func, re_index = re_index, re_index_func = re_index_func, + re_index_func_2 = re_index_func_2, is_vector = is_vector, ) diff --git a/test_polymatrix/test_polymatrix.py b/test_polymatrix/test_polymatrix.py index 1f0bb73..3b239bf 100644 --- a/test_polymatrix/test_polymatrix.py +++ b/test_polymatrix/test_polymatrix.py @@ -135,12 +135,12 @@ class TestPolyMatrix(unittest.TestCase): value = 2, ) - def test_param_matrix_const_vector_skew_symmetric(self): + def test_skew_symmetric_param_matrix_const_vector(self): """ param = [a11 a21 a31 a41] """ - def skew_symmetric(idx1, idx2): + def skew_symmetric(degree, idx1, idx2): if idx1 == idx2: return idx1, idx2, 0 elif idx2 < idx1: @@ -173,4 +173,65 @@ class TestPolyMatrix(unittest.TestCase): eq_idx = 1, row_idx = offset_dict[(mat, 0)] + 1, value = -1, - )
\ No newline at end of file + ) + + def test_const_matrix_param_gradient_vector(self): + """ + param = [v11 v12 v21 v22] + """ + + def gradient(degree, v_row, monom): + if degree == 1: + factor = sum(v_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 v_row <= monom[-1]: + n_v_row = v_row + n_monom = monom + + return n_v_row, n_monom, factor + + mat = init_poly_matrix( + subs={0: {(0, 0): 1, (1, 0): 1, (0, 1): 1, (1, 1): 1}}, + ) + vec = init_poly_vector( + degrees=(1,), + re_index_func_2=gradient, + ) + + eq = init_equation( + terms = [(mat, vec)], + n_var = 2, + ) + + terms, offset_dict = list(eq.create()) + + self.assert_term_in_eq( + terms = terms, + degree = 1, + monoms=(0,), + eq_idx = 0, + row_idx = offset_dict[(vec, 1)], + value = 2, + ) + + self.assert_term_in_eq( + terms = terms, + degree = 1, + monoms=(1,), + eq_idx = 0, + row_idx = offset_dict[(vec, 1)]+1, + value = 1, + ) + + self.assert_term_in_eq( + terms = terms, + degree = 1, + monoms=(0,), + eq_idx = 1, + row_idx = offset_dict[(vec, 1)]+1, + value = 1, + ) |