Remove sample.{get,set}

miplearn/features/extractor.py

@@ -39,7 +39,7 @@ class FeaturesExtractor:
         sample.put_vector("var_types", variables.types)
         sample.put_vector("var_upper_bounds", variables.upper_bounds)
         sample.put_vector("constr_names", constraints.names)
-        sample.put("constr_lhs", constraints.lhs)
+        # sample.put("constr_lhs", constraints.lhs)
         sample.put_vector("constr_rhs", constraints.rhs)
         sample.put_vector("constr_senses", constraints.senses)
         self._extract_user_features_vars(instance, sample)
@@ -49,13 +49,12 @@ class FeaturesExtractor:
         sample.put_vector_list(
             "var_features",
             self._combine(
-                sample,
                 [
-                    "var_features_AlvLouWeh2017",
-                    "var_features_user",
-                    "var_lower_bounds",
-                    "var_obj_coeffs",
-                    "var_upper_bounds",
+                    sample.get_vector_list("var_features_AlvLouWeh2017"),
+                    sample.get_vector_list("var_features_user"),
+                    sample.get_vector("var_lower_bounds"),
+                    sample.get_vector("var_obj_coeffs"),
+                    sample.get_vector("var_upper_bounds"),
                 ],
             ),
         )
@@ -85,45 +84,43 @@ class FeaturesExtractor:
         sample.put_vector_list(
             "lp_var_features",
             self._combine(
-                sample,
                 [
-                    "lp_var_features_AlvLouWeh2017",
-                    "lp_var_reduced_costs",
-                    "lp_var_sa_lb_down",
-                    "lp_var_sa_lb_up",
-                    "lp_var_sa_obj_down",
-                    "lp_var_sa_obj_up",
-                    "lp_var_sa_ub_down",
-                    "lp_var_sa_ub_up",
-                    "lp_var_values",
-                    "var_features_user",
-                    "var_lower_bounds",
-                    "var_obj_coeffs",
-                    "var_upper_bounds",
+                    sample.get_vector_list("lp_var_features_AlvLouWeh2017"),
+                    sample.get_vector("lp_var_reduced_costs"),
+                    sample.get_vector("lp_var_sa_lb_down"),
+                    sample.get_vector("lp_var_sa_lb_up"),
+                    sample.get_vector("lp_var_sa_obj_down"),
+                    sample.get_vector("lp_var_sa_obj_up"),
+                    sample.get_vector("lp_var_sa_ub_down"),
+                    sample.get_vector("lp_var_sa_ub_up"),
+                    sample.get_vector("lp_var_values"),
+                    sample.get_vector_list("var_features_user"),
+                    sample.get_vector("var_lower_bounds"),
+                    sample.get_vector("var_obj_coeffs"),
+                    sample.get_vector("var_upper_bounds"),
                 ],
             ),
         )
         sample.put_vector_list(
             "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",
+                    sample.get_vector_list("constr_features_user"),
+                    sample.get_vector("lp_constr_dual_values"),
+                    sample.get_vector("lp_constr_sa_rhs_down"),
+                    sample.get_vector("lp_constr_sa_rhs_up"),
+                    sample.get_vector("lp_constr_slacks"),
                 ],
             ),
         )
-        instance_features_user = sample.get("instance_features_user")
+        instance_features_user = sample.get_vector("instance_features_user")
         assert instance_features_user is not None
         sample.put_vector(
             "lp_instance_features",
             instance_features_user
             + [
-                sample.get("lp_value"),
-                sample.get("lp_wallclock_time"),
+                sample.get_scalar("lp_value"),
+                sample.get_scalar("lp_wallclock_time"),
             ],
         )
 
@@ -146,7 +143,7 @@ class FeaturesExtractor:
         user_features: List[Optional[List[float]]] = []
         var_features_dict = instance.get_variable_features()
         var_categories_dict = instance.get_variable_categories()
-        var_names = sample.get("var_names")
+        var_names = sample.get_vector("var_names")
         assert var_names is not None
         for (i, var_name) in enumerate(var_names):
             if var_name not in var_categories_dict:
