Remove sample.after_lp

miplearn/components/objective.py

@@ -77,20 +77,15 @@ class ObjectiveValueComponent(Component):
         _: Optional[Instance],
         sample: Sample,
     ) -> Tuple[Dict[Hashable, List[List[float]]], Dict[Hashable, List[List[float]]]]:
-        # Instance features
-        assert sample.after_load is not None
-        assert sample.after_load.instance is not None
-        f = sample.after_load.instance.to_list()
-
-        # LP solve features
-        if sample.after_lp is not None:
-            assert sample.after_lp.lp_solve is not None
-            f.extend(sample.after_lp.lp_solve.to_list())
+        lp_instance_features = sample.get("lp_instance_features")
+        if lp_instance_features is None:
+            lp_instance_features = sample.get("instance_features_user")
+        assert lp_instance_features is not None
 
         # Features
         x: Dict[Hashable, List[List[float]]] = {
-            "Upper bound": [f],
-            "Lower bound": [f],
+            "Upper bound": [lp_instance_features],
+            "Lower bound": [lp_instance_features],
         }
 
         # Labels

miplearn/components/primal.py

@@ -155,7 +155,15 @@ class PrimalSolutionComponent(Component):
         assert sample.after_load.variables is not None
         assert sample.after_load.variables.names is not None
         assert sample.after_load.variables.categories is not None
+
+        instance_features = sample.get("instance_features_user")
         mip_var_values = sample.get("mip_var_values")
+        var_features = sample.get("lp_var_features")
+        if var_features is None:
+            var_features = sample.get("var_features")
+
+        assert instance_features is not None
+        assert var_features is not None
 
         for (i, var_name) in enumerate(sample.after_load.variables.names):
             # Initialize categories
@@ -167,11 +175,8 @@ class PrimalSolutionComponent(Component):
                 y[category] = []
 
             # Features
-            features = list(sample.after_load.instance.to_list())
-            features.extend(sample.after_load.variables.to_list(i))
-            if sample.after_lp is not None:
-                assert sample.after_lp.variables is not None
-                features.extend(sample.after_lp.variables.to_list(i))
+            features = list(instance_features)
+            features.extend(var_features[i])
             x[category].append(features)
 
             # Labels

miplearn/components/static_lazy.py

@@ -204,17 +204,26 @@ class StaticLazyConstraintsComponent(Component):
         x: Dict[Hashable, List[List[float]]] = {}
         y: Dict[Hashable, List[List[float]]] = {}
         cids: Dict[Hashable, List[str]] = {}
-        assert sample.after_load is not None
-        constraints = sample.after_load.constraints
-        assert constraints is not None
-        assert constraints.names is not None
-        assert constraints.lazy is not None
-        assert constraints.categories is not None
-        for (cidx, cname) in enumerate(constraints.names):
+        instance_features = sample.get("instance_features_user")
+        constr_features = sample.get("lp_constr_features")
+        constr_names = sample.get("constr_names")
+        constr_categories = sample.get("constr_categories")
+        constr_lazy = sample.get("constr_lazy")
+        lazy_enforced = sample.get("lazy_enforced")
+        if constr_features is None:
+            constr_features = sample.get("constr_features_user")
+
+        assert instance_features is not None
+        assert constr_features is not None
+        assert constr_names is not None
+        assert constr_categories is not None
+        assert constr_lazy is not None
+
+        for (cidx, cname) in enumerate(constr_names):
             # Initialize categories
-            if not constraints.lazy[cidx]:
+            if not constr_lazy[cidx]:
                 continue
-            category = constraints.categories[cidx]
+            category = constr_categories[cidx]
             if category is None:
                 continue
             if category not in x:
@@ -223,18 +232,12 @@ class StaticLazyConstraintsComponent(Component):
                 cids[category] = []
 
             # Features
-            sf = sample.after_load
-            if sample.after_lp is not None:
-                sf = sample.after_lp
-            assert sf.instance is not None
-            assert sf.constraints is not None
-            features = list(sf.instance.to_list())
-            features.extend(sf.constraints.to_list(cidx))
+            features = list(instance_features)
+            features.extend(constr_features[cidx])
             x[category].append(features)
             cids[category].append(cname)
 
             # Labels
-            lazy_enforced = sample.get("lazy_enforced")
             if lazy_enforced is not None:
                 if cname in lazy_enforced:
                     y[category] += [[False, True]]

miplearn/features.py

@@ -147,14 +147,12 @@ class Sample:
     def __init__(
         self,
         after_load: Optional[Features] = None,
-        after_lp: Optional[Features] = None,
         data: Optional[Dict[str, Any]] = None,
     ) -> None:
         if data is None:
             data = {}
         self._data: Dict[str, Any] = data
         self.after_load = after_load
-        self.after_lp = after_lp
 
     def get(self, key: str) -> Optional[Any]:
         if key in self._data:
@@ -253,6 +251,29 @@ class FeaturesExtractor:
                 ],
             ),
         )
