#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from abc import ABC, abstractmethod
from typing import Any, Callable, List, Optional, Tuple, Union
import numpy as np
import torch
from ax.utils.common.docutils import copy_doc
from ax.utils.common.typeutils import checked_cast, not_none
from botorch.test_functions import synthetic as botorch_synthetic
# pyre-fixme[3]: Return annotation cannot be `Any`.
# pyre-fixme[24]: Generic type `Callable` expects 2 type parameters.
[docs]class SyntheticFunction(ABC):
# pyre-fixme[4]: Attribute must be annotated.
_required_dimensionality = None
# pyre-fixme[4]: Attribute must be annotated.
_domain = None
# pyre-fixme[4]: Attribute must be annotated.
_minimums = None
# pyre-fixme[4]: Attribute must be annotated.
_maximums = None
# pyre-fixme[4]: Attribute must be annotated.
_fmin = None
# pyre-fixme[4]: Attribute must be annotated.
_fmax = None
@property
@informative_failure_on_none
def name(self) -> str:
return f"{self.__class__.__name__}"
def __call__(
self,
*args: Union[int, float, np.ndarray],
**kwargs: Union[int, float, np.ndarray],
) -> Union[float, np.ndarray]:
"""Simplified way to call the synthetic function and pass the argument
numbers directly, e.g. `branin(2.0, 3.0)`.
"""
if kwargs:
if self.required_dimensionality:
assert (
len(kwargs) == self.required_dimensionality or len(kwargs) == 1
), (
f"Function {self.name} expected either "
f"{self.required_dimensionality} arguments "
"or a single numpy array argument."
)
assert not args, (
f"Function {self.name} expected either all anonymous "
"arguments or all keyword arguments."
)
args = list(kwargs.values()) # pyre-ignore[9]
for x in args:
if isinstance(x, np.ndarray):
return self.f(X=x)
assert np.isscalar(
x
), f"Expected numerical arguments or numpy arrays, got {type(x)}."
if isinstance(x, int):
x = float(x)
return checked_cast(float, self.f(np.array(args)))
[docs] def f(self, X: np.ndarray) -> Union[float, np.ndarray]:
"""Synthetic function implementation.
Args:
X (numpy.ndarray): an n by d array, where n represents the number
of observations and d is the dimensionality of the inputs.
Returns:
numpy.ndarray: an n-dimensional array.
"""
assert isinstance(X, np.ndarray), "X must be a numpy (nd)array."
if self.required_dimensionality:
if len(X.shape) == 1:
input_dim = X.shape[0]
elif len(X.shape) == 2:
input_dim = X.shape[1]
else:
raise ValueError(
"Synthetic function call expects input of either 1-d array or "
"n by d array, where n is number of observations and d is "
"dimensionality of the input."
)
assert input_dim == self.required_dimensionality, (
f"Input violates required dimensionality of {self.name}: "
f"{self.required_dimensionality}. Got {input_dim}."
)
X = X.astype(np.float64)
if len(X.shape) == 1:
return self._f(X=X)
else:
return np.array([self._f(X=x) for x in X])
@property
@informative_failure_on_none
def required_dimensionality(self) -> Optional[int]:
"""Required dimensionality of input to this function."""
return self._required_dimensionality
@property
@informative_failure_on_none
def domain(self) -> List[Tuple[float, ...]]:
"""Domain on which function is evaluated.
The list is of the same length as the dimensionality of the inputs,
where each element of the list is a tuple corresponding to the min
and max of the domain for that dimension.
"""
return self._domain
@property
@informative_failure_on_none
def minimums(self) -> List[Tuple[float, ...]]:
"""List of global minimums.
Each element of the list is a d-tuple, where d is the dimensionality
of the inputs. There may be more than one global minimums.
"""
return self._minimums
@property
@informative_failure_on_none
def maximums(self) -> List[Tuple[float, ...]]:
"""List of global minimums.
Each element of the list is a d-tuple, where d is the dimensionality
of the inputs. There may be more than one global minimums.
"""
return self._maximums
@property
@informative_failure_on_none
def fmin(self) -> float:
"""Value at global minimum(s)."""
return self._fmin
@property
@informative_failure_on_none
def fmax(self) -> float:
"""Value at global minimum(s)."""
return self._fmax
@classmethod
@abstractmethod
def _f(self, X: np.ndarray) -> float:
"""Implementation of the synthetic function. Must be implemented in subclass.
Args:
X (numpy.ndarray): an n by d array, where n represents the number
of observations and d is the dimensionality of the inputs.
Returns:
numpy.ndarray: an n-dimensional array.
"""
... # pragma: no cover
[docs]class FromBotorch(SyntheticFunction):
def __init__(
self, botorch_synthetic_function: botorch_synthetic.SyntheticTestFunction
) -> None:
self._botorch_function = botorch_synthetic_function
# pyre-fixme[4]: Attribute must be annotated.
self._required_dimensionality = self._botorch_function.dim
# pyre-fixme[4]: Attribute must be annotated.
self._domain = self._botorch_function._bounds
# pyre-fixme[4]: Attribute must be annotated.
self._fmin = self._botorch_function._optimal_value
@property
def name(self) -> str:
return f"{self.__class__.__name__}_{self._botorch_function.__class__.__name__}"
def _f(self, X: np.ndarray) -> float:
# TODO: support batch evaluation
return float(self._botorch_function(X=torch.from_numpy(X)).item())
[docs]def from_botorch(
botorch_synthetic_function: botorch_synthetic.SyntheticTestFunction,
) -> SyntheticFunction:
"""Utility to generate Ax synthetic functions from BoTorch synthetic functions."""
return FromBotorch(botorch_synthetic_function=botorch_synthetic_function)
[docs]class Hartmann6(SyntheticFunction):
"""Hartmann6 function (6-dimensional with 1 global minimum)."""
_required_dimensionality = 6
# pyre-fixme[4]: Attribute must be annotated.
_domain = [(0, 1) for i in range(6)]
_minimums = [(0.20169, 0.150011, 0.476874, 0.275332, 0.311652, 0.6573)]
# pyre-fixme[4]: Attribute must be annotated.
_fmin = -3.32237
_fmax = 0.0
# pyre-fixme[4]: Attribute must be annotated.
_alpha = np.array([1.0, 1.2, 3.0, 3.2])
# pyre-fixme[4]: Attribute must be annotated.
_A = np.array(
[
[10, 3, 17, 3.5, 1.7, 8],
[0.05, 10, 17, 0.1, 8, 14],
[3, 3.5, 1.7, 10, 17, 8],
[17, 8, 0.05, 10, 0.1, 14],
]
)
# pyre-fixme[4]: Attribute must be annotated.
_P = 10 ** (-4) * np.array(
[
[1312, 1696, 5569, 124, 8283, 5886],
[2329, 4135, 8307, 3736, 1004, 9991],
[2348, 1451, 3522, 2883, 3047, 6650],
[4047, 8828, 8732, 5743, 1091, 381],
]
)
@copy_doc(SyntheticFunction._f)
def _f(self, X: np.ndarray) -> float:
y = 0.0
for j, alpha_j in enumerate(self._alpha):
t = 0
for k in range(6):
t += self._A[j, k] * ((X[k] - self._P[j, k]) ** 2)
y -= alpha_j * np.exp(-t)
return float(y)
[docs]class Aug_Hartmann6(Hartmann6):
"""Augmented Hartmann6 function (7-dimensional with 1 global minimum)."""
_required_dimensionality = 7
# pyre-fixme[4]: Attribute must be annotated.
_domain = [(0, 1) for i in range(7)]
# pyre-fixme[15]: `_minimums` overrides attribute defined in `Hartmann6`
# inconsistently.
_minimums = [(0.20169, 0.150011, 0.476874, 0.275332, 0.311652, 0.6573, 1.0)]
# pyre-fixme[4]: Attribute must be annotated.
_fmin = -3.32237
_fmax = 0.0
@copy_doc(SyntheticFunction._f)
def _f(self, X: np.ndarray) -> float:
y = 0.0
alpha_0 = self._alpha[0] - 0.1 * (1 - X[-1])
for j, alpha_j in enumerate(self._alpha):
t = 0
for k in range(6):
t += self._A[j, k] * ((X[k] - self._P[j, k]) ** 2)
if j == 0:
y -= alpha_0 * np.exp(-t)
else:
y -= alpha_j * np.exp(-t)
return float(y)
[docs]class Branin(SyntheticFunction):
"""Branin function (2-dimensional with 3 global minima)."""
_required_dimensionality = 2
_domain = [(-5, 10), (0, 15)]
# pyre-fixme[4]: Attribute must be annotated.
_minimums = [(-np.pi, 12.275), (np.pi, 2.275), (9.42478, 2.475)]
_fmin = 0.397887
_fmax = 308.129
@copy_doc(SyntheticFunction._f)
def _f(self, X: np.ndarray) -> float:
x_1 = X[0]
x_2 = X[1]
return float(
(x_2 - 5.1 / (4 * np.pi**2) * x_1**2 + 5.0 / np.pi * x_1 - 6.0) ** 2
+ 10 * (1 - 1.0 / (8 * np.pi)) * np.cos(x_1)
+ 10
)
[docs]class Aug_Branin(SyntheticFunction):
"""Augmented Branin function (3-dimensional with infinitely many global minima)."""
_required_dimensionality = 3
_domain = [(-5, 10), (0, 15), (0, 1)]
# pyre-fixme[4]: Attribute must be annotated.
_minimums = [(-np.pi, 12.275, 1), (np.pi, 2.275, 1), (9.42478, 2.475, 1)]
_fmin = 0.397887
_fmax = 308.129
@copy_doc(SyntheticFunction._f)
def _f(self, X: np.ndarray) -> float:
x_1 = X[0]
x_2 = X[1]
return float(
(
x_2
- (5.1 / (4 * np.pi**2) - 0.1 * (1 - X[-1])) * x_1**2
+ 5.0 / np.pi * x_1
- 6.0
)
** 2
+ 10 * (1 - 1.0 / (8 * np.pi)) * np.cos(x_1)
+ 10
)
hartmann6 = Hartmann6()
aug_hartmann6 = Aug_Hartmann6()
branin = Branin()
aug_branin = Aug_Branin()
# Synthetic functions constructed from BoTorch.
ackley: SyntheticFunction = from_botorch(
botorch_synthetic_function=botorch_synthetic.Ackley()
)
