MIPLearn v0.3

tests/problems/test_stab.py

@@ -1,53 +1,30 @@
 #  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
-#  Copyright (C) 2020-2021, UChicago Argonne, LLC. All rights reserved.
+#  Copyright (C) 2020-2022, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
+from tempfile import NamedTemporaryFile
 
 import networkx as nx
 import numpy as np
-from scipy.stats import uniform, randint
 
-from miplearn.problems.stab import MaxWeightStableSetInstance
-from miplearn.solvers.learning import LearningSolver
+from miplearn.h5 import H5File
+from miplearn.problems.stab import (
+    MaxWeightStableSetData,
+    build_stab_model_pyomo,
+    build_stab_model_gurobipy,
+)
 
 
 def test_stab() -> None:
-    graph = nx.cycle_graph(5)
-    weights = np.array([1.0, 1.0, 1.0, 1.0, 1.0])
-    instance = MaxWeightStableSetInstance(graph, weights)
-    solver = LearningSolver()
-    stats = solver._solve(instance)
-    assert stats["mip_lower_bound"] == 2.0
-
-
-def test_stab_generator_fixed_graph() -> None:
-    np.random.seed(42)
-    from miplearn.problems.stab import MaxWeightStableSetGenerator
-
-    gen = MaxWeightStableSetGenerator(
-        w=uniform(loc=50.0, scale=10.0),
-        n=randint(low=10, high=11),
-        p=uniform(loc=0.05, scale=0.0),
-        fix_graph=True,
+    data = MaxWeightStableSetData(
+        graph=nx.cycle_graph(5),
+        weights=np.array([1.0, 1.0, 1.0, 1.0, 1.0]),
     )
-    data = gen.generate(1_000)
-    weights = np.array([d.weights for d in data])
-    weights_avg_actual = np.round(np.average(weights, axis=0))
-    weights_avg_expected = [55.0] * 10
-    assert list(weights_avg_actual) == weights_avg_expected
-
-
-def test_stab_generator_random_graph() -> None:
-    np.random.seed(42)
-    from miplearn.problems.stab import MaxWeightStableSetGenerator
-
-    gen = MaxWeightStableSetGenerator(
-        w=uniform(loc=50.0, scale=10.0),
-        n=randint(low=30, high=41),
-        p=uniform(loc=0.5, scale=0.0),
-        fix_graph=False,
-    )
-    data = gen.generate(1_000)
-    n_nodes = [d.graph.number_of_nodes() for d in data]
-    n_edges = [d.graph.number_of_edges() for d in data]
-    assert np.round(np.mean(n_nodes)) == 35.0
-    assert np.round(np.mean(n_edges), -1) == 300.0
+    for model in [
+        build_stab_model_pyomo(data),
+        build_stab_model_gurobipy(data),
+    ]:
+        with NamedTemporaryFile() as tempfile:
+            with H5File(tempfile.name) as h5:
+                model.optimize()
+                model.extract_after_mip(h5)
+                assert h5.get_scalar("mip_obj_value") == -2.0