update text of separating eq and ineq and switch to logical indexing

python/prima/_linear_constraints.py

@@ -42,29 +42,31 @@ def process_multiple_linear_constraints(constraints):
 
 
 def separate_LC_into_eq_and_ineq(linear_constraint):
-    # Two things:
-    # 1. PRIMA prefers A <= b as opposed to lb <= A <= ub
-    # 2. PRIMA has both A_eq and A_ineq (and b_eq and b_ineq)
+    # The Python interface receives linear constraints lb <= A*x <= ub, but the
+    # Fortran backend of PRIMA expects that the linear constraints are specified
+    # as A_eq*x = b_eq, A_ineq*x <= b_ineq.
     # As such, we must:
-    # 1. Convert lb <= A <= ub to A <= b
-    # 2. Split A <= b into A_eq == b_eq and A_ineq <= b_ineq
-    # Fortunately we can do both at the same time
-    A_eq = []
-    b_eq = []
-    A_ineq = []
-    b_ineq = []
-    for i in range(len(linear_constraint.lb)):
-        if linear_constraint.lb[i] == linear_constraint.ub[i]:
-            A_eq.append(linear_constraint.A[i])
-            b_eq.append(linear_constraint.lb[i])
-        else:
-            A_ineq.append(linear_constraint.A[i])
-            b_ineq.append(linear_constraint.ub[i])
-            # Flip the lb to to take format preferred by PRIMA, as long as it's not -inf
-            if linear_constraint.lb[i] != -np.inf:
-                A_ineq.append(-linear_constraint.A[i])
-                b_ineq.append(-linear_constraint.lb[i])
-    # Convert to numpy arrays, or set to None if empty
+    # 1. for constraints with lb == ub, rewrite them as A_eq*x = lb;
+    # 2. for constraints with lb < ub, rewrite them as A_ineq*x <= b_ineq.
+
+    # We suppose lb == ub if ub <= lb + 2*epsilon, assuming that the preprocessing
+    # ensures lb <= ub.
+    epsilon = np.finfo(np.float64).eps
+
+    eq_indices = linear_constraint.ub <= linear_constraint.lb + 2*epsilon
+    A_eq = linear_constraint.A[eq_indices]
+    b_eq = (linear_constraint.lb[eq_indices] + linear_constraint.ub[eq_indices])/2.0
+
+    ineq_ub_indices = linear_constraint.ub < np.inf
+    A_ineq_ub = linear_constraint.A[~eq_indices & ineq_ub_indices]
+    b_ineq_ub = linear_constraint.ub[~eq_indices & ineq_ub_indices]
+    ineq_lb_indices = linear_constraint.lb > -np.inf
+    A_ineq_lb = -linear_constraint.A[~eq_indices & ineq_lb_indices]
+    b_ineq_lb = -linear_constraint.lb[~eq_indices & ineq_lb_indices]
+    A_ineq = np.concatenate((A_ineq_ub, A_ineq_lb))
+    b_ineq = np.concatenate((b_ineq_ub, b_ineq_lb))
+
+    # Ensure dtype is float64, or set to None if empty
     A_eq = np.array(A_eq, dtype=np.float64) if len(A_eq) > 0 else None
     b_eq = np.array(b_eq, dtype=np.float64) if len(b_eq) > 0 else None
     A_ineq = np.array(A_ineq, dtype=np.float64) if len(A_ineq) > 0 else None