Implement solver.save and solver.load; update README

miplearn/solvers.py

@@ -7,6 +7,7 @@ from .warmstart import KnnWarmStartPredictor
 import pyomo.environ as pe
 import numpy as np
 from copy import deepcopy
+import pickle
 
 
 class LearningSolver:
@@ -87,6 +88,23 @@ class LearningSolver:
             y_train = y_train_dict[category]
             self.ws_predictors[category] = deepcopy(self.ws_predictor_prototype)
             self.ws_predictors[category].fit(x_train, y_train)
+            
+    def save(self, filename):
+        with open(filename, "wb") as file:
+            pickle.dump({
+                "version": 1,
+                "x_train": self.x_train,
+                "y_train": self.y_train,
+                "ws_predictors": self.ws_predictors,
+            }, file)
+
+    def load(self, filename):
+        with open(filename, "rb") as file:
+            data = pickle.load(file)
+            assert data["version"] == 1
+            self.x_train = data["x_train"]
+            self.y_train = data["y_train"]
+            self.ws_predictors = self.ws_predictors
 
     def _solve(self, model, tee=False):
         if hasattr(self.parent_solver, "set_instance"):

miplearn/tests/test_solver.py

@@ -14,3 +14,20 @@ def test_solver():
     solver.solve(instance)
     solver.fit()
     solver.solve(instance)
+
+def test_solve_save_load():
+    instance = KnapsackInstance2(weights=[23., 26., 20., 18.],
+                                 prices=[505., 352., 458., 220.],
+                                 capacity=67.)
+    solver = LearningSolver()
+    solver.solve(instance)
+    solver.fit()
+    solver.save("/tmp/knapsack_train.bin")
+    prev_x_train_len = len(solver.x_train)
+    prev_y_train_len = len(solver.y_train)
+    
+    solver = LearningSolver()
+    solver.load("/tmp/knapsack_train.bin")
+    assert len(solver.x_train) == prev_x_train_len
+    assert len(solver.y_train) == prev_y_train_len
+