Rewrite DynamicLazy.sample_xy

tests/components/test_dynamic_lazy.py

@@ -16,14 +16,16 @@ from miplearn.features import (
     TrainingSample,
     Features,
     InstanceFeatures,
+    Sample,
 )
 from miplearn.instance.base import Instance
+from miplearn.solvers.tests import assert_equals
 
 E = 0.1
 
 
 @pytest.fixture
-def training_instances() -> List[Instance]:
+def training_instances2() -> List[Instance]:
     instances = [cast(Instance, Mock(spec=Instance)) for _ in range(2)]
     instances[0].features = Features(
         instance=InstanceFeatures(
@@ -76,11 +78,64 @@ def training_instances() -> List[Instance]:
     return instances
 
 
-def test_fit(training_instances: List[Instance]) -> None:
+@pytest.fixture
+def training_instances() -> List[Instance]:
+    instances = [cast(Instance, Mock(spec=Instance)) for _ in range(2)]
+    instances[0].samples = [
+        Sample(
+            after_lp=Features(
+                instance=InstanceFeatures(),
+            ),
+            after_mip=Features(extra={"lazy_enforced": {"c1", "c2"}}),
+        )
+    ]
+    instances[0].samples[0].after_lp.instance.to_list = Mock(  # type: ignore
+        return_value=[5.0]
+    )
+    instances[0].get_constraint_category = Mock(  # type: ignore
+        side_effect=lambda cid: {
+            "c1": "type-a",
+            "c2": "type-a",
+            "c3": "type-b",
+            "c4": "type-b",
+        }[cid]
+    )
+    instances[0].get_constraint_features = Mock(  # type: ignore
+        side_effect=lambda cid: {
+            "c1": [1.0, 2.0, 3.0],
+            "c2": [4.0, 5.0, 6.0],
+            "c3": [1.0, 2.0],
+            "c4": [3.0, 4.0],
+        }[cid]
+    )
+
+    return instances
+
+
+def test_sample_xy(training_instances: List[Instance]) -> None:
+    comp = DynamicLazyConstraintsComponent()
+    comp.dynamic.known_cids = ["c1", "c2", "c3", "c4"]
+    x_expected = {
+        "type-a": [[5.0, 1.0, 2.0, 3.0], [5.0, 4.0, 5.0, 6.0]],
+        "type-b": [[5.0, 1.0, 2.0], [5.0, 3.0, 4.0]],
+    }
+    y_expected = {
+        "type-a": [[False, True], [False, True]],
+        "type-b": [[True, False], [True, False]],
+    }
+    x_actual, y_actual = comp.sample_xy(
+        training_instances[0],
+        training_instances[0].samples[0],
+    )
+    assert_equals(x_actual, x_expected)
+    assert_equals(y_actual, y_expected)
+
+
+def test_fit(training_instances2: List[Instance]) -> None:
     clf = Mock(spec=Classifier)
     clf.clone = Mock(side_effect=lambda: Mock(spec=Classifier))
     comp = DynamicLazyConstraintsComponent(classifier=clf)
-    comp.fit(training_instances)
+    comp.fit(training_instances2)
     assert clf.clone.call_count == 2
 
     assert "type-a" in comp.classifiers
@@ -142,7 +197,7 @@ def test_fit(training_instances: List[Instance]) -> None:
     )
 
 
-def test_sample_predict_evaluate(training_instances: List[Instance]) -> None:
+def test_sample_predict_evaluate(training_instances2: List[Instance]) -> None:
     comp = DynamicLazyConstraintsComponent()
     comp.known_cids.extend(["c1", "c2", "c3", "c4"])
     comp.thresholds["type-a"] = MinProbabilityThreshold([0.5, 0.5])
@@ -156,13 +211,13 @@ def test_sample_predict_evaluate(training_instances: List[Instance]) -> None:
         side_effect=lambda _: np.array([[0.9, 0.1], [0.1, 0.9]])
     )
     pred = comp.sample_predict(
-        training_instances[0],
-        training_instances[0].training_data[0],
+        training_instances2[0],
+        training_instances2[0].training_data[0],
     )
     assert pred == ["c1", "c4"]
     ev = comp.sample_evaluate_old(
-        training_instances[0],
-        training_instances[0].training_data[0],
+        training_instances2[0],
+        training_instances2[0].training_data[0],
     )
     print(ev)
     assert ev == {

tests/components/test_objective.py

@@ -88,7 +88,7 @@ def test_sample_xy(sample: Sample) -> None:
         "Lower bound": [[1.0]],
         "Upper bound": [[2.0]],
     }
-    xy = ObjectiveValueComponent().sample_xy(sample)
+    xy = ObjectiveValueComponent().sample_xy(None, sample)
     assert xy is not None
     x_actual, y_actual = xy
     assert x_actual == x_expected

tests/components/test_primal.py

@@ -82,7 +82,7 @@ def test_xy(sample: Sample) -> None:
             [True, False],
         ]
     }
-    xy = PrimalSolutionComponent().sample_xy(sample)
+    xy = PrimalSolutionComponent().sample_xy(None, sample)
     assert xy is not None
     x_actual, y_actual = xy
     assert x_actual == x_expected

tests/components/test_static_lazy.py

@@ -292,7 +292,7 @@ def test_sample_xy(sample: Sample) -> None:
         "type-a": [[False, True], [False, True], [True, False]],
         "type-b": [[False, True]],
     }
-    xy = StaticLazyConstraintsComponent().sample_xy(sample)
+    xy = StaticLazyConstraintsComponent().sample_xy(None, sample)
     assert xy is not None
     x_actual, y_actual = xy
     assert x_actual == x_expected