From 8fc9979b375505b503677647129f28c7b64da1ec Mon Sep 17 00:00:00 2001
From: "Alinson S. Xavier" <git@axavier.org>
Date: Wed, 31 Mar 2021 09:21:34 -0500
Subject: [PATCH] Use instance.features in LazyStatic and Objective

---
 miplearn/components/lazy_static.py   |  5 +--
 miplearn/components/objective.py     |  2 +-
 tests/components/test_lazy_static.py | 52 ++++++++++++++++------------
 tests/components/test_objective.py   |  8 +++--
 4 files changed, 38 insertions(+), 29 deletions(-)

diff --git a/miplearn/components/lazy_static.py b/miplearn/components/lazy_static.py
index d419aac..302ac09 100644
--- a/miplearn/components/lazy_static.py
+++ b/miplearn/components/lazy_static.py
@@ -215,13 +215,14 @@ class StaticLazyConstraintsComponent(Component):
         if "LazyStatic: All" not in sample:
             return x, y
         for cid in sorted(sample["LazyStatic: All"]):
-            category = instance.get_constraint_category(cid)
+            cfeatures = instance.features["Constraints"][cid]
+            category = cfeatures["Category"]
             if category is None:
                 continue
             if category not in x:
                 x[category] = []
                 y[category] = []
-            x[category] += [instance.get_constraint_features(cid)]
+            x[category] += [cfeatures["User features"]]
             if cid in sample["LazyStatic: Enforced"]:
                 y[category] += [[False, True]]
             else:
diff --git a/miplearn/components/objective.py b/miplearn/components/objective.py
index ae0bbae..63cd8b2 100644
--- a/miplearn/components/objective.py
+++ b/miplearn/components/objective.py
@@ -171,7 +171,7 @@ class ObjectiveValueComponent(Component):
         y: Dict = {}
         if "Lower bound" not in sample:
             return x, y
-        features = instance.get_instance_features()
+        features = instance.features["Instance"]["User features"]
         if "LP value" in sample and sample["LP value"] is not None:
             features += [sample["LP value"]]
         x["Lower bound"] = [features]
diff --git a/tests/components/test_lazy_static.py b/tests/components/test_lazy_static.py
index 40832bf..c212c2e 100644
--- a/tests/components/test_lazy_static.py
+++ b/tests/components/test_lazy_static.py
@@ -239,33 +239,39 @@ def test_xy_sample() -> None:
         "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]
-    )
+    instance.features = {
+        "Constraints": {
+            "c1": {
+                "Category": "type-a",
+                "User features": [1.0, 1.0],
+            },
+            "c2": {
+                "Category": "type-a",
+                "User features": [1.0, 2.0],
+            },
+            "c3": {
+                "Category": "type-a",
+                "User features": [1.0, 3.0],
+            },
+            "c4": {
+                "Category": "type-b",
+                "User features": [1.0, 4.0, 0.0],
+            },
+            "c5": {
+                "Category": "type-b",
+                "User features": [1.0, 5.0, 0.0],
+            },
+        }
+    }
     x_expected = {
         "type-a": [
-            [1, 1],
-            [1, 2],
-            [1, 3],
+            [1.0, 1.0],
+            [1.0, 2.0],
+            [1.0, 3.0],
         ],
         "type-b": [
-            [1, 4, 0],
-            [1, 5, 0],
+            [1.0, 4.0, 0.0],
+            [1.0, 5.0, 0.0],
         ],
     }
     y_expected = {
diff --git a/tests/components/test_objective.py b/tests/components/test_objective.py
index e751bbe..7932052 100644
--- a/tests/components/test_objective.py
+++ b/tests/components/test_objective.py
@@ -17,9 +17,11 @@ from tests.fixtures.knapsack import get_test_pyomo_instances
 
 def test_xy_sample() -> None:
     instance = cast(Instance, Mock(spec=Instance))
-    instance.get_instance_features = Mock(  # type: ignore
-        return_value=[1.0, 2.0],
-    )
+    instance.features = {
+        "Instance": {
+            "User features": [1.0, 2.0],
+        }
+    }
     sample: TrainingSample = {
         "Lower bound": 1.0,
         "Upper bound": 2.0,