Use compact variable features everywhere

miplearn/components/primal.py

@@ -104,7 +104,7 @@ class PrimalSolutionComponent(Component):
 
     def sample_predict(self, sample: Sample) -> Solution:
         assert sample.after_load is not None
-        assert sample.after_load.variables_old is not None
+        assert sample.after_load.variables is not None
 
         # Compute y_pred
         x, _ = self.sample_xy(None, sample)
@@ -125,10 +125,12 @@ class PrimalSolutionComponent(Component):
             ).T
 
         # Convert y_pred into solution
-        solution: Solution = {v: None for v in sample.after_load.variables_old.keys()}
+        assert sample.after_load.variables.names is not None
+        assert sample.after_load.variables.categories is not None
+        solution: Solution = {v: None for v in sample.after_load.variables.names}
         category_offset: Dict[Hashable, int] = {cat: 0 for cat in x.keys()}
-        for (var_name, var_features) in sample.after_load.variables_old.items():
-            category = var_features.category
+        for (i, var_name) in enumerate(sample.after_load.variables.names):
+            category = sample.after_load.variables.categories[i]
             if category not in category_offset:
                 continue
             offset = category_offset[category]
@@ -150,10 +152,13 @@ class PrimalSolutionComponent(Component):
         y: Dict = {}
         assert sample.after_load is not None
         assert sample.after_load.instance is not None
-        assert sample.after_load.variables_old is not None
-        for (var_name, var) in sample.after_load.variables_old.items():
+        assert sample.after_load.variables is not None
+        assert sample.after_load.variables.names is not None
+        assert sample.after_load.variables.categories is not None
+
+        for (i, var_name) in enumerate(sample.after_load.variables.names):
             # Initialize categories
-            category = var.category
+            category = sample.after_load.variables.categories[i]
             if category is None:
                 continue
             if category not in x.keys():
@@ -162,17 +167,17 @@ class PrimalSolutionComponent(Component):
 
             # Features
             features = list(sample.after_load.instance.to_list())
-            features.extend(sample.after_load.variables_old[var_name].to_list())
+            features.extend(sample.after_load.variables.to_list(i))
             if sample.after_lp is not None:
-                assert sample.after_lp.variables_old is not None
-                features.extend(sample.after_lp.variables_old[var_name].to_list())
+                assert sample.after_lp.variables is not None
+                features.extend(sample.after_lp.variables.to_list(i))
             x[category].append(features)
 
             # Labels
             if sample.after_mip is not None:
-                assert sample.after_mip.variables_old is not None
-                assert sample.after_mip.variables_old[var_name] is not None
-                opt_value = sample.after_mip.variables_old[var_name].value
+                assert sample.after_mip.variables is not None
+                assert sample.after_mip.variables.values is not None
+                opt_value = sample.after_mip.variables.values[i]
                 assert opt_value is not None
                 assert 0.0 - 1e-5 <= opt_value <= 1.0 + 1e-5, (
                     f"Variable {var_name} has non-binary value {opt_value} in the "
@@ -190,15 +195,18 @@ class PrimalSolutionComponent(Component):
         sample: Sample,
     ) -> Dict[Hashable, Dict[str, float]]:
         assert sample.after_mip is not None
-        assert sample.after_mip.variables_old is not None
+        assert sample.after_mip.variables is not None
+        assert sample.after_mip.variables.values is not None
+        assert sample.after_mip.variables.names is not None
 
-        solution_actual = sample.after_mip.variables_old
+        solution_actual = {
+            var_name: sample.after_mip.variables.values[i]
+            for (i, var_name) in enumerate(sample.after_mip.variables.names)
+        }
         solution_pred = self.sample_predict(sample)
         vars_all, vars_one, vars_zero = set(), set(), set()
         pred_one_positive, pred_zero_positive = set(), set()
-        for (var_name, var) in solution_actual.items():
-            assert var.value is not None
-            value_actual = var.value
+        for (var_name, value_actual) in solution_actual.items():
             vars_all.add(var_name)
             if value_actual > 0.5:
                 vars_one.add(var_name)

miplearn/features.py

