From 61c912d56b4e435105da9e95dbdb31ae1303a807 Mon Sep 17 00:00:00 2001
From: Daniel Weindl <dweindl@users.noreply.github.com>
Date: Thu, 26 Sep 2024 14:03:32 +0200
Subject: [PATCH] Use cloudpickle for serializing NegLogParameterPriors (#1467)

Cloudpickle is able to handle more complex objects than pickle.

See #1465
Closes #1465
---
 pypesto/objective/priors.py    |  9 +++++++++
 test/optimize/test_optimize.py | 16 ++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/pypesto/objective/priors.py b/pypesto/objective/priors.py
index e460c1572..4ffcdaf6a 100644
--- a/pypesto/objective/priors.py
+++ b/pypesto/objective/priors.py
@@ -3,6 +3,7 @@
 from collections.abc import Sequence
 from typing import Callable, Union
 
+import cloudpickle
 import numpy as np
 
 from .. import C
@@ -67,6 +68,14 @@ def __init__(
         self.prior_list = prior_list
         super().__init__(x_names)
 
+    def __getstate__(self):
+        """Get state using cloudpickle."""
+        return cloudpickle.dumps(self.__dict__)
+
+    def __setstate__(self, state):
+        """Set state using cloudpickle."""
+        self.__dict__.update(cloudpickle.loads(state))
+
     def call_unprocessed(
         self,
         x: np.ndarray,
diff --git a/test/optimize/test_optimize.py b/test/optimize/test_optimize.py
index 3a9098662..160e5baa9 100644
--- a/test/optimize/test_optimize.py
+++ b/test/optimize/test_optimize.py
@@ -524,6 +524,22 @@ def test_ess_multiprocess(problem, request):
 
     from pypesto.optimize.ess import ESSOptimizer, FunctionEvaluatorMP, RefSet
 
+    # augment objective with parameter prior to check it's copyable
+    #  https://github.com/ICB-DCM/pyPESTO/issues/1465
+    #  https://github.com/ICB-DCM/pyPESTO/pull/1467
+    problem.objective = pypesto.objective.AggregatedObjective(
+        [
+            problem.objective,
+            pypesto.objective.NegLogParameterPriors(
+                [
+                    pypesto.objective.get_parameter_prior_dict(
+                        0, "uniform", [0, 1], "lin"
+                    )
+                ]
+            ),
+        ]
+    )
+
     ess = ESSOptimizer(
         max_iter=20,
         # also test passing a callable as local_optimizer