Make factorize_linear_constraints faster

desc/objectives/utils.py

@@ -97,38 +97,25 @@ def factorize_linear_constraints(objective, constraint, x_scale="auto"):  # noqa
     # which are duplicate rows of A that also have duplicate entries of b,
     # if the entries of b aren't the same then the constraints are actually
     # incompatible and so we will leave those to be caught later.
-    A_augmented = np.hstack([A, np.reshape(b, (A.shape[0], 1))])
-    row_idx_to_delete = np.array([], dtype=int)
-    for row_idx in range(A_augmented.shape[0]):
-        # find all rows equal to this row
-        rows_equal_to_this_row = np.where(
-            np.all(A_augmented[row_idx, :] == A_augmented, axis=1)
-        )[0]
-        # find the rows equal to this row that are not this row
-        rows_equal_to_this_row_but_not_this_row = rows_equal_to_this_row[
-            rows_equal_to_this_row != row_idx
-        ]
-        # if there are rows equal to this row that aren't this row, AND this particular
-        # row has not already been detected as a duplicate of an earlier one and slated
-        # for deletion, add the duplicate row indices to the array of
-        # rows to be deleted
-        if (
-            rows_equal_to_this_row_but_not_this_row.size
-            and row_idx not in row_idx_to_delete
-        ):
-            row_idx_to_delete = np.append(row_idx_to_delete, rows_equal_to_this_row[1:])
-    # delete the affected rows, and also the corresponding rows of b
-    A_augmented = np.delete(A_augmented, row_idx_to_delete, axis=0)
+    A_augmented = jnp.hstack([A, jnp.reshape(b, (A.shape[0], 1))])
+
+    # Find unique rows of A_augmented
+    unique_rows, unique_indices = jnp.unique(A_augmented, axis=0, return_index=True)
+
+    # Sort the indices to preserve the order of appearance
+    unique_indices = jnp.sort(unique_indices)
+
+    # while loop has problems updating JAX arrays, convert them to numpy arrays
+    A_augmented = np.array(A_augmented)
+    # Extract the unique rows
+    A_augmented = A_augmented[unique_indices]
     A = A_augmented[:, :-1]
     b = np.atleast_1d(A_augmented[:, -1].squeeze())
 
     # will store the global index of the unfixed rows, idx
     indices_row = np.arange(A.shape[0])
     indices_idx = np.arange(A.shape[1])
 
-    # while loop has problems updating JAX arrays, convert them to numpy arrays
-    A = np.array(A)
-    b = np.array(b)
     while len(np.where(np.count_nonzero(A, axis=1) == 1)[0]):
         # fixed just means there is a single element in A, so A_ij*x_j = b_i
         fixed_rows = np.where(np.count_nonzero(A, axis=1) == 1)[0]