Extract instance, var and constr features into sample

tests/test_features.py

@@ -9,6 +9,7 @@ from miplearn.features import (
     InstanceFeatures,
     VariableFeatures,
     ConstraintFeatures,
+    Sample,
 )
 from miplearn.solvers.gurobi import GurobiSolver
 from miplearn.solvers.tests import assert_equals
@@ -21,71 +22,113 @@ def test_knapsack() -> None:
     instance = solver.build_test_instance_knapsack()
     model = instance.to_model()
     solver.set_instance(instance, model)
+    extractor = FeaturesExtractor()
+    sample = Sample()
+
+    # after-load
+    # -------------------------------------------------------
+    extractor.extract_after_load_features(instance, solver, sample)
+    assert_equals(sample.get("var_names"), ["x[0]", "x[1]", "x[2]", "x[3]", "z"])
+    assert_equals(sample.get("var_lower_bounds"), [0.0, 0.0, 0.0, 0.0, 0.0])
+    assert_equals(sample.get("var_obj_coeffs"), [505.0, 352.0, 458.0, 220.0, 0.0])
+    assert_equals(sample.get("var_types"), ["B", "B", "B", "B", "C"])
+    assert_equals(sample.get("var_upper_bounds"), [1.0, 1.0, 1.0, 1.0, 67.0])
+    assert_equals(
+        sample.get("var_categories"),
+        ["default", "default", "default", "default", None],
+    )
+    assert_equals(
+        sample.get("var_features_user"),
+        [[23.0, 505.0], [26.0, 352.0], [20.0, 458.0], [18.0, 220.0], None],
+    )
+    assert_equals(
+        sample.get("var_features_AlvLouWeh2017"),
+        [
+            [1.0, 0.32899, 0.0],
+            [1.0, 0.229316, 0.0],
+            [1.0, 0.298371, 0.0],
+            [1.0, 0.143322, 0.0],
+            [0.0, 0.0, 0.0],
+        ],
+    )
+    assert sample.get("var_features") is not None
+    assert_equals(sample.get("constr_names"), ["eq_capacity"])
+    assert_equals(
+        sample.get("constr_lhs"),
+        [
+            [
+                ("x[0]", 23.0),
+                ("x[1]", 26.0),
+                ("x[2]", 20.0),
+                ("x[3]", 18.0),
+                ("z", -1.0),
+            ],
+        ],
+    )
+    assert_equals(sample.get("constr_rhs"), [0.0])
+    assert_equals(sample.get("constr_senses"), ["="])
+    assert_equals(sample.get("constr_features_user"), [None])
+    assert_equals(sample.get("constr_categories"), ["eq_capacity"])
+    assert_equals(sample.get("constr_lazy"), [False])
+    assert_equals(sample.get("instance_features_user"), [67.0, 21.75])
+    assert_equals(sample.get("static_lazy_count"), 0)
+
+    # after-lp
+    # -------------------------------------------------------
     solver.solve_lp()
-
-    features = FeaturesExtractor().extract(instance, solver)
-    assert features.variables is not None
-    assert features.instance is not None
-
+    extractor.extract_after_lp_features(solver, sample)
     assert_equals(
-        features.variables,
-        VariableFeatures(
-            names=["x[0]", "x[1]", "x[2]", "x[3]", "z"],
-            basis_status=["U", "B", "U", "L", "U"],
-            categories=["default", "default", "default", "default", None],
-            lower_bounds=[0.0, 0.0, 0.0, 0.0, 0.0],
-            obj_coeffs=[505.0, 352.0, 458.0, 220.0, 0.0],
-            reduced_costs=[193.615385, 0.0, 187.230769, -23.692308, 13.538462],
-            sa_lb_down=[-inf, -inf, -inf, -0.111111, -inf],
-            sa_lb_up=[1.0, 0.923077, 1.0, 1.0, 67.0],
-            sa_obj_down=[311.384615, 317.777778, 270.769231, -inf, -13.538462],
-            sa_obj_up=[inf, 570.869565, inf, 243.692308, inf],
-            sa_ub_down=[0.913043, 0.923077, 0.9, 0.0, 43.0],
-            sa_ub_up=[2.043478, inf, 2.2, inf, 69.0],
-            types=["B", "B", "B", "B", "C"],
-            upper_bounds=[1.0, 1.0, 1.0, 1.0, 67.0],
-            user_features=[
-                [23.0, 505.0],
-                [26.0, 352.0],
-                [20.0, 458.0],
-                [18.0, 220.0],
-                None,
-            ],
-            values=[1.0, 0.923077, 1.0, 0.0, 67.0],
-            alvarez_2017=[
-                [1.0, 0.32899, 0.0, 0.0, 1.0, 1.0, 5.265874, 46.051702],
-                [1.0, 0.229316, 0.0, 0.076923, 1.0, 1.0, 3.532875, 5.388476],
-                [1.0, 0.298371, 0.0, 0.0, 1.0, 1.0, 5.232342, 46.051702],
-                [1.0, 0.143322, 0.0, 0.0, 1.0, -1.0, 46.051702, 3.16515],
-                [0.0, 0.0, 0.0, 0.0, 1.0, -1.0, 0.0, 0.0],
-            ],
-        ),
+        sample.get("lp_var_basis_status"),
+        ["U", "B", "U", "L", "U"],
     )
     assert_equals(
-        features.constraints,
-        ConstraintFeatures(
-            basis_status=["N"],
-            categories=["eq_capacity"],
-            dual_values=[13.538462],
-            names=["eq_capacity"],
-            lazy=[False],
-            lhs=[
-                [
-                    ("x[0]", 23.0),
-                    ("x[1]", 26.0),
-                    ("x[2]", 20.0),
-                    ("x[3]", 18.0),
-                    ("z", -1.0),
-                ],
-            ],
-            rhs=[0.0],
-            sa_rhs_down=[-24.0],
-            sa_rhs_up=[2.0],
-            senses=["="],
-            slacks=[0.0],
-            user_features=[None],
-        ),
+        sample.get("lp_var_reduced_costs"),
+        [193.615385, 0.0, 187.230769, -23.692308, 13.538462],
     )
