Refactor ObjectiveComponent

miplearn/components/objective.py

@@ -3,7 +3,7 @@
 #  Released under the modified BSD license. See COPYING.md for more details.
 
 import logging
-from copy import deepcopy
+from typing import List, Dict, Union, Callable, Optional, Any, TYPE_CHECKING
 
 import numpy as np
 from sklearn.linear_model import LinearRegression
@@ -17,7 +17,12 @@ from sklearn.metrics import (
 
 from miplearn.classifiers import Regressor
 from miplearn.components.component import Component
-from miplearn.extractors import InstanceFeaturesExtractor, ObjectiveValueExtractor
+from miplearn.extractors import InstanceIterator
+from miplearn.instance import Instance
+from miplearn.types import MIPSolveStats, TrainingSample, LearningSolveStats
+
+if TYPE_CHECKING:
+    from miplearn.solvers.learning import LearningSolver
 
 logger = logging.getLogger(__name__)
 
@@ -29,58 +34,102 @@ class ObjectiveValueComponent(Component):
 
     def __init__(
         self,
-        regressor: Regressor = LinearRegression(),
+        lb_regressor: Callable[[], Regressor] = LinearRegression,
+        ub_regressor: Callable[[], Regressor] = LinearRegression,
     ) -> None:
-        self.ub_regressor = None
-        self.lb_regressor = None
-        self.regressor_prototype = regressor
+        self.ub_regressor: Optional[Regressor] = None
+        self.lb_regressor: Optional[Regressor] = None
+        self.lb_regressor_factory = lb_regressor
+        self.ub_regressor_factory = ub_regressor
+        self._predicted_ub: Optional[float] = None
+        self._predicted_lb: Optional[float] = None
 
-    def before_solve(self, solver, instance, model):
+    def before_solve(
+        self,
+        solver: "LearningSolver",
+        instance: Instance,
+        model: Any,
+    ) -> None:
         if self.ub_regressor is not None:
             logger.info("Predicting optimal value...")
-            lb, ub = self.predict([instance])[0]
-            instance.predicted_ub = ub
-            instance.predicted_lb = lb
-            logger.info("Predicted values: lb=%.2f, ub=%.2f" % (lb, ub))
+            pred = self.predict([instance])
+            self._predicted_lb = pred["Upper bound"][0]
+            self._predicted_ub = pred["Lower bound"][0]
+            logger.info(
+                "Predicted values: lb=%.2f, ub=%.2f"
+                % (
+                    self._predicted_lb,
+                    self._predicted_ub,
+                )
+            )
 
     def after_solve(
         self,
-        solver,
-        instance,
-        model,
-        stats,
-        training_data,
-    ):
-        if self.ub_regressor is not None:
-            stats["Predicted UB"] = instance.predicted_ub
-            stats["Predicted LB"] = instance.predicted_lb
-        else:
-            stats["Predicted UB"] = None
-            stats["Predicted LB"] = None
+        solver: "LearningSolver",
+        instance: Instance,
+        model: Any,
+        stats: LearningSolveStats,
+        training_data: TrainingSample,
+    ) -> None:
+        if self._predicted_ub is not None:
+            stats["Objective: predicted UB"] = self._predicted_ub
+        if self._predicted_lb is not None:
+            stats["Objective: predicted LB"] = self._predicted_lb
 
-    def fit(self, training_instances):
+    def fit(self, training_instances: Union[List[str], List[Instance]]) -> None:
+        self.lb_regressor = self.lb_regressor_factory()
+        self.ub_regressor = self.ub_regressor_factory()
         logger.debug("Extracting features...")
-        features = InstanceFeaturesExtractor().extract(training_instances)
-        ub = ObjectiveValueExtractor(kind="upper bound").extract(training_instances)
-        lb = ObjectiveValueExtractor(kind="lower bound").extract(training_instances)
-        assert ub.shape == (len(training_instances), 1)
-        assert lb.shape == (len(training_instances), 1)
-        self.ub_regressor = deepcopy(self.regressor_prototype)
-        self.lb_regressor = deepcopy(self.regressor_prototype)
-        logger.debug("Fitting ub_regressor...")
-        self.ub_regressor.fit(features, ub.ravel())
+        x_train = self.x(training_instances)
+        y_train = self.y(training_instances)
+        logger.debug("Fitting lb_regressor...")
+        self.lb_regressor.fit(x_train, y_train["Lower bound"])
         logger.debug("Fitting ub_regressor...")
