Use np.ndarray in Constraints.{basis_status,senses}

miplearn/components/static_lazy.py

@@ -206,7 +206,7 @@ class StaticLazyConstraintsComponent(Component):
         cids: Dict[str, List[str]] = {}
         instance_features = sample.get_vector("static_instance_features")
         constr_features = sample.get_vector_list("lp_constr_features")
-        constr_names = sample.get_vector("static_constr_names")
+        constr_names = sample.get_array("static_constr_names")
         constr_categories = sample.get_vector("static_constr_categories")
         constr_lazy = sample.get_vector("static_constr_lazy")
         lazy_enforced = sample.get_set("mip_constr_lazy_enforced")

miplearn/features/extractor.py

@@ -38,10 +38,10 @@ class FeaturesExtractor:
         sample.put_array("static_var_obj_coeffs", variables.obj_coeffs)
         sample.put_array("static_var_types", variables.types)
         sample.put_array("static_var_upper_bounds", variables.upper_bounds)
-        sample.put_vector("static_constr_names", constraints.names)
+        sample.put_array("static_constr_names", constraints.names)
         # sample.put("static_constr_lhs", constraints.lhs)
         sample.put_array("static_constr_rhs", constraints.rhs)
-        sample.put_vector("static_constr_senses", constraints.senses)
+        sample.put_array("static_constr_senses", constraints.senses)
         vars_features_user, var_categories = self._extract_user_features_vars(
             instance, sample
         )
@@ -78,7 +78,7 @@ class FeaturesExtractor:
         sample.put_array("lp_var_sa_ub_down", variables.sa_ub_down)
         sample.put_array("lp_var_sa_ub_up", variables.sa_ub_up)
         sample.put_array("lp_var_values", variables.values)
-        sample.put_vector("lp_constr_basis_status", constraints.basis_status)
+        sample.put_array("lp_constr_basis_status", constraints.basis_status)
         sample.put_array("lp_constr_dual_values", constraints.dual_values)
         sample.put_array("lp_constr_sa_rhs_down", constraints.sa_rhs_down)
         sample.put_array("lp_constr_sa_rhs_up", constraints.sa_rhs_up)

miplearn/solvers/gurobi.py

@@ -104,12 +104,14 @@ class GurobiSolver(InternalSolver):
             lhs = self.gp.quicksum(
                 self._varname_to_var[varname] * coeff for (varname, coeff) in cf.lhs[i]
             )
-            if sense == "=":
+            if sense == b"=":
                 self.model.addConstr(lhs == cf.rhs[i], name=cf.names[i])
-            elif sense == "<":
+            elif sense == b"<":
                 self.model.addConstr(lhs <= cf.rhs[i], name=cf.names[i])
-            else:
+            elif sense == b">":
                 self.model.addConstr(lhs >= cf.rhs[i], name=cf.names[i])
+            else:
+                raise Exception(f"Unknown sense: {sense}")
         self.model.update()
         self._dirty = True
         self._has_lp_solution = False
@@ -218,7 +220,7 @@ class GurobiSolver(InternalSolver):
 
         if with_static:
             rhs = np.array(model.getAttr("rhs", gp_constrs), dtype=float)
-            senses = model.getAttr("sense", gp_constrs)
+            senses = np.array(model.getAttr("sense", gp_constrs), dtype="S")
             if with_lhs:
                 lhs = [None for _ in gp_constrs]
                 for (i, gp_constr) in enumerate(gp_constrs):
@@ -230,11 +232,9 @@ class GurobiSolver(InternalSolver):
 
         if self._has_lp_solution:
             dual_value = np.array(model.getAttr("pi", gp_constrs), dtype=float)
