Merge branch 'feature/tsp'

2025-12-06 09:28:51 -06:00 · 2020-02-25 22:33:02 -06:00
parent 477e7677ed 0b04fa93da
commit 9a2dd52566
37 changed files with 1617 additions and 643 deletions
--- a/benchmark/Makefile
+++ b/benchmark/Makefile
@@ -5,7 +5,8 @@
 CHALLENGES := \
    stab/ChallengeA \
-    knapsack/ChallengeA
+    knapsack/ChallengeA \
    tsp/ChallengeA
 main: $(addsuffix /performance.png, $(CHALLENGES))
--- a/benchmark/benchmark.py
+++ b/benchmark/benchmark.py
@@ -20,9 +20,16 @@ from miplearn import (LearningSolver, BenchmarkRunner)
 from numpy import median
 import pyomo.environ as pe
 import pickle
 import logging
 import sys
 logging.basicConfig(format='%(asctime)s %(levelname).1s %(name)s: %(message)12s',
                    datefmt='%H:%M:%S',
                    level=logging.INFO,
                    stream=sys.stdout)
 logging.getLogger('pyomo.core').setLevel(logging.ERROR)
 logging.getLogger('miplearn').setLevel(logging.INFO)
 logger = logging.getLogger("benchmark")
 n_jobs = 10
 time_limit = 900
@@ -34,7 +41,7 @@ pathlib.Path(basepath).mkdir(parents=True, exist_ok=True)
 def save(obj, filename):
-    print("Writing %s..." % filename)
+    logger.info("Writing %s..." % filename)
    with open(filename, "wb") as file:
        pickle.dump(obj, file)
@@ -55,7 +62,6 @@ def train():
                            solver=internal_solver,
                            components={})
    solver.parallel_solve(train_instances, n_jobs=n_jobs)
    solver.fit(n_jobs=n_jobs)
    save(train_instances, "%s/train_instances.bin" % basepath)
    save(test_instances, "%s/test_instances.bin" % basepath)
@@ -65,6 +71,7 @@ def test_baseline():
    solvers = {
        "baseline": LearningSolver(
            time_limit=time_limit,
            solver=internal_solver,
            components={},
        ),
    }
@@ -74,22 +81,28 @@ def test_baseline():
 def test_ml():
    logger.info("Loading instances...")
    train_instances = load("%s/train_instances.bin" % basepath)
    test_instances = load("%s/test_instances.bin" % basepath)
    solvers = {
        "ml-exact": LearningSolver(
-           time_limit=time_limit,
+            time_limit=time_limit,
            solver=internal_solver,
        ),
        "ml-heuristic": LearningSolver(
            time_limit=time_limit,
            solver=internal_solver,
            mode="heuristic",
        ),
    }
    benchmark = BenchmarkRunner(solvers)
    logger.info("Loading results...")
    benchmark.load_results("%s/benchmark_baseline.csv" % basepath)
    logger.info("Fitting...")
    benchmark.fit(train_instances)
    logger.info("Solving...")
    benchmark.parallel_solve(test_instances, n_jobs=n_jobs)
-    benchmark.save_results("%s/benchmark_ml.csv" % basepath)    
+    benchmark.save_results("%s/benchmark_ml.csv" % basepath)
 def charts():
--- a/benchmark/tsp/ChallengeA/benchmark_baseline.csv
+++ b/benchmark/tsp/ChallengeA/benchmark_baseline.csv
@@ -0,0 +1,51 @@
 ,Solver,Instance,Wallclock Time,Lower Bound,Upper Bound,Gap,Nodes,Mode,Relative Lower Bound,Relative Upper Bound,Relative Wallclock Time,Relative Gap,Relative Nodes
 0,baseline,0,88.44052076339722,13538.699999999997,13540.0,9.602103599333102e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 1,baseline,1,74.86838150024414,13565.666666666666,13567.0,9.828734304744377e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 2,baseline,2,113.96127772331238,13560.699999999997,13562.0,9.586525769340157e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 3,baseline,3,91.72307801246643,13520.666666666666,13522.0,9.861446674231594e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 4,baseline,4,75.7019145488739,13532.98571428571,13534.0,7.494914542165241e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 5,baseline,5,148.34671473503113,13530.670398009952,13532.0,9.82657880901199e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 6,baseline,6,128.75406980514526,13533.647058823528,13535.0,9.99687054488398e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 7,baseline,7,89.82294702529907,13611.833333333334,13613.0,8.57097378506001e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 8,baseline,8,163.10344243049622,13578.666666666664,13580.0,9.819324430497046e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 9,baseline,9,110.7302086353302,13582.666666666664,13584.0,9.816432708371643e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 10,baseline,10,969.3387920856476,13576.642857142855,13578.0,9.996159370362496e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 11,baseline,11,74.61696720123291,13575.0,13575.0,0.0,1,exact,1.0,1.0,1.0,,1.0
 12,baseline,12,83.92108988761902,13542.649999999998,13544.0,9.968506902284139e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 13,baseline,13,71.3746600151062,13532.75,13534.0,9.2368513421145e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 14,baseline,14,95.41746068000793,13549.66666666667,13551.0,9.840340475758086e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 15,baseline,15,119.16796040534973,13592.724074074078,13594.0,9.386830181858343e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 16,baseline,16,236.61669611930847,13592.65,13594.0,9.931838162539047e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 17,baseline,17,151.28878140449524,13542.65,13544.0,9.968506902270707e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 18,baseline,18,85.78852319717407,13523.833333333332,13525.0,8.62674537545725e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 19,baseline,19,101.29387998580933,13562.666666666662,13564.0,9.830908375965233e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 20,baseline,20,158.05654454231262,13567.666666666666,13569.0,9.827285458078813e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 21,baseline,21,142.137060880661,13564.750000000002,13566.0,9.215061095841655e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 22,baseline,22,75.79312753677368,13563.714285714286,13565.0,9.479072314785627e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 23,baseline,23,125.0184965133667,13578.651041666662,13580.0,9.934406070221867e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 24,baseline,24,110.1647527217865,13541.666666666666,13543.0,9.846153846158324e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 25,baseline,25,107.13047480583191,13540.75,13542.0,9.231394125140778e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 26,baseline,26,86.15372657775879,13530.672413793101,13532.0,9.811679466463352e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 27,baseline,27,94.40602779388428,13520.7,13522.0,9.614886803192678e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 28,baseline,28,65.10137605667114,13569.749999999998,13571.0,9.21166565339685e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 29,baseline,29,62.560155391693115,13593.777777777777,13595.0,8.991041653046851e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 30,baseline,30,123.58262610435486,13575.666666666668,13577.0,9.821494340354953e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 31,baseline,31,77.88054895401001,13580.75,13582.0,9.204204480606741e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 32,baseline,32,91.31177544593811,13523.0,13524.0,7.394808844191378e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 33,baseline,33,125.20542597770691,13546.724137931033,13548.0,9.418233190376748e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 34,baseline,34,93.4059066772461,13555.999999999996,13557.0,7.376807317819698e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 35,baseline,35,133.77592086791992,13566.671171171172,13568.0,9.794803839953082e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 36,baseline,36,105.60492086410522,13552.674418604645,13554.0,9.780958019141685e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 37,baseline,37,106.68410634994507,13530.666666666666,13532.0,9.854158454872436e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 38,baseline,38,76.45212483406067,13512.666666666666,13514.0,9.86728501653219e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 39,baseline,39,66.65553593635559,13536.666666666666,13538.0,9.849790691952276e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 40,baseline,40,72.2125403881073,13578.0,13578.0,0.0,1,exact,1.0,1.0,1.0,,1.0
 41,baseline,41,68.3276731967926,13526.0,13526.0,0.0,1,exact,1.0,1.0,1.0,,1.0
 42,baseline,42,91.24437737464905,13527.7,13529.0,9.609911514886287e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 43,baseline,43,197.3597583770752,13563.65151515152,13565.0,9.941901315984782e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 44,baseline,44,68.33869886398315,13551.749999999995,13553.0,9.223900972239434e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 45,baseline,45,88.19205832481384,13519.66666666667,13521.0,9.862176089131646e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 46,baseline,46,85.30370616912842,13541.745070422534,13543.0,9.267118609455138e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 47,baseline,47,80.86965203285217,13562.678333333333,13564.0,9.744879545055802e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 48,baseline,48,88.52302312850952,13550.6875,13552.0,9.685855422464727e-05,1,exact,1.0,1.0,1.0,1.0,1.0
 49,baseline,49,130.87233448028564,13523.000000000002,13524.0,7.394808844177925e-05,1,exact,1.0,1.0,1.0,1.0,1.0
