Source code for ax.metrics.botorch_test_problem

# 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 typing import Any, Optional

import pandas as pd
from ax.core.base_trial import BaseTrial
from import Data
from ax.core.metric import Metric, MetricFetchE, MetricFetchResult
from ax.utils.common.result import Err, Ok

[docs]class BotorchTestProblemMetric(Metric): """A Metric for retriving information from a BotorchTestProblemRunner. A BotorchTestProblemRunner will attach the result of a call to BaseTestProblem.forward per Arm on a given trial, and this Metric will extract the proper value from the resulting tensor given its index. """ def __init__( self, name: str, noise_sd: Optional[float] = None, index: int = 0, ) -> None: super().__init__(name=name) self.noise_sd = noise_sd self.index = index
[docs] def fetch_trial_data(self, trial: BaseTrial, **kwargs: Any) -> MetricFetchResult: try: mean = [ trial.run_metadata["Ys"][name][self.index] for name, arm in trial.arms_by_name.items() ] df = pd.DataFrame( { "arm_name": [name for name, _ in trial.arms_by_name.items()], "metric_name":, "mean": mean, # If no noise_std is returned then Botorch evaluated the true # function "sem": self.noise_sd, "trial_index": trial.index, } ) return Ok(value=Data(df=df)) except Exception as e: return Err( MetricFetchE(message=f"Failed to fetch {}", exception=e) )