Source code for ax.modelbridge.transforms.rounding

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

import math
import random
from copy import copy

import numpy as np
import numpy.typing as npt
from ax.core.parameter_constraint import OrderConstraint
from ax.core.search_space import SearchSpace
from ax.core.types import TParameterization


[docs] def randomized_round(x: float) -> int: """Randomized round of x""" z = math.floor(x) return int(z + float(random.random() <= (x - z)))
[docs] def randomized_onehot_round(x: npt.NDArray) -> npt.NDArray: """Randomized rounding of x to a one-hot vector. x should be 0 <= x <= 1. If x includes negative values, they will be rounded to zero. """ neg_x = x < 0 x[neg_x] = 0 if len(x) == 1: return np.array([randomized_round(x[0])]) if sum(x) == 0: x = np.ones_like(x) x[neg_x] = 0 w = x / sum(x) hot = np.random.choice(len(w), size=1, p=w)[0] z = np.zeros_like(x) z[hot] = 1 return z
[docs] def strict_onehot_round(x: npt.NDArray) -> npt.NDArray: """Round x to a one-hot vector by selecting the max element. Ties broken randomly.""" if len(x) == 1: return np.round(x) argmax = x == max(x) x[argmax] = 1 x[~argmax] = 0 return randomized_onehot_round(x)
[docs] def contains_constrained_integer( search_space: SearchSpace, transform_parameters: set[str] ) -> bool: """Check if any integer parameters are present in parameter_constraints. Order constraints are ignored since strict rounding preserves ordering. """ for constraint in search_space.parameter_constraints: if isinstance(constraint, OrderConstraint): continue constraint_params = set(constraint.constraint_dict.keys()) if constraint_params.intersection(transform_parameters): return True return False
[docs] def randomized_round_parameters( parameters: TParameterization, transform_parameters: set[str] ) -> TParameterization: rounded_parameters = copy(parameters) for p_name in transform_parameters: # pyre: param is declared to have type `float` but is used as # pyre-fixme[9]: type `Optional[typing.Union[bool, float, str]]`. param: float = parameters.get(p_name) rounded_parameters[p_name] = randomized_round(param) return rounded_parameters