Use np.ndarray for var_types, basis_status

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(
-                map(
-                    _parse_gurobi_vbasis,
-                    model.getAttr("vbasis", self._gp_vars),
-                )
+            basis_status = np.array(
+                [
+                    _parse_gurobi_vbasis(b)
+                    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"
-            assert var_types[i] in ["B", "C"], (
+                var_types[i] = b"B"
+            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]
                 ),