Use np.ndarray for var_types, basis_status

miplearn/features/extractor.py

@@ -36,7 +36,7 @@ class FeaturesExtractor:
         sample.put_array("static_var_lower_bounds", variables.lower_bounds)
         sample.put_array("static_var_names", variables.names)
         sample.put_array("static_var_obj_coeffs", variables.obj_coeffs)
-        sample.put_vector("static_var_types", variables.types)
+        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("static_constr_lhs", constraints.lhs)
@@ -69,7 +69,7 @@ class FeaturesExtractor:
     ) -> None:
         variables = solver.get_variables(with_static=False, with_sa=self.with_sa)
         constraints = solver.get_constraints(with_static=False, with_sa=self.with_sa)
-        sample.put_vector("lp_var_basis_status", variables.basis_status)
+        sample.put_array("lp_var_basis_status", variables.basis_status)
         sample.put_array("lp_var_reduced_costs", variables.reduced_costs)
         sample.put_array("lp_var_sa_lb_down", variables.sa_lb_down)
         sample.put_array("lp_var_sa_lb_up", variables.sa_lb_up)

miplearn/solvers/gurobi.py

@@ -79,7 +79,7 @@ class GurobiSolver(InternalSolver):
         self._gp_constrs: List["gurobipy.Constr"] = []
         self._var_names: np.ndarray = np.empty(0)
         self._constr_names: List[str] = []
-        self._var_types: List[str] = []
+        self._var_types: np.ndarray = np.empty(0)
         self._var_lbs: np.ndarray = np.empty(0)
         self._var_ubs: np.ndarray = np.empty(0)
         self._var_obj_coeffs: np.ndarray = np.empty(0)
@@ -322,8 +322,9 @@ class GurobiSolver(InternalSolver):
             else:
                 raise Exception(f"unknown vbasis: {basis_status}")
 
+        basis_status: Optional[np.ndarray] = None
         upper_bounds, lower_bounds, types, values = None, None, None, None
-        obj_coeffs, reduced_costs, basis_status = None, None, None
+        obj_coeffs, reduced_costs = None, None
         sa_obj_up, sa_ub_up, sa_lb_up = None, None, None
         sa_obj_down, sa_ub_down, sa_lb_down = None, None, None
 
@@ -335,11 +336,12 @@ class GurobiSolver(InternalSolver):
 
         if self._has_lp_solution:
             reduced_costs = np.array(model.getAttr("rc", self._gp_vars), dtype=float)
-            basis_status = list(
+            basis_status = np.array(
-                map(
+                [
-                    _parse_gurobi_vbasis,
+                    _parse_gurobi_vbasis(b)
-                    model.getAttr("vbasis", self._gp_vars),
+                    for b in model.getAttr("vbasis", self._gp_vars)
-                )
+                ],
+                dtype="S",
             )
 
             if with_sa:
@@ -513,7 +515,7 @@ class GurobiSolver(InternalSolver):
         self._apply_params(streams)
         assert self.model is not None
         for (i, var) in enumerate(self._gp_vars):
-            if self._var_types[i] == "B":
+            if self._var_types[i] == b"B":
                 var.vtype = self.gp.GRB.CONTINUOUS
                 var.lb = 0.0
                 var.ub = 1.0
@@ -521,7 +523,7 @@ class GurobiSolver(InternalSolver):
             self.model.optimize()
             self._dirty = False
         for (i, var) in enumerate(self._gp_vars):
-            if self._var_types[i] == "B":
+            if self._var_types[i] == b"B":
                 var.vtype = self.gp.GRB.BINARY
         log = streams[0].getvalue()
         self._has_lp_solution = self.model.solCount > 0
@@ -590,7 +592,10 @@ class GurobiSolver(InternalSolver):
             self.model.getAttr("varName", gp_vars),
             dtype="S",
         )
-        var_types: List[str] = self.model.getAttr("vtype", gp_vars)
+        var_types: np.ndarray = np.array(
+            self.model.getAttr("vtype", gp_vars),
+            dtype="S",
+        )
         var_ubs: np.ndarray = np.array(
             self.model.getAttr("ub", gp_vars),
             dtype=float,
@@ -611,7 +616,7 @@ class GurobiSolver(InternalSolver):
                 f"Duplicated variable name detected: {var_names[i]}. "
                 f"Unique variable names are currently required."
             )