@@ -177,7 +174,7 @@ class FeaturesExtractor:
                     )
                 user_features_i = list(user_features_i)
             user_features.append(user_features_i)
-        sample.put("var_categories", categories)
+        sample.put_vector("var_categories", categories)
         sample.put_vector_list("var_features_user", user_features)
 
     def _extract_user_features_constrs(
@@ -191,7 +188,7 @@ class FeaturesExtractor:
         lazy: List[bool] = []
         constr_categories_dict = instance.get_constraint_categories()
         constr_features_dict = instance.get_constraint_features()
-        constr_names = sample.get("constr_names")
+        constr_names = sample.get_vector("constr_names")
         assert constr_names is not None
 
         for (cidx, cname) in enumerate(constr_names):
@@ -229,7 +226,7 @@ class FeaturesExtractor:
                 lazy.append(False)
         sample.put_vector_list("constr_features_user", user_features)
         sample.put_vector("constr_lazy", lazy)
-        sample.put("constr_categories", categories)
+        sample.put_vector("constr_categories", categories)
 
     def _extract_user_features_instance(
         self,
@@ -248,7 +245,7 @@ class FeaturesExtractor:
                 f"Instance features must be a list of numbers. "
                 f"Found {type(v).__name__} instead."
             )
-        constr_lazy = sample.get("constr_lazy")
+        constr_lazy = sample.get_vector("constr_lazy")
         assert constr_lazy is not None
         sample.put_vector("instance_features_user", user_features)
         sample.put_scalar("static_lazy_count", sum(constr_lazy))
@@ -260,10 +257,10 @@ class FeaturesExtractor:
         sample: Sample,
         prefix: str = "",
     ) -> None:
-        obj_coeffs = sample.get("var_obj_coeffs")
-        obj_sa_down = sample.get("lp_var_sa_obj_down")
-        obj_sa_up = sample.get("lp_var_sa_obj_up")
-        values = sample.get(f"lp_var_values")
+        obj_coeffs = sample.get_vector("var_obj_coeffs")
+        obj_sa_down = sample.get_vector("lp_var_sa_obj_down")
+        obj_sa_up = sample.get_vector("lp_var_sa_obj_up")
+        values = sample.get_vector(f"lp_var_values")
         assert obj_coeffs is not None
 
         pos_obj_coeff_sum = 0.0
@@ -335,12 +332,10 @@ class FeaturesExtractor:
 
     def _combine(
         self,
-        sample: Sample,
-        attrs: List[str],
+        items: List,
     ) -> List[List[float]]:
         combined: List[List[float]] = []
-        for attr in attrs:
-            series = sample.get(attr)
+        for series in items:
             if series is None:
                 continue
             if len(combined) == 0:

miplearn/features/sample.py

@@ -4,14 +4,22 @@
 
 from abc import ABC, abstractmethod
 from copy import deepcopy
-from typing import Dict, Optional, Any, Union, List, Tuple, cast
+from typing import Dict, Optional, Any, Union, List, Tuple, cast, Set
 
 import h5py
 import numpy as np
+from h5py import Dataset
 from overrides import overrides
 
 Scalar = Union[None, bool, str, int, float]
