Add variables to model features

miplearn/features.py

@@ -21,10 +21,11 @@ class ModelFeaturesExtractor:
         constraints: Dict[str, ConstraintFeatures] = {}
         for cid in self.solver.get_constraint_ids():
             constraints[cid] = {
-                "rhs": self.solver.get_constraint_rhs(cid),
-                "lhs": self.solver.get_constraint_lhs(cid),
-                "sense": self.solver.get_constraint_sense(cid),
+                "RHS": self.solver.get_constraint_rhs(cid),
+                "LHS": self.solver.get_constraint_lhs(cid),
+                "Sense": self.solver.get_constraint_sense(cid),
             }
         return {
-            "constraints": constraints,
+            "Constraints": constraints,
+            "Variables": self.solver.get_empty_solution(),
         }

miplearn/types.py

@@ -76,9 +76,9 @@ LearningSolveStats = TypedDict(
 ConstraintFeatures = TypedDict(
     "ConstraintFeatures",
     {
-        "rhs": float,
-        "lhs": Dict[str, float],
-        "sense": str,
+        "RHS": float,
+        "LHS": Dict[str, float],
+        "Sense": str,
     },
     total=False,
 )
@@ -86,7 +86,8 @@ ConstraintFeatures = TypedDict(
 ModelFeatures = TypedDict(
     "ModelFeatures",
     {
-        "constraints": Dict[str, ConstraintFeatures],
+        "Variables": Solution,
+        "Constraints": Dict[str, ConstraintFeatures],
     },
     total=False,
 )

tests/test_features.py

@@ -21,11 +21,19 @@ def test_knapsack() -> None:
 
         # Test constraint features
         print(solver, features)
-        assert features["constraints"]["eq_capacity"]["lhs"] == {
+        assert features["Variables"] == {
+            "x": {
+                0: None,
+                1: None,
+                2: None,
+                3: None,
+            }
+        }
+        assert features["Constraints"]["eq_capacity"]["LHS"] == {
             "x[0]": 23.0,
             "x[1]": 26.0,
             "x[2]": 20.0,
             "x[3]": 18.0,
         }
-        assert features["constraints"]["eq_capacity"]["sense"] == "<"
-        assert features["constraints"]["eq_capacity"]["rhs"] == 67.0
+        assert features["Constraints"]["eq_capacity"]["Sense"] == "<"
+        assert features["Constraints"]["eq_capacity"]["RHS"] == 67.0