Remove sample.after_load

miplearn/components/dynamic_common.py

@@ -52,6 +52,8 @@ class DynamicConstraintsComponent(Component):
         cids: Dict[Hashable, List[str]] = {}
         constr_categories_dict = instance.get_constraint_categories()
         constr_features_dict = instance.get_constraint_features()
+        instance_features = sample.get("instance_features_user")
+        assert instance_features is not None
         for cid in self.known_cids:
             # Initialize categories
             if cid in constr_categories_dict:
@@ -66,10 +68,8 @@ class DynamicConstraintsComponent(Component):
                 cids[category] = []
 
             # Features
-            features = []
-            assert sample.after_load is not None
-            assert sample.after_load.instance is not None
-            features.extend(sample.after_load.instance.to_list())
+            features: List[float] = []
+            features.extend(instance_features)
             if cid in constr_features_dict:
                 features.extend(constr_features_dict[cid])
             for ci in features:

miplearn/components/primal.py

@@ -103,8 +103,10 @@ class PrimalSolutionComponent(Component):
         )
 
     def sample_predict(self, sample: Sample) -> Solution:
-        assert sample.after_load is not None
-        assert sample.after_load.variables is not None
+        var_names = sample.get("var_names")
+        var_categories = sample.get("var_categories")
+        assert var_names is not None
+        assert var_categories is not None
 
         # Compute y_pred
         x, _ = self.sample_xy(None, sample)
@@ -125,12 +127,10 @@ class PrimalSolutionComponent(Component):
             ).T
 
         # Convert y_pred into solution
-        assert sample.after_load.variables.names is not None
-        assert sample.after_load.variables.categories is not None
-        solution: Solution = {v: None for v in sample.after_load.variables.names}
+        solution: Solution = {v: None for v in var_names}
         category_offset: Dict[Hashable, int] = {cat: 0 for cat in x.keys()}
-        for (i, var_name) in enumerate(sample.after_load.variables.names):
-            category = sample.after_load.variables.categories[i]
+        for (i, var_name) in enumerate(var_names):
+            category = var_categories[i]
             if category not in category_offset:
                 continue
             offset = category_offset[category]
@@ -150,24 +150,21 @@ class PrimalSolutionComponent(Component):
     ) -> Tuple[Dict[Category, List[List[float]]], Dict[Category, List[List[float]]]]:
         x: Dict = {}
         y: Dict = {}
-        assert sample.after_load is not None
-        assert sample.after_load.instance is not None
-        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")
+        var_names = sample.get("var_names")
+        var_categories = sample.get("var_categories")
         if var_features is None:
             var_features = sample.get("var_features")
-
         assert instance_features is not None
         assert var_features is not None
+        assert var_names is not None
+        assert var_categories is not None
 
-        for (i, var_name) in enumerate(sample.after_load.variables.names):
+        for (i, var_name) in enumerate(var_names):
             # Initialize categories
-            category = sample.after_load.variables.categories[i]
+            category = var_categories[i]
             if category is None:
                 continue
             if category not in x.keys():

miplearn/components/static_lazy.py

@@ -74,16 +74,15 @@ class StaticLazyConstraintsComponent(Component):
         sample: Sample,
     ) -> None:
         assert solver.internal_solver is not None
-        assert sample.after_load is not None
-        assert sample.after_load.instance is not None
+        static_lazy_count = sample.get("static_lazy_count")
+        assert static_lazy_count is not None
 
         logger.info("Predicting violated (static) lazy constraints...")
-        if sample.after_load.instance.lazy_constraint_count == 0:
+        if static_lazy_count == 0:
             logger.info("Instance does not have static lazy constraints. Skipping.")
         self.enforced_cids = set(self.sample_predict(sample))
         logger.info("Moving lazy constraints to the pool...")
