Implement more compact get_variables

miplearn/features.py

@@ -6,7 +6,7 @@ import collections
 import numbers
 from dataclasses import dataclass
 from math import log, isfinite
-from typing import TYPE_CHECKING, Dict, Optional, List, Hashable
+from typing import TYPE_CHECKING, Dict, Optional, List, Hashable, Tuple
 
 import numpy as np
 
@@ -30,6 +30,26 @@ class InstanceFeatures:
         return features
 
 
+@dataclass
+class VariableFeatures:
+    names: Optional[Tuple[str, ...]] = None
+    basis_status: Optional[Tuple[str, ...]] = None
+    categories: Optional[Tuple[Hashable, ...]] = None
+    lower_bounds: Optional[Tuple[float, ...]] = None
+    obj_coeffs: Optional[Tuple[float, ...]] = None
+    reduced_costs: Optional[Tuple[float, ...]] = None
+    sa_lb_down: Optional[Tuple[float, ...]] = None
+    sa_lb_up: Optional[Tuple[float, ...]] = None
+    sa_obj_down: Optional[Tuple[float, ...]] = None
+    sa_obj_up: Optional[Tuple[float, ...]] = None
+    sa_ub_down: Optional[Tuple[float, ...]] = None
+    sa_ub_up: Optional[Tuple[float, ...]] = None
+    types: Optional[Tuple[str, ...]] = None
+    upper_bounds: Optional[Tuple[float, ...]] = None
+    user_features: Optional[Tuple[Tuple[float, ...]]] = None
+    values: Optional[Tuple[float, ...]] = None
+
+
 @dataclass
 class Variable:
     basis_status: Optional[str] = None
@@ -142,7 +162,7 @@ class FeaturesExtractor:
         with_static: bool = True,
     ) -> Features:
         features = Features()
