From ea620f0f45f0d208fbd7bad425e4cf447a1c7e1e Mon Sep 17 00:00:00 2001 From: y0z <30489874+y0z@users.noreply.github.com> Date: Fri, 17 Jul 2026 13:12:49 +0900 Subject: [PATCH 1/3] Reduce memory consumption --- cmaes/_cmawm.py | 97 +++++++++++++++++++++++++------------ tests/test_cmawm.py | 113 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 179 insertions(+), 31 deletions(-) diff --git a/cmaes/_cmawm.py b/cmaes/_cmawm.py index 2278c14..c519f27 100644 --- a/cmaes/_cmawm.py +++ b/cmaes/_cmawm.py @@ -119,14 +119,21 @@ 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 +143,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 +234,50 @@ 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 +290,9 @@ 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)) - ] + self.m_z_lim_low, self.m_z_lim_up = self._get_discrete_param_limits( + 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] # 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)) diff --git a/tests/test_cmawm.py b/tests/test_cmawm.py index 32dd9fb..48bc7f3 100644 --- a/tests/test_cmawm.py +++ b/tests/test_cmawm.py @@ -2,10 +2,50 @@ import numpy as np from numpy.testing import assert_almost_equal +from numpy.testing import assert_allclose from unittest import TestCase from cmaes import CMA, CMAwM +class _DenseCMAwM(CMAwM): + """CMAwM using the dense discretization tables from the previous implementation.""" + + def __init__(self, *args, **kwargs): + bounds = kwargs["bounds"] + steps = kwargs["steps"] + 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 discrete_idx + ] + max_discrete = max([len(discrete) for discrete in discrete_list], default=0) + self._dense_z_space = np.full((len(discrete_idx), max_discrete), np.nan) + for i, discrete in enumerate(discrete_list): + self._dense_z_space[i, : len(discrete)] = discrete + self._dense_z_lim = ( + self._dense_z_space[:, 1:] + self._dense_z_space[:, :-1] + ) / 2 + for i in range(len(discrete_idx)): + self._dense_z_space[i][np.isnan(self._dense_z_space[i])] = np.nanmax( + self._dense_z_space[i] + ) + self._dense_z_lim[i][np.isnan(self._dense_z_lim[i])] = np.nanmax( + self._dense_z_lim[i] + ) + super().__init__(*args, **kwargs) + + def _get_discrete_param_indices(self, values): + return np.array( + [np.searchsorted(self._dense_z_lim[i], values[i]) for i in range(len(values))] + ) + + def _get_discrete_param_values(self, indices): + return self._dense_z_space[np.arange(len(indices)), indices] + + def _get_discrete_param_limit_values(self, indices): + return self._dense_z_lim[np.arange(len(indices)), indices] + + class TestCMAwM(TestCase): def test_no_discrete_spaces(self): mean = np.zeros(2) @@ -31,3 +71,76 @@ def test_no_discrete_spaces(self): solutions.append((cma_x, objective)) cma_optimizer.tell(solutions) cmawm_optimizer.tell(solutions) + + def test_sampling_is_equivalent_to_dense_discretization(self): + mean = np.array([0.0, 0.0, 0.0]) + bounds = np.array([[-30.0, 30.0], [-4.0, 4.0], [-5.0, 5.0]]) + steps = np.array([0.001, 0.5, 0.0]) + kwargs = { + "mean": mean, + "sigma": 1.3, + "bounds": bounds, + "steps": steps, + "seed": 1, + } + optimizer = CMAwM(**kwargs) + dense_optimizer = _DenseCMAwM(**kwargs) + + for _ in range(30): + solutions = [] + dense_solutions = [] + for _ in range(optimizer.population_size): + x_for_eval, x_for_tell = optimizer.ask() + dense_x_for_eval, dense_x_for_tell = dense_optimizer.ask() + + assert_allclose(x_for_eval, dense_x_for_eval, rtol=0, atol=1e-12) + assert_allclose(x_for_tell, dense_x_for_tell, rtol=0, atol=1e-12) + + # Pass the same objective value to both optimizers so this test isolates + # differences in sampling and discretization from objective round-off. + value = float(np.sum(x_for_eval**2)) + solutions.append((x_for_tell, value)) + dense_solutions.append((x_for_tell, value)) + + optimizer.tell(solutions) + dense_optimizer.tell(dense_solutions) + assert_allclose(optimizer.mean, dense_optimizer.mean, rtol=0, atol=1e-12) + assert_allclose(optimizer._A, dense_optimizer._A, rtol=0, atol=1e-12) + + def test_discrete_encoding_is_equivalent_to_dense_discretization(self): + bounds = np.array([[-30.0, 30.0], [-4.0, 4.0], [-1.0, 1.0]]) + steps = np.array([0.001, 0.5, 0.3]) + kwargs = { + "mean": np.zeros(3), + "sigma": 1.3, + "bounds": bounds, + "steps": steps, + "seed": 1, + } + optimizer = CMAwM(**kwargs) + dense_optimizer = _DenseCMAwM(**kwargs) + rng = np.random.RandomState(0) + + for _ in range(1000): + values = rng.uniform(bounds[:, 0] - 1, bounds[:, 1] + 1) + indices = optimizer._get_discrete_param_indices(values) + dense_indices = dense_optimizer._get_discrete_param_indices(values) + assert_allclose( + optimizer._get_discrete_param_values(indices), + dense_optimizer._get_discrete_param_values(dense_indices), + rtol=0, + atol=1e-12, + ) + + for dimension, size in enumerate(optimizer._discrete_space_size): + for limit_index in (0, size // 2, size - 2): + values = np.zeros(3) + values[dimension] = dense_optimizer._dense_z_lim[dimension, limit_index] + indices = optimizer._get_discrete_param_indices(values) + dense_indices = dense_optimizer._get_discrete_param_indices(values) + assert_allclose( + optimizer._get_discrete_param_values(indices), + dense_optimizer._get_discrete_param_values(dense_indices), + rtol=0, + atol=1e-12, + ) From 1af96fab1cf055b0303f72cbf5053a59e6076dcf Mon Sep 17 00:00:00 2001 From: y0z <30489874+y0z@users.noreply.github.com> Date: Thu, 30 Jul 2026 18:42:12 +0900 Subject: [PATCH 2/3] Remove DenseCMAwM --- tests/test_cmawm.py | 113 -------------------------------------------- 1 file changed, 113 deletions(-) diff --git a/tests/test_cmawm.py b/tests/test_cmawm.py index 48bc7f3..32dd9fb 100644 --- a/tests/test_cmawm.py +++ b/tests/test_cmawm.py @@ -2,50 +2,10 @@ import numpy as np from numpy.testing import assert_almost_equal -from numpy.testing import assert_allclose from unittest import TestCase from cmaes import CMA, CMAwM -class _DenseCMAwM(CMAwM): - """CMAwM using the dense discretization tables from the previous implementation.""" - - def __init__(self, *args, **kwargs): - bounds = kwargs["bounds"] - steps = kwargs["steps"] - 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 discrete_idx - ] - max_discrete = max([len(discrete) for discrete in discrete_list], default=0) - self._dense_z_space = np.full((len(discrete_idx), max_discrete), np.nan) - for i, discrete in enumerate(discrete_list): - self._dense_z_space[i, : len(discrete)] = discrete - self._dense_z_lim = ( - self._dense_z_space[:, 1:] + self._dense_z_space[:, :-1] - ) / 2 - for i in range(len(discrete_idx)): - self._dense_z_space[i][np.isnan(self._dense_z_space[i])] = np.nanmax( - self._dense_z_space[i] - ) - self._dense_z_lim[i][np.isnan(self._dense_z_lim[i])] = np.nanmax( - self._dense_z_lim[i] - ) - super().__init__(*args, **kwargs) - - def _get_discrete_param_indices(self, values): - return np.array( - [np.searchsorted(self._dense_z_lim[i], values[i]) for i in range(len(values))] - ) - - def _get_discrete_param_values(self, indices): - return self._dense_z_space[np.arange(len(indices)), indices] - - def _get_discrete_param_limit_values(self, indices): - return self._dense_z_lim[np.arange(len(indices)), indices] - - class TestCMAwM(TestCase): def test_no_discrete_spaces(self): mean = np.zeros(2) @@ -71,76 +31,3 @@ def test_no_discrete_spaces(self): solutions.append((cma_x, objective)) cma_optimizer.tell(solutions) cmawm_optimizer.tell(solutions) - - def test_sampling_is_equivalent_to_dense_discretization(self): - mean = np.array([0.0, 0.0, 0.0]) - bounds = np.array([[-30.0, 30.0], [-4.0, 4.0], [-5.0, 5.0]]) - steps = np.array([0.001, 0.5, 0.0]) - kwargs = { - "mean": mean, - "sigma": 1.3, - "bounds": bounds, - "steps": steps, - "seed": 1, - } - optimizer = CMAwM(**kwargs) - dense_optimizer = _DenseCMAwM(**kwargs) - - for _ in range(30): - solutions = [] - dense_solutions = [] - for _ in range(optimizer.population_size): - x_for_eval, x_for_tell = optimizer.ask() - dense_x_for_eval, dense_x_for_tell = dense_optimizer.ask() - - assert_allclose(x_for_eval, dense_x_for_eval, rtol=0, atol=1e-12) - assert_allclose(x_for_tell, dense_x_for_tell, rtol=0, atol=1e-12) - - # Pass the same objective value to both optimizers so this test isolates - # differences in sampling and discretization from objective round-off. - value = float(np.sum(x_for_eval**2)) - solutions.append((x_for_tell, value)) - dense_solutions.append((x_for_tell, value)) - - optimizer.tell(solutions) - dense_optimizer.tell(dense_solutions) - assert_allclose(optimizer.mean, dense_optimizer.mean, rtol=0, atol=1e-12) - assert_allclose(optimizer._A, dense_optimizer._A, rtol=0, atol=1e-12) - - def test_discrete_encoding_is_equivalent_to_dense_discretization(self): - bounds = np.array([[-30.0, 30.0], [-4.0, 4.0], [-1.0, 1.0]]) - steps = np.array([0.001, 0.5, 0.3]) - kwargs = { - "mean": np.zeros(3), - "sigma": 1.3, - "bounds": bounds, - "steps": steps, - "seed": 1, - } - optimizer = CMAwM(**kwargs) - dense_optimizer = _DenseCMAwM(**kwargs) - rng = np.random.RandomState(0) - - for _ in range(1000): - values = rng.uniform(bounds[:, 0] - 1, bounds[:, 1] + 1) - indices = optimizer._get_discrete_param_indices(values) - dense_indices = dense_optimizer._get_discrete_param_indices(values) - assert_allclose( - optimizer._get_discrete_param_values(indices), - dense_optimizer._get_discrete_param_values(dense_indices), - rtol=0, - atol=1e-12, - ) - - for dimension, size in enumerate(optimizer._discrete_space_size): - for limit_index in (0, size // 2, size - 2): - values = np.zeros(3) - values[dimension] = dense_optimizer._dense_z_lim[dimension, limit_index] - indices = optimizer._get_discrete_param_indices(values) - dense_indices = dense_optimizer._get_discrete_param_indices(values) - assert_allclose( - optimizer._get_discrete_param_values(indices), - dense_optimizer._get_discrete_param_values(dense_indices), - rtol=0, - atol=1e-12, - ) From 6ba95474bcbaf8c479394602d1c73305ac8311bb Mon Sep 17 00:00:00 2001 From: y0z <30489874+y0z@users.noreply.github.com> Date: Thu, 30 Jul 2026 18:56:40 +0900 Subject: [PATCH 3/3] Apply ruff --- cmaes/_cmawm.py | 18 ++++-------------- 1 file changed, 4 insertions(+), 14 deletions(-) diff --git a/cmaes/_cmawm.py b/cmaes/_cmawm.py index c519f27..e06ebd2 100644 --- a/cmaes/_cmawm.py +++ b/cmaes/_cmawm.py @@ -122,11 +122,7 @@ def __init__( 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 - ) + (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. @@ -247,9 +243,7 @@ def _get_discrete_param_indices(self, values: np.ndarray) -> np.ndarray: # ``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) - ) + 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. @@ -268,9 +262,7 @@ def _get_discrete_param_limit_values(self, indices: np.ndarray) -> np.ndarray: 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]: + 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) @@ -290,9 +282,7 @@ def tell(self, solutions: list[tuple[np.ndarray, float]]) -> None: return # margin correction updated_m_integer = mean[self._discrete_idx] - self.m_z_lim_low, self.m_z_lim_up = self._get_discrete_param_limits( - updated_m_integer - ) + 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))