Add mip_stats to after-mip features

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:

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