Reorganize imports; start moving data to instance.training_data

.mypy.ini

@@ -3,3 +3,4 @@ ignore_missing_imports = True
 #disallow_untyped_defs = True
 disallow_untyped_calls = True
 disallow_incomplete_defs = True
+pretty = True

miplearn/__init__.py

@@ -2,37 +2,31 @@
 #  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 
+from .benchmark import BenchmarkRunner
+from .classifiers import Classifier, Regressor
+from .classifiers.adaptive import AdaptiveClassifier
+from .classifiers.threshold import MinPrecisionThreshold
+from .components.component import Component
+from .components.cuts import UserCutsComponent
+from .components.lazy_dynamic import DynamicLazyConstraintsComponent
+from .components.lazy_static import StaticLazyConstraintsComponent
+from .components.objective import ObjectiveValueComponent
+from .components.primal import PrimalSolutionComponent
+from .components.relaxation import RelaxationComponent
+from .components.steps.convert_tight import ConvertTightIneqsIntoEqsStep
+from .components.steps.drop_redundant import DropRedundantInequalitiesStep
+from .components.steps.relax_integrality import RelaxIntegralityStep
 from .extractors import (
     SolutionExtractor,
     InstanceFeaturesExtractor,
     ObjectiveValueExtractor,
     VariableFeaturesExtractor,
 )
-
-from .components.component import Component
-from .components.objective import ObjectiveValueComponent
-from .components.lazy_dynamic import DynamicLazyConstraintsComponent
-from .components.lazy_static import StaticLazyConstraintsComponent
-from .components.cuts import UserCutsComponent
-from .components.primal import PrimalSolutionComponent
-from .components.relaxation import RelaxationComponent
-from .components.steps.convert_tight import ConvertTightIneqsIntoEqsStep
-from .components.steps.relax_integrality import RelaxIntegralityStep
-from .components.steps.drop_redundant import DropRedundantInequalitiesStep
-
-from .classifiers import Classifier, Regressor
-from .classifiers.adaptive import AdaptiveClassifier
-from .classifiers.threshold import MinPrecisionThreshold
-
-from .benchmark import BenchmarkRunner
-
 from .instance import Instance
-
-from .solvers.pyomo.base import BasePyomoSolver
-from .solvers.pyomo.cplex import CplexPyomoSolver
-from .solvers.pyomo.gurobi import GurobiPyomoSolver
+from .log import setup_logger
 from .solvers.gurobi import GurobiSolver
 from .solvers.internal import InternalSolver
 from .solvers.learning import LearningSolver
-
-from .log import setup_logger
+from .solvers.pyomo.base import BasePyomoSolver
+from .solvers.pyomo.cplex import CplexPyomoSolver
+from .solvers.pyomo.gurobi import GurobiPyomoSolver

miplearn/benchmark.py

@@ -2,15 +2,14 @@
 #  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 
+import logging
+import os
 from copy import deepcopy
 
 import pandas as pd
-import numpy as np
-import logging
 from tqdm.auto import tqdm
-import os
 
-from .solvers.learning import LearningSolver
+from miplearn.solvers.learning import LearningSolver
 
 
 class BenchmarkRunner:

miplearn/classifiers/adaptive.py

@@ -5,14 +5,15 @@
 import logging
 from copy import deepcopy
 
-from miplearn.classifiers import Classifier
-from miplearn.classifiers.counting import CountingClassifier
-from miplearn.classifiers.evaluator import ClassifierEvaluator
 from sklearn.linear_model import LogisticRegression
 from sklearn.neighbors import KNeighborsClassifier
 from sklearn.pipeline import make_pipeline
 from sklearn.preprocessing import StandardScaler
 
+from miplearn.classifiers import Classifier
+from miplearn.classifiers.counting import CountingClassifier
+from miplearn.classifiers.evaluator import ClassifierEvaluator
+
 logger = logging.getLogger(__name__)
 
 

miplearn/classifiers/counting.py

@@ -2,9 +2,10 @@
 #  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 
-from miplearn.classifiers import Classifier
 import numpy as np
 
