From 9cf28f3cdca05c7eadbb82227d64d81dd11a0bc1 Mon Sep 17 00:00:00 2001
From: "Alinson S. Xavier" <git@axavier.org>
Date: Tue, 30 Mar 2021 21:29:33 -0500
Subject: [PATCH] Add variables to model features

---
 miplearn/features.py   |  9 +++++----
 miplearn/types.py      |  9 +++++----
 tests/test_features.py | 14 +++++++++++---
 3 files changed, 21 insertions(+), 11 deletions(-)

diff --git a/miplearn/features.py b/miplearn/features.py
index bf35193..f9e4cb2 100644
--- a/miplearn/features.py
+++ b/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(),
         }
diff --git a/miplearn/types.py b/miplearn/types.py
index 5728410..88d6283 100644
--- a/miplearn/types.py
+++ b/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,
 )
diff --git a/tests/test_features.py b/tests/test_features.py
index 5f597e9..883d92e 100644
--- a/tests/test_features.py
+++ b/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