-        self.lb_regressor.fit(features, lb.ravel())
+        self.ub_regressor.fit(x_train, y_train["Upper bound"])
 
-    def predict(self, instances):
-        features = InstanceFeaturesExtractor().extract(instances)
-        lb = self.lb_regressor.predict(features)
-        ub = self.ub_regressor.predict(features)
-        assert lb.shape == (len(instances),)
-        assert ub.shape == (len(instances),)
-        return np.array([lb, ub]).T
+    def predict(
+        self,
+        instances: Union[List[str], List[Instance]],
+    ) -> Dict[str, List[float]]:
+        assert self.lb_regressor is not None
+        assert self.ub_regressor is not None
+        x_test = self.x(instances)
+        (n_samples, n_features) = x_test.shape
+        lb = self.lb_regressor.predict(x_test)
+        ub = self.ub_regressor.predict(x_test)
+        assert lb.shape == (n_samples, 1)
+        assert ub.shape == (n_samples, 1)
+        return {
+            "Lower bound": lb.ravel().tolist(),
+            "Upper bound": ub.ravel().tolist(),
+        }
+
+    @staticmethod
+    def x(instances: Union[List[str], List[Instance]]) -> np.ndarray:
+        result = []
+        for instance in InstanceIterator(instances):
+            for _ in instance.training_data:
+                instance_features = instance.get_instance_features()
+                result.append(instance_features)
+        return np.array(result)
 
-    def evaluate(self, instances):
+    @staticmethod
+    def y(instances: Union[List[str], List[Instance]]) -> Dict[str, np.ndarray]:
+        ub: List[List[float]] = []
+        lb: List[List[float]] = []
+        for instance in InstanceIterator(instances):
+            for sample in instance.training_data:
+                lb.append([sample["Lower bound"]])
+                ub.append([sample["Upper bound"]])
+        return {
+            "Lower bound": np.array(lb),
+            "Upper bound": np.array(ub),
+        }
+
+    def evaluate(
+        self,
+        instances: Union[List[str], List[Instance]],
+    ) -> Dict[str, Dict[str, float]]:
         y_pred = self.predict(instances)
         y_true = np.array(
             [
@@ -88,11 +137,12 @@ class ObjectiveValueComponent(Component):
                     inst.training_data[0]["Lower bound"],
                     inst.training_data[0]["Upper bound"],
                 ]
-                for inst in instances
+                for inst in InstanceIterator(instances)
             ]
         )
