Use np.ndarray for constraint methods in Instance

miplearn/components/dynamic_common.py

@@ -8,12 +8,14 @@ from typing import Dict, List, Tuple, Optional, Any, Set
 import numpy as np
 from overrides import overrides
 
+from miplearn.features.extractor import FeaturesExtractor
 from miplearn.classifiers import Classifier
 from miplearn.classifiers.threshold import Threshold
 from miplearn.components import classifier_evaluation_dict
 from miplearn.components.component import Component
 from miplearn.features.sample import Sample
 from miplearn.instance.base import Instance
+from miplearn.types import ConstraintCategory, ConstraintName
 
 logger = logging.getLogger(__name__)
 
@@ -32,9 +34,9 @@ class DynamicConstraintsComponent(Component):
         assert isinstance(classifier, Classifier)
         self.threshold_prototype: Threshold = threshold
         self.classifier_prototype: Classifier = classifier
-        self.classifiers: Dict[str, Classifier] = {}
+        self.classifiers: Dict[ConstraintCategory, Classifier] = {}
-        self.thresholds: Dict[str, Threshold] = {}
+        self.thresholds: Dict[ConstraintCategory, Threshold] = {}
-        self.known_cids: List[str] = []
+        self.known_cids: List[ConstraintName] = []
         self.attr = attr
 
     def sample_xy_with_cids(
@@ -42,51 +44,45 @@ class DynamicConstraintsComponent(Component):
         instance: Optional[Instance],
         sample: Sample,
     ) -> Tuple[
-        Dict[str, List[List[float]]],
+        Dict[ConstraintCategory, List[List[float]]],
-        Dict[str, List[List[bool]]],
+        Dict[ConstraintCategory, List[List[bool]]],
-        Dict[str, List[str]],
+        Dict[ConstraintCategory, List[ConstraintName]],
     ]:
+        if len(self.known_cids) == 0:
+            return {}, {}, {}
         assert instance is not None
-        x: Dict[str, List[List[float]]] = {}
+        x: Dict[ConstraintCategory, List[List[float]]] = {}
-        y: Dict[str, List[List[bool]]] = {}
+        y: Dict[ConstraintCategory, List[List[bool]]] = {}
-        cids: Dict[str, List[str]] = {}
+        cids: Dict[ConstraintCategory, List[ConstraintName]] = {}
-        constr_categories_dict = instance.get_constraint_categories()
+        known_cids = np.array(self.known_cids, dtype="S")
-        constr_features_dict = instance.get_constraint_features()
+
+        # Get user-provided constraint features
+        (
+            constr_features,
+            constr_categories,
+            constr_lazy,
+        ) = FeaturesExtractor._extract_user_features_constrs(instance, known_cids)
+
+        # Augment with instance features
         instance_features = sample.get_array("static_instance_features")
         assert instance_features is not None
-        for cid in self.known_cids:
+        constr_features = np.hstack(
-            # Initialize categories
+            [
-            if cid in constr_categories_dict:
+                instance_features.reshape(1, -1).repeat(len(known_cids), axis=0),
-                category = constr_categories_dict[cid]
+                constr_features,
-            else:
+            ]
-                category = cid
+        )
-            if category is None:
+        assert len(known_cids) == constr_features.shape[0]
-                continue
+
-            if category not in x:
+        categories = np.unique(constr_categories)
-                x[category] = []
+        for c in categories:
-                y[category] = []
+            x[c] = constr_features[constr_categories == c].tolist()
-                cids[category] = []
+            cids[c] = known_cids[constr_categories == c].tolist()
-
+            enforced_cids = np.array(list(sample.get_set(self.attr)), dtype="S")
-            # Features
-            features: List[float] = []
-            features.extend(instance_features)
-            if cid in constr_features_dict:
-                features.extend(constr_features_dict[cid])
-            for ci in features:
-                assert isinstance(ci, float), (
-                    f"Constraint features must be a list of floats. "
-                    f"Found {ci.__class__.__name__} instead."
-                )
-            x[category].append(features)
-            cids[category].append(cid)
-
-            # Labels
-            enforced_cids = sample.get_set(self.attr)
             if enforced_cids is not None:
-                if cid in enforced_cids:
+                tmp = np.isin(cids[c], enforced_cids).reshape(-1, 1)
-                    y[category] += [[False, True]]
+                y[c] = np.hstack([~tmp, tmp]).tolist()  # type: ignore
-                else:
+
-                    y[category] += [[True, False]]
         return x, y, cids
 
     @overrides
@@ -111,8 +107,8 @@ class DynamicConstraintsComponent(Component):
         self,
         instance: Instance,
         sample: Sample,
-    ) -> List[str]:
+    ) -> List[ConstraintName]:
-        pred: List[str] = []
+        pred: List[ConstraintName] = []
         if len(self.known_cids) == 0:
             logger.info("Classifiers not fitted. Skipping.")
             return pred
@@ -137,8 +133,8 @@ class DynamicConstraintsComponent(Component):
     @overrides
     def fit_xy(
         self,
-        x: Dict[str, np.ndarray],
+        x: Dict[ConstraintCategory, np.ndarray],
-        y: Dict[str, np.ndarray],
+        y: Dict[ConstraintCategory, np.ndarray],
     ) -> None:
         for category in x.keys():
             self.classifiers[category] = self.classifier_prototype.clone()
@@ -153,40 +149,20 @@ class DynamicConstraintsComponent(Component):
         self,
         instance: Instance,
         sample: Sample,
-    ) -> Dict[str, Dict[str, float]]:
+    ) -> Dict[str, float]:
         actual = sample.get_set(self.attr)
         assert actual is not None
         pred = set(self.sample_predict(instance, sample))
-        tp: Dict[str, int] = {}
+        tp, tn, fp, fn = 0, 0, 0, 0
-        tn: Dict[str, int] = {}
-        fp: Dict[str, int] = {}
-        fn: Dict[str, int] = {}
-        constr_categories_dict = instance.get_constraint_categories()
         for cid in self.known_cids:
-            if cid not in constr_categories_dict:
-                continue
-            category = constr_categories_dict[cid]
-            if category not in tp.keys():
-                tp[category] = 0
-                tn[category] = 0
-                fp[category] = 0
-                fn[category] = 0
             if cid in pred:
                 if cid in actual:
-                    tp[category] += 1
+                    tp += 1
                 else:
-                    fp[category] += 1
+                    fp += 1
             else:
                 if cid in actual:
-                    fn[category] += 1
+                    fn += 1
                 else:
-                    tn[category] += 1
+                    tn += 1
-        return {
+        return classifier_evaluation_dict(tp=tp, tn=tn, fp=fp, fn=fn)
-            category: classifier_evaluation_dict(
-                tp=tp[category],
-                tn=tn[category],
-                fp=fp[category],
-                fn=fn[category],
-            )
-            for category in tp.keys()
-        }

miplearn/components/dynamic_lazy.py

@@ -15,7 +15,7 @@ from miplearn.components.component import Component
 from miplearn.components.dynamic_common import DynamicConstraintsComponent
 from miplearn.features.sample import Sample
 from miplearn.instance.base import Instance
-from miplearn.types import LearningSolveStats
+from miplearn.types import LearningSolveStats, ConstraintName, ConstraintCategory
 
 logger = logging.getLogger(__name__)
 
@@ -41,11 +41,11 @@ class DynamicLazyConstraintsComponent(Component):
         self.classifiers = self.dynamic.classifiers
         self.thresholds = self.dynamic.thresholds
         self.known_cids = self.dynamic.known_cids
-        self.lazy_enforced: Set[str] = set()
+        self.lazy_enforced: Set[ConstraintName] = set()
 
     @staticmethod
     def enforce(
-        cids: List[str],
+        cids: List[ConstraintName],
         instance: Instance,
         model: Any,
         solver: "LearningSolver",
@@ -117,7 +117,7 @@ class DynamicLazyConstraintsComponent(Component):
         self,
         instance: Instance,
         sample: Sample,
-    ) -> List[str]:
+    ) -> List[ConstraintName]:
         return self.dynamic.sample_predict(instance, sample)
 
     @overrides
