Make sample_ method accept instance

tests/components/test_lazy_static.py

@@ -23,6 +23,14 @@ from miplearn.features import (
 )
 
 
+@pytest.fixture
+def instance(features: Features) -> Instance:
+    instance = Mock(spec=Instance)
+    instance.features = features
+    instance.has_static_lazy_constraints = Mock(return_value=True)
+    return instance
+
+
 @pytest.fixture
 def sample() -> TrainingSample:
     return TrainingSample(
@@ -67,7 +75,7 @@ def features() -> Features:
     )
 
 
-def test_usage_with_solver(features: Features) -> None:
+def test_usage_with_solver(instance: Instance) -> None:
     solver = Mock(spec=LearningSolver)
     solver.use_lazy_cb = False
     solver.gap_tolerance = 1e-4
@@ -76,9 +84,6 @@ def test_usage_with_solver(features: Features) -> None:
     internal.extract_constraint = Mock(side_effect=lambda cid: "<%s>" % cid)
     internal.is_constraint_satisfied = Mock(return_value=False)
 
-    instance = Mock(spec=Instance)
-    instance.has_static_lazy_constraints = Mock(return_value=True)
-
     component = StaticLazyConstraintsComponent(violation_tolerance=1.0)
     component.thresholds["type-a"] = MinProbabilityThreshold([0.5, 0.5])
     component.thresholds["type-b"] = MinProbabilityThreshold([0.5, 0.5])
@@ -112,7 +117,7 @@ def test_usage_with_solver(features: Features) -> None:
         instance=instance,
         model=None,
         stats=stats,
-        features=features,
+        features=instance.features,
         training_data=sample,
     )
 
@@ -149,7 +154,7 @@ def test_usage_with_solver(features: Features) -> None:
         instance=instance,
         model=None,
         stats=stats,
-        features=features,
+        features=instance.features,
         training_data=sample,
     )
 
@@ -164,7 +169,7 @@ def test_usage_with_solver(features: Features) -> None:
 
 
 def test_sample_predict(
-    features: Features,
+    instance: Instance,
     sample: TrainingSample,
 ) -> None:
     comp = StaticLazyConstraintsComponent()
@@ -184,7 +189,7 @@ def test_sample_predict(
             [0.0, 1.0],  # c4
         ]
     )
-    pred = comp.sample_predict(features, sample)
+    pred = comp.sample_predict(instance, sample)
     assert pred == ["c1", "c2", "c4"]
 
 
