Multi-dimensional learning rate (sigma0) for AdamW

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Features
 
+- [#462](https://github.com/pybop-team/PyBOP/pull/462) - Enables multidimensional learning rate for `pybop.AdamW` with updated (more robust) integration testing. Fixes bug in `Minkowski` and `SumofPower` cost functions for gradient-based optimisers.
 - [#411](https://github.com/pybop-team/PyBOP/pull/411) - Updates notebooks with README in `examples/` directory, removes kaleido dependency and moves to nbviewer rendering, displays notebook figures with `notebook_connected` plotly renderer
 - [#6](https://github.com/pybop-team/PyBOP/issues/6) - Adds Monte Carlo functionality, with methods based on Pints' algorithms. A base class is added `BaseSampler`, in addition to `PintsBaseSampler`.
 - [#353](https://github.com/pybop-team/PyBOP/issues/353) - Allow user-defined check_params functions to enforce nonlinear constraints, and enable SciPy constrained optimisation methods

examples/scripts/spm_IRPropMin.py

@@ -54,4 +54,4 @@
 
 # Plot the cost landscape with optimisation path
 bounds = np.asarray([[0.5, 0.8], [0.4, 0.7]])
-pybop.plot2d(optim, bounds=bounds, steps=15)
+pybop.plot2d(optim, gradient=True, bounds=bounds, steps=15)

pybop/costs/_likelihoods.py

@@ -102,7 +102,7 @@ class GaussianLogLikelihood(BaseLikelihood):
     def __init__(
         self,
         problem: BaseProblem,
-        sigma0: Union[float, list[float], list[Parameter]] = 0.002,
+        sigma0: Union[float, list[float], list[Parameter]] = 1e-2,
         dsigma_scale: float = 1.0,
     ):
         super().__init__(problem)

pybop/costs/fitting_costs.py

@@ -122,8 +122,8 @@ def compute(
         r = np.asarray([y[signal] - self._target[signal] for signal in self.signal])
         e = np.sum(np.sum(np.abs(r) ** 2, axis=0), axis=0)
 
-        if calculate_grad is True:
-            de = 2 * np.sum(np.sum((r * dy.T), axis=2), axis=1)
+        if calculate_grad:
+            de = 2 * np.sum((r * dy.T), axis=(1, 2))
             return e, de
 
         return e
@@ -204,9 +204,9 @@ def compute(
         r = np.asarray([y[signal] - self._target[signal] for signal in self.signal])
         e = np.sum(np.abs(r) ** self.p) ** (1 / self.p)
 
-        if calculate_grad is True:
+        if calculate_grad:
             de = np.sum(
-                np.sum(r ** (self.p - 1) * dy.T, axis=2)
+                np.sum(np.sign(r) * np.abs(r) ** (self.p - 1) * dy.T, axis=2)
                 / (e ** (self.p - 1) + np.finfo(float).eps),
                 axis=1,
             )
@@ -287,8 +287,10 @@ def compute(
         r = np.asarray([y[signal] - self._target[signal] for signal in self.signal])
         e = np.sum(np.abs(r) ** self.p)
 
-        if calculate_grad is True:
-            de = self.p * np.sum(np.sum(r ** (self.p - 1) * dy.T, axis=2), axis=1)
+        if calculate_grad:
+            de = self.p * np.sum(
+                np.sign(r) * np.abs(r) ** (self.p - 1) * dy.T, axis=(1, 2)
+            )
             return e, de
 
         return e

pybop/optimisers/_adamw.py

@@ -86,8 +86,7 @@ def __init__(self, x0, sigma0=0.015, boundaries=None):
         self._b2 = 0.999
 
         # Step size
-        self._alpha = np.min(self._sigma0)
-
+        self._alpha = self._sigma0
         # Weight decay rate
         self._lam = 0.01
 

tests/integration/test_eis_parameterisation.py

@@ -55,11 +55,8 @@ def init_soc(self, request):
 
     @pytest.fixture(
         params=[
-            pybop.GaussianLogLikelihoodKnownSigma,
             pybop.GaussianLogLikelihood,
-            pybop.RootMeanSquaredError,
             pybop.SumSquaredError,
-            pybop.SumofPower,
             pybop.Minkowski,
             pybop.MAP,
         ]
@@ -82,8 +79,6 @@ def noise(self, sigma, values):
             pybop.SciPyDifferentialEvolution,
             pybop.CMAES,
             pybop.CuckooSearch,
-            pybop.NelderMead,
-            pybop.SNES,
             pybop.XNES,
         ]
     )
@@ -131,6 +126,9 @@ def optim(self, optimiser, model, parameters, cost, init_soc):
             "max_iterations": 250,
             "absolute_tolerance": 1e-6,
             "max_unchanged_iterations": 35,
+            "sigma0": [0.05, 0.05, 1e-3]
+            if isinstance(cost, pybop.GaussianLogLikelihood)
+            else 0.05,
         }
 
         if isinstance(cost, pybop.MAP):
@@ -139,10 +137,8 @@ def optim(self, optimiser, model, parameters, cost, init_soc):
                     0.2, 2.0
                 )  # Increase range to avoid prior == np.inf
 
-        # Set sigma0 and create optimiser
-        sigma0 = 0.05 if isinstance(cost, pybop.MAP) else None
-        optim = optimiser(sigma0=sigma0, **common_args)
-
+        # Create optimiser
+        optim = optimiser(**common_args)
         return optim
 
     @pytest.mark.integration

tests/integration/test_spm_parameterisations.py

@@ -112,7 +112,7 @@ def optim(self, optimiser, model, parameters, cost, init_soc):
                 problem, pybop.GaussianLogLikelihoodKnownSigma, sigma0=self.sigma0
             )
         elif cost in [pybop.SumofPower, pybop.Minkowski]:
-            cost = cost(problem, p=2)
+            cost = cost(problem, p=2.5)
         else:
             cost = cost(problem)
 
@@ -122,6 +122,9 @@ def optim(self, optimiser, model, parameters, cost, init_soc):
             "max_iterations": 250,
             "absolute_tolerance": 1e-6,
             "max_unchanged_iterations": 55,
+            "sigma0": [0.05, 0.05, 1e-3]
+            if isinstance(cost, pybop.GaussianLogLikelihood)
+            else 0.05,
         }
 
         if isinstance(cost, pybop.MAP):
@@ -131,16 +134,7 @@ def optim(self, optimiser, model, parameters, cost, init_soc):
                 )  # Increase range to avoid prior == np.inf
 
         # Set sigma0 and create optimiser
-        sigma0 = 0.05 if isinstance(cost, pybop.MAP) else None
-        optim = optimiser(sigma0=sigma0, **common_args)
-
-        # AdamW will use lowest sigma0 for learning rate, so allow more iterations
-        if issubclass(optimiser, (pybop.AdamW, pybop.IRPropMin)) and isinstance(
-            cost, pybop.GaussianLogLikelihood
-        ):
-            common_args["max_unchanged_iterations"] = 75
-            optim = optimiser(**common_args)
-
+        optim = optimiser(**common_args)
         return optim
 
     @pytest.mark.integration
@@ -221,6 +215,9 @@ def test_multiple_signals(self, multi_optimiser, spm_two_signal_cost):
             "max_iterations": 250,
             "absolute_tolerance": 1e-6,
             "max_unchanged_iterations": 55,
+            "sigma0": [0.035, 0.035, 6e-3, 6e-3]
+            if spm_two_signal_cost is pybop.GaussianLogLikelihood
+            else None,
         }
 
         # Test each optimiser
@@ -230,6 +227,9 @@ def test_multiple_signals(self, multi_optimiser, spm_two_signal_cost):
         if isinstance(spm_two_signal_cost, pybop.GaussianLogLikelihood):
             self.ground_truth = np.concatenate((self.ground_truth, combined_sigma0))
 
+        if issubclass(multi_optimiser, pybop.BasePintsOptimiser):
+            optim.set_max_unchanged_iterations(iterations=35, absolute_tolerance=1e-6)
+
         initial_cost = optim.cost(optim.parameters.initial_value())
         x, final_cost = optim.run()