from __future__ import annotations import functools import numpy as np from typing import cast from typing import Optional from cmaes import CMA from cmaes._cma import _is_valid_bounds try: from scipy import stats chi2_ppf = functools.partial(stats.chi2.ppf, df=1) norm_cdf = stats.norm.cdf except ImportError: from cmaes._stats import chi2_ppf # type: ignore from cmaes._stats import norm_cdf class CMAwM: """CMA-ES with Margin class with ask-and-tell interface. The code is adapted from https://github.com/EvoConJP/CMA-ES_with_Margin. Example: .. code:: import numpy as np from cmaes import CMAwM def ellipsoid_onemax(x, n_zdim): n = len(x) n_rdim = n - n_zdim ellipsoid = sum([(1000 ** (i / (n_rdim - 1)) * x[i]) ** 2 for i in range(n_rdim)]) onemax = n_zdim - (0. < x[(n - n_zdim):]).sum() return ellipsoid + 10 * onemax binary_dim, continuous_dim = 10, 10 dim = binary_dim + continuous_dim bounds = np.concatenate( [ np.tile([0, 1], (binary_dim, 1)), np.tile([-np.inf, np.inf], (continuous_dim, 1)), ] ) steps = np.concatenate([np.ones(binary_dim), np.zeros(continuous_dim)]) optimizer = CMAwM(mean=np.zeros(dim), sigma=2.0, bounds=bounds, steps=steps) evals = 0 while True: solutions = [] for _ in range(optimizer.population_size): x_for_eval, x_for_tell = optimizer.ask() value = ellipsoid_onemax(x_for_eval, binary_dim) evals += 1 solutions.append((x_for_tell, value)) optimizer.tell(solutions) if optimizer.should_stop(): break Args: mean: Initial mean vector of multi-variate gaussian distributions. sigma: Initial standard deviation of covariance matrix. bounds: Lower and upper domain boundaries for each parameter. steps: Each value represents a step of discretization for each dimension. Zero (or negative value) means a continuous space. n_max_resampling: A maximum number of resampling parameters (default: 100). If all sampled parameters are infeasible, the last sampled one will be clipped with lower and upper bounds. seed: A seed number (optional). population_size: A population size (optional). cov: A covariance matrix (optional). margin: A margin parameter (optional). """ # Paper: https://arxiv.org/abs/2205.13482 def __init__( self, mean: np.ndarray, sigma: float, bounds: np.ndarray, steps: np.ndarray, n_max_resampling: int = 100, seed: Optional[int] = None, population_size: Optional[int] = None, cov: Optional[np.ndarray] = None, margin: Optional[float] = None, ): # initialize `CMA` self._cma = CMA( mean, sigma, bounds, n_max_resampling, seed, population_size, cov ) n_dim = self._cma.dim population_size = self._cma.population_size self._n_max_resampling = n_max_resampling # split discrete space and continuous space assert len(bounds) == len(steps), "bounds and steps must be the same length" assert not np.isnan(steps).any(), "steps should not include NaN" self._discrete_idx = np.where(steps > 0)[0] discrete_list = [ np.arange(bounds[i][0], bounds[i][1] + steps[i] / 2, steps[i]) for i in self._discrete_idx ] max_discrete = max([len(discrete) for discrete in discrete_list], default=0) discrete_space = np.full((len(self._discrete_idx), max_discrete), np.nan) for i, discrete in enumerate(discrete_list): discrete_space[i, : len(discrete)] = discrete # continuous_space contains low and high of each parameter. self._continuous_idx = np.where(steps <= 0)[0] self._continuous_space = bounds[self._continuous_idx] assert _is_valid_bounds( self._continuous_space, mean[self._continuous_idx] ), "invalid bounds" # discrete_space self._n_zdim = len(discrete_space) if self._n_zdim == 0: return self.margin = margin if margin is not None else 1 / (n_dim * population_size) assert self.margin > 0, "margin must be non-zero positive value." self.z_space = discrete_space self.z_lim = (self.z_space[:, 1:] + self.z_space[:, :-1]) / 2 for i in range(self._n_zdim): self.z_space[i][np.isnan(self.z_space[i])] = np.nanmax(self.z_space[i]) self.z_lim[i][np.isnan(self.z_lim[i])] = np.nanmax(self.z_lim[i]) self.z_lim_low = np.concatenate( [self.z_lim.min(axis=1).reshape([self._n_zdim, 1]), self.z_lim], 1 ) self.z_lim_up = np.concatenate( [self.z_lim, self.z_lim.max(axis=1).reshape([self._n_zdim, 1])], 1 ) m_z = self._cma._mean[self._discrete_idx].reshape(([self._n_zdim, 1])) # m_z_lim_low ->| mean vector |<- m_z_lim_up self.m_z_lim_low = ( self.z_lim_low * np.where(np.sort(np.concatenate([self.z_lim, m_z], 1)) == m_z, 1, 0) ).sum(axis=1) self.m_z_lim_up = ( self.z_lim_up * np.where(np.sort(np.concatenate([self.z_lim, m_z], 1)) == m_z, 1, 0) ).sum(axis=1) self._A = np.full(n_dim, 1.0) @property def dim(self) -> int: """A number of dimensions""" return self._cma.dim @property def population_size(self) -> int: """A population size""" return self._cma.population_size @property def generation(self) -> int: """Generation number which is monotonically incremented when multi-variate gaussian distribution is updated.""" return self._cma.generation @property def mean(self) -> np.ndarray: """Mean Vector""" return self._cma.mean @property def _rng(self) -> np.random.RandomState: return self._cma._rng def reseed_rng(self, seed: int) -> None: self._cma.reseed_rng(seed) def ask(self) -> tuple[np.ndarray, np.ndarray]: """Sample a parameter and return (i) encoded x and (ii) raw x. The encoded x is used for the evaluation. The raw x is used for updating the distribution.""" for i in range(self._n_max_resampling): x = self._cma._sample_solution() if self._is_continuous_feasible(x[self._continuous_idx]): x_encoded = x.copy() if self._n_zdim > 0: x_encoded[self._discrete_idx] = self._encode_discrete_params( x[self._discrete_idx] ) return x_encoded, x x = self._cma._sample_solution() x[self._continuous_idx] = self._repair_continuous_params( x[self._continuous_idx] ) x_encoded = x.copy() if self._n_zdim > 0: x_encoded[self._discrete_idx] = self._encode_discrete_params( x[self._discrete_idx] ) return x_encoded, x def _is_continuous_feasible(self, continuous_param: np.ndarray) -> bool: if self._continuous_space is None: return True return cast( bool, np.all(continuous_param >= self._continuous_space[:, 0]) and np.all(continuous_param <= self._continuous_space[:, 1]), ) # Cast bool_ to bool. def _repair_continuous_params(self, continuous_param: np.ndarray) -> np.ndarray: if self._continuous_space is None: return continuous_param # clip with lower and upper bound. param = np.where( continuous_param < self._continuous_space[:, 0], self._continuous_space[:, 0], continuous_param, ) param = np.where( param > self._continuous_space[:, 1], self._continuous_space[:, 1], param ) return param def _encode_discrete_params(self, discrete_param: np.ndarray) -> np.ndarray: """Encode the values into discrete domain.""" mean = self._cma._mean x = (discrete_param - mean[self._discrete_idx]) * self._A[ self._discrete_idx ] + mean[self._discrete_idx] x = x.reshape([self._n_zdim, 1]) x_enc = ( self.z_space * np.where(np.sort(np.concatenate((self.z_lim, x), axis=1)) == x, 1, 0) ).sum(axis=1) return x_enc.reshape(self._n_zdim) def tell(self, solutions: list[tuple[np.ndarray, float]]) -> None: """Tell evaluation values""" self._cma.tell(solutions) mean = self._cma._mean sigma = self._cma._sigma C = self._cma._C if self._n_zdim == 0: return # margin correction updated_m_integer = mean[self._discrete_idx, np.newaxis] self.z_lim_low = np.concatenate( [self.z_lim.min(axis=1).reshape([self._n_zdim, 1]), self.z_lim], 1 ) self.z_lim_up = np.concatenate( [self.z_lim, self.z_lim.max(axis=1).reshape([self._n_zdim, 1])], 1 ) self.m_z_lim_low = ( self.z_lim_low * np.where( np.sort(np.concatenate([self.z_lim, updated_m_integer], 1)) == updated_m_integer, 1, 0, ) ).sum(axis=1) self.m_z_lim_up = ( self.z_lim_up * np.where( np.sort(np.concatenate([self.z_lim, updated_m_integer], 1)) == updated_m_integer, 1, 0, ) ).sum(axis=1) # calculate probability low_cdf := Pr(X <= m_z_lim_low) and up_cdf := Pr(m_z_lim_up < X) # sig_z_sq_Cdiag = self.model.sigma * self.model.A * np.sqrt(np.diag(self.model.C)) z_scale = ( sigma * self._A[self._discrete_idx] * np.sqrt(np.diag(C)[self._discrete_idx]) ) updated_m_integer = updated_m_integer.flatten() low_cdf = norm_cdf(self.m_z_lim_low, loc=updated_m_integer, scale=z_scale) up_cdf = 1.0 - norm_cdf(self.m_z_lim_up, loc=updated_m_integer, scale=z_scale) mid_cdf = 1.0 - (low_cdf + up_cdf) # edge case edge_mask = np.maximum(low_cdf, up_cdf) > 0.5 # otherwise side_mask = np.maximum(low_cdf, up_cdf) <= 0.5 if np.any(edge_mask): # modify mask (modify or not) modify_mask = np.minimum(low_cdf, up_cdf) < self.margin # modify sign modify_sign = np.sign(mean[self._discrete_idx] - self.m_z_lim_up) # distance from m_z_lim_up dist = ( sigma * self._A[self._discrete_idx] * np.sqrt( chi2_ppf(q=1.0 - 2.0 * self.margin) * np.diag(C)[self._discrete_idx] ) ) # modify mean vector mean[self._discrete_idx] = mean[ self._discrete_idx ] + modify_mask * edge_mask * ( self.m_z_lim_up + modify_sign * dist - mean[self._discrete_idx] ) # correct probability low_cdf = np.maximum(low_cdf, self.margin / 2.0) up_cdf = np.maximum(up_cdf, self.margin / 2.0) modified_low_cdf = low_cdf + (1.0 - low_cdf - up_cdf - mid_cdf) * ( low_cdf - self.margin / 2 ) / (low_cdf + mid_cdf + up_cdf - 3.0 * self.margin / 2) modified_up_cdf = up_cdf + (1.0 - low_cdf - up_cdf - mid_cdf) * ( up_cdf - self.margin / 2 ) / (low_cdf + mid_cdf + up_cdf - 3.0 * self.margin / 2) modified_low_cdf = np.clip(modified_low_cdf, 1e-10, 0.5 - 1e-10) modified_up_cdf = np.clip(modified_up_cdf, 1e-10, 0.5 - 1e-10) # modify mean vector and A (with sigma and C fixed) chi_low_sq = np.sqrt(chi2_ppf(q=1.0 - 2 * modified_low_cdf)) chi_up_sq = np.sqrt(chi2_ppf(q=1.0 - 2 * modified_up_cdf)) C_diag_sq = np.sqrt(np.diag(C))[self._discrete_idx] # simultaneous equations self._A[self._discrete_idx] = self._A[self._discrete_idx] + side_mask * ( (self.m_z_lim_up - self.m_z_lim_low) / ((chi_low_sq + chi_up_sq) * sigma * C_diag_sq) - self._A[self._discrete_idx] ) mean[self._discrete_idx] = mean[self._discrete_idx] + side_mask * ( (self.m_z_lim_low * chi_up_sq + self.m_z_lim_up * chi_low_sq) / (chi_low_sq + chi_up_sq) - mean[self._discrete_idx] ) def should_stop(self) -> bool: return self._cma.should_stop()