Convert VariableFeatures into dataclass

miplearn/components/primal.py

@@ -136,7 +136,7 @@ class PrimalSolutionComponent(Component):
         category_offset: Dict[Hashable, int] = {cat: 0 for cat in x.keys()}
         for (var_name, var_dict) in features.variables.items():
             for (idx, var_features) in var_dict.items():
-                category = var_features["Category"]
+                category = var_features.category
                 offset = category_offset[category]
                 category_offset[category] += 1
                 if y_pred[category][offset, 0]:
@@ -159,15 +159,15 @@ class PrimalSolutionComponent(Component):
             solution = sample["Solution"]
         for (var_name, var_dict) in features.variables.items():
             for (idx, var_features) in var_dict.items():
-                category = var_features["Category"]
+                category = var_features.category
                 if category is None:
                     continue
                 if category not in x.keys():
                     x[category] = []
                     y[category] = []
                 f: List[float] = []
-                assert var_features["User features"] is not None
-                f += var_features["User features"]
+                assert var_features.user_features is not None
+                f += var_features.user_features
                 if "LP solution" in sample and sample["LP solution"] is not None:
                     lp_value = sample["LP solution"][var_name][idx]
                     if lp_value is not None:

miplearn/features.py

@@ -4,9 +4,15 @@
 
 import numbers
 import collections
-from typing import TYPE_CHECKING, Dict
+from typing import TYPE_CHECKING, Dict, Hashable
 
-from miplearn.types import Features, ConstraintFeatures, InstanceFeatures
+from miplearn.types import (
+    Features,
+    ConstraintFeatures,
+    InstanceFeatures,
+    VariableFeatures,
+    VarIndex,
+)
 
 if TYPE_CHECKING:
     from miplearn import InternalSolver, Instance
@@ -24,9 +30,14 @@ class FeaturesExtractor:
         instance.features.constraints = self._extract_constraints(instance)
         instance.features.instance = self._extract_instance(instance, instance.features)
 
-    def _extract_variables(self, instance: "Instance") -> Dict:
-        variables = self.solver.get_empty_solution()
-        for (var_name, var_dict) in variables.items():
+    def _extract_variables(
+        self,
+        instance: "Instance",
+    ) -> Dict[str, Dict[VarIndex, VariableFeatures]]:
+        result: Dict[str, Dict[VarIndex, VariableFeatures]] = {}
+        empty_solution = self.solver.get_empty_solution()
+        for (var_name, var_dict) in empty_solution.items():
+            result[var_name] = {}
             for idx in var_dict.keys():
                 user_features = None
                 category = instance.get_variable_category(var_name, idx)
@@ -47,11 +58,11 @@ class FeaturesExtractor:
                             f"Found {type(v).__name__} instead "
                             f"for var={var_name}[{idx}]."
                         )
-                var_dict[idx] = {
-                    "Category": category,
-                    "User features": user_features,
-                }
-        return variables
+                result[var_name][idx] = VariableFeatures(
+                    category=category,
+                    user_features=user_features,
+                )
+        return result
 
     def _extract_constraints(
         self,

miplearn/solvers/internal.py

@@ -274,7 +274,7 @@ class InternalSolver(ABC):
         pass
 
     @abstractmethod
-    def get_empty_solution(self) -> Dict:
+    def get_empty_solution(self) -> Dict[str, Dict[VarIndex, Optional[float]]]:
         """
         Returns a dictionary with the same shape as the one produced by
         `get_solution`, but with all values set to None. This method is

miplearn/types.py

@@ -87,14 +87,12 @@ InstanceFeatures = TypedDict(
     total=False,
 )
 
-VariableFeatures = TypedDict(
-    "VariableFeatures",
-    {
-        "Category": Optional[Hashable],
-        "User features": Optional[List[float]],
-    },
-    total=False,
-)
+
+@dataclass
+class VariableFeatures:
+    category: Optional[Hashable] = None
+    user_features: Optional[List[float]] = None
+
 
 ConstraintFeatures = TypedDict(
     "ConstraintFeatures",