Fix bug caused by non-scalar variables

1 files changed, 21 insertions, 15 deletions

diff --git a/sumofsquares/solver/cvxopt.py b/sumofsquares/solver/cvxopt.py
index 0e5ce37..c4e1db3 100644
--- a/sumofsquares/solver/cvxopt.py
+++ b/sumofsquares/solver/cvxopt.py
@@ -3,6 +3,7 @@ Solve sumofsquares problems using CVXOPT
 """
 import cvxopt
 import numpy as np
+import math
 
 from collections import UserDict
 from numpy.typing import NDArray
@@ -65,17 +66,15 @@ def solve_sos_cone(prob: Problem, verbose: bool = False,
     G_rows, h_rows= [], []
     l_dims, q_dims, s_dims = [], [], []
 
-    # Collect variable indices in dictionary and sort them by value
-    # i.e. we want variable_indices.values() to be sorted
-    variable_indices: dict[OptVariable, VariableIndex] = dict(sorted({
-        # FIXME there can be more than one indices associated to v
-        v: next(prob.state.get_indices_as_variable_index(v))
-        for v in prob.variables
-    }.items(), key=lambda item: item[1]))
+    # indices of variables in the optimization problem
+    variable_indices = sum(sorted(tuple(prob.state.get_indices_as_variable_index(v))
+                                  for v in prob.variables), ())
 
     # cost linear term
-    cost_coeffs = dict(cost.affine_coefficients_by_degrees(variable_indices.values()))
-    q = cost_coeffs[1].T
+    cost_coeffs = dict(cost.affine_coefficients_by_degrees(variable_indices))
+    q = cost_coeffs.get(1) 
+    if q is None:
+        q = np.zeros((len(variable_indices), 1))
 
     # cost quadratic term
     if is_qp:
@@ -114,10 +113,9 @@ def solve_sos_cone(prob: Problem, verbose: bool = False,
 
             # Linear term
             l_stacked_rows = []
-            for v in variable_indices.values():
+            for v in variable_indices:
                 # Get the affine coefficient associated to this variable
-                linear = MonomialIndex((v,))
-                v_coeff = constr.affine_coefficient(linear)
+                v_coeff = constr.affine_coefficient(v)
 
                 # stack the columns of coeff matrix into a fat row vector (column major order).
                 v_stacked = v_coeff.T.reshape((nrows * ncols,))
@@ -145,7 +143,7 @@ def solve_sos_cone(prob: Problem, verbose: bool = False,
                                       "into the canonical form of CVXOPT.")
 
     # Convert to CVXOPT matrices
-    q = cvxopt.matrix(q)
+    q = cvxopt.matrix(q.T)
     G = cvxopt.matrix(np.hstack(G_rows).T) if G_rows else None
     h = cvxopt.matrix(np.hstack(h_rows)) if h_rows else None
     A = cvxopt.matrix(np.hstack(A_rows).T) if A_rows else None
@@ -177,7 +175,15 @@ def solve_sos_cone(prob: Problem, verbose: bool = False,
                                      *args, **kwargs)
 
     results = {}
-    for i, v in enumerate(variable_indices.keys()):
-        results[v] = info['x'][i]
+
+    i = 0
+    for variable in prob.variables:
+        num_indices = math.prod(variable.shape)
+        values = np.array(info['x'][i:i+num_indices]).reshape(variable.shape)
+        if values.shape == (1, 1):
+            values = values[0, 0]
+
+        results[variable] = values
+        i += num_indices 
 
     return results, CVXOPTInfo(info)