Make cmaes a bit faster

CyberAgentAILab · Dec 11, 2023 · 6ef7f35 · 6ef7f35
1 parent 07d26c9
commit 6ef7f35
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/cmaes/_cma.py b/cmaes/_cma.py
@@ -378,9 +378,13 @@ def tell(self, solutions: list[tuple[np.ndarray, float]]) -> None:
 
         # (eq.47)
         rank_one = np.outer(self._pc, self._pc)
+
         rank_mu = np.sum(
-            np.array([w * np.outer(y, y) for w, y in zip(w_io, y_k)]), axis=0
+            w_io.reshape(-1, 1, 1) * np.einsum("...i,...j->...ij", y_k, y_k), axis=0
         )
+        # The above line is equivalent to:
+        # rank_mu = np.sum(np.array([w * np.outer(y, y) for w, y in zip(w_io, y_k)]), axis=0)
+
         self._C = (
             (
                 1