diff --git a/miplearn/features.py b/miplearn/features.py index ea0ea90..c25f0f5 100644 --- a/miplearn/features.py +++ b/miplearn/features.py @@ -13,7 +13,7 @@ import numpy as np from miplearn.types import Solution, Category if TYPE_CHECKING: - from miplearn.solvers.internal import InternalSolver, LPSolveStats + from miplearn.solvers.internal import InternalSolver, LPSolveStats, MIPSolveStats from miplearn.instance.base import Instance @@ -103,6 +103,7 @@ class Features: variables: Optional[Dict[str, Variable]] = None constraints: Optional[Dict[str, Constraint]] = None lp_solve: Optional["LPSolveStats"] = None + mip_solve: Optional["MIPSolveStats"] = None class FeaturesExtractor: diff --git a/miplearn/solvers/learning.py b/miplearn/solvers/learning.py index 8b05cb0..5420b41 100644 --- a/miplearn/solvers/learning.py +++ b/miplearn/solvers/learning.py @@ -171,6 +171,7 @@ class LearningSolver: # Solve root LP relaxation # ------------------------------------------------------- + lp_stats = None if self.solve_lp: logger.debug("Running before_solve_lp callbacks...") for component in self.components.values(): @@ -252,6 +253,8 @@ class LearningSolver: # ------------------------------------------------------- logger.info("Extracting features (after-mip)...") features = FeaturesExtractor(self.internal_solver).extract(instance) + features.lp_solve = lp_stats + features.mip_solve = mip_stats instance.features_after_mip.append(features) # Add some information to training_sample