Repository: Helmholtz-AI-Energy/propulate (language: python)

Issue:
Remove ABC from surrogate class
We do not want to have explicit ABCs in Python, as they are relatively slow. Instead, all abstract methods shall just raise `NotImplementedError`.


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 4919c8b2..3f0675ff 100644
--- a/README.md
+++ b/README.md
@@ -8,7 +8,7 @@
 ![PyPI](https://img.shields.io/pypi/v/propulate)
 ![PyPI - Downloads](https://img.shields.io/pypi/dm/propulate)
 [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
-[![](https://img.shields.io/badge/Python-3.6+-blue.svg)](https://www.python.org/downloads/)
+[![](https://img.shields.io/badge/Python-3.9+-blue.svg)](https://www.python.org/downloads/)
 [![OpenSSF Best Practices](https://www.bestpractices.dev/projects/7785/badge)](https://www.bestpractices.dev/projects/7785)
 [![](https://img.shields.io/badge/Contact-propulate%40lists.kit.edu-orange)](mailto:propulate@lists.kit.edu)
 [![Documentation Status](https://readthedocs.org/projects/propulate/badge/?version=latest)](https://propulate.readthedocs.io/en/latest/?badge=latest)
@@ -37,9 +37,19 @@ soon as it becomes available.
 When a unit is finished evaluating a candidate and communicating the resulting fitness, it breeds the candidate for the
 next generation using the fitness values of all candidates it evaluated and received from other units so far.
 
-`Propulate` can be used for hyperparameter optimization and neural architecture search.
+`Propulate` can be used for hyperparameter optimization and neural architecture search at scale.
 It was already successfully applied in several accepted scientific publications. Applications include grid load
-forecasting, remote sensing, and structural molecular biology.
+forecasting, remote sensing, and structural molecular biology:
+
+> D. Coquelin, R. Sedona, M. Riedel, and M. Götz. **Evolutionary Optimization of Neural Architectures in Remote Sensing
+> Classification Problems**. IEEE International Geoscience and Remote Sensing Symposium IGARSS, Brussels, Belgium,
+> pp. 1587-1590 (2021). https://doi.org/10.1109/IGARSS47720.2021.9554309
+
+> O. Taubert, F. von der Lehr, A. Bazarova, et al. **RNA contact prediction by data efficient deep learning**. Commun
+> Biol 6, 913 (2023). https://doi.org/10.1038/s42003-023-05244-9
+
+> D. Coquelin, K. Flügel, M. Weiel, et al. **Harnessing Orthogonality to Train Low-Rank Neural Networks**. arXiv
+> preprint (2023). https://doi.org/10.48550/arXiv.2401.08505
 
 ## In more technical terms
 
@@ -75,9 +85,13 @@ discussions](https://github.com/Helmholtz-AI-Energy/propulate/discussions) :octo
 
 ## Installation
 
-You can install the latest stable release from PyPI: ``pip install propulate``
-If you need the latest updates, you can also install ``Propulate`` directly from the master branch at you own risk.
-Pull and run ``pip install -e .`` or ``python setup.py develop``.
+- You can install the **latest stable release** from PyPI: ``pip install propulate``
+- If you need the **latest updates**, you can also install ``Propulate`` directly from the master branch.
+Pull and run ``pip install .``.
+- If you want to run the **tutorials**, you can install the required dependencies via: ``pip install ."[tutorials]"``
+- If you want to **contribute** to ``Propulate`` as a developer, you need to install the required dependencies with the package:
+``pip install -e ."[dev]"``.
+
 ``Propulate`` depends on [``mpi4py``](https://mpi4py.readthedocs.io/en/stable/) and requires an MPI implementation under
 the hood. Currently, it is only tested with [OpenMPI](https://www.open-mpi.org/).
 
diff --git a/coverage.svg b/coverage.svg
index ae1c3500..7a18c7f4 100644
--- a/coverage.svg
+++ b/coverage.svg
@@ -9,13 +9,13 @@
     </mask>
     <g mask="url(#a)">
         <path fill="#555" d="M0 0h63v20H0z"/>
-        <path fill="#fe7d37" d="M63 0h36v20H63z"/>
+        <path fill="#a4a61d" d="M63 0h36v20H63z"/>
         <path fill="url(#b)" d="M0 0h99v20H0z"/>
     </g>
     <g fill="#fff" text-anchor="middle" font-family="DejaVu Sans,Verdana,Geneva,sans-serif" font-size="11">
         <text x="31.5" y="15" fill="#010101" fill-opacity=".3">coverage</text>
         <text x="31.5" y="14">coverage</text>
-        <text x="80" y="15" fill="#010101" fill-opacity=".3">54%</text>
-        <text x="80" y="14">54%</text>
+        <text x="80" y="15" fill="#010101" fill-opacity=".3">80%</text>
+        <text x="80" y="14">80%</text>
     </g>
 </svg>
diff --git a/propulate/population.py b/propulate/population.py
index 253a70dd..c0791db3 100644
--- a/propulate/population.py
+++ b/propulate/population.py
@@ -59,8 +59,8 @@ def __eq__(self, other) -> bool:
 
         Parameters
         ----------
-        other: Individual
-            other individual to compare individual under consideration to
+        other : Individual
+            Other individual to compare individual under consideration to.
 
         Returns
         -------
diff --git a/propulate/propulator.py b/propulate/propulator.py
index d0ffad8d..c6821026 100644
--- a/propulate/propulator.py
+++ b/propulate/propulator.py
@@ -153,6 +153,10 @@ def __init__(
         self.worker_sub_comm = (
             worker_sub_comm  # Sub communicator for each (multi rank) worker
         )
+
+        # Always initialize the ``Surrogate`` as the class attribute has to be set for ``None`` checks later.
+        self.surrogate = None if surrogate_factory is None else surrogate_factory()
+
         if self.propulate_comm is None:  # Exit early for sub-worker only ranks.
             # These ranks are not used for anything aside from the calculation of the user-defined loss function.
             return
@@ -167,9 +171,6 @@ def __init__(
         self.emigration_propagator = emigration_propagator  # Emigration propagator
         self.rng = rng
 
-        # Always initialize the ``Surrogate`` as the class attribute has to be set for ``None`` checks later.
-        self.surrogate = None if surrogate_factory is None else surrogate_factory()
-
         # Load initial population of evaluated individuals from checkpoint if exists.
         load_ckpt_file = self.checkpoint_path / f"island_{self.island_idx}_ckpt.pickle"
         if not os.path.isfile(load_ckpt_file):  # If not exists, check for backup file.
@@ -287,14 +288,14 @@ def loss_fn(individual):
 
         # Check if ``loss_fn`` is generator, prerequisite for surrogate model.
         if inspect.isgeneratorfunction(self.loss_fn):
-            last: float = 0.0
+            last = float("inf")
             for last in loss_fn(ind):
                 if self.surrogate is not None:
                     if self.surrogate.cancel(last):  # Check cancel for each yield.
                         break
-            ind.loss = last  # Set final loss as individual's loss.
+            ind.loss = float(last)  # Set final loss as individual's loss.
         else:
-            ind.loss = loss_fn(ind)  # Evaluate its loss.
+            ind.loss = float(loss_fn(ind))  # Evaluate its loss.
 
         # Add final value to surrogate.
         if self.surrogate is not None:
diff --git a/propulate/surrogate.py b/propulate/surrogate.py
index 3854db04..fd0d41c8 100644
--- a/propulate/surrogate.py
+++ b/propulate/surrogate.py
@@ -1,12 +1,16 @@
-from abc import ABC, abstractmethod
-from typing import Generic, TypeVar
+import random
+from typing import Dict, Tuple, TypeVar, Union
+
+import GPy
+import numpy as np
+from mpi4py import MPI
 
 from propulate.population import Individual
 
 T = TypeVar("T")
 
 
-class Surrogate(ABC, Generic[T]):
+class Surrogate:
     """
     Abstract Surrogate Model for Propulator.
 
@@ -22,7 +26,6 @@ class Surrogate(ABC, Generic[T]):
     2. Merge is commutative.
     """
 
-    @abstractmethod
     def __init__(self) -> None:
         """
         Initialize a new surrogate model.
@@ -30,10 +33,14 @@ def __init__(self) -> None:
         Propulate passes down a surrogate factory, as defined by the user,
         to each Propulator instance. That means ``__init__`` can be overwritten
         to take additional arguments.
+
+        Raises
+        ------
+        NotImplementedError
+            Not implemented (abstract base class method needs to be overridden).
         """
-        pass
+        raise NotImplementedError()
 
-    @abstractmethod
     def start_run(self, ind: Individual) -> None:
         """
         Signal that a new run is about to start.
@@ -46,10 +53,14 @@ def start_run(self, ind: Individual) -> None:
         ----------
         ind : propulate.Individual
             The individual to be evaluated.
+
+        Raises
+        ------
+        NotImplementedError
+            Not implemented (abstract base class method needs to be overridden).
         """
-        pass
+        raise NotImplementedError()
 
-    @abstractmethod
     def update(self, loss: float) -> None:
         """
         Update the surrogate model with the final loss.
@@ -60,10 +71,14 @@ def update(self, loss: float) -> None:
         ----------
         loss : float
             The final loss of the current run.
+
+        Raises
+        ------
+        NotImplementedError
+            Not implemented (abstract base class method needs to be overridden).
         """
-        pass
+        raise NotImplementedError()
 
-    @abstractmethod
     def cancel(self, loss: float) -> bool:
         """
         Evaluate surrogate to check if the current run should be cancelled.
@@ -80,10 +95,14 @@ def cancel(self, loss: float) -> bool:
         bool
            If the surrogate model determines that the current run
            will not result in a lower loss than previous runs.
+
+        Raises
+        ------
+        NotImplementedError
+            Not implemented (abstract base class method needs to be overridden).
         """
-        return False
+        raise NotImplementedError()
 
-    @abstractmethod
     def merge(self, data: T) -> None:
         """
         Merge the results of another surrogate model into itself.
@@ -96,10 +115,14 @@ def merge(self, data: T) -> None:
         ----------
         data : T
             All relevant information to update its model to the same state as the origin of the data.
+
+        Raises
+        ------
+        NotImplementedError
+            Not implemented (abstract base class method needs to be overridden).
         """
-        pass
+        raise NotImplementedError()
 
-    @abstractmethod
     def data(self) -> T:
         """
         Return all relevant information about the surrogate model for merging with another surrogate.
@@ -110,5 +133,533 @@ def data(self) -> T:
         -------
         T
             All relevant information to convey the current state of the surrogate model.
+
+        Raises
+        ------
+        NotImplementedError
+            Not implemented (abstract base class method needs to be overridden).
+        """
+        raise NotImplementedError()
+
+
+class StaticSurrogate(Surrogate):
+    """
+    Surrogate model using the best known run as baseline.
+
+    After the first run, each subsequent loss is compared to the baseline. Every run with a loss outside the margin of
+    the baseline is cancelled.
+
+    This model creates an internal index for the yielded losses during each run. That means the yield order and
+    frequency has to be consistent between runs. Otherwise, the indices of the baseline run won't match.
+
+    Loosely based on the paper:
+    Use of Static Surrogates in Hyperparameter Optimization
+    https://link.springer.com/article/10.1007/s43069-022-00128-w
+
+    Attributes
+    ----------
+    synthetic_id : int
+        An internal running index to keep track of the current runs loss index.
+    margin : float
+        A margin to be multiplied with the new incoming losses to compare them to the baseline.
+    first_run : bool
+        Flag for the first complete run.
+    baseline : np.ndarray
+        Chronological list of the losses of the best known run.
+    current_run : np.ndarray
+        Array to store the losses of the current run.
+
+    Methods
+    -------
+    __init__()
+        Overwrites the parents class's constructor to include the error margin.
+    start_run()
+        Reset the internal index and the current run array.
+    update()
+        Replace the baseline with the current run if the final loss is better.
+    cancel()
+        Cancel the current run if the loss is outside the margin of the baseline.
+    merge()
+        Merge an incoming baseline.
+    data()
+        Return the loss series of the best known run so far.
+
+    Notes
+    -----
+    The ``StaticSurrogate`` class implements all methods from the ``Surrogate`` class.
+
+    See Also
+    --------
+    :class:`Surrogate` : The parent class.
+    """
+
+    def __init__(self, margin: float = 0.8) -> None:
+        """
+        Initialize a static surrogate with a synthetic id, a margin, and empty arrays for baseline and current run.
+
+        Parameters
+        ----------
+        margin : float, optional
+            A margin on top of incoming losses for comparison with the baseline.
+            The default is 0.8.
+        """
+        self.synthetic_id: int = 0
+        self.margin: float = margin
+
+        # Canceling with ``cancel`` is only allowed after the first complete run.
+        # ``first_run`` keeps track of that.
+        self.first_run: bool = True
+
+        # Baseline is the best known run.
+        self.baseline: np.ndarray = np.zeros((0), dtype=float)
+        self.current_run: np.ndarray = np.zeros((0), dtype=float)
+
+    def start_run(self, ind: Individual) -> None:
+        """
+        Reset the internal index and the current run array.
+
+        Parameters
+        ----------
+        ind : Individual
+            The individual containing the current configuration.
+        """
+        # Reset to new run.
+        self.synthetic_id = 0
+        # Reset current run with correct size.
+        self.current_run = np.zeros((self.baseline.size), dtype=float)
+
+    def update(self, loss: float) -> None:
+        """
+        Replace the baseline with the current run if the final loss is better, or if there is no prior run.
+
+        Parameters
+        ----------
+        loss : float
+            The (unused) final loss of the current run.
+        """
+        if self.first_run:
+            self.baseline = self.current_run.copy()
+            self.first_run = False
+            return
+
+        # Check if current run is better than baseline.
+        if self.baseline[-1] > self.current_run[-1]:
+            self.baseline = self.current_run.copy()
+
+    def cancel(self, loss: float) -> bool:
+        """
+        Cancel the current run if the loss is outside the margin of the baseline.
+
+        Parameters
+        ----------
+        loss : float
+            The next interim loss of the current run.
+
+        Returns
+        -------
+        bool
+            True if the current run is cancelled, False otherwise.
+        """
+        self.synthetic_id += 1
+
+        # Cancel is only allowed after the first complete run.
+        if self.first_run:
+            self.current_run = np.append(self.current_run, loss)
+            return False
+
+        # Append loss to current run.
+        self.current_run[self.synthetic_id - 1] = loss
+
+        # Cancel if current run is outside margin of baseline.
+        if self.baseline[self.synthetic_id - 1] < loss * self.margin:
+            return True
+
+        return False
+
+    def merge(self, data: np.ndarray) -> None:
+        """
+        Replace the baseline with the incoming run if the final loss is better, or if there is no prior run.
+
+        Parameters
+        ----------
+        data : np.ndarray
+            The loss series of the incoming run.
+        """
+        # No prior data to merge with
+        if self.first_run:
+            self.baseline = data.copy()
+            self.first_run = False
+            return
+
+        # Merged run's final loss is better than baseline.
+        if self.baseline[-1] < data[-1]:
+            self.baseline = data.copy()
+
+    def data(self) -> np.ndarray:
+        """
+        Return the loss series of the best known run so far.
+
+        Returns
+        -------
+        np.ndarray
+            The loss series of the best known run so far.
+        """
+        # Return best run so far.
+        return self.baseline
+
+
+class DynamicSurrogate(Surrogate):
+    """
+    Surrogate model using a two-step Gaussian Process Regression to predict the final loss of a configuration.
+
+    The first global model is trained on the configurations with their final losses.
+    The second local model is trained on the interim losses of the current run.
+    The local model uses the global model's prediction as the mean function.
+
+    Loosely based on the paper:
+    Freeze-Thaw Bayesian Optimization
+    https://arxiv.org/abs/1406.3896
+
+    Attributes
+    ----------
+    limits : Union[Dict[str, Tuple[float, float]], Dict[str, Tuple[int, int]], Dict[str, Tuple[str, ...]]]
+        The hyperparameter configuration space's limits.
+    encodings : Dict[str, Dict[str, int]]
+        The encoding of categorical limit values to int.
+    history_X : np.ndarray
+        The encoded configurations.
+    history_Y : np.ndarray
+        The final losses corresponding to the encoded configurations.
+    current_run_data : np.ndarray
+        An array containing the interim losses of the current run.
+    current_encoding : np.ndarray
+        The encoded configuration of the current run.
+    global_kernel : GPy.kern.Kern
+        The kernel for the global Gaussian Process Regression model.
+    local_kernel : GPy.kern.Kern
+        The kernel for the local Gaussian Process Regression model.
+    mean_function : GPy.mappings.Mapping
+        The calculated mean function for the local model.
+    global_gpr : GPy.models.GPRegression
+        The global Gaussian Process Regression model.
+    local_gpr : GPy.models.GPRegression
+        The local Gaussian Process Regression model.
+    first_run : bool
+        A boolean indicating if this is the first run.
+    max_idx : int
+        The maximum index, i.e. the number of interim loss values for every run.
+    allowed_loss_margin : float
+        The allowed loss margin for cancelling a run.
+
+    Methods
+    -------
+    __init__()
+        Override the parent class's constructor to include the configuration space limits.
+    start_run()
+        Initialize the current run. Create the mean function for the local model.
+    update()
+        Update the global gaussian process model with the final loss of the current run.
+    cancel()
+        Cancel the current run if the local model predicts a higher final loss than the best so far.
+    merge()
+        Merge a configuration and loss tuple into the history arrays.
+    data()
+        Return the latest configuration and loss tuple from the history arrays.
+
+    Notes
+    -----
+    The ``DynamicSurrogate`` class implements all methods from the ``Surrogate`` class.
+
+    See Also
+    --------
+    :class:`Surrogate` : The parent class.
+    """
+
+    def __init__(
+        self,
+        limits: Union[
+            Dict[str, Tuple[float, float]],
+            Dict[str, Tuple[int, int]],
+            Dict[str, Tuple[str, ...]],
+        ],
+    ) -> None:
+        """
+        Initialize a dynamic surrogate with the configuration space limits.
+
+        Set the global and local kernels for Gaussian process regression. All other needed attributes are initialized as
+        empty arrays or None.
+
+        Parameters
+        ----------
+        limits : Union[Dict[str, Tuple[float, float]], Dict[str, Tuple[int, int]], Dict[str, Tuple[str, ...]]
+            The hyperparameter configuration space's limits.
+        """
+        self.limits = limits
+        self.encodings = self._create_encoding(limits)
+
+        # History arrays to store (encoded configuration, final loss) pairs
+        self.history_X: np.ndarray = np.array([[]])
+        self.history_Y: np.ndarray = np.array([[]])
+        # Arrays to store interim loss for the current run
+        self.current_run_data: np.ndarray = np.array([[]])
+        self.current_encoding: np.ndarray = np.array([[]])
+
+        self.global_kernel = GPy.kern.Matern32(input_dim=len(limits)) + GPy.kern.White(
+            input_dim=len(limits)
+        )
+        self.local_kernel = GPy.kern.Matern52(
+            1, variance=1.0, lengthscale=1.0, ARD=True
+        ) + GPy.kern.White(1, variance=1e-5)
+
+        self.mean_function = None
+
+        # Mean loss
+        self.global_gpr: GPy.models.GPRegression = None
+        # Newly created for every run
+        self.local_gpr: GPy.models.GPRegression = None
+
+        self.first_run: bool = True
+        self.max_idx: int = 0
+
+        self.allowed_loss_margin: float = 0.8
+
+        # Set seed for reproducibility.
+        rank = MPI.COMM_WORLD.Get_rank()
+        np.random.seed(42 * rank)
+        random.seed(42 * rank)
+
+    def start_run(self, ind: Individual) -> None:
+        """
+        Encode the config. given as individual and create the local GP's mean function from the global GP's prediction.
+
+        Parameters
+        ----------
+        ind : propulate.Individual
+            The individual containing the current configuration.
+        """
+        self.current_encoding = self.encode_configuration(ind)
+
+        if self.first_run:
+            return
+
+        # Use global model's prediction as the mean function for the local model.
+        self.mean_function = GPy.mappings.Constant(input_dim=1, output_dim=1)
+        mean, variance = self.global_gpr.predict(self.current_encoding)
+
+        self.mean_function.C = mean[0]
+
+        # Reset the local model and run data.
+        self.local_gpr = None
+        self.current_run_data = np.array([[]])
+
+    def update(self, loss: float) -> None:
+        """
+        Update the model with the current run's final loss and retrain the global Gaussian Process with the new data.
+
+        Parameters
+        ----------
+        loss : float
+            The final loss of the current run.
+        """
+        # Append the final loss to the history.
+        # Only if this is the first run:
+        if self.first_run:
+            if self.history_X.size == 0:
+                self.history_X = self.current_encoding
+            else:
+                self.history_X = np.vstack([self.history_X, self.current_encoding])
+            if self.history_Y.size == 0:
+                self.history_Y = np.array([[loss]])
+            else:
+                self.history_Y = np.vstack([self.history_Y, [loss]])
+
+            # Create global model and fit with the new data.
+            self.global_gpr = GPy.models.GPRegression(
+                self.history_X, self.history_Y, self.global_kernel
+            )
+            self.global_gpr.optimize()
+
+            self.first_run = False
+            return
+
+        # Or append if the run is finished.
+        if self.max_idx == len(self.current_encoding):
+            if self.history_X.size == 0:
+                self.history_X = self.current_encoding
+            else:
+                self.history_X = np.vstack([self.history_X, self.current_encoding])
+            if self.history_Y.size == 0:
+                self.history_Y = np.array([[loss]])
+            else:
+                self.history_Y = np.vstack([self.history_Y, [loss]])
+
+            # Fit global model with the new data.
+            self.global_gpr.set_XY(self.history_X, self.history_Y)
+            self.global_gpr.optimize()
+            return
+
+    def cancel(self, loss: float) -> bool:
+        """
+        Cancel the run if the local GP predicts a final loss higher than the best known loss with an allowed margin.
+
+        Parameters
+        ----------
+        loss : float
+            The next interim loss of the current run.
+
+        Returns
+        -------
+        bool
+            True if the run should be cancelled, False otherwise.
+        """
+        # The first run is never cancelled and has to run till the end.
+        if self.first_run:
+            # Here we also count how often cancel is called for later reference.
+            self.max_idx += 1
+            return False
+
+        # Append interim loss to the current run data.
+        if self.current_run_data.size == 0:
+            self.current_run_data = np.array([[loss]])
+        else:
+            self.current_run_data = np.vstack([self.current_run_data, [loss]])
+
+        if self.local_gpr is None:
+            self.local_gpr = GPy.models.GPRegression(
+                np.arange(1, self.current_run_data.shape[0] + 1).reshape(-1, 1),
+                self.current_run_data,
+                self.local_kernel,
+                mean_function=self.mean_function,
+            )
+        else:
+            # Retrain the local model with X consisting of synthetic indices.
+            self.local_gpr.set_XY(
+                np.arange(1, self.current_run_data.shape[0] + 1).reshape(-1, 1),
+                self.current_run_data,
+            )
+
+        self.local_gpr.optimize()
+
+        # Predict the final loss.
+        final_loss, final_variance = self.local_gpr.predict(np.array([[self.max_idx]]))
+        # Check if the final loss is lower than the best loss with margin so far.
+        if ((final_loss - final_variance) * self.allowed_loss_margin) > max(
+            self.history_Y
+        )[0]:
+            return True
+
+        return False
+
+    def merge(self, data: Tuple[np.ndarray, float]) -> None:
+        """
+        Merge a configuration and loss tuple into the history arrays and retrain the global GP with the new data.
+
+        Parameters
+        ----------
+        data : Tuple[np.ndarray, float]
+            The configurations and final losses of the merged run.
+        """
+        if len(data) != 2:
+            raise ValueError("Data must be a tuple of (configuration, loss)")
+        if data[1] == 0:
+            # Ignore 0 loss as this means some kind of unexpected behavior.
+            return
+
+        # Append new data loss to the existing.
+        if self.history_X.size == 0:
+            self.history_X = np.array([data[0]])
+        else:
+            self.history_X = np.vstack([self.history_X, [data[0]]])
+        if self.history_Y.size == 0:
+            self.history_Y = np.array([data[1]])
+        else:
+            self.history_Y = np.vstack([self.history_Y, [data[1]]])
+
+        # Fit global model with the new data.
+        if self.global_gpr is None:
+            self.global_gpr = GPy.models.GPRegression(
+                self.history_X, self.history_Y, self.global_kernel
+            )
+        else:
+            self.global_gpr.set_XY(self.history_X, self.history_Y)
+
+        self.global_gpr.optimize()
+
+    def data(self) -> Tuple[np.ndarray, float]:
+        """
+        Return the latest configuration and loss from the history arrays.
+
+        Returns
+        -------
+        Tuple[np.ndarray, float]
+            The latest configuration and loss.
+        """
+        # Return empty array if no data is available.
+        if self.history_X.size == 0 or self.history_Y.size == 0:
+            return (np.array([]), 0)
+
+        # Return the latest loss value as all other values are already shared.
+        return (self.history_X[-1], self.history_Y[-1])
+
+    def _create_encoding(
+        self,
+        limits: Union[
+            Dict[str, Tuple[float, float]],
+            Dict[str, Tuple[int, int]],
+            Dict[str, Tuple[str, ...]],
+        ],
+    ) -> Dict[str, Dict[str, int]]:
+        """
+        Create a mapping from categorical values to integers. Gaussian Process Regression only accepts numerical values.
+
+        Parameters
+        ----------
+        limits : Union[Dict[str, Tuple[float, float]], Dict[str, Tuple[int, int]], Dict[str, Tuple[str, ...]]
+            The hyperparameter configuration space's limits.
+
+        Returns
+        -------
+        Dict[str, Dict[str, int]]
+            The input dict but with the categorical values transformed to integers.
         """
-        pass
+        encodings = {}
+        for key, values in limits.items():
+            if isinstance(values, tuple) and isinstance(values[0], str):
+                # create a mapping for categorical values
+                encodings[key] = {v: i for i, v in enumerate(values)}
+        return encodings
+
+    def encode_configuration(
+        self,
+        config: Union[
+            Dict[str, float],
+            Dict[str, int],
+            Dict[str, str],
+        ],
+    ) -> np.ndarray:
+        """
+        Encode a configuration dictionary into a 2D array (required by the Gaussian Process Regression model).
+
+        Parameters
+        ----------
+        config : Union[Dict[str, float], Dict[str, int], Dict[str, str]]
+            The configuration dictionary to encode.
+
+        Returns
+        -------
+        np.ndarray
+            The encoded configuration.
+        """
+        # Initialize an empty array to ``limits`` size.
+        encoded_array = np.zeros((1, len(self.limits)), dtype=float)
+
+        # Fill the array with encoded values.
+        for i, key in enumerate(self.limits):
+            value = config[key]
+            if key in self.encodings:
+                # Ignore type hint as value can only be a string.
+                encoded_array[0, i] = self.encodings[key][value]
+            else:
+                encoded_array[0, i] = value
+
+        return encoded_array
diff --git a/propulate/utils.py b/propulate/utils/__init__.py
similarity index 94%
rename from propulate/utils.py
rename to propulate/utils/__init__.py
index 52b05436..c99c6a5b 100644
--- a/propulate/utils.py
+++ b/propulate/utils/__init__.py
@@ -9,8 +9,8 @@
 import numpy as np
 from mpi4py import MPI
 
-from .population import Individual, Particle
-from .propagators import (
+from ..population import Individual, Particle
+from ..propagators import (
     Compose,
     Conditional,
     CrossoverUniform,
@@ -21,6 +21,14 @@
     SelectMin,
     SelectUniform,
 )
+from . import benchmark_functions
+
+__all__ = [
+    "benchmark_functions",
+    "get_default_propagator",
+    "set_logger_config",
+    "make_particle",
+]
 
 
 def get_default_propagator(
@@ -148,6 +156,9 @@ def set_logger_config(
     if log_to_stdout:
         base_logger.addHandler(std_handler)
     if log_file is not None:
+        log_file = Path(log_file)
+        log_dir = log_file.parents[0]
+        log_dir.mkdir(parents=True, exist_ok=True)
         file_handler = logging.FileHandler(filename=log_file)
         file_handler.setFormatter(simple_formatter)
         base_logger.addHandler(file_handler)
diff --git a/tutorials/function_benchmark.py b/propulate/utils/benchmark_functions.py
similarity index 99%
rename from tutorials/function_benchmark.py
rename to propulate/utils/benchmark_functions.py
index 05d277b2..34d8230e 100644
--- a/tutorials/function_benchmark.py
+++ b/propulate/utils/benchmark_functions.py
@@ -381,7 +381,7 @@ def get_function_search_space(
 
     Parameters
     ----------
-    fname: str
+    fname : str
         The function name.
 
     Returns
@@ -412,8 +412,8 @@ def get_function_search_space(
     elif fname == "keane":
         function = keane
         limits = {
-            "a": (-10.0, 10.0),
-            "b": (-10.0, 10.0),
+            "a": (0.0, 10.0),
+            "b": (0.0, 10.0),
         }
     elif fname == "leon":
         function = leon
diff --git a/pyproject.toml b/pyproject.toml
index 37c2abab..260d0c61 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -21,23 +21,37 @@ classifiers = [
 ]
 dependencies = [
 	"deepdiff",
+    "matplotlib",
 	"mpi4py",
 	"numpy",
     "colorlog",
     "Gpy",
+    "h5py",
 ]
 
 [project.optional-dependencies]
 dev = [
     "pre-commit",
     "black",
+    "coverage",
     "ruff",
     "pytest",
     "pytest-cov",
+    "pytest-mpi",
     "sphinx-autoapi",
     "sphinx-rtd-theme",
     "sphinxcontrib-napoleon",
     "sphinxemoji",
+    "torch",
+    "torchvision",
+    "torchmetrics",
+]
+
+tutorials = [
+    "torch",
+    "torchvision",
+    "torchmetrics",
+    "lightning",
 ]
 
 [project.urls]
@@ -132,3 +146,6 @@ convention = "numpy"
 testpaths = [
     "tests",
 ]
+
+[tool.coverage.run]
+parallel = true
diff --git a/requirements.txt b/requirements.txt
index 8bd0597e..19551918 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,6 @@
-deepdiff>=5.8.0
-mpi4py==3.0.3
+deepdiff
+mpi4py
 numpy
 colorlog
 GPy
+h5py
diff --git a/tutorials/cmaes_example.py b/tutorials/cmaes_example.py
index 02a1628b..62d42a3d 100644
--- a/tutorials/cmaes_example.py
+++ b/tutorials/cmaes_example.py
@@ -2,12 +2,15 @@
 import pathlib
 import random
 
-from function_benchmark import get_function_search_space, parse_arguments
 from mpi4py import MPI
 
 from propulate import Propulator
 from propulate.propagators import ActiveCMA, BasicCMA, CMAPropagator
 from propulate.utils import set_logger_config
+from propulate.utils.benchmark_functions import (
+    get_function_search_space,
+    parse_arguments,
+)
 
 if __name__ == "__main__":
     comm = MPI.COMM_WORLD
diff --git a/tutorials/islands_example.py b/tutorials/islands_example.py
index 9799c45a..b23df370 100755
--- a/tutorials/islands_example.py
+++ b/tutorials/islands_example.py
@@ -3,12 +3,15 @@
 import random
 
 import numpy as np
-from function_benchmark import get_function_search_space, parse_arguments
 from mpi4py import MPI
 
 from propulate import Islands
 from propulate.propagators import SelectMax, SelectMin
 from propulate.utils import get_default_propagator, set_logger_config
+from propulate.utils.benchmark_functions import (
+    get_function_search_space,
+    parse_arguments,
+)
 
 if __name__ == "__main__":
     comm = MPI.COMM_WORLD
diff --git a/tutorials/multi_rank_workers_example.py b/tutorials/multi_rank_workers_example.py
index c72937d9..fd13e088 100644
--- a/tutorials/multi_rank_workers_example.py
+++ b/tutorials/multi_rank_workers_example.py
@@ -10,12 +10,12 @@
 from typing import Dict
 
 import numpy as np
-from function_benchmark import parse_arguments
 from mpi4py import MPI
 
 from propulate import Islands
 from propulate.propagators import SelectMax, SelectMin
 from propulate.utils import get_default_propagator, set_logger_config
+from propulate.utils.benchmark_functions import parse_arguments
 
 
 def parallel_sphere(params: Dict[str, float], comm: MPI.Comm = MPI.COMM_SELF) -> float:
diff --git a/tutorials/propulator_example.py b/tutorials/propulator_example.py
index c58d7853..d9437981 100755
--- a/tutorials/propulator_example.py
+++ b/tutorials/propulator_example.py
@@ -2,11 +2,14 @@
 import pathlib
 import random
 
-from function_benchmark import get_function_search_space, parse_arguments
 from mpi4py import MPI
 
 from propulate import Propulator
 from propulate.utils import get_default_propagator, set_logger_config
+from propulate.utils.benchmark_functions import (
+    get_function_search_space,
+    parse_arguments,
+)
 
 if __name__ == "__main__":
     comm = MPI.COMM_WORLD
diff --git a/tutorials/pso_example.py b/tutorials/pso_example.py
index 51015ec7..e32cde25 100644
--- a/tutorials/pso_example.py
+++ b/tutorials/pso_example.py
@@ -7,7 +7,6 @@
 import pathlib
 import random
 
-from function_benchmark import get_function_search_space, parse_arguments
 from mpi4py import MPI
 
 from propulate import Propulator, set_logger_config
@@ -19,6 +18,10 @@
     InitUniformPSO,
     VelocityClampingPSO,
 )
+from propulate.utils.benchmark_functions import (
+    get_function_search_space,
+    parse_arguments,
+)
 
 if __name__ == "__main__":
     comm = MPI.COMM_WORLD
diff --git a/tutorials/surrogate/cifar10_example.py b/tutorials/surrogate/cifar10_example.py
index 7db00c99..084eef35 100755
--- a/tutorials/surrogate/cifar10_example.py
+++ b/tutorials/surrogate/cifar10_example.py
@@ -13,8 +13,7 @@
 from torchvision.datasets import CIFAR10
 from torchvision.transforms import Compose, Normalize, ToTensor
 
-import tutorials.surrogate
-from propulate import Islands
+from propulate import Islands, surrogate
 from propulate.utils import get_default_propagator
 
 GPUS_PER_NODE: int = 1
@@ -301,6 +300,7 @@ def get_data_loaders(batch_size: int) -> Tuple[DataLoader, DataLoader]:
             shuffle=True,  # Shuffle data.
         )
 
+    # NOTE barrier only called, when dataset has not been downloaded yet
     if not hasattr(get_data_loaders, "barrier_called"):
         MPI.COMM_WORLD.Barrier()
 
@@ -472,10 +472,8 @@ def set_seeds(seed_value=42):
         generations=num_generations,  # Number of generations per worker
         num_islands=1,  # Number of islands
         checkpoint_path=log_path,
-        # surrogate_factory=lambda: tutorials.surrogate.MockSurrogate(),
-        # surrogate_factory=lambda: tutorials.surrogate.LogSurrogate(),
-        surrogate_factory=lambda: tutorials.surrogate.StaticSurrogate(),
-        # surrogate_factory=lambda: tutorials.surrogate.DynamicSurrogate(limits),
+        surrogate_factory=lambda: surrogate.StaticSurrogate(),
+        # surrogate_factory=lambda: surrogate.DynamicSurrogate(limits),
     )
     islands.evolve(  # Run evolutionary optimization.
         top_n=1,  # Print top-n best individuals on each island in summary.
diff --git a/tutorials/surrogate/dynamic_surrogate.py b/tutorials/surrogate/dynamic_surrogate.py
deleted file mode 100644
index b94e84b9..00000000
--- a/tutorials/surrogate/dynamic_surrogate.py
+++ /dev/null
@@ -1,368 +0,0 @@
-import random
-from typing import Dict, Tuple, Union
-
-import GPy
-import numpy as np
-from mpi4py import MPI
-
-from propulate import Surrogate
-from propulate.population import Individual
-
-
-# Gaussian Process Regression with an Exponential Decay Kernel Surrogate
-class DynamicSurrogate(Surrogate):
-    """
-    Surrogate model using a two-step Gaussian Process Regression to predict the final loss of a configuration.
-
-    The first global model is trained on the configurations with their final losses.
-    The second local model is trained on the interim losses of the current run.
-    The local model uses the global model's prediction as the mean function.
-
-    Loosely based on the paper:
-    Freeze-Thaw Bayesian Optimization
-    https://arxiv.org/abs/1406.3896
-
-    Attributes
-    ----------
-    limits : Union[Dict[str, Tuple[float, float]], Dict[str, Tuple[int, int]], Dict[str, Tuple[str, ...]]]
-        The hyperparameter configuration space's limits.
-    encodings : Dict[str, Dict[str, int]]
-        The encoding of categorical limit values to int.
-    history_X : np.ndarray
-        The encoded configurations.
-    history_Y : np.ndarray
-        The final losses corresponding to the encoded configurations.
-    current_run_data : np.ndarray
-        An array containing the interim losses of the current run.
-    current_encoding : np.ndarray
-        The encoded configuration of the current run.
-    global_kernel : GPy.kern.Kern
-        The kernel for the global Gaussian Process Regression model.
-    local_kernel : GPy.kern.Kern
-        The kernel for the local Gaussian Process Regression model.
-    mean_function : GPy.mappings.Mapping
-        The calculated mean function for the local model.
-    global_gpr : GPy.models.GPRegression
-        The global Gaussian Process Regression model.
-    local_gpr : GPy.models.GPRegression
-        The local Gaussian Process Regression model.
-    first_run : bool
-        A boolean indicating if this is the first run.
-    max_idx : int
-        The maximum index, i.e. the number of interim loss values for every run.
-    allowed_loss_margin : float
-        The allowed loss margin for cancelling a run.
-
-    Methods
-    -------
-    __init__()
-        Override the parent class's constructor to include the configuration space limits.
-    start_run()
-        Initialize the current run. Create the mean function for the local model.
-    update()
-        Update the global gaussian process model with the final loss of the current run.
-    cancel()
-        Cancel the current run if the local model predicts a higher final loss than the best so far.
-    merge()
-        Merge a configuration and loss tuple into the history arrays.
-    data()
-        Return the latest configuration and loss tuple from the history arrays.
-
-
-    Notes
-    -----
-    The ``DynamicSurrogate`` class implements all methods from the ``Surrogate`` class.
-
-    See Also
-    --------
-    :class:`Surrogate` : The parent class.
-    """
-
-    def __init__(
-        self,
-        limits: Union[
-            Dict[str, Tuple[float, float]],
-            Dict[str, Tuple[int, int]],
-            Dict[str, Tuple[str, ...]],
-        ],
-    ) -> None:
-        """
-        Initialize a dynamic surrogate with the configuration space limits.
-
-        Set the global and local kernels for Gaussian process regression. All other needed attributes are initialized as
-        empty arrays or None.
-
-        Parameters
-        ----------
-        limits : Union[Dict[str, Tuple[float, float]], Dict[str, Tuple[int, int]], Dict[str, Tuple[str, ...]]
-            The hyperparameter configuration space's limits.
-        """
-        self.limits = limits
-        self.encodings = self._create_encoding(limits)
-
-        # History arrays to store (encoded configuration, final loss) pairs
-        self.history_X: np.ndarray = np.array([[]])
-        self.history_Y: np.ndarray = np.array([[]])
-        # Arrays to store interim loss for the current run
-        self.current_run_data: np.ndarray = np.array([[]])
-        self.current_encoding: np.ndarray = np.array([[]])
-
-        self.global_kernel = GPy.kern.Matern32(input_dim=len(limits)) + GPy.kern.White(
-            input_dim=len(limits)
-        )
-        self.local_kernel = GPy.kern.Matern52(
-            1, variance=1.0, lengthscale=1.0, ARD=True
-        ) + GPy.kern.White(1, variance=1e-5)
-
-        self.mean_function = None
-
-        # Mean loss
-        self.global_gpr: GPy.models.GPRegression = None
-        # Newly created for every run
-        self.local_gpr: GPy.models.GPRegression = None
-
-        self.first_run: bool = True
-        self.max_idx: int = 0
-
-        self.allowed_loss_margin: float = 0.8
-
-        # Set seed for reproducibility.
-        rank = MPI.COMM_WORLD.Get_rank()
-        np.random.seed(42 * rank)
-        random.seed(42 * rank)
-
-    def start_run(self, ind: Individual) -> None:
-        """
-        Encode the config. given as individual and create the local GP's mean function from the global GP's prediction.
-
-        Parameters
-        ----------
-        ind : propulate.Individual
-            The individual containing the current configuration.
-        """
-        self.current_encoding = self.encode_configuration(ind)
-
-        if self.first_run:
-            return
-
-        # Use global model's prediction as the mean function for the local model.
-        self.mean_function = GPy.mappings.Constant(input_dim=1, output_dim=1)
-        mean, variance = self.global_gpr.predict(self.current_encoding)
-
-        self.mean_function.C = mean[0]
-
-        # Reset the local model and run data.
-        self.local_gpr = None
-        self.current_run_data = np.array([[]])
-
-    def update(self, loss: float) -> None:
-        """
-        Update the model with the current run's final loss and retrain the global Gaussian Process with the new data.
-
-        Parameters
-        ----------
-        loss : float
-            The final loss of the current run.
-        """
-        # Append the final loss to the history.
-        # Only if this is the first run:
-        if self.first_run:
-            if self.history_X.size == 0:
-                self.history_X = self.current_encoding
-            else:
-                self.history_X = np.vstack([self.history_X, self.current_encoding])
-            if self.history_Y.size == 0:
-                self.history_Y = np.array([[loss]])
-            else:
-                self.history_Y = np.vstack([self.history_Y, [loss]])
-
-            # Create global model and fit with the new data.
-            self.global_gpr = GPy.models.GPRegression(
-                self.history_X, self.history_Y, self.global_kernel
-            )
-            self.global_gpr.optimize()
-
-            self.first_run = False
-            return
-
-        # Or append if the run is finished.
-        if self.max_idx == len(self.current_encoding):
-            if self.history_X.size == 0:
-                self.history_X = self.current_encoding
-            else:
-                self.history_X = np.vstack([self.history_X, self.current_encoding])
-            if self.history_Y.size == 0:
-                self.history_Y = np.array([[loss]])
-            else:
-                self.history_Y = np.vstack([self.history_Y, [loss]])
-
-            # Fit global model with the new data.
-            self.global_gpr.set_XY(self.history_X, self.history_Y)
-            self.global_gpr.optimize()
-            return
-
-    def cancel(self, loss: float) -> bool:
-        """
-        Cancel the run if the local GP predicts a final loss higher than the best known loss with an allowed margin.
-
-        Parameters
-        ----------
-        loss : float
-            The next interim loss of the current run.
-
-        Returns
-        -------
-        bool
-            True if the run should be cancelled, False otherwise.
-        """
-        # The first run is never cancelled and has to run till the end.
-        if self.first_run:
-            # Here we also count how often cancel is called for later reference.
-            self.max_idx += 1
-            return False
-
-        # Append interim loss to the current run data.
-        if self.current_run_data.size == 0:
-            self.current_run_data = np.array([[loss]])
-        else:
-            self.current_run_data = np.vstack([self.current_run_data, [loss]])
-
-        if self.local_gpr is None:
-            self.local_gpr = GPy.models.GPRegression(
-                np.arange(1, self.current_run_data.shape[0] + 1).reshape(-1, 1),
-                self.current_run_data,
-                self.local_kernel,
-                mean_function=self.mean_function,
-            )
-        else:
-            # Retrain the local model with X consisting of synthetic indices.
-            self.local_gpr.set_XY(
-                np.arange(1, self.current_run_data.shape[0] + 1).reshape(-1, 1),
-                self.current_run_data,
-            )
-
-        self.local_gpr.optimize()
-
-        # Predict the final loss.
-        final_loss, final_variance = self.local_gpr.predict(np.array([[self.max_idx]]))
-        # Check if the final loss is lower than the best loss with margin so far.
-        if ((final_loss - final_variance) * self.allowed_loss_margin) > max(
-            self.history_Y
-        )[0]:
-            return True
-
-        return False
-
-    def merge(self, data: Tuple[np.ndarray, float]) -> None:
-        """
-        Merge a configuration and loss tuple into the history arrays and retrain the global GP with the new data.
-
-        Parameters
-        ----------
-        data : Tuple[np.ndarray, float]
-            The configurations and final losses of the merged run.
-        """
-        if len(data) != 2:
-            raise ValueError("Data must be a tuple of (configuration, loss)")
-        if data[1] == 0:
-            # Ignore 0 loss as this means some kind of unexpected behavior.
-            return
-
-        # Append new data loss to the existing.
-        if self.history_X.size == 0:
-            self.history_X = np.array([data[0]])
-        else:
-            self.history_X = np.vstack([self.history_X, [data[0]]])
-        if self.history_Y.size == 0:
-            self.history_Y = np.array([data[1]])
-        else:
-            self.history_Y = np.vstack([self.history_Y, [data[1]]])
-
-        # Fit global model with the new data.
-        if self.global_gpr is None:
-            self.global_gpr = GPy.models.GPRegression(
-                self.history_X, self.history_Y, self.global_kernel
-            )
-        else:
-            self.global_gpr.set_XY(self.history_X, self.history_Y)
-
-        self.global_gpr.optimize()
-
-    def data(self) -> Tuple[np.ndarray, float]:
-        """
-        Return the latest configuration and loss from the history arrays.
-
-        Returns
-        -------
-        Tuple[np.ndarray, float]
-            The latest configuration and loss.
-        """
-        # Return empty array if no data is available.
-        if self.history_X.size == 0 or self.history_Y.size == 0:
-            return (np.array([]), 0)
-
-        # Return the latest loss value as all other values are already shared.
-        return (self.history_X[-1], self.history_Y[-1])
-
-    def _create_encoding(
-        self,
-        limits: Union[
-            Dict[str, Tuple[float, float]],
-            Dict[str, Tuple[int, int]],
-            Dict[str, Tuple[str, ...]],
-        ],
-    ) -> Dict[str, Dict[str, int]]:
-        """
-        Create a mapping from categorical values to integers. Gaussian Process Regression only accepts numerical values.
-
-        Parameters
-        ----------
-        limits : Union[Dict[str, Tuple[float, float]], Dict[str, Tuple[int, int]], Dict[str, Tuple[str, ...]]
-            The hyperparameter configuration space's limits.
-
-        Returns
-        -------
-        Dict[str, Dict[str, int]]
-            The input dict but with the categorical values transformed to integers.
-        """
-        encodings = {}
-        for key, values in limits.items():
-            if isinstance(values, tuple) and isinstance(values[0], str):
-                # create a mapping for categorical values
-                encodings[key] = {v: i for i, v in enumerate(values)}
-        return encodings
-
-    def encode_configuration(
-        self,
-        config: Union[
-            Dict[str, float],
-            Dict[str, int],
-            Dict[str, str],
-        ],
-    ) -> np.ndarray:
-        """
-        Encode a configuration dictionary into a 2D array (required by the Gaussian Process Regression model).
-
-        Parameters
-        ----------
-        config : Union[Dict[str, float], Dict[str, int], Dict[str, str]]
-            The configuration dictionary to encode.
-
-        Returns
-        -------
-        np.ndarray
-            The encoded configuration.
-        """
-        # Initialize an empty array to ``limits`` size.
-        encoded_array = np.zeros((1, len(self.limits)), dtype=float)
-
-        # Fill the array with encoded values.
-        for i, key in enumerate(self.limits):
-            value = config[key]
-            if key in self.encodings:
-                # Ignore type hint as value can only be a string.
-                encoded_array[0, i] = self.encodings[key][value]
-            else:
-                encoded_array[0, i] = value
-
-        return encoded_array
diff --git a/tutorials/surrogate/mnist_example.py b/tutorials/surrogate/mnist_example.py
index a444aa65..a33e1d7f 100755
--- a/tutorials/surrogate/mnist_example.py
+++ b/tutorials/surrogate/mnist_example.py
@@ -12,8 +12,7 @@
 from torchvision.datasets import MNIST
 from torchvision.transforms import Compose, Normalize, ToTensor
 
-import tutorials.surrogate
-from propulate import Islands
+from propulate import Islands, surrogate
 from propulate.utils import get_default_propagator
 
 GPUS_PER_NODE: int = 1
@@ -173,6 +172,7 @@ def get_data_loaders(batch_size: int) -> Tuple[DataLoader, DataLoader]:
             shuffle=True,  # Shuffle data.
         )
 
+    # NOTE barrier only called, when dataset has not been downloaded yet
     if not hasattr(get_data_loaders, "barrier_called"):
         MPI.COMM_WORLD.Barrier()
 
@@ -336,10 +336,8 @@ def set_seeds(seed_value: int = 42) -> None:
         generations=num_generations,  # Number of generations per worker
         num_islands=1,  # Number of islands
         checkpoint_path=log_path,
-        # surrogate_factory=lambda: tutorials.surrogate.MockSurrogate(),
-        # surrogate_factory=lambda: tutorials.surrogate.LogSurrogate(),
-        surrogate_factory=lambda: tutorials.surrogate.StaticSurrogate(),
-        # surrogate_factory=lambda: tutorials.surrogate.DynamicSurrogate(limits),
+        surrogate_factory=lambda: surrogate.StaticSurrogate(),
+        # surrogate_factory=lambda: surrogate.DynamicSurrogate(limits),
     )
     islands.evolve(  # Run evolutionary optimization.
         top_n=1,  # Print top-n best individuals on each island in summary.
diff --git a/tutorials/surrogate/requirements.txt b/tutorials/surrogate/requirements.txt
deleted file mode 100644
index 2863da3e..00000000
--- a/tutorials/surrogate/requirements.txt
+++ /dev/null
@@ -1,5 +0,0 @@
-torch
-torchvision
-torchmetrics
-GPy
-mpi4py==3.0.3
diff --git a/tutorials/surrogate/static_surrogate.py b/tutorials/surrogate/static_surrogate.py
deleted file mode 100644
index a429f6e1..00000000
--- a/tutorials/surrogate/static_surrogate.py
+++ /dev/null
@@ -1,172 +0,0 @@
-import numpy as np
-
-from propulate import Surrogate
-from propulate.population import Individual
-
-
-class StaticSurrogate(Surrogate):
-    """
-    Surrogate model using the best known run as baseline.
-
-    After the first run, each subsequent loss is compared to the baseline. Every run with a loss outside the margin of
-    the baseline is cancelled.
-
-    This model creates an internal index for the yielded losses during each run. That means the yield order and
-    frequency has to be consistent between runs. Otherwise, the indices of the baseline run won't match.
-
-    Loosely based on the paper:
-    Use of Static Surrogates in Hyperparameter Optimization
-    https://link.springer.com/article/10.1007/s43069-022-00128-w
-
-    Attributes
-    ----------
-    synthetic_id : int
-        An internal running index to keep track of the current runs loss index.
-    margin : float
-        A margin to be multiplied with the new incoming losses to compare them to the baseline.
-    first_run : bool
-        Flag for the first complete run.
-    baseline : np.ndarray
-        Chronological list of the losses of the best known run.
-    current_run : np.ndarray
-        Array to store the losses of the current run.
-
-    Methods
-    -------
-    __init__()
-        Overwrites the parents class's constructor to include the error margin.
-    start_run()
-        Reset the internal index and the current run array.
-    update()
-        Replace the baseline with the current run if the final loss is better.
-    cancel()
-        Cancel the current run if the loss is outside the margin of the baseline.
-    merge()
-        Merge an incoming baseline.
-    data()
-        Return the loss series of the best known run so far.
-
-    Notes
-    -----
-    The ``StaticSurrogate`` class implements all methods from the ``Surrogate`` class.
-
-    See Also
-    --------
-    :class:`Surrogate` : The parent class.
-    """
-
-    def __init__(self, margin: float = 0.8) -> None:
-        """
-        Initialize a static surrogate with a synthetic id, a margin, and empty arrays for baseline and current run.
-
-        Parameters
-        ----------
-        margin : float, optional
-            A margin on top of incoming losses for comparison with the baseline.
-            The default is 0.8.
-        """
-        super().__init__()
-
-        self.synthetic_id: int = 0
-        self.margin: float = margin
-
-        # Canceling with ``cancel`` is only allowed after the first complete run.
-        # ``first_run`` keeps track of that.
-        self.first_run: bool = True
-
-        # Baseline is the best known run.
-        self.baseline: np.ndarray = np.zeros((0), dtype=float)
-        self.current_run: np.ndarray = np.zeros((0), dtype=float)
-
-    def start_run(self, ind: Individual) -> None:
-        """
-        Reset the internal index and the current run array.
-
-        Parameters
-        ----------
-        ind : Individual
-            The individual containing the current configuration.
-        """
-        # Reset to new run.
-        self.synthetic_id = 0
-        # Reset current run with correct size.
-        self.current_run = np.zeros((self.baseline.size), dtype=float)
-
-    def update(self, loss: float) -> None:
-        """
-        Replace the baseline with the current run if the final loss is better, or if there is no prior run.
-
-        Parameters
-        ----------
-        loss : float
-            The (unused) final loss of the current run.
-        """
-        if self.first_run:
-            self.baseline = self.current_run.copy()
-            self.first_run = False
-            return
-
-        # Check if current run is better than baseline.
-        if self.baseline[-1] > self.current_run[-1]:
-            self.baseline = self.current_run.copy()
-
-    def cancel(self, loss: float) -> bool:
-        """
-        Cancel the current run if the loss is outside the margin of the baseline.
-
-        Parameters
-        ----------
-        loss : float
-            The next interim loss of the current run.
-
-        Returns
-        -------
-        bool
-            True if the current run is cancelled, False otherwise.
-        """
-        self.synthetic_id += 1
-
-        # Cancel is only allowed after the first complete run.
-        if self.first_run:
-            self.current_run = np.append(self.current_run, loss)
-            return False
-
-        # Append loss to current run.
-        self.current_run[self.synthetic_id - 1] = loss
-
-        # Cancel if current run is outside margin of baseline.
-        if self.baseline[self.synthetic_id - 1] < loss * self.margin:
-            return True
-
-        return False
-
-    def merge(self, data: np.ndarray) -> None:
-        """
-        Replace the baseline with the incoming run if the final loss is better, or if there is no prior run.
-
-        Parameters
-        ----------
-        data : np.ndarray
-            The loss series of the incoming run.
-        """
-        # No prior data to merge with
-        if self.first_run:
-            self.baseline = data.copy()
-            self.first_run = False
-            return
-
-        # Merged run's final loss is better than baseline.
-        if self.baseline[-1] < data[-1]:
-            self.baseline = data.copy()
-
-    def data(self) -> np.ndarray:
-        """
-        Return the loss series of the best known run so far.
-
-        Returns
-        -------
-        np.ndarray
-            The loss series of the best known run so far.
-        """
-        # Return best run so far.
-        return self.baseline
diff --git a/tutorials/torch_example.py b/tutorials/torch_example.py
index 850e9317..49a27b60 100755
--- a/tutorials/torch_example.py
+++ b/tutorials/torch_example.py
@@ -233,7 +233,12 @@ def get_data_loaders(batch_size: int) -> Tuple[DataLoader, DataLoader]:
             shuffle=True,  # Shuffle data.
         )
 
-    MPI.COMM_WORLD.barrier()
+    # NOTE barrier only called, when dataset has not been downloaded yet
+    if not hasattr(get_data_loaders, "barrier_called"):
+        MPI.COMM_WORLD.Barrier()
+
+        setattr(get_data_loaders, "barrier_called", True)
+
     if MPI.COMM_WORLD.rank != 0:
         train_loader = DataLoader(
             dataset=MNIST(


Test patch:
diff --git a/.github/workflows/python-test.yml b/.github/workflows/python-test.yml
index 7e37d3d1..27823142 100644
--- a/.github/workflows/python-test.yml
+++ b/.github/workflows/python-test.yml
@@ -6,33 +6,33 @@ jobs:
     build:
         runs-on: ubuntu-latest
         steps:
-            - uses: actions/checkout@v3
+            - uses: actions/checkout@v4
             - name: Set up MPI
               uses: mpi4py/setup-mpi@v1
               with:
                 mpi: 'openmpi'
             - name: Set up Python
-              uses: actions/setup-python@v4
+              uses: actions/setup-python@v5
               with:
                   python-version: '3.9'
             - name: Install dependencies
               run: |
                   python -m pip install --upgrade pip
-                  pip install ruff pytest pytest-cov
-                  pip install coverage-badge
-                  pip install -r requirements.txt
-                  pip install .
+                  pip install ."[dev]"
             - name: Lint with ruff
               run: |
-                  # stop the build if there are Python syntax errors or undefined names
-                  ruff --output-format=github --select=E9,F63,F7,F82 --target-version=py39 .
-                  # default set of ruff rules with GitHub Annotations
-                  ruff --output-format=github --target-version=py39 .
+                  # Stop the build if there are Python syntax errors or undefined names.
+                  ruff check --output-format=github --select=E9,F63,F7,F82 --target-version=py39 .
+                  # Default set of ruff rules with GitHub Annotations
+                  ruff check --output-format=github --target-version=py39 .
             - name: Test with pytest
               run: |
-                pytest --cov=propulate
-                coverage-badge -fo coverage.svg
-
+                coverage run --rcfile=./pyproject.toml -m pytest
+                mpirun -n 8 coverage run --rcfile=./pyproject.toml -m mpi4py -m pytest --with-mpi
+                coverage combine
+                coverage report -m
+            - name: Coverage Badge
+              uses: tj-actions/coverage-badge-py@v2
             - name: Verify Changed files
               uses: tj-actions/verify-changed-files@v16
               id: verify-changed-files
diff --git a/pytest.ini b/pytest.ini
deleted file mode 100644
index 5ee64771..00000000
--- a/pytest.ini
+++ /dev/null
@@ -1,2 +0,0 @@
-[pytest]
-testpaths = tests
diff --git a/tests/test_cmaes.py b/tests/test_cmaes.py
index 34638d39..d43d29b4 100644
--- a/tests/test_cmaes.py
+++ b/tests/test_cmaes.py
@@ -1,54 +1,65 @@
+import pathlib
 import random
-import tempfile
-from typing import Dict
+from typing import Tuple
 
-import numpy as np
+import pytest
 
 from propulate import Propulator
 from propulate.propagators import BasicCMA, CMAPropagator
+from propulate.utils.benchmark_functions import get_function_search_space
 
 
-def sphere(params: Dict[str, float]) -> float:
+@pytest.fixture(
+    params=[
+        ("rosenbrock", 0.0),
+        ("step", -25.0),
+        ("quartic", 0.0),
+        ("rastrigin", 0.0),
+        ("griewank", 0.0),
+        ("schwefel", 0.0),
+        ("bisphere", 0.0),
+        ("birastrigin", 0.0),
+        ("bukin", 0.0),
+        ("eggcrate", -1.0),
+        ("himmelblau", 0.0),
+        ("keane", 0.6736675),
+        ("leon", 0.0),
+        ("sphere", 0.0),  # (fname, expected)
+    ]
+)
+def function_parameters(request):
+    """Define benchmark function parameter sets as used in tests."""
+    return request.param
+
+
+def test_cmaes(
+    function_parameters: Tuple[str, float], mpi_tmp_path: pathlib.Path
+) -> None:
     """
-    Sphere function: continuous, convex, separable, differentiable, unimodal.
+    Test Propulator to optimize a benchmark function using a CMA-ES propagator.
 
-    Input domain: -5.12 <= x, y <= 5.12
-    Global minimum 0 at (x, y) = (0, 0)
+    This test is run both sequentially and in parallel.
 
     Parameters
     ----------
-    params : Dict[str, float]
-        The function parameters.
-
-    Returns
-    -------
-    float
-        The function value.
+    function_parameters : Tuple[str, float]
+        The tuple containing each function name along with its global minimum.
+    mpi_tmp_path : pathlib.Path
+        The temporary checkpoint directory.
     """
-    return np.sum(np.array(list(params.values())) ** 2).item()
-
-
-def test_cmaes():
-    """Test single worker using Propulator to optimize sphere using a CMA-ES propagator."""
     rng = random.Random(42)  # Separate random number generator for optimization.
-    limits = {
-        "a": (-5.12, 5.12),
-        "b": (-5.12, 5.12),
-    }
-    with tempfile.TemporaryDirectory() as checkpoint_path:
-        # Set up evolutionary operator.
-        adapter = BasicCMA()
-        propagator = CMAPropagator(adapter, limits, rng=rng)
-
-        # Set up Propulator performing actual optimization.
-        propulator = Propulator(
-            loss_fn=sphere,
-            propagator=propagator,
-            rng=rng,
-            generations=100,
-            checkpoint_path=checkpoint_path,
-        )
-        # Run optimization and print summary of results.
-        propulator.propulate()
-        best = propulator.summarize(top_n=1, debug=2)
-        assert best[0][0].loss < 10**-1
+    function, limits = get_function_search_space(function_parameters[0])
+    # Set up evolutionary operator.
+    adapter = BasicCMA()
+    propagator = CMAPropagator(adapter, limits, rng=rng)
+
+    # Set up Propulator performing actual optimization.
+    propulator = Propulator(
+        loss_fn=function,
+        propagator=propagator,
+        rng=rng,
+        generations=100,
+        checkpoint_path=mpi_tmp_path,
+    )
+    # Run optimization and print summary of results.
+    propulator.propulate()
diff --git a/tests/test_island.py b/tests/test_island.py
new file mode 100644
index 00000000..a89fab65
--- /dev/null
+++ b/tests/test_island.py
@@ -0,0 +1,363 @@
+import copy
+import pathlib
+import random
+from typing import Tuple
+
+import deepdiff
+import numpy as np
+import pytest
+from mpi4py import MPI
+
+from propulate import Islands
+from propulate.utils import get_default_propagator, set_logger_config
+from propulate.utils.benchmark_functions import get_function_search_space
+
+
+@pytest.fixture(
+    params=[
+        ("rosenbrock", 0.0),
+        ("step", -25.0),
+        ("quartic", 0.0),
+        ("rastrigin", 0.0),
+        ("griewank", 0.0),
+        ("schwefel", 0.0),
+        ("bisphere", 0.0),
+        ("birastrigin", 0.0),
+        ("bukin", 0.0),
+        ("eggcrate", -1.0),
+        ("himmelblau", 0.0),
+        ("keane", 0.6736675),
+        ("leon", 0.0),
+        ("sphere", 0.0),  # (fname, expected)
+    ]
+)
+def function_parameters(request):
+    """Define benchmark function parameter sets as used in tests."""
+    return request.param
+
+
+@pytest.mark.mpi(min_size=4)
+def test_islands(
+    function_parameters: Tuple[str, float], mpi_tmp_path: pathlib.Path
+) -> None:
+    """
+    Test basic island functionality (only run in parallel with at least four processes).
+
+    Parameters
+    ----------
+    function_parameters : Tuple[str, float]
+        The tuple containing each function name along with its global minimum.
+    mpi_tmp_path : pathlib.Path
+        The temporary checkpoint directory.
+    """
+    rng = random.Random(
+        42 + MPI.COMM_WORLD.rank
+    )  # Separate random number generator for optimization
+    function, limits = get_function_search_space(function_parameters[0])
+    set_logger_config(log_file=mpi_tmp_path / "log.log")
+
+    # Set up evolutionary operator.
+    propagator = get_default_propagator(
+        pop_size=4,
+        limits=limits,
+        rng=rng,
+    )
+
+    # Set up island model.
+    islands = Islands(
+        loss_fn=function,
+        propagator=propagator,
+        rng=rng,
+        generations=100,
+        num_islands=2,
+        migration_probability=0.9,
+        pollination=False,
+        checkpoint_path=mpi_tmp_path,
+    )
+
+    # Run actual optimization.
+    islands.evolve(
+        top_n=1,
+        logging_interval=10,
+        debug=2,
+    )
+
+
+@pytest.mark.mpi(min_size=4)
+def test_checkpointing_isolated(
+    function_parameters: Tuple[str, float], mpi_tmp_path: pathlib.Path
+) -> None:
+    """
+    Test isolated island checkpointing without migration (only run in parallel with at least four processes).
+
+    Parameters
+    ----------
+    function_parameters : Tuple[str, float]
+        The tuple containing each function name along with its global minimum.
+    mpi_tmp_path : pathlib.Path
+        The temporary checkpoint directory.
+    """
+    rng = random.Random(
+        42 + MPI.COMM_WORLD.rank
+    )  # Separate random number generator for optimization
+    function, limits = get_function_search_space(function_parameters[0])
+    set_logger_config(log_file=mpi_tmp_path / "log.log")
+
+    # Set up evolutionary operator.
+    propagator = get_default_propagator(
+        pop_size=4,
+        limits=limits,
+        rng=rng,
+    )
+
+    # Set up island model.
+    islands = Islands(
+        loss_fn=function,
+        propagator=propagator,
+        rng=rng,
+        generations=100,
+        num_islands=2,
+        migration_probability=0.0,
+        checkpoint_path=mpi_tmp_path,
+    )
+
+    # Run actual optimization.
+    islands.evolve(
+        top_n=1,
+        logging_interval=10,
+        debug=2,
+    )
+
+    old_population = copy.deepcopy(islands.propulator.population)
+    del islands
+
+    islands = Islands(
+        loss_fn=function,
+        propagator=propagator,
+        rng=rng,
+        generations=100,
+        num_islands=2,
+        migration_probability=0.0,
+        checkpoint_path=mpi_tmp_path,
+    )
+
+    assert (
+        len(
+            deepdiff.DeepDiff(
+                old_population, islands.propulator.population, ignore_order=True
+            )
+        )
+        == 0
+    )
+
+
+@pytest.mark.mpi(min_size=4)
+def test_checkpointing_migration(
+    function_parameters: Tuple[str, float], mpi_tmp_path: pathlib.Path
+) -> None:
+    """
+    Test island checkpointing with migration (only run in parallel with at least four processes).
+
+    Parameters
+    ----------
+    function_parameters : Tuple[str, float]
+        The tuple containing each function name along with its global minimum.
+    mpi_tmp_path : pathlib.Path
+        The temporary checkpoint directory.
+    """
+    rng = random.Random(
+        42 + MPI.COMM_WORLD.rank
+    )  # Separate random number generator for optimization
+    function, limits = get_function_search_space(function_parameters[0])
+    set_logger_config(log_file=mpi_tmp_path / "log.log")
+
+    # Set up evolutionary operator.
+    propagator = get_default_propagator(
+        pop_size=4,
+        limits=limits,
+        rng=rng,
+    )
+
+    # Set up island model.
+    islands = Islands(
+        loss_fn=function,
+        propagator=propagator,
+        rng=rng,
+        generations=100,
+        num_islands=2,
+        migration_probability=0.9,
+        pollination=False,  # TODO fixtureize
+        checkpoint_path=mpi_tmp_path,
+    )
+
+    # Run actual optimization.
+    islands.evolve(
+        top_n=1,
+        logging_interval=10,
+        debug=2,
+    )
+
+    old_population = copy.deepcopy(islands.propulator.population)
+    del islands
+
+    islands = Islands(
+        loss_fn=function,
+        propagator=propagator,
+        rng=rng,
+        generations=100,
+        num_islands=2,
+        migration_probability=0.9,
+        pollination=False,  # TODO fixtureize
+        checkpoint_path=mpi_tmp_path,
+    )
+
+    assert (
+        len(
+            deepdiff.DeepDiff(
+                old_population, islands.propulator.population, ignore_order=True
+            )
+        )
+        == 0
+    )
+
+
+@pytest.mark.mpi(min_size=4)
+def test_checkpointing_pollination(
+    function_parameters: Tuple[str, float], mpi_tmp_path: pathlib.Path
+) -> None:
+    """
+    Test island checkpointing with pollination (only run in parallel with at least four processes).
+
+    Parameters
+    ----------
+    function_parameters : Tuple[str, float]
+        The tuple containing each function name along with its global minimum.
+    mpi_tmp_path : pathlib.Path
+        The temporary checkpoint directory.
+    """
+    rng = random.Random(
+        42 + MPI.COMM_WORLD.rank
+    )  # Separate random number generator for optimization
+    function, limits = get_function_search_space(function_parameters[0])
+    set_logger_config(log_file=mpi_tmp_path / "log.log")
+
+    # Set up evolutionary operator.
+    propagator = get_default_propagator(
+        pop_size=4,
+        limits=limits,
+        rng=rng,
+    )
+
+    # Set up island model.
+    islands = Islands(
+        loss_fn=function,
+        propagator=propagator,
+        rng=rng,
+        generations=100,
+        num_islands=2,
+        migration_probability=0.9,
+        pollination=False,  # TODO fixtureize
+        checkpoint_path=mpi_tmp_path,
+    )
+
+    # Run actual optimization.
+    islands.evolve(
+        top_n=1,
+        logging_interval=10,
+        debug=2,
+    )
+
+    old_population = copy.deepcopy(islands.propulator.population)
+    del islands
+
+    islands = Islands(
+        loss_fn=function,
+        propagator=propagator,
+        rng=rng,
+        generations=100,
+        num_islands=2,
+        migration_probability=0.9,
+        pollination=True,  # TODO fixtureize
+        checkpoint_path=mpi_tmp_path,
+    )
+
+    assert (
+        len(
+            deepdiff.DeepDiff(
+                old_population, islands.propulator.population, ignore_order=True
+            )
+        )
+        == 0
+    )
+
+
+@pytest.mark.mpi(min_size=8)
+def test_checkpointing_unequal_populations(
+    function_parameters: Tuple[str, float], mpi_tmp_path: pathlib.Path
+) -> None:
+    """
+    Test island checkpointing for inhomogeneous island sizes (only run in parallel with at least eight processes).
+
+    Parameters
+    ----------
+    function_parameters : Tuple[str, float]
+        The tuple containing each function name along with its global minimum.
+    mpi_tmp_path : pathlib.Path
+        The temporary checkpoint directory.
+    """
+    rng = random.Random(
+        42 + MPI.COMM_WORLD.rank
+    )  # Separate random number generator for optimization
+    function, limits = get_function_search_space(function_parameters[0])
+    set_logger_config(log_file=mpi_tmp_path / "log.log")
+
+    # Set up evolutionary operator.
+    propagator = get_default_propagator(
+        pop_size=4,
+        limits=limits,
+        rng=rng,
+    )
+
+    # Set up island model.
+    islands = Islands(
+        loss_fn=function,
+        propagator=propagator,
+        rng=rng,
+        generations=100,
+        num_islands=2,
+        island_sizes=np.array([3, 5]),
+        migration_probability=0.9,
+        pollination=False,  # TODO fixtureize
+        checkpoint_path=mpi_tmp_path,
+    )
+
+    # Run actual optimization.
+    islands.evolve(
+        top_n=1,
+        logging_interval=10,
+        debug=2,
+    )
+
+    old_population = copy.deepcopy(islands.propulator.population)
+    del islands
+
+    islands = Islands(
+        loss_fn=function,
+        propagator=propagator,
+        rng=rng,
+        generations=100,
+        num_islands=2,
+        island_sizes=np.array([3, 5]),
+        migration_probability=0.9,
+        pollination=True,  # TODO fixtureize
+        checkpoint_path=mpi_tmp_path,
+    )
+
+    assert (
+        len(
+            deepdiff.DeepDiff(
+                old_population, islands.propulator.population, ignore_order=True
+            )
+        )
+        == 0
+    )
diff --git a/tests/test_multi_rank_workers.py b/tests/test_multi_rank_workers.py
new file mode 100644
index 00000000..83ee6f1f
--- /dev/null
+++ b/tests/test_multi_rank_workers.py
@@ -0,0 +1,86 @@
+import pathlib
+import random
+from typing import Dict
+
+import numpy as np
+import pytest
+from mpi4py import MPI
+
+from propulate import Islands
+from propulate.utils import get_default_propagator, set_logger_config
+
+
+def parallel_sphere(params: Dict[str, float], comm: MPI.Comm = MPI.COMM_SELF) -> float:
+    """
+    Parallel sphere function to showcase using multi-rank workers in Propulate.
+
+    Sphere function: continuous, convex, separable, differentiable, unimodal
+
+    Input domain: -5.12 <= x, y <= 5.12
+    Global minimum 0 at (x, y) = (0, 0)
+
+    Parameters
+    ----------
+    params : Dict[str, float]
+        The function parameters.
+    comm : MPI.Comm
+        The communicator of the worker.
+
+    Returns
+    -------
+    float
+        The function value.
+    """
+    if comm != MPI.COMM_SELF:
+        term = (
+            list(params.values())[comm.rank] ** 2
+        )  # Each rank squares one of the inputs.
+        return comm.allreduce(term)  # Return the sum over all squared inputs.
+    else:
+        return np.sum(np.array(list(params.values())) ** 2).item()
+
+
+@pytest.mark.mpi(min_size=8)
+def test_multi_rank_workers(mpi_tmp_path: pathlib.Path) -> None:
+    """
+    Test multi-rank workers. Two islands with at least two workers with two ranks each.
+
+    Parameters
+    ----------
+    mpi_tmp_path : pathlib.Path
+        The temporary checkpoint directory.
+    """
+    full_world_comm = MPI.COMM_WORLD  # Get full world communicator.
+    set_logger_config(log_file=mpi_tmp_path / "log.log")
+
+    rng = random.Random(42 + full_world_comm.rank)
+    limits = {
+        "a": (-5.12, 5.12),
+        "b": (-5.12, 5.12),
+    }
+    propagator = get_default_propagator(  # Get default evolutionary operator.
+        pop_size=2,  # Breeding pool size
+        limits=limits,  # Search-space limits
+        rng=rng,  # Separate random number generator for Propulate optimization
+    )
+
+    # Set up island model.
+    islands = Islands(
+        loss_fn=parallel_sphere,  # Loss function to be minimized
+        propagator=propagator,  # Propagator, i.e., evolutionary operator to be used
+        rng=rng,  # Separate random number generator for Propulate optimization
+        generations=100,  # Overall number of generations
+        num_islands=2,  # Number of islands
+        migration_probability=0.9,  # Migration probability
+        pollination=False,  # Whether to use pollination or migration
+        checkpoint_path=mpi_tmp_path,  # Checkpoint path
+        # ----- SPECIFIC FOR MULTI-RANK UCS ----
+        ranks_per_worker=2,  # Number of ranks per (multi rank) worker
+    )
+
+    # Run actual optimization.
+    islands.evolve(
+        top_n=1,  # Print top-n best individuals on each island in summary.
+        logging_interval=10,  # Logging interval
+        debug=1,  # Debug level
+    )
diff --git a/tests/test_propulator.py b/tests/test_propulator.py
new file mode 100644
index 00000000..caff7e50
--- /dev/null
+++ b/tests/test_propulator.py
@@ -0,0 +1,123 @@
+import copy
+import pathlib
+import random
+from typing import Tuple
+
+import deepdiff
+import pytest
+from mpi4py import MPI
+
+from propulate import Propulator
+from propulate.utils import get_default_propagator, set_logger_config
+from propulate.utils.benchmark_functions import get_function_search_space
+
+
+@pytest.fixture(
+    params=[
+        ("rosenbrock", 0.0),
+        ("step", -25.0),
+        ("quartic", 0.0),
+        ("rastrigin", 0.0),
+        ("griewank", 0.0),
+        ("schwefel", 0.0),
+        ("bisphere", 0.0),
+        ("birastrigin", 0.0),
+        ("bukin", 0.0),
+        ("eggcrate", -1.0),
+        ("himmelblau", 0.0),
+        ("keane", 0.6736675),
+        ("leon", 0.0),
+        ("sphere", 0.0),  # (fname, expected)
+    ]
+)
+def function_parameters(request):
+    """Define benchmark function parameter sets as used in tests."""
+    return request.param
+
+
+def test_propulator(
+    function_parameters: Tuple[str, float], mpi_tmp_path: pathlib.Path
+) -> None:
+    """
+    Test standard Propulator to optimize the benchmark functions using the default genetic propagator.
+
+    This test is run both sequentially and in parallel.
+
+    Parameters
+    ----------
+    function_parameters : Tuple[str, float]
+        The tuple containing each function name along with its global minimum.
+    mpi_tmp_path : pathlib.Path
+        The temporary checkpoint directory.
+    """
+    rng = random.Random(
+        42 + MPI.COMM_WORLD.rank
+    )  # Random number generator for optimization
+    function, limits = get_function_search_space(function_parameters[0])
+    set_logger_config(log_file=mpi_tmp_path / "log.log")
+    propagator = get_default_propagator(
+        pop_size=4,
+        limits=limits,
+        rng=rng,
+    )  # Set up evolutionary operator.
+    propulator = Propulator(
+        loss_fn=function,
+        propagator=propagator,
+        rng=rng,
+        generations=100,
+        checkpoint_path=mpi_tmp_path,
+    )  # Set up propulator performing actual optimization.
+    propulator.propulate()  # Run optimization and print summary of results.
+
+
+def test_propulator_checkpointing(mpi_tmp_path: pathlib.Path) -> None:
+    """
+    Test standard Propulator checkpointing for the sphere benchmark function.
+
+    This test is run both sequentially and in parallel.
+
+    Parameters
+    ----------
+    mpi_tmp_path : pathlib.Path
+        The temporary checkpoint directory.
+    """
+    rng = random.Random(
+        42 + MPI.COMM_WORLD.rank
+    )  # Separate random number generator for optimization
+    function, limits = get_function_search_space("sphere")
+
+    propagator = get_default_propagator(  # Get default evolutionary operator.
+        pop_size=4,  # Breeding pool size
+        limits=limits,  # Search-space limits
+        rng=rng,  # Random number generator
+    )
+    propulator = Propulator(
+        loss_fn=function,
+        propagator=propagator,
+        generations=1000,
+        checkpoint_path=mpi_tmp_path,
+        rng=rng,
+    )  # Set up propulator performing actual optimization.
+
+    propulator.propulate()  # Run optimization and print summary of results.
+
+    old_population = copy.deepcopy(
+        propulator.population
+    )  # Save population list from the last run.
+    del propulator  # Delete propulator object.
+    MPI.COMM_WORLD.barrier()  # Synchronize all processes.
+
+    propulator = Propulator(
+        loss_fn=function,
+        propagator=propagator,
+        generations=20,
+        checkpoint_path=mpi_tmp_path,
+        rng=rng,
+    )  # Set up new propulator starting from checkpoint.
+
+    # As the number of requested generations is smaller than the number of generations from the run before,
+    # no new evaluations are performed. Thus, the length of both Propulators' populations must be equal.
+    assert (
+        len(deepdiff.DeepDiff(old_population, propulator.population, ignore_order=True))
+        == 0
+    )
diff --git a/tests/test_propulator_sphere.py b/tests/test_propulator_sphere.py
deleted file mode 100644
index d7b8e97a..00000000
--- a/tests/test_propulator_sphere.py
+++ /dev/null
@@ -1,70 +0,0 @@
-import logging
-import random
-import tempfile
-from typing import Dict
-
-import numpy as np
-
-from propulate import Propulator
-from propulate.utils import get_default_propagator, set_logger_config
-
-
-def sphere(params: Dict[str, float]) -> float:
-    """
-    Sphere function: continuous, convex, separable, differentiable, unimodal.
-
-    Input domain: -5.12 <= x, y <= 5.12
-    Global minimum 0 at (x, y) = (0, 0)
-
-    Parameters
-    ----------
-    params : Dict[str, float]
-        The function parameters.
-
-    Returns
-    -------
-    float
-        The function value.
-    """
-    return np.sum(np.array(list(params.values())) ** 2).item()
-
-
-def test_propulator():
-    """Test single worker using Propulator to optimize sphere."""
-    rng = random.Random(42)  # Separate random number generator for optimization.
-    limits = {
-        "a": (-5.12, 5.12),
-        "b": (-5.12, 5.12),
-    }
-    with tempfile.TemporaryDirectory() as checkpoint_path:
-        set_logger_config(
-            level=logging.INFO,
-            log_file=checkpoint_path + "/propulate.log",
-            log_to_stdout=True,
-            log_rank=False,
-            colors=True,
-        )
-        # Set up evolutionary operator.
-        propagator = get_default_propagator(
-            pop_size=4,
-            limits=limits,
-            crossover_prob=0.7,
-            mutation_prob=9.0,
-            random_init_prob=0.1,
-            rng=rng,
-        )
-
-        # Set up propulator performing actual optimization.
-        propulator = Propulator(
-            loss_fn=sphere,
-            propagator=propagator,
-            rng=rng,
-            generations=100,
-            checkpoint_path=checkpoint_path,
-        )
-
-        # Run optimization and print summary of results.
-        propulator.propulate()
-        best = propulator.summarize()
-
-        assert best[0][0].loss < 0.8
diff --git a/tests/test_pso.py b/tests/test_pso.py
index a19da15d..e3049a76 100644
--- a/tests/test_pso.py
+++ b/tests/test_pso.py
@@ -1,65 +1,73 @@
+import pathlib
 import random
-import tempfile
-from typing import Dict
+from typing import Tuple
 
-import numpy as np
+import pytest
 from mpi4py import MPI
 
 from propulate import Propulator
 from propulate.propagators import Conditional
 from propulate.propagators.pso import BasicPSO, InitUniformPSO
+from propulate.utils.benchmark_functions import get_function_search_space
 
 
-def sphere(params: Dict[str, float]) -> float:
-    """
-    Sphere function: continuous, convex, separable, differentiable, unimodal.
+@pytest.fixture(
+    params=[
+        ("rosenbrock", 0.0),
+        ("step", -25.0),
+        ("quartic", 0.0),
+        ("rastrigin", 0.0),
+        ("griewank", 0.0),
+        ("schwefel", 0.0),
+        ("bisphere", 0.0),
+        ("birastrigin", 0.0),
+        ("bukin", 0.0),
+        ("eggcrate", -1.0),
+        ("himmelblau", 0.0),
+        ("keane", 0.6736675),
+        ("leon", 0.0),
+        ("sphere", 0.0),  # (fname, expected)
+    ]
+)
+def function_parameters(request):
+    """Define benchmark function parameter sets as used in tests."""
+    return request.param
+
 
-    Input domain: -5.12 <= x, y <= 5.12
-    Global minimum 0 at (x, y) = (0, 0)
+@pytest.mark.mpi
+def test_pso(function_parameters: Tuple[str, float, float], mpi_tmp_path: pathlib.Path):
+    """
+    Test single worker using Propulator to optimize a benchmark function using the default genetic propagator.
 
     Parameters
     ----------
-    params : Dict[str, float]
-        The function parameters.
-
-    Returns
-    -------
-    float
-        The function value.
+    function_parameters : Tuple
+        The tuple containing each function name along with its global minimum.
+    mpi_tmp_path : pathlib.Path
+        The temporary checkpoint directory.
     """
-    return np.sum(np.array(list(params.values())) ** 2).item()
-
-
-def test_pso():
-    """Test single worker using Propulator to optimize sphere using a PSO propagator."""
     rng = random.Random(42)  # Separate random number generator for optimization.
-    limits = {
-        "a": (-5.12, 5.12),
-        "b": (-5.12, 5.12),
-    }
-    with tempfile.TemporaryDirectory() as checkpoint_path:
-        # Set up evolutionary operator.
-        pso_propagator = BasicPSO(
-            inertia=0.729,
-            c_cognitive=1.49334,
-            c_social=1.49445,
-            rank=MPI.COMM_WORLD.rank,  # MPI rank TODO fix when implemented proper MPI parallel tests
-            limits=limits,
-            rng=rng,
-        )
-        init = InitUniformPSO(limits, rng=rng, rank=MPI.COMM_WORLD.rank)
-        propagator = Conditional(1, pso_propagator, init)  # TODO MPIify
+    function, limits = get_function_search_space(function_parameters[0])
+    # Set up evolutionary operator.
+    pso_propagator = BasicPSO(
+        inertia=0.729,
+        c_cognitive=1.49334,
+        c_social=1.49445,
+        rank=MPI.COMM_WORLD.rank,  # MPI rank
+        limits=limits,
+        rng=rng,
+    )
+    init = InitUniformPSO(limits, rng=rng, rank=MPI.COMM_WORLD.rank)
+    propagator = Conditional(1, pso_propagator, init)
 
-        # Set up propulator performing actual optimization.
-        propulator = Propulator(
-            loss_fn=sphere,
-            propagator=propagator,
-            rng=rng,
-            generations=100,
-            checkpoint_path=checkpoint_path,
-        )
+    # Set up propulator performing actual optimization.
+    propulator = Propulator(
+        loss_fn=function,
+        propagator=propagator,
+        rng=rng,
+        generations=100,
+        checkpoint_path=mpi_tmp_path,
+    )
 
-        # Run optimization and print summary of results.
-        propulator.propulate()
-        best = propulator.summarize(top_n=1, debug=2)
-        assert best[0][0].loss < 30.0
+    # Run optimization and print summary of results.
+    propulator.propulate()
diff --git a/tests/test_surrogate.py b/tests/test_surrogate.py
new file mode 100644
index 00000000..8a08b6a9
--- /dev/null
+++ b/tests/test_surrogate.py
@@ -0,0 +1,381 @@
+import logging
+import os
+import random
+from functools import partial
+from pathlib import Path
+from typing import Dict, Generator, Tuple, Union
+
+import pytest
+import torch
+from mpi4py import MPI
+from torch import nn
+from torch.utils.data import DataLoader, TensorDataset
+from torchmetrics import Accuracy
+from torchvision.datasets import MNIST
+from torchvision.transforms import Compose, Normalize, ToTensor
+
+from propulate import Islands, surrogate
+from propulate.utils import get_default_propagator, set_logger_config
+
+log = logging.getLogger(__name__)  # Get logger instance.
+set_logger_config()
+
+
+class Net(nn.Module):
+    """Convolutional neural network class."""
+
+    def __init__(
+        self, conv_layers: int, activation: nn.Module, lr: float, loss_fn: nn.Module
+    ) -> None:
+        """
+        Set up neural network.
+
+        Parameters
+        ----------
+        conv_layers : int
+            The number of convolutional layers.
+        activation : torch.nn.modules.activation
+            The activation function to use.
+        lr : float
+            The learning rate.
+        loss_fn : torch.nn.modules.loss
+            The loss function.
+        """
+        super(Net, self).__init__()
+
+        self.lr = lr  # Set learning rate.
+        self.loss_fn = loss_fn  # Set the loss function used for training the model.
+        self.best_accuracy = 0.0  # Initialize the model's best validation accuracy.
+        layers = []  # Set up the model architecture (depending on number of convolutional layers specified).
+        layers += [
+            nn.Sequential(
+                nn.Conv2d(in_channels=1, out_channels=10, kernel_size=3, padding=1),
+                activation(),
+            ),
+        ]
+        layers += [
+            nn.Sequential(
+                nn.Conv2d(in_channels=10, out_channels=10, kernel_size=3, padding=1),
+                activation(),
+            )
+            for _ in range(conv_layers - 1)
+        ]
+
+        self.fc = nn.Linear(in_features=7840, out_features=10)  # MNIST has 10 classes.
+        self.conv_layers = nn.Sequential(*layers)
+        self.val_acc = Accuracy("multiclass", num_classes=10)
+        self.train_acc = Accuracy("multiclass", num_classes=10)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            The data sample.
+
+        Returns
+        -------
+        torch.Tensor
+            The model's predictions for input data sample.
+        """
+        b, c, w, h = x.size()
+        x = self.conv_layers(x)
+        x = x.view(b, 10 * 28 * 28)
+        x = self.fc(x)
+        return x
+
+    def training_step(self, batch: Tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
+        """
+        Calculate loss for training step in Lightning train loop.
+
+        Parameters
+        ----------
+        batch : Tuple[torch.Tensor, torch.Tensor]
+            The input batch.
+
+        Returns
+        -------
+        torch.Tensor
+            The training loss for the input batch.
+        """
+        x, y = batch
+        pred = self(x)
+        loss_val = self.loss_fn(pred, y)
+        return loss_val
+
+    def validation_step(self, batch: Tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
+        """
+        Calculate loss for validation step in Lightning validation loop during training.
+
+        Parameters
+        ----------
+        batch : Tuple[torch.Tensor, torch.Tensor]
+            The current batch.
+
+        Returns
+        -------
+        torch.Tensor
+            The validation loss for the input batch.
+        """
+        x, y = batch
+        pred = self(x)
+        loss_val = self.loss_fn(pred, y)
+        return loss_val
+
+    def configure_optimizers(self) -> torch.optim.SGD:
+        """
+        Configure the sgd optimizer.
+
+        Returns
+        -------
+        torch.optim.sgd.SGD
+            A stochastic gradient descent optimizer.
+        """
+        return torch.optim.SGD(self.parameters(), lr=self.lr)
+
+
+def get_data_loaders(batch_size: int, root=Path) -> Tuple[DataLoader, DataLoader]:
+    """
+    Get MNIST train and validation dataloaders.
+
+    Parameters
+    ----------
+    batch_size : int
+        The batch size.
+
+    Returns
+    -------
+    DataLoader
+        The training dataloader.
+    DataLoader
+        The validation dataloader.
+    """
+    data_transform = Compose([ToTensor(), Normalize((0.1307,), (0.3081,))])
+
+    # Set empty DataLoader.
+    train_loader = DataLoader(
+        dataset=TensorDataset(torch.empty(0), torch.empty(0)),
+        batch_size=batch_size,
+        shuffle=False,
+    )
+
+    if MPI.COMM_WORLD.Get_rank() == 0:  # Only root downloads data.
+        train_loader = DataLoader(
+            dataset=MNIST(
+                download=True, root=root, transform=data_transform, train=True
+            ),  # Use MNIST training dataset.
+            batch_size=batch_size,  # Batch size
+            shuffle=True,  # Shuffle data.
+        )
+
+    if not hasattr(get_data_loaders, "barrier_called"):
+        MPI.COMM_WORLD.Barrier()
+
+        setattr(get_data_loaders, "barrier_called", True)
+
+    if MPI.COMM_WORLD.Get_rank() != 0:
+        train_loader = DataLoader(
+            dataset=MNIST(
+                download=False, root=root, transform=data_transform, train=True
+            ),  # Use MNIST training dataset.
+            batch_size=batch_size,  # Batch size
+            shuffle=True,  # Shuffle data.
+        )
+    val_loader = DataLoader(
+        dataset=MNIST(
+            download=False, root=root, transform=data_transform, train=False
+        ),  # Use MNIST testing dataset.
+        batch_size=1,  # Batch size
+        shuffle=False,  # Do not shuffle data.
+    )
+    return train_loader, val_loader
+
+
+def ind_loss(
+    params: Dict[str, Union[int, float, str]], root: Path
+) -> Generator[float, None, None]:
+    """
+    Loss function for evolutionary optimization with Propulate. Minimize the model's negative validation accuracy.
+
+    Parameters
+    ----------
+    params : Dict[str, int | float | str]
+        The parameters to be optimized.
+
+    Returns
+    -------
+    Generator[float, None, None]
+        Yields the negative validation accuracy in regular intervals during training of the model.
+    """
+    # Extract hyperparameter combination to test from input dictionary.
+    conv_layers = int(params["conv_layers"])  # Number of convolutional layers
+    activation = str(params["activation"])  # Activation function
+    lr = float(params["lr"])  # Learning rate
+
+    epochs: int = 2  # Number of epochs to train
+
+    rank: int = MPI.COMM_WORLD.Get_rank()  # Get rank of current worker
+
+    num_gpus = torch.cuda.device_count()  # Number of GPUs available
+    if num_gpus == 0:
+        device = torch.device("cpu")
+    else:
+        device_index = rank % num_gpus
+        device = torch.device(
+            f"cuda:{device_index}" if torch.cuda.is_available() else "cpu"
+        )
+
+    log.info(f"Rank: {rank}, Using device: {device}")
+
+    activations = {
+        "relu": nn.ReLU,
+        "sigmoid": nn.Sigmoid,
+        "tanh": nn.Tanh,
+    }  # Define activation function mapping.
+    activation = activations[activation]  # Get activation function.
+    loss_fn = (
+        torch.nn.CrossEntropyLoss()
+    )  # Use cross-entropy loss for multi-class classification.
+
+    model = Net(conv_layers, activation, lr, loss_fn).to(
+        device
+    )  # Set up neural network with specified hyperparameters.
+    model.best_accuracy = 0.0  # Initialize the model's best validation accuracy.
+
+    train_loader, val_loader = get_data_loaders(
+        batch_size=8, root=root
+    )  # Get training and validation data loaders.
+
+    # Configure optimizer.
+    optimizer = model.configure_optimizers()
+
+    for epoch in range(epochs):
+        model.train()
+        total_train_loss = 0
+        # Training loop
+        for batch_idx, (data, target) in enumerate(train_loader):
+            data, target = data.to(device), target.to(device)
+            # Zero out gradients.
+            optimizer.zero_grad()
+            # Forward + backward pass and optimizer step to update parameters.
+            loss = model.training_step((data, target))
+            loss.backward()
+            optimizer.step()
+            # Update loss.
+            total_train_loss += loss.item()
+
+        avg_train_loss = total_train_loss / len(train_loader)
+        log.info(f"Epoch {epoch+1}: Avg Training Loss: {avg_train_loss}")
+
+        # Validation loop
+        model.eval()
+        total_val_loss = 0
+        with torch.no_grad():
+            for batch_idx, (data, target) in enumerate(val_loader):
+                data, target = data.to(device), target.to(device)
+                # Forward pass
+                loss = model.validation_step((data, target))
+                # Update loss.
+                total_val_loss += loss.item()
+
+        avg_val_loss = total_val_loss / len(val_loader)
+        log.info(f"Epoch {epoch+1}: Avg Validation Loss: {avg_val_loss}")
+
+        yield avg_val_loss
+
+
+def set_seeds(seed_value: int = 42) -> None:
+    """
+    Set seed for reproducibility.
+
+    Parameters
+    ----------
+    seed_value : int, optional
+        The seed to use. Default is 42.
+    """
+    random.seed(seed_value)  # Python random module
+    torch.manual_seed(seed_value)  # PyTorch random number generator for CPU
+    torch.cuda.manual_seed(seed_value)  # PyTorch random number generator for all GPUs
+    torch.cuda.manual_seed_all(
+        seed_value
+    )  # PyTorch random number generator for multi-GPU
+    torch.backends.cudnn.deterministic = True  # Use deterministic algorithms.
+    torch.backends.cudnn.benchmark = False  # Disable to be deterministic.
+    os.environ["PYTHONHASHSEED"] = str(seed_value)  # Python hash seed
+
+
+@pytest.mark.mpi(min_size=4)
+def test_mnist_static(mpi_tmp_path):
+    """Test static surrogate using a torch convolutional network on the MNIST dataset."""
+    num_generations = 3  # Number of generations
+    pop_size = 2 * MPI.COMM_WORLD.size  # Breeding population size
+    limits = {
+        "conv_layers": (2, 3),
+        "activation": ("relu", "sigmoid", "tanh"),
+        "lr": (0.01, 0.0001),
+    }  # Define search space.
+    rng = random.Random(
+        MPI.COMM_WORLD.rank
+    )  # Set up separate random number generator for evolutionary optimizer.
+    set_seeds(42 * MPI.COMM_WORLD.Get_rank())  # set seed for torch
+    propagator = get_default_propagator(  # Get default evolutionary operator.
+        pop_size=pop_size,  # Breeding population size
+        limits=limits,  # Search space
+        crossover_prob=0.7,  # Crossover probability
+        mutation_prob=0.4,  # Mutation probability
+        random_init_prob=0.1,  # Random-initialization probability
+        rng=rng,  # Random number generator for evolutionary optimizer
+    )
+    islands = Islands(  # Set up island model.
+        loss_fn=partial(ind_loss, root=mpi_tmp_path),  # Loss function to optimize
+        propagator=propagator,  # Evolutionary operator
+        rng=rng,  # Random number generator
+        generations=num_generations,  # Number of generations per worker
+        num_islands=1,  # Number of islands
+        checkpoint_path=mpi_tmp_path,
+        surrogate_factory=lambda: surrogate.StaticSurrogate(),
+    )
+    islands.evolve(  # Run evolutionary optimization.
+        top_n=1,  # Print top-n best individuals on each island in summary.
+        logging_interval=1,  # Logging interval
+        debug=2,  # Verbosity level
+    )
+    MPI.COMM_WORLD.barrier()
+    delattr(get_data_loaders, "barrier_called")
+
+
+@pytest.mark.mpi(min_size=4)
+def test_mnist_dynamic(mpi_tmp_path):
+    """Test static surrogate using a torch convolutional network on the MNIST dataset."""
+    num_generations = 3  # Number of generations
+    pop_size = 2 * MPI.COMM_WORLD.size  # Breeding population size
+    limits = {
+        "conv_layers": (2, 3),
+        "activation": ("relu", "sigmoid", "tanh"),
+        "lr": (0.01, 0.0001),
+    }  # Define search space.
+    rng = random.Random(
+        MPI.COMM_WORLD.rank
+    )  # Set up separate random number generator for evolutionary optimizer.
+    set_seeds(42 + MPI.COMM_WORLD.rank)  # set seed for torch
+    propagator = get_default_propagator(  # Get default evolutionary operator.
+        pop_size=pop_size,  # Breeding population size
+        limits=limits,  # Search space
+        rng=rng,  # Random number generator for evolutionary optimizer
+    )
+    islands = Islands(  # Set up island model.
+        loss_fn=partial(ind_loss, root=mpi_tmp_path),  # Loss function to optimize
+        propagator=propagator,  # Evolutionary operator
+        rng=rng,  # Random number generator
+        generations=num_generations,  # Number of generations per worker
+        num_islands=1,  # Number of islands
+        checkpoint_path=mpi_tmp_path,
+        surrogate_factory=lambda: surrogate.DynamicSurrogate(limits),
+    )
+    islands.evolve(  # Run evolutionary optimization.
+        top_n=1,  # Print top-n best individuals on each island in summary.
+        logging_interval=1,  # Logging interval
+        debug=2,  # Verbosity level
+    )