-        features.variables = self.solver.get_variables(
+        features.variables = self.solver.get_variables_old(
             with_static=with_static,
         )
         features.constraints = self.solver.get_constraints(

miplearn/solvers/gurobi.py

@@ -6,11 +6,11 @@ import re
 import sys
 from io import StringIO
 from random import randint
-from typing import List, Any, Dict, Optional, Hashable
+from typing import List, Any, Dict, Optional, Hashable, Tuple, cast, TYPE_CHECKING
 
 from overrides import overrides
 
-from miplearn.features import Constraint, Variable
+from miplearn.features import Constraint, Variable, VariableFeatures
 from miplearn.instance.base import Instance
 from miplearn.solvers import _RedirectOutput
 from miplearn.solvers.internal import (
@@ -27,6 +27,9 @@ from miplearn.types import (
     Solution,
 )
 
+if TYPE_CHECKING:
+    import gurobipy
+
 logger = logging.getLogger(__name__)
 
 
@@ -69,12 +72,12 @@ class GurobiSolver(InternalSolver):
         self._has_mip_solution = False
 
         self._varname_to_var: Dict[str, "gurobipy.Var"] = {}
-        self._gp_vars: List["gurobipy.Var"] = []
-        self._var_names: List[str] = []
-        self._var_types: List[str] = []
-        self._var_lbs: List[float] = []
-        self._var_ubs: List[float] = []
-        self._var_obj_coeffs: List[float] = []
+        self._gp_vars: Tuple["gurobipy.Var", ...] = tuple()
+        self._var_names: Tuple[str, ...] = tuple()
+        self._var_types: Tuple[str, ...] = tuple()
+        self._var_lbs: Tuple[float, ...] = tuple()
+        self._var_ubs: Tuple[float, ...] = tuple()
+        self._var_obj_coeffs: Tuple[float, ...] = tuple()
 
         if self.lazy_cb_frequency == 1:
             self.lazy_cb_where = [self.gp.GRB.Callback.MIPSOL]
@@ -267,10 +270,27 @@ class GurobiSolver(InternalSolver):
             "upper_bound",
             "user_features",
             "value",
+            # new attributes
+            "names",
+            "basis_status",
+            "categories",
+            "lower_bounds",
+            "obj_coeffs",
+            "reduced_costs",
+            "sa_lb_down",
+            "sa_lb_up",
+            "sa_obj_down",
+            "sa_obj_up",
+            "sa_ub_down",
+            "sa_ub_up",
+            "types",
+            "upper_bounds",
+            "user_features",
+            "values",
         ]
 
     @overrides
-    def get_variables(
+    def get_variables_old(
         self,
         with_static: bool = True,
         with_sa: bool = True,
@@ -352,6 +372,76 @@ class GurobiSolver(InternalSolver):
             variables[names[i]] = var
         return variables
 
+    @overrides
+    def get_variables(
+        self,
+        with_static: bool = True,
+        with_sa: bool = True,
+    ) -> VariableFeatures:
+        model = self.model
+        assert model is not None
+
+        def _parse_gurobi_vbasis(b: int) -> str:
+            if b == 0:
+                return "B"
+            elif b == -1:
+                return "L"
+            elif b == -2:
+                return "U"
+            elif b == -3:
+                return "S"
+            else:
+                raise Exception(f"unknown vbasis: {basis_status}")
+
+        names, upper_bounds, lower_bounds, types, values = None, None, None, None, None
+        obj_coeffs, reduced_costs, basis_status = None, 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
+
+        if with_static:
+            names = self._var_names
+            upper_bounds = self._var_ubs
+            lower_bounds = self._var_lbs
+            types = self._var_types
+            obj_coeffs = self._var_obj_coeffs
+
+        if self._has_lp_solution:
+            reduced_costs = tuple(model.getAttr("rc", self._gp_vars))
+            basis_status = tuple(
+                map(
+                    _parse_gurobi_vbasis,
+                    model.getAttr("vbasis", self._gp_vars),
+                )
+            )
+
+            if with_sa:
+                sa_obj_up = tuple(model.getAttr("saobjUp", self._gp_vars))
+                sa_obj_down = tuple(model.getAttr("saobjLow", self._gp_vars))
+                sa_ub_up = tuple(model.getAttr("saubUp", self._gp_vars))
+                sa_ub_down = tuple(model.getAttr("saubLow", self._gp_vars))
+                sa_lb_up = tuple(model.getAttr("salbUp", self._gp_vars))
+                sa_lb_down = tuple(model.getAttr("salbLow", self._gp_vars))
+
+        if model.solCount > 0:
+            values = tuple(model.getAttr("x", self._gp_vars))
+
+        return VariableFeatures(
+            names=names,
+            upper_bounds=upper_bounds,
+            lower_bounds=lower_bounds,
+            types=types,
+            obj_coeffs=obj_coeffs,
+            reduced_costs=reduced_costs,
+            basis_status=basis_status,
+            sa_obj_up=sa_obj_up,
+            sa_obj_down=sa_obj_down,
+            sa_ub_up=sa_ub_up,
+            sa_ub_down=sa_ub_down,
+            sa_lb_up=sa_lb_up,
+            sa_lb_down=sa_lb_down,
+            values=values,
+        )
+
     @overrides
     def is_constraint_satisfied(self, constr: Constraint, tol: float = 1e-6) -> bool:
         assert constr.lhs is not None
@@ -589,12 +679,12 @@ class GurobiSolver(InternalSolver):
 
     def _update_vars(self) -> None:
         assert self.model is not None
-        gp_vars = self.model.getVars()
-        var_names = self.model.getAttr("varName", gp_vars)
-        var_types = self.model.getAttr("vtype", gp_vars)
-        var_ubs = self.model.getAttr("ub", gp_vars)
-        var_lbs = self.model.getAttr("lb", gp_vars)
-        var_obj_coeffs = self.model.getAttr("obj", gp_vars)
+        gp_vars: List["gurobipy.Var"] = self.model.getVars()
+        var_names: List[str] = self.model.getAttr("varName", gp_vars)
+        var_types: List[str] = self.model.getAttr("vtype", gp_vars)
+        var_ubs: List[float] = self.model.getAttr("ub", gp_vars)
+        var_lbs: List[float] = self.model.getAttr("lb", gp_vars)
+        var_obj_coeffs: List[float] = self.model.getAttr("obj", gp_vars)
         varname_to_var: Dict = {}
         for (i, gp_var) in enumerate(gp_vars):
             assert var_names[i] not in varname_to_var, (
@@ -617,12 +707,12 @@ class GurobiSolver(InternalSolver):
             )
             varname_to_var[var_names[i]] = gp_var
         self._varname_to_var = varname_to_var
-        self._gp_vars = gp_vars
-        self._var_names = var_names
-        self._var_types = var_types
-        self._var_lbs = var_lbs
-        self._var_ubs = var_ubs
-        self._var_obj_coeffs = var_obj_coeffs
+        self._gp_vars = tuple(gp_vars)
+        self._var_names = tuple(var_names)
+        self._var_types = tuple(var_types)
+        self._var_lbs = tuple(var_lbs)
+        self._var_ubs = tuple(var_ubs)
+        self._var_obj_coeffs = tuple(var_obj_coeffs)
 
     def __getstate__(self) -> Dict:
         return {

miplearn/solvers/internal.py

@@ -9,7 +9,7 @@ from typing import Any, Dict, List, Optional
 
 from overrides import EnforceOverrides
 
-from miplearn.features import Constraint, Variable
+from miplearn.features import Constraint, Variable, VariableFeatures
 from miplearn.instance.base import Instance
 from miplearn.types import (
     IterationCallback,
@@ -237,7 +237,15 @@ class InternalSolver(ABC, EnforceOverrides):
         return False
 
     @abstractmethod
-    def get_variables(self, with_static: bool = True) -> Dict[str, Variable]:
+    def get_variables_old(self, with_static: bool = True) -> Dict[str, Variable]:
+        pass
+
+    @abstractmethod
+    def get_variables(
+        self,
+        with_static: bool = True,
+        with_sa: bool = True,
+    ) -> VariableFeatures:
         pass
 
     @abstractmethod

miplearn/solvers/pyomo/base.py

@@ -6,7 +6,7 @@ import logging
 import re
 import sys
 from io import StringIO
-from typing import Any, List, Dict, Optional
+from typing import Any, List, Dict, Optional, Tuple
 
 import numpy as np
 import pyomo
@@ -19,7 +19,7 @@ from pyomo.core.expr.numeric_expr import SumExpression, MonomialTermExpression
 from pyomo.opt import TerminationCondition
 from pyomo.opt.base.solvers import SolverFactory
 
-from miplearn.features import Variable
+from miplearn.features import Variable, VariableFeatures
 from miplearn.instance.base import Instance
 from miplearn.solvers import _RedirectOutput
 from miplearn.solvers.internal import (
@@ -176,7 +176,7 @@ class BasePyomoSolver(InternalSolver):
         return solution
 
     @overrides
-    def get_variables(self, with_static: bool = True) -> Dict[str, Variable]:
+    def get_variables_old(self, with_static: bool = True) -> Dict[str, Variable]:
         assert self.model is not None
         variables = {}
         for var in self.model.component_objects(pyomo.core.Var):
@@ -190,6 +190,97 @@ class BasePyomoSolver(InternalSolver):
                 )
         return variables
 
+    @overrides
+    def get_variables(
+        self,
+        with_static: bool = True,
+        with_sa: bool = True,
+    ) -> VariableFeatures:
+        assert self.model is not None
+
+        names: List[str] = []
+        types: List[str] = []
+        upper_bounds: List[float] = []
+        lower_bounds: List[float] = []
+        obj_coeffs: List[float] = []
+        reduced_costs: List[float] = []
+        values: List[float] = []
+
+        for (i, var) in enumerate(self.model.component_objects(pyomo.core.Var)):
+            for idx in var:
+                v = var[idx]
+
+                if with_static:
+                    # Variable name
+                    if idx is None:
+                        names.append(str(var))
+                    else:
+                        names.append(f"{var}[{idx}]")
+
+                    # Variable type
+                    if v.domain == pyomo.core.Binary:
+                        types.append("B")
+                    elif v.domain in [
+                        pyomo.core.Reals,
+                        pyomo.core.NonNegativeReals,
+                        pyomo.core.NonPositiveReals,
+                        pyomo.core.NegativeReals,
+                        pyomo.core.PositiveReals,
+                    ]:
+                        types.append("C")
+                    else:
+                        raise Exception(f"unknown variable domain: {v.domain}")
+
+                    # Bounds
+                    lb, ub = v.bounds
+                    upper_bounds.append(float(ub))
+                    lower_bounds.append(float(lb))
+
+                    # Objective coefficient
+                    if v.name in self._obj:
+                        obj_coeffs.append(self._obj[v.name])
+                    else:
+                        obj_coeffs.append(0.0)
+
+                # Reduced costs
+                if self._has_lp_solution:
+                    reduced_costs.append(self.model.rc[v])
+
+                # Values
+                if self._has_lp_solution or self._has_mip_solution:
+                    values.append(v.value)
+
+        names_t: Optional[Tuple[str, ...]] = None
+        types_t: Optional[Tuple[str, ...]] = None
+        upper_bounds_t: Optional[Tuple[float, ...]] = None
+        lower_bounds_t: Optional[Tuple[float, ...]] = None
+        obj_coeffs_t: Optional[Tuple[float, ...]] = None
+        reduced_costs_t: Optional[Tuple[float, ...]] = None
+        values_t: Optional[Tuple[float, ...]] = None
+
+        if with_static:
+            names_t = tuple(names)
+            types_t = tuple(types)
+            upper_bounds_t = tuple(upper_bounds)
+            lower_bounds_t = tuple(lower_bounds)
+            obj_coeffs_t = tuple(obj_coeffs)
+
+        if self._has_lp_solution:
+            reduced_costs_t = tuple(reduced_costs)
+
+        if self._has_lp_solution or self._has_mip_solution:
+            values_t = tuple(values)
+
+        return VariableFeatures(
+            names=names_t,
+            types=types_t,
+            upper_bounds=upper_bounds_t,
+            lower_bounds=lower_bounds_t,
+            obj_coeffs=obj_coeffs_t,
+            reduced_costs=reduced_costs_t,
+            values=values_t,
+        )
+
     @overrides
     def get_variable_attrs(self) -> List[str]:
         return [
@@ -206,6 +297,23 @@ class BasePyomoSolver(InternalSolver):
             "type",
             "upper_bound",
             "value",
+            # new attributes
+            "names",
+            # "basis_status",
+            "categories",
+            "lower_bounds",
+            "obj_coeffs",
+            "reduced_costs",
+            # "sa_lb_down",
+            # "sa_lb_up",
+            # "sa_obj_down",
+            # "sa_obj_up",
+            # "sa_ub_down",
+            # "sa_ub_up",
+            "types",
+            "upper_bounds",
+            "user_features",
+            "values",
         ]
 
     @overrides

miplearn/solvers/tests/__init__.py

@@ -2,10 +2,10 @@
 #  Copyright (C) 2020-2021, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 
-from typing import Any, Dict
+from typing import Any, Dict, List
 import numpy as np
 
-from miplearn.features import Constraint, Variable
+from miplearn.features import Constraint, Variable, VariableFeatures
 from miplearn.solvers.internal import InternalSolver
 
 inf = float("inf")
@@ -44,6 +44,30 @@ def _round_variables(vars: Dict[str, Variable]) -> Dict[str, Variable]:
     return vars
 
 
+def _round(obj: Any) -> Any:
+    if isinstance(obj, tuple):
+        if obj is None:
+            return None
+        return tuple([round(v, 6) for v in obj])
+    if isinstance(obj, VariableFeatures):
+        obj.reduced_costs = _round(obj.reduced_costs)
+        obj.sa_obj_up = _round(obj.sa_obj_up)
+        obj.sa_obj_down = _round(obj.sa_obj_down)
+        obj.sa_lb_up = _round(obj.sa_lb_up)
+        obj.sa_lb_down = _round(obj.sa_lb_down)
+        obj.sa_ub_up = _round(obj.sa_ub_up)
+        obj.sa_ub_down = _round(obj.sa_ub_down)
+        obj.values = _round(obj.values)
+    return obj
+
+
+def _filter_attrs(allowed_keys: List[str], obj: Any) -> Any:
+    for key in obj.__dict__.keys():
+        if key not in allowed_keys:
+            setattr(obj, key, None)
+    return obj
+
+
 def _remove_unsupported_constr_attrs(
     solver: InternalSolver,
     constraints: Dict[str, Constraint],
@@ -58,20 +82,6 @@ def _remove_unsupported_constr_attrs(
     return constraints
 
 
-def _remove_unsupported_var_attrs(
-    solver: InternalSolver,
-    variables: Dict[str, Variable],
-) -> Dict[str, Variable]:
-    for (cname, c) in variables.items():
-        to_remove = []
-        for k in c.__dict__.keys():
-            if k not in solver.get_variable_attrs():
-                to_remove.append(k)
-        for k in to_remove:
-            setattr(c, k, None)
-    return variables
-
-
 def run_internal_solver_tests(solver: InternalSolver) -> None:
     run_basic_usage_tests(solver.clone())
     run_warm_start_tests(solver.clone())
@@ -89,41 +99,13 @@ def run_basic_usage_tests(solver: InternalSolver) -> None:
 
     # Fetch variables (after-load)
     assert_equals(
-        _round_variables(solver.get_variables()),
-        _remove_unsupported_var_attrs(
-            solver,
-            {
-                "x[0]": Variable(
-                    lower_bound=0.0,
-                    obj_coeff=505.0,
-                    type="B",
-                    upper_bound=1.0,
-                ),
-                "x[1]": Variable(
-                    lower_bound=0.0,
-                    obj_coeff=352.0,
-                    type="B",
-                    upper_bound=1.0,
-                ),
-                "x[2]": Variable(
-                    lower_bound=0.0,
-                    obj_coeff=458.0,
-                    type="B",
-                    upper_bound=1.0,
-                ),
-                "x[3]": Variable(
-                    lower_bound=0.0,
-                    obj_coeff=220.0,
-                    type="B",
-                    upper_bound=1.0,
-                ),
-                "z": Variable(
-                    lower_bound=0.0,
-                    obj_coeff=0.0,
-                    type="C",
-                    upper_bound=67.0,
-                ),
-            },
+        solver.get_variables(),
+        VariableFeatures(
+            names=("x[0]", "x[1]", "x[2]", "x[3]", "z"),
+            lower_bounds=(0.0, 0.0, 0.0, 0.0, 0.0),
+            upper_bounds=(1.0, 1.0, 1.0, 1.0, 67.0),
+            types=("B", "B", "B", "B", "C"),
+            obj_coeffs=(505.0, 352.0, 458.0, 220.0, 0.0),
         ),
     )
 
@@ -150,88 +132,22 @@ def run_basic_usage_tests(solver: InternalSolver) -> None:
     assert lp_stats.lp_wallclock_time is not None
     assert lp_stats.lp_wallclock_time > 0
 
-    # Fetch variables (after-load)
+    # Fetch variables (after-lp)
     assert_equals(
-        _round_variables(solver.get_variables()),
-        _remove_unsupported_var_attrs(
-            solver,
-            {
-                "x[0]": Variable(
-                    basis_status="U",
-                    lower_bound=0.0,
-                    obj_coeff=505.0,
-                    reduced_cost=193.615385,
-                    sa_lb_down=-inf,
-                    sa_lb_up=1.0,
-                    sa_obj_down=311.384615,
-                    sa_obj_up=inf,
-                    sa_ub_down=0.913043,
-                    sa_ub_up=2.043478,
-                    type="B",
-                    upper_bound=1.0,
-                    value=1.0,
-                ),
-                "x[1]": Variable(
-                    basis_status="B",
-                    lower_bound=0.0,
-                    obj_coeff=352.0,
-                    reduced_cost=0.0,
-                    sa_lb_down=-inf,
-                    sa_lb_up=0.923077,
-                    sa_obj_down=317.777778,
-                    sa_obj_up=570.869565,
-                    sa_ub_down=0.923077,
-                    sa_ub_up=inf,
-                    type="B",
-                    upper_bound=1.0,
-                    value=0.923077,
+        _round(solver.get_variables(with_static=False)),
+        _filter_attrs(
+            solver.get_variable_attrs(),
+            VariableFeatures(
+                basis_status=("U", "B", "U", "L", "U"),
+                reduced_costs=(193.615385, 0.0, 187.230769, -23.692308, 13.538462),
+                sa_lb_down=(-inf, -inf, -inf, -0.111111, -inf),
+                sa_lb_up=(1.0, 0.923077, 1.0, 1.0, 67.0),
+                sa_obj_down=(311.384615, 317.777778, 270.769231, -inf, -13.538462),
+                sa_obj_up=(inf, 570.869565, inf, 243.692308, inf),
+                sa_ub_down=(0.913043, 0.923077, 0.9, 0.0, 43.0),
+                sa_ub_up=(2.043478, inf, 2.2, inf, 69.0),
+                values=(1.0, 0.923077, 1.0, 0.0, 67.0),
             ),
-                "x[2]": Variable(
-                    basis_status="U",
-                    lower_bound=0.0,
-                    obj_coeff=458.0,
-                    reduced_cost=187.230769,
-                    sa_lb_down=-inf,
-                    sa_lb_up=1.0,
-                    sa_obj_down=270.769231,
-                    sa_obj_up=inf,
-                    sa_ub_down=0.9,
-                    sa_ub_up=2.2,
-                    type="B",
-                    upper_bound=1.0,
-                    value=1.0,
-                ),
-                "x[3]": Variable(
-                    basis_status="L",
-                    lower_bound=0.0,
-                    obj_coeff=220.0,
-                    reduced_cost=-23.692308,
-                    sa_lb_down=-0.111111,
-                    sa_lb_up=1.0,
-                    sa_obj_down=-inf,
-                    sa_obj_up=243.692308,
-                    sa_ub_down=0.0,
-                    sa_ub_up=inf,
-                    type="B",
-                    upper_bound=1.0,
-                    value=0.0,
-                ),
-                "z": Variable(
-                    basis_status="U",
-                    lower_bound=0.0,
-                    obj_coeff=0.0,
-                    reduced_cost=13.538462,
-                    sa_lb_down=-inf,
-                    sa_lb_up=67.0,
-                    sa_obj_down=-13.538462,
-                    sa_obj_up=inf,
-                    sa_ub_down=43.0,
-                    sa_ub_up=69.0,
-                    type="C",
-                    upper_bound=67.0,
-                    value=67.0,
-                ),
-            },
         ),
     )
 
@@ -281,16 +197,10 @@ def run_basic_usage_tests(solver: InternalSolver) -> None:
 
     # Fetch variables (after-mip)
     assert_equals(
-        _round_variables(solver.get_variables(with_static=False)),
-        _remove_unsupported_var_attrs(
-            solver,
-            {
-                "x[0]": Variable(value=1.0),
-                "x[1]": Variable(value=0.0),
-                "x[2]": Variable(value=1.0),
-                "x[3]": Variable(value=1.0),
-                "z": Variable(value=61.0),
-            },
+        _round(solver.get_variables(with_static=False)),
+        _filter_attrs(
+            solver.get_variable_attrs(),
+            VariableFeatures(values=(1.0, 0.0, 1.0, 1.0, 61.0)),
         ),
     )