Modularize LearningSolver into components; implement branch-priority

miplearn/warmstart.py

@@ -2,7 +2,11 @@
 # Copyright (C) 2019-2020 Argonne National Laboratory. All rights reserved.
 # Written by Alinson S. Xavier <axavier@anl.gov>
 
+from . import Component
+from .transformers import PerVariableTransformer
+
 from abc import ABC, abstractmethod
+from copy import deepcopy
 import numpy as np
 from sklearn.pipeline import make_pipeline
 from sklearn.linear_model import LogisticRegression
@@ -10,6 +14,7 @@ from sklearn.preprocessing import StandardScaler
 from sklearn.model_selection import cross_val_score
 from sklearn.neighbors import KNeighborsClassifier
 
+
 class WarmStartPredictor(ABC):
     def __init__(self, thr_clip=[0.50, 0.50]):
         self.models = [None, None]
@@ -105,4 +110,79 @@ class KnnWarmStartPredictor(WarmStartPredictor):
         
         knn = KNeighborsClassifier(n_neighbors=self.k)
         knn.fit(x_train, y_train)
-        return knn
+        return knn
+    
+    
+class WarmStartComponent(Component):
+    def __init__(self,
+                 predictor_prototype=LogisticWarmStartPredictor(),
+                 mode="exact",
+                ):
+        self.mode = mode
+        self.transformer = PerVariableTransformer()
+        self.x_train = {}
+        self.y_train = {}
+        self.predictors = {}
+        self.predictor_prototype = predictor_prototype
+    
+    def before_solve(self, solver, instance, model):
+        var_split = self.transformer.split_variables(instance, model)
+        x_test = {}
+        
+        # Collect training data (x_train) and build x_test
+        for category in var_split.keys():
+            var_index_pairs = var_split[category]
+            x = self.transformer.transform_instance(instance, var_index_pairs)
+            x_test[category] = x
+            if category not in self.x_train.keys():
+                self.x_train[category] = x
+            else:
+                assert x.shape[1] == self.x_train[category].shape[1]
+                self.x_train[category] = np.vstack([self.x_train[category], x])
+        
+        # Predict solutions
+        for category in var_split.keys():
+            var_index_pairs = var_split[category]
+            if category in self.predictors.keys():
+                ws = self.predictors[category].predict(x_test[category])
+                assert ws.shape == (len(var_index_pairs), 2)
+                for i in range(len(var_index_pairs)):
+                    var, index = var_index_pairs[i]
+                    if self.mode == "heuristic":
+                        if ws[i,0] == 1:
+                            var[index].fix(0)
+                        elif ws[i,1] == 1:
+                            var[index].fix(1)
+                    else:
+                        if ws[i,0] == 1:
+                            var[index].value = 0
+                        elif ws[i,1] == 1:
+                            var[index].value = 1
+
+    def after_solve(self, solver, instance, model):
+        var_split = self.transformer.split_variables(instance, model)
+        for category in var_split.keys():
+            var_index_pairs = var_split[category]
+            y = self.transformer.transform_solution(var_index_pairs)
+            if category not in self.y_train.keys():
+                self.y_train[category] = y
+            else:
+                self.y_train[category] = np.vstack([self.y_train[category], y])
+                
+    def fit(self, solver):
+        for category in self.x_train.keys():
+            x_train = self.x_train[category]
+            y_train = self.y_train[category]
+            self.predictors[category] = deepcopy(self.predictor_prototype)
+            self.predictors[category].fit(x_train, y_train)
+
+    def merge(self, other):
+        for c in other.x_train.keys():
+            if c not in self.x_train:
+                self.x_train[c] = other.x_train[c]
+                self.y_train[c] = other.y_train[c]
+            else:
+                self.x_train[c] = np.vstack([self.x_train[c], other.x_train[c]])
+                self.y_train[c] = np.vstack([self.y_train[c], other.y_train[c]])
+            if (c in other.predictors.keys()) and (c not in self.predictors.keys()):
+                self.predictors[c] = other.predictors[c]