Rename Variables and Constraints; move to internal.py

miplearn/solvers/internal.py

@@ -5,9 +5,8 @@
 import logging
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Any, List, Optional, List
+from typing import Any, Optional, List, Tuple, TYPE_CHECKING
 
-from miplearn.features import VariableFeatures, ConstraintFeatures
 from miplearn.instance.base import Instance
 from miplearn.types import (
     IterationCallback,
@@ -18,6 +17,9 @@ from miplearn.types import (
 
 logger = logging.getLogger(__name__)
 
+if TYPE_CHECKING:
+    from miplearn.features import Sample
+
 
 @dataclass
 class LPSolveStats:
@@ -44,20 +46,90 @@ class MIPSolveStats:
     mip_warm_start_value: Optional[float] = None
 
 
+@dataclass
+class Variables:
+    names: Optional[List[str]] = None
+    basis_status: Optional[List[str]] = None
+    lower_bounds: Optional[List[float]] = None
+    obj_coeffs: Optional[List[float]] = None
+    reduced_costs: Optional[List[float]] = None
+    sa_lb_down: Optional[List[float]] = None
+    sa_lb_up: Optional[List[float]] = None
+    sa_obj_down: Optional[List[float]] = None
+    sa_obj_up: Optional[List[float]] = None
+    sa_ub_down: Optional[List[float]] = None
+    sa_ub_up: Optional[List[float]] = None
+    types: Optional[List[str]] = None
+    upper_bounds: Optional[List[float]] = None
+    values: Optional[List[float]] = None
+
+
+@dataclass
+class Constraints:
+    basis_status: Optional[List[str]] = None
+    dual_values: Optional[List[float]] = 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
+    senses: Optional[List[str]] = None
+    slacks: Optional[List[float]] = None
+
+    @staticmethod
+    def from_sample(sample: "Sample") -> "Constraints":
+        return Constraints(
+            basis_status=sample.get("lp_constr_basis_status"),
+            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"),
+        )
+
+    def __getitem__(self, selected: List[bool]) -> "Constraints":
+        return Constraints(
+            basis_status=self._filter(self.basis_status, selected),
+            dual_values=self._filter(self.dual_values, selected),
+            names=self._filter(self.names, selected),
+            lazy=self._filter(self.lazy, selected),
+            lhs=self._filter(self.lhs, selected),
+            rhs=self._filter(self.rhs, selected),
+            sa_rhs_down=self._filter(self.sa_rhs_down, selected),
+            sa_rhs_up=self._filter(self.sa_rhs_up, selected),
+            senses=self._filter(self.senses, selected),
+            slacks=self._filter(self.slacks, selected),
+        )
+
+    def _filter(
+        self,
+        obj: Optional[List],
+        selected: List[bool],
+    ) -> Optional[List]:
+        if obj is None:
+            return None
+        return [obj[i] for (i, selected_i) in enumerate(selected) if selected_i]
+
+
 class InternalSolver(ABC):
     """
     Abstract class representing the MIP solver used internally by LearningSolver.
     """
 
     @abstractmethod
-    def add_constraints(self, cf: ConstraintFeatures) -> None:
+    def add_constraints(self, cf: Constraints) -> None:
         """Adds the given constraints to the model."""
         pass
 
     @abstractmethod
     def are_constraints_satisfied(
         self,
-        cf: ConstraintFeatures,
+        cf: Constraints,
         tol: float = 1e-5,
     ) -> List[bool]:
         """
@@ -133,7 +205,7 @@ class InternalSolver(ABC):
         with_static: bool = True,
         with_sa: bool = True,
         with_lhs: bool = True,
-    ) -> ConstraintFeatures:
+    ) -> Constraints:
         pass
 
     @abstractmethod
@@ -149,7 +221,7 @@ class InternalSolver(ABC):
         self,
         with_static: bool = True,
         with_sa: bool = True,
-    ) -> VariableFeatures:
+    ) -> Variables:
         """
         Returns a description of the decision variables in the problem.