Make cuts component compatible with Pyomo+Gurobi

tests/components/cuts/test_mem.py

@@ -5,62 +5,69 @@
 from typing import Any, List, Dict
 from unittest.mock import Mock
 
-from sklearn.dummy import DummyClassifier
-from sklearn.neighbors import KNeighborsClassifier
-
 from miplearn.components.cuts.mem import MemorizingCutsComponent
 from miplearn.extractors.abstract import FeaturesExtractor
-from miplearn.problems.stab import build_stab_model
+from miplearn.problems.stab import build_stab_model_gurobipy, build_stab_model_pyomo
 from miplearn.solvers.learning import LearningSolver
+from sklearn.dummy import DummyClassifier
+from sklearn.neighbors import KNeighborsClassifier
+from typing import Callable
 
 
-def test_mem_component(
-    stab_h5: List[str],
+def test_mem_component_gp(
+    stab_gp_h5: List[str],
+    stab_pyo_h5: List[str],
     default_extractor: FeaturesExtractor,
 ) -> None:
-    clf = Mock(wraps=DummyClassifier())
-    comp = MemorizingCutsComponent(clf=clf, extractor=default_extractor)
-    comp.fit(stab_h5)
+    for h5 in [stab_pyo_h5, stab_gp_h5]:
+        clf = Mock(wraps=DummyClassifier())
+        comp = MemorizingCutsComponent(clf=clf, extractor=default_extractor)
+        comp.fit(h5)
 
-    # Should call fit method with correct arguments
-    clf.fit.assert_called()
-    x, y = clf.fit.call_args.args
-    assert x.shape == (3, 50)
-    assert y.shape == (3, 388)
-    y = y.tolist()
-    assert y[0][:20] == [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
-    assert y[1][:20] == [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 0, 1, 1]
-    assert y[2][:20] == [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1]
+        # Should call fit method with correct arguments
+        clf.fit.assert_called()
+        x, y = clf.fit.call_args.args
+        assert x.shape == (3, 50)
+        assert y.shape == (3, 415)
+        y = y.tolist()
+        assert y[0][:20] == [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
+        assert y[1][:20] == [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
+        assert y[2][:20] == [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 0, 0, 1, 1, 1, 1, 1, 1]
 
-    # Should store violations
-    assert comp.constrs_ is not None
-    assert comp.n_features_ == 50
-    assert comp.n_targets_ == 388
-    assert len(comp.constrs_) == 388
+        # Should store violations
+        assert comp.constrs_ is not None
+        assert comp.n_features_ == 50
+        assert comp.n_targets_ == 415
+        assert len(comp.constrs_) == 415
 
-    # Call before-mip
-    stats: Dict[str, Any] = {}
-    model = Mock()
-    comp.before_mip(stab_h5[0], model, stats)
+        # Call before-mip
+        stats: Dict[str, Any] = {}
+        model = Mock()
+        comp.before_mip(h5[0], model, stats)
 
-    # Should call predict with correct args
-    clf.predict.assert_called()
-    (x_test,) = clf.predict.call_args.args
-    assert x_test.shape == (1, 50)
+        # Should call predict with correct args
+        clf.predict.assert_called()
+        (x_test,) = clf.predict.call_args.args
+        assert x_test.shape == (1, 50)
 
-    # Should set cuts_aot_
-    assert model.cuts_aot_ is not None
-    assert len(model.cuts_aot_) == 243
+        # Should set cuts_aot_
+        assert model.cuts_aot_ is not None
+        assert len(model.cuts_aot_) == 285
 
 
 def test_usage_stab(
-    stab_h5: List[str],
+    stab_gp_h5: List[str],
+    stab_pyo_h5: List[str],
     default_extractor: FeaturesExtractor,
 ) -> None:
-    data_filenames = [f.replace(".h5", ".pkl.gz") for f in stab_h5]
-    clf = KNeighborsClassifier(n_neighbors=1)
-    comp = MemorizingCutsComponent(clf=clf, extractor=default_extractor)
-    solver = LearningSolver(components=[comp])
-    solver.fit(data_filenames)
-    stats = solver.optimize(data_filenames[0], build_stab_model)
-    assert stats["Cuts: AOT"] > 0
+    for (h5, build_model) in [
+        (stab_pyo_h5, build_stab_model_pyomo),
+        (stab_gp_h5, build_stab_model_gurobipy),
+    ]:
+        data_filenames = [f.replace(".h5", ".pkl.gz") for f in h5]
+        clf = KNeighborsClassifier(n_neighbors=1)
+        comp = MemorizingCutsComponent(clf=clf, extractor=default_extractor)
+        solver = LearningSolver(components=[comp])
+        solver.fit(data_filenames)
+        stats = solver.optimize(data_filenames[0], build_model)  # type: ignore
+        assert stats["Cuts: AOT"] > 0

tests/conftest.py

@@ -52,8 +52,13 @@ def tsp_h5(request: Any) -> List[str]:
 
 
 @pytest.fixture()
-def stab_h5(request: Any) -> List[str]:
-    return _h5_fixture("stab*.h5", request)
+def stab_gp_h5(request: Any) -> List[str]:
+    return _h5_fixture("stab-gp*.h5", request)
+
+
+@pytest.fixture()
+def stab_pyo_h5(request: Any) -> List[str]:
+    return _h5_fixture("stab-pyo*.h5", request)
 
 
 @pytest.fixture()

tests/fixtures/gen_stab.py

@@ -7,9 +7,11 @@ from miplearn.collectors.basic import BasicCollector
 from miplearn.io import write_pkl_gz
 from miplearn.problems.stab import (
     MaxWeightStableSetGenerator,
-    build_stab_model,
+    build_stab_model_gurobipy,
+    build_stab_model_pyomo,
 )
 
+
 np.random.seed(42)
 gen = MaxWeightStableSetGenerator(
     w=uniform(10.0, scale=1.0),
@@ -18,6 +20,25 @@ gen = MaxWeightStableSetGenerator(
     fix_graph=True,
 )
 data = gen.generate(3)
-data_filenames = write_pkl_gz(data, dirname(__file__), prefix="stab-n50-")
+
+params = {"seed": 42, "threads": 1}
+
+# Gurobipy
+data_filenames = write_pkl_gz(data, dirname(__file__), prefix="stab-gp-n50-")
 collector = BasicCollector()
-collector.collect(data_filenames, build_stab_model)
+collector.collect(
+    data_filenames,
+    lambda data: build_stab_model_gurobipy(data, params=params),
+    progress=True,
+    verbose=True,
+)
+
+# Pyomo
+data_filenames = write_pkl_gz(data, dirname(__file__), prefix="stab-pyo-n50-")
+collector = BasicCollector()
+collector.collect(
+    data_filenames,
+    lambda model: build_stab_model_pyomo(model, params=params),
+    progress=True,
+    verbose=True,
+)

tests/problems/test_stab.py

@@ -9,7 +9,8 @@ import numpy as np
 from miplearn.h5 import H5File
 from miplearn.problems.stab import (
     MaxWeightStableSetData,
-    build_stab_model,
+    build_stab_model_gurobipy,
+    build_stab_model_pyomo,
 )
 from miplearn.solvers.abstract import AbstractModel
 
@@ -20,7 +21,8 @@ def test_stab() -> None:
         weights=np.array([1.0, 1.0, 1.0, 1.0, 1.0]),
     )
     for model in [
-        build_stab_model(data),
+        build_stab_model_gurobipy(data),
+        build_stab_model_pyomo(data),
     ]:
         assert isinstance(model, AbstractModel)
         with NamedTemporaryFile() as tempfile: