Use np.ndarray for constraint methods in Instance

tests/components/test_dynamic_lazy.py

@@ -11,7 +11,7 @@ from miplearn.classifiers import Classifier
 from miplearn.classifiers.threshold import MinProbabilityThreshold
 from miplearn.components import classifier_evaluation_dict
 from miplearn.components.dynamic_lazy import DynamicLazyConstraintsComponent
-from miplearn.features.sample import Sample, MemorySample
+from miplearn.features.sample import MemorySample
 from miplearn.instance.base import Instance
 from miplearn.solvers.tests import assert_equals
 
@@ -24,71 +24,71 @@ def training_instances() -> List[Instance]:
     samples_0 = [
         MemorySample(
             {
-                "mip_constr_lazy_enforced": {"c1", "c2"},
-                "static_instance_features": [5.0],
+                "mip_constr_lazy_enforced": {b"c1", b"c2"},
+                "static_instance_features": np.array([5.0]),
             },
         ),
         MemorySample(
             {
-                "mip_constr_lazy_enforced": {"c2", "c3"},
-                "static_instance_features": [5.0],
+                "mip_constr_lazy_enforced": {b"c2", b"c3"},
+                "static_instance_features": np.array([5.0]),
             },
         ),
     ]
     instances[0].get_samples = Mock(return_value=samples_0)  # type: ignore
     instances[0].get_constraint_categories = Mock(  # type: ignore
-        return_value={
-            "c1": "type-a",
-            "c2": "type-a",
-            "c3": "type-b",
-            "c4": "type-b",
-        }
+        return_value=np.array(["type-a", "type-a", "type-b", "type-b"], dtype="S")
     )
     instances[0].get_constraint_features = Mock(  # type: ignore
-        return_value={
-            "c1": [1.0, 2.0, 3.0],
-            "c2": [4.0, 5.0, 6.0],
-            "c3": [1.0, 2.0],
-            "c4": [3.0, 4.0],
-        }
+        return_value=np.array(
+            [
+                [1.0, 2.0, 3.0],
+                [4.0, 5.0, 6.0],
+                [1.0, 2.0, 0.0],
+                [3.0, 4.0, 0.0],
+            ]
+        )
+    )
+    instances[0].are_constraints_lazy = Mock(  # type: ignore
+        return_value=np.zeros(4, dtype=bool)
     )
     samples_1 = [
         MemorySample(
             {
-                "mip_constr_lazy_enforced": {"c3", "c4"},
-                "static_instance_features": [8.0],
+                "mip_constr_lazy_enforced": {b"c3", b"c4"},
+                "static_instance_features": np.array([8.0]),
             },
         )
     ]
     instances[1].get_samples = Mock(return_value=samples_1)  # type: ignore
     instances[1].get_constraint_categories = Mock(  # type: ignore
-        return_value={
-            "c1": None,
-            "c2": "type-a",
-            "c3": "type-b",
-            "c4": "type-b",
-        }
+        return_value=np.array(["", "type-a", "type-b", "type-b"], dtype="S")
     )
     instances[1].get_constraint_features = Mock(  # type: ignore
-        return_value={
-            "c2": [7.0, 8.0, 9.0],
-            "c3": [5.0, 6.0],
-            "c4": [7.0, 8.0],
-        }
+        return_value=np.array(
+            [
+                [7.0, 8.0, 9.0],
+                [5.0, 6.0, 0.0],
+                [7.0, 8.0, 0.0],
+            ]
+        )
+    )
+    instances[1].are_constraints_lazy = Mock(  # type: ignore
+        return_value=np.zeros(4, dtype=bool)
     )
     return instances
 
 
 def test_sample_xy(training_instances: List[Instance]) -> None:
     comp = DynamicLazyConstraintsComponent()
-    comp.pre_fit([{"c1", "c2", "c3", "c4"}])
+    comp.pre_fit([{b"c1", b"c2", b"c3", b"c4"}])
     x_expected = {
-        "type-a": [[5.0, 1.0, 2.0, 3.0], [5.0, 4.0, 5.0, 6.0]],
-        "type-b": [[5.0, 1.0, 2.0], [5.0, 3.0, 4.0]],
+        b"type-a": np.array([[5.0, 1.0, 2.0, 3.0], [5.0, 4.0, 5.0, 6.0]]),
+        b"type-b": np.array([[5.0, 1.0, 2.0, 0.0], [5.0, 3.0, 4.0, 0.0]]),
     }
     y_expected = {
-        "type-a": [[False, True], [False, True]],
-        "type-b": [[True, False], [True, False]],
+        b"type-a": np.array([[False, True], [False, True]]),
+        b"type-b": np.array([[True, False], [True, False]]),
     }
     x_actual, y_actual = comp.sample_xy(
         training_instances[0],
@@ -98,95 +98,26 @@ def test_sample_xy(training_instances: List[Instance]) -> None:
     assert_equals(y_actual, y_expected)
 
 
-# def test_fit(training_instances: List[Instance]) -> None:
-#     clf = Mock(spec=Classifier)
-#     clf.clone = Mock(side_effect=lambda: Mock(spec=Classifier))
-#     comp = DynamicLazyConstraintsComponent(classifier=clf)
-#     comp.fit(training_instances)
-#     assert clf.clone.call_count == 2
-#
-#     assert "type-a" in comp.classifiers
-#     clf_a = comp.classifiers["type-a"]
-#     assert clf_a.fit.call_count == 1  # type: ignore
-#     assert_array_equal(
-#         clf_a.fit.call_args[0][0],  # type: ignore
-#         np.array(
-#             [
-#                 [5.0, 1.0, 2.0, 3.0],
-#                 [5.0, 4.0, 5.0, 6.0],
-#                 [5.0, 1.0, 2.0, 3.0],
-#                 [5.0, 4.0, 5.0, 6.0],
-#                 [8.0, 7.0, 8.0, 9.0],
-#             ]
-#         ),
-#     )
-#     assert_array_equal(
-#         clf_a.fit.call_args[0][1],  # type: ignore
-#         np.array(
-#             [
-#                 [False, True],
-#                 [False, True],
-#                 [True, False],
-#                 [False, True],
-#                 [True, False],
-#             ]
-#         ),
-#     )
-#
-#     assert "type-b" in comp.classifiers
-#     clf_b = comp.classifiers["type-b"]
-#     assert clf_b.fit.call_count == 1  # type: ignore
-#     assert_array_equal(
-#         clf_b.fit.call_args[0][0],  # type: ignore
-#         np.array(
-#             [
-#                 [5.0, 1.0, 2.0],
-#                 [5.0, 3.0, 4.0],
-#                 [5.0, 1.0, 2.0],
-#                 [5.0, 3.0, 4.0],
-#                 [8.0, 5.0, 6.0],
-#                 [8.0, 7.0, 8.0],
-#             ]
-#         ),
-#     )
-#     assert_array_equal(
-#         clf_b.fit.call_args[0][1],  # type: ignore
-#         np.array(
-#             [
-#                 [True, False],
-#                 [True, False],
-#                 [False, True],
-#                 [True, False],
-#                 [False, True],
-#                 [False, True],
-#             ]
-#         ),
-#     )
-
-
 def test_sample_predict_evaluate(training_instances: List[Instance]) -> None:
     comp = DynamicLazyConstraintsComponent()
-    comp.known_cids.extend(["c1", "c2", "c3", "c4"])
-    comp.thresholds["type-a"] = MinProbabilityThreshold([0.5, 0.5])
-    comp.thresholds["type-b"] = MinProbabilityThreshold([0.5, 0.5])
-    comp.classifiers["type-a"] = Mock(spec=Classifier)
-    comp.classifiers["type-b"] = Mock(spec=Classifier)
-    comp.classifiers["type-a"].predict_proba = Mock(  # type: ignore
+    comp.known_cids.extend([b"c1", b"c2", b"c3", b"c4"])
+    comp.thresholds[b"type-a"] = MinProbabilityThreshold([0.5, 0.5])
+    comp.thresholds[b"type-b"] = MinProbabilityThreshold([0.5, 0.5])
+    comp.classifiers[b"type-a"] = Mock(spec=Classifier)
+    comp.classifiers[b"type-b"] = Mock(spec=Classifier)
+    comp.classifiers[b"type-a"].predict_proba = Mock(  # type: ignore
         side_effect=lambda _: np.array([[0.1, 0.9], [0.8, 0.2]])
     )
-    comp.classifiers["type-b"].predict_proba = Mock(  # type: ignore
+    comp.classifiers[b"type-b"].predict_proba = Mock(  # type: ignore
         side_effect=lambda _: np.array([[0.9, 0.1], [0.1, 0.9]])
     )
     pred = comp.sample_predict(
         training_instances[0],
         training_instances[0].get_samples()[0],
     )
-    assert pred == ["c1", "c4"]
+    assert pred == [b"c1", b"c4"]
     ev = comp.sample_evaluate(
         training_instances[0],
         training_instances[0].get_samples()[0],
     )
-    assert ev == {
-        "type-a": classifier_evaluation_dict(tp=1, fp=0, tn=0, fn=1),
-        "type-b": classifier_evaluation_dict(tp=0, fp=1, tn=1, fn=0),
-    }
+    assert ev == classifier_evaluation_dict(tp=1, fp=1, tn=1, fn=1)

tests/components/test_dynamic_user_cuts.py

@@ -17,6 +17,7 @@ from miplearn.components.dynamic_user_cuts import UserCutsComponent
 from miplearn.instance.base import Instance
 from miplearn.solvers.gurobi import GurobiSolver
 from miplearn.solvers.learning import LearningSolver
+from miplearn.types import ConstraintName, ConstraintCategory
 
 logger = logging.getLogger(__name__)
 
@@ -40,13 +41,13 @@ class GurobiStableSetProblem(Instance):
         return True
 
     @overrides
-    def find_violated_user_cuts(self, model: Any) -> List[str]:
+    def find_violated_user_cuts(self, model: Any) -> List[ConstraintName]:
         assert isinstance(model, gp.Model)
         vals = model.cbGetNodeRel(model.getVars())
         violations = []
         for clique in nx.find_cliques(self.graph):
             if sum(vals[i] for i in clique) > 1:
-                violations.append(",".join([str(i) for i in clique]))
+                violations.append(",".join([str(i) for i in clique]).encode())
         return violations
 
     @overrides
@@ -54,9 +55,9 @@ class GurobiStableSetProblem(Instance):
         self,
         solver: InternalSolver,
         model: Any,
-        cid: str,
+        cid: ConstraintName,
     ) -> Any:
-        clique = [int(i) for i in cid.split(",")]
+        clique = [int(i) for i in cid.decode().split(",")]
         x = model.getVars()
         model.addConstr(gp.quicksum([x[i] for i in clique]) <= 1)
 

tests/components/test_static_lazy.py

@@ -17,6 +17,7 @@ from miplearn.solvers.internal import InternalSolver, Constraints
 from miplearn.solvers.learning import LearningSolver
 from miplearn.types import (
     LearningSolveStats,
+    ConstraintCategory,
 )
 
 
@@ -25,11 +26,11 @@ def sample() -> Sample:
     sample = MemorySample(
         {
             "static_constr_categories": [
-                "type-a",
-                "type-a",
-                "type-a",
-                "type-b",
-                "type-b",
+                b"type-a",
+                b"type-a",
+                b"type-a",
+                b"type-b",
+                b"type-b",
             ],
             "static_constr_lazy": np.array([True, True, True, True, False]),
             "static_constr_names": np.array(["c1", "c2", "c3", "c4", "c5"], dtype="S"),
@@ -68,13 +69,13 @@ def test_usage_with_solver(instance: Instance) -> None:
     )
 
     component = StaticLazyConstraintsComponent(violation_tolerance=1.0)
-    component.thresholds["type-a"] = MinProbabilityThreshold([0.5, 0.5])
-    component.thresholds["type-b"] = MinProbabilityThreshold([0.5, 0.5])
+    component.thresholds[b"type-a"] = MinProbabilityThreshold([0.5, 0.5])
+    component.thresholds[b"type-b"] = MinProbabilityThreshold([0.5, 0.5])
     component.classifiers = {
-        "type-a": Mock(spec=Classifier),
-        "type-b": Mock(spec=Classifier),
+        b"type-a": Mock(spec=Classifier),
+        b"type-b": Mock(spec=Classifier),
     }
-    component.classifiers["type-a"].predict_proba = Mock(  # type: ignore
+    component.classifiers[b"type-a"].predict_proba = Mock(  # type: ignore
         return_value=np.array(
             [
                 [0.00, 1.00],  # c1
@@ -83,7 +84,7 @@ def test_usage_with_solver(instance: Instance) -> None:
             ]
         )
     )
-    component.classifiers["type-b"].predict_proba = Mock(  # type: ignore
+    component.classifiers[b"type-b"].predict_proba = Mock(  # type: ignore
         return_value=np.array(
             [
                 [0.02, 0.98],  # c4
@@ -105,8 +106,8 @@ def test_usage_with_solver(instance: Instance) -> None:
     )
 
     # Should ask ML to predict whether each lazy constraint should be enforced
-    component.classifiers["type-a"].predict_proba.assert_called_once()
-    component.classifiers["type-b"].predict_proba.assert_called_once()
+    component.classifiers[b"type-a"].predict_proba.assert_called_once()
+    component.classifiers[b"type-b"].predict_proba.assert_called_once()
 
     # Should ask internal solver to remove some constraints
     assert internal.remove_constraints.call_count == 1
@@ -153,18 +154,18 @@ def test_usage_with_solver(instance: Instance) -> None:
 
 def test_sample_predict(sample: Sample) -> None:
     comp = StaticLazyConstraintsComponent()
-    comp.thresholds["type-a"] = MinProbabilityThreshold([0.5, 0.5])
-    comp.thresholds["type-b"] = MinProbabilityThreshold([0.5, 0.5])
-    comp.classifiers["type-a"] = Mock(spec=Classifier)
-    comp.classifiers["type-a"].predict_proba = lambda _: np.array(  # type:ignore
+    comp.thresholds[b"type-a"] = MinProbabilityThreshold([0.5, 0.5])
+    comp.thresholds[b"type-b"] = MinProbabilityThreshold([0.5, 0.5])
+    comp.classifiers[b"type-a"] = Mock(spec=Classifier)
+    comp.classifiers[b"type-a"].predict_proba = lambda _: np.array(  # type:ignore
         [
             [0.0, 1.0],  # c1
             [0.0, 0.9],  # c2
             [0.9, 0.1],  # c3
         ]
     )
-    comp.classifiers["type-b"] = Mock(spec=Classifier)
-    comp.classifiers["type-b"].predict_proba = lambda _: np.array(  # type:ignore
+    comp.classifiers[b"type-b"] = Mock(spec=Classifier)
+    comp.classifiers[b"type-b"].predict_proba = lambda _: np.array(  # type:ignore
         [
             [0.0, 1.0],  # c4
         ]
@@ -175,17 +176,17 @@ def test_sample_predict(sample: Sample) -> None:
 
 def test_fit_xy() -> None:
     x = cast(
-        Dict[str, np.ndarray],
+        Dict[ConstraintCategory, np.ndarray],
         {
-            "type-a": np.array([[1.0, 1.0], [1.0, 2.0], [1.0, 3.0]]),
-            "type-b": np.array([[1.0, 4.0, 0.0]]),
+            b"type-a": np.array([[1.0, 1.0], [1.0, 2.0], [1.0, 3.0]]),
+            b"type-b": np.array([[1.0, 4.0, 0.0]]),
         },
     )
     y = cast(
-        Dict[str, np.ndarray],
+        Dict[ConstraintCategory, np.ndarray],
         {
-            "type-a": np.array([[False, True], [False, True], [True, False]]),
-            "type-b": np.array([[False, True]]),
+            b"type-a": np.array([[False, True], [False, True], [True, False]]),
+            b"type-b": np.array([[False, True]]),
         },
     )
     clf: Classifier = Mock(spec=Classifier)
@@ -198,15 +199,15 @@ def test_fit_xy() -> None:
     )
     comp.fit_xy(x, y)
     assert clf.clone.call_count == 2
-    clf_a = comp.classifiers["type-a"]
-    clf_b = comp.classifiers["type-b"]
+    clf_a = comp.classifiers[b"type-a"]
+    clf_b = comp.classifiers[b"type-b"]
     assert clf_a.fit.call_count == 1  # type: ignore
     assert clf_b.fit.call_count == 1  # type: ignore
-    assert_array_equal(clf_a.fit.call_args[0][0], x["type-a"])  # type: ignore
-    assert_array_equal(clf_b.fit.call_args[0][0], x["type-b"])  # type: ignore
+    assert_array_equal(clf_a.fit.call_args[0][0], x[b"type-a"])  # type: ignore
+    assert_array_equal(clf_b.fit.call_args[0][0], x[b"type-b"])  # type: ignore
     assert thr.clone.call_count == 2
-    thr_a = comp.thresholds["type-a"]
-    thr_b = comp.thresholds["type-b"]
+    thr_a = comp.thresholds[b"type-a"]
+    thr_b = comp.thresholds[b"type-b"]
     assert thr_a.fit.call_count == 1  # type: ignore
     assert thr_b.fit.call_count == 1  # type: ignore
     assert thr_a.fit.call_args[0][0] == clf_a  # type: ignore
@@ -215,12 +216,12 @@ def test_fit_xy() -> None:
 
 def test_sample_xy(sample: Sample) -> None:
     x_expected = {
-        "type-a": [[5.0, 1.0, 1.0], [5.0, 1.0, 2.0], [5.0, 1.0, 3.0]],
-        "type-b": [[5.0, 1.0, 4.0, 0.0]],
+        b"type-a": [[5.0, 1.0, 1.0], [5.0, 1.0, 2.0], [5.0, 1.0, 3.0]],
+        b"type-b": [[5.0, 1.0, 4.0, 0.0]],
     }
     y_expected = {
-        "type-a": [[False, True], [False, True], [True, False]],
-        "type-b": [[False, True]],
+        b"type-a": [[False, True], [False, True], [True, False]],
+        b"type-b": [[False, True]],
     }
     xy = StaticLazyConstraintsComponent().sample_xy(None, sample)
     assert xy is not None

tests/features/test_extractor.py

@@ -32,27 +32,45 @@ def test_knapsack() -> None:
     # -------------------------------------------------------
     extractor.extract_after_load_features(instance, solver, sample)
     assert_equals(
-        sample.get_vector("static_var_names"),
+        sample.get_array("static_instance_features"),
+        np.array([67.0, 21.75]),
+    )
+    assert_equals(
+        sample.get_array("static_var_names"),
         np.array(["x[0]", "x[1]", "x[2]", "x[3]", "z"], dtype="S"),
     )
     assert_equals(
-        sample.get_vector("static_var_lower_bounds"), [0.0, 0.0, 0.0, 0.0, 0.0]
+        sample.get_array("static_var_lower_bounds"),
+        np.array([0.0, 0.0, 0.0, 0.0, 0.0]),
     )
     assert_equals(
-        sample.get_array("static_var_obj_coeffs"), [505.0, 352.0, 458.0, 220.0, 0.0]
+        sample.get_array("static_var_obj_coeffs"),
+        np.array([505.0, 352.0, 458.0, 220.0, 0.0]),
     )
     assert_equals(
         sample.get_array("static_var_types"),
         np.array(["B", "B", "B", "B", "C"], dtype="S"),
     )
     assert_equals(
-        sample.get_vector("static_var_upper_bounds"), [1.0, 1.0, 1.0, 1.0, 67.0]
+        sample.get_array("static_var_upper_bounds"),
+        np.array([1.0, 1.0, 1.0, 1.0, 67.0]),
     )
     assert_equals(
         sample.get_array("static_var_categories"),
         np.array(["default", "default", "default", "default", ""], dtype="S"),
     )
-    assert sample.get_vector_list("static_var_features") is not None
+    assert_equals(
+        sample.get_vector_list("static_var_features"),
+        np.array(
+            [
+                [23.0, 505.0, 1.0, 0.32899, 0.0, 0.0, 505.0, 1.0],
+                [26.0, 352.0, 1.0, 0.229316, 0.0, 0.0, 352.0, 1.0],
+                [20.0, 458.0, 1.0, 0.298371, 0.0, 0.0, 458.0, 1.0],
+                [18.0, 220.0, 1.0, 0.143322, 0.0, 0.0, 220.0, 1.0],
+                [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 67.0],
+            ]
+        ),
+    )
     assert_equals(
         sample.get_array("static_constr_names"),
         np.array(["eq_capacity"], dtype="S"),
@@ -69,18 +87,30 @@ def test_knapsack() -> None:
     #         ],
     #     ],
     # )
-    assert_equals(sample.get_vector("static_constr_rhs"), [0.0])
+    assert_equals(
+        sample.get_vector("static_constr_rhs"),
+        np.array([0.0]),
+    )
     assert_equals(
         sample.get_array("static_constr_senses"),
         np.array(["="], dtype="S"),
     )
-    assert_equals(sample.get_vector("static_constr_features"), [None])
+    assert_equals(
+        sample.get_vector("static_constr_features"),
+        np.array([[0.0]]),
+    )
     assert_equals(
         sample.get_vector("static_constr_categories"),
         np.array(["eq_capacity"], dtype="S"),
     )
-    assert_equals(sample.get_array("static_constr_lazy"), np.array([False]))
-    assert_equals(sample.get_vector("static_instance_features"), [67.0, 21.75])
+    assert_equals(
+        sample.get_array("static_constr_lazy"),
+        np.array([False]),
+    )
+    assert_equals(
+        sample.get_vector("static_instance_features"),
+        np.array([67.0, 21.75]),
+    )
     assert_equals(sample.get_scalar("static_constr_lazy_count"), 0)
 
     # after-lp
@@ -112,26 +142,187 @@ def test_knapsack() -> None:
         [inf, 570.869565, inf, 243.692308, inf],
     )
     assert_equals(
-        sample.get_array("lp_var_sa_ub_down"), [0.913043, 0.923077, 0.9, 0.0, 43.0]
+        sample.get_array("lp_var_sa_ub_down"),
+        np.array([0.913043, 0.923077, 0.9, 0.0, 43.0]),
+    )
+    assert_equals(
+        sample.get_array("lp_var_sa_ub_up"),
+        np.array([2.043478, inf, 2.2, inf, 69.0]),
+    )
+    assert_equals(
+        sample.get_array("lp_var_values"),
+        np.array([1.0, 0.923077, 1.0, 0.0, 67.0]),
+    )
+    assert_equals(
+        sample.get_vector_list("lp_var_features"),
+        np.array(
+            [
+                [
+                    23.0,
+                    505.0,
+                    1.0,
+                    0.32899,
+                    0.0,
+                    0.0,
+                    505.0,
+                    1.0,
+                    1.0,
+                    0.32899,
+                    0.0,
+                    0.0,
+                    1.0,
+                    1.0,
+                    5.265874,
+                    46.051702,
+                    193.615385,
+                    -inf,
+                    1.0,
+                    311.384615,
+                    inf,
+                    0.913043,
+                    2.043478,
+                    1.0,
+                ],
+                [
+                    26.0,
+                    352.0,
+                    1.0,
+                    0.229316,
+                    0.0,
+                    0.0,
+                    352.0,
+                    1.0,
+                    1.0,
+                    0.229316,
+                    0.0,
+                    0.076923,
+                    1.0,
+                    1.0,
+                    3.532875,
+                    5.388476,
+                    0.0,
+                    -inf,
+                    0.923077,
+                    317.777778,
+                    570.869565,
+                    0.923077,
+                    inf,
+                    0.923077,
+                ],
+                [
+                    20.0,
+                    458.0,
+                    1.0,
+                    0.298371,
+                    0.0,
+                    0.0,
+                    458.0,
+                    1.0,
+                    1.0,
+                    0.298371,
+                    0.0,
+                    0.0,
+                    1.0,
+                    1.0,
+                    5.232342,
+                    46.051702,
+                    187.230769,
+                    -inf,
+                    1.0,
+                    270.769231,
+                    inf,
+                    0.9,
+                    2.2,
+                    1.0,
+                ],
+                [
+                    18.0,
+                    220.0,
+                    1.0,
+                    0.143322,
+                    0.0,
+                    0.0,
+                    220.0,
+                    1.0,
+                    1.0,
+                    0.143322,
+                    0.0,
+                    0.0,
+                    1.0,
+                    -1.0,
+                    46.051702,
+                    3.16515,
+                    -23.692308,
+                    -0.111111,
+                    1.0,
+                    -inf,
+                    243.692308,
+                    0.0,
+                    inf,
+                    0.0,
+                ],
+                [
+                    0.0,
+                    0.0,
+                    0.0,
+                    0.0,
+                    0.0,
+                    0.0,
+                    0.0,
+                    67.0,
+                    0.0,
+                    0.0,
+                    0.0,
+                    0.0,
+                    1.0,
+                    -1.0,
+                    0.0,
+                    0.0,
+                    13.538462,
+                    -inf,
+                    67.0,
+                    -13.538462,
+                    inf,
+                    43.0,
+                    69.0,
+                    67.0,
+                ],
+            ]
+        ),
     )
-    assert_equals(sample.get_array("lp_var_sa_ub_up"), [2.043478, inf, 2.2, inf, 69.0])
-    assert_equals(sample.get_array("lp_var_values"), [1.0, 0.923077, 1.0, 0.0, 67.0])
-    assert sample.get_vector_list("lp_var_features") is not None
     assert_equals(
         sample.get_array("lp_constr_basis_status"),
         np.array(["N"], dtype="S"),
     )
-    assert_equals(sample.get_array("lp_constr_dual_values"), [13.538462])
-    assert_equals(sample.get_array("lp_constr_sa_rhs_down"), [-24.0])
-    assert_equals(sample.get_array("lp_constr_sa_rhs_up"), [2.0])
-    assert_equals(sample.get_array("lp_constr_slacks"), [0.0])
+    assert_equals(
+        sample.get_array("lp_constr_dual_values"),
+        np.array([13.538462]),
+    )
+    assert_equals(
+        sample.get_array("lp_constr_sa_rhs_down"),
+        np.array([-24.0]),
+    )
+    assert_equals(
+        sample.get_array("lp_constr_sa_rhs_up"),
+        np.array([2.0]),
+    )
+    assert_equals(
+        sample.get_array("lp_constr_slacks"),
+        np.array([0.0]),
+    )
+    assert_equals(
+        sample.get_array("lp_constr_features"),
+        np.array([[0.0, 13.538462, -24.0, 2.0, 0.0]]),
+    )
 
     # after-mip
     # -------------------------------------------------------
     solver.solve()
     extractor.extract_after_mip_features(solver, sample)
-    assert_equals(sample.get_array("mip_var_values"), [1.0, 0.0, 1.0, 1.0, 61.0])
-    assert_equals(sample.get_array("mip_constr_slacks"), [0.0])
+    assert_equals(
+        sample.get_array("mip_var_values"), np.array([1.0, 0.0, 1.0, 1.0, 61.0])
+    )
+    assert_equals(sample.get_array("mip_constr_slacks"), np.array([0.0]))
 
 
 def test_constraint_getindex() -> None:
@@ -200,7 +391,7 @@ class MpsInstance(Instance):
         return gp.read(self.filename)
 
 
-if __name__ == "__main__":
+def main() -> None:
     solver = GurobiSolver()
     instance = MpsInstance(sys.argv[1])
     solver.set_instance(instance)
@@ -213,3 +404,7 @@ if __name__ == "__main__":
         extractor.extract_after_lp_features(solver, sample, lp_stats)
 
     cProfile.run("run()", filename="tmp/prof")
+
+
+if __name__ == "__main__":
+    main()