Source code for ax.metrics.branin

#!/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.

# pyre-strict

import numpy.typing as npt
from ax.metrics.noisy_function import NoisyFunctionMetric
from ax.utils.common.typeutils import checked_cast
from ax.utils.measurement.synthetic_functions import aug_branin, branin



[docs]
class BraninMetric(NoisyFunctionMetric):

[docs]
    def f(self, x: npt.NDArray) -> float:
        x1, x2 = x
        return checked_cast(float, branin(x1=x1, x2=x2))





[docs]
class NegativeBraninMetric(BraninMetric):

[docs]
    def f(self, x: npt.NDArray) -> float:
        fpos = super().f(x)
        return -fpos





[docs]
class AugmentedBraninMetric(NoisyFunctionMetric):

[docs]
    def f(self, x: npt.NDArray) -> float:
        return checked_cast(float, aug_branin(x))