-            if var_types[i] == "I":
+            if var_types[i] == b"I":
                 assert var_ubs[i] == 1.0, (
                     "Only binary and continuous variables are currently supported. "
                     f"Integer variable {var_names[i]} has upper bound {var_ubs[i]}."
@@ -620,8 +625,8 @@ class GurobiSolver(InternalSolver):
                     "Only binary and continuous variables are currently supported. "
                     f"Integer variable {var_names[i]} has lower bound {var_ubs[i]}."
                 )
-                var_types[i] = "B"
+                var_types[i] = b"B"
-            assert var_types[i] in ["B", "C"], (
+            assert var_types[i] in [b"B", b"C"], (
                 "Only binary and continuous variables are currently supported. "
                 f"Variable {var_names[i]} has type {var_types[i]}."
             )

miplearn/solvers/internal.py

@@ -51,7 +51,7 @@ class MIPSolveStats:
 @dataclass
 class Variables:
     names: Optional[np.ndarray] = None
-    basis_status: Optional[List[str]] = None
+    basis_status: Optional[np.ndarray] = None
     lower_bounds: Optional[np.ndarray] = None
     obj_coeffs: Optional[np.ndarray] = None
     reduced_costs: Optional[np.ndarray] = None
@@ -61,7 +61,7 @@ class Variables:
     sa_obj_up: Optional[np.ndarray] = None
     sa_ub_down: Optional[np.ndarray] = None
     sa_ub_up: Optional[np.ndarray] = None
-    types: Optional[List[str]] = None
+    types: Optional[np.ndarray] = None
     upper_bounds: Optional[np.ndarray] = None
     values: Optional[np.ndarray] = None
 

miplearn/solvers/pyomo/base.py

@@ -327,7 +327,7 @@ class BasePyomoSolver(InternalSolver):
 
         return Variables(
             names=_none_if_empty(np.array(names, dtype="S")),
-            types=_none_if_empty(types),
+            types=_none_if_empty(np.array(types, dtype="S")),
             upper_bounds=_none_if_empty(np.array(upper_bounds, dtype=float)),
             lower_bounds=_none_if_empty(np.array(lower_bounds, dtype=float)),
             obj_coeffs=_none_if_empty(np.array(obj_coeffs, dtype=float)),

miplearn/solvers/tests/__init__.py

@@ -44,7 +44,7 @@ def run_basic_usage_tests(solver: InternalSolver) -> None:
             names=np.array(["x[0]", "x[1]", "x[2]", "x[3]", "z"], dtype="S"),
             lower_bounds=np.array([0.0, 0.0, 0.0, 0.0, 0.0]),
             upper_bounds=np.array([1.0, 1.0, 1.0, 1.0, 67.0]),
-            types=["B", "B", "B", "B", "C"],
+            types=np.array(["B", "B", "B", "B", "C"], dtype="S"),
             obj_coeffs=np.array([505.0, 352.0, 458.0, 220.0, 0.0]),
         ),
     )
@@ -85,7 +85,7 @@ def run_basic_usage_tests(solver: InternalSolver) -> None:
             solver.get_variable_attrs(),
             Variables(
                 names=np.array(["x[0]", "x[1]", "x[2]", "x[3]", "z"], dtype="S"),
-                basis_status=["U", "B", "U", "L", "U"],
+                basis_status=np.array(["U", "B", "U", "L", "U"], dtype="S"),
                 reduced_costs=np.array(
                     [193.615385, 0.0, 187.230769, -23.692308, 13.538462]
                 ),

tests/features/test_extractor.py

@@ -41,7 +41,10 @@ def test_knapsack() -> None:
     assert_equals(
         sample.get_vector("static_var_obj_coeffs"), [505.0, 352.0, 458.0, 220.0, 0.0]
     )
-    assert_equals(sample.get_vector("static_var_types"), ["B", "B", "B", "B", "C"])
+    assert_equals(
+        sample.get_array("static_var_types"),
+        np.array(["B", "B", "B", "B", "C"], dtype="S"),
+    )
     assert_equals(
         sample.get_vector("static_var_upper_bounds"), [1.0, 1.0, 1.0, 1.0, 67.0]
     )
@@ -76,8 +79,8 @@ def test_knapsack() -> None:
     solver.solve_lp()
     extractor.extract_after_lp_features(solver, sample)
     assert_equals(
-        sample.get_vector("lp_var_basis_status"),
+        sample.get_array("lp_var_basis_status"),
-        ["U", "B", "U", "L", "U"],
+        np.array(["U", "B", "U", "L", "U"], dtype="S"),
     )
     assert_equals(
         sample.get_vector("lp_var_reduced_costs"),