Convert InstanceFeatures into dataclass

miplearn/components/lazy_static.py

@@ -64,7 +64,7 @@ class StaticLazyConstraintsComponent(Component):
         assert features.instance is not None
         assert features.constraints is not None
 
-        if not features.instance["Lazy constraint count"] == 0:
+        if not features.instance.lazy_constraint_count == 0:
             logger.info("Instance does not have static lazy constraints. Skipping.")
         logger.info("Predicting required lazy constraints...")
         self.enforced_cids = set(self.sample_predict(features, training_data))

miplearn/components/objective.py

@@ -80,7 +80,7 @@ class ObjectiveValueComponent(Component):
         assert features.instance is not None
         x: Dict[Hashable, List[List[float]]] = {}
         y: Dict[Hashable, List[List[float]]] = {}
-        f = list(features.instance["User features"])
+        f = list(features.instance.user_features)
         if "LP value" in sample and sample["LP value"] is not None:
             f += [sample["LP value"]]
         for c in ["Upper bound", "Lower bound"]:

miplearn/features.py

@@ -120,7 +120,7 @@ class FeaturesExtractor:
         for (cid, cdict) in features.constraints.items():
             if cdict["Lazy"]:
                 lazy_count += 1
-        return {
-            "User features": user_features,
-            "Lazy constraint count": lazy_count,
-        }
+        return InstanceFeatures(
+            user_features=user_features,
+            lazy_constraint_count=lazy_count,
+        )

miplearn/types.py

@@ -78,14 +78,11 @@ LearningSolveStats = TypedDict(
     total=False,
 )
 
-InstanceFeatures = TypedDict(
-    "InstanceFeatures",
-    {
-        "User features": List[float],
-        "Lazy constraint count": int,
-    },
-    total=False,
-)
+
+@dataclass
+class InstanceFeatures:
+    user_features: List[float]
+    lazy_constraint_count: int = 0
 
 
 @dataclass