+        sample.put(
+            "lp_constr_features",
+            self._combine(
+                sample,
+                [
+                    "constr_features_user",
+                    "lp_constr_dual_values",
+                    "lp_constr_sa_rhs_down",
+                    "lp_constr_sa_rhs_up",
+                    "lp_constr_slacks",
+                ],
+            ),
+        )
+        instance_features_user = sample.get("instance_features_user")
+        assert instance_features_user is not None
+        sample.put(
+            "lp_instance_features",
+            instance_features_user
+            + [
+                sample.get("lp_value"),
+                sample.get("lp_wallclock_time"),
+            ],
+        )
 
     def extract_after_mip_features(
         self,

miplearn/solvers/learning.py

@@ -210,20 +210,13 @@ class LearningSolver:
             # -------------------------------------------------------
             logger.info("Extracting features (after-lp)...")
             initial_time = time.time()
+            for (k, v) in lp_stats.__dict__.items():
+                sample.put(k, v)
             self.extractor.extract_after_lp_features(self.internal_solver, sample)
-            features = self.extractor.extract(
-                instance,
-                self.internal_solver,
-                with_static=False,
-            )
             logger.info(
                 "Features (after-lp) extracted in %.2f seconds"
                 % (time.time() - initial_time)
             )
-            for (k, v) in lp_stats.__dict__.items():
-                sample.put(k, v)
-            features.lp_solve = lp_stats
-            sample.after_lp = features
 
         # Callback wrappers
         # -------------------------------------------------------
@@ -285,9 +278,9 @@ class LearningSolver:
         # -------------------------------------------------------
         logger.info("Extracting features (after-mip)...")
         initial_time = time.time()
-        self.extractor.extract_after_mip_features(self.internal_solver, sample)
         for (k, v) in mip_stats.__dict__.items():
             sample.put(k, v)
+        self.extractor.extract_after_mip_features(self.internal_solver, sample)
         logger.info(
             "Features (after-mip) extracted in %.2f seconds"
             % (time.time() - initial_time)

tests/components/test_objective.py

@@ -19,19 +19,12 @@ from miplearn.solvers.pyomo.gurobi import GurobiPyomoSolver
 @pytest.fixture
 def sample() -> Sample:
     sample = Sample(
-        after_load=Features(
-            instance=InstanceFeatures(),
-        ),
-        after_lp=Features(
-            lp_solve=LPSolveStats(),
-        ),
         data={
             "mip_lower_bound": 1.0,
             "mip_upper_bound": 2.0,
+            "lp_instance_features": [1.0, 2.0, 3.0],
         },
     )
-    sample.after_load.instance.to_list = Mock(return_value=[1.0, 2.0])  # type: ignore
-    sample.after_lp.lp_solve.to_list = Mock(return_value=[3.0])  # type: ignore
     return sample
 
 

tests/components/test_primal.py

@@ -33,12 +33,23 @@ def sample() -> Sample:
                 categories=["default", None, "default", "default"],
             ),
         ),
-        after_lp=Features(
-            variables=VariableFeatures(),
-        ),
         data={
             "var_names": ["x[0]", "x[1]", "x[2]", "x[3]"],
+            "var_categories": ["default", None, "default", "default"],
             "mip_var_values": [0.0, 1.0, 1.0, 0.0],
+            "instance_features_user": [5.0],
+            "var_features": [
+                [0.0, 0.0],
+                None,
+                [1.0, 0.0],
+                [1.0, 1.0],
+            ],
+            "lp_var_features": [
+                [0.0, 0.0, 2.0, 2.0],
+                None,
+                [1.0, 0.0, 3.0, 2.0],
+                [1.0, 1.0, 3.0, 3.0],
+            ],
         },
     )
     sample.after_load.instance.to_list = Mock(return_value=[5.0])  # type: ignore
@@ -50,14 +61,6 @@ def sample() -> Sample:
             [1.0, 1.0],
         ][i]
     )
-    sample.after_lp.variables.to_list = Mock(  # type:ignore
-        side_effect=lambda i: [
-            [2.0, 2.0],
-            None,
-            [3.0, 2.0],
-            [3.0, 3.0],
-        ][i]
-    )
     return sample
 
 

tests/components/test_static_lazy.py

@@ -44,24 +44,28 @@ def sample() -> Sample:
                 lazy=[True, True, True, True, False],
             ),
         ),
-        after_lp=Features(
-            instance=InstanceFeatures(),
-            constraints=ConstraintFeatures(names=["c1", "c2", "c3", "c4", "c5"]),
-        ),
         data={
+            "constr_categories": [
+                "type-a",
+                "type-a",
+                "type-a",
+                "type-b",
+                "type-b",
+            ],
+            "constr_lazy": [True, True, True, True, False],
+            "constr_names": ["c1", "c2", "c3", "c4", "c5"],
+            "instance_features_user": [5.0],
             "lazy_enforced": {"c1", "c2", "c4"},
+            "lp_constr_features": [
+                [1.0, 1.0],
+                [1.0, 2.0],
+                [1.0, 3.0],
+                [1.0, 4.0, 0.0],
+                None,
+            ],
+            "static_lazy_count": 4,
         },
     )
-    sample.after_lp.instance.to_list = Mock(return_value=[5.0])  # type: ignore
-    sample.after_lp.constraints.to_list = Mock(  # type: ignore
-        side_effect=lambda idx: {
-            0: [1.0, 1.0],
-            1: [1.0, 2.0],
-            2: [1.0, 3.0],
-            3: [1.0, 4.0, 0.0],
-            4: None,
-        }[idx]
-    )
     return sample