+from miplearn.classifiers import Classifier
+
 
 class CountingClassifier(Classifier):
     """

miplearn/classifiers/cv.py

@@ -2,15 +2,15 @@
 #  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 
+import logging
 from copy import deepcopy
 
 import numpy as np
-from miplearn.classifiers import Classifier
 from sklearn.dummy import DummyClassifier
 from sklearn.linear_model import LogisticRegression
 from sklearn.model_selection import cross_val_score
 
-import logging
+from miplearn.classifiers import Classifier
 
 logger = logging.getLogger(__name__)
 

miplearn/classifiers/tests/test_counting.py

@@ -1,11 +1,12 @@
 #  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.classifiers.counting import CountingClassifier
 
 import numpy as np
 from numpy.linalg import norm
 
+from miplearn.classifiers.counting import CountingClassifier
+
 E = 0.1
 
 

miplearn/classifiers/tests/test_cv.py

@@ -3,11 +3,12 @@
 #  Released under the modified BSD license. See COPYING.md for more details.
 
 import numpy as np
-from miplearn.classifiers.cv import CrossValidatedClassifier
 from numpy.linalg import norm
 from sklearn.preprocessing import StandardScaler
 from sklearn.svm import SVC
 
+from miplearn.classifiers.cv import CrossValidatedClassifier
+
 E = 0.1
 
 

miplearn/classifiers/tests/test_evaluator.py

@@ -3,9 +3,10 @@
 #  Released under the modified BSD license. See COPYING.md for more details.
 
 import numpy as np
-from miplearn.classifiers.evaluator import ClassifierEvaluator
 from sklearn.neighbors import KNeighborsClassifier
 
+from miplearn.classifiers.evaluator import ClassifierEvaluator
+
 
 def test_evaluator():
     clf_a = KNeighborsClassifier(n_neighbors=1)

miplearn/classifiers/tests/test_threshold.py

@@ -5,6 +5,7 @@
 from unittest.mock import Mock
 
 import numpy as np
+
 from miplearn.classifiers import Classifier
 from miplearn.classifiers.threshold import MinPrecisionThreshold
 

miplearn/components/composite.py

@@ -2,7 +2,7 @@
 #  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 
-from miplearn import Component
+from miplearn.components.component import Component
 
 
 class CompositeComponent(Component):

miplearn/components/cuts.py

@@ -2,14 +2,17 @@
 #  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 
+import logging
 import sys
 from copy import deepcopy
 
+import numpy as np
+from tqdm.auto import tqdm
+
 from miplearn.classifiers.counting import CountingClassifier
 from miplearn.components import classifier_evaluation_dict
-
-from .component import Component
-from ..extractors import *
+from miplearn.components.component import Component
+from miplearn.extractors import InstanceFeaturesExtractor
 
 logger = logging.getLogger(__name__)
 

miplearn/components/lazy_dynamic.py

@@ -2,14 +2,17 @@
 #  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 
+import logging
 import sys
 from copy import deepcopy
 
+import numpy as np
+from tqdm.auto import tqdm
+
 from miplearn.classifiers.counting import CountingClassifier
 from miplearn.components import classifier_evaluation_dict
-
-from .component import Component
-from ..extractors import *
+from miplearn.components.component import Component
+from miplearn.extractors import InstanceFeaturesExtractor, InstanceIterator
 
 logger = logging.getLogger(__name__)
 

miplearn/components/lazy_static.py

@@ -2,12 +2,15 @@
 #  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 
+import logging
 import sys
 from copy import deepcopy
 
+import numpy as np
+from tqdm.auto import tqdm
+
 from miplearn.classifiers.counting import CountingClassifier
-from .component import Component
-from ..extractors import *
+from miplearn.components.component import Component
 
 logger = logging.getLogger(__name__)
 

miplearn/components/objective.py

@@ -1,6 +1,12 @@
 #  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.
+
+import logging
+from copy import deepcopy
+
+import numpy as np
+from sklearn.linear_model import LinearRegression
 from sklearn.metrics import (
     mean_squared_error,
     explained_variance_score,
@@ -9,11 +15,8 @@ from sklearn.metrics import (
     r2_score,
 )
 
-from .. import Component, InstanceFeaturesExtractor, ObjectiveValueExtractor
-from sklearn.linear_model import LinearRegression
-from copy import deepcopy
-import numpy as np
-import logging
+from miplearn.components.component import Component
+from miplearn.extractors import InstanceFeaturesExtractor, ObjectiveValueExtractor
 
 logger = logging.getLogger(__name__)
 

miplearn/components/relaxation.py

@@ -4,8 +4,8 @@
 
 import logging
 
-from miplearn import Component
 from miplearn.classifiers.counting import CountingClassifier
+from miplearn.components.component import Component
 from miplearn.components.composite import CompositeComponent
 from miplearn.components.steps.convert_tight import ConvertTightIneqsIntoEqsStep
 from miplearn.components.steps.drop_redundant import DropRedundantInequalitiesStep

miplearn/components/steps/convert_tight.py

@@ -3,17 +3,17 @@
 #  Released under the modified BSD license. See COPYING.md for more details.
 
 import logging
+import random
 from copy import deepcopy
 
 import numpy as np
 from tqdm import tqdm
-import random
 
-from ... import Component
-from ...classifiers.counting import CountingClassifier
-from ...components import classifier_evaluation_dict
-from ...extractors import InstanceIterator
-from .drop_redundant import DropRedundantInequalitiesStep
+from miplearn.classifiers.counting import CountingClassifier
+from miplearn.components import classifier_evaluation_dict
+from miplearn.components.component import Component
+from miplearn.components.steps.drop_redundant import DropRedundantInequalitiesStep
+from miplearn.extractors import InstanceIterator
 
 logger = logging.getLogger(__name__)
 

miplearn/components/steps/drop_redundant.py

@@ -8,9 +8,9 @@ from copy import deepcopy
 import numpy as np
 from tqdm import tqdm
 
-from miplearn import Component
 from miplearn.classifiers.counting import CountingClassifier
 from miplearn.components import classifier_evaluation_dict
+from miplearn.components.component import Component
 from miplearn.components.lazy_static import LazyConstraint
 from miplearn.extractors import InstanceIterator
 

miplearn/components/steps/relax_integrality.py

@@ -4,7 +4,7 @@
 
 import logging
 
-from miplearn import Component
+from miplearn.components.component import Component
 
 logger = logging.getLogger(__name__)
 

miplearn/components/steps/tests/test_convert_tight.py

@@ -1,9 +1,12 @@
-from miplearn import LearningSolver, GurobiSolver, Instance, Classifier
+from unittest.mock import Mock
+
+from miplearn.classifiers import Classifier
 from miplearn.components.steps.convert_tight import ConvertTightIneqsIntoEqsStep
 from miplearn.components.steps.relax_integrality import RelaxIntegralityStep
+from miplearn.instance import Instance
 from miplearn.problems.knapsack import GurobiKnapsackInstance
-
-from unittest.mock import Mock
+from miplearn.solvers.gurobi import GurobiSolver
+from miplearn.solvers.learning import LearningSolver
 
 
 def test_convert_tight_usage():
@@ -40,7 +43,6 @@ def test_convert_tight_usage():
 class TestInstance(Instance):
     def to_model(self):
         import gurobipy as grb
-        from gurobipy import GRB
 
         m = grb.Model("model")
         x1 = m.addVar(name="x1")

miplearn/components/steps/tests/test_drop_redundant.py

@@ -2,21 +2,15 @@
 #  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 
-import numpy as np
 from unittest.mock import Mock, call
 
-from miplearn import (
-    LearningSolver,
-    Instance,
-    InternalSolver,
-    GurobiSolver,
-)
+import numpy as np
+
 from miplearn.classifiers import Classifier
-from miplearn.components.relaxation import (
-    DropRedundantInequalitiesStep,
-    RelaxIntegralityStep,
-)
-from miplearn.problems.knapsack import GurobiKnapsackInstance
+from miplearn.components.relaxation import DropRedundantInequalitiesStep
+from miplearn.instance import Instance
+from miplearn.solvers.internal import InternalSolver
+from miplearn.solvers.learning import LearningSolver
 
 
 def _setup():

miplearn/components/tests/test_composite.py

@@ -4,8 +4,10 @@
 
 from unittest.mock import Mock, call
 
-from miplearn import Component, LearningSolver, Instance
+from miplearn.components.component import Component
 from miplearn.components.composite import CompositeComponent
+from miplearn.instance import Instance
+from miplearn.solvers.learning import LearningSolver
 
 
 def test_composite():

miplearn/components/tests/test_lazy_dynamic.py

@@ -5,11 +5,14 @@
 from unittest.mock import Mock
 
 import numpy as np
-from miplearn import DynamicLazyConstraintsComponent, LearningSolver, InternalSolver
-from miplearn.classifiers import Classifier
-from miplearn.tests import get_test_pyomo_instances
 from numpy.linalg import norm
 
+from miplearn.classifiers import Classifier
+from miplearn.components.lazy_dynamic import DynamicLazyConstraintsComponent
+from miplearn.solvers.internal import InternalSolver
+from miplearn.solvers.learning import LearningSolver
+from miplearn.tests import get_test_pyomo_instances
+
 E = 0.1
 
 

miplearn/components/tests/test_lazy_static.py

@@ -4,13 +4,11 @@
 
 from unittest.mock import Mock, call
 
-from miplearn import (
-    StaticLazyConstraintsComponent,
-    LearningSolver,
-    Instance,
-    InternalSolver,
-)
 from miplearn.classifiers import Classifier
+from miplearn.components.lazy_static import StaticLazyConstraintsComponent
+from miplearn.instance import Instance
+from miplearn.solvers.internal import InternalSolver
+from miplearn.solvers.learning import LearningSolver
 
 
 def test_usage_with_solver():
@@ -49,7 +47,9 @@ def test_usage_with_solver():
     )
 
     component = StaticLazyConstraintsComponent(
-        threshold=0.90, use_two_phase_gap=False, violation_tolerance=1.0
+        threshold=0.90,
+        use_two_phase_gap=False,
+        violation_tolerance=1.0,
     )
     component.classifiers = {
         "type-a": Mock(spec=Classifier),

miplearn/components/tests/test_objective.py

@@ -5,8 +5,9 @@
 from unittest.mock import Mock
 
 import numpy as np
-from miplearn import ObjectiveValueComponent
+
 from miplearn.classifiers import Regressor
+from miplearn.components.objective import ObjectiveValueComponent
 from miplearn.tests import get_test_pyomo_instances
 
 

miplearn/components/tests/test_primal.py

@@ -5,8 +5,9 @@
 from unittest.mock import Mock
 
 import numpy as np
-from miplearn import PrimalSolutionComponent
+
 from miplearn.classifiers import Classifier
+from miplearn.components.primal import PrimalSolutionComponent
 from miplearn.tests import get_test_pyomo_instances
 
 

miplearn/instance.py

@@ -5,10 +5,12 @@
 import gzip
 import json
 from abc import ABC, abstractmethod
-from typing import Any
+from typing import Any, List
 
 import numpy as np
 
+from miplearn.types import TrainingSample
+
 
 class Instance(ABC):
     """