@@ -10,8 +10,6 @@ from typing import TYPE_CHECKING, Dict, Optional, List, Hashable, Tuple
 
 import numpy as np
 
-from miplearn.types import Category
-
 if TYPE_CHECKING:
     from miplearn.solvers.internal import InternalSolver, LPSolveStats, MIPSolveStats
     from miplearn.instance.base import Instance
@@ -49,49 +47,31 @@ class VariableFeatures:
     user_features: Optional[Tuple[Optional[Tuple[float, ...]], ...]] = None
     values: Optional[Tuple[float, ...]] = None
 
-
-@dataclass
-class Variable:
-    basis_status: Optional[str] = None
-    category: Optional[Hashable] = None
-    lower_bound: Optional[float] = None
-    obj_coeff: Optional[float] = None
-    reduced_cost: Optional[float] = None
-    sa_lb_down: Optional[float] = None
-    sa_lb_up: Optional[float] = None
-    sa_obj_down: Optional[float] = None
-    sa_obj_up: Optional[float] = None
-    sa_ub_down: Optional[float] = None
-    sa_ub_up: Optional[float] = None
-    type: Optional[str] = None
-    upper_bound: Optional[float] = None
-    user_features: Optional[List[float]] = None
-    value: Optional[float] = None
-
     # Alvarez, A. M., Louveaux, Q., & Wehenkel, L. (2017). A machine learning-based
     # approximation of strong branching. INFORMS Journal on Computing, 29(1), 185-195.
-    alvarez_2017: Optional[List[float]] = None
+    alvarez_2017: Optional[List[List[float]]] = None
 
-    def to_list(self) -> List[float]:
+    def to_list(self, index: int) -> List[float]:
         features: List[float] = []
         for attr in [
-            "lower_bound",
-            "obj_coeff",
-            "reduced_cost",
+            "lower_bounds",
+            "obj_coeffs",
+            "reduced_costs",
             "sa_lb_down",
             "sa_lb_up",
             "sa_obj_down",
             "sa_obj_up",
             "sa_ub_down",
             "sa_ub_up",
-            "upper_bound",
-            "value",
+            "upper_bounds",
+            "values",
         ]:
             if getattr(self, attr) is not None:
-                features.append(getattr(self, attr))
+                features.append(getattr(self, attr)[index])
         for attr in ["user_features", "alvarez_2017"]:
             if getattr(self, attr) is not None:
-                features.extend(getattr(self, attr))
+                if getattr(self, attr)[index] is not None:
+                    features.extend(getattr(self, attr)[index])
         _clip(features)
         return features
 
@@ -136,7 +116,6 @@ class Constraint:
 class Features:
     instance: Optional[InstanceFeatures] = None
     variables: Optional[VariableFeatures] = None
-    variables_old: Optional[Dict[str, Variable]] = None
     constraints: Optional[Dict[str, Constraint]] = None
     lp_solve: Optional["LPSolveStats"] = None
     mip_solve: Optional["MIPSolveStats"] = None
@@ -169,51 +148,16 @@ class FeaturesExtractor:
             with_static=with_static,
             with_sa=self.with_sa,
         )
-        features.variables_old = self.solver.get_variables_old(
-            with_static=with_static,
-        )
         features.constraints = self.solver.get_constraints(
             with_static=with_static,
         )
         if with_static:
             self._extract_user_features_vars(instance, features)
-            self._extract_user_features_vars_old(instance, features)
             self._extract_user_features_constrs(instance, features)
             self._extract_user_features_instance(instance, features)
             self._extract_alvarez_2017(features)
         return features
 
