changing integer and boolean indices to use unravel instead of np.ndi…

…ndex (cvxpy#2546)

Co-authored-by: William Zijie Zhang <willizz@Williams-MacBook-Pro.local>

cvxpy/expressions/leaf.py

@@ -124,20 +124,24 @@ def __init__(
                            'complex': complex, 'imag': imag,
                            'symmetric': symmetric, 'diag': diag,
                            'PSD': PSD, 'NSD': NSD,
-                           'hermitian': hermitian, 'boolean': bool(boolean),
+                           'hermitian': hermitian, 'boolean': boolean,
                            'integer':  integer, 'sparsity': sparsity, 'bounds': bounds}
-
-        if boolean:
-            self.boolean_idx = boolean if not isinstance(boolean, bool) else list(
-                np.ndindex(max(shape, (1,))))
-        else:
+        if boolean is True:
+            shape = max(shape, (1,))
+            flat_idx = np.arange(np.prod(shape))
+            self.boolean_idx = np.unravel_index(flat_idx, shape, order='F')
+        elif boolean is False:
             self.boolean_idx = []
-
-        if integer:
-            self.integer_idx = integer if not isinstance(integer, bool) else list(
-                np.ndindex(max(shape, (1,))))
         else:
+            self.boolean_idx = boolean
+        if integer is True:
+            shape = max(shape, (1,))
+            flat_idx = np.arange(np.prod(shape))
+            self.integer_idx = np.unravel_index(flat_idx, shape, order='F')
+        elif integer is False:
             self.integer_idx = []
+        else:
+            self.integer_idx = integer
 
         # Only one attribute be True (except can be boolean and integer).
         true_attr = sum(1 for k, v in self.attributes.items() if v)

cvxpy/reductions/matrix_stuffing.py

@@ -77,28 +77,22 @@ def extract_upper_bounds(variables: list, var_size: int) -> Optional[np.ndarray]
 
 
 def extract_mip_idx(variables) -> Tuple[List[int], List[int]]:
-    """Coalesces bool, int indices for variables.
-
-       The indexing scheme assumes that the variables will be coalesced into
-       a single one-dimensional variable, with each variable being reshaped
-       in Fortran order.
     """
-    def ravel_multi_index(multi_index, x, vert_offset):
-        """Ravel a multi-index and add a vertical offset to it.
-        """
-        ravel_idx = np.ravel_multi_index(multi_index, max(x.shape, (1,)), order='F')
-        return [(vert_offset + idx,) for idx in ravel_idx]
-    boolean_idx = []
-    integer_idx = []
-    vert_offset = 0
+    Coalesces bool, int indices for variables.
+    The indexing scheme assumes that the variables will be coalesced into
+    a single one-dimensional variable, with each variable being reshaped
+    in Fortran order.
+    """
+    boolean_idx, integer_idx, offset = [], [], 0
     for x in variables:
+        ravel_shape = max(x.shape, (1,))
         if x.boolean_idx:
-            multi_index = list(zip(*x.boolean_idx))
-            boolean_idx += ravel_multi_index(multi_index, x, vert_offset)
+            ravel_idx = np.ravel_multi_index(x.boolean_idx, ravel_shape, order='F')
+            boolean_idx += [(idx + offset,) for idx in ravel_idx]
         if x.integer_idx:
-            multi_index = list(zip(*x.integer_idx))
-            integer_idx += ravel_multi_index(multi_index, x, vert_offset)
-        vert_offset += x.size
+            ravel_idx = np.ravel_multi_index(x.integer_idx, ravel_shape, order='F')
+            integer_idx += [(idx + offset,) for idx in ravel_idx]
+        offset += x.size
     return boolean_idx, integer_idx