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 = [