Initial commit

33 files changed, 1069 insertions, 0 deletions

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..d7f42da
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,110 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+.static_storage/
+.media/
+local_settings.py
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+
+.idea/
+
+.vscode
+.pytest_cache
+
diff --git a/README.md b/README.md
new file mode 100644
index 0000000..5d159c7
--- /dev/null
+++ b/README.md
@@ -0,0 +1 @@
+# Sum of Sqaures
\ No newline at end of file
diff --git a/setup.py b/setup.py
new file mode 100644
index 0000000..fb413e4
--- /dev/null
+++ b/setup.py
@@ -0,0 +1,7 @@
+from setuptools import setup, find_packages
+
+setup(
+    name='sumofsquares',
+    packages=find_packages(include=['sumofsquares', 'sumofsquares.*']),
+    version='0.0.0',
+)
\ No newline at end of file
diff --git a/sumofsquares/__init__.py b/sumofsquares/__init__.py
new file mode 100644
index 0000000..b44da40
--- /dev/null
+++ b/sumofsquares/__init__.py
@@ -0,0 +1,6 @@
+from sumofsquares.abc.sosexpr import ParamSOSExpr, SOSExpr
+
+from sumofsquares.init.initsosexpr import init_sos_expr, init_param_expr, init_param_expr_from_reference
+from sumofsquares.init.initenergyfunction import init_energy_function
+
+from sumofsquares.cvxopt import solve_cone, solve_sos_problem
diff --git a/sumofsquares/abc/__init__.py b/sumofsquares/abc/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/sumofsquares/abc/__init__.py
diff --git a/sumofsquares/abc/sosexpr.py b/sumofsquares/abc/sosexpr.py
new file mode 100644
index 0000000..03e0b5d
--- /dev/null
+++ b/sumofsquares/abc/sosexpr.py
@@ -0,0 +1,12 @@
+import abc
+
+from sumofsquares.mixins.parametermixin import ParamSOSExprMixin
+from sumofsquares.mixins.sosexpropmixin import SOSExprOPMixin
+
+
+class SOSExpr(SOSExprOPMixin, abc.ABC):
+    pass
+
+
+class ParamSOSExpr(SOSExprOPMixin, ParamSOSExprMixin, abc.ABC):
+    pass
diff --git a/sumofsquares/cvxopt.py b/sumofsquares/cvxopt.py
new file mode 100644
index 0000000..72dd869
--- /dev/null
+++ b/sumofsquares/cvxopt.py
@@ -0,0 +1,216 @@
+import dataclasses
+import cvxopt
+import polymatrix
+import numpy as np
+import math
+from sumofsquares.abc.sosexpr import ParamSOSExpr
+
+from sumofsquares.sosconstraint.abc import SOSConstraint
+
+
+@dataclasses.dataclass
+class CVXOptConeQPResult:
+    x: np.ndarray
+    y: np.ndarray
+    s: np.ndarray
+    z: np.ndarray
+    status: str
+    gap: float
+    relative_gap: float
+    primal_objective: float
+    dual_objective: float
+    primal_infeasibility: float
+    dual_infeasibility: float
+    primal_slack: float
+    dual_slack: float
+    iterations: int
+    x0: dict
+
+
+
+def solve_cone(state, variables, cost, s_inequality, l_inequality = tuple(), previous_result=None, print_info=False, print_matrix=False):
+    state, matrix_repr = polymatrix.to_matrix_repr(
+        (cost,) + l_inequality + s_inequality,
+        polymatrix.v_stack(variables),
+    ).apply(state)
+    
+    # maximum degree of cost function must be 2
+    max_degree_cost = matrix_repr.data[0].get_max_degree()
+    assert max_degree_cost <= 2, f'{max_degree_cost=}'
+
+    # maximum degree of constraint must be 1
+    max_degree_constraints = max(data.get_max_degree() for data in matrix_repr.data[1:])
+    assert max_degree_constraints == 1, f'{max_degree_constraints=}'
+
+    cost = matrix_repr.data[0]
+
+    # cost[1] is a 1 x n vector
+    q = cost[1] / 2
+        
+    Gl = matrix_repr.data[1:1+len(l_inequality)]     
+    l = sum(G[0].shape[0] for G in Gl)
+
+    Gs = matrix_repr.data[1+len(l_inequality):]  
+    
+    def gen_square_matrix_dim():
+        for Gs_raw in Gs:
+            dim = np.sqrt(Gs_raw[0].shape[0])
+
+            assert math.isclose(int(dim), dim), f'{dim=}'
+            
+            yield int(dim)
+    
+    s = list(gen_square_matrix_dim())
+    
+    G = np.vstack(tuple(-v[1] for v in matrix_repr.data[1:]))
+    h = np.vstack(tuple(v[0] for v in matrix_repr.data[1:]))
+
+    if print_info:
+        print(f'number of variables: {G.shape[1]}')
+        print(f'PSD matrix dimensions: {s}')
+
+    if print_matrix:
+        print(f'cost={q}')
+        print(f'G={G}')
+        print(f'h={h}')
+    
+    if previous_result is None:
+        primalstart = None
+        
+    else:
+        primalstart = {
+            'x': cvxopt.matrix(previous_result.x),
+            'y': cvxopt.matrix(previous_result.y),
+            's': cvxopt.matrix(previous_result.s),
+            'z': cvxopt.matrix(previous_result.z),
+        }
+
+    if max_degree_cost == 1:
+        return_val = cvxopt.solvers.conelp(
+            c=cvxopt.matrix(q.reshape(-1, 1)),
+            G=cvxopt.matrix(G), 
+            h=cvxopt.matrix(h),
+            dims={'l': 0, 'q': [], 's': s},
+            primalstart=primalstart,
+        )
+        
+    else:
+        # cost[2] is a 1 x n*n vector
+        nP = int(np.sqrt(cost[2].shape[1]))
+        P = cost[2].reshape(nP, nP).toarray()
+        
+        return_val = cvxopt.solvers.coneqp(
+            P=cvxopt.matrix(P), 
+            q=cvxopt.matrix(q.reshape(-1, 1)),
+            G=cvxopt.matrix(G), 
+            h=cvxopt.matrix(h),
+            dims={'l': l, 'q': [], 's': s,},
+            initvals=primalstart,
+        )
+        
+    x = np.array(return_val['x']).reshape(-1)
+
+    if x[0] == None:
+        x0 = {}
+    else:
+        x0 = {var: matrix_repr.get_value(var, x) for var in variables}
+    
+    result = CVXOptConeQPResult(
+        x=x,
+        y=np.array(return_val['y']).reshape(-1),
+        s=np.array(return_val['s']).reshape(-1),
+        z=np.array(return_val['z']).reshape(-1),
+        status=return_val['status'],
+        gap=return_val['gap'],
+        relative_gap=return_val['relative gap'],
+        primal_objective=return_val['primal objective'],
+        dual_objective=return_val['dual objective'],
+        primal_infeasibility=return_val['primal infeasibility'],
+        dual_infeasibility=return_val['dual infeasibility'],
+        primal_slack=return_val['primal slack'],
+        dual_slack=return_val['dual slack'],
+        iterations=return_val['iterations'],
+        x0=x0
+    )
+
+    return state, result
+
+    
+def solve_sos_problem(
+    cost: tuple[polymatrix.Expression], 
+    sos_constraints: tuple[SOSConstraint],
+    state: polymatrix.ExpressionState,
+    free_param: tuple[ParamSOSExpr] = None, 
+    x0: dict[ParamSOSExpr, np.ndarray] = None,
+):
+    if x0 is None:
+        x0 = {}
+    
+    def gen_all_param_expr():
+        for sos_constraint in sos_constraints:
+            for param_expr in sos_constraint.dependence:
+                yield param_expr
+
+    all_param_expr = tuple(set(gen_all_param_expr()))
+
+    if free_param is None:
+        free_param = all_param_expr
+    
+    sub_vals = tuple((param_expr.param, x0[param_expr.param]) for param_expr in all_param_expr if param_expr not in free_param)
+    
+    def gen_inequality():
+        for sos_constraint in sos_constraints:
+            if any(param_expr in free_param for param_expr in sos_constraint.dependence):
+                yield sos_constraint.constraint.eval(sub_vals)
+    
+    inequality = tuple(gen_inequality())
+    
+    state, result = solve_cone(
+        state,
+        variables=tuple(param_expr.param for param_expr in free_param),
+        cost=sum(cost),
+        s_inequality=inequality,
+    )
+        
+    return state, result
+
+
+# def solve_sos_problem(
+#     cost: tuple[polymatrix.Expression], 
+#     sos_constraints: tuple[SOSConstraint],
+#     state: polymatrix.ExpressionState,
+#     subs: tuple[ParamSOSExpr] = None, 
+#     x0: dict[ParamSOSExpr, np.ndarray] = None,
+# ):
+#     if x0 is None:
+#         x0 = {}
+    
+#     if subs is None:
+#         subs = tuple()
+        
+#     def gen_free_param_expr():
+#         for sos_constraint in sos_constraints:
+#             for param_expr in sos_constraint.dependence:
+#                 if param_expr not in subs:
+#                     yield param_expr
+                    
+#     free_param_expr = tuple(set(gen_free_param_expr()))
+    
+#     sub_vals = tuple((sub.param, x0[sub.param]) for sub in subs)
+    
+#     def gen_inequality():
+#         for sos_constraint in sos_constraints:
+#             if any(param_expr in free_param_expr for param_expr in sos_constraint.dependence):
+#                 # print(tuple(d.name for d in sos_constraint.dependence))
+#                 yield sos_constraint.constraint.eval(sub_vals)
+    
+#     inequality = tuple(gen_inequality())
+    
+#     state, result = solve_cone(
+#         state,
+#         variables=tuple(param_expr.param for param_expr in free_param_expr),
+#         cost=sum(cost),
+#         s_inequality=inequality,
+#     )
+        
+#     return state, result
diff --git a/sumofsquares/impl/__init__.py b/sumofsquares/impl/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/sumofsquares/impl/__init__.py
diff --git a/sumofsquares/impl/energyfunctionimpl.py b/sumofsquares/impl/energyfunctionimpl.py
new file mode 100644
index 0000000..058ae58
--- /dev/null
+++ b/sumofsquares/impl/energyfunctionimpl.py
@@ -0,0 +1,19 @@
+import dataclassabc
+import polymatrix
+
+from sumofsquares.abc.sosexpr import ParamSOSExpr, SOSExpr
+from sumofsquares.sosconstraint.abc import SOSConstraint
+from sumofsquares.mixins.energyfunctionmixin import EnergyFunctionMixin
+
+
+@dataclassabc.dataclassabc(frozen=True)
+class EnergyFunctionImpl(EnergyFunctionMixin):
+    name: str
+    h: SOSExpr
+    h_dot: SOSExpr
+    epsilon: ParamSOSExpr
+    gamma_b: ParamSOSExpr
+    gamma_roi: ParamSOSExpr | None
+    gamma_w_inner: ParamSOSExpr | None
+    gamma_w_outer: ParamSOSExpr | None
+    sos_constraints: tuple[SOSConstraint]
diff --git a/sumofsquares/impl/sosexprimpl.py b/sumofsquares/impl/sosexprimpl.py
new file mode 100644
index 0000000..c48143b
--- /dev/null
+++ b/sumofsquares/impl/sosexprimpl.py
@@ -0,0 +1,27 @@
+import dataclassabc
+
+from sumofsquares.abc.sosexpr import ParamSOSExpr, SOSExpr
+from sumofsquares.sosexprbase.abc import ParamSOSExprBase, SOSExprBase
+
+
+@dataclassabc.dataclassabc(frozen=True)
+class SOSExprImpl(SOSExpr):
+    underlying: SOSExprBase
+
+
+@dataclassabc.dataclassabc(frozen=True)
+class ParamSOSExprImpl(ParamSOSExpr):
+    underlying: ParamSOSExprBase
+
+    def __eq__(self, other):
+        if isinstance(other, ParamSOSExprImpl):
+            return self.underlying == other.underlying
+        
+        elif isinstance(other, ParamSOSExprBase):
+            return self.underlying == other
+        
+        else:
+            return False
+    
+    def __hash__(self):
+        return hash(self.underlying)
diff --git a/sumofsquares/init/initenergyfunction.py b/sumofsquares/init/initenergyfunction.py
new file mode 100644
index 0000000..8b2f1bf
--- /dev/null
+++ b/sumofsquares/init/initenergyfunction.py
@@ -0,0 +1,101 @@
+import polymatrix
+
+from sumofsquares.abc.sosexpr import SOSExpr
+from sumofsquares.impl.energyfunctionimpl import EnergyFunctionImpl
+from sumofsquares.init.initsosexpr import init_param_expr, init_param_expr_from_reference, init_sos_expr
+
+
+def init_energy_function(
+    name: str,
+    h: SOSExpr,
+    x_dot: SOSExpr,
+    positive_b: bool,
+    epsilon_min: SOSExpr | None = None,
+    roi: SOSExpr | polymatrix.Expression | None = None,
+    decrease_rate: SOSExpr | polymatrix.Expression | None = None,
+    w_outer: SOSExpr | polymatrix.Expression | None = None,
+    w_inner: SOSExpr | polymatrix.Expression | None = None,
+):
+    sos_constraints = tuple()
+    
+    nh = h.diff().T
+    h_dot = (nh.T @ x_dot).cache()
+
+    gamma_b = init_param_expr_from_reference(
+        name=f'gamma_b_{name}',
+        reference=h_dot, 
+        multiplicand=h,
+    )
+    if positive_b:
+        sos_constraints += gamma_b.sos_constraints
+
+    energy_expr = -h_dot
+    energy_expr -= gamma_b * h
+
+    if epsilon_min is None:
+        epsilon = None
+
+    else:
+        epsilon = init_param_expr(
+            name=f'epsilon_{name}',
+            variables=h.variables,
+        )
+
+        sos_constraints += (epsilon - epsilon_min).sos_constraints
+        energy_expr += epsilon
+
+    if roi is None:
+        gamma_roi = None,
+
+    else:
+        gamma_roi = init_param_expr_from_reference(
+            name=f'gamma_roi_{name}',
+            reference=h_dot, 
+            multiplicand=roi,
+        )
+        sos_constraints += gamma_roi.sos_constraints
+
+        energy_expr += gamma_roi * roi
+
+    if decrease_rate is not None:    
+        energy_expr -= decrease_rate
+
+    sos_constraints += energy_expr.sos_constraints
+
+    if w_outer is None:
+        gamma_w_outer = None
+
+    else:
+        gamma_w_outer = init_param_expr_from_reference(
+            name=f'gamma_w_outer_{name}',
+            reference=h, 
+            multiplicand=w_outer,
+        )
+
+        sos_constraints += gamma_w_outer.sos_constraints
+        sos_constraints += (h - gamma_w_outer * w_outer).sos_constraints
+
+    if w_inner is None:
+        gamma_w_inner = None
+
+    else:
+        gamma_w_inner = init_param_expr_from_reference(
+            name=f'gamma_w_inner_{name}',
+            reference=h, 
+            multiplicand=w_inner,
+        )
+
+        sos_constraints += gamma_w_inner.sos_constraints
+        sos_constraints += (gamma_w_inner * w_inner - h).sos_constraints
+
+    return EnergyFunctionImpl(
+        name=name,
+        h=h,
+        h_dot=h_dot,
+        epsilon=epsilon,
+        gamma_roi=gamma_roi,
+        gamma_b=gamma_b,
+        gamma_w_outer=gamma_w_outer,
+        gamma_w_inner=gamma_w_inner,
+        sos_constraints=sos_constraints,
+    )
diff --git a/sumofsquares/init/initsosexpr.py b/sumofsquares/init/initsosexpr.py
new file mode 100644
index 0000000..c858549
--- /dev/null
+++ b/sumofsquares/init/initsosexpr.py
@@ -0,0 +1,83 @@
+import polymatrix
+from sumofsquares.abc.sosexpr import SOSExpr
+
+from sumofsquares.impl.sosexprimpl import ParamSOSExprImpl, SOSExprImpl
+from sumofsquares.mixins.sosexprmixin import SOSExprMixin
+from sumofsquares.sosexprbase.init.initsosexprbase import init_param_sos_expr_base, init_sos_expr_base
+from sumofsquares.sosexprbase.mixins.parametermixin import ParameterMixin
+
+
+def init_sos_expr(
+    expr: polymatrix.Expression, 
+    variables: polymatrix.Expression,
+    dependence: tuple[ParameterMixin],
+):
+    return SOSExprImpl(
+        underlying=init_sos_expr_base(
+            expr=expr,
+            variables=variables,
+            dependence=dependence,
+        ),
+    )
+
+
+def init_param_expr(
+    name: str, 
+    variables: polymatrix.Expression,
+    monom: polymatrix.Expression | None = None,
+    n_rows: int | None = None,
+):
+    return ParamSOSExprImpl(
+        underlying=init_param_sos_expr_base(
+            name=name,
+            monom=monom,
+            variables=variables,
+            n_rows=n_rows,
+        ),
+    )
+
+
+def init_param_expr_from_reference(
+    name: str, 
+    reference: SOSExpr,
+    # variables: polymatrix.Expression,
+    multiplicand: SOSExpr | polymatrix.Expression | None = None,
+):
+    variables = reference.variables
+
+    if multiplicand is None:
+        multiplicand_expr = polymatrix.from_(1)
+
+    elif isinstance(multiplicand, polymatrix.Expression):
+        multiplicand_expr = multiplicand
+
+    elif isinstance(multiplicand, SOSExprMixin):
+        assert multiplicand.variables == variables, f'{multiplicand.variables=}, {variables=}'
+
+        multiplicand_expr = multiplicand.expr
+
+    else:
+        multiplicand_expr = polymatrix.from_(multiplicand)
+
+    m_sos_monom = multiplicand_expr.quadratic_monomials(variables)
+    
+    max_degree = m_sos_monom.max_degree().T.max()
+    
+    m_max_monom = m_sos_monom.filter(
+        m_sos_monom.max_degree() - max_degree, 
+        inverse=True,
+    )
+    
+    sos_monom = reference.expr.quadratic_monomials(variables).subtract_monomials(m_max_monom)
+
+    expr = (sos_monom @ sos_monom.T).reshape(1, -1).sum()
+
+    monom = expr.linear_monomials(variables).cache()
+
+    return init_param_expr(
+        name=name,
+        monom=monom,
+        variables=variables,
+    )
+
+
diff --git a/sumofsquares/mixins/__init__.py b/sumofsquares/mixins/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/sumofsquares/mixins/__init__.py
diff --git a/sumofsquares/mixins/energyfunctionmixin.py b/sumofsquares/mixins/energyfunctionmixin.py
new file mode 100644
index 0000000..f73f225
--- /dev/null
+++ b/sumofsquares/mixins/energyfunctionmixin.py
@@ -0,0 +1,53 @@
+import abc
+import polymatrix
+
+from sumofsquares.abc.sosexpr import ParamSOSExpr, SOSExpr
+from sumofsquares.mixins.getsosconstraintmixin import GetSOSConstraintMixin
+from sumofsquares.sosconstraint.abc import SOSConstraint
+
+
+class EnergyFunctionMixin(GetSOSConstraintMixin, abc.ABC):
+    @property
+    @abc.abstractmethod
+    def name(self) -> str:
+        ...
+
+    @property
+    @abc.abstractmethod
+    def h(self) -> SOSExpr:
+        ...
+
+    @property
+    @abc.abstractmethod
+    def h_dot(self) -> SOSExpr:
+        ...
+
+    # @property
+    # @abc.abstractmethod
+    # def variables(self) -> polymatrix.Expression:
+    #     ...
+
+    @property
+    @abc.abstractmethod
+    def epsilon(self) -> ParamSOSExpr:
+        ...
+
+    @property
+    @abc.abstractmethod
+    def gamma_b(self) -> ParamSOSExpr:
+        ...
+
+    @property
+    @abc.abstractmethod
+    def gamma_roi(self) -> ParamSOSExpr | None:
+        ...
+
+    @property
+    @abc.abstractmethod
+    def gamma_w_outer(self) -> ParamSOSExpr | None:
+        ...
+
+    @property
+    @abc.abstractmethod
+    def gamma_w_inner(self) -> ParamSOSExpr | None:
+        ...
diff --git a/sumofsquares/mixins/getsosconstraintmixin.py b/sumofsquares/mixins/getsosconstraintmixin.py
new file mode 100644
index 0000000..23f3813
--- /dev/null
+++ b/sumofsquares/mixins/getsosconstraintmixin.py
@@ -0,0 +1,10 @@
+import abc
+
+from sumofsquares.sosconstraint.abc import SOSConstraint
+
+
+class GetSOSConstraintMixin(abc.ABC):
+    @property
+    @abc.abstractmethod
+    def sos_constraints(self) -> tuple[SOSConstraint]:
+        ...
diff --git a/sumofsquares/mixins/parametermixin.py b/sumofsquares/mixins/parametermixin.py
new file mode 100644
index 0000000..5281f35
--- /dev/null
+++ b/sumofsquares/mixins/parametermixin.py
@@ -0,0 +1,25 @@
+import abc
+import polymatrix
+
+from sumofsquares.mixins.sosexprmixin import SOSExprMixin
+from sumofsquares.sosexprbase.abc import ParamSOSExprBase
+from sumofsquares.sosexprbase.mixins.parametermixin import ParameterMixin
+
+
+class ParamSOSExprMixin(ParameterMixin, SOSExprMixin):
+    @property
+    @abc.abstractmethod
+    def underlying(self) -> ParamSOSExprBase:
+        ...
+
+    @property
+    def name(self) -> polymatrix.Expression:
+        return self.underlying.name
+
+    @property
+    def param(self) -> polymatrix.Expression:
+        return self.underlying.param
+
+    @property
+    def monom(self) -> polymatrix.Expression:
+        return self.underlying.monom
diff --git a/sumofsquares/mixins/sosexprmixin.py b/sumofsquares/mixins/sosexprmixin.py
new file mode 100644
index 0000000..d0c0d17
--- /dev/null
+++ b/sumofsquares/mixins/sosexprmixin.py
@@ -0,0 +1,37 @@
+import abc
+import polymatrix
+
+from sumofsquares.sosconstraint.abc import SOSConstraint
+from sumofsquares.sosconstraint.initsosconstraint import init_sos_constraint
+from sumofsquares.mixins.getsosconstraintmixin import GetSOSConstraintMixin
+from sumofsquares.sosexprbase.abc import SOSExprBase
+
+
+class SOSExprMixin(GetSOSConstraintMixin, abc.ABC):
+    @property
+    @abc.abstractmethod
+    def underlying(self) -> SOSExprBase:
+        ...
+
+    @property
+    def expr(self) -> polymatrix.Expression:
+        return self.underlying.expr
+
+    @property
+    def variables(self) -> polymatrix.Expression:
+        return self.underlying.variables
+
+    @property
+    def dependence(self) -> polymatrix.Expression:
+        return self.underlying.dependence
+
+    @property
+    def sos_matrix(self) -> polymatrix.Expression:
+        return self.underlying.sos_matrix
+
+    @property
+    def sos_constraints(self) -> tuple[SOSConstraint]:
+        return (init_sos_constraint(
+            dependence=self.dependence,
+            constraint=self.underlying.sos_matrix_as_vector,
+        ),)
diff --git a/sumofsquares/mixins/sosexpropmixin.py b/sumofsquares/mixins/sosexpropmixin.py
new file mode 100644
index 0000000..a345284
--- /dev/null
+++ b/sumofsquares/mixins/sosexpropmixin.py
@@ -0,0 +1,133 @@
+import dataclasses
+import typing
+
+import polymatrix
+import polymatrix.expression.init.initfromsympyexpr
+from polymatrix.expression.expression import Expression, ExpressionImpl
+
+from sumofsquares.mixins.sosexprmixin import SOSExprMixin
+from sumofsquares.sosexprbase.init.initsosexprbase import init_sos_expr_base
+
+
+class SOSExprOPMixin(SOSExprMixin):
+    @staticmethod
+    def _binary(
+        op, 
+        left: 'SOSExprOPMixin', 
+        right: typing.Union[polymatrix.Expression, 'SOSExprOPMixin'],
+    ) -> 'SOSExprOPMixin':
+
+        if not isinstance(left, SOSExprOPMixin):
+            left = polymatrix.expression.init.initfromsympyexpr.init_from_expr_or_none(left)
+
+            if left is None:
+                return NotImplemented
+
+            underlying = init_sos_expr_base(
+                expr=op(polymatrix.from_(left), right.expr),
+                variables=right.variables,
+                dependence=right.dependence,
+            )
+
+            return dataclasses.replace(
+                right,
+                underlying=underlying,
+            )
+
+        elif not isinstance(right, SOSExprOPMixin):
+            right = polymatrix.expression.init.initfromsympyexpr.init_from_expr_or_none(right)
+
+            if right is None:
+                return NotImplemented
+            
+            underlying = init_sos_expr_base(
+                expr=op(left.expr, polymatrix.from_(right)),
+                variables=left.variables,
+                dependence=left.dependence,
+            )
+
+        else:
+            assert left.variables == right.variables, f'{left.variables=}, {right.variables=}'
+
+            underlying=init_sos_expr_base(
+                expr=op(left.expr, right.expr),
+                variables=left.variables,
+                dependence=tuple(set(left.dependence + right.dependence)),
+            )
+
+        return dataclasses.replace(
+            left,
+            underlying=underlying,
+        )
+
+    @staticmethod
+    def _unary(op, expr: 'SOSExprOPMixin') -> 'SOSExprOPMixin':
+        return dataclasses.replace(
+            expr,
+            underlying=init_sos_expr_base(
+                expr=op(expr.expr),
+                variables=expr.variables,
+                dependence=expr.dependence,
+            ),
+        )
+
+    def __add__(self, other: typing.Union[polymatrix.Expression, 'SOSExprOPMixin']) -> 'SOSExprOPMixin':
+        return self._binary(Expression.__add__, self, other)
+
+    def __matmul__(self, other: typing.Union[polymatrix.Expression, 'SOSExprOPMixin']) -> 'SOSExprOPMixin':
+        return self._binary(Expression.__matmul__, self, other)
+
+    def __mul__(self, other: typing.Union[polymatrix.Expression, 'SOSExprOPMixin']) -> 'SOSExprOPMixin':
+        return self._binary(Expression.__mul__, self, other)
+    
+    def __neg__(self) -> 'SOSExprOPMixin':
+        return self._unary(Expression.__neg__, self)
+
+    def __radd__(self, other):
+        return self._binary(Expression.__add__, other, self)
+
+    def __rmatmul__(self, other: typing.Union[polymatrix.Expression, 'SOSExprOPMixin']) -> 'SOSExprOPMixin':
+        return self._binary(Expression.__matmul__, other, self)
+    
+    def __rmul__(self, other):
+        return self._binary(Expression.__mul__, other, self)
+
+    def __rsub__(self, other: typing.Union[polymatrix.Expression, 'SOSExprOPMixin']) -> 'SOSExprOPMixin':
+        return self._binary(Expression.__sub__, other, self)
+
+    def __sub__(self, other: typing.Union[polymatrix.Expression, 'SOSExprOPMixin']) -> 'SOSExprOPMixin':
+        return self._binary(Expression.__sub__, self, other)
+
+    def cache(self) -> 'SOSExprOPMixin':
+        return self._unary(Expression.cache, self)
+
+    def diff(
+        self, 
+        variables: polymatrix.Expression | None = None,
+    ) -> 'SOSExprOPMixin':
+        if variables is None:
+            variables = self.variables
+
+        return dataclasses.replace(
+            self,
+            underlying=init_sos_expr_base(
+                expr=self.expr.diff(variables),
+                variables=self.variables,
+                dependence=self.dependence,
+            ),
+        )
+
+    @property
+    def T(self):
+        return self._unary(Expression.transpose, self)
+
+    def substitute(self, substitutions, variables):
+        return dataclasses.replace(
+            self,
+            underlying=init_sos_expr_base(
+                expr=self.expr.substitute(substitutions),
+                # variables=self.variables.substitute(substitutions),
+                variables=variables,
+                dependence=self.dependence,
+            ),
+        )
diff --git a/sumofsquares/sosconstraint/__init__.py b/sumofsquares/sosconstraint/__init__.py
new file mode 100644
index 0000000..8b13789
--- /dev/null
+++ b/sumofsquares/sosconstraint/__init__.py
@@ -0,0 +1 @@
+
diff --git a/sumofsquares/sosconstraint/abc.py b/sumofsquares/sosconstraint/abc.py
new file mode 100644
index 0000000..5ce0b47
--- /dev/null
+++ b/sumofsquares/sosconstraint/abc.py
@@ -0,0 +1,5 @@
+from sumofsquares.sosconstraint.sosconstraintmixin import SOSConstraintMixin
+
+
+class SOSConstraint(SOSConstraintMixin):
+    pass
diff --git a/sumofsquares/sosconstraint/impl.py b/sumofsquares/sosconstraint/impl.py
new file mode 100644
index 0000000..e0d6cec
--- /dev/null
+++ b/sumofsquares/sosconstraint/impl.py
@@ -0,0 +1,11 @@
+import dataclassabc
+import polymatrix
+
+from sumofsquares.sosexprbase.mixins.parametermixin import ParameterMixin
+from sumofsquares.sosconstraint.abc import SOSConstraint
+
+
+@dataclassabc.dataclassabc
+class SOSConstraintImpl(SOSConstraint):
+    dependence: tuple[ParameterMixin]
+    constraint: polymatrix.Expression
\ No newline at end of file
diff --git a/sumofsquares/sosconstraint/initsosconstraint.py b/sumofsquares/sosconstraint/initsosconstraint.py
new file mode 100644
index 0000000..68a2f25
--- /dev/null
+++ b/sumofsquares/sosconstraint/initsosconstraint.py
@@ -0,0 +1,14 @@
+import polymatrix
+
+from sumofsquares.sosexprbase.mixins.parametermixin import ParameterMixin
+from sumofsquares.sosconstraint.impl import SOSConstraintImpl
+
+
+def init_sos_constraint(
+    dependence: tuple[ParameterMixin],
+    constraint: polymatrix.Expression
+):
+    return SOSConstraintImpl(
+        dependence=dependence,
+        constraint=constraint,
+    )
diff --git a/sumofsquares/sosconstraint/sosconstraintmixin.py b/sumofsquares/sosconstraint/sosconstraintmixin.py
new file mode 100644
index 0000000..be49620
--- /dev/null
+++ b/sumofsquares/sosconstraint/sosconstraintmixin.py
@@ -0,0 +1,16 @@
+import abc
+import polymatrix
+
+from sumofsquares.sosexprbase.mixins.parametermixin import ParameterMixin
+
+
+class SOSConstraintMixin(abc.ABC):
+    @property
+    @abc.abstractmethod
+    def dependence(self) -> tuple[ParameterMixin]:
+        ...
+
+    @property
+    @abc.abstractmethod
+    def constraint(self) -> polymatrix.Expression:
+        ...
diff --git a/sumofsquares/sosexprbase/__init__.py b/sumofsquares/sosexprbase/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/sumofsquares/sosexprbase/__init__.py
diff --git a/sumofsquares/sosexprbase/abc.py b/sumofsquares/sosexprbase/abc.py
new file mode 100644
index 0000000..c979919
--- /dev/null
+++ b/sumofsquares/sosexprbase/abc.py
@@ -0,0 +1,12 @@
+import abc
+from sumofsquares.sosexprbase.mixins.selfdependencemixin import SelfDependenceMixin
+from sumofsquares.sosexprbase.mixins.sosexprbasemixin import SOSExprBaseMixin
+from sumofsquares.sosexprbase.mixins.exprfrommonommixin import ExprFromMonomMixin
+
+
+class SOSExprBase(SOSExprBaseMixin, abc.ABC):
+    pass
+
+
+class ParamSOSExprBase(SelfDependenceMixin, SOSExprBaseMixin, ExprFromMonomMixin, abc.ABC):
+    pass
diff --git a/sumofsquares/sosexprbase/impl.py b/sumofsquares/sosexprbase/impl.py
new file mode 100644
index 0000000..02f220f
--- /dev/null
+++ b/sumofsquares/sosexprbase/impl.py
@@ -0,0 +1,21 @@
+import dataclassabc
+import polymatrix
+
+from sumofsquares.sosexprbase.abc import ParamSOSExprBase, SOSExprBase
+from sumofsquares.sosexprbase.mixins.parametermixin import ParameterMixin
+
+
+@dataclassabc.dataclassabc(frozen=True)
+class SOSExprBaseImpl(SOSExprBase):
+    expr: polymatrix.Expression
+    variables: polymatrix.Expression
+    dependence: tuple[ParameterMixin]
+
+
+@dataclassabc.dataclassabc(frozen=True)
+class ParamSOSExprBaseImpl(ParamSOSExprBase):
+    name: str
+    param: polymatrix.Expression
+    monom: polymatrix.Expression
+    variables: polymatrix.Expression
+    param_matrix: polymatrix.Expression
diff --git a/sumofsquares/sosexprbase/init/initsosexprbase.py b/sumofsquares/sosexprbase/init/initsosexprbase.py
new file mode 100644
index 0000000..6b7f299
--- /dev/null
+++ b/sumofsquares/sosexprbase/init/initsosexprbase.py
@@ -0,0 +1,56 @@
+import polymatrix
+
+from sumofsquares.sosexprbase.impl import ParamSOSExprBaseImpl, SOSExprBaseImpl
+from sumofsquares.sosexprbase.mixins.parametermixin import ParameterMixin
+
+
+def init_sos_expr_base(
+    expr: polymatrix.Expression, 
+    variables: polymatrix.Expression,
+    dependence: tuple[ParameterMixin] | None = None,
+):
+    
+    if not isinstance(expr, polymatrix.Expression):
+        expr = polymatrix.from_(expr)
+
+    # if variables is None:
+    #     variables = polymatrix.from_(1)
+
+    if dependence is None:
+        dependence = tuple()
+
+    return SOSExprBaseImpl(
+        expr=expr,
+        variables=variables,
+        dependence=dependence,
+    )
+
+
+def init_param_sos_expr_base(
+    name: str, 
+    variables: polymatrix.Expression,
+    monom: polymatrix.Expression | None = None,
+    n_rows: int | None = None,
+):
+    if monom is None:
+        monom = polymatrix.from_(1)
+
+    if n_rows is None:
+        n_rows = 1
+
+    if n_rows is 1:
+        param = monom.parametrize(f'{name}')
+        param_matrix = param.T
+
+    else:
+        params = tuple(monom.parametrize(f'{name}_{row+1}') for row in range(n_rows))
+        param = polymatrix.v_stack(params)
+        param_matrix = polymatrix.v_stack(tuple(param.T for param in params))
+
+    return ParamSOSExprBaseImpl(
+        name=name,
+        param=param,
+        monom=monom,
+        variables=variables,
+        param_matrix=param_matrix,
+    )
diff --git a/sumofsquares/sosexprbase/mixins/dependencemixin.py b/sumofsquares/sosexprbase/mixins/dependencemixin.py
new file mode 100644
index 0000000..fa60f19
--- /dev/null
+++ b/sumofsquares/sosexprbase/mixins/dependencemixin.py
@@ -0,0 +1,10 @@
+import abc
+
+from sumofsquares.sosexprbase.mixins.parametermixin import ParameterMixin
+
+class DependenceMixin(abc.ABC):
+    @property
+    @abc.abstractmethod
+    def dependence(self) -> tuple[ParameterMixin]:
+        ...
+
diff --git a/sumofsquares/sosexprbase/mixins/exprbasemixin.py b/sumofsquares/sosexprbase/mixins/exprbasemixin.py
new file mode 100644
index 0000000..9d1e30a
--- /dev/null
+++ b/sumofsquares/sosexprbase/mixins/exprbasemixin.py
@@ -0,0 +1,16 @@
+import abc
+import polymatrix
+
+from sumofsquares.sosexprbase.mixins.dependencemixin import DependenceMixin
+
+
+class ExprBaseMixin(DependenceMixin):
+    @property
+    @abc.abstractmethod
+    def expr(self) -> polymatrix.Expression:
+        ...
+
+    @property
+    @abc.abstractmethod
+    def variables(self) -> polymatrix.Expression:
+        ...
diff --git a/sumofsquares/sosexprbase/mixins/exprfrommonommixin.py b/sumofsquares/sosexprbase/mixins/exprfrommonommixin.py
new file mode 100644
index 0000000..354f46f
--- /dev/null
+++ b/sumofsquares/sosexprbase/mixins/exprfrommonommixin.py
@@ -0,0 +1,16 @@
+import abc
+import polymatrix
+
+from sumofsquares.sosexprbase.mixins.exprbasemixin import ExprBaseMixin
+from sumofsquares.sosexprbase.mixins.parametermixin import ParameterMixin
+
+
+class ExprFromMonomMixin(ParameterMixin, ExprBaseMixin):
+    @property
+    @abc.abstractmethod
+    def param_matrix(self) -> polymatrix.Expression:
+        ...
+
+    @property
+    def expr(self) -> polymatrix.Expression:
+        return (self.param_matrix @ self.monom).cache()
diff --git a/sumofsquares/sosexprbase/mixins/parametermixin.py b/sumofsquares/sosexprbase/mixins/parametermixin.py
new file mode 100644
index 0000000..3df6ab4
--- /dev/null
+++ b/sumofsquares/sosexprbase/mixins/parametermixin.py
@@ -0,0 +1,19 @@
+import abc
+import polymatrix
+
+
+class ParameterMixin(abc.ABC):
+    @property
+    @abc.abstractmethod
+    def name(self) -> str:
+        ...
+
+    @property
+    @abc.abstractmethod
+    def param(self) -> polymatrix.Expression:
+        ...
+
+    @property
+    @abc.abstractmethod
+    def monom(self) -> polymatrix.Expression:
+        ...
diff --git a/sumofsquares/sosexprbase/mixins/selfdependencemixin.py b/sumofsquares/sosexprbase/mixins/selfdependencemixin.py
new file mode 100644
index 0000000..8227e48
--- /dev/null
+++ b/sumofsquares/sosexprbase/mixins/selfdependencemixin.py
@@ -0,0 +1,9 @@
+import abc
+
+from sumofsquares.sosexprbase.mixins.dependencemixin import DependenceMixin
+from sumofsquares.sosexprbase.mixins.parametermixin import ParameterMixin
+
+class SelfDependenceMixin(DependenceMixin):
+    @property
+    def dependence(self) -> set[ParameterMixin]:
+        return (self,)
diff --git a/sumofsquares/sosexprbase/mixins/sosexprbasemixin.py b/sumofsquares/sosexprbase/mixins/sosexprbasemixin.py
new file mode 100644
index 0000000..c8cee19
--- /dev/null
+++ b/sumofsquares/sosexprbase/mixins/sosexprbasemixin.py
@@ -0,0 +1,23 @@
+import functools
+import polymatrix
+
+from sumofsquares.sosexprbase.mixins.exprbasemixin import ExprBaseMixin
+
+
+class SOSExprBaseMixin(ExprBaseMixin):
+    @functools.cached_property
+    def sos_monom(self) -> polymatrix.Expression:
+        return self.expr.quadratic_monomials(self.variables).cache()
+
+    @functools.cached_property
+    def sos_matrix(self) -> polymatrix.Expression:
+        sos_matrix = self.expr.quadratic_in(
+            variables=self.variables, 
+            monomials=self.sos_monom,
+        ).symmetric().cache()
+
+        return sos_matrix
+
+    @property
+    def sos_matrix_as_vector(self) -> polymatrix.Expression:
+        return self.sos_matrix.reshape(-1, 1)