Remove sample.after_load

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)