--- a/benchmark/tsp/ChallengeA/benchmark_ml.csv
+++ b/benchmark/tsp/ChallengeA/benchmark_ml.csv
@@ -0,0 +1,151 @@
 ,Solver,Instance,Wallclock Time,Lower Bound,Upper Bound,Gap,Nodes,Mode,Relative Lower Bound,Relative Upper Bound,Relative Wallclock Time,Relative Gap,Relative Nodes
 0,baseline,0,88.44052076339722,13538.699999999995,13540.0,9.602103599333102e-05,1,exact,0.999990151785604,1.0,41.72369503515307,1.1142966881214733,1.0
 1,baseline,1,74.86838150024414,13565.666666666664,13567.0,9.828734304744377e-05,1,exact,0.9999017223164048,1.0,76.61356087521061,inf,1.0
 2,baseline,2,113.96127772331238,13560.699999999995,13562.0,9.586525769340156e-05,1,exact,0.9999235465104994,1.0,76.24000609297043,1.0400038346080955,1.0
 3,baseline,3,91.72307801246643,13520.666666666664,13522.0,9.861446674231594e-05,1,exact,0.9999013952571116,1.0,94.94335769471294,inf,1.0
 4,baseline,4,75.70191454887392,13532.98571428571,13534.0,7.49491454216524e-05,1,exact,0.9999250564715317,1.0,41.48575715983321,inf,1.0
 5,baseline,5,148.34671473503113,13530.670398009952,13532.0,9.82657880901199e-05,1,exact,0.9999978122351358,1.0,119.90309626530366,1.0227729991614376,1.0
 6,baseline,6,128.75406980514526,13533.647058823528,13535.0,9.99687054488398e-05,1,exact,0.9999000412872943,1.0,89.50928932403158,inf,1.0
 7,baseline,7,89.82294702529906,13611.833333333336,13613.0,8.570973785060009e-05,1,exact,0.9999142976076791,1.0,58.998624886131445,inf,1.0
 8,baseline,8,163.10344243049622,13578.666666666664,13580.0,9.819324430497046e-05,1,exact,0.9999999999999999,1.0,89.45505955167431,1.0000000000013642,1.0
 9,baseline,9,110.7302086353302,13582.666666666664,13584.0,9.816432708371644e-05,1,exact,0.9999018453082056,1.0,69.69178446747482,inf,1.0
 10,baseline,10,969.3387920856476,13576.642857142857,13578.0,9.996159370362496e-05,1,exact,0.9999000483976179,1.0,515.7574545488682,inf,1.0
 11,baseline,11,74.61696720123291,13575.0,13575.0,0.0,1,exact,1.0,1.0,89.50190791176391,,1.0
 12,baseline,12,83.92108988761902,13542.649999999998,13544.0,9.968506902284139e-05,1,exact,0.9999003248670997,1.0,55.69203928988848,inf,1.0
 13,baseline,13,71.3746600151062,13532.75,13534.0,9.236851342114499e-05,1,exact,0.9999076400177331,1.0,52.71414831739878,inf,1.0
 14,baseline,14,95.41746068000792,13549.66666666667,13551.0,9.840340475758086e-05,1,exact,0.9999016062775197,1.0,64.56994873575405,inf,1.0
 15,baseline,15,119.16796040534972,13592.724074074076,13594.0,9.386830181858343e-05,1,exact,0.9999061405086124,1.0,65.40523211338186,inf,1.0
 16,baseline,16,236.61669611930847,13592.65,13594.0,9.931838162539048e-05,1,exact,0.9999742514529537,1.0,92.34020750352192,1.3500347614339328,1.0
 17,baseline,17,151.28878140449527,13542.65,13544.0,9.968506902270708e-05,1,exact,0.9999741563907554,1.0000738388835562,91.4641415949715,1.350034889774522,1.0
 18,baseline,18,85.78852319717406,13523.833333333332,13525.0,8.62674537545725e-05,1,exact,0.9999137399876771,1.0,96.1413272100395,inf,1.0
 19,baseline,19,101.29387998580931,13562.666666666662,13564.0,9.830908375965231e-05,1,exact,0.9999017005799663,1.0,75.05436661902205,inf,1.0
 20,baseline,20,158.05654454231262,13567.666666666664,13569.0,9.827285458078812e-05,1,exact,0.999901736802024,1.0,90.31845612470958,inf,1.0
 21,baseline,21,142.13706088066098,13564.750000000002,13566.0,9.215061095841657e-05,1,exact,0.9999078578799943,1.0,145.5161989355379,inf,1.0
 22,baseline,22,75.7931275367737,13563.714285714286,13565.0,9.479072314785627e-05,1,exact,0.9999052182612816,1.0,103.67687261190537,inf,1.0
 23,baseline,23,125.0184965133667,13578.651041666662,13580.0,9.934406070221867e-05,1,exact,0.9999006658075599,1.0,76.27885738292231,inf,1.0
 24,baseline,24,110.1647527217865,13541.666666666664,13543.0,9.846153846158324e-05,1,exact,0.9999015481552583,1.0,137.78651586217427,inf,1.0
 25,baseline,25,107.13047480583192,13540.75,13542.0,9.231394125140777e-05,1,exact,0.9999815375526179,1.0,81.76214710854182,1.2500230784875868,1.0
 26,baseline,26,86.15372657775879,13530.672413793101,13532.0,9.811679466463352e-05,1,exact,0.9999018928312963,1.0,95.4669269129848,inf,1.0
 27,baseline,27,94.40602779388428,13520.7,13522.0,9.614886803192679e-05,1,exact,0.9999778122919903,1.0,54.71143091363273,1.3000288446596822,1.0
 28,baseline,28,65.10137605667114,13569.749999999998,13571.0,9.21166565339685e-05,1,exact,0.9999078918281629,1.0,50.24970031074932,inf,1.0
 29,baseline,29,62.560155391693115,13593.777777777776,13595.0,8.991041653046851e-05,1,exact,0.9999100976666256,1.0,46.92879617539241,inf,1.0
 30,baseline,30,123.58262610435486,13575.666666666668,13577.0,9.821494340354953e-05,1,exact,0.9999754468670203,1.0,71.13568329179179,1.3333660716465885,1.0
 31,baseline,31,77.88054895401002,13580.75,13582.0,9.20420448060674e-05,1,exact,0.9999079664261522,1.0,101.52119482807224,inf,1.0
 32,baseline,32,91.31177544593812,13523.0,13524.0,7.394808844191378e-05,1,exact,1.0,1.0,58.01274392159489,1.0,1.0
 33,baseline,33,125.20542597770692,13546.724137931033,13548.0,9.418233190376748e-05,1,exact,0.9999796366672351,1.0,91.02421056827991,1.275888050298017,1.0
 34,baseline,34,93.4059066772461,13555.999999999995,13557.0,7.376807317819698e-05,1,exact,0.9999262373681489,1.0,75.47322729615996,inf,1.0
 35,baseline,35,133.77592086791992,13566.671171171172,13568.0,9.794803839953082e-05,1,exact,0.9999020615544791,1.0,84.38612997732044,inf,1.0
 36,baseline,36,105.60492086410522,13552.674418604644,13554.0,9.780958019141684e-05,1,exact,0.9999759771714486,1.0,58.58183671223716,1.3256132403342726,1.0
 37,baseline,37,106.68410634994508,13530.666666666664,13532.0,9.854158454872437e-05,1,exact,0.9999014681249382,1.0,136.8255659357326,inf,1.0
 38,baseline,38,76.45212483406067,13512.666666666664,13514.0,9.867285016532191e-05,1,exact,0.9999901328123457,1.0,84.32910693002859,1.1111220747609611,1.0
 39,baseline,39,66.65553593635559,13536.666666666664,13538.0,9.849790691952276e-05,1,exact,0.9999015117939625,1.0,85.20374060262657,inf,1.0
 40,baseline,40,72.2125403881073,13578.0,13578.0,0.0,1,exact,1.0,1.0,102.32184109577935,,1.0
 41,baseline,41,68.3276731967926,13526.0,13526.0,0.0,1,exact,1.0,1.0,36.988942309740004,,1.0
 42,baseline,42,91.24437737464905,13527.7,13529.0,9.609911514886287e-05,1,exact,0.9999039101190037,1.0,60.45086955649546,inf,1.0
 43,baseline,43,197.3597583770752,13563.651515151521,13565.0,9.941901315984781e-05,1,exact,0.9999743081061281,1.0,97.77650283861117,1.3485194945001757,1.0
 44,baseline,44,68.33869886398314,13551.749999999995,13553.0,9.223900972239434e-05,1,exact,0.9999907761841069,1.0,92.29526354225831,1.1111213598948202,1.0
 45,baseline,45,88.19205832481384,13519.66666666667,13521.0,9.862176089131646e-05,1,exact,0.9999013879644013,1.0,113.90345047382237,inf,1.0
 46,baseline,46,85.30370616912842,13541.745070422534,13543.0,9.267118609455138e-05,1,exact,0.9999073374010584,1.0,67.75852498047483,inf,1.0
 47,baseline,47,80.86965203285217,13562.678333333333,13564.0,9.744879545055802e-05,1,exact,0.9999025606998919,1.0,116.23732479169651,inf,1.0
 48,baseline,48,88.52302312850952,13550.6875,13552.0,9.685855422464728e-05,1,exact,0.9999031508264463,1.0,122.53937560747,inf,1.0
 49,baseline,49,130.87233448028564,13523.000000000002,13524.0,7.394808844177925e-05,1,exact,0.9999260573794737,1.0,97.84851972743084,inf,1.0
 50,ml-exact,0,9.014199018478394,13538.699999999999,13540.0,9.602103599319664e-05,1,exact,0.9999901517856042,1.0,4.252639939099353,1.114296688119914,1.0
 51,ml-exact,1,5.4113500118255615,13565.709677419356,13567.0,9.511648202170943e-05,1,exact,0.999904892564263,1.0,5.537488392836666,inf,1.0
 52,ml-exact,2,7.311800241470337,13560.749999999995,13562.0,9.217779252662703e-05,1,exact,0.9999272333539017,1.0,4.891588670265028,1.0,1.0
 53,ml-exact,3,22.324601411819458,13520.72222222222,13522.0,9.450514231252222e-05,1,exact,0.9999055037880654,1.0,23.108389547788498,inf,1.0
 54,ml-exact,4,7.276817560195923,13534.0,13534.0,0.0,1,exact,1.0,1.0,3.987802527818481,,1.0
 55,ml-exact,5,6.725529432296753,13530.699999999999,13532.0,9.607780824355662e-05,1,exact,1.0,1.0,5.435993674657236,1.0,1.0
 56,ml-exact,6,8.695855855941772,13533.665041782733,13535.0,9.863981509409236e-05,1,exact,0.9999013699137593,1.0,6.045322519960155,inf,1.0
 57,ml-exact,7,6.677775144577026,13612.0,13613.0,7.346459006758743e-05,1,exact,0.9999265408065819,1.0,4.386179299125712,inf,1.0
 58,ml-exact,8,11.08420991897583,13578.642857142859,13580.0,9.994687034775287e-05,1,exact,0.999998246549209,1.0,6.079201294649372,1.0178589276210523,1.0
 59,ml-exact,9,16.643534421920776,13582.66666666667,13584.0,9.816432708331462e-05,1,exact,0.9999018453082059,1.0,10.475168682554177,inf,1.0
 60,ml-exact,10,16.841771364212036,13576.691176470587,13578.0,9.640224649735783e-05,1,exact,0.999903607046,1.0,8.961024978903895,inf,1.0
 61,ml-exact,11,4.076904058456421,13573.75,13575.0,9.208951100469657e-05,1,exact,0.9999079189686925,1.0,4.89018389907214,inf,1.0
 62,ml-exact,12,6.3809661865234375,13542.750000000002,13544.0,9.230030828289534e-05,1,exact,0.9999077082102777,1.0,4.234561539217345,inf,1.0
 63,ml-exact,13,4.139940500259399,13532.999999999998,13534.0,7.389344565150514e-05,1,exact,0.9999261120141864,1.0,3.0575758610926025,inf,1.0
 64,ml-exact,14,7.238400936126709,13549.66666666667,13551.0,9.840340475758086e-05,1,exact,0.9999016062775197,1.0,4.898298215480208,inf,1.0
 65,ml-exact,15,7.4390175342559814,13592.6875,13594.0,9.655927129936593e-05,1,exact,0.9999034500514933,1.0,4.082898346741258,inf,1.0
 66,ml-exact,16,45.386595487594604,13592.666666666664,13594.0,9.809210849005043e-05,1,exact,0.9999754775742414,1.0,17.712222822558893,1.3333660307052555,1.0
 67,ml-exact,17,7.782220363616943,13541.750000000002,13543.0,9.230712426371635e-05,1,exact,0.9999077013955551,1.0,4.704870372100049,1.2501153839035106,1.0
 68,ml-exact,18,3.7182631492614746,13525.0,13525.0,0.0,1,exact,1.0,1.0,4.166976429522626,,1.0
 69,ml-exact,19,5.698674201965332,13564.0,13564.0,0.0,1,exact,1.0,1.0,4.222470132021685,,1.0
 70,ml-exact,20,27.741767644882202,13567.699999999997,13569.0,9.581579781414025e-05,1,exact,0.9999041933819734,1.0,15.852514244896831,inf,1.0
 71,ml-exact,21,5.77320122718811,13566.0,13566.0,0.0,1,exact,1.0,1.0,5.91045215839762,,1.0
 72,ml-exact,22,3.9558210372924805,13563.871192621193,13565.0,8.32216232944913e-05,1,exact,0.9999167853019678,1.0,5.41113907141357,inf,1.0
 73,ml-exact,23,23.586190462112427,13578.64814814815,13580.0,9.955717514005507e-05,1,exact,0.9999004527355044,1.0,14.390891817144492,inf,1.0
 74,ml-exact,24,4.234185695648193,13543.0,13543.0,0.0,1,exact,1.0,1.0,5.2958290206504826,,1.0
 75,ml-exact,25,4.053018093109131,13541.0,13542.0,7.384978952809984e-05,1,exact,0.9999999999999999,1.0,3.0932697924001986,1.0000000000018192,1.0
 76,ml-exact,26,4.456820011138916,13532.0,13532.0,0.0,1,exact,1.0,1.0,4.9386013486451725,,1.0
 77,ml-exact,27,6.604680299758911,13520.892857142859,13522.0,8.188385699369056e-05,1,exact,0.9999920758185681,1.0,3.827631755843238,1.10715163041169,1.0
 78,ml-exact,28,4.3237597942352295,13569.666666666666,13571.0,9.825837038497186e-05,1,exact,0.9999017512833738,1.0,3.3373739087612315,inf,1.0
 79,ml-exact,29,4.937336444854736,13593.833333333334,13595.0,8.582322866981795e-05,1,exact,0.9999141841363247,1.0,3.7036873425140273,inf,1.0
 80,ml-exact,30,7.1004478931427,13575.670776255707,13577.0,9.791219647266045e-05,1,exact,0.9999757495768788,1.0,4.087105351928353,1.3292559793128385,1.0
 81,ml-exact,31,4.129274129867554,13582.0,13582.0,0.0,1,exact,1.0,1.0,5.3827155697676625,,1.0
 82,ml-exact,32,4.985649824142456,13522.750000000002,13524.0,9.243681943378239e-05,1,exact,0.9999815129778896,1.0,3.1675128987275314,1.2500231092030394,1.0
 83,ml-exact,33,26.542813777923584,13546.649999999996,13548.0,9.9655634419137e-05,1,exact,0.9999741640215545,1.0,19.29659718442664,1.350034879473593,1.0
 84,ml-exact,34,5.016777992248535,13557.0,13557.0,0.0,1,exact,1.0,1.0,4.053624007009211,,1.0
 85,ml-exact,35,35.91903018951416,13566.666666666662,13568.0,9.828009828041115e-05,1,exact,0.9999017295597481,1.0,22.657799180648382,inf,1.0
 86,ml-exact,36,7.446932315826416,13552.654459753448,13554.0,9.928241368119515e-05,1,exact,0.9999745045195491,1.0,4.1310098938875015,1.345574552621238,1.0
 87,ml-exact,37,7.494789361953735,13530.666666666668,13532.0,9.854158454858992e-05,1,exact,0.9999014681249385,1.0,9.612292131450713,inf,1.0
 88,ml-exact,38,5.067941904067993,13512.8,13514.0,8.880468888762712e-05,1,exact,1.0,1.0,5.59009989154572,1.0,1.0
 89,ml-exact,39,6.453751564025879,13536.75,13538.0,9.234121927345928e-05,1,exact,0.99990766730684,1.0,8.249633979390692,inf,1.0
 90,ml-exact,40,5.8958964347839355,13578.0,13578.0,0.0,1,exact,1.0,1.0,8.354213476977858,,1.0
 91,ml-exact,41,4.61811900138855,13524.9455782313,13526.0,7.79612577811263e-05,1,exact,0.9999220448197028,1.0,2.5000022586733377,inf,1.0
 92,ml-exact,42,7.025639057159424,13527.725,13529.0,9.425088106090537e-05,1,exact,0.9999057580013305,1.0,4.654599027527176,inf,1.0
 93,ml-exact,43,10.472574234008789,13563.657608695652,13565.0,9.89697132643173e-05,1,exact,0.9999747573500186,1.0,5.188350921887166,1.3424251907171998,1.0
 94,ml-exact,44,4.061512231826782,13551.75,13553.0,9.223900972199163e-05,1,exact,0.9999907761841074,1.0,5.485301125832847,1.1111213598899692,1.0
 95,ml-exact,45,4.208792209625244,13521.0,13521.0,0.0,1,exact,1.0,1.0,5.435817738123648,,1.0
 96,ml-exact,46,4.613062381744385,13541.897435897436,13543.0,8.141873085242141e-05,1,exact,0.9999185878976177,1.0,3.664252312908731,inf,1.0
 97,ml-exact,47,4.933764457702637,13564.0,13564.0,0.0,1,exact,1.0,1.0,7.091505494332267,,1.0
 98,ml-exact,48,4.755006313323975,13550.955555555562,13552.0,7.707533540025109e-05,1,exact,0.9999229306047492,1.0,6.582191661014823,inf,1.0
 99,ml-exact,49,5.486748695373535,13522.728991596638,13524.0,9.3990525444403e-05,1,exact,0.9999060183079442,1.0,4.102243916490891,inf,1.0
 100,ml-heuristic,0,2.11967134475708,13538.833333333334,13540.0,8.61718759616949e-05,1,heuristic,1.0,1.0,1.0,1.0,1.0
 101,ml-heuristic,1,0.9772210121154785,13567.0,13567.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
 102,ml-heuristic,2,1.4947700500488281,13561.736842105263,13563.0,9.314130700537644e-05,1,heuristic,1.0,1.000073735437251,1.0,1.0104527831741152,1.0
 103,ml-heuristic,3,0.9660820960998535,13522.0,13522.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
 104,ml-heuristic,4,1.8247687816619873,13533.000000000002,13534.0,7.38934456512363e-05,1,heuristic,0.9999261120141866,1.0,1.0,inf,1.0
 105,ml-heuristic,5,1.2372217178344727,13530.666666666668,13532.0,9.854158454858992e-05,1,heuristic,0.9999975364664555,1.0,1.0,1.0256435523465277,1.0
 106,ml-heuristic,6,1.438443660736084,13535.0,13535.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
 107,ml-heuristic,7,1.5224583148956299,13613.0,13613.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
 108,ml-heuristic,8,1.8233003616333008,13578.666666666666,13580.0,9.819324430483649e-05,1,heuristic,1.0,1.0,1.0,1.0,1.0
 109,ml-heuristic,9,1.5888559818267822,13584.0,13584.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
 110,ml-heuristic,10,1.8794469833374023,13578.0,13578.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
 111,ml-heuristic,11,0.8336913585662842,13574.0,13575.0,7.367025195226167e-05,1,heuristic,0.999926335174954,1.0,1.0,inf,1.0
 112,ml-heuristic,12,1.5068776607513428,13544.0,13544.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
 113,ml-heuristic,13,1.353994369506836,13534.0,13534.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
 114,ml-heuristic,14,1.4777379035949707,13551.0,13551.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
 115,ml-heuristic,15,1.8219943046569824,13594.0,13594.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
 116,ml-heuristic,16,2.5624449253082275,13593.0,13594.0,7.356727727506805e-05,1,heuristic,1.0,1.0,1.0,1.0,1.0
 117,ml-heuristic,17,1.6540775299072266,13543.0,13544.0,7.383888355608063e-05,1,heuristic,1.0,1.0000738388835562,1.0,1.0,1.0
 118,ml-heuristic,18,0.8923168182373047,13525.0,13525.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
 119,ml-heuristic,19,1.3496067523956299,13563.0,13564.0,7.37300007373e-05,1,heuristic,0.9999262754349749,1.0,1.0,inf,1.0
 120,ml-heuristic,20,1.7499916553497314,13569.0,13569.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
 121,ml-heuristic,21,0.9767782688140869,13566.0,13566.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
 122,ml-heuristic,22,0.7310514450073242,13565.0,13565.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
 123,ml-heuristic,23,1.6389665603637695,13580.0,13580.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
 124,ml-heuristic,24,0.7995321750640869,13543.0,13543.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
 125,ml-heuristic,25,1.310269832611084,13541.000000000002,13542.0,7.38497895279655e-05,1,heuristic,1.0,1.0,1.0,1.0,1.0
 126,ml-heuristic,26,0.9024457931518555,13532.0,13532.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
 127,ml-heuristic,27,1.7255265712738037,13521.0,13522.0,7.395902669920864e-05,1,heuristic,1.0,1.0,1.0,1.0,1.0
 128,ml-heuristic,28,1.2955574989318848,13571.0,13571.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
 129,ml-heuristic,29,1.3330867290496826,13595.0,13595.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
 130,ml-heuristic,30,1.7372803688049316,13576.0,13577.0,7.365939893930465e-05,1,heuristic,1.0,1.0,1.0,1.0,1.0
 131,ml-heuristic,31,0.7671358585357666,13582.0,13582.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
 132,ml-heuristic,32,1.5739951133728027,13522.999999999998,13524.0,7.39480884420483e-05,1,heuristic,0.9999999999999999,1.0,1.0,1.0000000000018192,1.0
 133,ml-heuristic,33,1.3755178451538086,13546.999999999998,13548.0,7.381708127274076e-05,1,heuristic,1.0,1.0,1.0,1.0,1.0
 134,ml-heuristic,34,1.2376031875610352,13557.0,13557.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
 135,ml-heuristic,35,1.5852832794189453,13568.0,13568.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
 136,ml-heuristic,36,1.8026905059814453,13553.0,13554.0,7.378440197742197e-05,1,heuristic,1.0,1.0,1.0,1.0,1.0
 137,ml-heuristic,37,0.7797088623046875,13532.0,13532.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
 138,ml-heuristic,38,0.9065923690795898,13512.75,13514.0,9.250522654529981e-05,1,heuristic,0.9999962998046298,1.0,1.0,1.0416705210504744,1.0
 139,ml-heuristic,39,0.7823076248168945,13538.0,13538.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
 140,ml-heuristic,40,0.7057392597198486,13576.75,13578.0,9.206916235476089e-05,1,heuristic,0.9999079393135956,1.0,1.0,inf,1.0
 141,ml-heuristic,41,1.8472459316253662,13526.0,13526.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
 142,ml-heuristic,42,1.509397268295288,13529.0,13529.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
 143,ml-heuristic,43,2.0184783935546875,13564.0,13565.0,7.372456502506635e-05,1,heuristic,1.0,1.0,1.0,1.0,1.0
 144,ml-heuristic,44,0.7404356002807617,13551.875,13553.0,8.301434303371305e-05,1,heuristic,1.0,1.0,1.0,1.0,1.0
 145,ml-heuristic,45,0.7742702960968018,13520.0,13521.0,7.396449704142012e-05,1,heuristic,0.9999260409733008,1.0,1.0,inf,1.0
 146,ml-heuristic,46,1.258936882019043,13543.0,13543.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
 147,ml-heuristic,47,0.6957287788391113,13564.0,13564.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
 148,ml-heuristic,48,0.7224047183990479,13552.0,13552.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
 149,ml-heuristic,49,1.3374993801116943,13524.0,13524.0,0.0,1,heuristic,1.0,1.0,1.0,,1.0