+    assert_equals(
+        sample.get("lp_var_sa_lb_down"),
+        [-inf, -inf, -inf, -0.111111, -inf],
+    )
+    assert_equals(
+        sample.get("lp_var_sa_lb_up"),
+        [1.0, 0.923077, 1.0, 1.0, 67.0],
+    )
+    assert_equals(
+        sample.get("lp_var_sa_obj_down"),
+        [311.384615, 317.777778, 270.769231, -inf, -13.538462],
+    )
+    assert_equals(
+        sample.get("lp_var_sa_obj_up"),
+        [inf, 570.869565, inf, 243.692308, inf],
+    )
+    assert_equals(sample.get("lp_var_sa_ub_down"), [0.913043, 0.923077, 0.9, 0.0, 43.0])
+    assert_equals(sample.get("lp_var_sa_ub_up"), [2.043478, inf, 2.2, inf, 69.0])
+    assert_equals(sample.get("lp_var_values"), [1.0, 0.923077, 1.0, 0.0, 67.0])
+    assert_equals(
+        sample.get("lp_var_features_AlvLouWeh2017"),
+        [
+            [1.0, 0.32899, 0.0, 0.0, 1.0, 1.0, 5.265874, 46.051702],
+            [1.0, 0.229316, 0.0, 0.076923, 1.0, 1.0, 3.532875, 5.388476],
+            [1.0, 0.298371, 0.0, 0.0, 1.0, 1.0, 5.232342, 46.051702],
+            [1.0, 0.143322, 0.0, 0.0, 1.0, -1.0, 46.051702, 3.16515],
+            [0.0, 0.0, 0.0, 0.0, 1.0, -1.0, 0.0, 0.0],
+        ],
+    )
+    assert sample.get("lp_var_features") is not None
+    assert_equals(sample.get("lp_constr_basis_status"), ["N"])
+    assert_equals(sample.get("lp_constr_dual_values"), [13.538462])
+    assert_equals(sample.get("lp_constr_sa_rhs_down"), [-24.0])
+    assert_equals(sample.get("lp_constr_sa_rhs_up"), [2.0])
+    assert_equals(sample.get("lp_constr_slacks"), [0.0])
+
+    # after-mip
+    # -------------------------------------------------------
+    solver.solve()
+    extractor.extract_after_mip_features(solver, sample)
+    assert_equals(sample.get("mip_var_values"), [1.0, 0.0, 1.0, 1.0, 61.0])
+    assert_equals(sample.get("mip_constr_slacks"), [0.0])
+
+    features = extractor.extract(instance, solver)
     assert_equals(
         features.instance,
         InstanceFeatures(
@@ -138,10 +181,7 @@ def test_constraint_getindex() -> None:
 def test_assert_equals() -> None:
     assert_equals("hello", "hello")
     assert_equals([1.0, 2.0], [1.0, 2.0])
-    assert_equals(
-        np.array([1.0, 2.0]),
-        np.array([1.0, 2.0]),
-    )
+    assert_equals(np.array([1.0, 2.0]), np.array([1.0, 2.0]))
     assert_equals(
         np.array([[1.0, 2.0], [3.0, 4.0]]),
         np.array([[1.0, 2.0], [3.0, 4.0]]),
@@ -150,9 +190,6 @@ def test_assert_equals() -> None:
         VariableFeatures(values=np.array([1.0, 2.0])),  # type: ignore
         VariableFeatures(values=np.array([1.0, 2.0])),  # type: ignore
     )
-    assert_equals(
-        np.array([True, True]),
-        [True, True],
-    )
+    assert_equals(np.array([True, True]), [True, True])
     assert_equals((1.0,), (1.0,))
     assert_equals({"x": 10}, {"x": 10})