-    def _extract_user_features_vars_old(
-        self,
-        instance: "Instance",
-        features: Features,
-    ) -> None:
-        assert features.variables_old is not None
-        for (var_name, var) in features.variables_old.items():
-            user_features: Optional[List[float]] = None
-            category: Category = instance.get_variable_category(var_name)
-            if category is not None:
-                assert isinstance(category, collections.Hashable), (
-                    f"Variable category must be be hashable. "
-                    f"Found {type(category).__name__} instead for var={var_name}."
-                )
-                user_features = instance.get_variable_features(var_name)
-                if isinstance(user_features, np.ndarray):
-                    user_features = user_features.tolist()
-                assert isinstance(user_features, list), (
-                    f"Variable features must be a list. "
-                    f"Found {type(user_features).__name__} instead for "
-                    f"var={var_name}."
-                )
-                for v in user_features:
-                    assert isinstance(v, numbers.Real), (
-                        f"Variable features must be a list of numbers. "
-                        f"Found {type(v).__name__} instead "
-                        f"for var={var_name}."
-                    )
-            var.category = category
-            var.user_features = user_features
-
     def _extract_user_features_vars(
         self,
         instance: "Instance",
@@ -312,72 +256,80 @@ class FeaturesExtractor:
         )
 
     def _extract_alvarez_2017(self, features: Features) -> None:
-        assert features.variables_old is not None
+        assert features.variables is not None
+        assert features.variables.names is not None
+
+        obj_coeffs = features.variables.obj_coeffs
+        obj_sa_down = features.variables.sa_obj_down
+        obj_sa_up = features.variables.sa_obj_up
+        values = features.variables.values
 
         pos_obj_coeff_sum = 0.0
         neg_obj_coeff_sum = 0.0
-        for (varname, var) in features.variables_old.items():
-            if var.obj_coeff is not None:
-                if var.obj_coeff > 0:
-                    pos_obj_coeff_sum += var.obj_coeff
-                if var.obj_coeff < 0:
-                    neg_obj_coeff_sum += -var.obj_coeff
+        if obj_coeffs is not None:
+            for coeff in obj_coeffs:
+                if coeff > 0:
+                    pos_obj_coeff_sum += coeff
+                if coeff < 0:
+                    neg_obj_coeff_sum += -coeff
 
-        for (varname, var) in features.variables_old.items():
-            assert isinstance(var, Variable)
+        features.variables.alvarez_2017 = []
+        for i in range(len(features.variables.names)):
             f: List[float] = []
-            if var.obj_coeff is not None:
+            if obj_coeffs is not None:
                 # Feature 1
-                f.append(np.sign(var.obj_coeff))
+                f.append(np.sign(obj_coeffs[i]))
 
                 # Feature 2
                 if pos_obj_coeff_sum > 0:
-                    f.append(abs(var.obj_coeff) / pos_obj_coeff_sum)
+                    f.append(abs(obj_coeffs[i]) / pos_obj_coeff_sum)
                 else:
                     f.append(0.0)
 
                 # Feature 3
                 if neg_obj_coeff_sum > 0:
-                    f.append(abs(var.obj_coeff) / neg_obj_coeff_sum)
+                    f.append(abs(obj_coeffs[i]) / neg_obj_coeff_sum)
                 else:
                     f.append(0.0)
 
-            if var.value is not None:
+            if values is not None:
                 # Feature 37
                 f.append(
                     min(
-                        var.value - np.floor(var.value),
-                        np.ceil(var.value) - var.value,
+                        values[i] - np.floor(values[i]),
+                        np.ceil(values[i]) - values[i],
                     )
                 )
 
-            if var.sa_obj_up is not None:
-                assert var.obj_coeff is not None
-                assert var.sa_obj_down is not None
+            if obj_sa_up is not None:
+                assert obj_sa_down is not None
+                assert obj_coeffs is not None
+
                 # Convert inf into large finite numbers
-                sa_obj_down = max(-1e20, var.sa_obj_down)
-                sa_obj_up = min(1e20, var.sa_obj_up)
+                sd = max(-1e20, obj_sa_down[i])
+                su = min(1e20, obj_sa_up[i])
+                obj = obj_coeffs[i]
 
                 # Features 44 and 46
-                f.append(np.sign(var.sa_obj_up))
-                f.append(np.sign(var.sa_obj_down))
+                f.append(np.sign(obj_sa_up[i]))
+                f.append(np.sign(obj_sa_down[i]))
 
                 # Feature 47
