Add StaticLazyComponent.xy

miplearn/components/lazy_static.py

@@ -5,12 +5,14 @@
 import logging
 import sys
 from copy import deepcopy
+from typing import Any, Dict, Tuple
 
 import numpy as np
 from tqdm.auto import tqdm
 
 from miplearn.classifiers.counting import CountingClassifier
 from miplearn.components.component import Component
+from miplearn.types import TrainingSample
 
 logger = logging.getLogger(__name__)
 
@@ -202,3 +204,26 @@ class StaticLazyConstraintsComponent(Component):
                     else:
                         result[category].append([1, 0])
         return result
+
+    @staticmethod
+    def xy_sample(
+        instance: Any,
+        sample: TrainingSample,
+    ) -> Tuple[Dict, Dict]:
+        x: Dict = {}
+        y: Dict = {}
+        if "LazyStatic: All" not in sample:
+            return x, y
+        for cid in sorted(sample["LazyStatic: All"]):
+            category = instance.get_constraint_category(cid)
+            if category is None:
+                continue
+            if category not in x:
+                x[category] = []
+                y[category] = []
+            x[category] += [instance.get_constraint_features(cid)]
+            if cid in sample["LazyStatic: Enforced"]:
+                y[category] += [[False, True]]
+            else:
+                y[category] += [[True, False]]
+        return x, y

miplearn/types.py

@@ -2,7 +2,7 @@
 #  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 
-from typing import Optional, Dict, Callable, Any, Union, Tuple, List
+from typing import Optional, Dict, Callable, Any, Union, Tuple, List, Set
 
 from mypy_extensions import TypedDict
 
@@ -16,6 +16,8 @@ TrainingSample = TypedDict(
         "LP log": str,
         "LP solution": Optional[Solution],
         "LP value": Optional[float],
+        "LazyStatic: All": Set[str],
+        "LazyStatic: Enforced": Set[str],
         "Lower bound": Optional[float],
         "MIP log": str,
         "Solution": Optional[Solution],

tests/components/test_lazy_static.py

@@ -9,6 +9,7 @@ from miplearn.components.lazy_static import StaticLazyConstraintsComponent
 from miplearn.instance import Instance
 from miplearn.solvers.internal import InternalSolver
 from miplearn.solvers.learning import LearningSolver
+from miplearn.types import TrainingSample
 
 
 def test_usage_with_solver():
@@ -230,3 +231,54 @@ def test_fit():
         expected_x["type-b"],
         expected_y["type-b"],
     )
+
+
+def test_xy_sample() -> None:
+    instance = Mock(spec=Instance)
+    sample: TrainingSample = {
+        "LazyStatic: Enforced": {"c1", "c2", "c4", "c5"},
+        "LazyStatic: All": {"c1", "c2", "c3", "c4", "c5"},
+    }
+    instance.get_constraint_category = Mock(
+        side_effect=lambda cid: {
+            "c1": "type-a",
+            "c2": "type-a",
+            "c3": "type-a",
+            "c4": "type-b",
+            "c5": "type-b",
+        }[cid]
+    )
+    instance.get_constraint_features = Mock(
+        side_effect=lambda cid: {
+            "c1": [1, 1],
+            "c2": [1, 2],
+            "c3": [1, 3],
+            "c4": [1, 4, 0],
+            "c5": [1, 5, 0],
+        }[cid]
+    )
+    x_expected = {
+        "type-a": [
+            [1, 1],
+            [1, 2],
+            [1, 3],
+        ],
+        "type-b": [
+            [1, 4, 0],
+            [1, 5, 0],
+        ],
+    }
+    y_expected = {
+        "type-a": [
+            [False, True],
+            [False, True],
+            [True, False],
+        ],
+        "type-b": [
+            [False, True],
+            [False, True],
+        ],
+    }
+    x_actual, y_actual = StaticLazyConstraintsComponent.xy_sample(instance, sample)
+    assert x_actual == x_expected
+    assert y_actual == y_expected