@@ -229,7 +234,7 @@ def test_fit_xy() -> None:
 
 
 def test_sample_xy(
-    features: Features,
+    instance: Instance,
     sample: TrainingSample,
 ) -> None:
     x_expected = {
@@ -240,7 +245,7 @@ def test_sample_xy(
         "type-a": [[False, True], [False, True], [True, False]],
         "type-b": [[False, True]],
     }
-    xy = StaticLazyConstraintsComponent.sample_xy(features, sample)
+    xy = StaticLazyConstraintsComponent.sample_xy(instance, sample)
     assert xy is not None
     x_actual, y_actual = xy
     assert x_actual == x_expected

tests/components/test_objective.py

@@ -7,7 +7,7 @@ from unittest.mock import Mock
 import pytest
 from numpy.testing import assert_array_equal
 
-from miplearn import GurobiPyomoSolver, LearningSolver, Regressor
+from miplearn import GurobiPyomoSolver, LearningSolver, Regressor, Instance
 from miplearn.components.objective import ObjectiveValueComponent
 from miplearn.features import TrainingSample, InstanceFeatures, Features
 from tests.fixtures.knapsack import get_knapsack_instance
@@ -15,6 +15,13 @@ from tests.fixtures.knapsack import get_knapsack_instance
 import numpy as np
 
 
+@pytest.fixture
+def instance(features: Features) -> Instance:
+    instance = Mock(spec=Instance)
+    instance.features = features
+    return instance
+
+
 @pytest.fixture
 def features() -> Features:
     return Features(
@@ -50,7 +57,7 @@ def sample_without_ub() -> TrainingSample:
 
 
 def test_sample_xy(
-    features: Features,
+    instance: Instance,
     sample: TrainingSample,
 ) -> None:
     x_expected = {
@@ -61,7 +68,7 @@ def test_sample_xy(
         "Lower bound": [[1.0]],
         "Upper bound": [[2.0]],
     }
-    xy = ObjectiveValueComponent.sample_xy(features, sample)
+    xy = ObjectiveValueComponent.sample_xy(instance, sample)
     assert xy is not None
     x_actual, y_actual = xy
     assert x_actual == x_expected
@@ -69,7 +76,7 @@ def test_sample_xy(
 
 
 def test_sample_xy_without_lp(
-    features: Features,
+    instance: Instance,
     sample_without_lp: TrainingSample,
 ) -> None:
     x_expected = {
@@ -80,7 +87,7 @@ def test_sample_xy_without_lp(
         "Lower bound": [[1.0]],
         "Upper bound": [[2.0]],
     }
-    xy = ObjectiveValueComponent.sample_xy(features, sample_without_lp)
+    xy = ObjectiveValueComponent.sample_xy(instance, sample_without_lp)
     assert xy is not None
     x_actual, y_actual = xy
     assert x_actual == x_expected
@@ -88,7 +95,7 @@ def test_sample_xy_without_lp(
 
 
 def test_sample_xy_without_ub(
-    features: Features,
+    instance: Instance,
     sample_without_ub: TrainingSample,
 ) -> None:
     x_expected = {
@@ -96,7 +103,7 @@ def test_sample_xy_without_ub(
         "Upper bound": [[1.0, 2.0, 3.0]],
     }
     y_expected = {"Lower bound": [[1.0]]}
-    xy = ObjectiveValueComponent.sample_xy(features, sample_without_ub)
+    xy = ObjectiveValueComponent.sample_xy(instance, sample_without_ub)
     assert xy is not None
     x_actual, y_actual = xy
     assert x_actual == x_expected
@@ -170,10 +177,10 @@ def test_fit_xy_without_ub() -> None:
 
 
 def test_sample_predict(
-    features: Features,
+    instance: Instance,
     sample: TrainingSample,
 ) -> None:
-    x, y = ObjectiveValueComponent.sample_xy(features, sample)
+    x, y = ObjectiveValueComponent.sample_xy(instance, sample)
     comp = ObjectiveValueComponent()
     comp.regressors["Lower bound"] = Mock(spec=Regressor)
     comp.regressors["Upper bound"] = Mock(spec=Regressor)
@@ -183,7 +190,7 @@ def test_sample_predict(
     comp.regressors["Upper bound"].predict = Mock(  # type: ignore
         side_effect=lambda _: np.array([[60.0]])
     )
-    pred = comp.sample_predict(features, sample)
+    pred = comp.sample_predict(instance, sample)
     assert pred == {
         "Lower bound": 50.0,
         "Upper bound": 60.0,
@@ -199,16 +206,16 @@ def test_sample_predict(
 
 
 def test_sample_predict_without_ub(
-    features: Features,
+    instance: Instance,
     sample_without_ub: TrainingSample,
 ) -> None:
-    x, y = ObjectiveValueComponent.sample_xy(features, sample_without_ub)
+    x, y = ObjectiveValueComponent.sample_xy(instance, sample_without_ub)
     comp = ObjectiveValueComponent()
     comp.regressors["Lower bound"] = Mock(spec=Regressor)
     comp.regressors["Lower bound"].predict = Mock(  # type: ignore
         side_effect=lambda _: np.array([[50.0]])
     )
-    pred = comp.sample_predict(features, sample_without_ub)
+    pred = comp.sample_predict(instance, sample_without_ub)
     assert pred == {
         "Lower bound": 50.0,
     }
@@ -218,13 +225,13 @@ def test_sample_predict_without_ub(
     )
 
 
-def test_sample_evaluate(features: Features, sample: TrainingSample) -> None:
+def test_sample_evaluate(instance: Instance, sample: TrainingSample) -> None:
     comp = ObjectiveValueComponent()
     comp.regressors["Lower bound"] = Mock(spec=Regressor)
     comp.regressors["Lower bound"].predict = lambda _: np.array([[1.05]])  # type: ignore
     comp.regressors["Upper bound"] = Mock(spec=Regressor)
     comp.regressors["Upper bound"].predict = lambda _: np.array([[2.50]])  # type: ignore
-    ev = comp.sample_evaluate(features, sample)
+    ev = comp.sample_evaluate(instance, sample)
     assert ev == {
         "Lower bound": {
             "Actual value": 1.0,

tests/components/test_primal.py

@@ -8,7 +8,7 @@ import numpy as np
 from numpy.testing import assert_array_equal
 from scipy.stats import randint
 
-from miplearn import Classifier, LearningSolver
+from miplearn import Classifier, LearningSolver, Instance
 from miplearn.classifiers.threshold import Threshold
 from miplearn.components import classifier_evaluation_dict
 from miplearn.components.primal import PrimalSolutionComponent
@@ -38,6 +38,8 @@ def test_xy() -> None:
             }
         }
     )
+    instance = Mock(spec=Instance)
+    instance.features = features
     sample = TrainingSample(
         solution={
             "x": {
@@ -70,7 +72,7 @@ def test_xy() -> None:
             [True, False],
         ]
     }
-    xy = PrimalSolutionComponent.sample_xy(features, sample)
+    xy = PrimalSolutionComponent.sample_xy(instance, sample)
     assert xy is not None
     x_actual, y_actual = xy
     assert x_actual == x_expected
@@ -99,6 +101,8 @@ def test_xy_without_lp_solution() -> None:
             }
         }
     )
+    instance = Mock(spec=Instance)
+    instance.features = features
     sample = TrainingSample(
         solution={
             "x": {
@@ -123,7 +127,7 @@ def test_xy_without_lp_solution() -> None:
             [True, False],
         ]
     }
-    xy = PrimalSolutionComponent.sample_xy(features, sample)
+    xy = PrimalSolutionComponent.sample_xy(instance, sample)
     assert xy is not None
     x_actual, y_actual = xy
     assert x_actual == x_expected
@@ -161,6 +165,8 @@ def test_predict() -> None:
             }
         }
     )
+    instance = Mock(spec=Instance)
+    instance.features = features
     sample = TrainingSample(
         lp_solution={
             "x": {
@@ -170,11 +176,11 @@ def test_predict() -> None:
             }
         }
     )
-    x, _ = PrimalSolutionComponent.sample_xy(features, sample)
+    x, _ = PrimalSolutionComponent.sample_xy(instance, sample)
     comp = PrimalSolutionComponent()
     comp.classifiers = {"default": clf}
     comp.thresholds = {"default": thr}
-    solution_actual = comp.sample_predict(features, sample)
+    solution_actual = comp.sample_predict(instance, sample)
     clf.predict_proba.assert_called_once()
     assert_array_equal(x["default"], clf.predict_proba.call_args[0][0])
     thr.predict.assert_called_once()
@@ -243,7 +249,7 @@ def test_evaluate() -> None:
             4: 1.0,
         }
     }
-    features = Features(
+    features: Features = Features(
         variables={
             "x": {
                 0: VariableFeatures(),
@@ -254,7 +260,9 @@ def test_evaluate() -> None:
             }
         }
     )
-    sample = TrainingSample(
+    instance = Mock(spec=Instance)
+    instance.features = features
+    sample: TrainingSample = TrainingSample(
         solution={
             "x": {
                 0: 1.0,
@@ -265,7 +273,7 @@ def test_evaluate() -> None:
             }
         }
     )
-    ev = comp.sample_evaluate(features, sample)
+    ev = comp.sample_evaluate(instance, sample)
     assert ev == {
         0: classifier_evaluation_dict(tp=1, fp=1, tn=3, fn=0),
         1: classifier_evaluation_dict(tp=2, fp=0, tn=1, fn=2),