Source code for ax.core.map_metric

#!/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 __future__ import annotations

from typing import Dict, Type

from import Data

from ax.core.map_data import MapData
from ax.core.metric import Metric, MetricFetchE, MetricFetchResult
from ax.utils.common.result import Ok, Result
from ax.utils.common.typeutils import checked_cast

MapMetricFetchResult = Result[MapData, MetricFetchE]

[docs]class MapMetric(Metric): """Base class for representing metrics that return `MapData`. The `fetch_trial_data` method is the essential method to override when subclassing, which specifies how to retrieve a Metric, for a given trial. A MapMetric must return a MapData object, which requires (at minimum) the following: Attributes: lower_is_better: Flag for metrics which should be minimized. properties: Properties specific to a particular metric. """ data_constructor: Type[MapData] = MapData @classmethod def _wrap_experiment_data(cls, data: Data) -> Dict[int, MetricFetchResult]: return { trial_index: Ok( value=MapData( df=data.true_df.loc[data.true_df["trial_index"] == trial_index], map_key_infos=checked_cast(MapData, data).map_key_infos, ) ) for trial_index in data.true_df["trial_index"] } @classmethod def _wrap_trial_data_multi(cls, data: Data) -> Dict[str, MetricFetchResult]: return { metric_name: Ok( value=MapData( df=data.true_df.loc[data.true_df["metric_name"] == metric_name], map_key_infos=checked_cast(MapData, data).map_key_infos, ) ) for metric_name in data.true_df["metric_name"] } @classmethod def _wrap_experiment_data_multi( cls, data: Data ) -> Dict[int, Dict[str, MetricFetchResult]]: # pyre-fixme[7] return { trial_index: { metric_name: Ok( value=MapData( df=data.true_df.loc[ (data.true_df["trial_index"] == trial_index) & (data.true_df["metric_name"] == metric_name) ], map_key_infos=checked_cast(MapData, data).map_key_infos, ) ) for metric_name in data.true_df["metric_name"] } for trial_index in data.true_df["trial_index"] }