added success eval, changed normal sampler and modularized cbo

cbx/dynamics/cbo.py

@@ -1,5 +1,8 @@
 from .pdyn import CBXDynamic
 
+
+def cbo_update(drift, lamda, dt, sigma, noise):
+    return -lamda * dt * drift + sigma * noise
 #%% CBO
 class CBO(CBXDynamic):
     r"""Consensus-based optimization (CBO) class
@@ -35,30 +38,17 @@ class CBO(CBXDynamic):
     def __init__(self, f, **kwargs) -> None:
         super().__init__(f, **kwargs)
 
-
-    def inner_step(self,) -> None:
-        r"""Performs one step of the CBO algorithm.
-
-        Parameters
-        ----------
-        None
-
-        Returns
-        -------
-        None
-        
-        """
-        # update, consensus point, drift and energy
-        self.consensus, energy = self.compute_consensus()
-        self.energy[self.consensus_idx] = energy
-        self.drift = self.x[self.particle_idx] - self.consensus
-
+    def cbo_step(self,):
+        # compute consensus, sets self.energy and self.consensus
+        self.compute_consensus()
+        # update drift and apply drift correction
+        self.drift = self.correction(self.x[self.particle_idx] - self.consensus)
+        # perform cbo update step
+        self.x[self.particle_idx] += cbo_update(
+            self.drift, self.lamda, self.dt, 
+            self.sigma, self.noise()
+        )
 
-        # compute noise
-        self.s = self.sigma * self.noise()
+    inner_step = cbo_step
 
-        #  update particle positions
-        self.x[self.particle_idx] = (
-            self.x[self.particle_idx] -
-            self.correction(self.lamda * self.dt * self.drift) +
-            self.s)
+

cbx/dynamics/cbo_memory.py

@@ -2,10 +2,10 @@
 from typing import Union
 #from scipy.special import logsumexp
 
-from .pdyn import CBXDynamic
+from .cbo import CBO, cbo_update
 
 #%% CBO_Memory
-class CBOMemory(CBXDynamic):
+class CBOMemory(CBO):
     r"""Consensus-based optimization with memory effects (CBOMemory) class
 
     This class implements the CBO algorithm with memory effects as described in [1]_ and [2]_. The algorithm
