diff options
Diffstat (limited to '')
| -rw-r--r-- | polymatrix/polysolver.py | 356 | ||||
| -rw-r--r-- | polymatrix/polystruct.py | 348 | ||||
| -rw-r--r-- | polymatrix/sympyutils.py | 128 | ||||
| -rw-r--r-- | test_polymatrix/test_polymatrix.py | 243 |
4 files changed, 769 insertions, 306 deletions
diff --git a/polymatrix/polysolver.py b/polymatrix/polysolver.py index f611260..24c7095 100644 --- a/polymatrix/polysolver.py +++ b/polymatrix/polysolver.py @@ -1,41 +1,351 @@ import numpy as np +import sympy import itertools +import scipy.linalg +import scipy.sparse.linalg +import collections +from scipy.special import binom +import more_itertools +import time -def quadratic_poly_solver(b, c, m): - n = b.shape[0] - nn = n*n - b_inv = np.linalg.inv(b) +def substitude_x_add_a(data, a): - p = b_inv @ np.ones((n, 1)) + def kron_eye_and_a(n, data, a, indices): + permutations = more_itertools.distinct_permutations(indices) + # permutations = set(itertools.permutations(indices)) + + if not scipy.sparse.issparse(data): + a_mat = np.reshape(a, (-1, 1)) + ops = (np.eye(n), a_mat) + + def acc_kron(perm): + *_, last = itertools.accumulate((ops[idx] for idx in perm), lambda acc, v: np.kron(acc, v)) + return last + + return sum(data @ acc_kron(perm) for perm in permutations) + + else: + # csr_data = data + csr_data = data.tocsr() + + n_eye = sum(1 for idx in indices if idx == 0) + n_col = n**n_eye + n_index = n_col * n**(len(indices) - n_eye) + + def gen_array_per_permuatation(): + + def gen_coord(perm): + n_row = csr_data.shape[0] + + def acc_col_idx_and_value(acc, v): + relindex, relrow, col, val = acc - def square_matrix(p): - return (p @ p.T).reshape(nn, 1) + n_relrow = relrow / n + grp_index = int(relindex / n_relrow) + n_relindex = int(relindex - grp_index * n_relrow) + + if v == 1: + n_col = col + n_val = val * a[grp_index] + else: + n_col = n * col + grp_index + n_val = val + + return n_relindex, n_relrow, n_col, n_val - def merge_matrix(p1, p2): - return (np.kron(np.eye(n), p1) + np.kron(p1, np.eye(n))) @ p2 + for row in range(n_row): + + pt = slice(csr_data.indptr[row], csr_data.indptr[row+1]) + + for inner_idx, array_val in zip(csr_data.indices[pt], csr_data.data[pt]): + + *_, last = itertools.accumulate( + perm, + acc_col_idx_and_value, + initial=(inner_idx, n_index, 0, array_val), + ) + _, _, col, val = last + + yield col, val + + for perm in permutations: + col, data = list(zip(*gen_coord(perm))) + + yield scipy.sparse.csr_array((data, col, csr_data.indptr), shape=(csr_data.shape[0], n_col)) + + return sum(gen_array_per_permuatation()) + + n = data[1].shape[0] + + subs_data = {key: val.copy() for key, val in data.items()} + + for degree, d_data in data.items(): + for i_degree in range(degree): + indices = i_degree*(0,) + (degree - i_degree)*(1,) + update = kron_eye_and_a(n, d_data, a, indices) + subs_data[i_degree] += update + + first = -1/np.squeeze(subs_data[0]) + def gen_subs_data(): + for key, d_data in subs_data.items(): + if scipy.sparse.issparse(d_data): + + csr_data = d_data.tocsr() + for row, scale in enumerate(first): + csr_data.data[csr_data.indptr[row] : csr_data.indptr[row + 1]] *= scale + yield key, csr_data + + else: + d_data *= first[:, None] + yield key, d_data + return dict(gen_subs_data()) + + +def solve_poly_system(data, m): + if isinstance(data[1], np.ndarray): + b_num_inv = np.linalg.inv(data[1]) + else: + b_num_inv = scipy.sparse.linalg.inv(data[1]) + + n = b_num_inv.shape[0] + p = b_num_inv @ np.ones((n, 1)) + + assert (data[0] + np.ones((n, 1)) < 0.01).all(), f'{data[0]=}, {data[0] + np.ones((n, 1))=}' def func(acc, _): + """ + acc: a nxk matrix whose ith column represent the ith parameter of the power series for each of the n variables + """ + p = acc - k = p.shape[1] - uneven = k % 2 - half = int((k - uneven) / 2) + def gen_indices(n, d): + """ + n=2,d=1 = [(0,2), (1,1)] + n=3,d=1 = [(0,3), (1,2)] + n=2,d=2 = [(0,0,2), (0,1,1)] + """ + + if d==0: + yield (n,) + + else: + for idx in range(int(n/(d+1))+1): + yield from ((idx,)+tuple(i+idx*(d-1) for i in indices) for indices in gen_indices( + n - (2*d-1)*idx, + d-1, + )) + + def acc_kron(l): + *_, last = itertools.accumulate(l, lambda acc, v: np.kron(acc, v)) + return last - def gen_pvector(): - if uneven: - yield square_matrix(p[:, half:half+1]) - - for i in range(half): - yield merge_matrix(p[:, i:i+1], p[:, k-i-1:k-i]) + def gen_p(): + for degree, d_data in data.items(): + if 1 < degree and degree-1 <= k: + indices_list = list(gen_indices(k-degree+1, degree-1)) + permutations = (perm for indices in indices_list for perm in more_itertools.distinct_permutations(indices)) + + if not scipy.sparse.issparse(data): + + def acc_kron(perm): + *_, last = itertools.accumulate(((p[:,idx:idx+1] for idx in perm)), lambda acc, v: np.kron(acc, v)) + return last + + yield from (d_data @ acc_kron(perm) for perm in permutations) + + else: + csr_data = d_data.tocsr() + + n_eye = sum(1 for idx in indices if idx == 0) + n_col = n**n_eye + n_index = n_col * n**(len(indices) - n_eye) + + def gen_array_per_permuatation(): + + def gen_coord(perm): + n_row = csr_data.shape[0] + + def acc_col_idx_and_value(acc, v): + relindex, relrow, val = acc + + n_relrow = relrow / n + grp_index = int(relindex / n_relrow) + n_relindex = int(relindex - grp_index * n_relrow) + n_val = val * p[grp_index, v:v+1] + + return n_relindex, n_relrow, n_val + + for row in range(n_row): + + pt = slice(csr_data.indptr[row], csr_data.indptr[row+1]) + + def gen_val_per_row(): + for inner_idx, array_val in zip(csr_data.indices[pt], csr_data.data[pt]): + + *_, last = itertools.accumulate( + perm, + acc_col_idx_and_value, + initial=(inner_idx, n_index, array_val), + ) + _, _, val = last + + yield val + + yield sum(gen_val_per_row()) + + for perm in permutations: + + array_values = list(gen_coord(perm)) + + yield scipy.sparse.csr_array((array_values, np.zeros(len(array_values)), csr_data.indptr), shape=(csr_data.shape[0], n_col)) + + return np.concatenate((p, -b_num_inv @ sum(gen_p())), axis=1) + + *_, sol_subs = itertools.accumulate(range(m-1), func, initial=p) + + # return np.asarray(sol_subs) + return sum(np.asarray(sol_subs).T) + + +def inner_smart_solve(data, irange=None, idegree=None, a_init=None): + + if irange is None: + irange = np.geomspace(20, 1.1, 8) + + if idegree is None: + idegree = 4 + + if a_init is None: + a_init = np.zeros(data[0].shape[0]) + + def acc_func(acc, v): + a, prev_err = acc + scale = v + + # print(f'iter: {index}') + + # print('substitude...') + subs_data = substitude_x_add_a(data, a) + + # print('modify scaling...') + def gen_data(): + for degree, d_data in subs_data.items(): + if 1 == degree: + yield degree, d_data * scale + else: + yield degree, d_data + scaled_data = dict(gen_data()) + + # print('solve...') + try: + sol = solve_poly_system(scaled_data, idegree) + except: + # return a, prev_err + raise + + err = np.max(np.abs(eval_solution(scaled_data, sol))) + + a = sol + a + + return a, err + + a, err = list(zip(*itertools.accumulate( + irange, + acc_func, + initial=(a_init, 0), + ))) + + return a, err + + +def outer_smart_solve(data, n_iter=10, a_max=1.0, irange=None, idegree=None): + def gen_a_err(): + n = data[1].shape[0] + + for idx_iter in range(n_iter): + a_init = a_max * (2 * np.random.rand(n) - 1) + + a, err = inner_smart_solve(data, irange=irange, idegree=idegree, a_init=a_init) + + subs_data = substitude_x_add_a(data, a[-1]) + sol = solve_poly_system(subs_data, 6) + error_index = np.max(np.abs(eval_solution(subs_data, sol))) + max(a[-1]) + # error_index = np.max(np.abs(eval_solution(subs_data, sol))) + + # print(f'{error_index=}') + + yield error_index, a, err + + # if error_index < 1.0: + # break - p_vector = np.add.reduce(list(gen_pvector())) - p_new = b_inv @ c @ p_vector - return np.concatenate((p, p_new), axis=1) + _, a, err = min(gen_a_err()) + + return a, err + + +def eval_solution(data, x=None): + # if x is None: + # x = sum(sol.T) - *_, sol_subs = itertools.accumulate(range(m-1), func, initial=p) + x = np.reshape(x, (-1, 1)) + + def gen_sum(): + n = data[1].shape[0] + + for degree, d_data in data.items(): + if 0 < degree: + # *_, last = itertools.accumulate(degree*(x,), lambda acc, v: np.kron(acc, v)) + # yield d_data @ last + + if not scipy.sparse.issparse(d_data): + *_, last = itertools.accumulate(degree*(x,), lambda acc, v: np.kron(acc, v)) + yield d_data @ last + + else: + csr_data = d_data.tocsr() + + def gen_coord(): + n_row = csr_data.shape[0] + n_index = n**degree + + def acc_col_idx_and_value(acc, v): + relindex, relrow, val = acc - return np.sum(-sol_subs, axis=1, keepdims=True) + n_relrow = relrow / n + grp_index = int(relindex / n_relrow) + n_relindex = int(relindex - grp_index * n_relrow) + + n_val = val * x[grp_index, 0] + + return n_relindex, n_relrow, n_val + + for row in range(n_row): + + pt = slice(csr_data.indptr[row], csr_data.indptr[row+1]) + + for inner_idx, array_val in zip(csr_data.indices[pt], csr_data.data[pt]): + + *_, last = itertools.accumulate( + range(degree), + acc_col_idx_and_value, + initial=(inner_idx, n_index, array_val), + ) + _, _, val = last + + yield val + + array_values = list(gen_coord()) + + yield scipy.sparse.csr_array((array_values, np.zeros(len(array_values)), csr_data.indptr), shape=(csr_data.shape[0], 1)) + + else: + yield d_data + + val = np.ravel(sum(gen_sum()).T) + return val diff --git a/polymatrix/polystruct.py b/polymatrix/polystruct.py index 202f6cc..ad6cb74 100644 --- a/polymatrix/polystruct.py +++ b/polymatrix/polystruct.py @@ -1,15 +1,15 @@ import abc import collections -from pickletools import int4 import typing import numpy as np import dataclasses import dataclass_abc -from scipy.special import binom +import scipy.sparse import itertools import functools +import more_itertools -from polymatrix.utils import variable_to_index +import polymatrix.utils DegreeType = int # CoordType = tuple[int, int] @@ -42,31 +42,25 @@ class PolyMatrixMixin(abc.ABC): @property @abc.abstractmethod - def is_vector(self) -> bool: + def shape(self) -> tuple[int, int]: ... - class EqualityConstraintMixin(abc.ABC): @property @abc.abstractmethod def terms(self) -> dict[int, dict[tuple[int, tuple, int], float]]: ... - # @property - # @abc.abstractmethod - # def offset_dict(self) -> dict[tuple[PolyMatrixMixin, int], int]: - # ... - - # @property - # @abc.abstractmethod - # def param_dict(self) -> dict[tuple[PolyMatrixMixin, int], int]: - # ... - @property @abc.abstractmethod def n_param(self) -> int: ... + @property + @abc.abstractmethod + def variable_to_index(self) -> typing.Callable[[int, tuple[int, ...]], int]: + ... + @functools.cached_property def eq_to_row_index(self): rows_to_eq = list(set((eq_idx, perm) for eq_tuple_degree in self.terms.values() for (eq_idx, perm, var) in eq_tuple_degree.keys())) @@ -77,10 +71,25 @@ class EqualityConstraintMixin(abc.ABC): def n_eq(self): return len(self.eq_to_row_index) - # @property - # @abc.abstractmethod - # def n_var(self) -> int: - # ... + def get_constraint_matrix(self): + def gen_sparse_matrices(): + for degree, degree_tuples in self.terms.items(): + def gen_row_col_data(): + for (idx_eq, perm, variables), value in degree_tuples.items(): + row = self.eq_to_row_index[(idx_eq, perm)] + col = self.variable_to_index(self.n_param, variables) + yield row, col, value + + row, col, data = list(zip(*gen_row_col_data())) + + data = np.array(data, dtype=np.float) + + if degree <= 1: + yield degree, scipy.sparse.coo_array((data, (row, col)), shape=(self.n_eq, self.n_param**degree)).toarray() + else: + yield degree, scipy.sparse.coo_array((data, (row, col)), shape=(self.n_eq, self.n_param**degree)) + + return dict(gen_sparse_matrices()) def get_constraint_func(self): def func(x): @@ -100,7 +109,7 @@ class EqualityConstraintMixin(abc.ABC): for (idx_eq, perm, variables), value in degree_tuples.items(): row_idx = self.eq_to_row_index[(idx_eq, perm)] - vector_val = vector[variable_to_index(self.n_param, variables)] + vector_val = vector[self.variable_to_index(self.n_param, variables)] mat[row_idx] += value * vector_val return mat @@ -131,7 +140,7 @@ class EqualityConstraintMixin(abc.ABC): for var_idx, var in enumerate(variables): other_variables = variables[:var_idx] + variables[var_idx+1:] - vector_val = vector[variable_to_index(self.n_param, other_variables)] + vector_val = vector[self.variable_to_index(self.n_param, other_variables)] # col_idx = variable_to_index(self.n_param, (var,)) jac_mat[row_idx, var] += value*vector_val @@ -168,7 +177,7 @@ class EqualityConstraintMixin(abc.ABC): for var_idx_y, var_y in enumerate(other_variables): other_variables_2 = variables[:var_idx_y] + variables[var_idx_y+1:] - vector_val = vector[variable_to_index(self.n_param, other_variables_2)] + vector_val = vector[self.variable_to_index(self.n_param, other_variables_2)] hess_mat[var_x, var_y] = v[eq_idx]*value*vector_val return hess_mat @@ -196,6 +205,11 @@ class PolyEquationMixin(abc.ABC): ... + @property + @abc.abstractmethod + def variable_to_index(self) -> typing.Callable[[int, tuple[int, ...]], int]: + ... + @functools.cached_property def _param_list(self) -> list[tuple[PolyMatrixMixin, int], int]: """ @@ -206,7 +220,6 @@ class PolyEquationMixin(abc.ABC): all_structs = set(indexed_poly_mat for term in self.terms for indexed_poly_mat in term) def gen_n_param_per_struct(): - # acc = 0 for struct in all_structs: @@ -214,11 +227,10 @@ class PolyEquationMixin(abc.ABC): continue for degree in struct.degrees: + number_of_terms = int(struct.shape[0] * struct.shape[1] * (self.variable_to_index(self.n_var, degree*(self.n_var-1,)) + 1)) + # number_of_terms = int(struct.shape[0] * struct.shape[1] * binom(self.n_var+degree-1, degree)) - 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)) + # print(f'{struct=}, {number_of_terms=}') yield (struct, degree), number_of_terms @@ -272,15 +284,11 @@ class PolyEquationMixin(abc.ABC): self, subs: dict[PolyMatrixMixin, dict[DegreeType, dict[int, int, float]]] = None, ) -> EqualityConstraintMixin: - n_var_2 = self.n_var**2 - if subs is None: added_subs = {} else: added_subs = subs - # binom = scipy.special.binom - # create parameter offset all_structs = set(indexed_poly_mat for term in self.terms for indexed_poly_mat in term) @@ -313,150 +321,117 @@ class PolyEquationMixin(abc.ABC): terms = collections.defaultdict(lambda: collections.defaultdict(float)) - for left, right in self.terms: + def gen_re_indexing(term, degrees, curr_rows_col, monoms, col_defaults, d_subs, d_offsets, perm): + for m, degree, (poly_row, poly_col), monom, col_default, subs, offset in zip( + term, degrees, itertools.pairwise(curr_rows_col), monoms, col_defaults, d_subs, d_offsets): - subs_1 = subs_dict[left] - subs_2 = subs_dict[right] + if m.re_index is not None: + re_index = m.re_index(degree, poly_row, poly_col, monom) + else: + re_index = None - for d1 in left.degrees: + if re_index is None: + col = col_default + new_poly_row = poly_row + new_poly_col = poly_col + factor = 1.0 - if subs_1 is not None and d1 in subs_1: - subs_1_d = subs_1[d1] else: - subs_1_d = None + new_poly_row, new_poly_col, new_monom, factor = re_index + col = self.variable_to_index(self.n_var, new_monom) + + if subs is not None: + try: + subs_val = subs[(new_poly_row, new_poly_col, col)] + except KeyError: + subs_val = None + else: + subs_val = None - offset_1 = self.offset_dict.get((left, d1), 0) + if subs_val is None: + # the coefficient is selected after the reindexing - for d2 in right.degrees: + row = new_poly_row + new_poly_col * m.shape[0] - if subs_2 is not None and d2 in subs_2: - subs_2_d = subs_2[d2] - else: - subs_2_d = None + # linearize parameter matrix + param_idx = int(offset + row + col * m.shape[0] * m.shape[1]) + else: + param_idx = None + + # if poly_row == 0 and perm == (0,) and monom == (0,): + # print(f'{new_poly_row=}') + # print(f'{new_poly_col=}') + # print(f'{col=}') + # print(f'{monom=}') + # print(f'{new_monom=}') + + yield poly_row, poly_col, param_idx, factor, subs_val + + for term in self.terms: - offset_2 = self.offset_dict.get((right, d2), 0) + term_subs = [subs_dict[m] for m in term] + # n_matrices = len(term) - total_degree = d1 + d2 + for degrees in itertools.product(*(m.degrees for m in term)): + total_degree = sum(degrees) + d_subs = [subs[degree] if subs is not None and degree in subs else None for degree, subs in zip(degrees, term_subs)] + d_offsets = [self.offset_dict.get((m, degree), 0) for m, degree in zip(term, degrees)] - for combination in itertools.combinations_with_replacement(range(self.n_var), total_degree): - # n_var = 2, degree = 3 - # cominations = [(0, 0, 0), (0, 0, 1), (0, 1, 1), (1, 1, 1)] + for combination in itertools.combinations_with_replacement(range(self.n_var), total_degree): + # n_var = 2, degree = 3 + # cominations = [(0, 0, 0), (0, 0, 1), (0, 1, 1), (1, 1, 1)] + + for perm in more_itertools.distinct_permutations(combination): + + def acc_func(acc, v): + last, _ = acc + new = last + v + + return new, perm[last:new] + + # perm=(1, 0, 1) -> monom1=(x2, x1), monom2=(x2) + monoms = list(monom for _, monom in itertools.accumulate(degrees, acc_func, initial=(0, None)))[1:] - # def gen_coord_for_single_equation(): - for perm in set(itertools.permutations(combination)): - - # grp1 = (1, 1, 0) -> x1 x2^2 - grp1 = perm[0:d1] - grp2 = perm[d1:] - - def non_increasing(seq): - return all(y <= x for x, y in zip(seq, seq[1:])) - - # (1,0) -> x2*x1 instead of (0,1)->x1*x2 - if non_increasing(grp1) and non_increasing(grp2): - - left_col_default = variable_to_index(self.n_var, grp1) - right_col_default = variable_to_index(self.n_var, grp2) - - # for each column of the poly matrix, and row of the poly vector - for left_poly_col in range(self.n_var): - - right_poly_row = left_poly_col - - if right.re_index is not None: - re_index_2 = right.re_index(d2, right_poly_row, 0, grp2) - else: - re_index_2 = None - - if re_index_2 is None: - n_right_poly_row, factor_2, right_col = right_poly_row, 1, right_col_default - else: - n_right_poly_row, _, n_grp2, factor_2 = re_index_2 - right_col = variable_to_index(self.n_var, n_grp2) - - key_2 = (n_right_poly_row, 0, right_col) - if subs_2_d is not None: - try: - subs_val_2 = subs_2_d[key_2] - except KeyError: - subs_val_2 = None - else: - subs_val_2 = None - - if factor_2 == 0 or subs_val_2 == 0: - continue - - right_row = n_right_poly_row - right_param_idx = int(offset_2 + right_row + right_col * self.n_var) - - # for each polynomial equation - for left_poly_row in range(self.n_var): - - if left.re_index is not None: - re_index_1 = left.re_index(d1, left_poly_row, left_poly_col, grp1) - else: - re_index_1 = None - - if re_index_1 is None: - n_left_poly_row, n_left_poly_col, factor_1 = left_poly_row, left_poly_col, 1 - left_col = left_col_default - else: - n_left_poly_row, n_left_poly_col, n_grp1, factor_1 = re_index_1 - left_col = variable_to_index(self.n_var, n_grp1) - - key_1 = (n_left_poly_row, n_left_poly_col, left_col) - if subs_1_d is not None: - try: - subs_val_1 = subs_1_d[key_1] - except KeyError: - subs_val_1 = None - else: - subs_val_1 = None - - if factor_1 == 0 or subs_val_1 == 0: - continue - - left_row = n_left_poly_row + n_left_poly_col * self.n_var - left_param_idx = int(offset_1 + left_row + left_col * n_var_2) - - # print(left_poly_row) - # print(f'{left_row=}') - # print(f'{left_param_idx=}') - - total_factor = factor_1 * factor_2 - - match (subs_val_1, subs_val_2): - case (None, None): - col_idx = (left_param_idx, right_param_idx) - degree = 2 - value = total_factor - - # 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): - col_idx = (right_param_idx,) - degree = 1 - value = subs_val_1*total_factor - - # terms[degree][(idx_eq, perm)][0] += value - - case (None, subs_val): - col_idx = (left_param_idx,) - degree = 1 - value = subs_val_2*total_factor - - # terms[degree][(idx_eq, perm)][0] += value - - case _: - degree, col_idx, value = 0, tuple(), subs_val_1*subs_val_2*total_factor - - terms[degree][left_poly_row, perm, col_idx] += value + def non_increasing(seq): + return all(y <= x for x, y in zip(seq, seq[1:])) + + # (1,0) -> x2*x1 instead of (0,1)->x1*x2 + if all(non_increasing(monom) for monom in monoms): + + col_defaults = [self.variable_to_index(self.n_var, monom) for monom in monoms] + + n_rows_col = itertools.chain((m.shape[0] for m in term), (term[-1].shape[1],)) + + for curr_rows_col in itertools.product(*[range(e) for e in n_rows_col]): + + data = tuple(gen_re_indexing(term, degrees, curr_rows_col, monoms, col_defaults, d_subs, d_offsets, perm)) + + total_factor = functools.reduce(lambda x, y: x*y, (d[3] for d in data)) + + if total_factor == 0: + continue + + value = functools.reduce(lambda x, y: x*y, (d[4] for d in data if d[4] is not None), 1) * total_factor + + if value == 0: + continue + + poly_row = data[0][0] + param_idx = tuple(d[2] for d in data if d[4] is None) + degree = len(param_idx) + + # if poly_row == 0 and perm == (0,): + # print(f'{param_idx=}') + # print(f'{data=}') + # print(f'{monoms=}') + # print(f'{total_factor=}') + + terms[degree][poly_row, perm, param_idx] += value return EqualityConstraintImpl( terms=terms, n_param=self.n_param, + variable_to_index=self.variable_to_index, ) def matrix_to_poly(self, struct, x, param, tol=None): @@ -526,47 +501,48 @@ class PolyMatrixImpl(PolyMatrix): subs: dict[int, dict[tuple[int, int], float]] re_index: typing.Callable[[int, int, int, tuple[int, ...]], tuple[int, int, int, tuple[int, ...], float]] is_constant: bool - is_vector: bool + shape: tuple[int, int] @dataclass_abc.dataclass_abc(frozen=True) class EqualityConstraintImpl(EqualityConstraintMixin): terms: dict[int, dict[tuple[int, tuple, int], float]] - # offset_dict: dict[tuple[PolyMatrixMixin, int], int] - # param_dict: dict[tuple[PolyMatrixMixin, int], int] + variable_to_index: typing.Callable[[int, tuple[int, ...]], int] n_param: int - # n_var: int @dataclass_abc.dataclass_abc(frozen=True) class EquationImpl(PolyEquation): terms: list[tuple[PolyMatrix, PolyMatrix]] + variable_to_index: typing.Callable[[int, tuple[int, ...]], int] n_var: int ######################################## # init functions ######################################## -def init_poly_vector( - degrees: list[int] = None, - subs: dict[int, dict[tuple[int, int], float]] = None, - re_index: typing.Callable[[int, int, int, tuple[int, ...]], tuple[int, int, int, tuple[int, ...], float]] = None, - is_constant: bool = None, -): - return init_poly_matrix( - degrees=degrees, - subs=subs, - re_index=re_index, - is_vector=True, - is_constant=is_constant, - ) +# def init_poly_vector( +# n_row: int, +# degrees: list[int] = None, +# subs: dict[int, dict[tuple[int, int], float]] = None, +# re_index: typing.Callable[[int, int, int, tuple[int, ...]], tuple[int, int, int, tuple[int, ...], float]] = None, +# is_constant: bool = None, +# ): +# return init_poly_matrix( +# degrees=degrees, +# subs=subs, +# re_index=re_index, +# shape=(n_row, , +# n_col=1, +# is_constant=is_constant, +# ) def init_poly_matrix( + shape: tuple[int, int], degrees: list[int] = None, subs: dict[int, dict[tuple[int, int], float]] = None, re_index: typing.Callable[[int, int, int, tuple[int, ...]], tuple[int, int, int, tuple[int, ...], float]] = None, - is_vector: bool = None, is_constant: bool = None, ): if degrees is None: @@ -579,25 +555,27 @@ def init_poly_matrix( else: is_constant = True - if is_vector is None: - is_vector = False - return PolyMatrixImpl( degrees=degrees, subs=subs, - re_index = re_index, - is_constant = is_constant, - is_vector = is_vector, + re_index=re_index, + is_constant=is_constant, + shape=shape, ) def init_equation( n_var: int, terms: list[tuple[PolyMatrix, PolyMatrix]], + variable_to_index: typing.Callable[[int, tuple[int, ...]], int] = None, ): - assert all(not left.is_vector and right.is_vector for left, right in terms) + # assert all(not left.is_vector and right.is_vector for left, right in terms) + + if variable_to_index is None: + variable_to_index = polymatrix.utils.variable_to_index return EquationImpl( n_var=n_var, terms=terms, + variable_to_index=variable_to_index, ) diff --git a/polymatrix/sympyutils.py b/polymatrix/sympyutils.py index 6fe7938..7413ba3 100644 --- a/polymatrix/sympyutils.py +++ b/polymatrix/sympyutils.py @@ -19,17 +19,13 @@ def poly_to_data_coord(poly_list, x, degree = None): if degree is None: degree = max(degree for inner_poly_list in sympy_poly_list for poly in inner_poly_list for degree in poly.degree_list()) - - # def all_equal(iterable): - # g = itertools.groupby(iterable) - # return next(g, True) and not next(g, False) - - # assert all_equal((p.gens for p in poly_list)), 'all polynomials need to have identical generators' + + n = len(x) def gen_power_mat(): # for all powers generate a matrix - for current_degree in range(degree + 1): + for i_degree in range(degree + 1): def gen_value_index(): @@ -41,45 +37,131 @@ def poly_to_data_coord(poly_list, x, degree = None): # a5 x1 x3**2 -> c=a5, m=(1, 0, 2) for c, m in zip(p.coeffs(), p.monoms()): - if sum(m) == current_degree: + if sum(m) == i_degree: + + indices = (idx for idx, p in enumerate(m) for _ in range(p)) + index = sum(idx*(n**level) for level, idx in enumerate(indices)) - index = variable_powers_to_index(m) yield (poly_row, poly_col, index), c data = dict(gen_value_index()) | collections.defaultdict(float) if len(data) > 0: - yield current_degree, data + yield i_degree, data return dict(gen_power_mat()) def poly_to_matrix(poly_list, x, power = None): - """ - - """ - data_coord_dict = poly_to_data_coord(poly_list, x, power) - # n_free_symbols = len(poly_list[0][0].gens) n_free_symbols = len(x) + n_poly = len(poly_list) + n_eq = sum((1 for row_eq in poly_list for e in row_eq)) def gen_power_mat(): + if 0 not in data_coord_dict: + yield 0, np.zeros((n_free_symbols, 1)) # for all powers generate a matrix - for current_degree, data_coord in data_coord_dict: - + for degree, data_coord in data_coord_dict.items(): + # empty matrix - shape = (len(poly_list), n_free_symbols**current_degree) - # m = np.zeros((len(poly_list), n_free_symbols**current_power)) + shape = (n_eq, n_free_symbols**degree) # fill matrix if len(data_coord) == 0: - yield np.zeros((len(poly_list), n_free_symbols**current_degree)) + yield np.zeros((len(poly_list), n_free_symbols**degree)) else: - rows, cols, data = list(zip(*data_coord)) + def gen_row_col_coord(): + for (p_row, p_col, index), value in data_coord.items(): + yield p_row + p_col * n_poly, index, value + + rows, cols, data = list(zip(*gen_row_col_coord())) + sparse_array = scipy.sparse.coo_matrix((data, (rows, cols)), dtype=np.double, shape=shape) + + if degree <= 1: + yield degree, sparse_array.toarray() + else: + yield degree, sparse_array.tocsr() + + return dict(gen_power_mat()) + + +# def poly_to_data_coord(poly_list, x, degree = None): +# """ +# poly_list = [ +# poly(x1*x3**2, x) +# ] +# power: up to which power +# """ + +# sympy_poly_list = tuple(tuple(sympy.poly(p, x) for p in inner_poly_list) for inner_poly_list in poly_list) + +# if degree is None: +# degree = max(degree for inner_poly_list in sympy_poly_list for poly in inner_poly_list for degree in poly.degree_list()) + +# # def all_equal(iterable): +# # g = itertools.groupby(iterable) +# # return next(g, True) and not next(g, False) + +# # assert all_equal((p.gens for p in poly_list)), 'all polynomials need to have identical generators' + +# def gen_power_mat(): + +# # for all powers generate a matrix +# for current_degree in range(degree + 1): + +# def gen_value_index(): + +# # each polynomial defines a row in the matrix +# for poly_row, inner_poly_list in enumerate(sympy_poly_list): + +# for poly_col, p in enumerate(inner_poly_list): + +# # a5 x1 x3**2 -> c=a5, m=(1, 0, 2) +# for c, m in zip(p.coeffs(), p.monoms()): + +# if sum(m) == current_degree: + +# index = variable_powers_to_index(m) +# yield (poly_row, poly_col, index), c + +# data = dict(gen_value_index()) | collections.defaultdict(float) + +# if len(data) > 0: +# yield current_degree, data + +# return dict(gen_power_mat()) + + +# def poly_to_matrix(poly_list, x, power = None): +# """ + +# """ + +# data_coord_dict = poly_to_data_coord(poly_list, x, power) + +# # n_free_symbols = len(poly_list[0][0].gens) +# n_free_symbols = len(x) + +# def gen_power_mat(): + +# # for all powers generate a matrix +# for current_degree, data_coord in data_coord_dict: + +# # empty matrix +# shape = (len(poly_list), n_free_symbols**current_degree) +# # m = np.zeros((len(poly_list), n_free_symbols**current_power)) + +# # fill matrix +# if len(data_coord) == 0: +# yield np.zeros((len(poly_list), n_free_symbols**current_degree)) + +# else: +# rows, cols, data = list(zip(*data_coord)) - yield scipy.sparse.coo_matrix((data, (rows, cols)), dtype=np.double, shape=shape) +# yield scipy.sparse.coo_matrix((data, (rows, cols)), dtype=np.double, shape=shape) - return list(gen_power_mat()) +# return list(gen_power_mat()) 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, ) |