Add training_data argument to after_solve

miplearn/components/steps/convert_tight.py

@@ -74,13 +74,21 @@ class ConvertTightIneqsIntoEqsStep(Component):
 
         logger.info(f"Converted {self.n_converted} inequalities")
 
-    def after_solve(self, solver, instance, model, results):
-        instance.slacks = solver.internal_solver.get_inequality_slacks()
-        results["ConvertTight: Kept"] = self.n_kept
-        results["ConvertTight: Converted"] = self.n_converted
-        results["ConvertTight: Restored"] = self.n_restored
-        results["ConvertTight: Inf iterations"] = self.n_infeasible_iterations
-        results["ConvertTight: Subopt iterations"] = self.n_suboptimal_iterations
+    def after_solve(
+        self,
+        solver,
+        instance,
+        model,
+        stats,
+        training_data,
+    ):
+        if "slacks" not in training_data.keys():
+            training_data["slacks"] = solver.internal_solver.get_inequality_slacks()
+        stats["ConvertTight: Kept"] = self.n_kept
+        stats["ConvertTight: Converted"] = self.n_converted
+        stats["ConvertTight: Restored"] = self.n_restored
+        stats["ConvertTight: Inf iterations"] = self.n_infeasible_iterations
+        stats["ConvertTight: Subopt iterations"] = self.n_suboptimal_iterations
 
     def fit(self, training_instances):
         logger.debug("Extracting x and y...")
@@ -108,7 +116,7 @@ class ConvertTightIneqsIntoEqsStep(Component):
             if constraint_ids is not None:
                 cids = constraint_ids
             else:
-                cids = instance.slacks.keys()
+                cids = instance.training_data[0]["slacks"].keys()
             for cid in cids:
                 category = instance.get_constraint_category(cid)
                 if category is None:
@@ -130,7 +138,7 @@ class ConvertTightIneqsIntoEqsStep(Component):
             desc="Extract (rlx:conv_ineqs:y)",
             disable=len(instances) < 5,
         ):
-            for (cid, slack) in instance.slacks.items():
+            for (cid, slack) in instance.training_data[0]["slacks"].items():
                 category = instance.get_constraint_category(cid)
                 if category is None:
                     continue

miplearn/components/steps/drop_redundant.py

@@ -76,12 +76,19 @@ class DropRedundantInequalitiesStep(Component):
                     self.total_kept += 1
         logger.info(f"Extracted {self.total_dropped} predicted constraints")
 
-    def after_solve(self, solver, instance, model, results):
+    def after_solve(
+        self,
+        solver,
+        instance,
+        model,
+        stats,
+        training_data,
+    ):
         instance.slacks = solver.internal_solver.get_inequality_slacks()
-        results["DropRedundant: Kept"] = self.total_kept
-        results["DropRedundant: Dropped"] = self.total_dropped
-        results["DropRedundant: Restored"] = self.total_restored
-        results["DropRedundant: Iterations"] = self.total_iterations
+        stats["DropRedundant: Kept"] = self.total_kept
+        stats["DropRedundant: Dropped"] = self.total_dropped
+        stats["DropRedundant: Restored"] = self.total_restored
+        stats["DropRedundant: Iterations"] = self.total_iterations
 
     def fit(self, training_instances):
         logger.debug("Extracting x and y...")

miplearn/components/steps/relax_integrality.py

@@ -17,3 +17,13 @@ class RelaxIntegralityStep(Component):
     def before_solve(self, solver, instance, _):
         logger.info("Relaxing integrality...")
         solver.internal_solver.relax()
+
+    def after_solve(
+        self,
+        solver,
+        instance,
+        model,
+        stats,
+        training_data,
+    ):
+        return

miplearn/components/steps/tests/convert_tight_test.py

@@ -25,8 +25,7 @@ def test_convert_tight_usage():
     original_upper_bound = instance.upper_bound
 
     # Should collect training data
-    assert hasattr(instance, "slacks")
-    assert instance.slacks["eq_capacity"] == 0.0
+    assert instance.training_data[0]["slacks"]["eq_capacity"] == 0.0
 
     # Fit and resolve
     solver.fit([instance])
@@ -53,21 +52,6 @@ class TestInstance(Instance):
         return m
 
 
-class TestInstanceMin(Instance):
-    def to_model(self):
-        import gurobipy as grb
-        from gurobipy import GRB
-
-        m = grb.Model("model")
-        x1 = m.addVar(name="x1")
-        x2 = m.addVar(name="x2")
-        m.setObjective(x1 + 2 * x2, grb.GRB.MAXIMIZE)
-        m.addConstr(x1 <= 2, name="c1")
-        m.addConstr(x2 <= 2, name="c2")
-        m.addConstr(x1 + x2 <= 3, name="c2")
-        return m
-
-
 def test_convert_tight_infeasibility():
     comp = ConvertTightIneqsIntoEqsStep()
     comp.classifiers = {