From 9a16a5637d1f4b1606a261d5bdb733edda7a49db Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Joseph=20Hopfm=C3=BCller?= Date: Sun, 24 Nov 2024 01:55:12 +0100 Subject: [PATCH] add optional parameter suggestion methods for Optuna trials --- src/single-core-regen/util/optuna_helpers.py | 360 +++++++++++++++++-- 1 file changed, 325 insertions(+), 35 deletions(-) diff --git a/src/single-core-regen/util/optuna_helpers.py b/src/single-core-regen/util/optuna_helpers.py index 123f61d..2067273 100644 --- a/src/single-core-regen/util/optuna_helpers.py +++ b/src/single-core-regen/util/optuna_helpers.py @@ -1,45 +1,335 @@ -def _optional_suggest(trial, name, range_or_value, log=False, step=None, type='int'): - # not a range - if not hasattr(range_or_value, '__iter__') or isinstance(range_or_value, str): - return range_or_value - +from typing import Any +from optuna import trial + + +def install_optional_suggests(): + trial.Trial.suggest_categorical_optional = suggest_categorical_optional_wrapper + trial.Trial.suggest_int_optional = suggest_int_optional_wrapper + trial.Trial.suggest_float_optional = suggest_float_optional_wrapper + + +def _is_listlike(obj: Any) -> bool: + return hasattr(obj, "__iter__") and not isinstance(obj, str) + + +def _optional_suggest( + *, + trial: trial.Trial, + name: str, + range_or_value: Any, + type: str, + log: bool = False, + step: int | float | None = None, + add_user: bool = False, + force: bool = False, + multiply: float | int = 1, + set_new: bool = True, +): + """ + Suggest a value for a parameter with more control over the process + + Parameters + ---------- + type : str + The type of the parameter + trial : optuna.trial.Trial + The trial object + name : str + The name of the parameter + range_or_value : Any + The range of values or a single value + log : bool, optional + Whether to use a logarithmic scale, by default False + step : int|float|None, optional + The step size, by default None + add_user : bool, optional + Whether to add the suggested value to the user attributes if not added as a parameter, by default False + force : bool, optional + Whether to force a single value to be suggested, by default False + multiply : float| int, optional + A multiplier to apply to the range or value, by default 1. Ignored for type "categorical". + set_new : bool, optional + Whether to override the parameter if it already exists, by default True + """ + + # value should be retrieved from trial + if not set_new and name in trial.params: + return trial.params[name] + + # value is not a list or tuple + if not _is_listlike(range_or_value): + range_or_value = (range_or_value,) + # range with only one value - if len(range_or_value) == 1: + if len(range_or_value) == 1 and not force: + if add_user: + trial.set_user_attr(name, range_or_value[0]) return range_or_value[0] - - if type == 'int': - step = step or 1 - return trial.suggest_int(name, *range_or_value, step=step, log=log) - if type == 'float': - return trial.suggest_float(name, *range_or_value, step=step, log=log) - - if type == 'categorical': + # normal operation + if type == "categorical": return trial.suggest_categorical(name, range_or_value) - + + # multiply range + range_or_value = tuple(multiply * x for x in range_or_value) + # + if len(range_or_value) > 2: + raise UserWarning("More than two values in range, using highest and lowest") + low = min(range_or_value) + high = max(range_or_value) + + if type == "float": + return trial.suggest_float(name, low, high, step=step, log=log) + + if type == "int": + step = step or 1 + lowi = int(low) + highi = int(high) + if lowi != low or highi != high: + raise ValueError(f"Range {low} to {high} (using multiplier {multiply}) is not valid for int") + return trial.suggest_int(name, lowi, highi, step=step, log=log) + raise ValueError(f"Unknown type: {type}") - -def optional_suggest_categorical(trial, name, choices_or_value): - return _optional_suggest(trial, name, choices_or_value, type='categorical') -def optional_suggest_int(trial, name, range_or_value, step=None, log=False): - return _optional_suggest(trial, name, range_or_value, step=step, log=log, type='int') +def suggest_categorical_optional( + trial: trial.Trial, + name: str, + choices_or_value: tuple[Any] | list[Any] | Any, + add_user: bool = False, + force: bool = False, + set_new: bool = True, +): + """ + Suggest a value for a categorical parameter with more control over the process -def optional_suggest_float(trial, name, range_or_value, step=None, log=False): - return _optional_suggest(trial, name, range_or_value, step=step, log=log, type='float') + Parameters + ---------- + trial : optuna.trial.Trial + The trial object + name : str + The name of the parameter + choices_or_value : tuple|list|Any + The choices or a single value + add_user : bool, optional + Whether to add the suggested value to the user attributes if not added as a parameter, by default False + force : bool, optional + Whether to suggest a single value as a parameter, by default False + set_new : bool, optional + Whether to override the parameter if it already exists, by default True + """ + return _optional_suggest( + trial=trial, name=name, range_or_value=choices_or_value, type="categorical", add_user=add_user, force=force, set_new=set_new + ) -def force_suggest_int(trial, name, range_or_value, step=1, log=False): - if not hasattr(range_or_value, '__iter__') or isinstance(range_or_value, str): - return trial.suggest_int(name, range_or_value, range_or_value, step=step, log=log) - return trial.suggest_int(name, *range_or_value, step=step, log=log) -def force_suggest_float(trial, name, range_or_value, step=None, log=False): - if not hasattr(range_or_value, '__iter__') or isinstance(range_or_value, str): - return trial.suggest_float(name, range_or_value, range_or_value, step=step, log=log) - return trial.suggest_float(name, *range_or_value, step=step, log=log) - -def force_suggest_categorical(trial, name, range_or_value): - if not hasattr(range_or_value, '__iter__') or isinstance(range_or_value, str): - return trial.suggest_categorical(name, [range_or_value]) - return trial.suggest_categorical(name, range_or_value) \ No newline at end of file +def suggest_int_optional( + trial: trial.Trial, + name: str, + range_or_value: tuple[int] | list[int] | int, + step: int = 1, + log: bool = False, + add_user: bool = False, + force: bool = False, + multiply: int = 1, + set_new: bool = True, +): + """ + Suggest a value for an integer parameter with more control over the process + + Parameters + ---------- + trial : optuna.trial.Trial + The trial object + name : str + The name of the parameter + range_or_value : tuple|list|int + The range of values or a single value. + step : int, optional + The step size, by default 1 + log : bool, optional + Whether to use a logarithmic scale, by default False + add_user : bool, optional + Whether to add the suggested value to the user attributes if not added as a parameter, by default False + force : bool, optional + Whether to suggest a single value as a parameter, by default False + """ + return _optional_suggest( + trial=trial, + name=name, + range_or_value=range_or_value, + step=step, + log=log, + type="int", + add_user=add_user, + force=force, + multiply=multiply, + set_new=set_new, + ) + + +def suggest_float_optional( + trial: trial.Trial, + name: str, + range_or_value: tuple[float] | list[float] | float, + step: float | None = None, + log: bool = False, + add_user: bool = False, + force: bool = False, + multiply: float = 1, + set_new: bool = True, +): + """ + Suggest a value for a float parameter with more control over the process + + Parameters + ---------- + trial : optuna.trial.Trial + The trial object + name : str + The name of the parameter + range_or_value : tuple|list|float + The range of values or a single value + step : float|None, optional + The step size, by default None + log : bool, optional + Whether to use a logarithmic scale, by default False + add_user : bool, optional + Whether to add the suggested value to the user attributes if not added as a parameter, by default False + force : bool, optional + Whether to suggest a single value as a parameter, by default False + multiply : float, optional + A multiplier to apply to the range or value, by default 1 + set_new : bool, optional + Whether to override the parameter if it already exists, by default True + """ + + return _optional_suggest( + trial=trial, + name=name, + range_or_value=range_or_value, + step=step, + log=log, + type="float", + add_user=add_user, + force=force, + multiply=multiply, + set_new=set_new, + ) + + +def suggest_categorical_optional_wrapper( + self: trial.Trial, + name: str, + choices_or_value: tuple[Any] | list[Any] | Any, + add_user: bool = False, + force: bool = False, + set_new: bool = True, +): + """ + Suggest a value for a categorical parameter with more control over the process + + Parameters + ---------- + name : str + The name of the parameter + choices_or_value : tuple|list|Any + The choices or a single value + add_user : bool, optional + Whether to add the suggested value to the user attributes if not added as a parameter, by default False + force : bool, optional + Whether to suggest a single value as a parameter, by default False + set_new : bool, optional + Whether to override the parameter if it already exists, by default True + """ + return suggest_categorical_optional( + trial=self, name=name, choices_or_value=choices_or_value, add_user=add_user, force=force, set_new=set_new + ) + + +def suggest_int_optional_wrapper( + self: trial.Trial, + name: str, + range_or_value: tuple[int] | list[int] | int, + step: int = 1, + log: bool = False, + add_user: bool = False, + force: bool = False, + multiply: int = 1, + set_new: bool = True, +): + """ + Suggest a value for an integer parameter with more control over the process + + Parameters + ---------- + name : str + The name of the parameter + range_or_value : tuple|list|int + The range of values or a single value. + step : int, optional + The step size, by default 1 + log : bool, optional + Whether to use a logarithmic scale, by default False + add_user : bool, optional + Whether to add the suggested value to the user attributes if not added as a parameter, by default False + force : bool, optional + Whether to suggest a single value as a parameter, by default False + """ + return suggest_int_optional( + trial=self, + name=name, + range_or_value=range_or_value, + step=step, + log=log, + add_user=add_user, + force=force, + multiply=multiply, + set_new=set_new, + ) + + +def suggest_float_optional_wrapper( + self: trial.Trial, + name: str, + range_or_value: tuple[float] | list[float] | float, + step: float | None = None, + log: bool = False, + add_user: bool = False, + force: bool = False, + multiply: float = 1, + set_new: bool = True, +): + """ + Suggest a value for a float parameter with more control over the process + + Parameters + ---------- + name : str + The name of the parameter + range_or_value : tuple|list|float + The range of values or a single value + step : float|None, optional + The step size, by default None + log : bool, optional + Whether to use a logarithmic scale, by default False + add_user : bool, optional + Whether to add the suggested value to the user attributes if not added as a parameter, by default False + force : bool, optional + Whether to suggest a single value as a parameter, by default False + multiply : float, optional + A multiplier to apply to the range or value, by default 1 + set_new : bool, optional + Whether to override the parameter if it already exists, by default True + """ + return suggest_float_optional( + trial=self, + name=name, + range_or_value=range_or_value, + step=step, + log=log, + add_user=add_user, + force=force, + multiply=multiply, + set_new=set_new, + )