--- a/benchmark/tsp/ChallengeA/performance.png
+++ b/benchmark/tsp/ChallengeA/performance.png
--- a/docs-src/css/custom.css
+++ b/docs-src/css/custom.css
--- a/docs-src/figures/benchmark_tsp_a.png
+++ b/docs-src/figures/benchmark_tsp_a.png
@@ -0,0 +1 @@
 ../../benchmark/tsp/ChallengeA/performance.png
--- a/docs-src/problems.md
+++ b/docs-src/problems.md
@@ -20,6 +20,8 @@ To illustrate the performance of `LearningSolver`, and to set a baseline for new
 All experiments presented here were performed on a Linux server (Ubuntu Linux 18.04 LTS) with Intel Xeon Gold 6230s (2 processors, 40 cores, 80 threads) and 256 GB RAM (DDR4, 2933 MHz). All solvers were restricted to use 4 threads, with no time limits, and 10 instances were solved simultaneously at a time.
 ## Maximum Weight Stable Set Problem
 ### Problem definition
@@ -45,6 +47,53 @@ MaxWeightStableSetGenerator(w=uniform(loc=100., scale=50.),
 ![alt](figures/benchmark_stab_a.png)
 ## Traveling Salesman Problem
 ### Problem definition
 Given a list of cities and the distance between each pair of cities, the problem asks for the
 shortest route starting at the first city, visiting each other city exactly once, then returning
 to the first city. This problem is a generalization of the Hamiltonian path problem, one of Karp's
 21 NP-complete problems.
 ### Random problem generator
 The class `TravelingSalesmanGenerator` can be used to generate random instances of this
 problem. Initially, the generator creates $n$ cities $(x_1,y_1),\ldots,(x_n,y_n) \in \mathbb{R}^2$,
 where $n, x_i$ and $y_i$ are sampled independently from the provided probability distributions `n`,
 `x` and `y`. For each pair of cities $(i,j)$, the distance $d_{i,j}$ between them is set to:
 $$
    d_{i,j} = \gamma_{i,j} \sqrt{(x_i-x_j)^2 + (y_i - y_j)^2}
 $$
 where $\gamma_{i,j}$ is sampled from the distribution `gamma`.
 If `fix_cities=True` is provided, the list of cities is kept the same for all generated instances.
 The $gamma$ values, and therefore also the distances, are still different.
 By default, all distances $d_{i,j}$ are rounded to the nearest integer.  If `round=False`
 is provided, this rounding will be disabled.
 ### Challenge A
 * Fixed list of 350 cities in the $[0, 1000]^2$ square
 * $\gamma_{i,j} \sim U(0.95, 1.05)$
 * 500 training instances, 50 test instances
 ```python
 TravelingSalesmanGenerator(x=uniform(loc=0.0, scale=1000.0),
                           y=uniform(loc=0.0, scale=1000.0),
                           n=randint(low=350, high=351),
                           gamma=uniform(loc=0.95, scale=0.1),
                           fix_cities=True,
                           round=True,
                          )
 ```
 ![alt](figures/benchmark_tsp_a.png)
 ## Multidimensional 0-1 Knapsack Problem
 ### Problem definition
@@ -115,3 +164,4 @@ MultiKnapsackGenerator(n=randint(low=250, high=251),
 ```
 ![alt](figures/benchmark_knapsack_a.png)
--- a/docs-src/usage.md
+++ b/docs-src/usage.md
@@ -35,16 +35,16 @@ Instances to be solved by `LearningSolver` must derive from the abstract class `
 * `instance.to_model()`, which returns a concrete Pyomo model corresponding to the instance;
 * `instance.get_instance_features()`, which returns a 1-dimensional Numpy array of (numerical) features describing the entire instance;
-* `instance.get_variable_features(var, index)`, which returns a 1-dimensional array of (numerical) features describing a particular decision variable.
+* `instance.get_variable_features(var_name, index)`, which returns a 1-dimensional array of (numerical) features describing a particular decision variable.
 The first method is used by `LearningSolver` to construct a concrete Pyomo model, which will be provided to the internal MIP solver. The user should keep a reference to this Pyomo model, in order to retrieve, for example, the optimal variable values.
 The second and third methods provide an encoding of the instance, which can be used by the ML models to make predictions. In the knapsack problem, for example, an implementation may decide to provide as instance features the average weights, average prices, number of items and the size of the knapsack. The weight and the price of each individual item could be provided as variable features. See `miplearn/problems/knapsack.py` for a concrete example.
-An optional method which can be implemented is `instance.get_variable_category(var, index)`, which returns a category (a string, an integer or any hashable type) for each decision variable. If two variables have the same category, `LearningSolver` will use the same internal ML model to predict the values of both variables. By default, all variables belong to the `"default"` category, and therefore only one ML model is used for all variables. If the returned category is `None`, ML predictors will ignore the variable.
+An optional method which can be implemented is `instance.get_variable_category(var_name, index)`, which returns a category (a string, an integer or any hashable type) for each decision variable. If two variables have the same category, `LearningSolver` will use the same internal ML model to predict the values of both variables. By default, all variables belong to the `"default"` category, and therefore only one ML model is used for all variables. If the returned category is `None`, ML predictors will ignore the variable.
-It is not necessary to have a one-to-one correspondence between features and problem instances. One important (and deliberate) limitation of MIPLearn, however, is that `get_instance_features()` must always return arrays of same length for all relevant instances of the problem. Similarly, `get_variable_features(var, index)` must also always return arrays of same length for all variables in each category. It is up to the user to decide how to encode variable-length characteristics of the problem into fixed-length vectors. In graph problems, for example, graph embeddings can be used to reduce the (variable-length) lists of nodes and edges into a fixed-length structure that still preserves some properties of the graph. Different instance encodings may have significant impact on performance.
+It is not necessary to have a one-to-one correspondence between features and problem instances. One important (and deliberate) limitation of MIPLearn, however, is that `get_instance_features()` must always return arrays of same length for all relevant instances of the problem. Similarly, `get_variable_features(var_name, index)` must also always return arrays of same length for all variables in each category. It is up to the user to decide how to encode variable-length characteristics of the problem into fixed-length vectors. In graph problems, for example, graph embeddings can be used to reduce the (variable-length) lists of nodes and edges into a fixed-length structure that still preserves some properties of the graph. Different instance encodings may have significant impact on performance.
 ### Obtaining heuristic solutions
--- a/docs/404.html
+++ b/docs/404.html
@@ -32,15 +32,6 @@
            <script src="https://cdn.jsdelivr.net/npm/respond.js@1.4.2/dest/respond.min.js"></script>
        <![endif]-->
    <script src="//ajax.googleapis.com/ajax/libs/webfont/1.6.26/webfont.js"></script>
    <script>
    WebFont.load({
        google: {
            families: ['Open Sans', 'PT Sans']
        }
    });
    </script>
@@ -142,9 +133,10 @@
        <hr>
        <p>
-        <small>Copyright © 2020, UChicago Argonne, LLC. All Rights Reserved.<br></small>
+        <small>Copyright © 2020, UChicago Argonne, LLC. All Rights Reserved.</small><br>
-        <small>Documentation built with <a href="http://www.mkdocs.org/">MkDocs</a>.</p></small>
+        <small>Documentation built with <a href="http://www.mkdocs.org/">MkDocs</a>.</small>
        </p>
--- a/docs/about/index.html
+++ b/docs/about/index.html
@@ -32,15 +32,6 @@
            <script src="https://cdn.jsdelivr.net/npm/respond.js@1.4.2/dest/respond.min.js"></script>
        <![endif]-->
    <script src="//ajax.googleapis.com/ajax/libs/webfont/1.6.26/webfont.js"></script>
    <script>
    WebFont.load({
        google: {
            families: ['Open Sans', 'PT Sans']
        }
    });
    </script>
@@ -200,9 +191,10 @@ POSSIBILITY OF SUCH DAMAGE.
        <hr>
        <p>
-        <small>Copyright © 2020, UChicago Argonne, LLC. All Rights Reserved.<br></small>
+        <small>Copyright © 2020, UChicago Argonne, LLC. All Rights Reserved.</small><br>
-        <small>Documentation built with <a href="http://www.mkdocs.org/">MkDocs</a>.</p></small>
+        <small>Documentation built with <a href="http://www.mkdocs.org/">MkDocs</a>.</small>
        </p>
--- a/docs/benchmark/index.html
+++ b/docs/benchmark/index.html
@@ -32,15 +32,6 @@
            <script src="https://cdn.jsdelivr.net/npm/respond.js@1.4.2/dest/respond.min.js"></script>
        <![endif]-->
    <script src="//ajax.googleapis.com/ajax/libs/webfont/1.6.26/webfont.js"></script>
    <script>
    WebFont.load({
        google: {
            families: ['Open Sans', 'PT Sans']
        }
    });
    </script>
@@ -209,9 +200,10 @@ benchmark.parallel_solve(test_instances)
        <hr>
        <p>
-        <small>Copyright © 2020, UChicago Argonne, LLC. All Rights Reserved.<br></small>
+        <small>Copyright © 2020, UChicago Argonne, LLC. All Rights Reserved.</small><br>
-        <small>Documentation built with <a href="http://www.mkdocs.org/">MkDocs</a>.</p></small>
+        <small>Documentation built with <a href="http://www.mkdocs.org/">MkDocs</a>.</small>
        </p>
--- a/docs/css/custom.css
+++ b/docs/css/custom.css
--- a/docs/customization/index.html
+++ b/docs/customization/index.html
@@ -32,15 +32,6 @@
            <script src="https://cdn.jsdelivr.net/npm/respond.js@1.4.2/dest/respond.min.js"></script>
        <![endif]-->
    <script src="//ajax.googleapis.com/ajax/libs/webfont/1.6.26/webfont.js"></script>
    <script>
    WebFont.load({
        google: {
            families: ['Open Sans', 'PT Sans']
        }
    });
    </script>
@@ -162,9 +153,10 @@ solver = LearningSolver(solver=&quot;cplex&quot;,
        <hr>
        <p>
-        <small>Copyright © 2020, UChicago Argonne, LLC. All Rights Reserved.<br></small>
+        <small>Copyright © 2020, UChicago Argonne, LLC. All Rights Reserved.</small><br>
-        <small>Documentation built with <a href="http://www.mkdocs.org/">MkDocs</a>.</p></small>
+        <small>Documentation built with <a href="http://www.mkdocs.org/">MkDocs</a>.</small>
        </p>
--- a/docs/figures/benchmark_tsp_a.png
+++ b/docs/figures/benchmark_tsp_a.png
--- a/docs/index.html
+++ b/docs/index.html
@@ -32,15 +32,6 @@
            <script src="https://cdn.jsdelivr.net/npm/respond.js@1.4.2/dest/respond.min.js"></script>
        <![endif]-->
    <script src="//ajax.googleapis.com/ajax/libs/webfont/1.6.26/webfont.js"></script>
    <script>
    WebFont.load({
        google: {
            families: ['Open Sans', 'PT Sans']
        }
    });
    </script>
@@ -188,9 +179,10 @@
        <hr>
        <p>
-        <small>Copyright © 2020, UChicago Argonne, LLC. All Rights Reserved.<br></small>
+        <small>Copyright © 2020, UChicago Argonne, LLC. All Rights Reserved.</small><br>
-        <small>Documentation built with <a href="http://www.mkdocs.org/">MkDocs</a>.</p></small>
+        <small>Documentation built with <a href="http://www.mkdocs.org/">MkDocs</a>.</small>
        </p>
@@ -276,6 +268,6 @@
 </html>
 <!--
-MkDocs version : 1.0.4
+MkDocs version : 1.1
-Build Date UTC : 2020-02-24 16:57:12
+Build Date UTC : 2020-02-26 04:29:16
 -->
--- a/docs/js/base.js
+++ b/docs/js/base.js
@@ -1,88 +1,78 @@
-function getSearchTerm()
+function getSearchTerm() {
 {
    var sPageURL = window.location.search.substring(1);
    var sURLVariables = sPageURL.split('&');
-    for (var i = 0; i < sURLVariables.length; i++)
+    for (var i = 0; i < sURLVariables.length; i++) {
    {
        var sParameterName = sURLVariables[i].split('=');
-        if (sParameterName[0] == 'q')
+        if (sParameterName[0] == 'q') {
        {
            return sParameterName[1];
        }
    }
 }
 $(document).ready(function() {
    /**
     * ------------------------------------------------------------------------
     * Cinder theme specific
     * ------------------------------------------------------------------------
     */
    hljs.initHighlightingOnLoad();
    /**
     * ------------------------------------------------------------------------
     * Taken from themes/mkdocs/js/base.js
     * ------------------------------------------------------------------------
     */
    var search_term = getSearchTerm(),
-    $search_modal = $('#mkdocs_search_modal'),
+        $search_modal = $('#mkdocs_search_modal'),
-    $keyboard_modal = $('#mkdocs_keyboard_modal');
+        $keyboard_modal = $('#mkdocs_keyboard_modal');
-    if(search_term){
+    if (search_term) {
        $search_modal.modal();
    }
    // make sure search input gets autofocus everytime modal opens.
-    $search_modal.on('shown.bs.modal', function () {
+    $search_modal.on('shown.bs.modal', function() {
        $search_modal.find('#mkdocs-search-query').focus();
    });
    // Close search modal when result is selected
    // The links get added later so listen to parent
    $('#mkdocs-search-results').click(function(e) {
-      if ($(e.target).is('a')) {
+        if ($(e.target).is('a')) {
-        $search_modal.modal('hide');
+            $search_modal.modal('hide');
-      }
+        }
    });
    if (typeof shortcuts !== 'undefined') {
-      // Populate keyboard modal with proper Keys
+        // Populate keyboard modal with proper Keys
-      $keyboard_modal.find('.help.shortcut kbd')[0].innerHTML = keyCodes[shortcuts.help];
+        $keyboard_modal.find('.help.shortcut kbd')[0].innerHTML = keyCodes[shortcuts.help];
-      $keyboard_modal.find('.prev.shortcut kbd')[0].innerHTML = keyCodes[shortcuts.previous];
+        $keyboard_modal.find('.prev.shortcut kbd')[0].innerHTML = keyCodes[shortcuts.previous];
-      $keyboard_modal.find('.next.shortcut kbd')[0].innerHTML = keyCodes[shortcuts.next];
+        $keyboard_modal.find('.next.shortcut kbd')[0].innerHTML = keyCodes[shortcuts.next];
-      $keyboard_modal.find('.search.shortcut kbd')[0].innerHTML = keyCodes[shortcuts.search];
+        $keyboard_modal.find('.search.shortcut kbd')[0].innerHTML = keyCodes[shortcuts.search];
-      // Keyboard navigation
+        // Keyboard navigation
-      document.addEventListener("keydown", function(e) {
+        document.addEventListener("keydown", function(e) {
-          if ($(e.target).is(':input')) return true;
+            if ($(e.target).is(':input')) return true;
-          var key = e.which || e.key || window.event && window.event.key;
+            var key = e.which || e.key || window.event && window.event.key;
-          var page;
+            var page;
-          switch (key) {
+            switch (key) {
-              case shortcuts.next:
+                case shortcuts.next:
-                  page = $('.navbar a[rel="next"]:first').prop('href');
+                    page = $('.navbar a[rel="next"]:first').prop('href');
-                  break;
+                    break;
-              case shortcuts.previous:
+                case shortcuts.previous:
-                  page = $('.navbar a[rel="prev"]:first').prop('href');
+                    page = $('.navbar a[rel="prev"]:first').prop('href');
-                  break;
+                    break;
-              case shortcuts.search:
+                case shortcuts.search:
-                  e.preventDefault();
+                    e.preventDefault();
-                  $keyboard_modal.modal('hide');
+                    $keyboard_modal.modal('hide');
-                  $search_modal.modal('show');
+                    $search_modal.modal('show');
-                  $search_modal.find('#mkdocs-search-query').focus();
+                    $search_modal.find('#mkdocs-search-query').focus();
-                  break;
+                    break;
-              case shortcuts.help:
+                case shortcuts.help:
-                  $search_modal.modal('hide');
+                    $search_modal.modal('hide');
-                  $keyboard_modal.modal('show');
+                    $keyboard_modal.modal('show');
-                  break;
+                    break;
-              default: break;
+                default:
-          }
+                    break;
-          if (page) {
+            }
-              $keyboard_modal.modal('hide');
+            if (page) {
-              window.location.href = page;
+                $keyboard_modal.modal('hide');
-          }
+                window.location.href = page;
-      });
+            }
        });
    }
    $('table').addClass('table table-striped table-hover');
@@ -121,115 +111,115 @@ $("li.disabled a").click(function() {
 // See https://www.cambiaresearch.com/articles/15/javascript-char-codes-key-codes
 // We only list common keys below. Obscure keys are omited and their use is discouraged.
 var keyCodes = {
-  8: 'backspace',
+    8: 'backspace',
-  9: 'tab',
+    9: 'tab',
-  13: 'enter',
+    13: 'enter',
-  16: 'shift',
+    16: 'shift',
-  17: 'ctrl',
+    17: 'ctrl',
-  18: 'alt',
+    18: 'alt',
-  19: 'pause/break',
+    19: 'pause/break',
-  20: 'caps lock',
+    20: 'caps lock',
-  27: 'escape',
+    27: 'escape',
-  32: 'spacebar',
+    32: 'spacebar',
-  33: 'page up',
+    33: 'page up',
-  34: 'page down',
+    34: 'page down',
-  35: 'end',
+    35: 'end',
-  36: 'home',
+    36: 'home',
-  37: '&larr;',
+    37: '&larr;',
-  38: '&uarr;',
+    38: '&uarr;',
-  39: '&rarr;',
+    39: '&rarr;',
-  40: '&darr;',
+    40: '&darr;',
-  45: 'insert',
+    45: 'insert',
-  46: 'delete',
+    46: 'delete',
-  48: '0',
+    48: '0',
-  49: '1',
+    49: '1',
-  50: '2',
+    50: '2',
-  51: '3',
+    51: '3',
-  52: '4',
+    52: '4',
-  53: '5',
+    53: '5',
-  54: '6',
+    54: '6',
-  55: '7',
+    55: '7',
-  56: '8',
+    56: '8',
-  57: '9',
+    57: '9',
-  65: 'a',
+    65: 'a',
-  66: 'b',
+    66: 'b',
-  67: 'c',
+    67: 'c',
-  68: 'd',
+    68: 'd',
-  69: 'e',
+    69: 'e',
-  70: 'f',
+    70: 'f',
-  71: 'g',
+    71: 'g',
-  72: 'h',
+    72: 'h',
-  73: 'i',
+    73: 'i',
-  74: 'j',
+    74: 'j',
-  75: 'k',
+    75: 'k',
-  76: 'l',
+    76: 'l',
-  77: 'm',
+    77: 'm',
-  78: 'n',
+    78: 'n',
-  79: 'o',
+    79: 'o',
-  80: 'p',
+    80: 'p',
-  81: 'q',
+    81: 'q',
-  82: 'r',
+    82: 'r',
-  83: 's',
+    83: 's',
-  84: 't',
+    84: 't',
-  85: 'u',
+    85: 'u',
-  86: 'v',
+    86: 'v',
-  87: 'w',
+    87: 'w',
-  88: 'x',
+    88: 'x',
-  89: 'y',
+    89: 'y',
-  90: 'z',
+    90: 'z',
-  91: 'Left Windows Key / Left ⌘',
+    91: 'Left Windows Key / Left ⌘',
-  92: 'Right Windows Key',
+    92: 'Right Windows Key',
-  93: 'Windows Menu / Right ⌘',
+    93: 'Windows Menu / Right ⌘',
-  96: 'numpad 0',
+    96: 'numpad 0',
-  97: 'numpad 1',
+    97: 'numpad 1',
-  98: 'numpad 2',
+    98: 'numpad 2',
-  99: 'numpad 3',
+    99: 'numpad 3',
-  100: 'numpad 4',
+    100: 'numpad 4',
-  101: 'numpad 5',
+    101: 'numpad 5',
-  102: 'numpad 6',
+    102: 'numpad 6',
-  103: 'numpad 7',
+    103: 'numpad 7',
-  104: 'numpad 8',
+    104: 'numpad 8',
-  105: 'numpad 9',
+    105: 'numpad 9',
-  106: 'multiply',
+    106: 'multiply',
-  107: 'add',
+    107: 'add',
-  109: 'subtract',
+    109: 'subtract',
-  110: 'decimal point',
+    110: 'decimal point',
-  111: 'divide',
+    111: 'divide',
-  112: 'f1',
+    112: 'f1',
-  113: 'f2',
+    113: 'f2',
-  114: 'f3',
+    114: 'f3',
-  115: 'f4',
+    115: 'f4',
-  116: 'f5',
+    116: 'f5',
-  117: 'f6',
+    117: 'f6',
-  118: 'f7',
+    118: 'f7',
-  119: 'f8',
+    119: 'f8',
-  120: 'f9',
+    120: 'f9',
-  121: 'f10',
+    121: 'f10',
-  122: 'f11',
+    122: 'f11',
-  123: 'f12',
+    123: 'f12',
-  124: 'f13',
+    124: 'f13',
-  125: 'f14',
+    125: 'f14',
-  126: 'f15',
+    126: 'f15',
-  127: 'f16',
+    127: 'f16',
-  128: 'f17',
+    128: 'f17',
-  129: 'f18',
+    129: 'f18',
-  130: 'f19',
+    130: 'f19',
-  131: 'f20',
+    131: 'f20',
-  132: 'f21',
+    132: 'f21',
-  133: 'f22',
+    133: 'f22',
-  134: 'f23',
+    134: 'f23',
-  135: 'f24',
+    135: 'f24',
-  144: 'num lock',
+    144: 'num lock',
-  145: 'scroll lock',
+    145: 'scroll lock',
-  186: '&semi;',
+    186: '&semi;',
-  187: '&equals;',
+    187: '&equals;',
-  188: '&comma;',
+    188: '&comma;',
-  189: '&hyphen;',
+    189: '&hyphen;',
-  190: '&period;',
+    190: '&period;',
-  191: '&quest;',
+    191: '&quest;',
-  192: '&grave;',
+    192: '&grave;',
-  219: '&lsqb;',
+    219: '&lsqb;',
-  220: '&bsol;',
+    220: '&bsol;',
-  221: '&rsqb;',
+    221: '&rsqb;',
-  222: '&apos;',
+    222: '&apos;',
-};
+};
--- a/docs/problems/index.html
+++ b/docs/problems/index.html
@@ -32,15 +32,6 @@
            <script src="https://cdn.jsdelivr.net/npm/respond.js@1.4.2/dest/respond.min.js"></script>
        <![endif]-->
    <script src="//ajax.googleapis.com/ajax/libs/webfont/1.6.26/webfont.js"></script>
    <script>
    WebFont.load({
        google: {
            families: ['Open Sans', 'PT Sans']
        }
    });
    </script>
@@ -149,11 +140,16 @@
                <li class="third-level"><a href="#problem-definition">Problem definition</a></li>
                <li class="third-level"><a href="#random-instance-generator">Random instance generator</a></li>
                <li class="third-level"><a href="#challenge-a">Challenge A</a></li>
-            <li class="second-level"><a href="#multidimensional-0-1-knapsack-problem">Multidimensional 0-1 Knapsack Problem</a></li>
+            <li class="second-level"><a href="#traveling-salesman-problem">Traveling Salesman Problem</a></li>
                <li class="third-level"><a href="#problem-definition_1">Problem definition</a></li>
-                <li class="third-level"><a href="#random-instance-generator_1">Random instance generator</a></li>
+                <li class="third-level"><a href="#random-problem-generator">Random problem generator</a></li>
                <li class="third-level"><a href="#challenge-a_1">Challenge A</a></li>
            <li class="second-level"><a href="#multidimensional-0-1-knapsack-problem">Multidimensional 0-1 Knapsack Problem</a></li>
                <li class="third-level"><a href="#problem-definition_2">Problem definition</a></li>
                <li class="third-level"><a href="#random-instance-generator_1">Random instance generator</a></li>
                <li class="third-level"><a href="#challenge-a_2">Challenge A</a></li>
    </ul>
 </div></div>
        <div class="col-md-9" role="main">
@@ -190,8 +186,43 @@
 </code></pre>
 <p><img alt="alt" src="../figures/benchmark_stab_a.png" /></p>
-<h2 id="multidimensional-0-1-knapsack-problem">Multidimensional 0-1 Knapsack Problem</h2>
+<h2 id="traveling-salesman-problem">Traveling Salesman Problem</h2>
 <h3 id="problem-definition_1">Problem definition</h3>
 <p>Given a list of cities and the distance between each pair of cities, the problem asks for the
 shortest route starting at the first city, visiting each other city exactly once, then returning
 to the first city. This problem is a generalization of the Hamiltonian path problem, one of Karp's
 21 NP-complete problems.</p>
 <h3 id="random-problem-generator">Random problem generator</h3>
 <p>The class <code>TravelingSalesmanGenerator</code> can be used to generate random instances of this
 problem. Initially, the generator creates $n$ cities $(x_1,y_1),\ldots,(x_n,y_n) \in \mathbb{R}^2$,
 where $n, x_i$ and $y_i$ are sampled independently from the provided probability distributions <code>n</code>,
 <code>x</code> and <code>y</code>. For each pair of cities $(i,j)$, the distance $d_{i,j}$ between them is set to:
 <script type="math/tex; mode=display">
    d_{i,j} = \gamma_{i,j} \sqrt{(x_i-x_j)^2 + (y_i - y_j)^2}
 </script>
 where $\gamma_{i,j}$ is sampled from the distribution <code>gamma</code>.</p>
 <p>If <code>fix_cities=True</code> is provided, the list of cities is kept the same for all generated instances.
 The $gamma$ values, and therefore also the distances, are still different.</p>
 <p>By default, all distances $d_{i,j}$ are rounded to the nearest integer.  If <code>round=False</code>
 is provided, this rounding will be disabled.</p>
 <h3 id="challenge-a_1">Challenge A</h3>
 <ul>
 <li>Fixed list of 350 cities in the $[0, 1000]^2$ square</li>
 <li>$\gamma_{i,j} \sim U(0.95, 1.05)$</li>
 <li>500 training instances, 50 test instances</li>
 </ul>
 <pre><code class="python">TravelingSalesmanGenerator(x=uniform(loc=0.0, scale=1000.0),
                           y=uniform(loc=0.0, scale=1000.0),
                           n=randint(low=350, high=351),
                           gamma=uniform(loc=0.95, scale=0.1),
                           fix_cities=True,
                           round=True,
                          )
 </code></pre>
 <p><img alt="alt" src="../figures/benchmark_tsp_a.png" /></p>
 <h2 id="multidimensional-0-1-knapsack-problem">Multidimensional 0-1 Knapsack Problem</h2>
 <h3 id="problem-definition_2">Problem definition</h3>
 <p>Given a set of $n$ items and $m$ types of resources (also called <em>knapsacks</em>), the problem is to find a subset of items that maximizes profit without consuming more resources than it is available. More precisely, the problem is:</p>
 <p>
 <script type="math/tex; mode=display">\begin{align*}
@@ -233,7 +264,7 @@ from the provided probability distributions <code>K</code> and <code>u</code>.</
 <li>Fréville, Arnaud. <em>The multidimensional 0–1 knapsack problem: An overview.</em> European Journal of Operational Research 155.1 (2004): 1-21.</li>
 </ul>
 </div>
-<h3 id="challenge-a_1">Challenge A</h3>
+<h3 id="challenge-a_2">Challenge A</h3>
 <ul>
 <li>250 variables, 10 constraints, fixed weights</li>
 <li>$w \sim U(0, 1000), \gamma \sim U(0.95, 1.05)$</li>
@@ -260,9 +291,10 @@ from the provided probability distributions <code>K</code> and <code>u</code>.</
        <hr>
        <p>
-        <small>Copyright © 2020, UChicago Argonne, LLC. All Rights Reserved.<br></small>
+        <small>Copyright © 2020, UChicago Argonne, LLC. All Rights Reserved.</small><br>
-        <small>Documentation built with <a href="http://www.mkdocs.org/">MkDocs</a>.</p></small>
+        <small>Documentation built with <a href="http://www.mkdocs.org/">MkDocs</a>.</small>
        </p>
--- a/docs/search/lunr.js
+++ b/docs/search/lunr.js
--- a/docs/search/search_index.json
+++ b/docs/search/search_index.json
--- a/docs/sitemap.xml
+++ b/docs/sitemap.xml
@@ -1,33 +1,27 @@
 <?xml version="1.0" encoding="UTF-8"?>
-<urlset xmlns="http://www.sitemaps.org/schemas/sitemap/0.9">
+<urlset xmlns="http://www.sitemaps.org/schemas/sitemap/0.9"><url>
    <url>
     <loc>None</loc>
-     <lastmod>2020-02-24</lastmod>
+     <lastmod>2020-02-25</lastmod>
     <changefreq>daily</changefreq>
-    </url>
+    </url><url>
    <url>
     <loc>None</loc>
-     <lastmod>2020-02-24</lastmod>
+     <lastmod>2020-02-25</lastmod>
     <changefreq>daily</changefreq>
-    </url>
+    </url><url>
    <url>
     <loc>None</loc>
-     <lastmod>2020-02-24</lastmod>
+     <lastmod>2020-02-25</lastmod>
     <changefreq>daily</changefreq>
-    </url>
+    </url><url>
    <url>
     <loc>None</loc>
-     <lastmod>2020-02-24</lastmod>
+     <lastmod>2020-02-25</lastmod>
     <changefreq>daily</changefreq>
-    </url>
+    </url><url>
    <url>
     <loc>None</loc>
-     <lastmod>2020-02-24</lastmod>
+     <lastmod>2020-02-25</lastmod>
     <changefreq>daily</changefreq>
-    </url>
+    </url><url>
    <url>
     <loc>None</loc>
-     <lastmod>2020-02-24</lastmod>
+     <lastmod>2020-02-25</lastmod>
     <changefreq>daily</changefreq>
    </url>
 </urlset>
--- a/docs/sitemap.xml.gz
+++ b/docs/sitemap.xml.gz
--- a/docs/usage/index.html
+++ b/docs/usage/index.html
@@ -32,15 +32,6 @@
            <script src="https://cdn.jsdelivr.net/npm/respond.js@1.4.2/dest/respond.min.js"></script>
        <![endif]-->
    <script src="//ajax.googleapis.com/ajax/libs/webfont/1.6.26/webfont.js"></script>
    <script>
    WebFont.load({
        google: {
            families: ['Open Sans', 'PT Sans']
        }
    });
    </script>
@@ -182,12 +173,12 @@ for instance in all_instances:
 <ul>
 <li><code>instance.to_model()</code>, which returns a concrete Pyomo model corresponding to the instance;</li>
 <li><code>instance.get_instance_features()</code>, which returns a 1-dimensional Numpy array of (numerical) features describing the entire instance;</li>
-<li><code>instance.get_variable_features(var, index)</code>, which returns a 1-dimensional array of (numerical) features describing a particular decision variable.</li>
+<li><code>instance.get_variable_features(var_name, index)</code>, which returns a 1-dimensional array of (numerical) features describing a particular decision variable.</li>
 </ul>
 <p>The first method is used by <code>LearningSolver</code> to construct a concrete Pyomo model, which will be provided to the internal MIP solver. The user should keep a reference to this Pyomo model, in order to retrieve, for example, the optimal variable values.</p>
 <p>The second and third methods provide an encoding of the instance, which can be used by the ML models to make predictions. In the knapsack problem, for example, an implementation may decide to provide as instance features the average weights, average prices, number of items and the size of the knapsack. The weight and the price of each individual item could be provided as variable features. See <code>miplearn/problems/knapsack.py</code> for a concrete example.</p>
-<p>An optional method which can be implemented is <code>instance.get_variable_category(var, index)</code>, which returns a category (a string, an integer or any hashable type) for each decision variable. If two variables have the same category, <code>LearningSolver</code> will use the same internal ML model to predict the values of both variables. By default, all variables belong to the <code>"default"</code> category, and therefore only one ML model is used for all variables. If the returned category is <code>None</code>, ML predictors will ignore the variable.</p>
+<p>An optional method which can be implemented is <code>instance.get_variable_category(var_name, index)</code>, which returns a category (a string, an integer or any hashable type) for each decision variable. If two variables have the same category, <code>LearningSolver</code> will use the same internal ML model to predict the values of both variables. By default, all variables belong to the <code>"default"</code> category, and therefore only one ML model is used for all variables. If the returned category is <code>None</code>, ML predictors will ignore the variable.</p>
-<p>It is not necessary to have a one-to-one correspondence between features and problem instances. One important (and deliberate) limitation of MIPLearn, however, is that <code>get_instance_features()</code> must always return arrays of same length for all relevant instances of the problem. Similarly, <code>get_variable_features(var, index)</code> must also always return arrays of same length for all variables in each category. It is up to the user to decide how to encode variable-length characteristics of the problem into fixed-length vectors. In graph problems, for example, graph embeddings can be used to reduce the (variable-length) lists of nodes and edges into a fixed-length structure that still preserves some properties of the graph. Different instance encodings may have significant impact on performance.</p>
+<p>It is not necessary to have a one-to-one correspondence between features and problem instances. One important (and deliberate) limitation of MIPLearn, however, is that <code>get_instance_features()</code> must always return arrays of same length for all relevant instances of the problem. Similarly, <code>get_variable_features(var_name, index)</code> must also always return arrays of same length for all variables in each category. It is up to the user to decide how to encode variable-length characteristics of the problem into fixed-length vectors. In graph problems, for example, graph embeddings can be used to reduce the (variable-length) lists of nodes and edges into a fixed-length structure that still preserves some properties of the graph. Different instance encodings may have significant impact on performance.</p>
 <h3 id="obtaining-heuristic-solutions">Obtaining heuristic solutions</h3>
 <p>By default, <code>LearningSolver</code> uses Machine Learning to accelerate the MIP solution process, while maintaining all optimality guarantees provided by the MIP solver. In the default mode of operation, for example, predicted optimal solutions are used only as MIP starts.</p>
 <p>For more significant performance benefits, <code>LearningSolver</code> can also be configured to place additional trust in the Machine Learning predictors, by using the <code>mode="heuristic"</code> constructor argument. When operating in this mode, if a ML model is statistically shown (through <em>stratified k-fold cross validation</em>) to have exceptionally high accuracy, the solver may decide to restrict the search space based on its predictions. The parts of the solution which the ML models cannot predict accurately will still be explored using traditional (branch-and-bound) methods.  For particular applications, this mode has been shown to quickly produce optimal or near-optimal solutions (see <a href="../about/#references">references</a> and <a href="../benchmark/">benchmark results</a>).</p>
@@ -245,9 +236,10 @@ solver.solve(test_instance)
        <hr>
        <p>
-        <small>Copyright © 2020, UChicago Argonne, LLC. All Rights Reserved.<br></small>
+        <small>Copyright © 2020, UChicago Argonne, LLC. All Rights Reserved.</small><br>
-        <small>Documentation built with <a href="http://www.mkdocs.org/">MkDocs</a>.</p></small>
+        <small>Documentation built with <a href="http://www.mkdocs.org/">MkDocs</a>.</small>
        </p>
--- a/miplearn/init.py
+++ b/miplearn/init.py
@@ -3,13 +3,13 @@
 #  Released under the modified BSD license. See COPYING.md for more details.
 from .extractors import (SolutionExtractor,
                         CombinedExtractor,
                         InstanceFeaturesExtractor,
                         ObjectiveValueExtractor,
                         VariableFeaturesExtractor,
                        )
 from .components.component import Component
 from .components.objective import ObjectiveValueComponent
 from .components.lazy import LazyConstraintsComponent
 from .components.primal import (PrimalSolutionComponent,
                                AdaptivePredictor,
                               )
--- a/miplearn/components/component.py
+++ b/miplearn/components/component.py
@@ -18,10 +18,6 @@ class Component(ABC):
    def after_solve(self, solver, instance, model):
        pass
    @abstractmethod
    def merge(self, other):
        pass
    @abstractmethod
    def fit(self, training_instances):
        pass
--- a/miplearn/components/lazy.py
+++ b/miplearn/components/lazy.py
@@ -0,0 +1,57 @@
 #  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
 #  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 from .component import Component
 from ..extractors import *
 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
 from sklearn.preprocessing import StandardScaler
 from sklearn.model_selection import cross_val_score
 from sklearn.metrics import roc_curve
 from sklearn.neighbors import KNeighborsClassifier
 from tqdm.auto import tqdm
 import pyomo.environ as pe
 import logging
 logger = logging.getLogger(__name__)
 class LazyConstraintsComponent(Component):
    """
    A component that predicts which lazy constraints to enforce.
    """
    def __init__(self):
        self.violations = set()
        self.count = {}
        self.n_samples = 0
    def before_solve(self, solver, instance, model):
        logger.info("Enforcing %d lazy constraints" % len(self.violations))
        for v in self.violations:
            if self.count[v] < self.n_samples * 0.05:
                continue
            cut = instance.build_lazy_constraint(model, v)
            solver.internal_solver.add_constraint(cut)
    def after_solve(self, solver, instance, model):
        pass
    def fit(self, training_instances):
        logger.debug("Fitting...")
        self.n_samples = len(training_instances)
        for instance in training_instances:
            if not hasattr(instance, "found_violations"):
                continue
            for v in instance.found_violations:
                self.violations.add(v)
                if v not in self.count.keys():
                    self.count[v] = 0
                self.count[v] += 1
    def predict(self, instance, model=None):
        return self.violations
--- a/miplearn/components/objective.py
+++ b/miplearn/components/objective.py
@@ -30,16 +30,16 @@ class ObjectiveValueComponent(Component):
    def after_solve(self, solver, instance, model):
        pass
    def merge(self, other):
        pass
    def fit(self, training_instances):
        logger.debug("Extracting features...")
        features = InstanceFeaturesExtractor().extract(training_instances)
        ub = ObjectiveValueExtractor(kind="upper bound").extract(training_instances)
        lb = ObjectiveValueExtractor(kind="lower bound").extract(training_instances)
        self.ub_regressor = deepcopy(self.regressor_prototype)
        self.lb_regressor = deepcopy(self.regressor_prototype)
        logger.debug("Fitting ub_regressor...")
        self.ub_regressor.fit(features, ub)
        logger.debug("Fitting ub_regressor...")
        self.lb_regressor.fit(features, lb)
    def predict(self, instances):
--- a/miplearn/components/primal.py
+++ b/miplearn/components/primal.py
@@ -129,7 +129,7 @@ class PrimalSolutionComponent(Component):
        self.dynamic_thresholds = dynamic_thresholds
    def before_solve(self, solver, instance, model):
-        solution = self.predict(instance, model)
+        solution = self.predict(instance)
        if self.mode == "heuristic":
            solver.internal_solver.fix(solution)
        else:
@@ -139,6 +139,7 @@ class PrimalSolutionComponent(Component):
        pass
    def fit(self, training_instances):
        logger.debug("Extracting features...")
        features = VariableFeaturesExtractor().extract(training_instances)
        solutions = SolutionExtractor().extract(training_instances)
@@ -180,12 +181,10 @@ class PrimalSolutionComponent(Component):
                self.thresholds[category, label] = thresholds[k]
-    def predict(self, instance, model=None):
+    def predict(self, instance):
-        if model is None:
+        x_test = VariableFeaturesExtractor().extract([instance])
            model = instance.to_model()
        x_test = VariableFeaturesExtractor().extract([instance], [model])
        solution = {}
-        var_split = Extractor.split_variables(instance, model)
+        var_split = Extractor.split_variables(instance)
        for category in var_split.keys():
            for (i, (var, index)) in enumerate(var_split[category]):
                if var not in solution.keys():
@@ -200,6 +199,3 @@ class PrimalSolutionComponent(Component):
                    if ws[i, 1] >= self.thresholds[category, label]:
                        solution[var][index] = label
        return solution
    def merge(self, other_components):
        pass
--- a/miplearn/components/tests/test_primal.py
+++ b/miplearn/components/tests/test_primal.py
@@ -27,29 +27,7 @@ def test_predict():
    instances, models = _get_instances()
    comp = PrimalSolutionComponent()
    comp.fit(instances)
-    solution = comp.predict(instances[0], models[0])
+    solution = comp.predict(instances[0])
-    assert models[0].x in solution.keys()
+    assert "x" in solution
    for idx in range(4):
-        assert idx in solution[models[0].x].keys()
+        assert idx in solution["x"]
 # def test_warm_start_save_load():
 #     state_file = tempfile.NamedTemporaryFile(mode="r")
 #     solver = LearningSolver(components={"warm-start": WarmStartComponent()})
 #     solver.parallel_solve(_get_instances(), n_jobs=2)
 #     solver.fit()
 #     comp = solver.components["warm-start"]
 #     assert comp.x_train["default"].shape == (8, 6)
 #     assert comp.y_train["default"].shape == (8, 2)
 #     assert ("default", 0) in comp.predictors.keys()
 #     assert ("default", 1) in comp.predictors.keys()
 #     solver.save_state(state_file.name)
 #     solver.solve(_get_instances()[0])
 #     solver = LearningSolver(components={"warm-start": WarmStartComponent()})
 #     solver.load_state(state_file.name)
 #     comp = solver.components["warm-start"]
 #     assert comp.x_train["default"].shape == (8, 6)
 #     assert comp.y_train["default"].shape == (8, 2)
 #     assert ("default", 0) in comp.predictors.keys()
 #     assert ("default", 1) in comp.predictors.keys()
--- a/miplearn/extractors.py
+++ b/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 var_name in instance.lp_solution:
-            for index in var:
+            for index in instance.lp_solution[var_name]:
-                category = instance.get_variable_category(var, index)
+                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,
+    def extract(self, instances):
                instances,
                models=None,
               ):
        result = {}
-        if models is None:
+        for instance in tqdm(instances,
-            models = [instance.to_model() for instance in instances]
+                             desc="Extract var features",
-        for (index, instance) in enumerate(instances):
+                             disable=len(instances) < 5):
            model = models[index]
            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:
+                for (var_name, index) in var_index_pairs:
-                    result[category] += [np.hstack([
+                    result[category] += [
-                        instance_features,
+                        instance_features.tolist() + \
-                        instance.get_variable_features(var, index),
+                        instance.get_variable_features(var_name, index).tolist() + \
-                        instance.lp_solution[str(var)][index],
+                        [instance.lp_solution[var_name][index]]
-                    ])]
+                    ]
-        for category in result.keys():
+        for category in result:
-            result[category] = np.vstack(result[category])
+            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:
+        for instance in tqdm(instances,
-            models = [instance.to_model() for instance in instances]
+                             desc="Extract solution",
-        for (index, instance) in enumerate(instances):
+                             disable=len(instances) < 5):
-            model = models[index]
+            var_split = self.split_variables(instance)
            var_split = self.split_variables(instance, model)
            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():
+        for category in result:
-            result[category] = np.vstack(result[category])            
+            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([
--- a/miplearn/instance.py
+++ b/miplearn/instance.py
@@ -65,12 +65,50 @@ class Instance(ABC):
    def get_variable_category(self, var, index):
        """
-        Returns a category (a string, an integer or any hashable type) for each decision variable.
+        Returns the category (a string, an integer or any hashable type) for each decision
        variable.
-        If two variables have the same category, LearningSolver will use the same internal ML model
+        If two variables have the same category, LearningSolver will use the same internal ML
-        to predict the values of both variables. By default, all variables belong to the "default"
+        model to predict the values of both variables. By default, all variables belong to the
-        category, and therefore only one ML model is used for all variables.
+        "default" category, and therefore only one ML model is used for all variables.
        If the returned category is None, ML models will ignore the variable.
        """
        return "default"
    def find_violations(self, model):
        """
        Returns lazy constraint violations found for the current solution.
        After solving a model, LearningSolver will ask the instance to identify which lazy
        constraints are violated by the current solution. For each identified violation,
        LearningSolver will then call the build_lazy_constraint, add the generated Pyomo
        constraint to the model, then resolve the problem. The process repeats until no further
        lazy constraint violations are found.
        Each "violation" is simply a string, a tuple or any other hashable type which allows the
        instance to identify unambiguously which lazy constraint should be generated. In the
        Traveling Salesman Problem, for example, a subtour violation could be a frozen set
        containing the cities in the subtour.
        For a concrete example, see TravelingSalesmanInstance.
        """
        return []
    def build_lazy_constraint(self, model, violation):
        """
        Returns a Pyomo constraint which fixes a given violation.
        This method is typically called immediately after find_violations. The violation object
        provided to this method is exactly the same object returned earlier by find_violations.
        After some training, LearningSolver may decide to proactively build some lazy constraints
        at the beginning of the optimization process, before a solution is even available. In this
        case, build_lazy_constraints will be called without a corresponding call to
        find_violations.
        The implementation should not directly add the constraint to the model. The constraint
        will be added by LearningSolver after the method returns.
        For a concrete example, see TravelingSalesmanInstance.
        """
        pass
--- a/miplearn/problems/tests/test_tsp.py
+++ b/miplearn/problems/tests/test_tsp.py
@@ -0,0 +1,68 @@
 #  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
 #  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 from miplearn import LearningSolver
 from miplearn.problems.tsp import TravelingSalesmanGenerator, TravelingSalesmanInstance
 import numpy as np
 from numpy.linalg import norm
 from scipy.spatial.distance import pdist, squareform
 from scipy.stats import uniform, randint
 def test_generator():
    instances = TravelingSalesmanGenerator(x=uniform(loc=0.0, scale=1000.0),
                                           y=uniform(loc=0.0, scale=1000.0),
                                           n=randint(low=100, high=101),
                                           gamma=uniform(loc=0.95, scale=0.1),
                                           fix_cities=True).generate(100)
    assert len(instances) == 100
    assert instances[0].n_cities == 100
    assert norm(instances[0].distances - instances[0].distances.T) < 1e-6
    d = [instance.distances[0,1] for instance in instances]
    assert np.std(d) > 0
 def test_instance():
    n_cities = 4
    distances = np.array([
        [0., 1., 2., 1.],
        [1., 0., 1., 2.],
        [2., 1., 0., 1.],
        [1., 2., 1., 0.],
    ])
    instance = TravelingSalesmanInstance(n_cities, distances)
    solver = LearningSolver()
    solver.solve(instance)
    x = instance.solution["x"]
    assert x[0,1] == 1.0
    assert x[0,2] == 0.0
    assert x[0,3] == 1.0
    assert x[1,2] == 1.0
    assert x[1,3] == 0.0
    assert x[2,3] == 1.0
    assert instance.lower_bound == 4.0
    assert instance.upper_bound == 4.0
 def test_subtour():
    n_cities = 6
    cities = np.array([
        [0., 0.],
        [1., 0.],
        [2., 0.],
        [3., 0.],
        [0., 1.],
        [3., 1.],
    ])
    distances = squareform(pdist(cities))
    instance = TravelingSalesmanInstance(n_cities, distances)
    solver = LearningSolver()
    solver.solve(instance)
    x = instance.solution["x"]
    assert x[0,1] == 1.0
    assert x[0,4] == 1.0
    assert x[1,2] == 1.0
    assert x[2,3] == 1.0
    assert x[3,5] == 1.0
    assert x[4,5] == 1.0
--- a/miplearn/problems/tsp.py
+++ b/miplearn/problems/tsp.py
@@ -0,0 +1,169 @@
 # MIPLearn, an extensible framework for Learning-Enhanced Mixed-Integer Optimization
 # Copyright (C) 2019-2020 Argonne National Laboratory. All rights reserved.
 # Written by Alinson S. Xavier <axavier@anl.gov>
 import numpy as np
 import pyomo.environ as pe
 from miplearn import Instance
 from scipy.stats import uniform, randint
 from scipy.spatial.distance import pdist, squareform
 from scipy.stats.distributions import rv_frozen
 import networkx as nx
 import random
 class ChallengeA:
    def __init__(self,
                 seed=42,
                 n_training_instances=500,
                 n_test_instances=50,
                ):
        np.random.seed(seed)
        self.generator = TravelingSalesmanGenerator(x=uniform(loc=0.0, scale=1000.0),
                                                    y=uniform(loc=0.0, scale=1000.0),
                                                    n=randint(low=350, high=351),
                                                    gamma=uniform(loc=0.95, scale=0.1),
                                                    fix_cities=True,
                                                    round=True,
                                                   )
        np.random.seed(seed + 1)
        self.training_instances = self.generator.generate(n_training_instances)
        np.random.seed(seed + 2)
        self.test_instances = self.generator.generate(n_test_instances)        
 class TravelingSalesmanGenerator:
    """Random generator for the Traveling Salesman Problem."""
    def __init__(self,
                 x=uniform(loc=0.0, scale=1000.0),
                 y=uniform(loc=0.0, scale=1000.0),
                 n=randint(low=100, high=101),
                 gamma=uniform(loc=1.0, scale=0.0),
                 fix_cities=True,
                 round=True,
                ):
        """Initializes the problem generator.
        Initially, the generator creates n cities (x_1,y_1),...,(x_n,y_n) where n, x_i and y_i are
        sampled independently from the provided probability distributions `n`, `x` and `y`. For each
        (unordered) pair of cities (i,j), the distance d[i,j] between them is set to:
            d[i,j] = gamma[i,j] \sqrt{(x_i - x_j)^2 + (y_i - y_j)^2}
        where gamma is sampled from the provided probability distribution `gamma`.
        If fix_cities=True, the list of cities is kept the same for all generated instances. The
        gamma values, and therefore also the distances, are still different.
        By default, all distances d[i,j] are rounded to the nearest integer.  If `round=False`
        is provided, this rounding will be disabled.
        Arguments
        ---------
        x: rv_continuous
            Probability distribution for the x-coordinate of each city.
        y: rv_continuous
            Probability distribution for the y-coordinate of each city.
        n: rv_discrete
            Probability distribution for the number of cities.
        fix_cities: bool
            If False, cities will be resampled for every generated instance. Otherwise, list of
            cities will be computed once, during the constructor.
        round: bool
            If True, distances are rounded to the nearest integer.
        """
        assert isinstance(x, rv_frozen), "x should be a SciPy probability distribution"
        assert isinstance(y, rv_frozen), "y should be a SciPy probability distribution"
        assert isinstance(n, rv_frozen), "n should be a SciPy probability distribution"
        assert isinstance(gamma, rv_frozen), "gamma should be a SciPy probability distribution"
        self.x = x
        self.y = y
        self.n = n
        self.gamma = gamma
        self.round = round
        if fix_cities:
            self.fixed_n, self.fixed_cities = self._generate_cities()
        else:
            self.fixed_n = None
            self.fixed_cities = None
    def generate(self, n_samples):
        def _sample():
            if self.fixed_cities is not None:
                n, cities = self.fixed_n, self.fixed_cities
            else:
                n, cities = self._generate_cities()
            distances = squareform(pdist(cities)) * self.gamma.rvs(size=(n, n))
            distances = np.tril(distances) + np.triu(distances.T, 1)
            if self.round:
                distances = distances.round()
            return TravelingSalesmanInstance(n, distances)
        return [_sample() for _ in range(n_samples)]
    def _generate_cities(self):
        n = self.n.rvs()
        cities = np.array([(self.x.rvs(), self.y.rvs()) for _ in range(n)])
        return n, cities
 class TravelingSalesmanInstance(Instance):
    """An instance ot the Traveling Salesman Problem.
    Given a list of cities and the distance between each pair of cities, the problem asks for the
    shortest route starting at the first city, visiting each other city exactly once, then
    returning to the first city. This problem is a generalization of the Hamiltonian path problem,
    one of Karp's 21 NP-complete problems.
    """
    def __init__(self, n_cities, distances):
        assert isinstance(distances, np.ndarray)
        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 = [(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)
        model.obj = pe.Objective(expr=sum(model.x[i,j] * self.distances[i,j]
                                          for (i,j) in edges),
                                 sense=pe.minimize)
        model.eq_degree = pe.ConstraintList()
        model.eq_subtour = pe.ConstraintList()
        for i in range(self.n_cities):
            model.eq_degree.add(sum(model.x[min(i,j), max(i,j)]
                                    for j in range(self.n_cities) if i != j) == 2)
        return model
    def get_instance_features(self):
        return np.array([1])
    def get_variable_features(self, var_name, index):
        return np.array([1])
    def get_variable_category(self, var_name, index):
        return index
    def find_violations(self, model):
        selected_edges = [e for e in model.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))]
        violations = []
        for c in components:
            if len(c) < self.n_cities:
                violations += [c]
        return violations
    def build_lazy_constraint(self, model, component):
        cut_edges = [e for e in model.edges
                     if (e[0] in component and e[1] not in component) or
                        (e[0] not in component and e[1] in component)]
        return model.eq_subtour.add(sum(model.x[e] for e in cut_edges) >= 2)
--- a/miplearn/solvers.py
+++ b/miplearn/solvers.py
@@ -2,45 +2,66 @@
 #  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
-from . import ObjectiveValueComponent, PrimalSolutionComponent
+from . import ObjectiveValueComponent, PrimalSolutionComponent, LazyConstraintsComponent
 import pyomo.environ as pe
 from pyomo.core import Var
 from copy import deepcopy
 import pickle
 from scipy.stats import randint
 from p_tqdm import p_map
 import numpy as np
 import logging
 logger = logging.getLogger(__name__)
 # Global memory for multiprocessing
 SOLVER = [None]
 INSTANCES = [None]
 def _parallel_solve(instance_idx):
    solver = deepcopy(SOLVER[0])
    instance = INSTANCES[0][instance_idx]
    results = solver.solve(instance)
    return {
        "Results": results,
        "Solution": instance.solution,
        "LP solution": instance.lp_solution,
        "LP value": instance.lp_value,
        "Upper bound": instance.upper_bound,
        "Lower bound": instance.lower_bound,
        "Violations": instance.found_violations,
    }
 class InternalSolver:
    def __init__(self):
        self.is_warm_start_available = False
        self.model = None
-        pass
+        self.var_name_to_var = {}
    def solve_lp(self, tee=False):
        self.solver.set_instance(self.model)
        # Relax domain
        from pyomo.core.base.set_types import Reals
-        original_domain = {}
+        original_domains = []
-        for var in self.model.component_data_objects(Var):
+        for (idx, var) in enumerate(self.model.component_data_objects(Var)):
-            original_domain[str(var)] = var.domain
+            original_domains += [var.domain]
            lb, ub = var.bounds
            var.setlb(lb)
            var.setub(ub)
            var.domain = Reals
            self.solver.update_var(var)
        # Solve LP relaxation
        self.solver.set_instance(self.model)
        results = self.solver.solve(tee=tee)
        # Restore domains
-        for var in self.model.component_data_objects(Var):
+        for (idx, var) in enumerate(self.model.component_data_objects(Var)):
-            var.domain = original_domain[str(var)]
+            var.domain = original_domains[idx]
            self.solver.update_var(var)
        # Reload original model
        self.solver.set_instance(self.model)
        return {
            "Optimal value": results["Problem"][0]["Lower bound"],
        }
@@ -58,36 +79,43 @@ class InternalSolver:
                solution[str(var)][index] = var[index].value
        return solution   
-    def set_warm_start(self, ws):
+    def set_warm_start(self, solution):
        self.is_warm_start_available = True
        self.clear_values()
        count_total, count_fixed = 0, 0
-        for var in ws.keys():
+        for var_name in solution:
-            for index in var:
+            var = self.var_name_to_var[var_name]
            for index in solution[var_name]:
                count_total += 1
-                var[index].value = ws[var][index]
+                var[index].value = solution[var_name][index]
-                if ws[var][index] is not None:
+                if solution[var_name][index] is not None:
                    count_fixed += 1
        logger.info("Setting start values for %d variables (out of %d)" %
                    (count_fixed, count_total))
    def set_model(self, model):
        self.model = model
        self.solver.set_instance(model)
        self.var_name_to_var = {}
        for var in model.component_objects(Var):
            self.var_name_to_var[var.name] = var
-    def fix(self, ws):
+    def fix(self, solution):
        count_total, count_fixed = 0, 0
-        for var in ws.keys():
+        for var_name in solution:
-            for index in var:
+            for index in solution[var_name]:
                var = self.var_name_to_var[var_name]
                count_total += 1
-                if ws[var][index] is None:
+                if solution[var_name][index] is None:
                    continue
                count_fixed += 1
-                var[index].fix(ws[var][index])
+                var[index].fix(solution[var_name][index])
                self.solver.update_var(var[index])        
        logger.info("Fixing values for %d variables (out of %d)" %
                    (count_fixed, count_total))
    def add_constraint(self, cut):
        self.solver.add_constraint(cut)
 class GurobiSolver(InternalSolver):
@@ -198,6 +226,7 @@ class LearningSolver:
            self.components = {
                "ObjectiveValue": ObjectiveValueComponent(),
                "PrimalSolution": PrimalSolutionComponent(),
                "LazyConstraints": LazyConstraintsComponent(),
            }
        assert self.mode in ["exact", "heuristic"]
@@ -231,27 +260,44 @@ class LearningSolver:
        self.internal_solver = self._create_internal_solver()
        self.internal_solver.set_model(model)
-        # Solve LP relaxation
+        logger.debug("Solving LP relaxation...")
        results = self.internal_solver.solve_lp(tee=tee)
        instance.lp_solution = self.internal_solver.get_solution()
        instance.lp_value = results["Optimal value"]
-        # Invoke before_solve callbacks
+        logger.debug("Running before_solve callbacks...")
        for component in self.components.values():
            component.before_solve(self, instance, model)
        if relaxation_only:
            return results
-        # Solver original MIP
+        total_wallclock_time = 0
-        results = self.internal_solver.solve(tee=tee)
+        instance.found_violations = []
        while True:
            logger.debug("Solving MIP...")
            results = self.internal_solver.solve(tee=tee)
            logger.debug("    %.2f s" % results["Wallclock time"])
            total_wallclock_time += results["Wallclock time"]
            if not hasattr(instance, "find_violations"):
                break
            logger.debug("Finding violated constraints...")    
            violations = instance.find_violations(model)
            if len(violations) == 0:
                break
            instance.found_violations += violations
            logger.debug("    %d violations found" % len(violations))
            for v in violations:
                cut = instance.build_lazy_constraint(model, v)
                self.internal_solver.add_constraint(cut)
        results["Wallclock time"] = total_wallclock_time
        # Read MIP solution and bounds
        instance.lower_bound = results["Lower bound"]
        instance.upper_bound = results["Upper bound"]
        instance.solution = self.internal_solver.get_solution()
-        # Invoke after_solve callbacks
+        logger.debug("Calling after_solve callbacks...")    
        for component in self.components.values():
            component.after_solve(self, instance, model)
@@ -266,40 +312,23 @@ class LearningSolver:
                       label="Solve",
                       collect_training_data=True,
                      ):
        self.internal_solver = None
-        
+        SOLVER[0] = self
-        def _process(instance):
+        INSTANCES[0] = instances
-            solver = deepcopy(self)
+        p_map_results = p_map(_parallel_solve,
-            results = solver.solve(instance)
+                              list(range(len(instances))),
-            solver.internal_solver = None
+                              num_cpus=n_jobs,
-            if not collect_training_data:
+                              desc=label)
                solver.components = {}
            return {
                "Solver": solver,
                "Results": results,
                "Solution": instance.solution,
                "LP solution": instance.lp_solution,
                "LP value": instance.lp_value,
                "Upper bound": instance.upper_bound,
                "Lower bound": instance.lower_bound,
            }
        p_map_results = p_map(_process, instances, num_cpus=n_jobs, desc=label)
        subsolvers = [p["Solver"] for p in p_map_results]
        results = [p["Results"] for p in p_map_results]
        for (idx, r) in enumerate(p_map_results):
            instances[idx].solution = r["Solution"]
            instances[idx].lp_solution = r["LP solution"]
            instances[idx].lp_value = r["LP value"]
            instances[idx].lower_bound = r["Lower bound"]
            instances[idx].upper_bound = r["Upper bound"]
-        
+            instances[idx].found_violations = r["Violations"]
        for (name, component) in self.components.items():
            subcomponents = [subsolver.components[name]
                             for subsolver in subsolvers
                             if name in subsolver.components.keys()]
            self.components[name].merge(subcomponents)
        return results
@@ -310,21 +339,3 @@ class LearningSolver:
            return
        for component in self.components.values():
            component.fit(training_instances)
    def save_state(self, filename):
        with open(filename, "wb") as file:
            pickle.dump({
                "version": 2,
                "components": self.components,
            }, file)
    def load_state(self, filename):
        with open(filename, "rb") as file:
            data = pickle.load(file)
            assert data["version"] == 2
            for (component_name, component) in data["components"].items():
                if component_name not in self.components.keys():
                    continue
                else:
                    self.components[component_name].merge([component])
--- a/miplearn/tests/test_benchmark.py
+++ b/miplearn/tests/test_benchmark.py
@@ -18,8 +18,6 @@ def test_benchmark():
    # Training phase...
    training_solver = LearningSolver()
    training_solver.parallel_solve(train_instances, n_jobs=10)
    training_solver.fit()
    training_solver.save_state("data.bin")
    # Test phase...
    test_solvers = {
@@ -27,7 +25,7 @@ def test_benchmark():
        "Strategy B": LearningSolver(),
    }
    benchmark = BenchmarkRunner(test_solvers)
-    benchmark.load_state("data.bin")
+    benchmark.fit(train_instances)
    benchmark.parallel_solve(test_instances, n_jobs=2, n_trials=2)
    assert benchmark.raw_results().values.shape == (12,13)
--- a/miplearn/tests/test_extractors.py
+++ b/miplearn/tests/test_extractors.py
@@ -5,7 +5,6 @@
 from miplearn.problems.knapsack import KnapsackInstance
 from miplearn import (LearningSolver,
                      SolutionExtractor,
                      CombinedExtractor,
                      InstanceFeaturesExtractor,
                      VariableFeaturesExtractor,
                     )
@@ -33,7 +32,7 @@ def _get_instances():
 def test_solution_extractor():
    instances, models = _get_instances()
-    features = SolutionExtractor().extract(instances, models)
+    features = SolutionExtractor().extract(instances)
    assert isinstance(features, dict)
    assert "default" in features.keys()
    assert isinstance(features["default"], np.ndarray)
@@ -48,17 +47,6 @@ def test_solution_extractor():
    ]
 def test_combined_extractor():
    instances, models = _get_instances()
    extractor = CombinedExtractor(extractors=[VariableFeaturesExtractor(),
                                              SolutionExtractor()])
    features = extractor.extract(instances, models)
    assert isinstance(features, dict)
    assert "default" in features.keys()
    assert isinstance(features["default"], np.ndarray)
    assert features["default"].shape == (6, 7)
 def test_instance_features_extractor():
    instances, models = _get_instances()
    features = InstanceFeaturesExtractor().extract(instances)
--- a/miplearn/tests/test_solver.py
+++ b/miplearn/tests/test_solver.py
@@ -41,29 +41,6 @@ def test_solver():
            solver.fit()
            solver.solve(instance)
 # def test_solve_save_load_state():
 #     instance = _get_instance()
 #     components_before = {
 #         "warm-start": WarmStartComponent(),
 #     }
 #     solver = LearningSolver(components=components_before)
 #     solver.solve(instance)
 #     solver.fit()
 #     solver.save_state("/tmp/knapsack_train.bin")
 #     prev_x_train_len = len(solver.components["warm-start"].x_train)
 #     prev_y_train_len = len(solver.components["warm-start"].y_train)
 #     components_after = {
 #         "warm-start": WarmStartComponent(),
 #     }
 #     solver = LearningSolver(components=components_after)
 #     solver.load_state("/tmp/knapsack_train.bin")
 #     assert len(solver.components.keys()) == 1
 #     assert len(solver.components["warm-start"].x_train) == prev_x_train_len
 #     assert len(solver.components["warm-start"].y_train) == prev_y_train_len
 def test_parallel_solve():
    instances = [_get_instance() for _ in range(10)]
    solver = LearningSolver()
		`@@ -0,0 +1 @@`
							`../../benchmark/tsp/ChallengeA/performance.png`