diff --git a/miplearn/components/lazy_static.py b/miplearn/components/lazy_static.py index 08ac01e..d419aac 100644 --- a/miplearn/components/lazy_static.py +++ b/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 diff --git a/miplearn/types.py b/miplearn/types.py index 6a5871d..5728410 100644 --- a/miplearn/types.py +++ b/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], diff --git a/tests/components/test_lazy_static.py b/tests/components/test_lazy_static.py index c8fd069..40832bf 100644 --- a/tests/components/test_lazy_static.py +++ b/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