@@ -64,7 +64,8 @@ def __init__(self,
             self.sigma_memory = sigma_memory
 
         self.energy = self.f(self.x)
-        self.num_f_eval += np.ones(self.M, dtype=int) * self.N # update number of function evaluations   
+        self.num_f_eval += np.ones(self.M, dtype=int) * self.N # update number of function evaluations
+        self.ergexp = tuple([Ellipsis] + [None for _ in range(self.x.ndim-2)])
 
     def pre_step(self,):
         # save old positions
@@ -74,6 +75,13 @@ def pre_step(self,):
         # set new batch indices
         self.set_batch_idx()
 
+    def memory_step(self,):
+        # add memory step, first define new drift
+        self.drift = self.x[self.particle_idx] - self.y[self.particle_idx]
+        self.x[self.particle_idx] += cbo_update(
+            self.drift, self.lamda_memory, self.dt, 
+            self.sigma_memory, self.noise()
+        )
 
     def inner_step(self,) -> None:
         r"""Performs one step of the CBOMemory algorithm.
@@ -87,45 +95,23 @@ def inner_step(self,) -> None:
         None
         
         """
-
-        mind = self.consensus_idx
-        ind = self.particle_idx
-        # first update
-        self.consensus = self.compute_consensus(self.y[mind], self.energy[mind])
 
-        # compute consensus and memory drift
-        consensus_drift = self.x[ind] - self.consensus
-        memory_drift = self.x[ind] - self.y[ind]
-
-        # perform noise steps
-        # **NOTE**: noise always uses the ``drift`` attribute of the dynamic. 
-        # We first use the actual drift here and 
-        # then the memory difference
-        self.drift = consensus_drift
-        self.s_consensus = self.sigma * self.noise()
-        self.drift = memory_drift
-        self.s_memory = self.sigma_memory * self.noise()
-
-        # dynamcis update
-        # momentaneous positions of particles
-        self.x[ind] = (
-            self.x[ind] -
-            self.correction(self.lamda * self.dt * consensus_drift) -
-            self.lamda_memory * self.dt * memory_drift +
-            self.s_consensus + 
-            self.s_memory)
+        # first perform regular cbo step
+        self.cbo_step()
+        self.memory_step()
 
         # evaluation of objective function on all particles
-        energy_new = self.f(self.x[ind])    
-        self.num_f_eval += np.ones(self.M, dtype=int) * self.x[ind].shape[-2] # update number of function evaluations   
+        energy_new = self.eval_f(self.x[self.particle_idx])  
 
-        # historical best positions of particles
-        energy_expand = tuple([Ellipsis] + [None for _ in range(self.x.ndim-2)]) 
-        self.y[ind] = self.y[ind] + ((self.energy>energy_new)[energy_expand]) * (self.x[ind] - self.y[ind])
+        # historical best positions of particles 
+        self.y[self.particle_idx] += (
+            ((self.energy>energy_new)[self.ergexp]) * 
+            (self.x[self.particle_idx] - self.y[self.particle_idx])
+        )
         self.energy = np.minimum(self.energy, energy_new)
 
 
-    def compute_consensus(self, x_batch, energy) -> None:
+    def compute_consensus(self,) -> None:
         r"""Updates the weighted mean of the particles.
 
         Parameters
@@ -137,8 +123,10 @@ def compute_consensus(self, x_batch, energy) -> None:
         None
 
         """
-        c, _ = self._compute_consensus(energy, self.x[self.consensus_idx], self.alpha[self.active_runs_idx, :])
-        return c
+        self.consensus, _ = self._compute_consensus(
+            self.energy, self.y[self.consensus_idx], 
+            self.alpha[self.active_runs_idx, :]
+        )
 
     def update_best_cur_particle(self,) -> None:
         self.f_min = self.energy.min(axis=-1)

cbx/dynamics/cbs.py

@@ -1,40 +1,31 @@
 import warnings
-#%%
-import numpy as np
-from .pdyn import CBXDynamic
+from .cbo import CBO
 from ..scheduler import scheduler
 #%%
-class CBS(CBXDynamic):
+class CBS(CBO):
     def __init__(self, f, mode='sampling', noise='covariance',
                  M=1,
                  track_args=None, 
                  **kwargs):
         track_args = track_args if track_args is not None else{'names':[]}
         super().__init__(f, track_args=track_args, noise=noise, M=M, **kwargs)
+        self.sigma = 1.
 
         if self.batched:
             raise NotImplementedError('Batched mode not implemented for CBS!')
         if self.x.ndim > 3:
             raise NotImplementedError('Multi dimensional domains not implemented for CBS! The particle should have the dimension M x N x d, where d is an integer!')
-
-        self.exp_dt = np.exp(-self.dt)
-        if mode == 'sampling':
-            self.lamda = 1/(1 + self.alpha)
-        elif mode == 'optimization':
-            self.lamda = 1
-        else:
-            raise NotImplementedError("Invalid mode")
-
+
         if noise not in ['covariance', 'sampling']:
             raise warnings.warn('For CBS usually covariance or sampling noise is used!', stacklevel=2)
 
+        self.noise_callable.mode = mode
 
-    def inner_step(self,):
-        self.consensus, energy = self.compute_consensus()
-        self.energy = energy
-        self.drift = self.x - self.consensus
-        self.update_covariance()
-        self.x = self.consensus + self.exp_dt * self.drift + self.noise()
+    # def inner_step(self,):
+    #     self.consensus, energy = self.compute_consensus()
+    #     self.energy = energy
+    #     self.drift = self.x - self.consensus
+    #     self.x = self.consensus + self.exp_dt * self.drift + self.noise()
 
     def run(self, sched = 'default'):
             if self.verbosity > 0:

cbx/dynamics/pdyn.py

@@ -1,7 +1,7 @@
 #%%
 from ..noise import get_noise
 from ..correction import get_correction
-from ..scheduler import scheduler, multiply
+from ..scheduler import scheduler, multiply, effective_sample_size
 from ..utils.termination import max_it_term
 from ..utils.history import track_x, track_energy, track_update_norm, track_consensus, track_drift, track_drift_mean
 from ..utils.particle_init import init_particles
@@ -112,8 +112,8 @@ class ParticleDynamic:
         A callable that copies an array. The default is ``np.copy``.
     norm : Callable
         A callable that computes the norm of an array. The default is ``np.linalg.norm``.
-    normal : Callable
-        A callable that generates an array of random numbers that are distributed according to a normal distribution. The default is ``np.random.normal``.
+    sampler : Callable
+        A callable that generates an array of random numbers. The default is ``np.random.normal``.
 
     verbosity : int, optional
         The verbosity level. The default is 1.
@@ -134,7 +134,7 @@ def __init__(
             verbosity: int = 1,
             copy: Callable = None,
             norm: Callable = None,
-            normal: Callable = None,
+            sampler: Callable = None,
             post_process: Callable = None,
             ) -> None:
 
@@ -143,7 +143,8 @@ def __init__(
         # set utilities
         self.copy = copy if copy is not None else np.copy 
         self.norm = norm if norm is not None else np.linalg.norm
-        self.normal = normal if normal is not None else np.random.normal
+        rng = np.random.default_rng(12345)
+        self.sampler = sampler if sampler is not None else rng.standard_normal
 
         # init particles    
         self.init_x(x, M, N, d, x_min, x_max)
@@ -219,7 +220,7 @@ def check_f_dims(self, check=True) -> None:
             None
         """
         if check: # check if f returns correct shape
-            x = self.normal(0., 1., self.x.shape)
+            x = self.sampler(size=self.x.shape)
             if self.f(x).shape != (self.M,self.N):
                 raise ValueError("The given objective function does not return the correct shape!")
             self.num_f_eval += self.N * np.ones((self.M,), dtype=int) # number of function evaluations
@@ -340,6 +341,8 @@ def optimize(self,
             sched = scheduler([])
         elif sched == 'default':
             sched = self.default_sched()
+        elif sched == 'effective':
+            sched = effective_sample_size()
         else:
             self.sched = sched
 
@@ -535,7 +538,7 @@ def compute_mat_sqrt(A):
     """
     B, V = np.linalg.eigh(A)
     B = np.maximum(B,0.)
-    return V@(np.sqrt(B)[...,None]*V.transpose(0,2,1))
+    return V@(np.sqrt(B)[...,None]*np.moveaxis(V, -1, -2))
 
 class compute_consensus_default:
     def __init__(self, check_coeffs = False):
@@ -829,7 +832,7 @@ def set_noise(self, noise) -> None:
         """
         # set noise model
         if isinstance(noise, str):
-            self.noise_callable = get_noise(noise, norm=self.norm, sampler=self.normal)
+            self.noise_callable = get_noise(noise, self)
         elif callable(noise):
             self.noise_callable = noise
         else:
@@ -920,7 +923,9 @@ def compute_consensus(self,) -> None:
         # evaluation of objective function on batch
 
         energy = self.eval_f(self.x[self.consensus_idx]) # update energy
-        return self._compute_consensus(energy, self.x[self.consensus_idx], self.alpha[self.active_runs_idx, :])
-
-
+        self.consensus, energy = self._compute_consensus(
+            energy, self.x[self.consensus_idx], 
+            self.alpha[self.active_runs_idx, :]
+        )
+        self.energy[self.consensus_idx] = energy
 

cbx/dynamics/polarcbo.py

@@ -3,6 +3,7 @@
 from numpy.typing import ArrayLike
 from scipy.special import logsumexp
 
+from .pdyn import compute_mat_sqrt
 from .cbo import CBO
 
 #%% Kernel for PolarCBO
@@ -238,6 +239,33 @@ def kernel_factor(self, ):
     def compute_consensus(self,):
         x = self.x[self.consensus_idx]
         energy = self.eval_f(x)
-        neg_log_eval = self.kernel.neg_log(x[:,None,...], x[:,:,None,...])
-        return self._compute_consensus(energy, x, neg_log_eval, alpha = self.alpha[self.active_runs_idx, :, None], kernel_factor = self.kernel_factor())
+        self.neg_log_eval = self.kernel.neg_log(x[:,None,...], x[:,:,None,...])
+
+
+
+        self.consensus, energy = self._compute_consensus(
+            energy, x, self.neg_log_eval, 
+            alpha = self.alpha[self.active_runs_idx, :, None], 
+            kernel_factor = self.kernel_factor()
+        )
+        self.energy[self.consensus_idx] = energy
+
+    def update_covariance(self,):
+        r"""Update the covariance matrix :math:`\mathsf{C}(x_i)` of the noise model
+    
+        Parameters
+        ----------
+        None
+    
+        Returns
+        -------
+        None.
+    
+        """
+        weights = - self.kernel_factor() * self.neg_log_eval - self.alpha * self.energy
+        coeffs = np.exp(weights - logsumexp(weights, axis=(-1,), keepdims=True))
+
+        D = self.drift[...,None] * self.drift[...,None,:]
+        D = np.sum(D * coeffs[..., None, None], axis = -3)
+        self.Cov_sqrt = compute_mat_sqrt(D)