Source code for ax.modelbridge.generation_strategy

#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.

from inspect import signature
from typing import Any, Callable, Dict, List, NamedTuple, Optional

import pandas as pd
from ax.core.data import Data
from ax.core.experiment import Experiment
from ax.core.generator_run import GeneratorRun
from ax.modelbridge.base import ModelBridge
from ax.modelbridge.factory import Models
from ax.utils.common.typeutils import checked_cast, not_none


TModelFactory = Callable[..., ModelBridge]
MAX_CONDITIONS_GENERATED = 10000


def _filter_kwargs(function: Callable, **kwargs: Any) -> Any:
    """Filter out kwargs that are not applicable for a given function.
    Return a copy of given kwargs dict with only the required kwargs."""
    return {k: v for k, v in kwargs.items() if k in signature(function).parameters}


[docs]class GenerationStep(NamedTuple): """One step in the generation strategy, corresponds to a single model. Describes the model, how many arms will be generated with this model, what minimum number of observations is required to proceed to the next model, etc. """ model: Models num_arms: int min_arms_observed: int = 0 max_async_parallelism: Optional[int] = None # TODO[drfreund, T41983558]: use enforce_num_arms: bool = True # Kwargs to pass into the Models factory function. model_kwargs: Dict[str, Any] = None # Kwargs to pass into the Model's `.gen` function. model_gen_kwargs: Dict[str, Any] = None index: Optional[int] = None # Index of this step, set internally.
[docs]class GenerationStrategy: """GenerationStrategy describes which model should be used to generate new points for which trials, enabling and automating use of different models throughout the optimization process. For instance, it allows to use one model for the initialization trials, and another one for all subsequent trials. In the general case, this allows to automate use of an arbitrary number of models to generate an arbitrary numbers of arms described in the `arms_per_model` argument. """ _name: Optional[str] _steps: List[GenerationStep] _generated: List[str] # Arms generated in the current generation step. _observed: List[str] # Arms in the current step for which we observed data. _model: Optional[ModelBridge] # Current model. _data: Data # All data this strategy has been updated with. _curr: GenerationStep # Current step in the strategy. def __init__(self, steps: List[GenerationStep], name: Optional[str] = None) -> None: self._name = name self._steps = steps assert isinstance(self._steps, list), "Steps must be a GenerationStep list." for idx, step in enumerate(self._steps): if step.num_arms == -1: if idx < len(self._steps) - 1: raise ValueError( "Only last step in generation strategy can have num_arms " "set to -1 to indicate that the model in the step should " "be used to generate new arms indefinitely." ) elif step.num_arms < 1: raise ValueError("`num_arms` must be positive or -1 for all models.") self._steps[idx] = step._replace(index=idx) self._generated = [] self._observed = [] self._model = None self._data = Data() self._curr = steps[0] @property def name(self) -> str: """Name of this generation strategy. Defaults to a combination of model names provided in generation steps.""" if self._name: return self._name # pyre-ignore[16]: "`Models` have to attribute `__name__`", but they do. factory_names = (checked_cast(str, step.model.__name__) for step in self._steps) # Trim the "get_" beginning of the factory function if it's there. factory_names = (n[4:] if n[:4] == "get_" else n for n in factory_names) return "+".join(factory_names) @property def generator_changes(self) -> List[int]: """List of arm indices where a transition happened from one model to another.""" gen_changes = [step.num_arms for step in self._steps] return [sum(gen_changes[: i + 1]) for i in range(len(gen_changes))][:-1] @property def model(self) -> Optional[ModelBridge]: """Current model in this strategy.""" return self._model # pragma: no cover
[docs] def gen( self, experiment: Experiment, new_data: Optional[Data] = None, # Take in just the new data. n: int = 1, **kwargs: Any, ) -> GeneratorRun: """Produce the next points in the experiment.""" # Get arm signatures for each entry in new_data that is indeed new. new_arms = self._get_new_arm_signatures( experiment=experiment, new_data=new_data ) enough_observed = ( len(self._observed) + len(new_arms) ) >= self._curr.min_arms_observed unlimited_arms = self._curr.num_arms == -1 enough_generated = ( not unlimited_arms and len(self._generated) >= self._curr.num_arms ) remaining_arms = self._curr.num_arms - len(self._generated) # Check that minimum observed_arms is satisfied if it's enforced. if self._curr.enforce_num_arms and enough_generated and not enough_observed: raise ValueError( "All trials for current model have been generated, but not enough " "data has been observed to fit next model. Try again when more data " "are available." ) if ( self._curr.enforce_num_arms and not unlimited_arms and 0 < remaining_arms < n ): raise ValueError( f"Cannot generate {n} new arms as there are only {remaining_arms} " "remaining arms to generate using the current model." ) all_data = ( Data.from_multiple_data(data=[self._data, new_data]) if new_data else self._data ) if self._model is None: # Instantiate the first model. self._set_current_model(experiment=experiment, data=all_data, **kwargs) elif enough_generated and enough_observed: # Change to the next model. self._change_model(experiment=experiment, data=all_data, **kwargs) elif new_data is not None: # We're sticking with the current model, but update with new data self._model.update(experiment=experiment, data=new_data) gen_run = not_none(self._model).gen(n=n, **(self._curr.model_gen_kwargs or {})) # If nothing failed, update known data, _generated, and _observed. self._data = all_data self._observed.extend(new_arms) self._generated.extend(a.signature for a in gen_run.arms) return gen_run
[docs] def clone_reset(self) -> "GenerationStrategy": """Copy this generation strategy without it's state.""" return GenerationStrategy(name=self.name, steps=self._steps)
def _set_current_model( self, experiment: Experiment, data: Data, **kwargs: Any ) -> None: """Instantiate the current model with all available data. """ self._model = self._curr.model( # pyre-ignore[29] T41922457 **_filter_kwargs( self._curr.model, experiment=experiment, data=data, search_space=experiment.search_space, **(self._curr.model_kwargs or {}), **kwargs, ) ) def _change_model(self, experiment: Experiment, data: Data, **kwargs: Any) -> None: """Get a new model for the next step. """ # Increment the model if len(self._steps) == not_none(self._curr.index) + 1: raise ValueError(f"Generation strategy {self.name} is completed.") self._curr = self._steps[not_none(self._curr.index) + 1] # New step => reset _generated and _observed. self._generated, self._observed = [], [] self._set_current_model(experiment=experiment, data=data, **kwargs) def _get_new_arm_signatures( self, experiment: Experiment, new_data: Optional[Data] ) -> List[str]: new_signatures = [] if new_data is not None: for _, row in new_data.df.iterrows(): # If a row with the same trial index, arm name, and metric name # has already been seen in this generation strategy, the # data passed into this function is not entirely new. if not self._data.df.empty: if not pd.merge( new_data.df, self._data.df, on=["arm_name", "metric_name", "trial_index"], ).empty: arm = row["arm_name"] trial = row["trial_index"] metric = row["metric_name"] raise ValueError( f"Data for arm {arm} in trial {trial} for metric " f"{metric} has already been seen. Please only pass " "new data to `GenerationStrategy.gen`." ) if ( row["arm_name"] in experiment.arms_by_name and not experiment.trials.get(row["trial_index"]).status.is_failed ): new_signatures.append( experiment.arms_by_name.get(row["arm_name"]).signature ) return new_signatures