@@ -127,8 +127,8 @@ class DynamicLazyConstraintsComponent(Component):
     @overrides
     def fit_xy(
         self,
-        x: Dict[str, np.ndarray],
+        x: Dict[ConstraintCategory, np.ndarray],
-        y: Dict[str, np.ndarray],
+        y: Dict[ConstraintCategory, np.ndarray],
     ) -> None:
         self.dynamic.fit_xy(x, y)
 
@@ -137,5 +137,5 @@ class DynamicLazyConstraintsComponent(Component):
         self,
         instance: Instance,
         sample: Sample,
-    ) -> Dict[str, Dict[str, float]]:
+    ) -> Dict[ConstraintCategory, Dict[str, float]]:
         return self.dynamic.sample_evaluate(instance, sample)

miplearn/components/dynamic_user_cuts.py

@@ -15,7 +15,7 @@ from miplearn.components.component import Component
 from miplearn.components.dynamic_common import DynamicConstraintsComponent
 from miplearn.features.sample import Sample
 from miplearn.instance.base import Instance
-from miplearn.types import LearningSolveStats
+from miplearn.types import LearningSolveStats, ConstraintName, ConstraintCategory
 
 logger = logging.getLogger(__name__)
 
@@ -34,7 +34,7 @@ class UserCutsComponent(Component):
             threshold=threshold,
             attr="mip_user_cuts_enforced",
         )
-        self.enforced: Set[str] = set()
+        self.enforced: Set[ConstraintName] = set()
         self.n_added_in_callback = 0
 
     @overrides
@@ -71,7 +71,7 @@ class UserCutsComponent(Component):
         for cid in cids:
             if cid in self.enforced:
                 continue
-            assert isinstance(cid, str)
+            assert isinstance(cid, ConstraintName)
             instance.enforce_user_cut(solver.internal_solver, model, cid)
             self.enforced.add(cid)
             self.n_added_in_callback += 1
@@ -110,7 +110,7 @@ class UserCutsComponent(Component):
         self,
         instance: "Instance",
         sample: Sample,
-    ) -> List[str]:
+    ) -> List[ConstraintName]:
         return self.dynamic.sample_predict(instance, sample)
 
     @overrides
@@ -120,8 +120,8 @@ class UserCutsComponent(Component):
     @overrides
     def fit_xy(
         self,
-        x: Dict[str, np.ndarray],
+        x: Dict[ConstraintCategory, np.ndarray],
-        y: Dict[str, np.ndarray],
+        y: Dict[ConstraintCategory, np.ndarray],
     ) -> None:
         self.dynamic.fit_xy(x, y)
 
@@ -130,5 +130,5 @@ class UserCutsComponent(Component):
         self,
         instance: "Instance",
         sample: Sample,
-    ) -> Dict[str, Dict[str, float]]:
+    ) -> Dict[ConstraintCategory, Dict[str, float]]:
         return self.dynamic.sample_evaluate(instance, sample)

miplearn/components/static_lazy.py

@@ -15,7 +15,7 @@ from miplearn.components.component import Component
 from miplearn.features.sample import Sample
 from miplearn.solvers.internal import Constraints
 from miplearn.instance.base import Instance
-from miplearn.types import LearningSolveStats
+from miplearn.types import LearningSolveStats, ConstraintName, ConstraintCategory
 
 logger = logging.getLogger(__name__)
 
@@ -24,7 +24,7 @@ if TYPE_CHECKING:
 
 
 class LazyConstraint:
-    def __init__(self, cid: str, obj: Any) -> None:
+    def __init__(self, cid: ConstraintName, obj: Any) -> None:
         self.cid = cid
         self.obj = obj
 
@@ -44,11 +44,11 @@ class StaticLazyConstraintsComponent(Component):
         assert isinstance(classifier, Classifier)
         self.classifier_prototype: Classifier = classifier
         self.threshold_prototype: Threshold = threshold
-        self.classifiers: Dict[str, Classifier] = {}
+        self.classifiers: Dict[ConstraintCategory, Classifier] = {}
-        self.thresholds: Dict[str, Threshold] = {}
+        self.thresholds: Dict[ConstraintCategory, Threshold] = {}
         self.pool: Constraints = Constraints()
         self.violation_tolerance: float = violation_tolerance
-        self.enforced_cids: Set[str] = set()
+        self.enforced_cids: Set[ConstraintName] = set()
         self.n_restored: int = 0
         self.n_iterations: int = 0
 
@@ -105,8 +105,8 @@ class StaticLazyConstraintsComponent(Component):
     @overrides
     def fit_xy(
         self,
-        x: Dict[str, np.ndarray],
+        x: Dict[ConstraintCategory, np.ndarray],
-        y: Dict[str, np.ndarray],
+        y: Dict[ConstraintCategory, np.ndarray],
     ) -> None:
         for c in y.keys():
             assert c in x
@@ -136,9 +136,9 @@ class StaticLazyConstraintsComponent(Component):
     ) -> None:
         self._check_and_add(solver)
 
-    def sample_predict(self, sample: Sample) -> List[str]:
+    def sample_predict(self, sample: Sample) -> List[ConstraintName]:
         x, y, cids = self._sample_xy_with_cids(sample)
-        enforced_cids: List[str] = []
+        enforced_cids: List[ConstraintName] = []
         for category in x.keys():
             if category not in self.classifiers:
                 continue
@@ -156,7 +156,10 @@ class StaticLazyConstraintsComponent(Component):
         self,
         _: Optional[Instance],
         sample: Sample,
-    ) -> Tuple[Dict[str, List[List[float]]], Dict[str, List[List[float]]]]:
+    ) -> Tuple[
+        Dict[ConstraintCategory, List[List[float]]],
+        Dict[ConstraintCategory, List[List[float]]],
+    ]:
         x, y, __ = self._sample_xy_with_cids(sample)
         return x, y
 
@@ -197,13 +200,13 @@ class StaticLazyConstraintsComponent(Component):
     def _sample_xy_with_cids(
         self, sample: Sample
     ) -> Tuple[
-        Dict[str, List[List[float]]],
+        Dict[ConstraintCategory, List[List[float]]],
-        Dict[str, List[List[float]]],
+        Dict[ConstraintCategory, List[List[float]]],
-        Dict[str, List[str]],
+        Dict[ConstraintCategory, List[ConstraintName]],
     ]:
-        x: Dict[str, List[List[float]]] = {}
+        x: Dict[ConstraintCategory, List[List[float]]] = {}
-        y: Dict[str, List[List[float]]] = {}
+        y: Dict[ConstraintCategory, List[List[float]]] = {}
-        cids: Dict[str, List[str]] = {}
+        cids: Dict[ConstraintCategory, List[ConstraintName]] = {}
         instance_features = sample.get_vector("static_instance_features")
         constr_features = sample.get_vector_list("lp_constr_features")
         constr_names = sample.get_array("static_constr_names")

miplearn/features/extractor.py

@@ -2,10 +2,8 @@
 #  Copyright (C) 2020-2021, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 
-import collections
-import numbers
 from math import log, isfinite
-from typing import TYPE_CHECKING, Dict, Optional, List, Any, Tuple, KeysView, cast
+from typing import TYPE_CHECKING, List, Tuple
 
 import numpy as np
 
@@ -34,6 +32,7 @@ class FeaturesExtractor:
     ) -> None:
         variables = solver.get_variables(with_static=True)
         constraints = solver.get_constraints(with_static=True, with_lhs=self.with_lhs)
+        assert constraints.names is not None
         sample.put_array("static_var_lower_bounds", variables.lower_bounds)
         sample.put_array("static_var_names", variables.names)
         sample.put_array("static_var_obj_coeffs", variables.obj_coeffs)
@@ -43,15 +42,30 @@ class FeaturesExtractor:
         # sample.put("static_constr_lhs", constraints.lhs)
         sample.put_array("static_constr_rhs", constraints.rhs)
         sample.put_array("static_constr_senses", constraints.senses)
-        vars_features_user, var_categories = self._extract_user_features_vars(
+
-            instance, sample
+        # Instance features
-        )
-        sample.put_array("static_var_categories", var_categories)
-        self._extract_user_features_constrs(instance, sample)
         self._extract_user_features_instance(instance, sample)
-        alw17 = self._extract_var_features_AlvLouWeh2017(sample)
 