-        constraints = sample.after_load.constraints
-        assert constraints is not None
+        constraints = ConstraintFeatures.from_sample(sample)
         assert constraints.lazy is not None
         assert constraints.names is not None
         selected = [

miplearn/features.py

@@ -82,9 +82,9 @@ class ConstraintFeatures:
     basis_status: Optional[List[str]] = None
     categories: Optional[List[Optional[Hashable]]] = None
     dual_values: Optional[List[float]] = None
-    names: Optional[List[str]] = None
     lazy: Optional[List[bool]] = None
     lhs: Optional[List[List[Tuple[str, float]]]] = None
+    names: Optional[List[str]] = None
     rhs: Optional[List[float]] = None
     sa_rhs_down: Optional[List[float]] = None
     sa_rhs_up: Optional[List[float]] = None
@@ -92,6 +92,23 @@ class ConstraintFeatures:
     slacks: Optional[List[float]] = None
     user_features: Optional[List[Optional[List[float]]]] = None
 
+    @staticmethod
+    def from_sample(sample: "Sample") -> "ConstraintFeatures":
+        return ConstraintFeatures(
+            basis_status=sample.get("lp_constr_basis_status"),
+            categories=sample.get("constr_categories"),
+            dual_values=sample.get("lp_constr_dual_values"),
+            lazy=sample.get("constr_lazy"),
+            lhs=sample.get("constr_lhs"),
+            names=sample.get("constr_names"),
+            rhs=sample.get("constr_rhs"),
+            sa_rhs_down=sample.get("lp_constr_sa_rhs_down"),
+            sa_rhs_up=sample.get("lp_constr_sa_rhs_up"),
+            senses=sample.get("constr_senses"),
+            slacks=sample.get("lp_constr_slacks"),
+            user_features=sample.get("constr_features_user"),
+        )
+
     def to_list(self, index: int) -> List[float]:
         features: List[float] = []
         for attr in [
@@ -146,13 +163,11 @@ class Features:
 class Sample:
     def __init__(
         self,
-        after_load: 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
 
     def get(self, key: str) -> Optional[Any]:
         if key in self._data:

miplearn/solvers/learning.py

@@ -176,7 +176,6 @@ class LearningSolver:
             "Features (after-load) extracted in %.2f seconds"
             % (time.time() - initial_time)
         )
-        sample.after_load = features
 
         callback_args = (
             self,

tests/components/test_dynamic_lazy.py

@@ -27,16 +27,18 @@ def training_instances() -> List[Instance]:
     instances = [cast(Instance, Mock(spec=Instance)) for _ in range(2)]
     samples_0 = [
         Sample(
-            after_load=Features(instance=InstanceFeatures()),
-            data={"lazy_enforced": {"c1", "c2"}},
+            {
+                "lazy_enforced": {"c1", "c2"},
+                "instance_features_user": [5.0],
+            },
         ),
         Sample(
-            after_load=Features(instance=InstanceFeatures()),
-            data={"lazy_enforced": {"c2", "c3"}},
+            {
+                "lazy_enforced": {"c2", "c3"},
+                "instance_features_user": [5.0],
+            },
         ),
     ]
-    samples_0[0].after_load.instance.to_list = Mock(return_value=[5.0])  # type: ignore
-    samples_0[1].after_load.instance.to_list = Mock(return_value=[5.0])  # type: ignore
     instances[0].get_samples = Mock(return_value=samples_0)  # type: ignore
     instances[0].get_constraint_categories = Mock(  # type: ignore
         return_value={
@@ -56,11 +58,12 @@ def training_instances() -> List[Instance]:
     )
     samples_1 = [
         Sample(
-            after_load=Features(instance=InstanceFeatures()),
-            data={"lazy_enforced": {"c3", "c4"}},
+            {
+                "lazy_enforced": {"c3", "c4"},
+                "instance_features_user": [8.0],
+            },
         )
     ]
-    samples_1[0].after_load.instance.to_list = Mock(return_value=[8.0])  # type: ignore
     instances[1].get_samples = Mock(return_value=samples_1)  # type: ignore
     instances[1].get_constraint_categories = Mock(  # type: ignore
         return_value={

tests/components/test_objective.py

@@ -10,8 +10,7 @@ from numpy.testing import assert_array_equal
 
 from miplearn.classifiers import Regressor
 from miplearn.components.objective import ObjectiveValueComponent
-from miplearn.features import InstanceFeatures, Features, Sample
-from miplearn.solvers.internal import MIPSolveStats, LPSolveStats
+from miplearn.features import Sample
 from miplearn.solvers.learning import LearningSolver
 from miplearn.solvers.pyomo.gurobi import GurobiPyomoSolver
 
@@ -19,7 +18,7 @@ from miplearn.solvers.pyomo.gurobi import GurobiPyomoSolver
 @pytest.fixture
 def sample() -> Sample:
     sample = Sample(
-        data={
+        {
             "mip_lower_bound": 1.0,
             "mip_upper_bound": 2.0,
             "lp_instance_features": [1.0, 2.0, 3.0],

tests/components/test_primal.py

@@ -26,14 +26,7 @@ from miplearn.solvers.tests import assert_equals
 @pytest.fixture
 def sample() -> Sample:
     sample = Sample(
-        after_load=Features(
-            instance=InstanceFeatures(),
-            variables=VariableFeatures(
-                names=["x[0]", "x[1]", "x[2]", "x[3]"],
-                categories=["default", None, "default", "default"],
-            ),
-        ),
-        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],
@@ -52,15 +45,6 @@ def sample() -> Sample:
             ],
         },
     )
-    sample.after_load.instance.to_list = Mock(return_value=[5.0])  # type: ignore
-    sample.after_load.variables.to_list = Mock(  # type:ignore
-        side_effect=lambda i: [
-            [0.0, 0.0],
-            None,
-            [1.0, 0.0],
-            [1.0, 1.0],
-        ][i]
-    )
     return sample
 
 

tests/components/test_static_lazy.py

@@ -28,23 +28,7 @@ from miplearn.types import (
 @pytest.fixture
 def sample() -> Sample:
     sample = Sample(
-        after_load=Features(
-            instance=InstanceFeatures(
-                lazy_constraint_count=4,
-            ),
-            constraints=ConstraintFeatures(
-                names=["c1", "c2", "c3", "c4", "c5"],
-                categories=[
-                    "type-a",
-                    "type-a",
-                    "type-a",
-                    "type-b",
-                    "type-b",
-                ],
-                lazy=[True, True, True, True, False],
-            ),
-        ),
-        data={
+        {
             "constr_categories": [
                 "type-a",
                 "type-a",
@@ -139,9 +123,7 @@ def test_usage_with_solver(instance: Instance) -> None:
 
     # Should ask internal solver to verify if constraints in the pool are
     # satisfied and add the ones that are not
-    assert sample.after_load is not None
-    assert sample.after_load.constraints is not None
-    c = sample.after_load.constraints[[False, False, True, False, False]]
+    c = ConstraintFeatures.from_sample(sample)[[False, False, True, False, False]]
     internal.are_constraints_satisfied.assert_called_once_with(c, tol=1.0)
     internal.are_constraints_satisfied.reset_mock()
     internal.add_constraints.assert_called_once_with(c)