diff --git a/src/SumOfSquares/SoS.py b/src/SumOfSquares/SoS.py index caaf6dd..b22cf66 100644 --- a/src/SumOfSquares/SoS.py +++ b/src/SumOfSquares/SoS.py @@ -1,3 +1,5 @@ +from __future__ import annotations # for type hinting + import math import picos as pic import sympy as sp @@ -6,6 +8,7 @@ from collections import defaultdict from operator import floordiv, and_ from itertools import combinations +from typing import List, Optional from .util import * from .basis import Basis, poly_variable @@ -24,11 +27,11 @@ def __init__(self, *args, **kwargs): self._sos_const_count = 0 self._pexpect_count = 0 - def __getitem__(self, sym): + def __getitem__(self, sym: sp.Symbol) -> pic.RealVariable: assert isinstance(sym, sp.Symbol), f'{sym} must be a sympy symbol!' return self.sym_to_var(sym) - def sym_to_var(self, sym): + def sym_to_var(self, sym: sp.Symbol) -> pic.RealVariable: '''Map between a sympy symbol to a unique picos variable. As sympy symbols are hashable, each symbol is assigned a unique picos variable. A new picos variable is created if it previously doesn't exist. @@ -37,7 +40,7 @@ def sym_to_var(self, sym): self._sym_var_map[sym] = pic.RealVariable(repr(sym)) return self._sym_var_map[sym] - def sp_to_picos(self, expr): + def sp_to_picos(self, expr: sp.Expr) -> pic.expressions.Expression: '''Converts a sympy affine expression to a picos expression, converting numeric values to floats, and sympy symbols to picos variables. ''' @@ -50,7 +53,7 @@ def sp_to_picos(self, expr): else: return pic.Constant(float(expr)) - def sp_mat_to_picos(self, mat): + def sp_mat_to_picos(self, mat: sp.Matrix) -> pic.expressions.Expression: '''Converts a sympy matrix a picos affine expression, converting numeric values to floats, and sympy symbols to picos variables. ''' @@ -59,7 +62,8 @@ def sp_mat_to_picos(self, mat): return reduce(floordiv, [reduce(and_, map(self.sp_to_picos, mat.row(r))) for r in range(num_rows)]) - def add_sos_constraint(self, expr, variables, name='', sparse=False): + def add_sos_constraint(self, expr: sp.Expr, variables: List[sp.Symbol], + name: str='', sparse: bool=False) -> SOSConstraint: '''Adds a constraint that the polynomial EXPR is a Sum-of-Squares. EXPR is a sympy expression treated as a polynomial in VARIABLES. Any symbols in EXPR not in VARIABLES are converted to picos variables @@ -86,7 +90,9 @@ def add_sos_constraint(self, expr, variables, name='', sparse=False): pic_const = self.add_constraint(Q >> 0) return SOSConstraint(pic_const, Q, basis, variables, deg) - def get_pexpect(self, variables, deg, hom=False, name=''): + def get_pexpect(self, variables: List[sp.Symbol], deg:int, + hom: bool=False, name: str='' + ) -> Callable[[sp.Expr], pic.expressions.Expression]: '''Returns a degree DEG pseudoexpectation operator. This operator is a function that takes in a polynomial of at most degree DEG in VARIABLES, and returns a picos affine expression. If HOM=True, this polynomial must @@ -121,7 +127,13 @@ def pexpect(p): return self.sp_mat_to_picos(basis.sos_sym_poly_repr(poly)) | X return pexpect -def poly_opt_prob(vars, obj, eqs=None, ineqs=None, ineq_prods=False, deg=None, sparse=False): +def poly_opt_prob(vars : List[sp.Symbol], + obj : sp.Expr, + eqs : Optional[List[sp.Expr]] = None, + ineqs : Optional[List[sp.Expr]] = None, + ineq_prods : bool = False, + deg : Optional[int] = None, + sparse : bool = False) -> SOSProblem: '''Formulates and returns a degree DEG Sum-of-Squares relaxation of a polynomial optimization problem in variables VARS that mininizes OBJ subject to equality constraints EQS (g(x) = 0) and inequality constraints INEQS (h(x) @@ -161,7 +173,13 @@ def poly_opt_prob(vars, obj, eqs=None, ineqs=None, ineq_prods=False, deg=None, s prob.set_objective('max', gamma_p) return prob -def poly_cert_prob(vars, poly, eqs=None, ineqs=None, ineq_prods=False, deg=None, sparse=False): +def poly_cert_prob(vars : List[sp.Symbol], + poly : sp.Expr, + eqs : Optional[List[sp.expr]] = None, + ineqs : Optional[List[sp.expr]] = None, + ineq_prods : bool = False, + deg : Optional[int] = None, + sparse : bool = False) -> SOSProblem: '''Formulates and returns a degree DEG Sum-of-Squares relaxation of a polynomial optimization problem in variables VARS that certifies POLY is a sum of squares on the set defined by EQS and INEQS. INEQ_PRODS determines if @@ -203,7 +221,12 @@ class SOSConstraint: and its dual, and allows one to compute the pseudoexpectation of any polynomial. ''' - def __init__(self, pic_const, Q, basis, symbols, deg): + def __init__(self, + pic_const: pic.constraints.Constraint, + Q : pic.expressions.variables.SymmetricVariable, + basis : Basis, + symbols : List[sp.Symbol], + deg : int): self.pic_const = pic_const self.Q = Q self.basis = basis @@ -220,21 +243,21 @@ def Qval(self): ' (is the problem solved?)') return self.Q.value - def get_chol_factor(self): + def get_chol_factor(self) -> np.array: '''Returns L, the Cholesky factorization of Q = LL^T. Adds a small multiple of identity to Q if it has small negative eigenvalues. ''' mineig = min(min(np.linalg.eigh(self.Qval)[0]), 0) return np.linalg.cholesky(self.Qval - np.eye(len(self.basis))*mineig*1.1) - def get_sos_decomp(self, precision=3): + def get_sos_decomp(self, precision: int=3) -> sp.Matrix: '''Returns a vector containing the sum of squares decompositon of this constraint''' L = sp.Matrix(self.get_chol_factor()) S = (L.T @ sp.Matrix(self.b_sym)).applyfunc(lambda x: x**2) return round_sympy_expr(S, precision) - def pexpect(self, expr): + def pexpect(self, expr: sp.Expr) -> sp.Expr: '''Computes the pseudoexpectation of a given polynomial EXPR''' poly = sp.poly(expr, self.symbols) self.basis.check_can_represent(poly) diff --git a/src/SumOfSquares/basis.py b/src/SumOfSquares/basis.py index 6bcc44c..9f7157f 100644 --- a/src/SumOfSquares/basis.py +++ b/src/SumOfSquares/basis.py @@ -1,11 +1,14 @@ +from __future__ import annotations # for type hinting Basis + import sympy as sp import numpy as np import math from collections import defaultdict +from typing import Iterable, Tuple, List from .util import * -def basis_hom(n, d): +def basis_hom(n: int, d: int) -> Iterable[Tuple[int]]: '''Generator for a homogeneous polynomial basis for n variables of degree d, represented as a list of tuples (same as sympy), so that: len(list(basis_hom(n,d))) == binom(n+d-1, d) @@ -19,7 +22,7 @@ def basis_hom(n, d): for b in basis_hom(n-1, di): yield b + (d-di,) -def basis_inhom(n, d): +def basis_inhom(n: int, d: int) -> Iterable[Tuple[int]]: '''Generator for an inhomogeneous polynomial basis for n variables of degree d, represented as a list of tuples (same as sympy), so that: len(list(basis(n,d))) == binom(n+d, d) @@ -28,7 +31,7 @@ def basis_inhom(n, d): yield b[:-1] class Basis(): - def __init__(self, monoms): + def __init__(self, monoms: List[Tuple[int]]): '''Initializes a basis using a sequence of tuples representing monomials ''' self.monoms = monoms @@ -43,17 +46,17 @@ def __init__(self, monoms): for j, bj in enumerate(self): self.sos_sym_entries[sum_tuple(bi, bj)].append((i, j)) - def __len__(self): + def __len__(self) -> int: return len(self.monoms) - def __iter__(self): + def __iter__(self) -> Iterable[Tuple[int]]: return iter(self.monoms) - def from_degree(nvars, deg, hom=False): + def from_degree(nvars: int, deg: int, hom: bool=False) -> Basis: '''Constructs a basis by specifying the number of variables and degree''' return Basis(list((basis_hom if hom else basis_inhom)(nvars, deg))) - def from_poly_lex(poly, sparse=True): + def from_poly_lex(poly: sp.Poly, sparse: bool=True) -> Basis: '''Returns a basis from a polynomial compatible with SoS, ordering monomials in lexicographic order''' poly_deg = poly.total_degree() @@ -75,14 +78,14 @@ def in_hull(pt): # Point lies in affine subspace and convex hull return Basis(list(filter(in_hull, full_basis.monoms))) return full_basis - def to_sym(self, syms): + def to_sym(self, syms: List[sp.Expr]) -> List[sp.Expr]: '''Convert basis to a list of symbolic monomials ''' if self.nvars != len(syms): raise ValueError('Mismatched basis size!') return [prod(s**m for s,m in zip(syms, mono)) for mono in self] - def check_can_represent(self, poly): + def check_can_represent(self, poly: sp.Poly): '''Check if sympy polynomial POLY can be represented by this basis. Raises an error otherwise.''' if poly.total_degree() > self.deg * 2: @@ -97,7 +100,7 @@ def check_can_represent(self, poly): ' are not in basis!') - def sos_sym_poly_repr(self, poly): + def sos_sym_poly_repr(self, poly: sp.Poly) -> sp.Matrix: '''Given a polynomial p, returns a SoS-symmetric representation Qp where p(x) = b(x)^T Qp b(x), in terms of this basis. Qp is returned as a sympy matrix. @@ -110,7 +113,8 @@ def sos_sym_poly_repr(self, poly): return Qp -def poly_variable(name, variables, deg, hom=False): +def poly_variable(name: str, variables: List[sp.Symbol], deg: int, + hom: bool=False) -> sp.Expr: '''Returns a (possibly homogeneous) degree DEG polynomial in VARIABLES, with a variable (a sympy symbol) named using NAME for each coefficient. Used in Sum-of-Squares relaxations for polynomial optimization. diff --git a/src/SumOfSquares/util.py b/src/SumOfSquares/util.py index 6e97f74..49977b2 100644 --- a/src/SumOfSquares/util.py +++ b/src/SumOfSquares/util.py @@ -4,37 +4,38 @@ from operator import mul from functools import reduce from itertools import combinations +from typing import Iterable, List, Optional, Tuple def prod(seq): return reduce(mul, seq, 1) -def factorial(n): +def factorial(n: int) -> int: return prod(range(1, n+1)) -def binom(n, d): +def binom(n: int, d: int) -> int: assert n >= d, f'invalid binom({n}, {d})!' return factorial(n)//factorial(n-d)//factorial(d) -def sum_tuple(t1, t2): +def sum_tuple(t1: tuple, t2: tuple) -> tuple: return tuple(a1+a2 for a1, a2 in zip(t1, t2)) -def is_hom(poly, deg): +def is_hom(poly: sp.Poly, deg: int) -> bool: '''Determines if a polynomial POLY is homogeneous of degree DEG''' return sum(sum(m) != deg for m in poly.monoms()) == 0 -def round_sympy_expr(expr, precision=3): +def round_sympy_expr(expr: sp.Expr, precision: int=3) -> sp.Expr: '''Rounds all numbers in a sympy expression to stated precision''' return expr.xreplace({n : round(n, precision) for n in expr.atoms(sp.Number)}) -def poly_degree(p, variables): +def poly_degree(p: sp.Expr, variables: List[sp.Symbol]) -> int: '''Returns the max degree of P when treated as a polynomial in VARIABLES''' return sp.poly(p, variables).total_degree() -def orth(M): +def orth(M: np.array) -> Tuple[np.array, np.array]: _, D, V = np.linalg.svd(M) return V[D >= 1e-9], V[D < 1e-9] -def get_poly_degree(vars, polys, deg=None): +def get_poly_degree(vars, polys: List[sp.Poly], deg: Optional[int]=None) -> int: '''Given a vector of polynomials POLY, return minimum degree to run sum of squares, or check if DEG is above such minimum degree if provided''' max_deg = max(map(lambda p: poly_degree(p, vars), polys))