Create superclass of OptVariableMixin for variables that are not expressions

1 files changed, 21 insertions, 7 deletions

diff --git a/sumofsquares/optvariable.py b/sumofsquares/optvariable.py
index 6e933b6..12df4df 100644
--- a/sumofsquares/optvariable.py
+++ b/sumofsquares/optvariable.py
@@ -1,3 +1,15 @@
+"""
+This module is an extension of the polymatrix package that adds a new type of
+variable, a decision variable, to formulate optimization problems. Speficically,
+this extends the following parts of polymatrix:
+
+    - polymatrix.variable, by adding a new OptVariable type
+    - polymatrix.expression, by adding an OptVariableMixin, so that there can
+      be expression with underlying objects that are OptVariables.
+"""
+
+from __future__ import annotations
+
 from itertools import product
 from typing_extensions import override
 from dataclassabc import dataclassabc
@@ -11,11 +23,16 @@ from polymatrix.polymatrix.index import PolyMatrixDict, PolyDict, MonomialIndex,
 from polymatrix.variable.abc import Variable
 
 
-class OptVariableMixin(ExpressionBaseMixin, Variable):
+class OptVariable(Variable):
     """ Optimization (decision) variable. """
+
+
+class OptVariableMixin(ExpressionBaseMixin, OptVariable):
+    """ Optimization (decision) variable mixin for expression object. """
+
     @override
     def apply(self, state: ExpressionState) -> tuple[ExpressionState, PolyMatrix]:
-        state = state.register(self) 
+        state = state.register(self)
         indices = state.get_indices(self)
         p = PolyMatrixDict()
 
@@ -42,13 +59,10 @@ class OptVariableImpl(OptVariableMixin):
 def init_opt_variable_expr(name, shape):
     return OptVariableImpl(name, shape)
 
-def init_opt_variable_expression(underlying: OptVariableMixin) -> VariableExpression:
-    return init_variable_expression(underlying=underlying)
-
 
-def from_names(names: str, shape: tuple[int, int] = (1,1)) -> tuple[VariableExpression] | VariableExpression:
+def from_names(names: str, shape: tuple[int, int] = (1, 1)) -> tuple[VariableExpression] | VariableExpression:
     """ Construct one or multiple variables from comma separated a list of names. """
-    variables = tuple(init_opt_variable_expression(
+    variables = tuple(init_variable_expression(
         underlying=init_opt_variable_expr(name.strip(), shape))
                       for name in names.split(","))