feat: New surrogate models API

syne_tune/optimizer/schedulers/searchers/bayesopt/models/estimator.py

@@ -174,31 +174,3 @@ def transform_state_to_data(
         targets = np.vstack([targets * np.ones((1, num_fantasy_samples))] + fanta_lst)
         features = np.vstack([features] + cand_lst)
     return TransformedData(features, targets, mean, std)
-
-
-class EstimatorFromTransformedData(Estimator):
-    def _fit(self, data: TransformedData, update_params: bool) -> Predictor:
-        """
-        Implements :meth:`fit_from_state`, given transformed data.
-
-        :param data: Transformed data (features, targets)
-        :param update_params: Should model (hyper)parameters be updated?
-        :return: Predictor, wrapping the posterior state
-        """
-        raise NotImplementedError()
-
-    @property
-    def normalize_targets(self) -> bool:
-        """
-        :return: Should targets in ``state`` be normalized before calling
-            :meth:`_fit`?
-        """
-        raise NotImplementedError()
-
-    def fit_from_state(self, state: TuningJobState, update_params: bool) -> Predictor:
-        return self._fit(
-            data=transform_state_to_data(
-                state=state, normalize_targets=self.normalize_targets
-            ),
-            update_params=update_params,
-        )

syne_tune/optimizer/schedulers/searchers/bayesopt/models/sklearn_estimator.py

@@ -0,0 +1,75 @@
+# Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# A copy of the License is located at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the "license" file accompanying this file. This file is distributed
+# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+# express or implied. See the License for the specific language governing
+# permissions and limitations under the License.
+from syne_tune.optimizer.schedulers.searchers.bayesopt.datatypes.tuning_job_state import (
+    TuningJobState,
+)
+from syne_tune.optimizer.schedulers.searchers.bayesopt.models.estimator import (
+    Estimator,
+    transform_state_to_data,
+)
+from syne_tune.optimizer.schedulers.searchers.bayesopt.sklearn.estimator import (
+    SklearnEstimator,
+)
+from syne_tune.optimizer.schedulers.searchers.bayesopt.models.sklearn_predictor import (
+    SklearnPredictorWrapper,
+)
+from syne_tune.optimizer.schedulers.searchers.bayesopt.tuning_algorithms.base_classes import (
+    Predictor,
+)
+
+
+class SklearnEstimatorWrapper(Estimator):
+    """
+    Wrapper class for the sklearn estimators to be used with BayesianOptimizationSearcher
+    """
+
+    def __init__(self, contributed_estimator: SklearnEstimator, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.contributed_estimator = contributed_estimator
+
+    @property
+    def normalize_targets(self) -> bool:
+        """
+        :return: Should targets in ``state`` be normalized before fitting the estimator
+        """
+        return self.contributed_estimator.normalize_targets
+
+    def fit_from_state(
+        self, state: TuningJobState, update_params: bool = False
+    ) -> Predictor:
+        """
+        Creates a
+        :class:`~syne_tune.optimizer.schedulers.searchers.bayesopt.tuning_algorithms.base_classes.Predictor`
+        object based on data in ``state``.
+
+        If the model also has hyperparameters, these are learned iff
+        ``update_params == True``. Otherwise, these parameters are not changed,
+        but only the posterior state is computed.
+
+        If your surrogate model is not Bayesian, or does not have hyperparameters,
+        you can ignore the ``update_params`` argument,
+
+        :param state: Current data model parameters are to be fit on, and the
+            posterior state is to be computed from
+        :param update_params: See above
+        :return: Predictor, wrapping the posterior state
+        """
+        data = transform_state_to_data(
+            state=state, normalize_targets=self.normalize_targets
+        )
+        contributed_predictor = self.contributed_estimator.fit(
+            data.features, data.targets, update_params=update_params
+        )
+        return SklearnPredictorWrapper(
+            contributed_predictor=contributed_predictor, state=state
+        )

syne_tune/optimizer/schedulers/searchers/bayesopt/models/sklearn_predictor.py

@@ -0,0 +1,83 @@
+# Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# A copy of the License is located at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the "license" file accompanying this file. This file is distributed
+# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+# express or implied. See the License for the specific language governing
+# permissions and limitations under the License.
+from typing import Optional, List, Dict
+
+import numpy as np
+
+from syne_tune.optimizer.schedulers.searchers.bayesopt.datatypes.tuning_job_state import (
+    TuningJobState,
+)
+from syne_tune.optimizer.schedulers.searchers.bayesopt.models.model_base import (
+    BasePredictor,
+)
+from syne_tune.optimizer.schedulers.searchers.bayesopt.sklearn.predictor import (
+    SklearnPredictor,
+)
+
+
+class SklearnPredictorWrapper(BasePredictor):
+    """
+    Wrapper class for the sklearn estimators to be used with ContributedEstimatorWrapper
+    """
+
+    def __init__(
+        self,
+        contributed_predictor: SklearnPredictor,
+        state: TuningJobState,
+        active_metric: Optional[str] = None,
+    ):
+        super().__init__(state, active_metric)
+        self.contributed_predictor = contributed_predictor
+
+    def predict(self, inputs: np.ndarray) -> List[Dict[str, np.ndarray]]:
+        """
+        Returns signals which are statistics of the predictive distribution at
+        input points ``inputs``. By default:
+
+        * "mean": Predictive means.
+        - "std": Predictive stddevs, shape ``(n,)``
+
+        This function relies on the assigned ContributedPredictor to execute the predictions
+
+        :param inputs: Input points, shape ``(n, d)``
+        :return: List of ``dict`` with keys :meth:`keys_predict`, of length 1
+        """
+
+        mean, std = self.contributed_predictor.predict(inputs)
+        outputs = {"mean": mean}
+        if std is not None:
+            outputs["std"] = std
+        return [outputs]
+
+    def backward_gradient(
+        self, input: np.ndarray, head_gradients: List[Dict[str, np.ndarray]]
+    ) -> List[np.ndarray]:
+        """
+        Computes the gradient :math:`\nabla f(x)` for an acquisition
+        function :math:`f(x)`, where :math:`x` is a single input point. This
+        is using reverse mode differentiation, the head gradients are passed
+        by the acquisition function. The head gradients are
+        :math:`\partial_k f`, where :math:`k` runs over the statistics
+        returned by :meth:`predict` for the single input point :math:`x`.
+        The shape of head gradients is the same as the shape of the
+        statistics.
+
+        :param input: Single input point :math:`x`, shape ``(d,)``
+        :param head_gradients: See above
+        :return: Gradient :math:`\nabla f(x)` (one-length list)
+        """
+        return [
+            self.contributed_predictor.backward_gradient(
+                input=input, head_gradients=head_gradients
+            )
+        ]

syne_tune/optimizer/schedulers/searchers/bayesopt/sklearn/__init__.py

@@ -0,0 +1,12 @@
+# Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# A copy of the License is located at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the "license" file accompanying this file. This file is distributed
+# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+# express or implied. See the License for the specific language governing
+# permissions and limitations under the License.

syne_tune/optimizer/schedulers/searchers/bayesopt/sklearn/estimator.py

@@ -0,0 +1,44 @@
+# Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# A copy of the License is located at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the "license" file accompanying this file. This file is distributed
+# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+# express or implied. See the License for the specific language governing
+# permissions and limitations under the License.
+import numpy as np
+
+from syne_tune.optimizer.schedulers.searchers.bayesopt.sklearn.predictor import (
+    SklearnPredictor,
+)
+
+
+class SklearnEstimator:
+    """
+    Base class for the sklearn Estimators
+    """
+
+    def fit(
+        self, X: np.ndarray, y: np.ndarray, update_params: bool
+    ) -> SklearnPredictor:
+        """
+        Implements :meth:`fit_from_state`, given transformed data.
+
+        :param X: Training data in ndarray of shape (n_samples, n_features)
+        :param y: Target values in ndarray of shape (n_samples,)
+        :param update_params: Should model (hyper)parameters be updated?
+        :return: Predictor, wrapping the posterior state
+        """
+        raise NotImplementedError()
+
+    @property
+    def normalize_targets(self) -> bool:
+        """
+        :return: Should targets in ``state`` be normalized before calling
+            :meth:`fit`?
+        """
+        return False

syne_tune/optimizer/schedulers/searchers/bayesopt/sklearn/predictor.py

@@ -0,0 +1,56 @@
+# Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# A copy of the License is located at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the "license" file accompanying this file. This file is distributed
+# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+# express or implied. See the License for the specific language governing
+# permissions and limitations under the License.
+from typing import Tuple, List, Dict
+
+import numpy as np
+
+
+class SklearnPredictor:
+    """
+    Base class for the sklearn predictors
+    """
+
+    def predict(
+        self, X: np.ndarray, return_std: bool = True
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Returns signals which are statistics of the predictive distribution at
+        input points ``inputs``.
+
+
+        :param inputs: Input points, shape ``(n, d)``
+        :return: Tuple with the following entries:
+            * "mean": Predictive means in shape of ``(n,)``
+            * "std": Predictive stddevs, shape ``(n,)``
+        """
+
+        raise NotImplementedError()
+
+    def backward_gradient(
+        self, input: np.ndarray, head_gradients: List[Dict[str, np.ndarray]]
+    ) -> np.ndarray:
+        """
+        Computes the gradient :math:`\nabla f(x)` for an acquisition
+        function :math:`f(x)`, where :math:`x` is a single input point. This
+        is using reverse mode differentiation, the head gradients are passed
+        by the acquisition function. The head gradients are
+        :math:`\partial_k f`, where :math:`k` runs over the statistics
+        returned by :meth:`predict` for the single input point :math:`x`.
+        The shape of head gradients is the same as the shape of the
+        statistics.
+
+        :param input: Single input point :math:`x`, shape ``(d,)``
+        :param head_gradients: See above
+        :return: Gradient :math:`\nabla f(x)`
+        """
+        raise NotImplementedError()

tst/schedulers/bayesopt/test_sklearn_surrogate.py

@@ -0,0 +1,88 @@
+# Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# A copy of the License is located at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the "license" file accompanying this file. This file is distributed
+# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+# express or implied. See the License for the specific language governing
+# permissions and limitations under the License.
+from typing import Tuple
+
+import numpy as np
+import pytest
+
+from syne_tune.config_space import uniform, choice
+from syne_tune.optimizer.schedulers.searchers.bayesopt.sklearn.estimator import (
+    SklearnEstimator,
+)
+from syne_tune.optimizer.schedulers.searchers.bayesopt.models.sklearn_estimator import (
+    SklearnEstimatorWrapper,
+)
+from syne_tune.optimizer.schedulers.searchers.bayesopt.sklearn.predictor import (
+    SklearnPredictor,
+)
+from syne_tune.optimizer.schedulers.searchers.bayesopt.models.sklearn_predictor import (
+    SklearnPredictorWrapper,
+)
+from syne_tune.optimizer.schedulers.searchers.bayesopt.datatypes.common import (
+    dictionarize_objective,
+)
+from syne_tune.optimizer.schedulers.searchers.bayesopt.datatypes.tuning_job_state import (
+    TuningJobState,
+)
+from syne_tune.optimizer.schedulers.searchers.bayesopt.utils.test_objects import (
+    create_tuning_job_state,
+)
+from syne_tune.optimizer.schedulers.searchers.utils.hp_ranges_factory import (
+    make_hyperparameter_ranges,
+)
+
+
+class TestPredictor(SklearnPredictor):
+    def predict(
+        self, X: np.ndarray, return_std: bool = True
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        nexamples = X.shape[0]
+        return np.ones_like(nexamples), np.zeros(nexamples)
+
+
+class TestEstimator(SklearnEstimator):
+    def fit(self, X: np.ndarray, y: np.ndarray, update_params: bool) -> TestPredictor:
+        # Assert the right data is passed to the fit
+        np.testing.assert_allclose(X[:, 0], np.array([0.2, 0.31, 0.15]))
+        np.testing.assert_allclose(y.ravel(), np.array([1.0, 2.0, 0.3]))
+        return TestPredictor()
+
+
+@pytest.fixture
+def tuning_job_state() -> TuningJobState:
+    hp_ranges1 = make_hyperparameter_ranges(
+        {"a1_hp_1": uniform(-5.0, 5.0), "a1_hp_2": choice(["a", "b", "c"])}
+    )
+    X1 = [(-3.0, "a"), (-1.9, "c"), (-3.5, "a")]
+    Y1 = [dictionarize_objective(y) for y in (1.0, 2.0, 0.3)]
+
+    return create_tuning_job_state(hp_ranges=hp_ranges1, cand_tuples=X1, metrics=Y1)
+
+
+def test_estimator_wrapper_interface(tuning_job_state):
+    estimator = SklearnEstimatorWrapper(contributed_estimator=TestEstimator())
+    predictor = estimator.fit_from_state(tuning_job_state)
+
+    assert isinstance(predictor, SklearnPredictorWrapper)
+    assert isinstance(predictor.contributed_predictor, TestPredictor)
+    assert isinstance(estimator, SklearnEstimatorWrapper)
+    assert isinstance(estimator.contributed_estimator, TestEstimator)
+
+
+def test_predictor_wrapper_interface(tuning_job_state):
+    estimator = SklearnEstimatorWrapper(contributed_estimator=TestEstimator())
+    predictor = estimator.fit_from_state(tuning_job_state)
+    predictions = predictor.predict(np.random.uniform(size=(10, 3)))
+
+    np.testing.assert_allclose(predictions[0]["mean"], np.ones(shape=10))
+    np.testing.assert_allclose(predictions[0]["std"], np.zeros(shape=10))