Refer to variables by varname instead of (vname, index)

miplearn/problems/knapsack.py

@@ -1,14 +1,16 @@
 #  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
 
 import numpy as np
 import pyomo.environ as pe
+from overrides import overrides
 from scipy.stats import uniform, randint
 from scipy.stats.distributions import rv_frozen
 
 from miplearn.instance.base import Instance
+from miplearn.types import VariableName
 
 
 class ChallengeA:
@@ -67,7 +69,9 @@ 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):
         model = pe.ConcreteModel()
         model.x = pe.Var(range(self.n), domain=pe.Binary)
@@ -84,10 +88,13 @@ class MultiKnapsackInstance(Instance):
 
         return model
 
+    @overrides
     def get_instance_features(self):
         return [np.mean(self.prices)] + list(self.capacities)
 
-    def get_variable_features(self, var, index):
+    @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])
 
 
@@ -237,7 +244,11 @@ class KnapsackInstance(Instance):
         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):
         model = pe.ConcreteModel()
         items = range(len(self.weights))
@@ -250,16 +261,19 @@ class KnapsackInstance(Instance):
         )
         return model
 
+    @overrides
     def get_instance_features(self):
         return [
             self.capacity,
             np.average(self.weights),
         ]
 
-    def get_variable_features(self, var, index):
+    @overrides
+    def get_variable_features(self, var_name):
+        item = self.varname_to_item[var_name]
         return [
-            self.weights[index],
-            self.prices[index],
+            self.weights[item],
+            self.prices[item],
         ]
 
 
@@ -277,6 +291,7 @@ class GurobiKnapsackInstance(KnapsackInstance):
     ) -> None:
         super().__init__(weights, prices, capacity)
 
+    @overrides
     def to_model(self):
         import gurobipy as gp
         from gurobipy import GRB

miplearn/problems/stab.py

@@ -5,6 +5,7 @@
 import networkx as nx
 import numpy as np
 import pyomo.environ as pe
+from overrides import overrides
 from scipy.stats import uniform, randint
 from scipy.stats.distributions import rv_frozen
 
@@ -104,32 +105,38 @@ class MaxWeightStableSetInstance(Instance):
         super().__init__()
         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):
-        nodes = list(self.graph.nodes)
         model = pe.ConcreteModel()
-        model.x = pe.Var(nodes, domain=pe.Binary)
+        model.x = pe.Var(self.nodes, domain=pe.Binary)
         model.OBJ = pe.Objective(
-            expr=sum(model.x[v] * self.weights[v] for v in nodes), sense=pe.maximize
+            expr=sum(model.x[v] * self.weights[v] for v in self.nodes),
+            sense=pe.maximize,
         )
         model.clique_eqs = pe.ConstraintList()
         for clique in nx.find_cliques(self.graph):
-            model.clique_eqs.add(sum(model.x[i] for i in clique) <= 1)
+            model.clique_eqs.add(sum(model.x[v] for v in clique) <= 1)
         return model
 
-    def get_variable_features(self, var, index):
+    @overrides
+    def get_variable_features(self, var_name):
+        v1 = self.varname_to_node[var_name]
         neighbor_weights = [0] * 15
         neighbor_degrees = [100] * 15
-        for n in self.graph.neighbors(index):
-            neighbor_weights += [self.weights[n] / self.weights[index]]
-            neighbor_degrees += [self.graph.degree(n) / self.graph.degree(index)]
+        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(index)]
+        features += [self.graph.degree(v1)]
         return features
 
-    def get_variable_category(self, var, index):
+    @overrides
+    def get_variable_category(self, var):
         return "default"

miplearn/problems/tsp.py

@@ -5,6 +5,7 @@
 import networkx as nx
 import numpy as np
 import pyomo.environ as pe
+from overrides import overrides
 from scipy.spatial.distance import pdist, squareform
 from scipy.stats import uniform, randint
 from scipy.stats.distributions import rv_frozen
@@ -133,15 +134,17 @@ class TravelingSalesmanInstance(Instance):
         assert distances.shape == (n_cities, n_cities)
         self.n_cities = n_cities
         self.distances = distances
-
-    def to_model(self):
-        model = pe.ConcreteModel()
-        model.edges = edges = [
+        self.edges = [
             (i, j) for i in range(self.n_cities) for j in range(i + 1, self.n_cities)
         ]
-        model.x = pe.Var(edges, domain=pe.Binary)
+        self.varname_to_index = {f"x[{e}]": e for e in self.edges}
+
+    @overrides
+    def to_model(self):
+        model = pe.ConcreteModel()
+        model.x = pe.Var(self.edges, domain=pe.Binary)
         model.obj = pe.Objective(
-            expr=sum(model.x[i, j] * self.distances[i, j] for (i, j) in edges),
+            expr=sum(model.x[i, j] * self.distances[i, j] for (i, j) in self.edges),
             sense=pe.minimize,
         )
         model.eq_degree = pe.ConstraintList()
@@ -157,17 +160,13 @@ class TravelingSalesmanInstance(Instance):
             )
         return model
 
-    def get_instance_features(self):
-        return [0.0]
-
-    def get_variable_features(self, var_name, index):
-        return [0.0]
-
-    def get_variable_category(self, var_name, index):
-        return index
+    @overrides
+    def get_variable_category(self, var_name):
+        return self.varname_to_index[var_name]
 
+    @overrides
     def find_violated_lazy_constraints(self, model):
-        selected_edges = [e for e in model.edges if model.x[e].value > 0.5]
+        selected_edges = [e for e in self.edges if model.x[e].value > 0.5]
         graph = nx.Graph()
         graph.add_edges_from(selected_edges)
         components = [frozenset(c) for c in list(nx.connected_components(graph))]
@@ -177,10 +176,11 @@ class TravelingSalesmanInstance(Instance):
                 violations += [c]
         return violations
 
+    @overrides
     def build_lazy_constraint(self, model, component):
         cut_edges = [
             e
-            for e in model.edges
+            for e in self.edges
             if (e[0] in component and e[1] not in component)
             or (e[0] not in component and e[1] in component)
         ]