-        # Build static_var_features
+        # Constraint features
+        (
+            constr_features,
+            constr_categories,
+            constr_lazy,
+        ) = FeaturesExtractor._extract_user_features_constrs(
+            instance,
+            constraints.names,
+        )
+        sample.put_array("static_constr_features", constr_features)
+        sample.put_array("static_constr_categories", constr_categories)
+        sample.put_array("static_constr_lazy", constr_lazy)
+        sample.put_scalar("static_constr_lazy_count", int(constr_lazy.sum()))
+
+        # Variable features
+        (
+            vars_features_user,
+            var_categories,
+        ) = self._extract_user_features_vars(instance, sample)
+        sample.put_array("static_var_categories", var_categories)
         assert variables.lower_bounds is not None
         assert variables.obj_coeffs is not None
         assert variables.upper_bounds is not None
@@ -60,7 +74,7 @@ class FeaturesExtractor:
             np.hstack(
                 [
                     vars_features_user,
-                    alw17,
+                    self._extract_var_features_AlvLouWeh2017(sample),
                     variables.lower_bounds.reshape(-1, 1),
                     variables.obj_coeffs.reshape(-1, 1),
                     variables.upper_bounds.reshape(-1, 1),
@@ -92,13 +106,12 @@ class FeaturesExtractor:
         sample.put_array("lp_constr_sa_rhs_down", constraints.sa_rhs_down)
         sample.put_array("lp_constr_sa_rhs_up", constraints.sa_rhs_up)
         sample.put_array("lp_constr_slacks", constraints.slacks)
-        alw17 = self._extract_var_features_AlvLouWeh2017(sample)
 
-        # Build lp_var_features
+        # Variable features
         lp_var_features_list = []
         for f in [
             sample.get_array("static_var_features"),
-            alw17,
+            self._extract_var_features_AlvLouWeh2017(sample),
         ]:
             if f is not None:
                 lp_var_features_list.append(f)
@@ -116,18 +129,20 @@ class FeaturesExtractor:
                 lp_var_features_list.append(f.reshape(-1, 1))
         sample.put_array("lp_var_features", np.hstack(lp_var_features_list))
 
-        sample.put_vector_list(
+        # Constraint features
-            "lp_constr_features",
+        lp_constr_features_list = []
-            self._combine(
+        for f in [sample.get_array("static_constr_features")]:
-                [
+            if f is not None:
-                    sample.get_vector_list("static_constr_features"),
+                lp_constr_features_list.append(f)
-                    sample.get_array("lp_constr_dual_values"),
+        for f in [
-                    sample.get_array("lp_constr_sa_rhs_down"),
+            sample.get_array("lp_constr_dual_values"),
-                    sample.get_array("lp_constr_sa_rhs_up"),
+            sample.get_array("lp_constr_sa_rhs_down"),
-                    sample.get_array("lp_constr_slacks"),
+            sample.get_array("lp_constr_sa_rhs_up"),
-                ],
+            sample.get_array("lp_constr_slacks"),
-            ),
+        ]:
-        )
+            if f is not None:
+                lp_constr_features_list.append(f.reshape(-1, 1))
+        sample.put_array("lp_constr_features", np.hstack(lp_constr_features_list))
 
         # Build lp_instance_features
         static_instance_features = sample.get_array("static_instance_features")
@@ -155,6 +170,7 @@ class FeaturesExtractor:
         sample.put_array("mip_var_values", variables.values)
         sample.put_array("mip_constr_slacks", constraints.slacks)
 
+    # noinspection DuplicatedCode
     def _extract_user_features_vars(
         self,
         instance: "Instance",
@@ -180,7 +196,7 @@ class FeaturesExtractor:
         )
         assert var_features.dtype.kind in ["f"], (
             f"Variable features must be floating point numbers. "
-            f"Found dtype: {var_features.dtype} instead."
+            f"Found {var_features.dtype} instead."
         )
 
         # Query variable categories
@@ -195,7 +211,7 @@ class FeaturesExtractor:
         )
         assert len(var_categories) == len(var_names), (
             f"Variable categories must have exactly {len(var_names)} elements. "
-            f"Found {var_features.shape[0]} elements instead."
+            f"Found {var_categories.shape[0]} elements instead."
         )
         assert var_categories.dtype.kind == "S", (
             f"Variable categories must be a numpy array with dtype='S'. "
@@ -203,58 +219,71 @@ class FeaturesExtractor:
         )
         return var_features, var_categories
 
+    # noinspection DuplicatedCode
+    @classmethod
     def _extract_user_features_constrs(
-        self,
+        cls,
         instance: "Instance",
-        sample: Sample,
+        constr_names: np.ndarray,
-    ) -> None:
+    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
-        has_static_lazy = instance.has_static_lazy_constraints()
+        # Query constraint features
-        user_features: List[Optional[List[float]]] = []
+        constr_features = instance.get_constraint_features(constr_names)
-        categories: List[Optional[bytes]] = []
+        assert isinstance(constr_features, np.ndarray), (
-        lazy: List[bool] = []
+            f"get_constraint_features must return a numpy array. "
-        constr_categories_dict = instance.get_constraint_categories()
+            f"Found {constr_features.__class__} instead."
-        constr_features_dict = instance.get_constraint_features()
+        )
-        constr_names = sample.get_array("static_constr_names")
+        assert len(constr_features.shape) == 2, (
-        assert constr_names is not None
+            f"get_constraint_features must return a 2-dimensional array. "
-
+            f"Found array with shape {constr_features.shape} instead."
-        for (cidx, cname) in enumerate(constr_names):
+        )
-            category: Optional[str] = cname
+        assert constr_features.shape[0] == len(constr_names), (
-            if cname in constr_categories_dict:
+            f"get_constraint_features must return an array with {len(constr_names)} "
-                category = constr_categories_dict[cname]
+            f"rows. Found {constr_features.shape[0]} rows instead."
-            if category is None:
+        )
-                user_features.append(None)
+        assert constr_features.dtype.kind in ["f"], (
-                categories.append(None)
+            f"get_constraint_features must return floating point numbers. "
-                continue
+            f"Found {constr_features.dtype} instead."
-            assert isinstance(category, bytes), (
+        )
-                f"Constraint category must be bytes. "
+
-                f"Found {type(category).__name__} instead for cname={cname}.",
+        # Query constraint categories
-            )
+        constr_categories = instance.get_constraint_categories(constr_names)
-            categories.append(category)
+        assert isinstance(constr_categories, np.ndarray), (
-            cf: Optional[List[float]] = None
+            f"get_constraint_categories must return a numpy array. "
-            if cname in constr_features_dict:
+            f"Found {constr_categories.__class__} instead."
-                cf = constr_features_dict[cname]
+        )
-                if isinstance(cf, np.ndarray):
+        assert len(constr_categories.shape) == 1, (
-                    cf = cf.tolist()
+            f"get_constraint_categories must return a vector. "
-                assert isinstance(cf, list), (
+            f"Found array with shape {constr_categories.shape} instead."
-                    f"Constraint features must be a list. "
+        )
-                    f"Found {type(cf).__name__} instead for cname={cname}."
+        assert len(constr_categories) == len(constr_names), (
-                )
+            f"get_constraint_categories must return a vector with {len(constr_names)} "
-                for f in cf:
+            f"elements. Found {constr_categories.shape[0]} elements instead."
-                    assert isinstance(f, numbers.Real), (
+        )
-                        f"Constraint features must be a list of numbers. "
+        assert constr_categories.dtype.kind == "S", (
-                        f"Found {type(f).__name__} instead for cname={cname}."
+            f"get_constraint_categories must return a numpy array with dtype='S'. "
-                    )
+            f"Found {constr_categories.dtype} instead."
-                cf = list(cf)
+        )
-            user_features.append(cf)
+
-            if has_static_lazy:
+        # Query constraint lazy attribute
-                lazy.append(instance.is_constraint_lazy(cname))
+        constr_lazy = instance.are_constraints_lazy(constr_names)
-            else:
+        assert isinstance(constr_lazy, np.ndarray), (
-                lazy.append(False)
+            f"are_constraints_lazy must return a numpy array. "
-        sample.put_vector_list("static_constr_features", user_features)
+            f"Found {constr_lazy.__class__} instead."
-        sample.put_array("static_constr_categories", np.array(categories, dtype="S"))
+        )
-        constr_lazy = np.array(lazy, dtype=bool)
+        assert len(constr_lazy.shape) == 1, (
-        sample.put_array("static_constr_lazy", constr_lazy)
+            f"are_constraints_lazy must return a vector. "
-        sample.put_scalar("static_constr_lazy_count", int(constr_lazy.sum()))
+            f"Found array with shape {constr_lazy.shape} instead."
+        )
+        assert constr_lazy.shape[0] == len(constr_names), (
+            f"are_constraints_lazy must return a vector with {len(constr_names)} "
+            f"elements. Found {constr_lazy.shape[0]} elements instead."
+        )
+        assert constr_lazy.dtype.kind == "b", (
+            f"are_constraints_lazy must return a boolean array. "
+            f"Found {constr_lazy.dtype} instead."
+        )
+
+        return constr_features, constr_categories, constr_lazy
 
     def _extract_user_features_instance(
         self,
@@ -272,7 +301,7 @@ class FeaturesExtractor:
         )
         assert features.dtype.kind in [
             "f"
-        ], f"Instance features have unsupported dtype: {features.dtype}"
+        ], f"Instance features have unsupported {features.dtype}"
         sample.put_array("static_instance_features", features)
 
     # Alvarez, A. M., Louveaux, Q., & Wehenkel, L. (2017). A machine learning-based
