Fixes #581 - population_size setter

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Features
 
+- [#582](https://github.com/pybop-team/PyBOP/pull/582) - Fixes `population_size` arg for Pints' based optimisers, reshapes `parameters.rvs` to be parameter instances.
 - [#570](https://github.com/pybop-team/PyBOP/pull/570) - Updates the contour and surface plots, adds mixed chain effective sample size computation, x0 to optim.log
 - [#566](https://github.com/pybop-team/PyBOP/pull/566) - Adds `UnitHyperCube` transformation class, fixes incorrect application of gradient transformation.
 - [#569](https://github.com/pybop-team/PyBOP/pull/569) - Adds parameter specific learning rate functionality to GradientDescent optimiser.

pybop/optimisers/base_pints_optimiser.py

@@ -6,6 +6,7 @@
 from pints import NelderMead as PintsNelderMead
 from pints import Optimiser as PintsOptimiser
 from pints import ParallelEvaluator as PintsParallelEvaluator
+from pints import PopulationBasedOptimiser
 from pints import PopulationBasedOptimiser as PintsPopulationBasedOptimiser
 from pints import RectangularBoundaries as PintsRectangularBoundaries
 from pints import SequentialEvaluator as PintsSequentialEvaluator
@@ -117,6 +118,11 @@ def _set_up_optimiser(self):
             if tol_key in self.unset_options:
                 max_unchanged_kwargs[tol_key] = self.unset_options.pop(tol_key)
 
+        # Set population size (if applicable)
+        if "population_size" in self.unset_options:
+            population_size = self.unset_options.pop("population_size")
+            self.set_population_size(population_size)
+
     def _sanitise_inputs(self):
         """
         Check and remove any duplicate optimiser options.
@@ -514,3 +520,10 @@ def set_threshold(self, threshold=None):
             self._threshold = None
         else:
             self._threshold = float(threshold)
+
+    def set_population_size(self, population_size=None):
+        """
+        Set the population size for population-based optimisers, if specified.
+        """
+        if isinstance(self.optimiser, PopulationBasedOptimiser):
+            self.optimiser.set_population_size(population_size)

pybop/parameters/parameter.py

@@ -399,6 +399,9 @@ def rvs(self, n_samples: int = 1, apply_transform: bool = False) -> np.ndarray:
             samples = param.rvs(n_samples, apply_transform=apply_transform)
             all_samples.append(samples)
 
+        if n_samples > 1:
+            return np.asarray(all_samples).T
+
         return np.concatenate(all_samples)
 
     def get_sigma0(self, apply_transform: bool = False) -> list:

tests/unit/test_optimisation.py

@@ -5,6 +5,7 @@
 
 import numpy as np
 import pytest
+from pints import PopulationBasedOptimiser
 
 import pybop
 
@@ -206,6 +207,13 @@ def check_bounds_handling(optim, expected_bounds, should_raise=False):
                 warnings.simplefilter("always")
                 optimiser(cost=cost, unrecognised=10)
             assert not optim.optimiser.running()
+
+            # Check population setter
+            if isinstance(optim.optimiser, PopulationBasedOptimiser):
+                optim = pybop.Optimisation(
+                    cost=cost, optimiser=optimiser, population_size=100
+                )
+                assert optim.optimiser.population_size() == 100
         else:
             bounds_list = [
                 (lower, upper) for lower, upper in zip(bounds["lower"], bounds["upper"])