Select: problem-specific minimize method for SaCeSS

pypesto/select/__init__.py

@@ -7,7 +7,7 @@
 """
 
 from . import postprocessors
-from .misc import model_to_pypesto_problem
+from .misc import SacessMinimizeMethod, model_to_pypesto_problem
 from .problem import Problem
 
 try:

pypesto/select/misc.py

@@ -2,6 +2,7 @@
 
 import logging
 from collections.abc import Iterable
+from pathlib import Path
 
 import pandas as pd
 import petab.v1 as petab
@@ -10,7 +11,13 @@
 from petab_select import Model, parameter_string_to_value
 from petab_select.constants import PETAB_PROBLEM
 
+from ..history import Hdf5History
 from ..objective import Objective
+from ..optimize import Optimizer
+from ..optimize.ess import (
+    SacessOptimizer,
+    get_default_ess_options,
+)
 from ..petab import PetabImporter
 from ..problem import Problem
 
@@ -164,3 +171,78 @@ def correct_x_guesses(
                 corrected_x_guess.append(corrected_value)
             corrected_x_guesses.append(corrected_x_guess)
     return corrected_x_guesses
+
+
+class SacessMinimizeMethod:
+    """Create a minimize method for SaCeSS that adapts to each problem.
+
+    When a pyPESTO SaCeSS optimizer is created, it takes the problem
+    dimension as input. Hence, an optimizer needs to be constructed for
+    each problem. Objects of this class act like a minimize method for model
+    selection, but a new problem-specific SaCeSS optimizer will be created
+    every time a model is minimized.
+
+    Instance attributes correspond to pyPESTO's SaCeSS optimizer, and are
+    documented there. Extra keyword arguments supplied to the constructor
+    will be passed on to the constructor of the SaCeSS optimizer, for example,
+    `max_walltime_s` can be specified in this way. If specified, `tmpdir` will
+    be treated as a parent directory.
+    """
+
+    def __init__(
+        self,
+        num_workers: int,
+        local_optimizer: Optimizer = None,
+        tmpdir: str | Path | None = None,
+        save_history: bool = False,
+        **optimizer_kwargs,
+    ):
+        """Construct a minimize-like object."""
+        self.num_workers = num_workers
+        self.local_optimizer = local_optimizer
+        self.optimizer_kwargs = optimizer_kwargs
+        self.save_history = save_history
+
+        self.tmpdir = tmpdir
+        if self.tmpdir is not None:
+            self.tmpdir = Path(self.tmpdir)
+
+        if self.save_history and self.tmpdir is None:
+            self.tmpdir = Path.cwd() / "sacess_tmpdir"
+
+    def __call__(self, problem: Problem, model_hash: str, **minimize_options):
+        """Create then run a problem-specific sacess optimizer."""
+        # create optimizer
+        ess_init_args = get_default_ess_options(
+            num_workers=self.num_workers,
+            dim=problem.dim,
+        )
+        for x in ess_init_args:
+            x["local_optimizer"] = self.local_optimizer
+        model_tmpdir = None
+        if self.tmpdir is not None:
+            model_tmpdir = self.tmpdir / model_hash
+            model_tmpdir.mkdir(exist_ok=False, parents=True)
+
+        ess = SacessOptimizer(
+            ess_init_args=ess_init_args,
+            tmpdir=model_tmpdir,
+            **self.optimizer_kwargs,
+        )
+
+        # optimize
+        result = ess.minimize(
+            problem=problem,
+            **minimize_options,
+        )
+
+        if self.save_history:
+            history_dir = model_tmpdir / "history"
+            history_dir.mkdir(exist_ok=False, parents=True)
+            for history_index, history in enumerate(ess.histories):
+                Hdf5History.from_history(
+                    other=history,
+                    file=history_dir / (str(history_index) + ".hdf5"),
+                    id_=history_index,
+                )
+        return result

pypesto/select/model_problem.py

@@ -9,7 +9,7 @@
 from ..optimize import minimize
 from ..problem import Problem
 from ..result import OptimizerResult, Result
-from .misc import model_to_pypesto_problem
+from .misc import SacessMinimizeMethod, model_to_pypesto_problem
 
 OBJECTIVE_CUSTOMIZER_TYPE = Callable[[ObjectiveBase], None]
 TYPE_POSTPROCESSOR = Callable[["ModelProblem"], None]  # noqa: F821
@@ -142,9 +142,16 @@ def __init__(
     def minimize(self) -> Result:
         """Optimize the model.
 
-        Returns:
+        Returns
+        -------
             The optimization result.
         """
+        if isinstance(self.minimize_method, SacessMinimizeMethod):
+            return self.minimize_method(
+                self.pypesto_problem,
+                model_hash=self.model.get_hash(),
+                **self.minimize_options,
+            )
         return self.minimize_method(
             self.pypesto_problem,
             **self.minimize_options,
@@ -195,6 +202,12 @@ def create_fake_pypesto_result_from_fval(
     ----------
     fval:
         The objective function value.
+    evaluation_time:
+        CPU time taken to compute the objective function value.
+
+    Returns
+    -------
+    The dummy result.
     """
     result = Result()
 

test/select/test_select.py

@@ -20,7 +20,7 @@
 import pypesto.engine
 import pypesto.select
 import pypesto.visualize.select
-from pypesto.select import model_problem
+from pypesto.select import SacessMinimizeMethod, model_problem
 from pypesto.select.misc import correct_x_guesses
 
 model_problem_options = {
@@ -496,3 +496,31 @@ def model_to_pypesto_problem(
     # The custom objective and gradient were returned.
     assert test_fun == expected_fun
     assert np.isclose(test_grad, expected_grad).all()
+
+
+def test_sacess_minimize_method(pypesto_select_problem, initial_models):
+    """Test `SacessMinimizeMethod`.
+
+    Only ensures that the pipeline runs.
+    """
+    predecessor_model = initial_models[1]
+
+    minimize_method = SacessMinimizeMethod(
+        num_workers=2,
+        max_walltime_s=1,
+    )
+
+    minimize_options = {
+        "startpoint_method": pypesto.startpoint.UniformStartpoints(),
+    }
+
+    model_problem_options = {
+        "minimize_method": minimize_method,
+        "minimize_options": minimize_options,
+    }
+
+    pypesto_select_problem.select_to_completion(
+        model_problem_options=model_problem_options,
+        method=petab_select.Method.FORWARD,
+        predecessor_model=predecessor_model,
+    )

tox.ini

@@ -120,7 +120,7 @@ description =
     Test hierarchical optimization module
 
 [testenv:select]
-extras = test,amici,petab,select
+extras = test,amici,petab,select,fides
 commands =
     pytest --cov=pypesto --cov-report=xml --cov-append \
         test/select