@@ -352,29 +381,3 @@ class FeaturesExtractor:
 
             features.append(f)
         return np.array(features, dtype=float)
-
-    def _combine(
-        self,
-        items: List,
-    ) -> List[List[float]]:
-        combined: List[List[float]] = []
-        for series in items:
-            if series is None:
-                continue
-            if len(combined) == 0:
-                for i in range(len(series)):
-                    combined.append([])
-            for (i, s) in enumerate(series):
-                if s is None:
-                    continue
-                elif isinstance(s, list):
-                    combined[i].extend([_clip(sj) for sj in s])
-                else:
-                    combined[i].append(_clip(s))
-        return combined
-
-
-def _clip(vi: float) -> float:
-    if not isfinite(vi):
-        return max(min(vi, 1e20), -1e20)
-    return vi

miplearn/features/sample.py

@@ -103,7 +103,7 @@ class Sample(ABC):
     def _assert_is_scalar(self, value: Any) -> None:
         if value is None:
             return
-        if isinstance(value, (str, bool, int, float, np.bytes_)):
+        if isinstance(value, (str, bool, int, float, bytes, np.bytes_)):
             return
         assert False, f"scalar expected; found instead: {value} ({value.__class__})"
 

miplearn/instance/base.py

@@ -9,6 +9,7 @@ from typing import Any, List, TYPE_CHECKING, Dict
 import numpy as np
 
 from miplearn.features.sample import Sample, MemorySample
+from miplearn.types import ConstraintName, ConstraintCategory
 
 logger = logging.getLogger(__name__)
 
@@ -97,26 +98,23 @@ class Instance(ABC):
         """
         return names
 
-    def get_constraint_features(self) -> Dict[str, List[float]]:
+    def get_constraint_features(self, names: np.ndarray) -> np.ndarray:
-        return {}
+        return np.zeros((len(names), 1))
-
-    def get_constraint_categories(self) -> Dict[str, str]:
-        return {}
 
-    def has_static_lazy_constraints(self) -> bool:
+    def get_constraint_categories(self, names: np.ndarray) -> np.ndarray:
-        return False
+        return names
 
     def has_dynamic_lazy_constraints(self) -> bool:
         return False
 
-    def is_constraint_lazy(self, cid: str) -> bool:
+    def are_constraints_lazy(self, names: np.ndarray) -> np.ndarray:
-        return False
+        return np.zeros(len(names), dtype=bool)
 
     def find_violated_lazy_constraints(
         self,
         solver: "InternalSolver",
         model: Any,
-    ) -> List[str]:
+    ) -> List[ConstraintName]:
         """
         Returns lazy constraint violations found for the current solution.
 
@@ -142,7 +140,7 @@ class Instance(ABC):
         self,
         solver: "InternalSolver",
         model: Any,
