Finish TSP implementation; improve performance of Extractors

miplearn/extractors.py

@@ -5,6 +5,10 @@
 import numpy as np
 from abc import ABC, abstractmethod
 from pyomo.core import Var
+from tqdm.auto import tqdm, trange
+from p_tqdm import p_map
+import logging
+logger = logging.getLogger(__name__)
 
 
 class Extractor(ABC):
@@ -13,59 +17,39 @@ class Extractor(ABC):
         pass
     
     @staticmethod
-    def split_variables(instance, model):
+    def split_variables(instance):
+        assert hasattr(instance, "lp_solution")
         result = {}
-        for var in model.component_objects(Var):
-            for index in var:
-                category = instance.get_variable_category(var, index)
+        for var_name in instance.lp_solution:
+            for index in instance.lp_solution[var_name]:
+                category = instance.get_variable_category(var_name, index)
                 if category is None:
                     continue
-                if category not in result.keys():
+                if category not in result:
                     result[category] = []
-                result[category] += [(var, index)]
+                result[category] += [(var_name, index)]
         return result
-    
-    @staticmethod
-    def merge(partial_results, vertical=False):
-        results = {}
-        all_categories = set()
-        for pr in partial_results:
-            all_categories |= pr.keys()
-        for category in all_categories:
-            results[category] = []
-            for pr in partial_results:
-                if category in pr.keys():
-                    results[category] += [pr[category]]
-            if vertical:
-                results[category] = np.vstack(results[category])
-            else:
-                results[category] = np.hstack(results[category])
-        return results
 
     
 class VariableFeaturesExtractor(Extractor):
-    def extract(self,
-                instances,
-                models=None,
-               ):
+    def extract(self, instances):
         result = {}
-        if models is None:
-            models = [instance.to_model() for instance in instances]
-        for (index, instance) in enumerate(instances):
-            model = models[index]
+        for instance in tqdm(instances,
+                             desc="Extract var features",
+                             disable=len(instances) < 5):
             instance_features = instance.get_instance_features()
-            var_split = self.split_variables(instance, model)
+            var_split = self.split_variables(instance)
             for (category, var_index_pairs) in var_split.items():
-                if category not in result.keys():
+                if category not in result:
                     result[category] = []
-                for (var, index) in var_index_pairs:
-                    result[category] += [np.hstack([
-                        instance_features,
-                        instance.get_variable_features(var, index),
-                        instance.lp_solution[str(var)][index],
-                    ])]
-        for category in result.keys():
-            result[category] = np.vstack(result[category])
+                for (var_name, index) in var_index_pairs:
+                    result[category] += [
+                        instance_features.tolist() + \
+                        instance.get_variable_features(var_name, index).tolist() + \
+                        [instance.lp_solution[var_name][index]]
+                    ]
+        for category in result:
+            result[category] = np.array(result[category])
         return result
 
 
@@ -73,39 +57,29 @@ class SolutionExtractor(Extractor):
     def __init__(self, relaxation=False):
         self.relaxation = relaxation
         
-    def extract(self, instances, models=None):
+    def extract(self, instances):
         result = {}
-        if models is None:
-            models = [instance.to_model() for instance in instances]
-        for (index, instance) in enumerate(instances):
-            model = models[index]
-            var_split = self.split_variables(instance, model)
+        for instance in tqdm(instances,
+                             desc="Extract solution",
+                             disable=len(instances) < 5):
+            var_split = self.split_variables(instance)
             for (category, var_index_pairs) in var_split.items():
-                if category not in result.keys():
+                if category not in result:
                     result[category] = []
-                for (var, index) in var_index_pairs:
+                for (var_name, index) in var_index_pairs:
                     if self.relaxation:
-                        v = instance.lp_solution[str(var)][index]
+                        v = instance.lp_solution[var_name][index]
                     else:
-                        v = instance.solution[str(var)][index]
+                        v = instance.solution[var_name][index]
                     if v is None:
                         result[category] += [[0, 0]]
                     else:
                         result[category] += [[1 - v, v]]
-        for category in result.keys():
-            result[category] = np.vstack(result[category])            
+        for category in result:
+            result[category] = np.array(result[category])            
         return result
     
     
-class CombinedExtractor(Extractor):
-    def __init__(self, extractors):
-        self.extractors = extractors
-    
-    def extract(self, instances, models):
-        return self.merge([ex.extract(instances, models)
-                           for ex in self.extractors])
-
-    
 class InstanceFeaturesExtractor(Extractor):
     def extract(self, instances, models=None):
         return np.vstack([