From b654695240440f1425f7da2dc484c9f8acb7120a Mon Sep 17 00:00:00 2001
From: Michael Clerx <michael.clerx@nottingham.ac.uk>
Date: Tue, 19 Dec 2023 11:49:13 +0000
Subject: [PATCH] Looking at xnes. Definitely some issue. Hard to debug as
 convergence is quite robust on simple problems (just slow).

---
 pints/_optimisers/_xnes.py | 16 ++++----
 pints/cptests/xnes.py      | 82 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 91 insertions(+), 7 deletions(-)
 create mode 100644 pints/cptests/xnes.py

diff --git a/pints/_optimisers/_xnes.py b/pints/_optimisers/_xnes.py
index a8dfbab32..137a9207c 100644
--- a/pints/_optimisers/_xnes.py
+++ b/pints/_optimisers/_xnes.py
@@ -47,7 +47,7 @@ def __init__(self, x0, sigma0=None, boundaries=None):
         self._bounded_ids = None  # Indices of those xs
 
         # Normalisation / distribution
-        self._mu = np.array(self._x0)   # Mean
+        self._mu = pints.vector(x0)     # Mean
         self._A = None                  # Covariance
 
         # Best solution seen
@@ -106,13 +106,13 @@ def _initialise(self):
         d = self._n_parameters
         n = self._population_size
 
-        # Learning rates
+        # Learning rates, see Table 1
         # TODO Allow changing before run() with method call
         self._eta_mu = 1
         # TODO Allow changing before run() with method call
         self._eta_A = 0.6 * (3 + np.log(d)) * d ** -1.5
 
-        # Pre-calculated utilities
+        # Pre-calculated utilities, see Table 1
         self._us = np.maximum(0, np.log(n / 2 + 1) - np.log(1 + np.arange(n)))
         self._us /= np.sum(self._us)
         self._us -= 1 / n
@@ -162,10 +162,12 @@ def tell(self, fx):
         self._mu += self._eta_mu * np.dot(self._A, Gd)
 
         # Update root of covariance matrix
-        Gm = np.dot(
-            np.array([np.outer(z, z).T - self._I for z in self._zs]).T,
-            self._us)
-        self._A *= scipy.linalg.expm(np.dot(0.5 * self._eta_A, Gm))
+        # Note that this is equation 11 (for the eta-sigma=eta-B case), not the
+        # more general equations 9&10 version given in Algorithm 1
+        Gm = 0.5 * np.sum([
+            u * (np.outer(z, z.T) - self._I)
+            for u, z in zip(self._us, self._zs)], axis=0)
+        self._A *= scipy.linalg.expm(0.5 * self._eta_A * Gm)
 
         # Update f_guessed on the assumption that the lowest value in our
         # sample approximates f(mu)
diff --git a/pints/cptests/xnes.py b/pints/cptests/xnes.py
new file mode 100644
index 000000000..25f1029cf
--- /dev/null
+++ b/pints/cptests/xnes.py
@@ -0,0 +1,82 @@
+#!/usr/bin/env python3
+#
+# Change point tests for XNES.
+#
+# This file is part of PINTS (https://github.com/pints-team/pints/) which is
+# released under the BSD 3-clause license. See accompanying LICENSE.md for
+# copyright notice and full license details.
+#
+import pints
+import pints.cptests as cpt
+
+
+def bounded_fitzhugh_nagumo(n_iterations=100):
+    """
+    Tests :class:`pints.XNES` on a bounded Fitzhugh-Nagumo model, and returns
+    a dictionary with ``error`` and ``distance``.
+
+    For details of the solved problem, see
+    :class:`pints.cptests.RunOptimiserOnBoundedUntransformedLogistic`.
+    """
+    problem = cpt.RunOptimiserOnBoundedFitzhughNagumo(
+        _method, n_iterations, _fguess)
+    return {
+        'error': problem.error(),
+        'distance': problem.distance()
+    }
+
+
+def bounded_untransformed_logistic(n_iterations=300):
+    """
+    Tests :class:`pints.XNES` on a bounded logistic model without
+    transformations, and returns a dictionary with ``error`` and ``distance``.
+
+    For details of the solved problem, see
+    :class:`pints.cptests.RunOptimiserOnBoundedUntransformedLogistic`.
+    """
+    problem = cpt.RunOptimiserOnBoundedUntransformedLogistic(
+        _method, n_iterations, _fguess)
+    return {
+        'error': problem.error(),
+        'distance': problem.distance()
+    }
+
+
+def rosenbrock(n_iterations=100):
+    """
+    Tests :class:`pints.XNES` on a Rosenbrock error and returns a dictionary
+    with ``error`` and ``distance``.
+
+    For details of the solved problem, see
+    :class:`pints.cptests.RunOptimiserOnRosenbrockError`.
+    """
+    problem = cpt.RunOptimiserOnRosenbrockError(_method, n_iterations, _fguess)
+    return {
+        'error': problem.error(),
+        'distance': problem.distance()
+    }
+
+
+def two_dim_parabola(n_iterations=50):
+    """
+    Tests :class:`pints.XNES` on a two-dimensional parabolic error and returns
+    a dictionary with entries ``error`` and ``distance``.
+
+    For details of the solved problem, see
+    :class:`pints.cptests.RunOptimiserOnTwoDimParabola`.
+    """
+    problem = cpt.RunOptimiserOnTwoDimParabola(_method, n_iterations, _fguess)
+    return {
+        'error': problem.error(),
+        'distance': problem.distance()
+    }
+
+
+_method = pints.XNES
+_fguess = True
+_change_point_tests = [
+    bounded_fitzhugh_nagumo,
+    bounded_untransformed_logistic,
+    rosenbrock,
+    two_dim_parabola,
+]