-        violation: str,
+        violation: ConstraintName,
     ) -> None:
         """
         Adds constraints to the model to ensure that the given violation is fixed.
@@ -168,14 +166,14 @@ class Instance(ABC):
     def has_user_cuts(self) -> bool:
         return False
 
-    def find_violated_user_cuts(self, model: Any) -> List[str]:
+    def find_violated_user_cuts(self, model: Any) -> List[ConstraintName]:
         return []
 
     def enforce_user_cut(
         self,
         solver: "InternalSolver",
         model: Any,
-        violation: str,
+        violation: ConstraintName,
     ) -> Any:
         return None
 

miplearn/instance/file.py

@@ -11,6 +11,7 @@ from overrides import overrides
 
 from miplearn.features.sample import Hdf5Sample, Sample
 from miplearn.instance.base import Instance
+from miplearn.types import ConstraintName, ConstraintCategory
 
 if TYPE_CHECKING:
     from miplearn.solvers.learning import InternalSolver
@@ -46,19 +47,14 @@ class FileInstance(Instance):
         return self.instance.get_variable_categories(names)
 
     @overrides
-    def get_constraint_features(self) -> Dict[str, List[float]]:
+    def get_constraint_features(self, names: np.ndarray) -> np.ndarray:
         assert self.instance is not None
-        return self.instance.get_constraint_features()
+        return self.instance.get_constraint_features(names)
 
     @overrides
-    def get_constraint_categories(self) -> Dict[str, str]:
+    def get_constraint_categories(self, names: np.ndarray) -> np.ndarray:
         assert self.instance is not None
-        return self.instance.get_constraint_categories()
+        return self.instance.get_constraint_categories(names)
-
-    @overrides
-    def has_static_lazy_constraints(self) -> bool:
-        assert self.instance is not None
-        return self.instance.has_static_lazy_constraints()
 
     @overrides
     def has_dynamic_lazy_constraints(self) -> bool:
@@ -66,16 +62,16 @@ class FileInstance(Instance):
         return self.instance.has_dynamic_lazy_constraints()
 
     @overrides
-    def is_constraint_lazy(self, cid: str) -> bool:
+    def are_constraints_lazy(self, names: np.ndarray) -> np.ndarray:
         assert self.instance is not None
-        return self.instance.is_constraint_lazy(cid)
+        return self.instance.are_constraints_lazy(names)
 
     @overrides
     def find_violated_lazy_constraints(
         self,
         solver: "InternalSolver",
         model: Any,
-    ) -> List[str]:
+    ) -> List[ConstraintName]:
         assert self.instance is not None
         return self.instance.find_violated_lazy_constraints(solver, model)
 
@@ -84,13 +80,13 @@ class FileInstance(Instance):
         self,
         solver: "InternalSolver",
         model: Any,
-        violation: str,
+        violation: ConstraintName,
     ) -> None:
         assert self.instance is not None
         self.instance.enforce_lazy_constraint(solver, model, violation)
 
     @overrides
-    def find_violated_user_cuts(self, model: Any) -> List[str]:
+    def find_violated_user_cuts(self, model: Any) -> List[ConstraintName]:
         assert self.instance is not None
         return self.instance.find_violated_user_cuts(model)
 
@@ -99,7 +95,7 @@ class FileInstance(Instance):
         self,
         solver: "InternalSolver",
         model: Any,
-        violation: str,
+        violation: ConstraintName,
     ) -> None:
         assert self.instance is not None
         self.instance.enforce_user_cut(solver, model, violation)

miplearn/instance/picklegz.py

@@ -13,6 +13,7 @@ from overrides import overrides
 
 from miplearn.features.sample import Sample
 from miplearn.instance.base import Instance
+from miplearn.types import ConstraintName, ConstraintCategory
 
 if TYPE_CHECKING:
     from miplearn.solvers.learning import InternalSolver
@@ -58,19 +59,14 @@ class PickleGzInstance(Instance):
         return self.instance.get_variable_categories(names)
 
     @overrides
-    def get_constraint_features(self) -> Dict[str, List[float]]:
+    def get_constraint_features(self, names: np.ndarray) -> np.ndarray:
         assert self.instance is not None
-        return self.instance.get_constraint_features()
+        return self.instance.get_constraint_features(names)
 
     @overrides
-    def get_constraint_categories(self) -> Dict[str, str]:
+    def get_constraint_categories(self, names: np.ndarray) -> np.ndarray:
         assert self.instance is not None
-        return self.instance.get_constraint_categories()
+        return self.instance.get_constraint_categories(names)
-
-    @overrides
-    def has_static_lazy_constraints(self) -> bool:
-        assert self.instance is not None
-        return self.instance.has_static_lazy_constraints()
 
     @overrides
     def has_dynamic_lazy_constraints(self) -> bool:
@@ -78,16 +74,16 @@ class PickleGzInstance(Instance):
         return self.instance.has_dynamic_lazy_constraints()
 
     @overrides
-    def is_constraint_lazy(self, cid: str) -> bool:
+    def are_constraints_lazy(self, names: np.ndarray) -> np.ndarray:
         assert self.instance is not None
-        return self.instance.is_constraint_lazy(cid)
+        return self.instance.are_constraints_lazy(names)
 
     @overrides
     def find_violated_lazy_constraints(
         self,
         solver: "InternalSolver",
         model: Any,
-    ) -> List[str]:
+    ) -> List[ConstraintName]:
         assert self.instance is not None
         return self.instance.find_violated_lazy_constraints(solver, model)
 
@@ -96,13 +92,13 @@ class PickleGzInstance(Instance):
         self,
         solver: "InternalSolver",
         model: Any,
-        violation: str,
+        violation: ConstraintName,
     ) -> None:
         assert self.instance is not None
         self.instance.enforce_lazy_constraint(solver, model, violation)
 
     @overrides
-    def find_violated_user_cuts(self, model: Any) -> List[str]:
+    def find_violated_user_cuts(self, model: Any) -> List[ConstraintName]:
         assert self.instance is not None
         return self.instance.find_violated_user_cuts(model)
 
@@ -111,7 +107,7 @@ class PickleGzInstance(Instance):
         self,
         solver: "InternalSolver",
         model: Any,
-        violation: str,
+        violation: ConstraintName,
     ) -> None:
         assert self.instance is not None
         self.instance.enforce_user_cut(solver, model, violation)

miplearn/problems/tsp.py

@@ -14,6 +14,7 @@ from scipy.stats.distributions import rv_frozen
 from miplearn.instance.base import Instance
 from miplearn.solvers.learning import InternalSolver
 from miplearn.solvers.pyomo.base import BasePyomoSolver
+from miplearn.types import ConstraintName
 
 
 class ChallengeA:
@@ -85,14 +86,14 @@ class TravelingSalesmanInstance(Instance):
         self,
         solver: InternalSolver,
         model: Any,
-    ) -> List[str]:
+    ) -> List[ConstraintName]:
         selected_edges = [e for e in self.edges if model.x[e].value > 0.5]
         graph = nx.Graph()
         graph.add_edges_from(selected_edges)
         violations = []
         for c in list(nx.connected_components(graph)):
             if len(c) < self.n_cities:
-                violations.append(",".join(map(str, c)))
+                violations.append(",".join(map(str, c)).encode())
         return violations
 
     @overrides
@@ -100,10 +101,10 @@ class TravelingSalesmanInstance(Instance):
         self,
         solver: InternalSolver,
         model: Any,
-        violation: str,
+        violation: ConstraintName,
     ) -> None:
         assert isinstance(solver, BasePyomoSolver)
-        component = [int(v) for v in violation.split(",")]
+        component = [int(v) for v in violation.decode().split(",")]
         cut_edges = [
             e
             for e in self.edges

miplearn/solvers/tests/__init__.py

@@ -260,7 +260,7 @@ def run_lazy_cb_tests(solver: InternalSolver) -> None:
         assert relsol is not None
         assert relsol[b"x[0]"] is not None
         if relsol[b"x[0]"] > 0:
-            instance.enforce_lazy_constraint(cb_solver, cb_model, "cut")
+            instance.enforce_lazy_constraint(cb_solver, cb_model, b"cut")
 
     solver.set_instance(instance, model)
     solver.solve(lazy_cb=lazy_cb)

miplearn/types.py

@@ -11,6 +11,8 @@ if TYPE_CHECKING:
     from miplearn.solvers.learning import InternalSolver
 
 Category = bytes
+ConstraintName = bytes
+ConstraintCategory = bytes
 IterationCallback = Callable[[], bool]
 LazyCallback = Callable[[Any, Any], None]
 SolverParams = Dict[str, Any]

tests/components/test_dynamic_lazy.py

@@ -11,7 +11,7 @@ from miplearn.classifiers import Classifier
 from miplearn.classifiers.threshold import MinProbabilityThreshold
 from miplearn.components import classifier_evaluation_dict
 from miplearn.components.dynamic_lazy import DynamicLazyConstraintsComponent
-from miplearn.features.sample import Sample, MemorySample
+from miplearn.features.sample import MemorySample
 from miplearn.instance.base import Instance
 from miplearn.solvers.tests import assert_equals
 
@@ -24,71 +24,71 @@ def training_instances() -> List[Instance]:
     samples_0 = [
         MemorySample(
             {
-                "mip_constr_lazy_enforced": {"c1", "c2"},
+                "mip_constr_lazy_enforced": {b"c1", b"c2"},
-                "static_instance_features": [5.0],
+                "static_instance_features": np.array([5.0]),
             },
         ),
         MemorySample(
             {
-                "mip_constr_lazy_enforced": {"c2", "c3"},
+                "mip_constr_lazy_enforced": {b"c2", b"c3"},
-                "static_instance_features": [5.0],
+                "static_instance_features": np.array([5.0]),
             },
         ),
     ]
     instances[0].get_samples = Mock(return_value=samples_0)  # type: ignore
     instances[0].get_constraint_categories = Mock(  # type: ignore
-        return_value={
+        return_value=np.array(["type-a", "type-a", "type-b", "type-b"], dtype="S")
-            "c1": "type-a",
-            "c2": "type-a",
-            "c3": "type-b",
-            "c4": "type-b",
-        }
     )
     instances[0].get_constraint_features = Mock(  # type: ignore
-        return_value={
+        return_value=np.array(
-            "c1": [1.0, 2.0, 3.0],
+            [
-            "c2": [4.0, 5.0, 6.0],
+                [1.0, 2.0, 3.0],
-            "c3": [1.0, 2.0],
+                [4.0, 5.0, 6.0],
-            "c4": [3.0, 4.0],
+                [1.0, 2.0, 0.0],
-        }
+                [3.0, 4.0, 0.0],
+            ]
+        )
+    )
+    instances[0].are_constraints_lazy = Mock(  # type: ignore
+        return_value=np.zeros(4, dtype=bool)
     )
     samples_1 = [
         MemorySample(
             {
-                "mip_constr_lazy_enforced": {"c3", "c4"},
+                "mip_constr_lazy_enforced": {b"c3", b"c4"},
-                "static_instance_features": [8.0],
+                "static_instance_features": np.array([8.0]),
             },
         )
     ]
     instances[1].get_samples = Mock(return_value=samples_1)  # type: ignore
     instances[1].get_constraint_categories = Mock(  # type: ignore
-        return_value={
+        return_value=np.array(["", "type-a", "type-b", "type-b"], dtype="S")
-            "c1": None,
-            "c2": "type-a",
-            "c3": "type-b",
-            "c4": "type-b",
-        }
     )
     instances[1].get_constraint_features = Mock(  # type: ignore
-        return_value={
+        return_value=np.array(
-            "c2": [7.0, 8.0, 9.0],
+            [
-            "c3": [5.0, 6.0],
+                [7.0, 8.0, 9.0],
-            "c4": [7.0, 8.0],
+                [5.0, 6.0, 0.0],
-        }
+                [7.0, 8.0, 0.0],
+            ]
+        )
+    )
+    instances[1].are_constraints_lazy = Mock(  # type: ignore
+        return_value=np.zeros(4, dtype=bool)
     )
     return instances
 
 
 def test_sample_xy(training_instances: List[Instance]) -> None:
     comp = DynamicLazyConstraintsComponent()
-    comp.pre_fit([{"c1", "c2", "c3", "c4"}])
+    comp.pre_fit([{b"c1", b"c2", b"c3", b"c4"}])
     x_expected = {
-        "type-a": [[5.0, 1.0, 2.0, 3.0], [5.0, 4.0, 5.0, 6.0]],
+        b"type-a": np.array([[5.0, 1.0, 2.0, 3.0], [5.0, 4.0, 5.0, 6.0]]),
-        "type-b": [[5.0, 1.0, 2.0], [5.0, 3.0, 4.0]],
+        b"type-b": np.array([[5.0, 1.0, 2.0, 0.0], [5.0, 3.0, 4.0, 0.0]]),
     }
     y_expected = {
-        "type-a": [[False, True], [False, True]],
+        b"type-a": np.array([[False, True], [False, True]]),
-        "type-b": [[True, False], [True, False]],
+        b"type-b": np.array([[True, False], [True, False]]),
     }
     x_actual, y_actual = comp.sample_xy(
         training_instances[0],
@@ -98,95 +98,26 @@ def test_sample_xy(training_instances: List[Instance]) -> None:
     assert_equals(y_actual, y_expected)
 
 
-# def test_fit(training_instances: List[Instance]) -> None:
-#     clf = Mock(spec=Classifier)
-#     clf.clone = Mock(side_effect=lambda: Mock(spec=Classifier))
-#     comp = DynamicLazyConstraintsComponent(classifier=clf)
-#     comp.fit(training_instances)
-#     assert clf.clone.call_count == 2
-#
-#     assert "type-a" in comp.classifiers
-#     clf_a = comp.classifiers["type-a"]
-#     assert clf_a.fit.call_count == 1  # type: ignore
-#     assert_array_equal(
-#         clf_a.fit.call_args[0][0],  # type: ignore
-#         np.array(
-#             [
-#                 [5.0, 1.0, 2.0, 3.0],
-#                 [5.0, 4.0, 5.0, 6.0],
-#                 [5.0, 1.0, 2.0, 3.0],
-#                 [5.0, 4.0, 5.0, 6.0],
-#                 [8.0, 7.0, 8.0, 9.0],
-#             ]
-#         ),
-#     )
-#     assert_array_equal(
-#         clf_a.fit.call_args[0][1],  # type: ignore
-#         np.array(
-#             [
-#                 [False, True],
-#                 [False, True],
-#                 [True, False],
-#                 [False, True],
-#                 [True, False],
-#             ]
-#         ),
-#     )
-#
-#     assert "type-b" in comp.classifiers
-#     clf_b = comp.classifiers["type-b"]
-#     assert clf_b.fit.call_count == 1  # type: ignore
-#     assert_array_equal(
-#         clf_b.fit.call_args[0][0],  # type: ignore
-#         np.array(
-#             [
-#                 [5.0, 1.0, 2.0],
-#                 [5.0, 3.0, 4.0],
-#                 [5.0, 1.0, 2.0],
-#                 [5.0, 3.0, 4.0],
-#                 [8.0, 5.0, 6.0],
-#                 [8.0, 7.0, 8.0],
-#             ]
-#         ),
-#     )
-#     assert_array_equal(
-#         clf_b.fit.call_args[0][1],  # type: ignore
-#         np.array(
-#             [
-#                 [True, False],
-#                 [True, False],
-#                 [False, True],
-#                 [True, False],
-#                 [False, True],
-#                 [False, True],
-#             ]
-#         ),
-#     )
-
-
 def test_sample_predict_evaluate(training_instances: List[Instance]) -> None:
     comp = DynamicLazyConstraintsComponent()
-    comp.known_cids.extend(["c1", "c2", "c3", "c4"])
+    comp.known_cids.extend([b"c1", b"c2", b"c3", b"c4"])
-    comp.thresholds["type-a"] = MinProbabilityThreshold([0.5, 0.5])
+    comp.thresholds[b"type-a"] = MinProbabilityThreshold([0.5, 0.5])
-    comp.thresholds["type-b"] = MinProbabilityThreshold([0.5, 0.5])
+    comp.thresholds[b"type-b"] = MinProbabilityThreshold([0.5, 0.5])
-    comp.classifiers["type-a"] = Mock(spec=Classifier)
+    comp.classifiers[b"type-a"] = Mock(spec=Classifier)
-    comp.classifiers["type-b"] = Mock(spec=Classifier)
+    comp.classifiers[b"type-b"] = Mock(spec=Classifier)
-    comp.classifiers["type-a"].predict_proba = Mock(  # type: ignore
+    comp.classifiers[b"type-a"].predict_proba = Mock(  # type: ignore
         side_effect=lambda _: np.array([[0.1, 0.9], [0.8, 0.2]])
     )
-    comp.classifiers["type-b"].predict_proba = Mock(  # type: ignore
+    comp.classifiers[b"type-b"].predict_proba = Mock(  # type: ignore
         side_effect=lambda _: np.array([[0.9, 0.1], [0.1, 0.9]])
     )
     pred = comp.sample_predict(
         training_instances[0],
         training_instances[0].get_samples()[0],
     )
-    assert pred == ["c1", "c4"]
+    assert pred == [b"c1", b"c4"]
     ev = comp.sample_evaluate(
         training_instances[0],
         training_instances[0].get_samples()[0],
     )
-    assert ev == {
+    assert ev == classifier_evaluation_dict(tp=1, fp=1, tn=1, fn=1)
-        "type-a": classifier_evaluation_dict(tp=1, fp=0, tn=0, fn=1),
-        "type-b": classifier_evaluation_dict(tp=0, fp=1, tn=1, fn=0),
-    }

tests/components/test_dynamic_user_cuts.py

@@ -17,6 +17,7 @@ from miplearn.components.dynamic_user_cuts import UserCutsComponent
 from miplearn.instance.base import Instance
 from miplearn.solvers.gurobi import GurobiSolver
 from miplearn.solvers.learning import LearningSolver
+from miplearn.types import ConstraintName, ConstraintCategory
 
 logger = logging.getLogger(__name__)
 
@@ -40,13 +41,13 @@ class GurobiStableSetProblem(Instance):
         return True
 
     @overrides
-    def find_violated_user_cuts(self, model: Any) -> List[str]:
+    def find_violated_user_cuts(self, model: Any) -> List[ConstraintName]:
         assert isinstance(model, gp.Model)
         vals = model.cbGetNodeRel(model.getVars())
         violations = []
         for clique in nx.find_cliques(self.graph):
             if sum(vals[i] for i in clique) > 1:
-                violations.append(",".join([str(i) for i in clique]))
+                violations.append(",".join([str(i) for i in clique]).encode())
         return violations
 
     @overrides
@@ -54,9 +55,9 @@ class GurobiStableSetProblem(Instance):
         self,
         solver: InternalSolver,
         model: Any,
-        cid: str,
+        cid: ConstraintName,
     ) -> Any:
-        clique = [int(i) for i in cid.split(",")]
+        clique = [int(i) for i in cid.decode().split(",")]
         x = model.getVars()
         model.addConstr(gp.quicksum([x[i] for i in clique]) <= 1)
 

tests/components/test_static_lazy.py

@@ -17,6 +17,7 @@ from miplearn.solvers.internal import InternalSolver, Constraints
 from miplearn.solvers.learning import LearningSolver
 from miplearn.types import (
     LearningSolveStats,
+    ConstraintCategory,
 )
 
 
@@ -25,11 +26,11 @@ def sample() -> Sample:
     sample = MemorySample(
         {
             "static_constr_categories": [
-                "type-a",
+                b"type-a",
-                "type-a",
+                b"type-a",
-                "type-a",
+                b"type-a",
-                "type-b",
+                b"type-b",
-                "type-b",
+                b"type-b",
             ],
             "static_constr_lazy": np.array([True, True, True, True, False]),
             "static_constr_names": np.array(["c1", "c2", "c3", "c4", "c5"], dtype="S"),
@@ -68,13 +69,13 @@ def test_usage_with_solver(instance: Instance) -> None:
     )
 
     component = StaticLazyConstraintsComponent(violation_tolerance=1.0)
-    component.thresholds["type-a"] = MinProbabilityThreshold([0.5, 0.5])
+    component.thresholds[b"type-a"] = MinProbabilityThreshold([0.5, 0.5])
-    component.thresholds["type-b"] = MinProbabilityThreshold([0.5, 0.5])
+    component.thresholds[b"type-b"] = MinProbabilityThreshold([0.5, 0.5])
     component.classifiers = {
-        "type-a": Mock(spec=Classifier),
+        b"type-a": Mock(spec=Classifier),
-        "type-b": Mock(spec=Classifier),
+        b"type-b": Mock(spec=Classifier),
     }
-    component.classifiers["type-a"].predict_proba = Mock(  # type: ignore
+    component.classifiers[b"type-a"].predict_proba = Mock(  # type: ignore
         return_value=np.array(
             [
                 [0.00, 1.00],  # c1
@@ -83,7 +84,7 @@ def test_usage_with_solver(instance: Instance) -> None:
             ]
         )
     )
-    component.classifiers["type-b"].predict_proba = Mock(  # type: ignore
+    component.classifiers[b"type-b"].predict_proba = Mock(  # type: ignore
         return_value=np.array(
             [
                 [0.02, 0.98],  # c4
@@ -105,8 +106,8 @@ def test_usage_with_solver(instance: Instance) -> None:
     )
 
     # Should ask ML to predict whether each lazy constraint should be enforced
-    component.classifiers["type-a"].predict_proba.assert_called_once()
+    component.classifiers[b"type-a"].predict_proba.assert_called_once()
-    component.classifiers["type-b"].predict_proba.assert_called_once()
+    component.classifiers[b"type-b"].predict_proba.assert_called_once()
 
     # Should ask internal solver to remove some constraints
     assert internal.remove_constraints.call_count == 1
@@ -153,18 +154,18 @@ def test_usage_with_solver(instance: Instance) -> None:
 
 def test_sample_predict(sample: Sample) -> None:
     comp = StaticLazyConstraintsComponent()
-    comp.thresholds["type-a"] = MinProbabilityThreshold([0.5, 0.5])
+    comp.thresholds[b"type-a"] = MinProbabilityThreshold([0.5, 0.5])
-    comp.thresholds["type-b"] = MinProbabilityThreshold([0.5, 0.5])
+    comp.thresholds[b"type-b"] = MinProbabilityThreshold([0.5, 0.5])
-    comp.classifiers["type-a"] = Mock(spec=Classifier)
+    comp.classifiers[b"type-a"] = Mock(spec=Classifier)
-    comp.classifiers["type-a"].predict_proba = lambda _: np.array(  # type:ignore
+    comp.classifiers[b"type-a"].predict_proba = lambda _: np.array(  # type:ignore
         [
             [0.0, 1.0],  # c1
             [0.0, 0.9],  # c2
             [0.9, 0.1],  # c3
         ]
     )
-    comp.classifiers["type-b"] = Mock(spec=Classifier)
+    comp.classifiers[b"type-b"] = Mock(spec=Classifier)
-    comp.classifiers["type-b"].predict_proba = lambda _: np.array(  # type:ignore
+    comp.classifiers[b"type-b"].predict_proba = lambda _: np.array(  # type:ignore
         [
             [0.0, 1.0],  # c4
         ]
@@ -175,17 +176,17 @@ def test_sample_predict(sample: Sample) -> None:
 
 def test_fit_xy() -> None:
     x = cast(
-        Dict[str, np.ndarray],
+        Dict[ConstraintCategory, np.ndarray],
         {
-            "type-a": np.array([[1.0, 1.0], [1.0, 2.0], [1.0, 3.0]]),
+            b"type-a": np.array([[1.0, 1.0], [1.0, 2.0], [1.0, 3.0]]),
-            "type-b": np.array([[1.0, 4.0, 0.0]]),
+            b"type-b": np.array([[1.0, 4.0, 0.0]]),
         },
     )
     y = cast(
-        Dict[str, np.ndarray],
+        Dict[ConstraintCategory, np.ndarray],
         {
-            "type-a": np.array([[False, True], [False, True], [True, False]]),
+            b"type-a": np.array([[False, True], [False, True], [True, False]]),
-            "type-b": np.array([[False, True]]),
+            b"type-b": np.array([[False, True]]),
         },
     )
     clf: Classifier = Mock(spec=Classifier)
@@ -198,15 +199,15 @@ def test_fit_xy() -> None:
     )
     comp.fit_xy(x, y)
     assert clf.clone.call_count == 2
-    clf_a = comp.classifiers["type-a"]
+    clf_a = comp.classifiers[b"type-a"]
-    clf_b = comp.classifiers["type-b"]
+    clf_b = comp.classifiers[b"type-b"]
     assert clf_a.fit.call_count == 1  # type: ignore
     assert clf_b.fit.call_count == 1  # type: ignore
-    assert_array_equal(clf_a.fit.call_args[0][0], x["type-a"])  # type: ignore
+    assert_array_equal(clf_a.fit.call_args[0][0], x[b"type-a"])  # type: ignore
-    assert_array_equal(clf_b.fit.call_args[0][0], x["type-b"])  # type: ignore
+    assert_array_equal(clf_b.fit.call_args[0][0], x[b"type-b"])  # type: ignore
     assert thr.clone.call_count == 2
-    thr_a = comp.thresholds["type-a"]
+    thr_a = comp.thresholds[b"type-a"]
-    thr_b = comp.thresholds["type-b"]
+    thr_b = comp.thresholds[b"type-b"]
     assert thr_a.fit.call_count == 1  # type: ignore
     assert thr_b.fit.call_count == 1  # type: ignore
     assert thr_a.fit.call_args[0][0] == clf_a  # type: ignore
@@ -215,12 +216,12 @@ def test_fit_xy() -> None:
 
 def test_sample_xy(sample: Sample) -> None:
     x_expected = {
-        "type-a": [[5.0, 1.0, 1.0], [5.0, 1.0, 2.0], [5.0, 1.0, 3.0]],
+        b"type-a": [[5.0, 1.0, 1.0], [5.0, 1.0, 2.0], [5.0, 1.0, 3.0]],
-        "type-b": [[5.0, 1.0, 4.0, 0.0]],
+        b"type-b": [[5.0, 1.0, 4.0, 0.0]],
     }
     y_expected = {
-        "type-a": [[False, True], [False, True], [True, False]],
+        b"type-a": [[False, True], [False, True], [True, False]],
-        "type-b": [[False, True]],
+        b"type-b": [[False, True]],
     }
     xy = StaticLazyConstraintsComponent().sample_xy(None, sample)
     assert xy is not None

tests/features/test_extractor.py

@@ -32,27 +32,45 @@ def test_knapsack() -> None:
     # -------------------------------------------------------
     extractor.extract_after_load_features(instance, solver, sample)
     assert_equals(
-        sample.get_vector("static_var_names"),
+        sample.get_array("static_instance_features"),
+        np.array([67.0, 21.75]),
+    )
+    assert_equals(
+        sample.get_array("static_var_names"),
         np.array(["x[0]", "x[1]", "x[2]", "x[3]", "z"], dtype="S"),
     )
     assert_equals(
-        sample.get_vector("static_var_lower_bounds"), [0.0, 0.0, 0.0, 0.0, 0.0]
+        sample.get_array("static_var_lower_bounds"),
+        np.array([0.0, 0.0, 0.0, 0.0, 0.0]),
     )
     assert_equals(
-        sample.get_array("static_var_obj_coeffs"), [505.0, 352.0, 458.0, 220.0, 0.0]
+        sample.get_array("static_var_obj_coeffs"),
+        np.array([505.0, 352.0, 458.0, 220.0, 0.0]),
     )
     assert_equals(
         sample.get_array("static_var_types"),
         np.array(["B", "B", "B", "B", "C"], dtype="S"),
     )
     assert_equals(
-        sample.get_vector("static_var_upper_bounds"), [1.0, 1.0, 1.0, 1.0, 67.0]
+        sample.get_array("static_var_upper_bounds"),
+        np.array([1.0, 1.0, 1.0, 1.0, 67.0]),
     )
     assert_equals(
         sample.get_array("static_var_categories"),
         np.array(["default", "default", "default", "default", ""], dtype="S"),
     )
-    assert sample.get_vector_list("static_var_features") is not None
+    assert_equals(
+        sample.get_vector_list("static_var_features"),
+        np.array(
+            [
+                [23.0, 505.0, 1.0, 0.32899, 0.0, 0.0, 505.0, 1.0],
+                [26.0, 352.0, 1.0, 0.229316, 0.0, 0.0, 352.0, 1.0],
+                [20.0, 458.0, 1.0, 0.298371, 0.0, 0.0, 458.0, 1.0],
+                [18.0, 220.0, 1.0, 0.143322, 0.0, 0.0, 220.0, 1.0],
+                [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 67.0],
+            ]
+        ),
+    )
     assert_equals(
         sample.get_array("static_constr_names"),
         np.array(["eq_capacity"], dtype="S"),
@@ -69,18 +87,30 @@ def test_knapsack() -> None:
     #         ],
     #     ],
     # )
-    assert_equals(sample.get_vector("static_constr_rhs"), [0.0])
+    assert_equals(
+        sample.get_vector("static_constr_rhs"),
+        np.array([0.0]),
+    )
     assert_equals(
         sample.get_array("static_constr_senses"),
         np.array(["="], dtype="S"),
     )
-    assert_equals(sample.get_vector("static_constr_features"), [None])
+    assert_equals(
+        sample.get_vector("static_constr_features"),
+        np.array([[0.0]]),
+    )
     assert_equals(
         sample.get_vector("static_constr_categories"),
         np.array(["eq_capacity"], dtype="S"),
     )
-    assert_equals(sample.get_array("static_constr_lazy"), np.array([False]))
+    assert_equals(
-    assert_equals(sample.get_vector("static_instance_features"), [67.0, 21.75])
+        sample.get_array("static_constr_lazy"),
+        np.array([False]),
+    )
+    assert_equals(
+        sample.get_vector("static_instance_features"),
+        np.array([67.0, 21.75]),
+    )
     assert_equals(sample.get_scalar("static_constr_lazy_count"), 0)
 
     # after-lp
@@ -112,26 +142,187 @@ def test_knapsack() -> None:
         [inf, 570.869565, inf, 243.692308, inf],
     )
     assert_equals(
-        sample.get_array("lp_var_sa_ub_down"), [0.913043, 0.923077, 0.9, 0.0, 43.0]
+        sample.get_array("lp_var_sa_ub_down"),
+        np.array([0.913043, 0.923077, 0.9, 0.0, 43.0]),
+    )
+    assert_equals(
+        sample.get_array("lp_var_sa_ub_up"),
+        np.array([2.043478, inf, 2.2, inf, 69.0]),
+    )
+    assert_equals(
+        sample.get_array("lp_var_values"),
+        np.array([1.0, 0.923077, 1.0, 0.0, 67.0]),
+    )
+    assert_equals(
+        sample.get_vector_list("lp_var_features"),
+        np.array(
+            [
+                [
+                    23.0,
+                    505.0,
+                    1.0,
+                    0.32899,
+                    0.0,
+                    0.0,
+                    505.0,
+                    1.0,
+                    1.0,
+                    0.32899,
+                    0.0,
+                    0.0,
+                    1.0,
+                    1.0,
+                    5.265874,
+                    46.051702,
+                    193.615385,
+                    -inf,
+                    1.0,
+                    311.384615,
+                    inf,
+                    0.913043,
+                    2.043478,
+                    1.0,
+                ],
+                [
+                    26.0,
+                    352.0,
+                    1.0,
+                    0.229316,
+                    0.0,
+                    0.0,
+                    352.0,
+                    1.0,
+                    1.0,
+                    0.229316,
+                    0.0,
+                    0.076923,
+                    1.0,
+                    1.0,
+                    3.532875,
+                    5.388476,
+                    0.0,
+                    -inf,
+                    0.923077,
+                    317.777778,
+                    570.869565,
+                    0.923077,
+                    inf,
+                    0.923077,
+                ],
+                [
+                    20.0,
+                    458.0,
+                    1.0,
+                    0.298371,
+                    0.0,
+                    0.0,
+                    458.0,
+                    1.0,
+                    1.0,
+                    0.298371,
+                    0.0,
+                    0.0,
+                    1.0,
+                    1.0,
+                    5.232342,
+                    46.051702,
+                    187.230769,
+                    -inf,
+                    1.0,
+                    270.769231,
+                    inf,
+                    0.9,
+                    2.2,
+                    1.0,
+                ],
+                [
+                    18.0,
+                    220.0,
+                    1.0,
+                    0.143322,
+                    0.0,
+                    0.0,
+                    220.0,
+                    1.0,
+                    1.0,
+                    0.143322,
+                    0.0,
+                    0.0,
+                    1.0,
+                    -1.0,
+                    46.051702,
+                    3.16515,
+                    -23.692308,
+                    -0.111111,
+                    1.0,
+                    -inf,
+                    243.692308,
+                    0.0,
+                    inf,
+                    0.0,
+                ],
+                [
+                    0.0,
+                    0.0,
+                    0.0,
+                    0.0,
+                    0.0,
+                    0.0,
+                    0.0,
+                    67.0,
+                    0.0,
+                    0.0,
+                    0.0,
+                    0.0,
+                    1.0,
+                    -1.0,
+                    0.0,
+                    0.0,
+                    13.538462,
+                    -inf,
+                    67.0,
+                    -13.538462,
+                    inf,
+                    43.0,
+                    69.0,
+                    67.0,
+                ],
+            ]
+        ),
     )
-    assert_equals(sample.get_array("lp_var_sa_ub_up"), [2.043478, inf, 2.2, inf, 69.0])
-    assert_equals(sample.get_array("lp_var_values"), [1.0, 0.923077, 1.0, 0.0, 67.0])
-    assert sample.get_vector_list("lp_var_features") is not None
     assert_equals(
         sample.get_array("lp_constr_basis_status"),
         np.array(["N"], dtype="S"),
     )
-    assert_equals(sample.get_array("lp_constr_dual_values"), [13.538462])
+    assert_equals(
-    assert_equals(sample.get_array("lp_constr_sa_rhs_down"), [-24.0])
+        sample.get_array("lp_constr_dual_values"),
-    assert_equals(sample.get_array("lp_constr_sa_rhs_up"), [2.0])
+        np.array([13.538462]),
-    assert_equals(sample.get_array("lp_constr_slacks"), [0.0])
+    )
+    assert_equals(
+        sample.get_array("lp_constr_sa_rhs_down"),
+        np.array([-24.0]),
+    )
+    assert_equals(
+        sample.get_array("lp_constr_sa_rhs_up"),
+        np.array([2.0]),
+    )
+    assert_equals(
+        sample.get_array("lp_constr_slacks"),
+        np.array([0.0]),
+    )
+    assert_equals(
+        sample.get_array("lp_constr_features"),
+        np.array([[0.0, 13.538462, -24.0, 2.0, 0.0]]),
+    )
 
     # after-mip
     # -------------------------------------------------------
     solver.solve()
     extractor.extract_after_mip_features(solver, sample)
-    assert_equals(sample.get_array("mip_var_values"), [1.0, 0.0, 1.0, 1.0, 61.0])
+    assert_equals(
-    assert_equals(sample.get_array("mip_constr_slacks"), [0.0])
+        sample.get_array("mip_var_values"), np.array([1.0, 0.0, 1.0, 1.0, 61.0])
+    )
+    assert_equals(sample.get_array("mip_constr_slacks"), np.array([0.0]))
 
 
 def test_constraint_getindex() -> None:
@@ -200,7 +391,7 @@ class MpsInstance(Instance):
         return gp.read(self.filename)
 
 
-if __name__ == "__main__":
+def main() -> None:
     solver = GurobiSolver()
     instance = MpsInstance(sys.argv[1])
     solver.set_instance(instance)
@@ -213,3 +404,7 @@ if __name__ == "__main__":
         extractor.extract_after_lp_features(solver, sample, lp_stats)
 
     cProfile.run("run()", filename="tmp/prof")
+
+
+if __name__ == "__main__":
+    main()