diff --git a/cmaes/_cmawm.py b/cmaes/_cmawm.py index 2278c14..e06ebd2 100644 --- a/cmaes/_cmawm.py +++ b/cmaes/_cmawm.py @@ -119,14 +119,17 @@ def __init__( 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 + self._discrete_space_low = bounds[self._discrete_idx, 0] + requested_steps = steps[self._discrete_idx] + self._discrete_space_size = np.ceil( + (bounds[self._discrete_idx, 1] + requested_steps / 2 - self._discrete_space_low) + / requested_steps + ).astype(int) + # np.arange uses dtype(start + step) - dtype(start) as its actual step. + # Retain that behavior without materializing every value in the range. + self._discrete_space_step = ( + self._discrete_space_low + requested_steps - self._discrete_space_low + ) # continuous_space contains low and high of each parameter. self._continuous_idx = np.where(steps <= 0)[0] @@ -136,23 +139,17 @@ def __init__( ) # discrete_space - self._n_zdim = len(discrete_space) + self._n_zdim = len(self._discrete_idx) if self._n_zdim == 0: return + assert np.all(self._discrete_space_size >= 2), ( + "each discrete parameter must have at least two choices" + ) 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]) m_z = self._cma._mean[self._discrete_idx] # m_z_lim_low ->| mean vector |<- m_z_lim_up - m_pos = np.array([np.searchsorted(self.z_lim[i], m_z[i]) for i in range(len(m_z))]) - z_lim_low_index = np.clip(m_pos - 1, 0, self.z_lim.shape[1] - 1) - z_lim_up_index = np.clip(m_pos, 0, self.z_lim.shape[1] - 1) - self.m_z_lim_low = self.z_lim[np.arange(len(self.z_lim)), z_lim_low_index] - self.m_z_lim_up = self.z_lim[np.arange(len(self.z_lim)), z_lim_up_index] + self.m_z_lim_low, self.m_z_lim_up = self._get_discrete_param_limits(m_z) self._A = np.full(n_dim, 1.0) @@ -233,9 +230,46 @@ def _encode_discrete_params(self, discrete_param: np.ndarray) -> np.ndarray: x = (discrete_param - mean[self._discrete_idx]) * self._A[self._discrete_idx] + mean[ self._discrete_idx ] - x_pos = np.array([np.searchsorted(self.z_lim[i], x[i]) for i in range(len(x))]) - x_enc = self.z_space[np.arange(len(self.z_space)), x_pos] - return x_enc + x_pos = self._get_discrete_param_indices(x) + return self._get_discrete_param_values(x_pos) + + def _get_discrete_param_indices(self, values: np.ndarray) -> np.ndarray: + """Return indices of the closest discrete values, preferring the lower value on ties.""" + indices = np.floor( + (values - self._discrete_space_low) / self._discrete_space_step + 0.5 + ).astype(int) + indices = np.clip(indices, 0, self._discrete_space_size - 1) + + # ``floor(x + 0.5)`` selects the upper value at an exact midpoint, while + # ``np.searchsorted`` used by the previous implementation selected the lower one. + has_lower_limit = indices > 0 + lower_limits = self._get_discrete_param_limit_values(np.maximum(indices - 1, 0)) + indices -= has_lower_limit & (values <= lower_limits) + + # Correct a possible one-position error caused by floating-point division. + has_upper_limit = indices < self._discrete_space_size - 1 + upper_limits = self._get_discrete_param_limit_values( + np.minimum(indices, self._discrete_space_size - 2) + ) + indices += has_upper_limit & (values > upper_limits) + return indices + + def _get_discrete_param_values(self, indices: np.ndarray) -> np.ndarray: + return self._discrete_space_low + indices * self._discrete_space_step + + def _get_discrete_param_limit_values(self, indices: np.ndarray) -> np.ndarray: + lower_values = self._get_discrete_param_values(indices) + upper_values = self._get_discrete_param_values(indices + 1) + return (lower_values + upper_values) / 2 + + def _get_discrete_param_limits(self, values: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + positions = self._get_discrete_param_indices(values) + lower_indices = np.clip(positions - 1, 0, self._discrete_space_size - 2) + upper_indices = np.clip(positions, 0, self._discrete_space_size - 2) + return ( + self._get_discrete_param_limit_values(lower_indices), + self._get_discrete_param_limit_values(upper_indices), + ) def tell(self, solutions: list[tuple[np.ndarray, float]]) -> None: """Tell evaluation values""" @@ -248,16 +282,7 @@ def tell(self, solutions: list[tuple[np.ndarray, float]]) -> None: return # margin correction updated_m_integer = mean[self._discrete_idx] - m_pos = np.array( - [ - np.searchsorted(self.z_lim[i], updated_m_integer[i]) - for i in range(len(updated_m_integer)) - ] - ) - z_lim_low_index = np.clip(m_pos - 1, 0, self.z_lim.shape[1] - 1) - z_lim_up_index = np.clip(m_pos, 0, self.z_lim.shape[1] - 1) - self.m_z_lim_low = self.z_lim[np.arange(len(self.z_lim)), z_lim_low_index] - self.m_z_lim_up = self.z_lim[np.arange(len(self.z_lim)), z_lim_up_index] + self.m_z_lim_low, self.m_z_lim_up = self._get_discrete_param_limits(updated_m_integer) # 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))