-        y_true_lb, y_true_ub = y_true[:, 0], y_true[:, 1]
-        y_pred_lb, y_pred_ub = y_pred[:, 1], y_pred[:, 1]
+        y_pred_lb = y_pred["Lower bound"]
+        y_pred_ub = y_pred["Upper bound"]
+        y_true_lb, y_true_ub = y_true[:, 1], y_true[:, 1]
         ev = {
             "Lower bound": {
                 "Mean squared error": mean_squared_error(y_true_lb, y_pred_lb),

miplearn/problems/stab.py

@@ -117,9 +117,6 @@ class MaxWeightStableSetInstance(Instance):
             model.clique_eqs.add(sum(model.x[i] for i in clique) <= 1)
         return model
 
-    def get_instance_features(self):
-        return np.ones(0)
-
     def get_variable_features(self, var, index):
         neighbor_weights = [0] * 15
         neighbor_degrees = [100] * 15

miplearn/types.py

@@ -65,6 +65,8 @@ LearningSolveStats = TypedDict(
         "Primal: free": int,
         "Primal: zero": int,
         "Primal: one": int,
+        "Objective: predicted LB": float,
+        "Objective: predicted UB": float,
     },
     total=False,
 )

tests/components/test_objective.py

@@ -1,33 +1,81 @@
 #  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
 #  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
-
+from typing import cast
 from unittest.mock import Mock
 
 import numpy as np
+from numpy.testing import assert_array_equal
 
+from miplearn.instance import Instance
 from miplearn.classifiers import Regressor
 from miplearn.components.objective import ObjectiveValueComponent
 from .. import get_test_pyomo_instances
 
 
-def test_usage():
-    instances, models = get_test_pyomo_instances()
-    comp = ObjectiveValueComponent()
-    comp.fit(instances)
-    assert instances[0].training_data[0]["Lower bound"] == 1183.0
-    assert instances[0].training_data[0]["Upper bound"] == 1183.0
-    assert np.round(comp.predict(instances), 2).tolist() == [
-        [1183.0, 1183.0],
-        [1070.0, 1070.0],
+def test_x_y_predict() -> None:
+    # Construct instance
+    instance = cast(Instance, Mock(spec=Instance))
+    instance.get_instance_features = Mock(  # type: ignore
+        return_value=[1.0, 2.0],
+    )
+    instance.training_data = [
+        {
+            "Lower bound": 1.0,
+            "Upper bound": 2.0,
+        },
+        {
+            "Lower bound": 1.5,
+            "Upper bound": 2.2,
+        },
     ]
 
+    # Construct mock regressors
+    lb_regressor = Mock(spec=Regressor)
+    lb_regressor.predict = Mock(return_value=np.array([[5.0], [6.0]]))
+    ub_regressor = Mock(spec=Regressor)
+    ub_regressor.predict = Mock(return_value=np.array([[3.0], [3.0]]))
+    comp = ObjectiveValueComponent(
+        lb_regressor=lambda: lb_regressor,
+        ub_regressor=lambda: ub_regressor,
+    )
+
+    # Should build x correctly
+    x_expected = np.array([[1.0, 2.0], [1.0, 2.0]])
+    assert_array_equal(comp.x([instance]), x_expected)
+
+    # Should build y correctly
+    y_actual = comp.y([instance])
+    y_expected_lb = np.array([[1.0], [1.5]])
+    y_expected_ub = np.array([[2.0], [2.2]])
+    assert_array_equal(y_actual["Lower bound"], y_expected_lb)
+    assert_array_equal(y_actual["Upper bound"], y_expected_ub)
+
+    # Should pass arrays to regressors
+    comp.fit([instance])
+    assert_array_equal(lb_regressor.fit.call_args[0][0], x_expected)
+    assert_array_equal(lb_regressor.fit.call_args[0][1], y_expected_lb)
+    assert_array_equal(ub_regressor.fit.call_args[0][0], x_expected)
+    assert_array_equal(ub_regressor.fit.call_args[0][1], y_expected_ub)
+
+    # Should return predictions
+    pred = comp.predict([instance])
+    assert_array_equal(lb_regressor.predict.call_args[0][0], x_expected)
+    assert_array_equal(ub_regressor.predict.call_args[0][0], x_expected)
+    assert pred == {
+        "Lower bound": [5.0, 6.0],
+        "Upper bound": [3.0, 3.0],
+    }
+
 
 def test_obj_evaluate():
     instances, models = get_test_pyomo_instances()
     reg = Mock(spec=Regressor)
-    reg.predict = Mock(return_value=np.array([1000.0, 1000.0]))
-    comp = ObjectiveValueComponent(regressor=reg)
+    reg.predict = Mock(return_value=np.array([[1000.0], [1000.0]]))
+    comp = ObjectiveValueComponent(
+        lb_regressor=lambda: reg,
+        ub_regressor=lambda: reg,
+    )
     comp.fit(instances)
     ev = comp.evaluate(instances)
     assert ev == {

tests/solvers/test_learning_solver.py

@@ -130,7 +130,7 @@ def test_simulate_perfect():
             simulate_perfect=True,
         )
         stats = solver.solve(tmp.name)
-        assert stats["Lower bound"] == stats["Predicted LB"]
+        assert stats["Lower bound"] == stats["Objective: predicted LB"]
 
 
 def test_gap():