@@ -20,6 +22,9 @@ class Instance(ABC):
     into arrays of features, which can be provided as inputs to machine learning models.
     """
 
+    def __init__(self):
+        self.training_data: List[TrainingSample] = []
+
     @abstractmethod
     def to_model(self) -> Any:
         """

miplearn/log.py

@@ -2,10 +2,9 @@
 #  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 
-from datetime import timedelta
 import logging
-import time
 import sys
+import time
 
 
 class TimeFormatter(logging.Formatter):

miplearn/problems/knapsack.py

@@ -2,13 +2,13 @@
 #  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 
-import miplearn
-from miplearn import Instance
 import numpy as np
 import pyomo.environ as pe
-from scipy.stats import uniform, randint, bernoulli
+from scipy.stats import uniform, randint
 from scipy.stats.distributions import rv_frozen
 
+from miplearn.instance import Instance
+
 
 class ChallengeA:
     """
@@ -56,6 +56,7 @@ class MultiKnapsackInstance(Instance):
     """
 
     def __init__(self, prices, capacities, weights):
+        super().__init__()
         assert isinstance(prices, np.ndarray)
         assert isinstance(capacities, np.ndarray)
         assert isinstance(weights, np.ndarray)
@@ -241,6 +242,7 @@ class KnapsackInstance(Instance):
     """
 
     def __init__(self, weights, prices, capacity):
+        super().__init__()
         self.weights = weights
         self.prices = prices
         self.capacity = capacity

miplearn/problems/stab.py

@@ -8,7 +8,7 @@ import pyomo.environ as pe
 from scipy.stats import uniform, randint
 from scipy.stats.distributions import rv_frozen
 
-from miplearn import Instance
+from miplearn.instance import Instance
 
 
 class ChallengeA:
@@ -101,6 +101,7 @@ class MaxWeightStableSetInstance(Instance):
     """
 
     def __init__(self, graph, weights):