-                csign = np.sign(var.obj_coeff)
-                if csign != 0 and ((var.obj_coeff - sa_obj_down) / csign) > 0.001:
-                    f.append(log((var.obj_coeff - sa_obj_down) / csign))
+                csign = np.sign(obj)
+                if csign != 0 and ((obj - sd) / csign) > 0.001:
+                    f.append(log((obj - sd) / csign))
                 else:
                     f.append(0.0)
 
                 # Feature 48
-                if csign != 0 and ((sa_obj_up - var.obj_coeff) / csign) > 0.001:
-                    f.append(log((sa_obj_up - var.obj_coeff) / csign))
+                if csign != 0 and ((su - obj) / csign) > 0.001:
+                    f.append(log((su - obj) / csign))
                 else:
                     f.append(0.0)
 
             for v in f:
                 assert isfinite(v), f"non-finite elements detected: {f}"
-            var.alvarez_2017 = f
+            features.variables.alvarez_2017.append(f)
 
 
 def _clip(v: List[float]) -> None:

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, Tuple, cast, TYPE_CHECKING
+from typing import List, Any, Dict, Optional, Hashable, Tuple, TYPE_CHECKING
 
 from overrides import overrides
 
-from miplearn.features import Constraint, Variable, VariableFeatures
+from miplearn.features import Constraint, VariableFeatures
 from miplearn.instance.base import Instance
 from miplearn.solvers import _RedirectOutput
 from miplearn.solvers.internal import (
@@ -289,89 +289,6 @@ class GurobiSolver(InternalSolver):
             "values",
         ]
 
-    @overrides
-    def get_variables_old(
-        self,
-        with_static: bool = True,
-        with_sa: bool = True,
-    ) -> Dict[str, Variable]:
-        assert self.model is not None
-
-        names = self._var_names
-        ub = self._var_ubs
-        lb = self._var_lbs
-        obj_coeff = self._var_obj_coeffs
-
-        values = None
-        rc = None
-        sa_obj_up = None
-        sa_obj_down = None
-        sa_ub_up = None
-        sa_ub_down = None
-        sa_lb_up = None
-        sa_lb_down = None
-        vbasis = None
-
-        if self.model.solCount > 0:
-            values = self.model.getAttr("x", self._gp_vars)
-
-        if self._has_lp_solution:
-            rc = self.model.getAttr("rc", self._gp_vars)
-            vbasis = self.model.getAttr("vbasis", self._gp_vars)
-            if with_sa:
-                sa_obj_up = self.model.getAttr("saobjUp", self._gp_vars)
-                sa_obj_down = self.model.getAttr("saobjLow", self._gp_vars)
-                sa_ub_up = self.model.getAttr("saubUp", self._gp_vars)
-                sa_ub_down = self.model.getAttr("saubLow", self._gp_vars)
-                sa_lb_up = self.model.getAttr("salbUp", self._gp_vars)
-                sa_lb_down = self.model.getAttr("salbLow", self._gp_vars)
-
-        variables = {}
-        for (i, gp_var) in enumerate(self._gp_vars):
-            assert len(names[i]) > 0, "Empty variable name detected."
-            assert (
-                names[i] not in variables
-            ), f"Duplicated variable name detected: {names[i]}"
-            var = Variable()
-            if with_static:
-                assert lb is not None
-                assert ub is not None
-                assert obj_coeff is not None
-                var.lower_bound = lb[i]
-                var.upper_bound = ub[i]
-                var.obj_coeff = obj_coeff[i]
-                var.type = self._var_types[i]
-            if values is not None:
-                var.value = values[i]
-            if rc is not None:
-                assert vbasis is not None
-                var.reduced_cost = rc[i]
-                if vbasis[i] == 0:
-                    var.basis_status = "B"
-                elif vbasis[i] == -1:
-                    var.basis_status = "L"
-                elif vbasis[i] == -2:
-                    var.basis_status = "U"
-                elif vbasis[i] == -3:
-                    var.basis_status = "S"
-                else:
-                    raise Exception(f"unknown vbasis: {vbasis}")
-                if with_sa:
-                    assert sa_obj_up is not None
-                    assert sa_obj_down is not None
-                    assert sa_ub_up is not None
-                    assert sa_ub_down is not None
-                    assert sa_lb_up is not None
-                    assert sa_lb_down is not None
-                    var.sa_obj_up = sa_obj_up[i]
-                    var.sa_obj_down = sa_obj_down[i]
-                    var.sa_ub_up = sa_ub_up[i]
-                    var.sa_ub_down = sa_ub_down[i]
-                    var.sa_lb_up = sa_lb_up[i]
-                    var.sa_lb_down = sa_lb_down[i]
-            variables[names[i]] = var
-        return variables
-
     @overrides
     def get_variables(
         self,
@@ -651,27 +568,6 @@ class GurobiSolver(InternalSolver):
                 "get_value cannot be called from cb_where=%s" % self.cb_where
             )
 