-Vector = Union[None, List[bool], List[str], List[int], List[float]]
+Vector = Union[
+    None,
+    List[bool],
+    List[str],
+    List[int],
+    List[float],
+    List[Optional[str]],
+]
 VectorList = Union[
     List[List[bool]],
     List[List[str]],
@@ -51,39 +59,16 @@ class Sample(ABC):
     def put_vector_list(self, key: str, value: VectorList) -> None:
         pass
 
-    @abstractmethod
-    def get(self, key: str) -> Optional[Any]:
-        pass
+    def get_set(self, key: str) -> Set:
+        v = self.get_vector(key)
+        if v:
+            return set(v)
+        else:
+            return set()
 
-    @abstractmethod
-    def put(self, key: str, value: Any) -> None:
-        """
-        Add a new key/value pair to the sample. If the key already exists,
-        the previous value is silently replaced.
-
-        Only the following data types are supported:
-        - str, bool, int, float
-        - List[str], List[bool], List[int], List[float]
-        """
-        pass
-
-    def _assert_supported(self, value: Any) -> None:
-        def _is_primitive(v: Any) -> bool:
-            if isinstance(v, (str, bool, int, float)):
-                return True
-            if v is None:
-                return True
-            return False
-
-        if _is_primitive(value):
-            return
-        if isinstance(value, list):
-            if _is_primitive(value[0]):
-                return
-            if isinstance(value[0], list):
-                if _is_primitive(value[0][0]):
-                    return
-        assert False, f"Value has unsupported type: {value}"
+    def put_set(self, key: str, value: Set) -> None:
+        v = list(value)
+        self.put_vector(key, v)
 
     def _assert_is_scalar(self, value: Any) -> None:
         if value is None:
@@ -118,42 +103,40 @@ class MemorySample(Sample):
 
     @overrides
     def get_scalar(self, key: str) -> Optional[Any]:
-        return self.get(key)
+        return self._get(key)
 
     @overrides
     def get_vector(self, key: str) -> Optional[Any]:
-        return self.get(key)
+        return self._get(key)
 
     @overrides
     def get_vector_list(self, key: str) -> Optional[Any]:
-        return self.get(key)
+        return self._get(key)
 
     @overrides
     def put_scalar(self, key: str, value: Scalar) -> None:
         self._assert_is_scalar(value)
-        self.put(key, value)
+        self._put(key, value)
 
     @overrides
     def put_vector(self, key: str, value: Vector) -> None:
         if value is None:
             return
         self._assert_is_vector(value)
-        self.put(key, value)
+        self._put(key, value)
 
     @overrides
     def put_vector_list(self, key: str, value: VectorList) -> None:
         self._assert_is_vector_list(value)
-        self.put(key, value)
+        self._put(key, value)
 
-    @overrides
-    def get(self, key: str) -> Optional[Any]:
+    def _get(self, key: str) -> Optional[Any]:
         if key in self._data:
             return self._data[key]
         else:
             return None
 
-    @overrides
-    def put(self, key: str, value: Any) -> None:
+    def _put(self, key: str, value: Any) -> None:
         self._data[key] = value
 
 
@@ -200,20 +183,18 @@ class Hdf5Sample(Sample):
     @overrides
     def put_scalar(self, key: str, value: Any) -> None:
         self._assert_is_scalar(value)
-        self.put(key, value)
+        self._put(key, value)
 
     @overrides
     def put_vector(self, key: str, value: Vector) -> None:
         if value is None:
             return
         self._assert_is_vector(value)
-        self.put(key, value)
+        self._put(key, value)
 
     @overrides
     def put_vector_list(self, key: str, value: VectorList) -> None:
         self._assert_is_vector_list(value)
-        if key in self.file:
-            del self.file[key]
         padded, lens = _pad(value)
         data = None
         for v in value:
@@ -227,22 +208,13 @@ class Hdf5Sample(Sample):
                 data = np.array(padded)
             break
         assert data is not None
-        ds = self.file.create_dataset(key, data=data)
+        ds = self._put(key, data)
         ds.attrs["lengths"] = lens
 
-    @overrides
-    def get(self, key: str) -> Optional[Any]:
-        ds = self.file[key]
-        if h5py.check_string_dtype(ds.dtype):
-            return ds.asstr()[:].tolist()
-        else:
-            return ds[:].tolist()
-
-    @overrides
-    def put(self, key: str, value: Any) -> None:
+    def _put(self, key: str, value: Any) -> Dataset:
         if key in self.file:
             del self.file[key]
-        self.file.create_dataset(key, data=value)
+        return self.file.create_dataset(key, data=value)
 
 
 def _pad(veclist: VectorList) -> Tuple[VectorList, List[int]]: