experiment.py

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

from __future__ import annotations

import inspect
import logging
import warnings
from collections import defaultdict
from collections.abc import Hashable, Iterable, Mapping, Sequence
from copy import deepcopy
from datetime import datetime
from functools import partial, reduce
from typing import Any, cast, Self, Union

import ax.core.observation as observation
import pandas as pd
from ax.core.arm import Arm
from ax.core.auxiliary import (
    AuxiliaryExperiment,
    AuxiliaryExperimentPurpose,
    AuxiliaryExperimentValidation,
    TransferLearningMetadata,
)
from ax.core.base_trial import BaseTrial
from ax.core.batch_trial import BatchTrial
from ax.core.data import combine_data_rows_favoring_recent, Data
from ax.core.derived_metric import DerivedMetric
from ax.core.experiment_status import ExperimentStatus
from ax.core.generator_run import GeneratorRun
from ax.core.llm_provider import LLMMessage
from ax.core.metric import Metric, MetricFetchE, MetricFetchResult
from ax.core.objective import Objective
from ax.core.optimization_config import (
    MultiObjectiveOptimizationConfig,
    OptimizationConfig,
)
from ax.core.parameter import DerivedParameter, Parameter
from ax.core.parameter_constraint import ParameterConstraint
from ax.core.runner import Runner
from ax.core.search_space import SearchSpace
from ax.core.trial import Trial
from ax.core.trial_status import (
    DEFAULT_STATUSES_TO_WARM_START,
    STATUSES_EXPECTING_DATA,
    TrialStatus,
)
from ax.core.types import ComparisonOp, TParameterization
from ax.exceptions.core import (
    AxError,
    OptimizationNotConfiguredError,
    RunnerNotFoundError,
    UnsupportedError,
    UserInputError,
)
from ax.utils.common.base import Base
from ax.utils.common.constants import Keys
from ax.utils.common.docutils import copy_doc
from ax.utils.common.executils import retry_on_exception
from ax.utils.common.logger import _round_floats_for_logging, get_logger
from ax.utils.common.result import Err, Ok
from pyre_extensions import assert_is_instance, none_throws

logger: logging.Logger = get_logger(__name__)

NO_RETRY_EXCEPTIONS: tuple[type[Exception], ...] = (
    cast(type[Exception], RunnerNotFoundError),
    cast(type[Exception], NotImplementedError),
    cast(type[Exception], UnsupportedError),
)

ROUND_FLOATS_IN_LOGS_TO_DECIMAL_PLACES: int = 6

round_floats_for_logging: partial[Any] = partial(
    _round_floats_for_logging,
    decimal_places=ROUND_FLOATS_IN_LOGS_TO_DECIMAL_PLACES,
)
METRIC_DF_COLNAMES: Mapping[Hashable, str] = {
    "name": "Name",
    "type": "Type",
    "goal": "Goal",
    "bound": "Bound",
    "lower_is_better": "Lower is Better",
}


class Experiment(Base):
    """Base class for defining an experiment."""

    def __init__(
        self,
        search_space: SearchSpace,
        name: str | None = None,
        optimization_config: OptimizationConfig | None = None,
        metrics: Sequence[Metric] | None = None,
        runner: Runner | None = None,
        status_quo: Arm | None = None,
        description: str | None = None,
        is_test: bool = False,
        experiment_type: str | None = None,
        properties: dict[str, Any] | None = None,
        default_data_type: Any = None,
        auxiliary_experiments_by_purpose: None
        | (dict[AuxiliaryExperimentPurpose, list[AuxiliaryExperiment]]) = None,
        default_trial_type: str | None = None,
        tracking_metrics: list[Metric] | None = None,
    ) -> None:
        """Inits Experiment.

        Args:
            search_space: Search space of the experiment.
            name: Name of the experiment.
            optimization_config: Optimization config of the experiment.
            runner: Default runner used for trials on this experiment.
            status_quo: Arm representing existing "control" arm.
            description: Description of the experiment.
            is_test: Mark experiment as test in metadata. This flag is meant purely
                for development and integration testing purposes. Leave as False for
                live experiments.
            experiment_type: The class of experiments this one belongs to.
            properties: Dictionary of this experiment's properties.  It is meant to
                only store primitives that pertain to Ax experiment state. Any trial
                deployment-related information and modeling-layer configuration
                should be stored elsewhere, e.g. in ``run_metadata`` of the trials.
            default_data_type: Deprecated and ignored.
            auxiliary_experiments_by_purpose: Dictionary of auxiliary experiments
                for different purposes (e.g., transfer learning).
            tracking_metrics: Deprecated. Use ``metrics`` directly instead.
        """
        if default_data_type is not None:
            warnings.warn(
                "default_data_type is deprecated and will be removed in a future "
                "release.",
                DeprecationWarning,
                stacklevel=2,
            )
        self._search_space: SearchSpace = search_space
        self._status_quo: Arm | None = None

        self._name = name
        self.description = description
        self._runner = runner
        self.is_test: bool = is_test

        self.data: Data = Data()
        self._experiment_type: str | None = experiment_type
        self._optimization_config: OptimizationConfig | None = None
        self._metrics: dict[str, Metric] = {}
        self._time_created: datetime = datetime.now()
        self._status: ExperimentStatus | None = None
        self._trials: dict[int, BaseTrial] = {}
        self._properties: dict[str, Any] = properties or {}

        # Initialize trial type to runner mapping
        self._default_trial_type = default_trial_type
        self._trial_type_to_runner: dict[str | None, Runner | None] = {
            default_trial_type: runner
        }
        # Maps each trial type to the set of metric names relevant to that type.
        # This is the complement of _trial_type_to_runner and is used for
        # multi-type experiments where different metrics apply to different
        # trial types.
        self._trial_type_to_metric_names: dict[str, set[str]] = {}
        # Used to keep track of whether any trials on the experiment
        # specify a TTL. Since trials need to be checked for their TTL's
        # expiration often, having this attribute helps avoid unnecessary
        # TTL checks for experiments that do not use TTL.
        self._trials_have_ttl = False
        # Make sure all statuses appear in this dict, to avoid key errors.
        self._trial_indices_by_status: dict[TrialStatus, set[int]] = {
            status: set() for status in TrialStatus
        }
        self._arms_by_signature: dict[str, Arm] = {}
        self._arms_by_name: dict[str, Arm] = {}

        # Used to keep track of auxiliary experiments that were removed.
        self._initial_auxiliary_experiments_by_purpose: dict[
            AuxiliaryExperimentPurpose, list[AuxiliaryExperiment]
        ] = auxiliary_experiments_by_purpose or {}

        # Only tracks active auxiliary experiments.
        self.auxiliary_experiments_by_purpose: dict[
            AuxiliaryExperimentPurpose, list[AuxiliaryExperiment]
        ] = {
            purpose: [
                auxiliary_experiment
                for auxiliary_experiment in auxiliary_experiments
                if auxiliary_experiment.is_active
            ]
            for (
                purpose,
                auxiliary_experiments,
            ) in self._initial_auxiliary_experiments_by_purpose.items()
        }

        # Attach metrics from both tracking_metrics and metrics, preferring metrics if
        # a naming collision occurs.
        for m in [*(tracking_metrics or []), *(metrics or [])]:
            self._metrics[m.name] = m
            if self._default_trial_type is not None:
                self._trial_type_to_metric_names.setdefault(
                    self._default_trial_type, set()
                ).add(m.name)

        # call setters defined below
        self.status_quo = status_quo
        if optimization_config is not None:
            # Auto-register placeholder Metric objects for any optimization
            # config metric names not already on the experiment. The
            # ``optimization_config`` setter (below) validates that every
            # referenced metric name exists in ``self._metrics``, so
            # placeholders are needed here to satisfy that check. These
            # placeholders are later replaced in two ways:
            #   1. During SQA decoding, the decoder overwrites them with
            #      properly typed metrics from the database
            #      (see ``sqa_store/decoder.py``).
            #   2. Any call to ``add_metric`` or ``update_metric`` with a
            #      real Metric of the same name will overwrite a placeholder.
            for name in optimization_config.metric_names:
                if name not in self._metrics:
                    self._metrics[name] = Metric(name=name)
            self.optimization_config = optimization_config

        # Keyed on tuple[trial_index, metric_name].
        self._metric_fetching_errors: dict[
            tuple[int, str], dict[str, Union[int, str]]
        ] = {}

    @property
    def has_name(self) -> bool:
        """Return true if experiment's name is not None."""
        return self._name is not None

    @property
    def name(self) -> str:
        """Get experiment name. Throws if name is None."""
        if self._name is None:
            raise ValueError("Experiment's name is None.")
        return self._name

    @name.setter
    def name(self, name: str) -> None:
        """Set experiment name."""
        self._name = name

    @property
    def time_created(self) -> datetime:
        """Creation time of the experiment."""
        return self._time_created

    @property
    def experiment_type(self) -> str | None:
        """The type of the experiment."""
        return self._experiment_type

    @experiment_type.setter
    def experiment_type(self, experiment_type: str | None) -> None:
        """Set the type of the experiment."""
        self._experiment_type = experiment_type

    @property
    def status(self) -> ExperimentStatus | None:
        """The current status of the experiment.

        Status tracks the high-level lifecycle phase of the experiment:
        DRAFT, INITIALIZATION, OPTIMIZATION, COMPLETED.

        For new experiments, status defaults to DRAFT. For legacy experiments
        that were created before the status field was added, status may be None.
        """
        return self._status

    @status.setter
    def status(self, status: ExperimentStatus | None) -> None:
        """Set the status of the experiment.

        Args:
            status: The new status for the experiment.
        """
        self._status = status

    @staticmethod
    def experiment_status_from_generator_runs(
        generator_runs: list[GeneratorRun],
    ) -> ExperimentStatus | None:
        """Extract and validate suggested experiment status from generator runs.

        Collects the suggested_experiment_status directly from the GeneratorRun
        objects, validates that all runs suggest the same status, and returns
        that status.

        Args:
            generator_runs: List of generator runs to extract statuses from.

        Returns:
            The suggested experiment status that all generator runs agree on,
            or None if no statuses were found or if there are conflicting statuses.
        """
        suggested_statuses: set[ExperimentStatus] = set()
        for gr in generator_runs:
            if gr.suggested_experiment_status is not None:
                suggested_statuses.add(gr.suggested_experiment_status)

        if len(suggested_statuses) > 1:
            # TODO: Consider making this invalid state an actual error once
            # related development is completed.
            logger.warning(
                "Multiple different suggested experiment statuses found: "
                f"{suggested_statuses}. "
                "All generator runs used in a single gen() call should suggest the "
                "same experiment status. Skipping updating experiment status."
            )
            return None

        if len(suggested_statuses) == 0:
            return None

        return suggested_statuses.pop()

    @property
    def search_space(self) -> SearchSpace:
        """The search space for this experiment.

        When setting a new search space, all parameter names and types
        must be preserved. However, if no trials have been created, all
        modifications are allowed.
        """
        # TODO: maybe return a copy here to guard against implicit changes
        return self._search_space

    @search_space.setter
    def search_space(self, search_space: SearchSpace) -> None:
        # Allow all modifications when no trials present.
        if not hasattr(self, "_search_space") or len(self.trials) < 1:
            self._search_space = search_space
            return

        # At least 1 trial is present.
        if self.immutable_search_space_and_opt_config:
            raise UnsupportedError(
                "Modifications of search space are disabled by the "
                f"`{Keys.IMMUTABLE_SEARCH_SPACE_AND_OPT_CONF.value}` "
                "property that is set to `True` on this experiment."
            )
        if len(search_space.parameters) < len(self._search_space.parameters):
            raise ValueError(
                "New search_space must contain all parameters in the existing."
            )
        for param_name, parameter in search_space.parameters.items():
            if param_name not in self._search_space.parameters:
                raise ValueError(
                    f"Cannot add new parameter `{param_name}` because "
                    "it is not defined in the existing search space."
                )
            elif (
                parameter.parameter_type
                != self._search_space.parameters[param_name].parameter_type
            ):
                raise ValueError(
                    f"Expected parameter `{param_name}` to be of type "
                    f"{self._search_space.parameters[param_name].parameter_type}, "
                    f"got {parameter.parameter_type}."
                )
        self._search_space = search_space

    def add_parameters_to_search_space(
        self,
        parameters: Sequence[Parameter],
        status_quo_values: TParameterization | None = None,
        parameter_constraints: Sequence[ParameterConstraint] | None = None,
    ) -> None:
        """
        Add new parameters to the experiment's search space. This allows extending
        the search space after the experiment has run some trials.

        Backfill values must be provided for all new parameters if the experiment has
        already run some trials. The backfill values represent the parameter values
        that were used in the existing trials.

        Args:
            parameters: A sequence of parameter configurations to add to the search
                space.
            status_quo_values: Optional parameter values for the new parameters to
                use in the status quo (baseline) arm, if one is defined.
            parameter_constraints: Optional sequence of typed ParameterConstraint
                objects to add to the search space after the parameters are added.
        """
        status_quo_values = status_quo_values or {}

        # Additional checks iff a trial exists
        if len(self.trials) != 0:
            if any(
                (parameter.backfill_value is None)
                and not isinstance(parameter, DerivedParameter)
                for parameter in parameters
            ):
                raise UserInputError(
                    "Must provide backfill values for all new parameters (except "
                    "DerivedParameters) when adding parameters to an experiment "
                    "with existing trials."
                )

        # Validate status quo values
        status_quo = self._status_quo
        if status_quo_values is not None and status_quo is None:
            logger.warning(
                "Status quo values specified, but experiment does not have a "
                "status quo. Ignoring provided status quo values."
            )
        if status_quo is not None:
            parameter_names = {parameter.name for parameter in parameters}
            status_quo_parameters = status_quo_values.keys()
            disabled_parameters = {
                parameter.name
                for parameter in self._search_space.parameters.values()
                if parameter.is_disabled
            }
            extra_status_quo_values = status_quo_parameters - parameter_names
            if extra_status_quo_values:
                logger.warning(
                    "Status quo value provided for parameters "
                    f"`{extra_status_quo_values}` which is are being added to "
                    "the search space. Ignoring provided status quo values."
                )
            missing_status_quo_values = (
                parameter_names - disabled_parameters - status_quo_parameters
            )
            if missing_status_quo_values:
                raise UserInputError(
                    "No status quo value provided for parameters "
                    f"`{missing_status_quo_values}` which are being added to "
                    "the search space."
                )
            for parameter_name, value in status_quo_values.items():
                status_quo._parameters[parameter_name] = value

        # Add parameters to search space
        self._search_space.add_parameters(parameters)

        # Add parameter constraints to search space
        if parameter_constraints:
            self._search_space.add_parameter_constraints(list(parameter_constraints))

    def disable_parameters_in_search_space(
        self, default_parameter_values: TParameterization
    ) -> None:
        """
        Disable parameters in the experiment. This allows narrowing the search space
        after the experiment has run some trials.

        When parameters are disabled, they are effectively removed from the search
        space for future trial generation. Existing trials remain valid, and the
        disabled parameters are replaced with fixed default values for all subsequent
        trials.

        Args:
            default_parameter_values: Fixed values to use for the disabled parameters
                in all future trials. These values will be used for the parameter in
                all subsequent trials.
        """
        self._search_space.disable_parameters(default_parameter_values)

    @property
    def status_quo(self) -> Arm | None:
        """The existing arm that new arms will be compared against."""
        return self._status_quo

    @status_quo.setter
    def status_quo(self, status_quo: Arm | None) -> None:
        # Make status_quo immutable once any trial has been created.
        if self._status_quo is not None and len(self.trials) > 0:
            raise UnsupportedError(
                "Modifications of status_quo are disabled after trials have been "
                "created."
            )
        if status_quo == self.status_quo:
            return  # No need to update the SQ arm.

        if status_quo is not None:
            self.search_space.check_types(
                parameterization=status_quo.parameters,
                allow_extra_params=False,
                raise_error=True,
            )
            self.search_space.check_all_parameters_present(
                parameterization=status_quo.parameters, raise_error=True
            )

            # Compute a unique name if "status_quo" is taken
            name = "status_quo"
            sq_idx = 0
            arms_by_name = self.arms_by_name
            while name in arms_by_name:
                name = f"status_quo_e{sq_idx}"
                sq_idx += 1
            self._name_and_store_arm_if_not_exists(arm=status_quo, proposed_name=name)

        self._status_quo = status_quo

    @property
    def runner(self) -> Runner | None:
        """Default runner used for trials on this experiment."""
        return self._runner

    @runner.setter
    def runner(self, runner: Runner | None) -> None:
        """Set the default runner and update trial type mapping."""
        self._runner = runner
        if runner is not None:
            self._trial_type_to_runner[self._default_trial_type] = runner
        else:
            self._trial_type_to_runner = {None: None}

    @runner.deleter
    def runner(self) -> None:
        """Delete the runner."""
        self._runner = None
        self._trial_type_to_runner = {None: None}

    @property
    def parameters(self) -> dict[str, Parameter]:
        """The parameters in the experiment's search space."""
        return self.search_space.parameters

    @property
    def arms_by_name(self) -> dict[str, Arm]:
        """The arms belonging to this experiment, by their name."""
        return self._arms_by_name

    @property
    def arms_by_signature(self) -> dict[str, Arm]:
        """The arms belonging to this experiment, by their signature."""
        return self._arms_by_signature

    @property
    def arms_by_signature_for_deduplication(self) -> dict[str, Arm]:
        """The arms belonging to this experiment that should be used for deduplication
        in ``GenerationStrategy``, by their signature.

        In its current form, this includes all arms except for those that are
        associated with a ``FAILED`` trial.
        - The ``CANDIDATE``, ``STAGED``, ``RUNNING``, and ``ABANDONED`` arms are
        included as pending points during generation, so they should be less likely
         to get suggested by the model again.
        - The ``EARLY_STOPPED`` and ``COMPLETED`` trials were already evaluated, so
        the model will have data for these and is unlikely to suggest them again.
        """
        arms_dict = self.arms_by_signature.copy()
        for trial in self.trials_by_status[TrialStatus.FAILED]:
            for arm in trial.arms:
                arms_dict.pop(arm.signature, None)
        return arms_dict

    @property
    def sum_trial_sizes(self) -> int:
        """Sum of numbers of arms attached to each trial in this experiment."""
        return reduce(lambda a, b: a + len(b.arms_by_name), self._trials.values(), 0)

    @property
    def num_abandoned_arms(self) -> int:
        """How many arms attached to this experiment are abandoned."""
        abandoned = set()
        for trial in self.trials.values():
            for x in trial.abandoned_arms:
                abandoned.add(x)
        return len(abandoned)

    @property
    def optimization_config(self) -> OptimizationConfig | None:
        """The experiment's optimization config."""
        return self._optimization_config

    @optimization_config.setter
    def optimization_config(self, optimization_config: OptimizationConfig) -> None:
        if (
            len(self.trials) > 0
            and getattr(self, "_optimization_config", None) is not None
            and self.immutable_search_space_and_opt_config
        ):
            raise UnsupportedError(
                "Modifications of optimization config are disabled by the "
                f"`{Keys.IMMUTABLE_SEARCH_SPACE_AND_OPT_CONF.value}` "
                "property that is set to `True` on this experiment."
            )
        for metric_name in optimization_config.metric_names:
            if metric_name not in self._metrics:
                raise ValueError(
                    f"Metric '{metric_name}' referenced in optimization config "
                    "but not found on experiment. Add it first with add_metric()."
                )
        self._optimization_config = optimization_config
        resolved_trial_type = self._resolve_trial_type(None)
        if resolved_trial_type is not None:
            for metric_name in optimization_config.metric_names:
                self._trial_type_to_metric_names.setdefault(
                    resolved_trial_type, set()
                ).add(metric_name)

    @property
    def is_moo_problem(self) -> bool:
        """Whether the experiment's optimization config contains multiple objectives."""
        if self.optimization_config is None:
            return False
        return none_throws(self.optimization_config).is_moo_problem

    @property
    def is_bope_problem(self) -> bool:
        """Whether this experiment is a BO with Preference Exploration (BOPE)
        experiment.

        An experiment is considered a preference learning experiment if:
        1. It has a PreferenceOptimizationConfig as its optimization config, OR
        2. It has a PE_EXPERIMENT (preference exploration) auxiliary experiment attached

        Returns:
            True if this is a preference learning experiment, False otherwise.
        """
        # Check if optimization config indicates preference learning
        if self.optimization_config is not None:
            if self.optimization_config.is_bope_problem:
                return True

        # Check if experiment has a PE_EXPERIMENT auxiliary experiment
        return bool(
            self.auxiliary_experiments_by_purpose.get(
                AuxiliaryExperimentPurpose.PE_EXPERIMENT, []
            )
        )

    @property
    def immutable_search_space_and_opt_config(self) -> bool:
        """Boolean representing whether search space and metrics on this experiment
        are mutable (by default they are).

        NOTE: For experiments with immutable search spaces and metrics, generator
        runs will not store copies of search space and metrics, which improves
        storage layer performance. Not keeping copies of those on generator runs
        also disables keeping track of changes to search space and metrics,
        thereby necessitating that those attributes be immutable on experiment.
        """
        return self._properties.get(Keys.IMMUTABLE_SEARCH_SPACE_AND_OPT_CONF, False)

    @property
    def llm_messages(self) -> list[LLMMessage]:
        """LLM conversation messages stored on this experiment.

        These messages capture the conversation history for LLM-integrated
        optimization workflows such as LILO (Language-in-the-Loop Optimization)
        and LLAMBO. They include system prompts, user feedback, and assistant
        responses that inform candidate generation.
        """
        return self._properties.get(Keys.LLM_MESSAGES, [])

    @llm_messages.setter
    def llm_messages(self, messages: list[LLMMessage]) -> None:
        self._properties[Keys.LLM_MESSAGES] = messages

    def get_llm_messages_with_experiment_summary(
        self,
    ) -> list[LLMMessage]:
        """Retrieve a copy of LLM messages with an experiment summary appended.

        Returns a deep copy of the stored LLM messages with a user message
        appended containing a markdown summary of the experiment's search
        space, optimization config, and observed trial data. This is useful
        for LLM-based workflows (e.g., LLAMBO, LILO) that condition on both
        the conversation history and observed trial outcomes.

        The summary always includes experiment metadata (search space and
        optimization config). Arm parameters and trial results are included
        when completed, early-stopped, or running trial data exists.

        Returns:
            A list of ``LLMMessage`` objects suitable for passing to an
            ``LLMProvider.generate()`` call.
        """
        messages = [
            LLMMessage(role=m.role, content=m.content, metadata=deepcopy(m.metadata))
            for m in self.llm_messages
        ]

        experiment_summary = self._format_experiment_summary_for_llm()
        if experiment_summary:
            messages.append(LLMMessage(role="user", content=experiment_summary))

        return messages

    def _format_experiment_summary_for_llm(self) -> str:
        """Format experiment summary as text for inclusion in LLM messages.

        Produces a human-readable markdown summary of the experiment including:
        - Experiment name
        - Search space parameter definitions
        - Optimization configuration (objectives and constraints)
        - Arm parameterizations (deduplicated), when data is available
        - Observed trial results with metric values
          (mean ± 95% CI when SEM is available), when data is available

        The experiment metadata (name, search space, optimization config) is
        always included. Arm parameters and trial results are appended only
        when completed, early-stopped, or running trial data exists.

        Returns:
            A markdown-formatted string summary of the experiment.
        """
        sections: list[str] = []

        # Experiment name
        sections.append(f"## Experiment: {self._name or 'Unnamed'}")

        # Search space parameters
        param_lines: list[str] = ["### Search Space"]
        for param in self.search_space.parameters.values():
            param_lines.append(f"- {param}")
        if self.search_space.is_hierarchical:
            ss_str = self.search_space.hierarchical_structure_str(
                parameter_names_only=True
            )
            param_lines.append(
                "\n**Hierarchical Structure:**\n"
                "Not all parameters are active at the same time. "
                "The active parameters depend on the values of other parameters "
                "as described in the tree below. Each top-level entry is a "
                "parameter. Indented entries in parentheses, e.g. (0), are "
                "possible values for that parameter. The parameters listed "
                "under a value are the ones that become active when the parent "
                "parameter takes that value. Only provide values for the "
                f"active parameters.\n{ss_str}"
            )
        sections.append("\n".join(param_lines))

        # Optimization config
        opt_config = self.optimization_config
        if opt_config is not None:
            opt_lines: list[str] = ["### Optimization Config"]

            objective = opt_config.objective
            for metric_name, weight in objective.metric_weights:
                direction = "minimize" if weight < 0 else "maximize"
                opt_lines.append(f"- Objective: {direction} `{metric_name}`")

            for constraint in opt_config.outcome_constraints:
                opt_lines.append(f"- Constraint: `{constraint.expression}` ")
            sections.append("\n".join(opt_lines))

        # Trial data
        data = self.lookup_data()
        if data.df.empty:
            return "\n\n".join(sections)

        # Build arm parameters table (deduplicated by arm_name).
        arm_param_rows: list[dict[str, str]] = []
        seen_arms: set[str] = set()
        # Build trial results table.
        result_rows: list[dict[str, str]] = []
        has_ci = False

        for trial_index, trial in self.trials.items():
            if trial.status not in (
                TrialStatus.COMPLETED,
                TrialStatus.EARLY_STOPPED,
                TrialStatus.RUNNING,
            ):
                continue
            trial_data = data.df[data.df["trial_index"] == trial_index]
            if trial_data.empty:
                continue
            for arm_name in trial_data["arm_name"].unique():
                arm = trial.arms_by_name.get(arm_name)
                if arm is None:
                    continue
                # Deduplicated arm parameters.
                if arm_name not in seen_arms:
                    seen_arms.add(arm_name)
                    arm_param_rows.append(
                        {
                            "arm_name": arm_name,
                            "parameters": str(arm.parameters),
                        }
                    )
                # Metric results for this trial/arm.
                arm_data = trial_data[trial_data["arm_name"] == arm_name]
                row: dict[str, str] = {
                    "trial_index": str(trial_index),
                    "arm_name": arm_name,
                    "trial_status": trial.status.name,
                }
                for _, metric_row in arm_data.iterrows():
                    metric_name = metric_row["metric_name"]
                    mean = metric_row["mean"]
                    sem = metric_row.get("sem")
                    if sem is not None and not pd.isna(sem):
                        ci = 1.96 * sem
                        row[metric_name] = f"{mean:.4g} ± {ci:.4g}"
                        has_ci = True
                    else:
                        row[metric_name] = str(mean)
                result_rows.append(row)

        if not result_rows:
            return "\n\n".join(sections)

        # Arm parameters section.
        params_df = pd.DataFrame(arm_param_rows)
        sections.append(f"### Arm Parameters\n\n{params_df.to_markdown(index=False)}")

        # Trial results section.
        results_df = pd.DataFrame(result_rows)
        ci_note = "\n\nMetric values with ± denote mean ± 95% CI." if has_ci else ""
        sections.append(
            f"### Trial Results{ci_note}\n\n{results_df.to_markdown(index=False)}"
        )

        return "\n\n".join(sections)

    @property
    def tracking_metrics(self) -> list[Metric]:
        """Metrics that are tracked but not part of the optimization config."""
        if self._optimization_config is None:
            return list(self._metrics.values())
        opt_metric_names = self._optimization_config.metric_names
        return [m for name, m in self._metrics.items() if name not in opt_metric_names]

    @property
    def optimization_config_metrics(self) -> list[Metric]:
        """Metrics that are part of the optimization config."""
        if self._optimization_config is None:
            return []
        opt_metric_names = self._optimization_config.metric_names
        return self.get_metrics(metric_names=list(opt_metric_names))

    def get_metric(self, name: str) -> Metric:
        """Get a single Metric by name.

        Args:
            name: The name of the metric to retrieve.

        Returns:
            The Metric object with the given name.

        Raises:
            KeyError: If no metric with the given name exists.
        """
        if name not in self._metrics:
            raise KeyError(
                f"Metric '{name}' not found. Available: {list(self._metrics.keys())}"
            )
        return self._metrics[name]

    def _resolve_trial_type(self, trial_type: str | None) -> str | None:
        """Resolve an explicit or default trial type and validate it.

        Returns ``trial_type`` if explicitly provided (after validating via
        ``supports_trial_type``), falls back to ``_default_trial_type`` when
        available, and raises ``ValueError`` if this experiment uses trial types
        (``_trial_type_to_metric_names`` is non-empty) but none could be
        resolved.

        Args:
            trial_type: The explicitly provided trial type, or ``None``.

        Returns:
            The resolved trial type, which may be ``None`` for single-type
            experiments.

        Raises:
            ValueError: If ``trial_type`` is provided but not supported, or if
                no trial type could be resolved for a multi-type experiment.
        """
        if trial_type is not None:
            if not self.supports_trial_type(trial_type):
                raise ValueError(f"`{trial_type}` is not a supported trial type.")
            return trial_type
        if self._default_trial_type is not None:
            return self._default_trial_type
        if self._trial_type_to_metric_names:
            raise ValueError(
                "This experiment has trial-type-aware metrics but no "
                "`trial_type` was specified and no `default_trial_type` is set. "
                "Please specify a `trial_type`."
            )
        return None

    def add_metric(self, metric: Metric, trial_type: str | None = None) -> Self:
        """Add a new metric to the experiment.

        Metrics that are not referenced by the experiment's optimization config
        (i.e. not part of an objective or outcome constraint) are automatically
        treated as tracking metrics.

        Args:
            metric: Metric to be added.
            trial_type: If provided, associates the metric with this trial type.
                When ``None`` and a ``default_trial_type`` is set, defaults to
                the default trial type.

        Raises:
            ValueError: If the metric already exists, the trial type is not
                supported, or trial types are in use but none could be resolved.
        """
        if metric.name in self._metrics:
            raise ValueError(
                f"Metric `{metric.name}` already defined on experiment. "
                "Use `update_metric` to update an existing metric definition."
            )
        trial_type = self._resolve_trial_type(trial_type)
        if trial_type is not None:
            self._trial_type_to_metric_names.setdefault(trial_type, set()).add(
                metric.name
            )
        self._metrics[metric.name] = metric
        return self

    def add_tracking_metric(
        self,
        metric: Metric,
        trial_type: str | None = None,
        canonical_name: str | None = None,
    ) -> Self:
        """*Deprecated.* Use ``add_metric`` instead."""
        warnings.warn(
            "add_tracking_metric is deprecated. Use add_metric instead.",
            DeprecationWarning,
            stacklevel=2,
        )
        return self.add_metric(metric, trial_type=trial_type)

    def add_tracking_metrics(
        self,
        metrics: list[Metric],
        metrics_to_trial_types: dict[str, str] | None = None,
        canonical_names: dict[str, str] | None = None,
    ) -> Experiment:
        """*Deprecated.* Use ``add_metric`` instead."""
        warnings.warn(
            "add_tracking_metrics is deprecated. Use add_metric instead.",
            DeprecationWarning,
            stacklevel=2,
        )
        metrics_to_trial_types = metrics_to_trial_types or {}
        for metric in metrics:
            canonical_name = (canonical_names or {}).get(metric.name)
            self.add_tracking_metric(
                metric=metric,
                trial_type=metrics_to_trial_types.get(metric.name),
                canonical_name=canonical_name,
            )
        return self

    def update_metric(self, metric: Metric, trial_type: str | None = None) -> Self:
        """Redefine a metric that already exists on the experiment.

        Args:
            metric: New metric definition.
            trial_type: If provided, reassociates the metric with this trial
                type. When ``None``, keeps the metric's existing trial type.
        """
        if metric.name not in self._metrics:
            raise ValueError(f"Metric `{metric.name}` doesn't exist on experiment.")
        if trial_type is not None:
            trial_type = self._resolve_trial_type(trial_type)
            # Remove from any existing trial type set
            for names in self._trial_type_to_metric_names.values():
                names.discard(metric.name)
            # Add to new trial type set
            self._trial_type_to_metric_names.setdefault(trial_type, set()).add(
                metric.name
            )
        self._metrics[metric.name] = metric
        return self

    def update_tracking_metric(
        self,
        metric: Metric,
        trial_type: str | None = None,
        canonical_name: str | None = None,
    ) -> Experiment:
        """*Deprecated.* Use ``update_metric`` instead."""
        warnings.warn(
            "update_tracking_metric is deprecated. Use update_metric instead.",
            DeprecationWarning,
            stacklevel=2,
        )
        return self.update_metric(metric, trial_type=trial_type)

    def remove_metric(self, metric_name: str) -> Self:
        """Remove a metric from the experiment.

        Args:
            metric_name: Unique name of metric to remove.

        Raises:
            ValueError: If the metric is referenced by the optimization config.
        """
        if metric_name not in self._metrics:
            raise ValueError(f"Metric `{metric_name}` doesn't exist on experiment.")
        if (
            self._optimization_config is not None
            and metric_name in self._optimization_config.metric_names
        ):
            raise ValueError(
                f"Metric `{metric_name}` is referenced by the optimization config "
                "and cannot be removed. Update the optimization config first."
            )
        # Clean up _trial_type_to_metric_names
        for names in self._trial_type_to_metric_names.values():
            names.discard(metric_name)
        del self._metrics[metric_name]
        return self

    def remove_tracking_metric(self, metric_name: str) -> Experiment:
        """*Deprecated.* Use ``remove_metric`` instead."""
        warnings.warn(
            "remove_tracking_metric is deprecated. Use remove_metric instead.",
            DeprecationWarning,
            stacklevel=2,
        )
        return self.remove_metric(metric_name)

    @property
    def metrics(self) -> dict[str, Metric]:
        """The metrics attached to the experiment."""
        return dict(self._metrics)

    @property
    def signature_to_metric(self) -> dict[str, Metric]:
        """Returns a dictionary of metrics attached to the experiment, keyed by
        their signature. Useful for cases that require accessing metric objects
        by their signature (e.g. plotting from observation data/adapter).
        """
        return {metric.signature: metric for metric in self.metrics.values()}

    def _metrics_by_class(
        self, metrics: Sequence[Metric] | None = None
    ) -> dict[type[Metric], list[Metric]]:
        metrics_by_class: dict[type[Metric], list[Metric]] = defaultdict(list)
        for metric in metrics or list(self.metrics.values()):
            # By default, all metrics are grouped by their class for fetch;
            # however, for some metrics, `fetch_trial_data_multi` of a
            # superclass is used for fetch the subclassing metrics' data. In
            # those cases, "fetch_multi_group_by_metric" property on metric
            # will be set to a class other than its own (likely a superclass).
            metrics_by_class[metric.fetch_multi_group_by_metric].append(metric)
        return metrics_by_class

    def get_metrics(self, metric_names: list[str] | None) -> list[Metric]:
        """Get a list of metrics from the experiment by name.

        Args:
            metric_names: List of metric names to retrieve. If None, returns all metrics
                defined on the experiment.

        Returns:
            List of Metric objects corresponding to the requested metric names,
            deduplicated so the same `Metric` does not occur twice.

        Raises:
            UserInputError: If any of the requested metric names are not found in the
                experiment.
        """
        if metric_names is None:
            return list(self.metrics.values())
        try:
            return [self.metrics[name] for name in metric_names]
        except KeyError as e:
            raise AxError(
                "One of the requested metrics was not present on the "
                f"experiment; original error: {e}."
            )

    def bulk_configure_metrics_of_class(
        self,
        metric_class: type[Metric],
        attributes_to_update: dict[str, Any],
    ) -> None:
        """Apply the same set of updates to all metrics of a specified type at once.

        Args:
            metric_class: the class of metric to update
            attributes_to_update: A dictionary that maps the parameter to be updated
                with it's new value.

        Raises:
            * If there are no metrics of the specified class on the experiment
            * If any of the specified parameters to update are not args in the
            metric class's initialization method

        Note: All metrics of the specified class will receive the same update. For
        custom updates to specific metrics, use `update_tracking_metric`
        """
        metrics_of_class = self._metrics_by_class(metrics=list(self.metrics.values()))[
            metric_class
        ]
        if len(metrics_of_class) == 0:
            raise UserInputError(
                f"No metrics of class {metric_class} found on experiment."
            )
        metric_attributes = set(
            inspect.signature(metric_class.__init__).parameters.keys()
        )
        if not set(attributes_to_update.keys()).issubset(metric_attributes):
            raise (
                UserInputError(
                    f"Metric class {metric_class} does not contain the requested "
                    "attributes to update. Requested updates to attributes: "
                    f"{set(attributes_to_update.keys())} but metric class defines "
                    f"{metric_attributes}."
                )
            )
        # Update each metric on the experiment with the specified new values
        # for each parameter to update
        for metric in metrics_of_class:
            new_metric = metric
            for param_name, param_value in attributes_to_update.items():
                setattr(new_metric, param_name, param_value)
            self._metrics[metric.name] = new_metric
        return

    def fetch_trials_data_results(
        self,
        trial_indices: Iterable[int] | None = None,
        metrics: list[Metric] | None = None,
        **kwargs: Any,
    ) -> dict[int, dict[str, MetricFetchResult]]:
        """Fetches data for specific trials on the experiment.

        If a metric fetch fails, the Exception will be captured in the
        MetricFetchResult along with a message.

        NOTE: For metrics that are not available while trial is running, the data
        may be retrieved from cache on the experiment. Data is cached on the experiment
        via calls to `experiment.attach_data` and whether a given metric class is
        available while trial is running is determined by the boolean returned from its
        `is_available_while_running` class method.

        Args:
            trial_indices: Indices of trials, for which to fetch data.
            metrics: If provided, fetch data for these metrics instead of the ones
                defined on the experiment.
            kwargs: keyword args to pass to underlying metrics' fetch data functions.

        Returns:
            A nested Dictionary from trial_index => metric_name => result
        """
        trials = (
            list(self.trials.values())
            if trial_indices is None
            else self.get_trials_by_indices(trial_indices=trial_indices)
        )
        return self._lookup_or_fetch_trials_results(
            trials=trials, metrics=metrics, **kwargs
        )

    def fetch_data(
        self,
        trial_indices: Iterable[int] | None = None,
        metrics: list[Metric] | None = None,
        **kwargs: Any,
    ) -> Data:
        """Fetches data for all trials on this experiment and for either the
        specified metrics or all metrics currently on the experiment, if `metrics`
        argument is not specified.

        NOTE: For metrics that are not available while trial is running, the data
        may be retrieved from cache on the experiment. Data is cached on the experiment
        via calls to `experiment.attach_data` and whether a given metric class is
        available while trial is running is determined by the boolean returned from its
        `is_available_while_running` class method.

        Args:
            metrics: If provided, fetch data for these metrics; otherwise, fetch
                data for all metrics defined on the experiment.
            kwargs: keyword args to pass to underlying metrics' fetch data functions.

        Returns:
            Data for the experiment.
        """
        results = self._lookup_or_fetch_trials_results(
            trials=list(self.trials.values())
            if trial_indices is None
            else self.get_trials_by_indices(trial_indices=trial_indices),
            metrics=metrics,
            **kwargs,
        )

        return Metric._unwrap_experiment_data_multi(
            results=results,
        )

    def _lookup_or_fetch_trials_results(
        self,
        trials: list[BaseTrial],
        metrics: Iterable[Metric] | None = None,
        **kwargs: Any,
    ) -> dict[int, dict[str, MetricFetchResult]]:
        if not self.metrics and not metrics:
            raise ValueError(
                "No metrics to fetch data for, as no metrics are defined for "
                "this experiment, and none were passed in to `fetch_data`."
            )
        if not any(t.status.expecting_data for t in trials):
            logger.debug("No trials are in a state expecting data.")
            return {}
        metrics_to_fetch = list(metrics or self.metrics.values())

        # Separate base metrics from derived metrics.
        # Derived metrics must be fetched after base metrics because they
        # depend on base metric data being available in the cache.
        base_metrics: list[Metric] = [
            m for m in metrics_to_fetch if not isinstance(m, DerivedMetric)
        ]
        derived_metrics: list[Metric] = [
            m for m in metrics_to_fetch if isinstance(m, DerivedMetric)
        ]

        results: dict[int, dict[str, MetricFetchResult]] = {}

        # Phase 1: Fetch all base (non-derived) metrics first.
        if base_metrics:
            base_results, base_new = self._fetch_metrics_by_class(
                trials=trials,
                metrics=base_metrics,
                **kwargs,
            )
            results = base_results

            # Attach base metric results to the cache so derived metrics
            # can access base data via lookup_data(), and so base data is
            # persisted even when there are no derived metrics.
            if base_new:
                self._try_attach_fetch_results(base_results)

        # Phase 2: Fetch derived metrics (they look up base data from cache).
        # Base data is already attached above, so we only need to attach
        # derived results here.
        if derived_metrics:
            derived_results, derived_new = self._fetch_metrics_by_class(
                trials=trials,
                metrics=derived_metrics,
                **kwargs,
            )
            for trial_index, trial_metrics in derived_results.items():
                results.setdefault(trial_index, {}).update(trial_metrics)
            if derived_new:
                self._try_attach_fetch_results(derived_results)

        return results

    def _try_attach_fetch_results(
        self,
        results: dict[int, dict[str, MetricFetchResult]],
    ) -> None:
        """Attach fetch results to the experiment cache, logging on error."""
        try:
            self.attach_fetch_results(results=results)
        except ValueError as e:
            logger.error(
                f"Encountered ValueError {e} while attaching results. "
                "Proceeding and returning results fetched without attaching."
            )

    def _fetch_metrics_by_class(
        self,
        trials: list[BaseTrial],
        metrics: list[Metric],
        **kwargs: Any,
    ) -> tuple[dict[int, dict[str, MetricFetchResult]], bool]:
        """Fetch metrics grouped by class.

        Args:
            trials: List of trials to fetch data for.
            metrics: List of metrics to fetch.
            **kwargs: Additional keyword arguments passed to fetch methods.

        Returns:
            A tuple of (results dict, contains_new_data bool).
        """
        metrics_by_class = self._metrics_by_class(metrics=metrics)

        results: dict[int, dict[str, MetricFetchResult]] = {}
        contains_new_data = False

        for _metric_cls, cls_metrics in metrics_by_class.items():
            first_metric_of_group = cls_metrics[0]
            (
                new_fetch_results,
                new_results_contains_new_data,
            ) = first_metric_of_group.fetch_data_prefer_lookup(
                experiment=self,
                metrics=cls_metrics,
                trials=trials,
                **kwargs,
            )

            contains_new_data = contains_new_data or new_results_contains_new_data

            # Merge in results
            results = {
                trial.index: {
                    **(
                        new_fetch_results[trial.index]
                        if trial.index in new_fetch_results
                        else {}
                    ),
                    **(results[trial.index] if trial.index in results else {}),
                }
                for trial in trials
            }

        return results, contains_new_data

    @copy_doc(BaseTrial.fetch_data)
    def _fetch_trial_data(
        self, trial_index: int, metrics: list[Metric] | None = None, **kwargs: Any
    ) -> dict[str, MetricFetchResult]:
        trial = self.trials[trial_index]

        trial_data = self._lookup_or_fetch_trials_results(
            trials=[trial], metrics=metrics, **kwargs
        )

        if trial_index in trial_data:
            return trial_data[trial_index]

        return {}

    def attach_data(self, data: Data, **kwargs: Any) -> None:
        """
        Update the experiment's `data` attribute.

        When a new observation is attached and ``self.data`` already has an
        observation for that arm name, metric, and (if present) step, the new
        observation replaces the old one.

        Args:
            data: Data to attach.
            kwargs: Deprecated arguments.
        """
        deprecated_arguments = ["combine_with_last_data", "overwrite_existing_data"]
        for arg in deprecated_arguments:
            if arg in kwargs:
                warnings.warn(
                    f"Passing {arg} to `attach_data` is deprecated. "
                    "`attach_data` will behave in a way similar to the old "
                    "`combine_with_last_data=True`, "
                    "`overwrite_existing_data=False`. behavior.",
                    DeprecationWarning,
                    stacklevel=2,
                )
        unexpected_args = set(kwargs.keys()) - set(deprecated_arguments)
        if unexpected_args:
            raise ValueError(
                f"Unexpected arguments {unexpected_args} passed to `attach_data`."
            )
        if data.empty:
            raise ValueError("Data to attach is empty.")
        metrics_not_on_exp = data.metric_names - set(self.metrics.keys())
        if metrics_not_on_exp:
            logger.debug(
                f"Attached data has some metrics ({metrics_not_on_exp}) that are "
                "not among the metrics on this experiment. Note that attaching data "
                "will not automatically add those metrics to the experiment. "
                "For these metrics to be automatically fetched by `experiment."
                "fetch_data`, add them via `experiment.add_tracking_metric` or update "
                "the experiment's optimization config."
            )
        unattached_trial_indices = data.trial_indices - set(self.trials.keys())
        if unattached_trial_indices:
            raise ValueError(
                f"Cannot attach data for trials {unattached_trial_indices} "
                "because they have not been attached to the experiment."
            )

        self.data = Data(
            data_rows=combine_data_rows_favoring_recent(
                last_rows=self.data._data_rows, new_rows=data._data_rows
            )
        )

    def attach_fetch_results(
        self,
        results: Mapping[int, Mapping[str, MetricFetchResult]],
    ) -> None:
        """
        UNSAFE: Prefer to use attach_data directly instead.

        Attach fetched data results to the Experiment so they will not have to be
        fetched again. Additionally caches any metric fetching errors that occurred
        to the experiment.

        NOTE: Any Errs in the results passed in will silently be dropped! This will
        cause the Experiment to fail to find them in ``self.data`` and
        attempt to refetch at fetch time. If this is not your intended behavior you
        MUST resolve your results first and use attach_data directly instead.
        """

        completed_trial_indices = self.trial_indices_by_status[TrialStatus.COMPLETED]
        oks: list[Ok[Data, MetricFetchE]] = []
        for trial_index, metrics in results.items():
            for metric_name, result in metrics.items():
                if isinstance(result, Ok):
                    oks.append(result)
                    self._metric_fetching_errors.pop((trial_index, metric_name), None)
                elif isinstance(result, Err):
                    msg = (
                        "Discovered Metric fetching Err while attaching data "
                        f"{result.err}. "
                        "Ignoring for now -- will retry query on next call to fetch."
                    )
                    self._cache_metric_fetch_error(
                        trial_index=trial_index,
                        metric_name=metric_name,
                        metric_fetch_e=result.err,
                    )
                    if trial_index in completed_trial_indices:
                        logger.error(msg)
                    else:
                        msg += (
                            f" Suppressing error for trial {trial_index} not in "
                            "COMPLETED state."
                        )
                        logger.debug(msg)

        if len(oks) < 1:
            return None

        data = Data.from_multiple_data(data=[ok.ok for ok in oks])
        self.attach_data(data=data)

    def lookup_data(
        self,
        trial_indices: Iterable[int] | None = None,
    ) -> Data:
        """
        Return ``Data`` for trials ``trial_indices``.

        For each trial, returns data present for this trial.  Returns
        metrics - to do that, use `fetch_data()` instead.

        Args:
            trial_indices: Indices of trials for which to fetch data. If omitted,
                lookup data for all trials on the experiment.

        Returns:
            Data for trials ``trial_indices`` on the experiment.
        """
        if trial_indices is None:
            return self.data
        trial_indices = list(trial_indices)
        if len(trial_indices) == 0:
            return Data()
        if set(trial_indices) == self.data.trial_indices:
            return self.data

        return self.data.filter(trial_indices=trial_indices)

    @property
    def num_trials(self) -> int:
        """How many trials are associated with this experiment."""
        return len(self._trials)

    @property
    def trials(self) -> dict[int, BaseTrial]:
        """The trials associated with the experiment.

        NOTE: If some trials on this experiment specify their TTL, `CANDIDATE` trials
        will be checked for whether their TTL elapsed during this call. Found past-
        TTL trials will be marked as `STALE`.
        """
        self._check_TTL_on_candidate_trials()
        return self._trials

    @property
    def trials_by_status(self) -> dict[TrialStatus, list[BaseTrial]]:
        """Trials associated with the experiment, grouped by trial status."""
        # Make sure all statuses appear in this dict, to avoid key errors.
        return {
            status: self.get_trials_by_indices(trial_indices=idcs)
            for status, idcs in self.trial_indices_by_status.items()
        }

    @property
    def trials_expecting_data(self) -> list[BaseTrial]:
        """list[BaseTrial]: the list of all trials for which data has arrived
        or is expected to arrive.
        """
        return [trial for trial in self.trials.values() if trial.status.expecting_data]

    @property
    def completed_trials(self) -> list[BaseTrial]:
        """list[BaseTrial]: the list of all trials for which data has arrived
        or is expected to arrive.
        """
        return self.trials_by_status[TrialStatus.COMPLETED]

    @property
    def trial_indices_by_status(self) -> dict[TrialStatus, set[int]]:
        """Indices of trials associated with the experiment, grouped by trial
        status.
        """
        self._check_TTL_on_candidate_trials()  # Marks past-TTL trials as stale.
        return self._trial_indices_by_status

    @property
    def running_trial_indices(self) -> set[int]:
        """Indices of running trials, associated with the experiment."""
        return self._trial_indices_by_status[TrialStatus.RUNNING]

    @property
    def trial_indices_expecting_data(self) -> set[int]:
        """Set of indices of trials, statuses of which indicate that we expect
        these trials to have data, either already or in the future.
        """
        return set.union(
            *(
                self.trial_indices_by_status[status]
                for status in STATUSES_EXPECTING_DATA
            )
        )

    def trial_indices_with_data(
        self, critical_metrics_only: bool | None = True
    ) -> set[int]:
        """Set of indices of trials for which we have data for either all metrics on
        the experiment, or all metrics in the optimization config. Helpful for
        determining which trials currently have data for modeling.

        Args:
            critical_metrics_only: If True, only return trials for which we have
            metrics for the optimization config. If False, return trials for which
            we have data for all metrics on the experiment, including tracking metrics
        """
        if critical_metrics_only:
            if self.optimization_config is None:
                raise OptimizationNotConfiguredError(
                    "Cannot find trials with data for optimization config metrics "
                    "because no optimization config has been defined."
                )
            metric_names = self.optimization_config.metric_names
        else:
            metric_names = set(self.metrics.keys())
            if len(metric_names) == 0:
                return set()

        exp_data = self.lookup_data().filter(metric_names=metric_names)
        trials_with_data = set()
        for trial_idx in self.trials.keys():
            metrics_in_trial_data = set(
                exp_data.df[exp_data.df["trial_index"] == trial_idx][
                    "metric_name"
                ].unique()
            )
            if metrics_in_trial_data == metric_names:
                trials_with_data.add(trial_idx)
            else:
                logger.debug(
                    f"Trial {trial_idx} does not have data for required metrics "
                    f"({metric_names}) on the experiment. "
                    f"Metrics present in trial data: {metrics_in_trial_data}"
                )

        return trials_with_data

    def new_trial(
        self,
        generator_run: GeneratorRun | None = None,
        trial_type: str | None = None,
        ttl_seconds: int | None = None,
    ) -> Trial:
        """Create a new trial associated with this experiment.

        Args:
            generator_run: GeneratorRun, associated with this trial.
                Trial has only one arm attached to it and this generator_run
                must therefore contain one arm. This arm can also be set later
                through `add_arm` or `add_generator_run`, but a trial's
                associated generator run is immutable once set.
            trial_type: Type of this trial, if used in MultiTypeExperiment.
            ttl_seconds: If specified, trials will be considered stale after
                this many seconds since the time the trial was ran, unless the
                trial is completed before then. Meant to be used to detect
                'dead' trials, for which the evaluation process might have
                crashed etc., and which should be considered stale after
                their 'time to live' has passed.
        """
        if ttl_seconds is not None:
            self._trials_have_ttl = True
        return Trial(
            experiment=self,
            trial_type=trial_type,
            generator_run=generator_run,
            ttl_seconds=ttl_seconds,
        )

    def new_batch_trial(
        self,
        generator_run: GeneratorRun | None = None,
        generator_runs: list[GeneratorRun] | None = None,
        should_add_status_quo_arm: bool | None = False,
        trial_type: str | None = None,
        ttl_seconds: int | None = None,
    ) -> BatchTrial:
        """Create a new batch trial associated with this experiment.

        Args:
            generator_run: GeneratorRun, associated with this trial. This can a
                also be set later through `add_arm` or `add_generator_run`, but a
                trial's associated generator run is immutable once set.
            generator_runs: GeneratorRuns, associated with this trial. This can a
                also be set later through `add_arm` or `add_generator_run`, but a
                trial's associated generator run is immutable once set.  This cannot
                be combined with the `generator_run` argument.
            should_add_status_quo_arm: If True, adds the status quo arm to the trial
                with a weight of 1.0. If False, the _status_quo is still set on the
                trial for tracking purposes, but without a weight it will not be an
                Arm present on the trial
            trial_type: Type of this trial, if used in MultiTypeExperiment.
            ttl_seconds: If specified, trials will be considered stale after
                this many seconds since the time the trial was ran, unless the
                trial is completed before then. Meant to be used to detect
                'dead' trials, for which the evaluation process might have
                crashed etc., and which should be considered stale after
                their 'time to live' has passed.
        """
        if ttl_seconds is not None:
            self._trials_have_ttl = True
        return BatchTrial(
            experiment=self,
            trial_type=trial_type,
            generator_run=generator_run,
            generator_runs=generator_runs,
            should_add_status_quo_arm=should_add_status_quo_arm,
            ttl_seconds=ttl_seconds,
        )

    def get_batch_trial(self, trial_index: int) -> BatchTrial:
        """
        Return a trial on experiment cast as BatchTrial
        Args:
            trial_index: The index of the trial to lookup data for.
        Returns:
            The requested trial cast as BatchTrial
        """
        return assert_is_instance(
            self.get_trials_by_indices(trial_indices=[trial_index])[0], BatchTrial
        )

    def get_trials_by_indices(self, trial_indices: Iterable[int]) -> list[BaseTrial]:
        """Grabs trials on this experiment by their indices."""
        trial_indices = list(trial_indices)
        try:
            trials = self.trials
            return [trials[idx] for idx in trial_indices]
        except KeyError:
            missing = set(trial_indices) - set(self.trials)
            raise ValueError(
                f"Trial indices {missing} are not associated with the experiment.\n"
                f"Trials indices available on this experiment: {list(self.trials)}."
            )

    def extract_relevant_trials(
        self,
        trial_indices: Sequence[int] | None = None,
        trial_statuses: Sequence[TrialStatus] | None = None,
    ) -> list[BaseTrial]:
        """
        Find the trials on this experiment which meet both the trial_indices
        filtering condition and the trial_statuses filtering condition. If None is
        provided for either condition, the condition is not applied.

        Args:
            trial_indices: Indices of trials to include. If None, include all trials.
                If empty list, include no trials.
            trial_statuses: Statuses to filter by. If None, no filtering by status.
        """
        trials = (
            [self.trials[i] for i in trial_indices]
            if trial_indices is not None
            else [*self.trials.values()]
        )

        filtered_by_status = [
            trial
            for trial in trials
            if (trial_statuses is None or trial.status in trial_statuses)
        ]

        return filtered_by_status

    @retry_on_exception(retries=3, no_retry_on_exception_types=NO_RETRY_EXCEPTIONS)
    def stop_trial_runs(
        self, trials: list[BaseTrial], reasons: list[str | None] | None = None
    ) -> None:
        """Stops the jobs that execute given trials.

        Used if, for example, TTL for a trial was specified and expired, or poor
        early results suggest the trial is not worth running to completion.

        Override default implementation on the ``Runner`` if its desirable to stop
        trials in bulk.

        Args:
            trials: Trials to be stopped.
            reasons: A list of strings describing the reasons for why the
                trials are to be stopped (in the same order).
        """
        if len(trials) == 0:
            return

        if reasons is None:
            reasons = [None] * len(trials)

        for trial, reason in zip(trials, reasons):
            runner = self.runner_for_trial_type(trial_type=trial.trial_type)
            if runner is None:
                raise RunnerNotFoundError(
                    "Unable to stop trial runs: Runner not configured "
                    "for experiment or trial."
                )
            runner.stop(trial=trial, reason=reason)
            trial.mark_early_stopped(reason=reason)

    def _attach_trial(self, trial: BaseTrial, index: int | None = None) -> int:
        """Attach a trial to this experiment.

        Should only be called within the trial constructor.

        Args:
            trial: The trial to be attached.
            index: If specified, the trial's index will be set accordingly.
                This should generally not be specified, as the index
                will be automatically determined based on the number
                of existing trials. This is only used for the purpose
                of loading from storage.

        Returns:
            The index of the trial within the experiment's trial list.
        """

        if trial.experiment is not self:
            raise ValueError("BatchTrial does not belong to this experiment.")

        for existing_trial in self._trials.values():
            if existing_trial is trial:
                raise ValueError("BatchTrial already attached to experiment.")

        if index is not None and index in self._trials:
            logger.debug(
                f"Trial index {index} already exists on the experiment. Overwriting."
            )
        index = (
            index
            if index is not None
            else (0 if len(self._trials) == 0 else max(self._trials.keys()) + 1)
        )
        self._trials[index] = trial
        return index

    def validate_trials(self, trials: Iterable[BaseTrial]) -> None:
        """Raise ValueError if any of the trials in the input are not from
        this experiment.
        """
        for t in trials:
            if t.experiment is not self:
                raise ValueError(
                    "All trials must be from the input experiment (name: "
                    f"{self.name}), but trial #{t.index} is from "
                    f"another experiment (name: {t.experiment.name})."
                )

    def warm_start_from_old_experiment(
        self,
        old_experiment: Experiment,
        copy_run_metadata_keys: list[str] | None = None,
        trial_statuses_to_copy: list[TrialStatus] | None = None,
        search_space_check_membership_raise_error: bool = True,
    ) -> list[Trial]:
        """Copy all completed trials with data from an old Ax experiment to this one.
        This function checks that the parameters of each trial are members of the
        current experiment's search_space.

        NOTE: Currently only handles experiments with 1-arm ``Trial``-s, not
        ``BatchTrial``-s as there has not yet been need for support of the latter.

        Args:
            old_experiment: The experiment from which to transfer trials and data
            copy_run_metadata_keys: A list of keys denoting which items to copy over
                from each trial's run_metadata. Defaults to copying all run_metadata.
            trial_statuses_to_copy: All trials with a status in this list will be
                copied. By default, copies all ``RUNNING``, ``COMPLETED``,
                ``ABANDONED``, and ``EARLY_STOPPED`` trials.
            search_space_check_membership_raise_error: Whether to raise an exception
                if the warm started trials being imported fall outside of the
                defined search space.

        Returns:
            List of trials successfully copied from old_experiment to this one
        """
        if len(self.trials) > 0:
            raise ValueError(
                f"Can only warm-start experiments that don't yet have trials. "
                f"Experiment {self._name} has {len(self.trials)} trials."
            )

        old_parameter_names = set(old_experiment.search_space.parameters.keys())
        parameter_names = set(self.search_space.parameters.keys())
        if old_parameter_names.symmetric_difference(parameter_names):
            raise ValueError(
                f"Cannot warm-start experiment '{self._name}' from experiment "
                f"'{old_experiment._name}' due to mismatch in search space parameters."
                f"Parameters in '{self._name}' but not in '{old_experiment._name}': "
                f"{old_parameter_names - parameter_names}. Vice-versa: "
                f"{parameter_names - old_parameter_names}."
            )

        trial_statuses_to_copy = (
            trial_statuses_to_copy
            if trial_statuses_to_copy is not None
            else DEFAULT_STATUSES_TO_WARM_START
        )

        warm_start_trials = [
            trial
            for trial in old_experiment.trials.values()
            if trial.status in trial_statuses_to_copy
        ]
        copied_trials = []
        for trial in warm_start_trials:
            if not isinstance(trial, Trial):
                raise NotImplementedError(
                    "Only experiments with 1-arm trials currently supported."
                )
            self.search_space.check_membership(
                none_throws(trial.arm).parameters,
                raise_error=search_space_check_membership_raise_error,
            )
            dat = old_experiment.lookup_data(trial_indices={trial.index})
            # Set trial index and arm name to their values in new trial.
            new_trial = self.new_trial()
            add_arm_and_prevent_naming_collision(
                new_trial=new_trial,
                old_trial=trial,
                old_experiment_name=old_experiment._name,
            )
            new_trial.mark_running(no_runner_required=True)
            new_trial._properties["source"] = (
                f"Warm start from Experiment: `{old_experiment._name}`, "
                f"trial: `{trial.index}`"
            )
            # Associates a generation_model_key to the new trial.
            generation_model_key = trial._properties.get("generation_model_key")
            if generation_model_key is None and trial.generator_run is not None:
                generation_model_key = trial.generator_run._generator_key or "Manual"
            new_trial._properties["generation_model_key"] = generation_model_key

            # Copy all run_metadata by default.
            if copy_run_metadata_keys is None:
                copy_run_metadata_keys = list(trial.run_metadata.keys())

            for run_metadata_field in copy_run_metadata_keys:
                new_trial.update_run_metadata(
                    {run_metadata_field: trial.run_metadata.get(run_metadata_field)}
                )
            # Trial has data, so we replicate it on the new experiment.
            has_data = not dat.df.empty
            if has_data:
                new_df = dat.full_df.copy()
                new_df["trial_index"] = new_df["trial_index"].replace(
                    {trial.index: new_trial.index}
                )
                new_df["arm_name"] = new_df["arm_name"].replace(
                    {none_throws(trial.arm).name: none_throws(new_trial.arm).name}
                )
                # Attach updated data to new trial on experiment.
                self.attach_data(data=Data(df=new_df))
            if trial.status is not TrialStatus.RUNNING:
                new_trial.mark_as(trial.status, reason=trial.status_reason)
            copied_trials.append(new_trial)

        if self._name is not None:
            logger.debug(
                f"Copied {len(copied_trials)} completed trials and their data "
                f"from {old_experiment._name} to {self._name}."
            )
        else:
            logger.debug(
                f"Copied {len(copied_trials)} completed trials and their data "
                f"from {old_experiment._name}."
            )

        return copied_trials

    def _name_and_store_arm_if_not_exists(
        self, arm: Arm, proposed_name: str, replace: bool = False
    ) -> None:
        """Tries to lookup arm with same signature, otherwise names and stores it.

        - Looks up if arm already exists on experiment
            - If so, name the input arm the same as the existing arm
            - else name the arm with given name and store in _arms_by_signature

        Args:
            arm: The arm object to name.
            proposed_name: The name to assign if it doesn't have one already.
            replace: If true, override arm w/ same name and different signature.
                If false, raise an error if this conflict occurs.
        """

        # If arm is identical to an existing arm, return that
        # so that the names match.
        if arm.signature in self.arms_by_signature:
            existing_arm = self.arms_by_signature[arm.signature]
            if arm.has_name:
                if arm.name != existing_arm.name:
                    raise ValueError(
                        f"Arm already exists with name {existing_arm.name}, "
                        f"which doesn't match given arm name of {arm.name}."
                    )
            else:
                arm.name = existing_arm.name
        else:
            if not arm.has_name:
                arm.name = proposed_name

            # Check for signature conflict by arm name/proposed name
            if (
                arm.name in self.arms_by_name
                and arm.signature != self.arms_by_name[arm.name].signature
            ):
                error_msg = (
                    f"Arm with name {arm.name} already exists on experiment "
                    + "with different signature."
                )
                if replace:
                    logger.warning(f"{error_msg} Replacing the existing arm. ")
                else:
                    raise AxError(error_msg)

            # Add the new arm
            self._register_arm(arm)

    def _register_arm(self, arm: Arm) -> None:
        """Add a new arm to the experiment, updating the relevant
        lookup dictionaries.

        Args:
            arm: Arm to add
        """
        self._arms_by_signature[arm.signature] = arm
        self._arms_by_name[arm.name] = arm

    def _check_TTL_on_candidate_trials(self) -> None:
        """Checks whether any past-TTL trials are still marked as `CANDIDATE`
        and marks them as stale if so.

        NOTE: this function just calls `trial.status` for each trial, as the
        computation of that property checks the TTL for trials.
        """
        if not self._trials_have_ttl:
            return

        # The trial status changes during TTL check modifies the original set
        # _trial_indices_by_status, hence create a copy of indices for iteration
        candidate_indices = self._trial_indices_by_status[TrialStatus.CANDIDATE].copy()

        for idx in candidate_indices:
            self._trials[idx].status  # `status` property checks TTL if applicable.

    def _cache_metric_fetch_error(
        self, trial_index: int, metric_name: str, metric_fetch_e: MetricFetchE | None
    ) -> None:
        """Caches a given metric fetch error to the experiment.
        Args:
            trial_index: Index of the trial that encountered the metric fetch error.
            metric_name: Name of the metric that failed to be fetched.
            metric_fetch_e: The metric fetch error to cache.
        Returns:
            None
        """
        error_data = {
            "trial_index": trial_index,
            "metric_name": metric_name,
            "reason": "",
            "timestamp": datetime.now().isoformat(),
            "traceback": "",
        }

        if metric_fetch_e is not None:
            reason_for_failure = metric_fetch_e.message
            if metric_fetch_e.exception is not None:
                exception_str = (
                    f"{type(metric_fetch_e.exception).__name__}: "
                    f"{metric_fetch_e.exception}"
                )
                reason_for_failure = (
                    f"Ran into the following exception: {exception_str}"
                )

            error_data["reason"] = reason_for_failure
            error_data["traceback"] = metric_fetch_e.tb_str() or ""

        self._metric_fetching_errors[(trial_index, metric_name)] = error_data

    def __repr__(self) -> str:
        return self.__class__.__name__ + f"({self._name})"

    # --- MultiTypeExperiment convenience functions ---
    #
    # Certain functionalities have special behavior for multi-type experiments.
    # This defines the base behavior for regular experiments that will be
    # overridden in the MultiTypeExperiment class.

    @property
    def default_trial_type(self) -> str | None:
        """Default trial type assigned to trials in this experiment.

        In the base experiment class this is always None. For experiments
        with multiple trial types, use the MultiTypeExperiment class.
        """
        return self._default_trial_type

    @property
    def trial_type_to_metric_names(self) -> dict[str, set[str]]:
        """Map from trial type to the set of metric names relevant to that
        type.

        Returns a shallow copy of the internal mapping.
        """
        return dict(self._trial_type_to_metric_names)

    @property
    def metric_to_trial_type(self) -> dict[str, str]:
        """Map each metric name to its associated trial type.

        Computed from ``_trial_type_to_metric_names``. When a
        ``default_trial_type`` is set and an ``optimization_config`` exists,
        optimization config metrics are pinned to the default trial type.
        """
        result: dict[str, str] = {}
        for trial_type, metric_names in self._trial_type_to_metric_names.items():
            for name in metric_names:
                result[name] = trial_type
        opt_config = self._optimization_config
        default_trial_type = self._default_trial_type
        if default_trial_type is not None and opt_config is not None:
            for metric_name in opt_config.metric_names:
                result[metric_name] = default_trial_type
        return result

    def metrics_for_trial_type(self, trial_type: str) -> list[Metric]:
        """Return the metrics associated with a given trial type.

        Args:
            trial_type: The trial type to look up metrics for.

        Raises:
            ValueError: If the trial type is not supported.
        """
        if not self.supports_trial_type(trial_type):
            raise ValueError(f"Trial type `{trial_type}` is not supported.")
        valid_names = self._trial_type_to_metric_names.get(trial_type, set())
        return [self._metrics[name] for name in valid_names if name in self._metrics]

    @property
    def default_trials(self) -> set[int]:
        """Return the indices for trials of the default type."""
        return {
            idx
            for idx, trial in self.trials.items()
            if trial.trial_type == self.default_trial_type
        }

    def add_trial_type(self, trial_type: str, runner: Runner | None = None) -> Self:
        """Add a new trial type to be supported by this experiment.

        Args:
            trial_type: The new trial type to be added.
            runner: The default runner for trials of this type.
        """
        if self.supports_trial_type(trial_type):
            raise ValueError(f"Experiment already contains trial_type `{trial_type}`")

        if runner is not None:
            self._trial_type_to_runner[trial_type] = runner

        return self

    def update_runner(self, trial_type: str, runner: Runner) -> Self:
        """Update the default runner for an existing trial type.

        Args:
            trial_type: The trial type whose runner should be updated.
            runner: The new runner for trials of this type.
        """
        if not self.supports_trial_type(trial_type):
            raise ValueError(f"Experiment does not contain trial_type `{trial_type}`")
        self._trial_type_to_runner[trial_type] = runner
        self._runner = runner
        return self

    def runner_for_trial_type(self, trial_type: str | None) -> Runner | None:
        """The default runner to use for a given trial type.

        Looks up the appropriate runner for this trial type in the trial_type_to_runner.
        """
        if not self.supports_trial_type(trial_type):
            raise ValueError(f"Trial type `{trial_type}` is not supported.")
        if (runner := self._trial_type_to_runner.get(trial_type)) is None:
            return self.runner  # return the default runner
        return runner

    def supports_trial_type(self, trial_type: str | None) -> bool:
        """Whether this experiment allows trials of the given type.

        For experiments with a ``default_trial_type`` (multi-type experiments),
        only trial types registered in ``_trial_type_to_runner`` are supported.
        For single-type experiments, ``None`` is always supported, along with
        ``SHORT_RUN`` and ``LONG_RUN`` for backward compatibility with
        generation strategies that use those trial types.
        """
        if self._default_trial_type is not None:
            return trial_type in self._trial_type_to_runner
        return (
            trial_type is None
            or trial_type == Keys.SHORT_RUN
            or trial_type == Keys.LONG_RUN
            or trial_type == Keys.LILO_LABELING
            or trial_type in self._trial_type_to_runner
        )

    def attach_trial(
        self,
        parameterizations: list[TParameterization],
        arm_names: list[str] | None = None,
        should_add_status_quo_arm: bool = False,
        ttl_seconds: int | None = None,
        run_metadata: dict[str, Any] | None = None,
    ) -> tuple[dict[str, TParameterization], int]:
        """Attach a new trial with the given parameterization to the experiment.

        Args:
            parameterizations: List of parameterization for the new trial. If
                only one is provided a single-arm Trial is created. If multiple
                arms are provided a BatchTrial is created.
            arm_names: Names of arm(s) in the new trial.
            should_add_status_quo_arm: If True, adds the status quo arm to the trial
                with a weight of 1.0. If False, the _status_quo is still set on the
                trial for tracking purposes, but without a weight it will not be an
                Arm present on the trial
            ttl_seconds: If specified, will consider the trial stale after this
                many seconds. Used to detect dead trials that did not complete.
            run_metadata: Metadata to attach to the trial.

        Returns:
            Tuple of arm name to parameterization dict, and trial index from
            newly created trial.
        """

        # If more than one parameterization is provided,
        # proceed with a batch trial
        is_batch = len(parameterizations) > 1

        # Validate search space membership for all parameterizations
        for parameterization in parameterizations:
            try:
                self.search_space.validate_membership(parameters=parameterization)
            except ValueError as e:
                # To not raise on out-of-design parameterizations
                if "is not a valid value for parameter" in str(e):
                    warnings.warn(
                        f"Parameterization {parameterization} is in out-of-design. "
                        "Ax will still attach the trial for use in candidate "
                        "generation.",
                        RuntimeWarning,
                        stacklevel=2,
                    )
                else:
                    raise e

        # Validate number of arm names if any arm names are provided.
        named_arms = False
        arm_names = arm_names or []
        if len(arm_names) > 0:
            named_arms = True
            if len(arm_names) != len(parameterizations):
                raise UserInputError(
                    f"Number of arm names ({len(arm_names)} "
                    "does not match number of parameterizations "
                    f"({len(parameterizations)})."
                )

        # Prepare arm(s) to be added to the trial created later
        arms = [
            Arm(
                parameters=parameterization,
                name=arm_names[i] if named_arms else None,
            )
            for i, parameterization in enumerate(parameterizations)
        ]

        # Create the trial and add arm(s)
        trial = None
        if is_batch:
            trial = self.new_batch_trial(
                ttl_seconds=ttl_seconds,
                should_add_status_quo_arm=should_add_status_quo_arm,
            ).add_arms_and_weights(arms=arms)

        else:
            # If search space is hierarchical, we need to store dummy values of
            # parameters that are not in the arm (but are in flattened
            # search space), as metadata, so later we are able to make the
            # data for this arm "complete" in the flattened search space.
            candidate_metadata = None
            if self.search_space.is_hierarchical:
                candidate_metadata = self.search_space.cast_observation_features(
                    observation_features=self.search_space.flatten_observation_features(
                        observation_features=observation.ObservationFeatures(
                            parameters=parameterizations[0]
                        ),
                        inject_dummy_values_to_complete_flat_parameterization=True,
                    )
                ).metadata

            trial = self.new_trial(ttl_seconds=ttl_seconds).add_arm(
                arms[0],
                candidate_metadata=candidate_metadata,
            )

        trial.mark_running(no_runner_required=True)

        logger.debug(
            "Attached custom parameterizations "
            f"{round_floats_for_logging(item=parameterizations)} "
            f"as trial {trial.index}."
        )

        if run_metadata is not None:
            trial.update_run_metadata(metadata=run_metadata)

        return {arm.name: arm.parameters for arm in trial.arms}, trial.index

    def clone_with(
        self,
        search_space: SearchSpace | None = None,
        name: str | None = None,
        optimization_config: OptimizationConfig | None = None,
        tracking_metrics: list[Metric] | None = None,
        runner: Runner | None = None,
        status_quo: Arm | None = None,
        description: str | None = None,
        is_test: bool | None = None,
        properties_to_keep: list[str] | None = None,
        trial_indices: list[int] | None = None,
        clear_trial_type: bool = False,
    ) -> Experiment:
        r"""
        Return a copy of this experiment with some attributes replaced.

        NOTE: This method only retains the latest data attached to the experiment.
        This is the same data that would be accessed using common APIs such as
        ``Experiment.lookup_data()``.

        Args:
            search_space: New search space. If None, it uses the cloned search space
                of the original experiment.
            name: New experiment name. If None, it adds cloned_experiment_  prefix
                to the original experiment name.
            optimization_config: New optimization config. If None, it clones the same
                optimization_config from the orignal experiment.
            tracking_metrics: New list of metrics to track. If None, it clones the
                tracking metrics already attached to the main experiment.
            runner: New runner. If None, it clones the existing runner.
            status_quo: New status quo arm. If None, it clones the existing status quo.
            description: New description. If None, it uses the same description.
            is_test: Whether the cloned experiment should be considered a test. If None,
                it uses the same value.
            properties_to_keep: List of property keys to retain in the cloned
                experiment. Defaults to ["owners"].
            trial_indices: If specified, only clones the specified trials. If None,
                clones all trials.
            clear_trial_type: If True, all cloned trials on the cloned experiment have
                `trial_type` set to `None`.
        """
        if properties_to_keep is None:
            properties_to_keep = ["owners"]
        search_space = (
            self.search_space.clone() if (search_space is None) else search_space
        )
        name = (
            "cloned_experiment_" + self.name
            if (name is None and self.name is not None)
            else name
        )
        optimization_config = (
            self.optimization_config.clone()
            if (optimization_config is None and self.optimization_config is not None)
            else optimization_config
        )
        tracking_metrics = (
            [m.clone() for m in self.metrics.values()]
            if tracking_metrics is None
            else tracking_metrics
        )

        runner = (
            self.runner.clone()
            if (runner is None and self.runner is not None)
            else runner
        )
        status_quo = (
            self.status_quo.clone()
            if (status_quo is None and self.status_quo is not None)
            else status_quo
        )
        description = self.description if description is None else description
        is_test = self.is_test if is_test is None else is_test

        properties = {
            k: v for k, v in self._properties.items() if k in properties_to_keep
        }
        dropped_keys = set(self._properties.keys()) - set(properties.keys())
        if dropped_keys:
            logger.warning(
                f"When cloning the experiment, the following fields were dropped from "
                f"properties: {', '.join(dropped_keys)}.",
            )

        cloned_experiment = Experiment(
            search_space=search_space,
            name=name,
            optimization_config=optimization_config,
            tracking_metrics=tracking_metrics,
            runner=runner,
            status_quo=status_quo,
            description=description,
            is_test=is_test,
            experiment_type=self.experiment_type,
            properties=properties,
        )

        # Clone only the specified trials.
        original_trial_indices = self.trials.keys()
        trial_indices_to_keep = (
            set(original_trial_indices) if trial_indices is None else set(trial_indices)
        )
        if trial_indices_diff := trial_indices_to_keep.difference(
            original_trial_indices
        ):
            warnings.warn(
                f"Trials indexed with {trial_indices_diff} are not a part "
                "of the original experiment. ",
                stacklevel=2,
            )

        old_index_to_new_index = {}
        intersection_indices = trial_indices_to_keep.intersection(
            original_trial_indices
        )
        for trial_index in intersection_indices:
            trial = self.trials[trial_index]
            if not isinstance(trial, (Trial, BatchTrial)):
                raise NotImplementedError(f"Cloning of {type(trial)} is not supported.")
            new_trial = trial.clone_to(
                cloned_experiment, clear_trial_type=clear_trial_type
            )
            new_index = new_trial.index
            old_index_to_new_index[trial_index] = new_index

        new_df = self.data.full_df.loc[
            lambda x: x["trial_index"].isin(intersection_indices)
        ]
        new_df.loc[:, "trial_index"] = new_df["trial_index"].map(old_index_to_new_index)
        cloned_experiment.data = self.data.__class__(df=new_df)
        return cloned_experiment

    @property
    def metric_config_summary_df(self) -> pd.DataFrame:
        """Creates a dataframe with information about each metric in the
        experiment. The resulting dataframe has one row per metric, and the
        following columns:
            - Name: the name of the metric.
            - Type: the metric subclass (e.g., Metric, BraninMetric).
            - Goal: the goal for this for this metric, based on the optimization
              config (minimize, maximize, constraint or track).
            - Bound: the bound of this metric (e.g., "<=10.0") if it is being used
              as part of an ObjectiveThreshold or OutcomeConstraint.
            - Lower is Better: whether the user prefers this metric to be lower,
              if provided.

        """
        records = {}
        for metric_name in self.metrics.keys():
            m = self.metrics[metric_name]
            records[m.name] = m.summary_dict
        if self.optimization_config is not None:
            opt_config = self.optimization_config
            objective = opt_config.objective
            if objective.is_multi_objective:
                parts = [p.strip() for p in objective.expression.split(",")]
                sub_objectives = [Objective(expression=part) for part in parts]
                for sub_obj in sub_objectives:
                    obj_name = sub_obj.metric_names[0]
                    if obj_name in records:
                        records[obj_name][METRIC_DF_COLNAMES["goal"]] = (
                            "minimize" if sub_obj.minimize else "maximize"
                        )
            elif objective.is_scalarized_objective:
                for metric_name, weight in objective.metric_weights:
                    if metric_name in records:
                        records[metric_name][METRIC_DF_COLNAMES["goal"]] = (
                            "minimize" if weight < 0 else "maximize"
                        )
            else:
                obj_name = objective.metric_names[0]
                if obj_name in records:
                    records[obj_name][METRIC_DF_COLNAMES["goal"]] = (
                        "minimize" if objective.minimize else "maximize"
                    )

            objective_threshold_names: set[str] = set()
            if isinstance(opt_config, MultiObjectiveOptimizationConfig):
                objective_threshold_names = {
                    ot.metric_names[0] for ot in opt_config.objective_thresholds
                }

            for constraint in opt_config.all_constraints:
                c_name = constraint.metric_names[0]
                if c_name not in objective_threshold_names:
                    if c_name in records:
                        records[c_name][METRIC_DF_COLNAMES["goal"]] = "constrain"
                op = ">= " if constraint.op == ComparisonOp.GEQ else "<= "
                relative = "%" if constraint.relative else ""
                if c_name in records:
                    records[c_name][METRIC_DF_COLNAMES["bound"]] = (
                        f"{op}{constraint.bound}{relative}"
                    )

        for metric in self.tracking_metrics or []:
            records[metric.name][METRIC_DF_COLNAMES["goal"]] = "track"

        # Sort metrics rows by purpose and name, then sort columns.
        df = pd.DataFrame(records.values()).fillna(value="None")
        df.rename(columns=METRIC_DF_COLNAMES, inplace=True)
        df[METRIC_DF_COLNAMES["goal"]] = pd.Categorical(
            df[METRIC_DF_COLNAMES["goal"]],
            categories=["minimize", "maximize", "constrain", "track", "None"],
            ordered=True,
        )
        df = df.sort_values(
            by=[METRIC_DF_COLNAMES["goal"], METRIC_DF_COLNAMES["name"]]
        ).reset_index()
        # Reorder columns.
        df = df[
            [
                colname
                for colname in METRIC_DF_COLNAMES.values()
                if colname in df.columns
            ]
        ]
        return df

    def to_df(
        self,
        trial_indices: Iterable[int] | None = None,
        trial_statuses: Sequence[TrialStatus] | None = None,
        omit_empty_columns: bool = True,
        relativize: bool = False,
    ) -> pd.DataFrame:
        """
        High-level summary of the Experiment with one row per arm. Any values missing at
        compute time will be represented as None. Columns where every value is None will
        be omitted by default.

        The DataFrame computed will contain one row per arm and the following columns:
            - trial_index: The trial index of the arm
            - arm_name: The name of the arm
            - trial_status: The status of the trial (e.g. RUNNING, SUCCEEDED, FAILED)
            - status_reason: The reason for the trial status (e.g., failure,
                abandonment, early stopping), if applicable
            - generation_node: The name of the ``GenerationNode`` that generated the arm
            - **METADATA: Any metadata associated with the trial, as specified by the
                Experiment's runner.run_metadata_report_keys field
            - **METRIC_NAME: The observed mean of the metric specified, for each metric
            - **PARAMETER_NAME: The parameter value for the arm, for each parameter

        Args:
            trial_indices: If specified, only include these trial indices.
            omit_empty_columns: If True, omit columns where every value is None.
            trial_status: If specified, only include trials with this status.
            relativize: If True and:
                * experiment has a status quo on all of its ``BatchTrial``-s
                * OR a status quo trial among its ``Trial``-s,
                , relativize metrics against the status quo.
        """

        records = []
        data = self.lookup_data(trial_indices=trial_indices)

        # Relativize metrics if requested
        if relativize:
            if self.status_quo is None:
                raise UserInputError(
                    "Attempting to relativize the experiment data, however, "
                    "the experiment status quo is None. Please set the experiment "
                    "status quo, or set `relativize` = False"
                )

            data_df = data.relativize(
                status_quo_name=self.status_quo.name,
                as_percent=True,
                include_sq=True,
            ).df
        else:
            data_df = data.df

        # Filter trials by trial_indices and trial_statuses
        trials = self.extract_relevant_trials(
            trial_indices=list(trial_indices) if trial_indices is not None else None,
            trial_statuses=trial_statuses,
        )
        # Iterate through trials, and for each trial, iterate through its arms
        # and add a record for each arm.
        for trial in trials:
            for arm in trial.arms:
                # Find the observed means for each metric, placing None if not found
                observed_means = {}
                for metric in self.metrics.keys():
                    try:
                        observed_means[metric] = data_df[
                            (data_df["trial_index"] == trial.index)
                            & (data_df["arm_name"] == arm.name)
                            & (data_df["metric_name"] == metric)
                        ]["mean"].item()
                    except (ValueError, KeyError):
                        # ValueError if there is no row for the (trial, arm, metric).
                        # KeyError if the df is empty and missing one of the columns.
                        observed_means[metric] = None

                # Find the arm's associated generation method from the trial via the
                # GeneratorRuns if possible
                grs = [gr for gr in trial.generator_runs if arm in gr.arms]
                generation_node = grs[0]._generation_node_name if len(grs) > 0 else None

                # Find other metadata from the trial to include from the trial based
                # on the experiment's runner
                metadata = (
                    {
                        key: value
                        for key, value in trial.run_metadata.items()
                        if key in none_throws(self.runner).run_metadata_report_keys
                    }
                    if self.runner is not None
                    else {}
                )

                # Construct the record
                record = {
                    "trial_index": trial.index,
                    "arm_name": arm.name,
                    "trial_status": trial.status.name,
                    "status_reason": trial.status_reason,
                    "generation_node": generation_node,
                    **metadata,
                    **observed_means,
                    **arm.parameters,
                }

                records.append(record)

        df = pd.DataFrame(records)
        if omit_empty_columns:
            df = df.loc[:, df.notnull().any()]

        # Format metric columns as percentages with 4 significant figures when
        # relativized
        if relativize:
            for metric_name in self.metrics.keys():
                if metric_name in df.columns:
                    df[metric_name] = df[metric_name].apply(
                        lambda x: (
                            f"{x:.4g}%"
                            if pd.notna(x) and x != 0.0
                            else ("0%" if pd.notna(x) else None)
                        )
                    )

        return df

    def add_auxiliary_experiment(
        self,
        purpose: AuxiliaryExperimentPurpose,
        auxiliary_experiment: AuxiliaryExperiment,
    ) -> None:
        """Add a (non-duplicated) auxiliary experiment to this experiment.

        This method adds the auxiliary experiment as the first element in the list
        of auxiliary experiments with the specified purpose. If the auxiliary is
        already present, it is moved to the first position in the list.

        Args:
            purpose: The purpose of the auxiliary experiment.
            auxiliary_experiment: The auxiliary experiment to add.
        """
        if purpose not in self.auxiliary_experiments_by_purpose:
            # if no aux experiment, make aux the first one
            self.auxiliary_experiments_by_purpose[purpose] = [auxiliary_experiment]
            return

        # Add or move auxiliary_experiment to be the first element
        # Adding to the first and use the order as a default tie-breaker when multiple
        # auxiliary experiments are present but only one is going to be used.
        self.auxiliary_experiments_by_purpose[purpose] = [auxiliary_experiment] + [
            item
            for item in self.auxiliary_experiments_by_purpose[purpose]
            if item != auxiliary_experiment
        ]

    def find_auxiliary_experiment_by_name(
        self,
        purpose: AuxiliaryExperimentPurpose,
        auxiliary_experiment_name: str,
        raise_if_not_found: bool = False,
    ) -> AuxiliaryExperiment | None:
        """Find the aux experiment with the given name and purpose in the experiment.

        Args:
            purpose: The purpose of the aux experiment.
            auxiliary_experiment_name: The name of the aux experiment.

        Returns:
            The aux experiment with the given name and purpose, or None if not found.
            if raise_if_not_found is True, raises a ValueError if not found.
        """
        found_aux_exp = None
        if purpose in self.auxiliary_experiments_by_purpose:
            for auxiliary_experiment in self.auxiliary_experiments_by_purpose[purpose]:
                if auxiliary_experiment.experiment.name == auxiliary_experiment_name:
                    found_aux_exp = auxiliary_experiment
                    break

        if raise_if_not_found:
            if found_aux_exp is None:
                raise UserInputError(
                    f"Auxiliary experiment {auxiliary_experiment_name} is not "
                    f"found for purpose {purpose}."
                )
        return found_aux_exp

    def validate_auxiliary_experiment(
        self,
        source_experiment: Experiment,
        purpose: AuxiliaryExperimentPurpose,
    ) -> AuxiliaryExperimentValidation:
        """Validate a source auxiliary experiment against the current experiment as a
           target experiment based on a given purpose.

        Args:
            source_experiment: The source experiment to validate.
            purpose: The purpose of the auxiliary experiment.

        Returns:
            An AuxiliaryExperimentValidation object containing the validation result.
        """

        match purpose:
            case AuxiliaryExperimentPurpose.TRANSFERABLE_EXPERIMENT:
                overlapping_parameters = (
                    source_experiment.search_space.get_overlapping_parameters(
                        self.search_space
                    )
                )
                is_valid = len(overlapping_parameters) > 0
                invalid_reason = None if is_valid else "No overlapping parameters."
                return AuxiliaryExperimentValidation(
                    is_valid=is_valid,
                    invalid_reason=invalid_reason,
                    metadata=TransferLearningMetadata(
                        overlap_parameters=overlapping_parameters,
                    ),
                )
            case _:
                return AuxiliaryExperimentValidation(
                    is_valid=False,
                    invalid_reason="Validation not supported for auxiliary "
                    f"experiment purpose: {purpose}",
                )

    @property
    def auxiliary_experiments_by_purpose_for_storage(
        self,
    ) -> dict[AuxiliaryExperimentPurpose, list[AuxiliaryExperiment]]:
        """Tracks removed auxiliary experiments to be stored as inactive auxiliary
        experiments."""
        # Start with the current active auxiliary experiments.
        result = {
            purpose: list(auxiliary_experiments)
            for (
                purpose,
                auxiliary_experiments,
            ) in self.auxiliary_experiments_by_purpose.items()
        }
        # Iterate through the auxiliary experiments that were loaded and mark any
        # deleted ones as inactive.
        for (
            purpose,
            prev_auxiliary_experiments,
        ) in self._initial_auxiliary_experiments_by_purpose.items():
            # If the purpose is not in the new auxiliary experiments, mark all
            # previous auxiliary experiments as inactive.
            if purpose not in result:
                for pre_experiment in prev_auxiliary_experiments:
                    pre_experiment.is_active = False
                result[purpose] = prev_auxiliary_experiments
                continue
            # Mark any removed auxiliary experiments as inactive.
            for prev_auxiliary_experiment in prev_auxiliary_experiments:
                if prev_auxiliary_experiment not in result[purpose]:
                    prev_auxiliary_experiment.is_active = False
                    result[purpose].append(prev_auxiliary_experiment)
        return result


def add_arm_and_prevent_naming_collision(
    new_trial: Trial, old_trial: Trial, old_experiment_name: str | None = None
) -> None:
    # Add all of an old trial's arms to a new trial. Rename any arm with auto-generated
    # naming format to prevent naming collisions during warm-start. If an old
    # experiment name is provided, append that to the original arm name. Else, clear
    # the arm name. Preserves all names not matching the automatic naming format.
    # experiment is not named, clear the arm's name.
    # `arm_index` is 0 since all trials are single-armed.
    old_arm_name = none_throws(old_trial.arm).name
    has_default_name = bool(old_arm_name == old_trial._get_default_name(arm_index=0))
    if has_default_name:
        new_arm = none_throws(old_trial.arm).clone(clear_name=True)
        if old_experiment_name is not None:
            new_arm.name = f"{old_arm_name}_{old_experiment_name}"
        new_trial.add_arm(new_arm)
    else:
        try:
            new_trial.add_arm(none_throws(old_trial.arm).clone(clear_name=False))
        except ValueError as e:
            warnings.warn(
                f"Attaching arm {old_trial.arm} to trial {new_trial} while preserving "
                f"its name failed with error: {e}. Retrying with `clear_name=True`.",
                stacklevel=2,
            )
            new_trial.add_arm(none_throws(old_trial.arm).clone(clear_name=True))