MIPLearn v0.3

miplearn/components/primal/mem.py

@@ -0,0 +1,167 @@
+#  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
+#  Copyright (C) 2020-2022, UChicago Argonne, LLC. All rights reserved.
+#  Released under the modified BSD license. See COPYING.md for more details.
+
+import logging
+from abc import ABC, abstractmethod
+from typing import List, Dict, Any, Optional, Tuple
+
+import numpy as np
+
+from . import _extract_bin_var_names_values
+from .actions import PrimalComponentAction
+from ...extractors.abstract import FeaturesExtractor
+from ...solvers.abstract import AbstractModel
+from ...h5 import H5File
+
+logger = logging.getLogger()
+
+
+class SolutionConstructor(ABC):
+    @abstractmethod
+    def construct(self, y_proba: np.ndarray, solutions: np.ndarray) -> np.ndarray:
+        pass
+
+
+class MemorizingPrimalComponent:
+    """
+    Component that memorizes all solutions seen during training, then fits a
+    single classifier to predict which of the memorized solutions should be
+    provided to the solver. Optionally combines multiple memorized solutions
+    into a single, partial one.
+    """
+
+    def __init__(
+        self,
+        clf: Any,
+        extractor: FeaturesExtractor,
+        constructor: SolutionConstructor,
+        action: PrimalComponentAction,
+    ) -> None:
+        assert clf is not None
+        self.clf = clf
+        self.extractor = extractor
+        self.constructor = constructor
+        self.solutions_: Optional[np.ndarray] = None
+        self.bin_var_names_: Optional[np.ndarray] = None
+        self.action = action
+
+    def fit(self, train_h5: List[str]) -> None:
+        logger.info("Reading training data...")
+        n_samples = len(train_h5)
+        solutions_ = []
+        self.bin_var_names_ = None
+        x, y, n_features = [], [], None
+        solution_to_idx: Dict[Tuple, int] = {}
+        for h5_filename in train_h5:
+            with H5File(h5_filename, "r") as h5:
+                bin_var_names, bin_var_values, _ = _extract_bin_var_names_values(h5)
+
+                # Store/check variable names
+                if self.bin_var_names_ is None:
+                    self.bin_var_names_ = bin_var_names
+                else:
+                    assert np.all(bin_var_names == self.bin_var_names_)
+
+                # Store solution
+                sol = tuple(np.where(bin_var_values)[0])
+                if sol not in solution_to_idx:
+                    solutions_.append(bin_var_values)
+                    solution_to_idx[sol] = len(solution_to_idx)
+                y.append(solution_to_idx[sol])
+
+                # Extract features
+                x_sample = self.extractor.get_instance_features(h5)
+                assert len(x_sample.shape) == 1
+                if n_features is None:
+                    n_features = len(x_sample)
+                else:
+                    assert len(x_sample) == n_features
+                x.append(x_sample)
+
+        logger.info("Constructing matrices...")
+        x_np = np.vstack(x)
+        y_np = np.array(y)
+        assert len(x_np.shape) == 2
+        assert x_np.shape[0] == n_samples
+        assert x_np.shape[1] == n_features
+        assert y_np.shape == (n_samples,)
+        self.solutions_ = np.array(solutions_)
+        n_classes = len(solution_to_idx)
+        logger.info(
+            f"Dataset has {n_samples:,d} samples, "
+            f"{n_features:,d} features and {n_classes:,d} classes"
+        )
+
+        logger.info("Training classifier...")
+        self.clf.fit(x_np, y_np)
+
+        logger.info("Done fitting.")
+
+    def before_mip(
+        self, test_h5: str, model: AbstractModel, stats: Dict[str, Any]
+    ) -> None:
+        assert self.solutions_ is not None
+        assert self.bin_var_names_ is not None
+
+        # Read features
+        with H5File(test_h5, "r") as h5:
+            x_sample = self.extractor.get_instance_features(h5)
+        assert len(x_sample.shape) == 1
+        x_sample = x_sample.reshape(1, -1)
+
+        # Predict optimal solution
+        logger.info("Predicting primal solution...")
+        y_proba = self.clf.predict_proba(x_sample)
+        assert len(y_proba.shape) == 2
+        assert y_proba.shape[0] == 1
+        assert y_proba.shape[1] == len(self.solutions_)
+
+        # Construct warm starts, based on prediction
+        starts = self.constructor.construct(y_proba[0, :], self.solutions_)
+        self.action.perform(model, self.bin_var_names_, starts, stats)
+
+
+class SelectTopSolutions(SolutionConstructor):
+    """
+    Warm start construction strategy that selects and returns the top k solutions.
+    """
+
+    def __init__(self, k: int) -> None:
+        self.k = k
+
+    def construct(self, y_proba: np.ndarray, solutions: np.ndarray) -> np.ndarray:
+        # Check arguments
+        assert len(y_proba.shape) == 1
+        assert len(solutions.shape) == 2
+        assert len(y_proba) == solutions.shape[0]
+
+        # Select top k solutions
+        ind = np.argsort(-y_proba, kind="stable")
+        selected = ind[: min(self.k, len(ind))]
+        return solutions[selected, :]
+
+
+class MergeTopSolutions(SolutionConstructor):
+    """
+    Warm start construction strategy that first selects the top k solutions,
+    then merges them into a single solution.
+
+    To merge the solutions, the strategy first computes the mean optimal value of each
+    decision variable, then: (i) sets the variable to zero if the mean is below
+    thresholds[0]; (ii) sets the variable to one if the mean is above thresholds[1];
+    (iii) leaves the variable free otherwise.
+    """
+
+    def __init__(self, k: int, thresholds: List[float]):
+        assert len(thresholds) == 2
+        self.k = k
+        self.thresholds = thresholds
+
+    def construct(self, y_proba: np.ndarray, solutions: np.ndarray) -> np.ndarray:
+        filtered = SelectTopSolutions(self.k).construct(y_proba, solutions)
+        mean = filtered.mean(axis=0)
+        start = np.full((1, solutions.shape[1]), float("nan"))
+        start[0, mean <= self.thresholds[0]] = 0
+        start[0, mean >= self.thresholds[1]] = 1
+        return start