from typing import Optional, Sequence, List, Dict, Union
from ..jaml import JAML, JAMLCompatible
if False:
from optuna.trial import Trial
[docs]class OptimizationParameter(JAMLCompatible):
"""Base class for all optimization parameters."""
def __init__(
self,
parameter_name: str = '',
executor_name: Optional[str] = None,
prefix: str = 'JINA',
jaml_variable: Optional[str] = None,
):
self.parameter_name = parameter_name
if jaml_variable is None:
self.jaml_variable = f'{prefix}_{executor_name}_{parameter_name}'.upper()
else:
self.jaml_variable = jaml_variable
def _suggest(self, trial: 'Trial'):
"""
Suggest an instance of the parameter for a given trial.
:param trial: An instance of an Optuna Trial object
# noqa: DAR201
"""
raise NotImplementedError
[docs] def update_trial_params(self, trial: 'Trial', trial_params: Dict):
"""
Update the trial parameters given the variables and the trials
:param trial: An instance of an Optuna Trial object
:param trial_params: the parameters to update
# noqa: DAR201
"""
trial_params[self.jaml_variable] = self._suggest(trial)
[docs]class IntegerParameter(OptimizationParameter):
"""
Used for optimizing integer parameters with the FlowOptimizer.
For detailed information about sampling and usage see
https://optuna.readthedocs.io/en/stable/reference/generated/optuna.trial.Trial.html#optuna.trial.Trial.suggest_int
"""
def __init__(
self,
low: int,
high: int,
step_size: int = 1,
log: bool = False,
*args,
**kwargs,
):
super().__init__(*args, **kwargs)
self.low = low
self.high = high
if log and step_size != 1:
raise ValueError(
'''The step_size != 1 and log arguments cannot be used at the same time. When setting log argument to
True, set the step argument to 1. '''
)
self.step_size = step_size
self.log = log
def _suggest(self, trial: 'Trial'):
"""
Suggest an instance of the parameter for a given trial.
:param trial: An instance of an Optuna Trial object
# noqa: DAR201
"""
return trial.suggest_int(
name=self.jaml_variable,
low=self.low,
high=self.high,
step=self.step_size,
log=self.log,
)
[docs]class CategoricalParameter(OptimizationParameter):
"""
Used for optimizing categorical parameters with the FlowOptimizer.
For detailed information about sampling and usage see
https://optuna.readthedocs.io/en/stable/reference/generated/optuna.trial.Trial.html#optuna.trial.Trial.suggest_categorical
"""
def __init__(
self, choices: Sequence[Optional[Union[bool, int, float, str]]], *args, **kwargs
):
super().__init__(*args, **kwargs)
self.choices = choices
def _suggest(self, trial: 'Trial'):
"""
Suggest an instance of the parameter for a given trial.
:param trial: An instance of an Optuna Trial object
# noqa: DAR201
"""
return trial.suggest_categorical(name=self.jaml_variable, choices=self.choices)
[docs]class ExecutorAlternativeParameter(CategoricalParameter):
"""
Used for optimizing alternative executor parameters.
It selects from choices the same way as :class::CategoricalParameter does. Plus adds some inner parameters
that expose different `OptimizationParameters` for each of the specific options.
.. highlight:: python
.. code-block:: python
choices = ['executorA', 'executorB']
inner_parameters = {
'executorA': [IntegerParameter(..., parameter_name='executorA_param1'),
CategoricalParameter(..., parameter_name='executorA_param2')],
'executorB': [IntegerParameter(..., parameter_name='executorB_param1'),
CategoricalParameter(..., parameter_name='executorB_param2')]
}
parameter = ExecutorAlternativeParameter(choices=choices, inner_parameters=inner_parameters)
"""
def __init__(
self,
inner_parameters: Dict[str, List[OptimizationParameter]],
*args,
**kwargs,
):
super().__init__(*args, **kwargs)
assert len(self.choices) == len(
inner_parameters.keys()
), 'Every executor alternative needs to have its set of parameters'
self.inner_parameters = inner_parameters
[docs] def update_trial_params(self, trial: 'Trial', trial_params: Dict):
"""
Update the trial parameters given the variables and the trials
:param trial: An instance of an Optuna Trial object
:param trial_params: the parameters to update
# noqa: DAR201
"""
super().update_trial_params(trial, trial_params)
for param in self.inner_parameters[trial_params[self.jaml_variable]]:
param.update_trial_params(trial, trial_params)
[docs]def load_optimization_parameters(filepath: str):
"""
Loads optimization parameters from a `.yml` file and parses it with the JAML parser.
:param filepath: Path to a file that contains optimization parameters.
:return:: The loaded :class:`OptimizationParameter` objects.
"""
with open(filepath, encoding='utf8') as fp:
return JAML.load(fp)