-            basis_status = list(
-                map(
-                    _parse_gurobi_cbasis,
-                    model.getAttr("cbasis", gp_constrs),
-                )
+            basis_status = np.array(
+                [_parse_gurobi_cbasis(c) for c in model.getAttr("cbasis", gp_constrs)],
+                dtype="S",
             )
             if with_sa:
                 sa_rhs_up = np.array(model.getAttr("saRhsUp", gp_constrs), dtype=float)

miplearn/solvers/internal.py

@@ -68,7 +68,7 @@ class Variables:
 
 @dataclass
 class Constraints:
-    basis_status: Optional[List[str]] = None
+    basis_status: Optional[np.ndarray] = None
     dual_values: Optional[np.ndarray] = None
     lazy: Optional[List[bool]] = None
     lhs: Optional[List[List[Tuple[bytes, float]]]] = None
@@ -76,13 +76,13 @@ class Constraints:
     rhs: Optional[np.ndarray] = None
     sa_rhs_down: Optional[np.ndarray] = None
     sa_rhs_up: Optional[np.ndarray] = None
-    senses: Optional[List[str]] = None
+    senses: Optional[np.ndarray] = None
     slacks: Optional[np.ndarray] = None
 
     @staticmethod
     def from_sample(sample: "Sample") -> "Constraints":
         return Constraints(
-            basis_status=sample.get_vector("lp_constr_basis_status"),
+            basis_status=sample.get_array("lp_constr_basis_status"),
             dual_values=sample.get_vector("lp_constr_dual_values"),
             lazy=sample.get_vector("static_constr_lazy"),
             # lhs=sample.get_vector("static_constr_lhs"),
@@ -90,13 +90,15 @@ class Constraints:
             rhs=sample.get_vector("static_constr_rhs"),
             sa_rhs_down=sample.get_vector("lp_constr_sa_rhs_down"),
             sa_rhs_up=sample.get_vector("lp_constr_sa_rhs_up"),
-            senses=sample.get_vector("static_constr_senses"),
+            senses=sample.get_array("static_constr_senses"),
             slacks=sample.get_vector("lp_constr_slacks"),
         )
 
     def __getitem__(self, selected: List[bool]) -> "Constraints":
         return Constraints(
-            basis_status=self._filter(self.basis_status, selected),
+            basis_status=(
+                None if self.basis_status is None else self.basis_status[selected]
+            ),
             dual_values=(
                 None if self.dual_values is None else self.dual_values[selected]
             ),
@@ -108,7 +110,7 @@ class Constraints:
                 None if self.sa_rhs_down is None else self.sa_rhs_down[selected]
             ),
             sa_rhs_up=(None if self.sa_rhs_up is None else self.sa_rhs_up[selected]),
-            senses=self._filter(self.senses, selected),
+            senses=(None if self.senses is None else self.senses[selected]),
             slacks=(None if self.slacks is None else self.slacks[selected]),
         )
 

miplearn/solvers/pyomo/base.py

@@ -89,12 +89,14 @@ class BasePyomoSolver(InternalSolver):
             for (varname, coeff) in cf.lhs[i]:
                 var = self._varname_to_var[varname]
                 lhs += var * coeff
-            if cf.senses[i] == "=":
+            if cf.senses[i] == b"=":
                 expr = lhs == cf.rhs[i]
-            elif cf.senses[i] == "<":
+            elif cf.senses[i] == b"<":
                 expr = lhs <= cf.rhs[i]
-            else:
+            elif cf.senses[i] == b">":
                 expr = lhs >= cf.rhs[i]
+            else:
+                raise Exception(f"Unknown sense: {cf.senses[i]}")
             cl = pe.Constraint(expr=expr, name=name)
             self.model.add_component(name.decode(), cl)
             self._pyomo_solver.add_constraint(cl)
@@ -235,7 +237,7 @@ class BasePyomoSolver(InternalSolver):
         return Constraints(
             names=_none_if_empty(np.array(names, dtype="S")),
             rhs=_none_if_empty(np.array(rhs, dtype=float)),
-            senses=_none_if_empty(senses),
+            senses=_none_if_empty(np.array(senses, dtype="S")),
             lhs=_none_if_empty(lhs),
             slacks=_none_if_empty(np.array(slacks, dtype=float)),
             dual_values=_none_if_empty(np.array(dual_values, dtype=float)),

miplearn/solvers/tests/__init__.py

@@ -64,7 +64,7 @@ def run_basic_usage_tests(solver: InternalSolver) -> None:
                     (b"z", -1.0),
                 ],
             ],