-    @staticmethod
-    def _parse_gurobi_var_lp(gp_var: Any, var: Variable) -> None:
-        var.reduced_cost = gp_var.rc
-        var.sa_obj_up = gp_var.saobjUp
-        var.sa_obj_down = gp_var.saobjLow
-        var.sa_ub_up = gp_var.saubUp
-        var.sa_ub_down = gp_var.saubLow
-        var.sa_lb_up = gp_var.salbUp
-        var.sa_lb_down = gp_var.salbLow
-        vbasis = gp_var.vbasis
-        if vbasis == 0:
-            var.basis_status = "B"
-        elif vbasis == -1:
-            var.basis_status = "L"
-        elif vbasis == -2:
-            var.basis_status = "U"
-        elif vbasis == -3:
-            var.basis_status = "S"
-        else:
-            raise Exception(f"unknown vbasis: {vbasis}")
-
     def _raise_if_callback(self) -> None:
         if self.cb_where is not None:
             raise Exception("method cannot be called from a callback")

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, VariableFeatures
+from miplearn.features import Constraint, VariableFeatures
 from miplearn.instance.base import Instance
 from miplearn.types import (
     IterationCallback,
@@ -17,7 +17,6 @@ from miplearn.types import (
     BranchPriorities,
     UserCutCallback,
     Solution,
-    VariableName,
 )
 
 logger = logging.getLogger(__name__)
@@ -236,10 +235,6 @@ class InternalSolver(ABC, EnforceOverrides):
         """
         return False
 
-    @abstractmethod
-    def get_variables_old(self, with_static: bool = True) -> Dict[str, Variable]:
-        pass
-
     @abstractmethod
     def get_variables(
         self,

miplearn/solvers/pyomo/base.py

@@ -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, VariableFeatures
+from miplearn.features import VariableFeatures
 from miplearn.instance.base import Instance
 from miplearn.solvers import _RedirectOutput
 from miplearn.solvers.internal import (
@@ -175,21 +175,6 @@ class BasePyomoSolver(InternalSolver):
                 solution[f"{var}[{index}]"] = var[index].value
         return solution
 
-    @overrides
-    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):
-            for idx in var:
-                varname = f"{var}[{idx}]"
-                if idx is None:
-                    varname = str(var)
-                variables[varname] = self._parse_pyomo_variable(
-                    var[idx],
-                    with_static=with_static,
-                )
-        return variables
-
     @overrides
     def get_variables(
         self,
@@ -495,49 +480,6 @@ class BasePyomoSolver(InternalSolver):
     def _get_warm_start_regexp(self) -> Optional[str]:
         return None
 
-    def _parse_pyomo_variable(
-        self,
-        pyomo_var: pyomo.core.Var,
-        with_static: bool = True,
-    ) -> Variable:
-        assert self.model is not None
-        variable = Variable()
-
-        if with_static:
-            # Variable type
-            vtype: Optional[str] = None
-            if pyomo_var.domain == pyomo.core.Binary:
-                vtype = "B"
-            elif pyomo_var.domain in [
-                pyomo.core.Reals,
-                pyomo.core.NonNegativeReals,
-                pyomo.core.NonPositiveReals,
-                pyomo.core.NegativeReals,
-                pyomo.core.PositiveReals,
-            ]:
-                vtype = "C"
-            if vtype is None:
-                raise Exception(f"unknown variable domain: {pyomo_var.domain}")
-            variable.type = vtype
-
-            # Bounds
-            lb, ub = pyomo_var.bounds
-            variable.upper_bound = float(ub)
-            variable.lower_bound = float(lb)
-
-            # Objective coefficient
-            obj_coeff = 0.0
-            if pyomo_var.name in self._obj:
-                obj_coeff = self._obj[pyomo_var.name]
-            variable.obj_coeff = obj_coeff
-
-        # Reduced costs
-        if pyomo_var in self.model.rc:
-            variable.reduced_cost = self.model.rc[pyomo_var]
-
-        variable.value = pyomo_var.value
-        return variable
-
     def _parse_pyomo_constraint(
         self,
         pyomo_constr: pyomo.core.Constraint,

miplearn/solvers/pyomo/gurobi.py

@@ -3,14 +3,12 @@
 #  Released under the modified BSD license. See COPYING.md for more details.
 
 import logging
-from typing import Optional, List, Dict
+from typing import Optional
 
 from overrides import overrides
 from pyomo import environ as pe
 from scipy.stats import randint
 
-from miplearn.features import Variable
-from miplearn.solvers.gurobi import GurobiSolver
 from miplearn.solvers.pyomo.base import BasePyomoSolver
 from miplearn.types import SolverParams, BranchPriorities
 

miplearn/solvers/tests/__init__.py

@@ -3,9 +3,8 @@
 #  Released under the modified BSD license. See COPYING.md for more details.
 
 from typing import Any, Dict, List
-import numpy as np
 
-from miplearn.features import Constraint, Variable, VariableFeatures
+from miplearn.features import Constraint, VariableFeatures
 from miplearn.solvers.internal import InternalSolver
 
 inf = float("inf")
@@ -22,33 +21,15 @@ def _round_constraints(constraints: Dict[str, Constraint]) -> Dict[str, Constrai
     return constraints
 
 
-def _round_variables(vars: Dict[str, Variable]) -> Dict[str, Variable]:
-    for (cname, c) in vars.items():
-        for attr in [
-            "upper_bound",
-            "lower_bound",
-            "obj_coeff",
-            "value",
-            "reduced_cost",
-            "sa_obj_up",
-            "sa_obj_down",
-            "sa_ub_up",
-            "sa_ub_down",
-            "sa_lb_up",
-            "sa_lb_down",
-        ]:
-            if getattr(c, attr) is not None:
-                setattr(c, attr, round(getattr(c, attr), 6))
-        if c.alvarez_2017 is not None:
-            c.alvarez_2017 = list(np.round(c.alvarez_2017, 6))
-    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, float):
+        return round(obj, 6)
+    if isinstance(obj, tuple):
+        return tuple([_round(v) for v in obj])
+    if isinstance(obj, list):
+        return [_round(v) for v in obj]
     if isinstance(obj, VariableFeatures):
         obj.reduced_costs = _round(obj.reduced_costs)
         obj.sa_obj_up = _round(obj.sa_obj_up)
@@ -58,6 +39,7 @@ def _round(obj: Any) -> Any:
         obj.sa_ub_up = _round(obj.sa_ub_up)
         obj.sa_ub_down = _round(obj.sa_ub_down)
         obj.values = _round(obj.values)
+        obj.alvarez_2017 = _round(obj.alvarez_2017)
     return obj
 
 

tests/components/test_primal.py

@@ -13,10 +13,10 @@ from miplearn.classifiers.threshold import Threshold
 from miplearn.components import classifier_evaluation_dict
 from miplearn.components.primal import PrimalSolutionComponent
 from miplearn.features import (
-    Variable,
     Features,
     Sample,
     InstanceFeatures,
+    VariableFeatures,
 )
 from miplearn.problems.tsp import TravelingSalesmanGenerator
 from miplearn.solvers.learning import LearningSolver
@@ -28,39 +28,37 @@ def sample() -> Sample:
     sample = Sample(
         after_load=Features(
             instance=InstanceFeatures(),
-            variables_old={
-                "x[0]": Variable(category="default"),
-                "x[1]": Variable(category=None),
-                "x[2]": Variable(category="default"),
-                "x[3]": Variable(category="default"),
-            },
+            variables=VariableFeatures(
+                names=("x[0]", "x[1]", "x[2]", "x[3]"),
+                categories=("default", None, "default", "default"),
+            ),
         ),
         after_lp=Features(
-            variables_old={
-                "x[0]": Variable(),
-                "x[1]": Variable(),
-                "x[2]": Variable(),
-                "x[3]": Variable(),
-            },
+            variables=VariableFeatures(),
         ),
         after_mip=Features(
-            variables_old={
-                "x[0]": Variable(value=0.0),
-                "x[1]": Variable(value=1.0),
-                "x[2]": Variable(value=1.0),
-                "x[3]": Variable(value=0.0),
-            }
+            variables=VariableFeatures(
+                names=("x[0]", "x[1]", "x[2]", "x[3]"),
+                values=(0.0, 1.0, 1.0, 0.0),
+            )
         ),
     )
     sample.after_load.instance.to_list = Mock(return_value=[5.0])  # type: ignore
-    sample.after_lp.variables_old["x[0]"].to_list = Mock(  # type: ignore
-        return_value=[0.0, 0.0]
+    sample.after_load.variables.to_list = Mock(  # type:ignore
+        side_effect=lambda i: [
+            [0.0, 0.0],
+            None,
+            [1.0, 0.0],
+            [1.0, 1.0],
+        ][i]
     )
-    sample.after_lp.variables_old["x[2]"].to_list = Mock(  # type: ignore
-        return_value=[1.0, 0.0]
-    )
-    sample.after_lp.variables_old["x[3]"].to_list = Mock(  # type: ignore
-        return_value=[1.0, 1.0]
+    sample.after_lp.variables.to_list = Mock(  # type:ignore
+        side_effect=lambda i: [
+            [2.0, 2.0],
+            None,
+            [3.0, 2.0],
+            [3.0, 3.0],
+        ][i]
     )
     return sample
 
@@ -68,9 +66,9 @@ def sample() -> Sample:
 def test_xy(sample: Sample) -> None:
     x_expected = {
         "default": [
-            [5.0, 0.0, 0.0],
-            [5.0, 1.0, 0.0],
-            [5.0, 1.0, 1.0],
+            [5.0, 0.0, 0.0, 2.0, 2.0],
+            [5.0, 1.0, 0.0, 3.0, 2.0],
+            [5.0, 1.0, 1.0, 3.0, 3.0],
         ]
     }
     y_expected = {

tests/problems/test_tsp.py

@@ -43,13 +43,8 @@ def test_instance() -> None:
     assert instance.samples[0].after_mip is not None
     features = instance.samples[0].after_mip
     assert features is not None
-    assert features.variables_old is not None
-    assert features.variables_old["x[(0, 1)]"].value == 1.0
-    assert features.variables_old["x[(0, 2)]"].value == 0.0
-    assert features.variables_old["x[(0, 3)]"].value == 1.0
-    assert features.variables_old["x[(1, 2)]"].value == 1.0
-    assert features.variables_old["x[(1, 3)]"].value == 0.0
-    assert features.variables_old["x[(2, 3)]"].value == 1.0
+    assert features.variables is not None
+    assert features.variables.values == (1.0, 0.0, 1.0, 1.0, 0.0, 1.0)
     assert features.mip_solve is not None
     assert features.mip_solve.mip_lower_bound == 4.0
     assert features.mip_solve.mip_upper_bound == 4.0
@@ -79,12 +74,23 @@ def test_subtour() -> None:
     lazy_enforced = features.extra["lazy_enforced"]
     assert lazy_enforced is not None
     assert len(lazy_enforced) > 0
-    assert features.variables_old is not None
-    assert features.variables_old["x[(0, 1)]"].value == 1.0
-    assert features.variables_old["x[(0, 4)]"].value == 1.0
-    assert features.variables_old["x[(1, 2)]"].value == 1.0
-    assert features.variables_old["x[(2, 3)]"].value == 1.0
-    assert features.variables_old["x[(3, 5)]"].value == 1.0
-    assert features.variables_old["x[(4, 5)]"].value == 1.0
+    assert features.variables is not None
+    assert features.variables.values == (
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0,
+        0.0,
+        1.0,
+        1.0,
+    )
     solver.fit([instance])
     solver.solve(instance)

tests/solvers/test_learning_solver.py

@@ -16,6 +16,7 @@ from miplearn.solvers.internal import InternalSolver
 from miplearn.solvers.learning import LearningSolver
 
 # noinspection PyUnresolvedReferences
+from miplearn.solvers.tests import _round
 from tests.solvers.test_internal_solver import internal_solvers
 
 logger = logging.getLogger(__name__)
@@ -39,12 +40,9 @@ def test_learning_solver(
 
             after_mip = sample.after_mip
             assert after_mip is not None
-            assert after_mip.variables_old is not None
+            assert after_mip.variables is not None
+            assert after_mip.variables.values == (1.0, 0.0, 1.0, 1.0, 61.0)
             assert after_mip.mip_solve is not None
-            assert after_mip.variables_old["x[0]"].value == 1.0
-            assert after_mip.variables_old["x[1]"].value == 0.0
-            assert after_mip.variables_old["x[2]"].value == 1.0
-            assert after_mip.variables_old["x[3]"].value == 1.0
             assert after_mip.mip_solve.mip_lower_bound == 1183.0
             assert after_mip.mip_solve.mip_upper_bound == 1183.0
             assert after_mip.mip_solve.mip_log is not None
@@ -52,16 +50,9 @@ def test_learning_solver(
 
             after_lp = sample.after_lp
             assert after_lp is not None
-            assert after_lp.variables_old is not None
+            assert after_lp.variables is not None
+            assert _round(after_lp.variables.values) == (1.0, 0.923077, 1.0, 0.0, 67.0)
             assert after_lp.lp_solve is not None
-            assert after_lp.variables_old["x[0]"].value is not None
-            assert after_lp.variables_old["x[1]"].value is not None
-            assert after_lp.variables_old["x[2]"].value is not None
-            assert after_lp.variables_old["x[3]"].value is not None
-            assert round(after_lp.variables_old["x[0]"].value, 3) == 1.000
-            assert round(after_lp.variables_old["x[1]"].value, 3) == 0.923
-            assert round(after_lp.variables_old["x[2]"].value, 3) == 1.000
-            assert round(after_lp.variables_old["x[3]"].value, 3) == 0.000
             assert after_lp.lp_solve.lp_value is not None
             assert round(after_lp.lp_solve.lp_value, 3) == 1287.923
             assert after_lp.lp_solve.lp_log is not None

tests/test_features.py

@@ -5,14 +5,12 @@
 from miplearn.features import (
     FeaturesExtractor,
     InstanceFeatures,
-    Variable,
     Constraint,
     VariableFeatures,
 )
 from miplearn.solvers.gurobi import GurobiSolver
 from miplearn.solvers.tests import (
     assert_equals,
-    _round_variables,
     _round_constraints,
     _round,
 )
@@ -28,7 +26,7 @@ def test_knapsack() -> None:
     solver.solve_lp()
 
     features = FeaturesExtractor(solver).extract(instance)
-    assert features.variables_old is not None
+    assert features.variables is not None
     assert features.constraints is not None
     assert features.instance is not None
 
@@ -57,6 +55,13 @@ def test_knapsack() -> None:
                 None,
             ),
             values=(1.0, 0.923077, 1.0, 0.0, 67.0),
+            alvarez_2017=[
+                [1.0, 0.32899, 0.0, 0.0, 1.0, 1.0, 5.265874, 46.051702],
+                [1.0, 0.229316, 0.0, 0.076923, 1.0, 1.0, 3.532875, 5.388476],
+                [1.0, 0.298371, 0.0, 0.0, 1.0, 1.0, 5.232342, 46.051702],
+                [1.0, 0.143322, 0.0, 0.0, 1.0, -1.0, 46.051702, 3.16515],
+                [0.0, 0.0, 0.0, 0.0, 1.0, -1.0, 0.0, 0.0],
+            ],
         ),
     )
     assert_equals(