Request variable features/categories in bulk

miplearn/problems/knapsack.py

@@ -48,13 +48,14 @@ class ChallengeA:
 class MultiKnapsackInstance(Instance):
     """Representation of the Multidimensional 0-1 Knapsack Problem.
 
-    Given a set of n items and m knapsacks, the problem is to find a subset of items S maximizing
-    sum(prices[i] for i in S). If selected, each item i occupies weights[i,j] units of space in
-    each knapsack j. Furthermore, each knapsack j has limited storage space, given by capacities[j].
+    Given a set of n items and m knapsacks, the problem is to find a subset of items
+    S maximizing sum(prices[i] for i in S). If selected, each item i occupies
+    weights[i,j] units of space in each knapsack j. Furthermore, each knapsack j has
+    limited storage space, given by capacities[j].
 
-    This implementation assigns a different category for each decision variable, and therefore
-    trains one ML model per variable. It is only suitable when training and test instances have
-    same size and items don't shuffle around.
+    This implementation assigns a different category for each decision variable,
+    and therefore trains one ML model per variable. It is only suitable when training
+    and test instances have same size and items don't shuffle around.
     """
 
     def __init__(
@@ -74,7 +75,6 @@ class MultiKnapsackInstance(Instance):
         self.prices = prices
         self.capacities = capacities
         self.weights = weights
-        self.varname_to_index = {f"x[{i}]": i for i in range(self.n)}
 
     @overrides
     def to_model(self) -> pe.ConcreteModel:
@@ -98,9 +98,11 @@ class MultiKnapsackInstance(Instance):
         return [float(np.mean(self.prices))] + list(self.capacities)
 
     @overrides
-    def get_variable_features(self, var_name: VariableName) -> List[float]:
-        index = self.varname_to_index[var_name]
-        return [self.prices[index]] + list(self.weights[:, index])
+    def get_variable_features(self) -> Dict[str, List[float]]:
+        return {
+            f"x[{i}]": [self.prices[i] + list(self.weights[:, i])]
+            for i in range(self.n)
+        }
 
 
 # noinspection PyPep8Naming
@@ -110,7 +112,7 @@ class MultiKnapsackGenerator:
         n: rv_frozen = randint(low=100, high=101),
         m: rv_frozen = randint(low=30, high=31),
         w: rv_frozen = randint(low=0, high=1000),
-        K: rv_frozen = randint(low=500, high=500),
+        K: rv_frozen = randint(low=500, high=501),
         u: rv_frozen = uniform(loc=0.0, scale=1.0),
         alpha: rv_frozen = uniform(loc=0.25, scale=0.0),
         fix_w: bool = False,
@@ -241,51 +243,3 @@ class MultiKnapsackGenerator:
             return MultiKnapsackInstance(p, b, w)
 
         return [_sample() for _ in range(n_samples)]
-
-
-class KnapsackInstance(Instance):
-    """
-    Simpler (one-dimensional) Knapsack Problem, used for testing.
-    """
-
-    def __init__(
-        self,
-        weights: List[float],
-        prices: List[float],
-        capacity: float,
-    ) -> None:
-        super().__init__()
-        self.weights = weights
-        self.prices = prices
-        self.capacity = capacity
-        self.varname_to_item: Dict[VariableName, int] = {
-            f"x[{i}]": i for i in range(len(self.weights))
-        }
-
-    @overrides
-    def to_model(self) -> pe.ConcreteModel:
-        model = pe.ConcreteModel()
-        items = range(len(self.weights))
-        model.x = pe.Var(items, domain=pe.Binary)
-        model.OBJ = pe.Objective(
-            expr=sum(model.x[v] * self.prices[v] for v in items), sense=pe.maximize
-        )
-        model.eq_capacity = pe.Constraint(
-            expr=sum(model.x[v] * self.weights[v] for v in items) <= self.capacity
-        )
-        return model
-
-    @overrides
-    def get_instance_features(self) -> List[float]:
-        return [
-            self.capacity,
-            np.average(self.weights),
-        ]
-
-    @overrides
-    def get_variable_features(self, var_name: VariableName) -> List[Category]:
-        item = self.varname_to_item[var_name]
-        return [
-            self.weights[item],
-            self.prices[item],
-        ]

miplearn/problems/stab.py

@@ -1,7 +1,7 @@
 #  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
 #  Copyright (C) 2020-2021, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
-from typing import List
+from typing import List, Dict, Hashable
 
 import networkx as nx
 import numpy as np
@@ -52,7 +52,6 @@ class MaxWeightStableSetInstance(Instance):
         self.graph = graph
         self.weights = weights
         self.nodes = list(self.graph.nodes)
-        self.varname_to_node = {f"x[{v}]": v for v in self.nodes}
 
     @overrides
     def to_model(self) -> pe.ConcreteModel:
@@ -68,24 +67,26 @@ class MaxWeightStableSetInstance(Instance):
         return model
 
     @overrides
-    def get_variable_features(self, var_name: VariableName) -> List[float]:
-        v1 = self.varname_to_node[var_name]
-        neighbor_weights = [0.0] * 15
-        neighbor_degrees = [100.0] * 15
-        for v2 in self.graph.neighbors(v1):
-            neighbor_weights += [self.weights[v2] / self.weights[v1]]
-            neighbor_degrees += [self.graph.degree(v2) / self.graph.degree(v1)]
-        neighbor_weights.sort(reverse=True)
-        neighbor_degrees.sort()
-        features = []
-        features += neighbor_weights[:5]
-        features += neighbor_degrees[:5]
-        features += [self.graph.degree(v1)]
+    def get_variable_features(self) -> Dict[str, List[float]]:
+        features = {}
+        for v1 in self.nodes:
+            neighbor_weights = [0.0] * 15
+            neighbor_degrees = [100.0] * 15
+            for v2 in self.graph.neighbors(v1):
+                neighbor_weights += [self.weights[v2] / self.weights[v1]]
+                neighbor_degrees += [self.graph.degree(v2) / self.graph.degree(v1)]
+            neighbor_weights.sort(reverse=True)
+            neighbor_degrees.sort()
+            f = []
+            f += neighbor_weights[:5]
+            f += neighbor_degrees[:5]
+            f += [self.graph.degree(v1)]
+            features[f"x[{v1}]"] = f
         return features
 
     @overrides
-    def get_variable_category(self, var: VariableName) -> Category:
-        return "default"
+    def get_variable_categories(self) -> Dict[str, Hashable]:
+        return {f"x[{v}]": "default" for v in self.nodes}
 
 
 class MaxWeightStableSetGenerator:

miplearn/problems/tsp.py

@@ -1,7 +1,7 @@
 #  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
 #  Copyright (C) 2020-2021, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
-from typing import List, Tuple, FrozenSet, Any, Optional, Hashable
+from typing import List, Tuple, FrozenSet, Any, Optional, Hashable, Dict
 
 import networkx as nx
 import numpy as np
@@ -59,7 +59,6 @@ class TravelingSalesmanInstance(Instance):
         self.edges = [
             (i, j) for i in range(self.n_cities) for j in range(i + 1, self.n_cities)
         ]
-        self.varname_to_index = {f"x[{e}]": e for e in self.edges}
 
     @overrides
     def to_model(self) -> pe.ConcreteModel:
@@ -83,8 +82,8 @@ class TravelingSalesmanInstance(Instance):
         return model
 
     @overrides
-    def get_variable_category(self, var_name: VariableName) -> Category:
-        return self.varname_to_index[var_name]
+    def get_variable_categories(self) -> Dict[str, Hashable]:
+        return {f"x[{e}]": f"x[{e}]" for e in self.edges}
 
     @overrides
     def find_violated_lazy_constraints(