-            senses=["="],
+            senses=np.array(["="], dtype="S"),
         ),
     )
 
@@ -108,7 +108,7 @@ def run_basic_usage_tests(solver: InternalSolver) -> None:
         _filter_attrs(
             solver.get_constraint_attrs(),
             Constraints(
-                basis_status=["N"],
+                basis_status=np.array(["N"], dtype="S"),
                 dual_values=np.array([13.538462]),
                 names=np.array(["eq_capacity"], dtype="S"),
                 sa_rhs_down=np.array([-24.0]),
@@ -164,7 +164,7 @@ def run_basic_usage_tests(solver: InternalSolver) -> None:
         names=np.array(["cut"], dtype="S"),
         lhs=[[(b"x[0]", 1.0)]],
         rhs=np.array([0.0]),
-        senses=["<"],
+        senses=np.array(["<"], dtype="S"),
     )
     assert_equals(solver.are_constraints_satisfied(cf), [False])
 
@@ -189,7 +189,7 @@ def run_basic_usage_tests(solver: InternalSolver) -> None:
                         (b"x[0]", 1.0),
                     ],
                 ],
-                senses=["=", "<"],
+                senses=np.array(["=", "<"], dtype="S"),
             ),
         ),
     )

tests/features/test_extractor.py

@@ -54,7 +54,7 @@ def test_knapsack() -> None:
     )
     assert sample.get_vector_list("static_var_features") is not None
     assert_equals(
-        sample.get_vector("static_constr_names"),
+        sample.get_array("static_constr_names"),
         np.array(["eq_capacity"], dtype="S"),
     )
     # assert_equals(
@@ -70,7 +70,10 @@ def test_knapsack() -> None:
     #     ],
     # )
     assert_equals(sample.get_vector("static_constr_rhs"), [0.0])
-    assert_equals(sample.get_vector("static_constr_senses"), ["="])
+    assert_equals(
+        sample.get_array("static_constr_senses"),
+        np.array(["="], dtype="S"),
+    )
     assert_equals(sample.get_vector("static_constr_features"), [None])
     assert_equals(
         sample.get_vector("static_constr_categories"),
@@ -114,7 +117,10 @@ def test_knapsack() -> None:
     assert_equals(sample.get_vector("lp_var_sa_ub_up"), [2.043478, inf, 2.2, inf, 69.0])
     assert_equals(sample.get_vector("lp_var_values"), [1.0, 0.923077, 1.0, 0.0, 67.0])
     assert sample.get_vector_list("lp_var_features") is not None
-    assert_equals(sample.get_vector("lp_constr_basis_status"), ["N"])
+    assert_equals(
+        sample.get_array("lp_constr_basis_status"),
+        np.array(["N"], dtype="S"),
+    )
     assert_equals(sample.get_vector("lp_constr_dual_values"), [13.538462])
     assert_equals(sample.get_vector("lp_constr_sa_rhs_down"), [-24.0])
     assert_equals(sample.get_vector("lp_constr_sa_rhs_up"), [2.0])
@@ -132,7 +138,7 @@ def test_constraint_getindex() -> None:
     cf = Constraints(
         names=np.array(["c1", "c2", "c3"], dtype="S"),
         rhs=np.array([1.0, 2.0, 3.0]),
-        senses=["=", "<", ">"],
+        senses=np.array(["=", "<", ">"], dtype="S"),
         lhs=[
             [
                 (b"x1", 1.0),
@@ -153,7 +159,7 @@ def test_constraint_getindex() -> None:
         Constraints(
             names=np.array(["c1", "c3"], dtype="S"),
             rhs=np.array([1.0, 3.0]),
-            senses=["=", ">"],
+            senses=np.array(["=", ">"], dtype="S"),
             lhs=[
                 [
                     (b"x1", 1.0),