import asyncio
from collections.abc import Awaitable
from datetime import datetime, timedelta
from typing import Any, Callable, Generic, Optional, Union
from ._runtime import Mode, blocking_property, operation
from .drivers import Driver
from .objects import Observation, Suggestion
from .requestor import NumParallelError, OptimizerRequestor, warn
[docs]
class Study(Generic[Mode]):
"""Configure and run an optimization study and retrieve its results.
Create a study with :func:`Client.create_study`, then provide a function
that evaluates each suggested set of parameters. For example::
def evaluate(study: Study, x1: float, x2: float) -> Observation:
observation = study.new_observation()
observation.add(x1**2 + x2**2)
return observation
study.configure(max_iter=30, num_parallel=3)
study.set_evaluator(evaluate)
study.run()
``run()`` requests suggestions, calls your function and submits its
observations automatically. To evaluate suggestions one at a time yourself::
while not study.is_done():
suggestion = study.get_suggestion()
observation = evaluate(study, **suggestion.kwargs)
study.add_observation(observation, suggestion.id)
With ``Client(asynchronous=True)``, use ``await`` for calls that contact the
server, including ``configure()`` and ``run()``. See :class:`Client` for
an example.
"""
def __init__(
self,
study_id: str,
project_id: str,
driver: Driver,
requestor: OptimizerRequestor,
extend_on_request: Optional[timedelta] = None,
extend_on_suggestion: Optional[timedelta] = None,
) -> None:
self._runtime = requestor.runtime
self.study_id = study_id
self.project_id = project_id
self._driver = driver
self._requestor = requestor
self._extend_on_request = extend_on_request
self._extend_on_suggestion = extend_on_suggestion
self.evaluator: Optional[
Callable[..., Union[Observation, Awaitable[Observation]]]
] = None
self._async_running = False
self.num_failed = 0
self.max_num_failed = 3
self.deleted = False
self.suggestions: dict[int, Suggestion] = {}
@property
def driver(self) -> Driver:
"""
The driver of the study. For a documentation see the :ref:`DriverReference`
of the corresponding driver.
"""
return self._driver
@property
def qualifier(self) -> str:
return self.project_id + "." + self.study_id
async def _post(
self,
purpose: str,
object: str,
operation: str,
data: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
result = await self._requestor._async_post(
purpose, object, operation, self.qualifier, data
)
await self._extend_shutdown_time(self._extend_on_request)
if purpose == "get suggestion":
await self._extend_shutdown_time(self._extend_on_suggestion)
return result
async def _run_task(self, purpose: str, task: str, data: dict[str, Any]) -> Any:
result = await self._requestor._async_run_task(
purpose, task, self.qualifier, data
)
await self._extend_shutdown_time(self._extend_on_request)
return result
async def _get(self, purpose: str, object: str, type: str) -> dict[str, Any]:
result = await self._requestor._async_get(purpose, object, type, self.qualifier)
await self._extend_shutdown_time(self._extend_on_request)
return result
async def _extend_shutdown_time(self, duration: Optional[timedelta]) -> None:
if duration is None:
return
r = await self._requestor._async_get("check server")
if "shutdown_at" not in r:
return
shutdown_at = datetime.fromisoformat(r["shutdown_at"])
now = datetime.now().astimezone()
if shutdown_at < now + duration:
new_shutdown_at = now + duration
await self._requestor._async_post(
"extend server shutdown time",
"server",
"update_shutdown_at",
data={"shutdown_at": new_shutdown_at.astimezone().isoformat()},
)
async def _delete_on_server(self) -> None:
"""Finish using this study through the client. Saved observations are kept.
Call ``study.close()`` with a synchronous client, or
``await study.close()`` with ``Client(asynchronous=True)``.
This does not stop the server. Closing the client also closes its studies,
so you usually do not need to call this method separately.
"""
if self.deleted:
return
try:
await self._post("delete study", "study", "delete")
except ConnectionError:
pass
self.deleted = True
def set_parameters(self, **kwargs: Any) -> None:
raise AttributeError(
"The method 'set_parameters()' is deprecated. Please, use the method 'configure()' instead."
)
async def _async_configure(
self,
num_parallel: int = 1,
max_iter: Optional[int] = None,
max_time: Optional[float] = None,
**kwargs: Any
) -> None:
"""Configures the study for its run. Example::
study.configure(max_iter=100, num_parallel=5)
Args:
num_parallel: Number of parallel observations of the evaluator function
(default: ``1``).
max_iter: Maximum number of evaluations of the evaluator 2
function (default: ``inf``).
max_time: Maximum optimization time in seconds (default: ``inf``).
.. note:: The full list of parameters depends on the chosen driver.
For a parameter description, see the the corresponding driver in the
:ref:`driver reference <DriverReference>`.
"""
data = kwargs
data["num_parallel"] = num_parallel
if max_iter is not None:
data["max_iter"] = max_iter
if max_time is not None:
data["max_time"] = max_time
await self._post("configure study", "study", "configure", data=data)
configure = operation(_async_configure, interruptible=False)
async def _async_get_configuration(self) -> dict[str, Any]:
"""Return the current configuration for the driver.
Use it to configure another study. Example::
config = study.get_configuration()
study2.configure(**config)
"""
answer = await self._post(
"get configuration", "driver", "configuration", data={}
)
return answer["configuration"]
get_configuration = operation(_async_get_configuration, interruptible=False)
@property
def configuration(self) -> dict[str, Any]:
"""The current study configuration.
With ``Client(asynchronous=True)``, use
``await study.get_configuration()`` instead of this property.
"""
return blocking_property(self, "get_configuration")
async def _async_start_clock(self) -> None:
"""The optimization stops after the time ``max_time``
(see :func:`configure`). This function resets the
clock to zero. Example::
study.start_clock()
.. note:: The clock is also set to zero by calling :func:`configure`.
"""
await self._post("start clock", "study", "start_clock")
start_clock = operation(_async_start_clock, interruptible=False)
async def _async_get_state(self, path: Optional[str] = None) -> dict[str, Any]:
"""Get state of the study. Example::
acquisition_time = study.get_state(
path="observation.acquisition_time"
)
Args:
path: A dot-separated path to a submodule or parameter.
If none, the full state is returned.
Returns: If path is None, a dictionary with the following entries
is returned:
:driver: Dictionary with information of driver state.
:observation: Dictionary with information of the latest observation.
:suggestion: Dictionary with information about the suggestion that
corresponds to the last observation
.. note:: A description of the meaning of each entry in the state
can be retrieved by :func:`describe`.
"""
answer = await self._post(
"get study state", "study", "get_state", data={"path": path}
)
return answer["state"]
get_state = operation(_async_get_state, interruptible=False)
async def _async_describe(self) -> dict[str, Any]:
"""Get description of all modules and their parameters that are used
by the study. Example::
description = study.describe()
print(description["observation"]["acquisition_time"])
print(description["driver"]["members"]["surrogates"]["0"])
Returns: A dictionary with the root entries:
:driver: Nested dictionary with description of submodules consisting
of a name and a descriptive text. If the entry describes a module,
it has an additional ``"members"`` entry with dictionaries describing
submodules and parameters.
:observation: Dictionary with a description of the parameters of
an observation.
:suggestion: Dictionary with a description of the parameters of
a suggestion of the driver.
"""
answer = await self._get("get description of the study", "study", "describe")
return answer["description"]
describe = operation(_async_describe, interruptible=False)
async def _async_historic_parameter_values(self, path: str) -> list[Any]:
"""Get the values of an internal parameter for each iteration of the study. Example::
acquisition_times = study.historic_parameter_values(
path="observation.acquisition_time")
Args:
path: A dot-separated path to the parameter.
.. note:: A description of the meaning of each parameter can be retrieved
by :func:`describe`.
"""
answer = await self._post(
"get historic parameter values",
"study",
"historic_parameter_values",
data={"path": path},
)
return answer["values"]
historic_parameter_values = operation(
_async_historic_parameter_values, interruptible=False
)
async def _async_get_observation_data(self) -> dict[str, Any]:
"""Get table with data of the observations. This can be used to copy
the data manually to another study. Example::
obs_data = study.get_observation_data()
other_study.add_observation_data(obs_data)
Returns: Dictionary, where each entry holds an equally long list of observation
data. The keys of the dictionary are:
:value: Observed value of black-box function
:derivative: Name of derivative parameter or None for non-derivative observation
:uncertainty: Uncertainty of observed value or None if no uncertainty
:model_name: Name of the surrogate model that is trained on the data or None
:design_value: Value of design parameters
:environment_value: Value of environment parameters or None if no environment
is specified
"""
answer = await self._get(
"get observation data", "study", "get_observation_data"
)
return answer["observation_data"]
get_observation_data = operation(_async_get_observation_data, interruptible=False)
async def _async_add_observation_data(self, data: dict[str, Any]) -> None:
"""Add data from another study to the study. Example::
obs_data = study.get_observation_data()
other_study.add_observation_data(obs_data)
Args:
data: Dict with observation data. See :func:`get_observation_data` for the details.
"""
observations: list[Observation] = []
design_values: list[list[Union[float, str]]] = []
obs = self.new_observation()
last_design_value = None
for idx in range(len(data["value"])):
design_value = data["design_value"][idx]
if last_design_value != design_value:
if last_design_value is not None:
observations.append(obs)
design_values.append(last_design_value)
obs = self.new_observation()
last_design_value = design_value
obs.add(
value=data["value"][idx],
derivative=data["derivative"][idx],
uncertainty=data["uncertainty"][idx],
model_name=data["model_name"][idx],
environment_value=data["environment_value"][idx],
)
if last_design_value is not None:
observations.append(obs)
design_values.append(last_design_value)
await self._async_add_many(observations, design_values)
add_observation_data = operation(_async_add_observation_data, interruptible=False)
async def _async_is_done(self) -> bool:
"""Checks if the study has finished. Example::
if study.is_done(): break
Returns: True if some stopping criterion set by
:func:`configure` was met.
.. note:: Before returning true, the function call waits until all open
suggestions have been added to the study.
"""
is_done = await self._is_done()
if is_done:
await self._wait_for_open_suggestions()
return is_done
is_done = operation(_async_is_done, interruptible=False)
async def _is_done(self) -> bool:
"""Check stopping criteria without waiting for open evaluations."""
answer = await self._get("get running status of study", "study", "is_done")
return bool(answer["is_done"])
async def _wait_for_open_suggestions(self) -> None:
while True:
state = await self._async_get_state()
num_open_suggestions = state["driver"]["num_open_suggestions"]
if num_open_suggestions > 0:
await asyncio.sleep(0.5)
else:
break
async def _async_get_suggestion(
self, environment_value: Optional[list[float]] = None
) -> Suggestion:
"""Get a new suggestion to be evaluated by the user.
Example::
def evaluate(study: Study, x1: float, x2: float) -> Observation:
obs = study.new_observation()
obs.add(x1**2 + x2**2)
return obs
suggestion = study.get_suggestion()
obs = evaluate(study, **suggestion.kwargs)
study.add_observation(observation=obs, suggestion_id=suggestion.id)
Args:
environment_value: If an environment is specified, this optional
argument specifies the list of variable environment parameter values,
for which a suggestion should be computed. E.g. ``[0.1, 1.2]``.
If an environment exists and no values are specified, the last known
environment values are used.
.. warning:: If ``num_parallel`` suggestions are already waiting for results,
this call waits. Submit an observation or clear an unfinished
suggestion before requesting another one. Use :func:`run` to
manage parallel evaluations automatically.
"""
from .async_study import _suggest
while True:
try:
return await _suggest(self, environment_value)
except NumParallelError:
await asyncio.sleep(0.2)
get_suggestion = operation(_async_get_suggestion, interruptible=True)
async def _get_suggestion(
self, environment_value: Optional[list[float]] = None
) -> Suggestion:
"""Request once, leaving capacity retries to the acquisition loop."""
answer = await self._post(
"get suggestion",
"suggestion",
"create",
data={"environment_value": environment_value},
)
s = Suggestion(sample=answer["sample"], id=answer["suggestion_id"])
self.suggestions[s.id] = s
return s
async def _async_clear_suggestion(
self, suggestion_id: int, message: str = ""
) -> None:
"""If the evaluation for a certain suggestion fails, the suggestion
can be cleared from the study. Example::
study.clear_suggestion(suggestion.id, 'Computation failed')
.. note:: The study only creates ``num_parallel`` suggestions (see
:func:`configure`) until it waits for an
observation to be added (see :func:`add_observation`)
or a suggestion to be cleared.
Args:
suggestion_id: Id of the suggestion to be cleared.
message: An optional message that is printed out.
"""
await self._post(
"clear suggestion",
"suggestion",
"remove",
data={"suggestion_id": suggestion_id, "message": message},
)
del self.suggestions[suggestion_id]
clear_suggestion = operation(_async_clear_suggestion, interruptible=False)
async def _async_clear_all_suggestions(self) -> None:
"""Clear all open suggestions. Example::
study.clear_all_suggestions()
.. note:: The study only creates ``num_parallel`` suggestions (see
:func:`configure`) until it waits for an
observation to be added (see :func:`add_observation`)
or a suggestion to be cleared.
"""
sids = list(self.suggestions.keys())
for sid in sids:
await self._async_clear_suggestion(sid)
clear_all_suggestions = operation(_async_clear_all_suggestions, interruptible=False)
[docs]
def new_observation(self) -> Observation:
"""Create a new :class:`Observation` object that allows
to add data via :func:`~Observation.add`. Example::
observation = study.new_observation()
observation.add(1.2)
observation.add(0.1, derivative='x1')
"""
return Observation()
async def _async_add_observation(
self,
observation: Observation,
suggestion_id: Optional[int] = None,
design_value: Optional[list[Union[float, str]]] = None,
environment_value: Optional[list[float]] = None,
acquisition_time: Optional[float] = None,
check: bool = True,
) -> None:
"""Adds an observation to the study. Example::
study.add_observation(observation, suggestion.id)
Args:
observation: :class:`Observation` object with added values
(see :func:`new_observation`)
suggestion_id: Id of the corresponding suggestion if it exists.
design_value: If the observation does not belong to an open suggestion,
the corresponding design value must be provided as a list of floats for
continuous and discrete parameters or string for categorial parameters.
E.g. ``[0.1, 2.0, 'cat1']``.
environment_value: If an environment parameters are specified, this
specifies the value of variable environment parameters as a list of
floats that is valid for all values added to the observation.
E.g. ``[1.0, 2.0]``.
Alternatively, one can also set different environment values for
each entry of the observation (see :func:`~Observation.add`).
acquisition_time: If the observation does not belong to an open
suggestion, it is possible to specify the time it took to retrieve
the observation (e.g. the computation time). This information can be
used to adapt the effort of computing the next suggestions.
check: If true, the validity of the observation is checked
"""
if not isinstance(observation, Observation):
raise TypeError(
"observation -> expected Observation object. "
+ "Check return value of evaluator function"
)
if suggestion_id is not None:
acquisition_time = (
observation.finished - self.suggestions[suggestion_id].created
)
await self._post(
"add observation",
"observation",
"create",
data={
"observation_data": observation.data,
"acquisition_time": acquisition_time,
"suggestion_id": suggestion_id,
"design_value": design_value,
"environment_value": environment_value,
"check": check,
},
)
if suggestion_id is not None:
del self.suggestions[suggestion_id]
add_observation = operation(_async_add_observation, interruptible=False)
async def _async_add_many(
self,
observations: list[Observation],
design_values: Optional[list[list[Union[float, str]]]] = None,
environment_values: Optional[list[list[float]]] = None,
acquisition_times: Optional[list[float]] = None,
check: bool = True,
suggestion_ids: Optional[list[int]] = None,
) -> None:
"""Adds many observations to the study. Example::
study.add_many(observations, design_values)
Args:
observations: List of :class:`Observation`
objects for each sample
(see :func:`new_observation`)
design_values: List of design values. This is required unless
``suggestion_ids`` is provided.
E.g. ``[[0.1, 1.0, 'cat1'], [0.2, 2.0, 'cat2']]``
environment_values: Optional list of environment values.
If not specified, the last known environment values are taken.
E.g. ``[[1.0, 2.0], [1.1, 2.3]]``
acquisition_times: Optional list of times required to acquire
each observation.
check: If true, the validity of the observation is checked
suggestion_ids: Optional list of ids of the corresponding open
suggestions. Their design and environment values are used and the
suggestions are consumed when the batch is added successfully. If
``design_values`` is also provided, the suggestion values take precedence.
.. note:: The surrogate is retrained once after the complete batch. Historic
variable and objective values are stored per observation, while other
driver-specific values represent the state after the complete batch.
"""
obs_data = []
for o in observations:
if not isinstance(o, Observation):
raise TypeError("observations -> expected Observation objects.")
obs_data.append(o.data)
if design_values is None and suggestion_ids is None:
raise ValueError("design_values or suggestion_ids must be provided.")
if suggestion_ids is not None:
if len(suggestion_ids) != len(observations):
raise ValueError(
"The number of suggestion_ids must equal the number of observations."
)
if len(set(suggestion_ids)) != len(suggestion_ids):
raise ValueError("suggestion_ids must be unique.")
if acquisition_times is None:
acquisition_times = [
observation.finished - self.suggestions[suggestion_id].created
for observation, suggestion_id in zip(
observations, suggestion_ids, strict=True
)
]
await self._post(
"add observations",
"observation",
"create_many",
data={
"observations": obs_data,
"design_values": design_values,
"suggestion_ids": suggestion_ids,
"environment_values": environment_values,
"acquisition_times": acquisition_times,
"check": check,
},
)
if suggestion_ids is not None:
for suggestion_id in suggestion_ids:
del self.suggestions[suggestion_id]
add_many = operation(_async_add_many, interruptible=False)
def set_objective(self, objective: Callable) -> None:
raise AttributeError(
"The method 'set_objective()' is deprecated. Please, use the method 'set_evaluator()' instead."
)
[docs]
def set_evaluator(
self, evaluator: Callable[..., Union[Observation, Awaitable[Observation]]]
) -> None:
"""Set the function that maps design parameters to an :class:`Observation`.
Example::
def evaluate(study: Study, x1: float, x2: float) -> Observation:
observation = study.new_observation()
observation.add(x1**2 + x2**2)
return observation
study.set_evaluator(evaluate)
Args:
evaluator: Function that evaluates the suggested parameters and returns an
:class:`Observation` with the result.
The function must accept a ``"study"`` argument as well as
an argument with the name of each design parameter and fixed environment
parameter.
You can define the function with ``def`` or ``async def``.
With ``num_parallel > 1``, several evaluations may run at once;
avoid changing shared data without protecting it from other calls.
"""
if self.evaluator is not None and evaluator != self.evaluator:
warn(
"The evaluator was already set before. Changing the evaluator function (that is the mapping between parameters and observations) can lead to undefined behaviour."
)
self.evaluator = evaluator
async def _async_run(self) -> None:
"""Evaluate suggestions until a stopping condition from :func:`configure` is met.
First register your evaluation function with :func:`set_evaluator`.
``run()`` requests suggestions, evaluates them and submits the observations.
``num_parallel`` controls how many evaluations can run at the same time.
The evaluation function can be defined with ``def`` or ``async def``.
If an evaluation fails, the study continues with a new suggestion. After
``max_num_failed`` consecutive failures (default: 3), it stops and reports
the last error. Errors when communicating with the server stop the run
immediately.
Interrupting the run with Ctrl-C, or cancelling an awaited run, stops the
study and clears its unfinished suggestions. Stopping may take time if a
server request is still being processed. Observations already submitted
are kept, so you can continue the study later.
"""
from .async_study import run
return await run(self)
run = operation(_async_run, interruptible=True)
close = operation(_delete_on_server, name="close")
async def _async_cancel_result(self, name: str, result: Any) -> None:
if name == "_async_get_suggestion" and result.id in self.suggestions:
await self._async_clear_suggestion(
result.id, "Suggestion request cancelled"
)