+        super().__init__()
         self.graph = graph
         self.weights = weights
 

miplearn/problems/tests/test_knapsack.py

@@ -2,10 +2,10 @@
 #  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.knapsack import MultiKnapsackGenerator, MultiKnapsackInstance
-from scipy.stats import uniform, randint
 import numpy as np
+from scipy.stats import uniform, randint
+
+from miplearn.problems.knapsack import MultiKnapsackGenerator
 
 
 def test_knapsack_generator():

miplearn/problems/tests/test_stab.py

@@ -4,10 +4,11 @@
 
 import networkx as nx
 import numpy as np
-from miplearn import LearningSolver
-from miplearn.problems.stab import MaxWeightStableSetInstance
 from scipy.stats import uniform, randint
 
+from miplearn.problems.stab import MaxWeightStableSetInstance
+from miplearn.solvers.learning import LearningSolver
+
 
 def test_stab():
     graph = nx.cycle_graph(5)

miplearn/problems/tests/test_tsp.py

@@ -2,13 +2,14 @@
 #  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
 
+from miplearn.problems.tsp import TravelingSalesmanGenerator, TravelingSalesmanInstance
+from miplearn.solvers.learning import LearningSolver
+
 
 def test_generator():
     instances = TravelingSalesmanGenerator(

miplearn/problems/tsp.py

@@ -2,14 +2,14 @@
 #  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 
+import networkx as nx
 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 import uniform, randint
 from scipy.stats.distributions import rv_frozen
-import networkx as nx
-import random
+
+from miplearn.instance import Instance
 
 
 class ChallengeA:

miplearn/solvers/gurobi.py

@@ -8,15 +8,15 @@ from io import StringIO
 from random import randint
 from typing import List, Any, Dict, Union, Tuple, Optional
 
-from . import RedirectOutput
-from .internal import (
+from miplearn.instance import Instance
+from miplearn.solvers import RedirectOutput
+from miplearn.solvers.internal import (
     InternalSolver,
     LPSolveStats,
     IterationCallback,
     LazyCallback,
     MIPSolveStats,
 )
-from .. import Instance
 
 logger = logging.getLogger(__name__)
 
@@ -181,6 +181,7 @@ class GurobiSolver(InternalSolver):
             sense = "max"
             lb = self.model.objVal
             ub = self.model.objBound
+        ws_value = self._extract_warm_start_value(log)
         stats: MIPSolveStats = {
             "Lower bound": lb,
             "Upper bound": ub,
@@ -188,10 +189,9 @@ class GurobiSolver(InternalSolver):
             "Nodes": total_nodes,
             "Sense": sense,
             "Log": log,
+            "Warm start value": ws_value,
+            "LP value": None,
         }
-        ws_value = self._extract_warm_start_value(log)
-        if ws_value is not None:
-            stats["Warm start value"] = ws_value
         return stats
 
     def get_solution(self) -> Dict:

miplearn/solvers/internal.py

@@ -4,11 +4,15 @@
 
 import logging
 from abc import ABC, abstractmethod
-from typing import Callable, Any, Dict, List
+from typing import Any, Dict, List
 
-from typing_extensions import TypedDict
-
-from ..instance import Instance
+from miplearn.instance import Instance
+from miplearn.types import (
+    LPSolveStats,
+    IterationCallback,
+    LazyCallback,
+    MIPSolveStats,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -21,33 +25,6 @@ class Constraint:
     pass
 
 
-LPSolveStats = TypedDict(
-    "LPSolveStats",
-    {
-        "Optimal value": float,
-        "Log": str,
-    },
-)
-
-MIPSolveStats = TypedDict(
-    "MIPSolveStats",
-    {
-        "Lower bound": float,
-        "Upper bound": float,
-        "Wallclock time": float,
-        "Nodes": float,
-        "Sense": str,
-        "Log": str,
-        "Warm start value": float,
-    },
-    total=False,
-)
-
-IterationCallback = Callable[[], bool]
-
-LazyCallback = Callable[[Any, Any], None]
-
-
 class InternalSolver(ABC):
     """
     Abstract class representing the MIP solver used internally by LearningSolver.

miplearn/solvers/learning.py

@@ -2,26 +2,24 @@
 #  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 
-import logging
-import pickle
-import os
-import tempfile
 import gzip
-
+import logging
+import os
+import pickle
+import tempfile
 from copy import deepcopy
-from typing import Optional, List
-from p_tqdm import p_map
-from tempfile import NamedTemporaryFile
+from typing import Optional, List, Any, IO, cast, BinaryIO, Union
 
-from . import RedirectOutput
-from .. import (
-    ObjectiveValueComponent,
-    PrimalSolutionComponent,
-    DynamicLazyConstraintsComponent,
-    UserCutsComponent,
-)
-from ..solvers.internal import InternalSolver
-from ..solvers.pyomo.gurobi import GurobiPyomoSolver
+from p_tqdm import p_map
+
+from miplearn.components.cuts import UserCutsComponent
+from miplearn.components.lazy_dynamic import DynamicLazyConstraintsComponent
+from miplearn.components.objective import ObjectiveValueComponent
+from miplearn.components.primal import PrimalSolutionComponent
+from miplearn.instance import Instance
+from miplearn.solvers import RedirectOutput
+from miplearn.solvers.pyomo.gurobi import GurobiPyomoSolver
+from miplearn.types import MIPSolveStats, TrainingSample
 
 logger = logging.getLogger(__name__)
 
@@ -192,46 +190,55 @@ class LearningSolver:
 
     def _solve(
         self,
-        instance,
-        model=None,
-        output="",
-        tee=False,
-    ):
+        instance: Instance,
+        model: Any = None,
+        output: str = "",
+        tee: bool = False,
+    ) -> MIPSolveStats:
+
+        # Load instance from file, if necessary
         filename = None
         fileformat = None
+        file: Union[BinaryIO, gzip.GzipFile]
         if isinstance(instance, str):
             filename = instance
             logger.info("Reading: %s" % filename)
             if filename.endswith(".gz"):
                 fileformat = "pickle-gz"
                 with gzip.GzipFile(filename, "rb") as file:
-                    instance = pickle.load(file)
+                    instance = pickle.load(cast(IO[bytes], file))
             else:
                 fileformat = "pickle"
                 with open(filename, "rb") as file:
-                    instance = pickle.load(file)
+                    instance = pickle.load(cast(IO[bytes], file))
 
+        # Generate model
         if model is None:
             with RedirectOutput([]):
                 model = instance.to_model()
 
+        # Initialize training data
+        training_sample: TrainingSample = {}
+
+        # Initialize internal solver
         self.tee = tee
         self.internal_solver = self.solver_factory()
         self.internal_solver.set_instance(instance, model)
 
+        # Solve linear relaxation
         if self.solve_lp_first:
             logger.info("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"]
-        else:
-            instance.lp_solution = self.internal_solver.get_empty_solution()
-            instance.lp_value = 0.0
+            stats = self.internal_solver.solve_lp(tee=tee)
+            training_sample["LP solution"] = self.internal_solver.get_solution()
+            training_sample["LP value"] = stats["Optimal value"]
+            training_sample["LP log"] = stats["Log"]
 
+        # Before-solve callbacks
         logger.debug("Running before_solve callbacks...")
         for component in self.components.values():
             component.before_solve(self, instance, model)
 
+        # Define wrappers
         def iteration_cb():
             should_repeat = False
             for comp in self.components.values():
@@ -247,29 +254,33 @@ class LearningSolver:
         if self.use_lazy_cb:
             lazy_cb = lazy_cb_wrapper
 
+        # Solve MILP
         logger.info("Solving MILP...")
         stats = self.internal_solver.solve(
             tee=tee,
             iteration_cb=iteration_cb,
             lazy_cb=lazy_cb,
         )
-        stats["LP value"] = instance.lp_value
+        if "LP value" in training_sample.keys():
+            stats["LP value"] = training_sample["LP value"]
 
         # Read MIP solution and bounds
-        instance.lower_bound = stats["Lower bound"]
-        instance.upper_bound = stats["Upper bound"]
-        instance.solver_log = stats["Log"]
-        instance.solution = self.internal_solver.get_solution()
+        training_sample["Lower bound"] = stats["Lower bound"]
+        training_sample["Upper bound"] = stats["Upper bound"]
+        training_sample["MIP log"] = stats["Log"]
+        training_sample["Solution"] = self.internal_solver.get_solution()
 
+        # After-solve callbacks
         logger.debug("Calling after_solve callbacks...")
-        training_data = {}
         for component in self.components.values():
-            component.after_solve(self, instance, model, stats, training_data)
+            component.after_solve(self, instance, model, stats, training_sample)
 
+        # Append training data
         if not hasattr(instance, "training_data"):
             instance.training_data = []
-        instance.training_data += [training_data]
+        instance.training_data += [training_sample]
 
+        # Write to file, if necessary
         if filename is not None and output is not None:
             output_filename = output
             if len(output) == 0:
@@ -277,11 +288,10 @@ class LearningSolver:
             logger.info("Writing: %s" % output_filename)
             if fileformat == "pickle":
                 with open(output_filename, "wb") as file:
-                    pickle.dump(instance, file)
+                    pickle.dump(instance, cast(IO[bytes], file))
             else:
                 with gzip.GzipFile(output_filename, "wb") as file:
-                    pickle.dump(instance, file)
-
+                    pickle.dump(instance, cast(IO[bytes], file))
         return stats
 
     def parallel_solve(

miplearn/solvers/pyomo/base.py

@@ -12,15 +12,15 @@ import pyomo
 from pyomo import environ as pe
 from pyomo.core import Var, Constraint
 
-from .. import RedirectOutput
-from ..internal import (
+from miplearn.instance import Instance
+from miplearn.solvers import RedirectOutput
+from miplearn.solvers.internal import (
     InternalSolver,
     LPSolveStats,
     IterationCallback,
     LazyCallback,
     MIPSolveStats,
 )
-from ...instance import Instance
 
 logger = logging.getLogger(__name__)
 
@@ -98,19 +98,18 @@ class BasePyomoSolver(InternalSolver):
             if not should_repeat:
                 break
         log = streams[0].getvalue()
+        node_count = self._extract_node_count(log)
+        ws_value = self._extract_warm_start_value(log)
         stats: MIPSolveStats = {
             "Lower bound": results["Problem"][0]["Lower bound"],
             "Upper bound": results["Problem"][0]["Upper bound"],
             "Wallclock time": total_wallclock_time,
             "Sense": self._obj_sense,
             "Log": log,
+            "Nodes": node_count,
+            "Warm start value": ws_value,
+            "LP value": None,
         }
-        node_count = self._extract_node_count(log)
-        ws_value = self._extract_warm_start_value(log)
-        if node_count is not None:
-            stats["Nodes"] = node_count
-        if ws_value is not None:
-            stats["Warm start value"] = ws_value
         return stats
 
     def get_solution(self) -> Dict:

miplearn/solvers/pyomo/cplex.py

@@ -5,7 +5,7 @@
 from pyomo import environ as pe
 from scipy.stats import randint
 
-from .base import BasePyomoSolver
+from miplearn.solvers.pyomo.base import BasePyomoSolver
 
 
 class CplexPyomoSolver(BasePyomoSolver):

miplearn/solvers/pyomo/gurobi.py

@@ -7,7 +7,7 @@ import logging
 from pyomo import environ as pe
 from scipy.stats import randint
 
-from .base import BasePyomoSolver
+from miplearn.solvers.pyomo.base import BasePyomoSolver
 
 logger = logging.getLogger(__name__)
 

miplearn/solvers/pyomo/xpress.py

@@ -7,7 +7,7 @@ import logging
 from pyomo import environ as pe
 from scipy.stats import randint
 
-from .base import BasePyomoSolver
+from miplearn.solvers.pyomo.base import BasePyomoSolver
 
 logger = logging.getLogger(__name__)
 

miplearn/solvers/tests/__init__.py

@@ -5,15 +5,18 @@
 from inspect import isclass
 from typing import List, Callable
 
-from miplearn import BasePyomoSolver, GurobiSolver, GurobiPyomoSolver, InternalSolver
 from miplearn.problems.knapsack import KnapsackInstance, GurobiKnapsackInstance
+from miplearn.solvers.gurobi import GurobiSolver
+from miplearn.solvers.internal import InternalSolver
+from miplearn.solvers.pyomo.base import BasePyomoSolver
+from miplearn.solvers.pyomo.gurobi import GurobiPyomoSolver
 from miplearn.solvers.pyomo.xpress import XpressPyomoSolver
 
 
 def _get_instance(solver):
-    def _is_subclass_or_instance(solver, parentClass):
-        return isinstance(solver, parentClass) or (
-            isclass(solver) and issubclass(solver, parentClass)
+    def _is_subclass_or_instance(obj, parent_class):
+        return isinstance(obj, parent_class) or (
+            isclass(obj) and issubclass(obj, parent_class)
         )
 
     if _is_subclass_or_instance(solver, BasePyomoSolver):

miplearn/solvers/tests/test_internal_solver.py

@@ -8,9 +8,10 @@ from warnings import warn
 
 import pyomo.environ as pe
 
-from miplearn import BasePyomoSolver, GurobiSolver
 from miplearn.solvers import RedirectOutput
-from . import _get_instance, _get_internal_solvers
+from miplearn.solvers.gurobi import GurobiSolver
+from miplearn.solvers.pyomo.base import BasePyomoSolver
+from miplearn.solvers.tests import _get_instance, _get_internal_solvers
 
 logger = logging.getLogger(__name__)
 
@@ -44,7 +45,7 @@ def test_internal_solver_warm_starts():
             }
         )
         stats = solver.solve(tee=True)
-        if "Warm start value" in stats:
+        if stats["Warm start value"] is not None:
             assert stats["Warm start value"] == 725.0
         else:
             warn(f"{solver_class.__name__} should set warm start value")
@@ -60,7 +61,7 @@ def test_internal_solver_warm_starts():
             }
         )
         stats = solver.solve(tee=True)
-        assert "Warm start value" not in stats
+        assert stats["Warm start value"] is None
 
         solver.fix(
             {

miplearn/solvers/tests/test_lazy_cb.py

@@ -4,8 +4,8 @@
 
 import logging
 
-from . import _get_instance
-from ... import GurobiSolver
+from miplearn.solvers.gurobi import GurobiSolver
+from miplearn.solvers.tests import _get_instance
 
 logger = logging.getLogger(__name__)
 

miplearn/solvers/tests/test_learning_solver.py

@@ -7,13 +7,9 @@ import pickle
 import tempfile
 import os
 
-from miplearn import (
-    LearningSolver,
-    GurobiSolver,
-    DynamicLazyConstraintsComponent,
-)
-
-from . import _get_instance, _get_internal_solvers
+from miplearn.solvers.gurobi import GurobiSolver
+from miplearn.solvers.learning import LearningSolver
+from miplearn.solvers.tests import _get_instance, _get_internal_solvers
 
 logger = logging.getLogger(__name__)
 

miplearn/tests/__init__.py

@@ -1,8 +1,9 @@
 #  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.knapsack import KnapsackInstance
+from miplearn.solvers.learning import LearningSolver
 
 
 def get_test_pyomo_instances():

miplearn/tests/test_benchmark.py

@@ -4,10 +4,12 @@
 
 import os.path
 
-from miplearn import LearningSolver, BenchmarkRunner
+from miplearn.benchmark import BenchmarkRunner
 from miplearn.problems.stab import MaxWeightStableSetGenerator
 from scipy.stats import randint
 
+from miplearn.solvers.learning import LearningSolver
+
 
 def test_benchmark():
     # Generate training and test instances

miplearn/tests/test_extractors.py

@@ -1,16 +1,15 @@
 #  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.
+import numpy as np
 
-from miplearn.problems.knapsack import KnapsackInstance
-from miplearn import (
-    LearningSolver,
+from miplearn.extractors import (
     SolutionExtractor,
     InstanceFeaturesExtractor,
     VariableFeaturesExtractor,
 )
-import numpy as np
-import pyomo.environ as pe
+from miplearn.problems.knapsack import KnapsackInstance
+from miplearn.solvers.learning import LearningSolver
 
 
 def _get_instances():

miplearn/types.py

@@ -0,0 +1,45 @@
+#  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 typing import TypedDict, Optional, Dict, Callable, Any
+
+TrainingSample = TypedDict(
+    "TrainingSample",
+    {
+        "LP log": Optional[str],
+        "LP solution": Optional[Dict],
+        "LP value": Optional[float],
+        "Lower bound": Optional[float],
+        "MIP log": Optional[str],
+        "Solution": Optional[Dict],
+        "Upper bound": Optional[float],
+    },
+    total=False,
+)
+
+LPSolveStats = TypedDict(
+    "LPSolveStats",
+    {
+        "Optimal value": float,
+        "Log": str,
+    },
+)
+
+MIPSolveStats = TypedDict(
+    "MIPSolveStats",
+    {
+        "Lower bound": Optional[float],
+        "Upper bound": Optional[float],
+        "Wallclock time": float,
+        "Nodes": Optional[int],
+        "Sense": str,
+        "Log": str,
+        "Warm start value": Optional[float],
+        "LP value": Optional[float],
+    },
+)
+
+IterationCallback = Callable[[], bool]
+
+LazyCallback = Callable[[Any, Any], None]