Reorganize imports; start moving data to instance.training_data

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__)