Added type hints

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)

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.

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))