Hierarchical optimization: Support for box constraints on offset and …

…scaling parameters (#1238)

* Initial working

Implemented scaling and offset bounds for the analytical and numerical solver of the hierarchical problem.

* Fix hierarchical test

* Add test_constrained_inner_solver

* Add tests for non coupled parameters

* intentional failure

* intentional failure

* intentional failure

* Test fix

It seems L-BFGS-B would not converge to optimum in some cases if it was provided with the big bounds [-1e20, 1e20]. It seems like it is a scipy issue? Not sure.
In any case, the numerical optimizer is better if the optimization is not provided with these dummy bounds. But in this case, we need to sample the bounds outside of the pypesto minimize function and provide it as x_guesses.

* Add notebook changes

* Update C.py

* Revert "Update C.py"

This reverts commit f099a2b.

* Fix sampling + test

* Cleanup

* Some updates

* Some more small updates

* Add inner bounds to parameter plot

* Fix import issues

* Fix hier test

* Fix parameter plot test

* Remove figures

* DIlan&Daniel review changes

* Update parameters.py

* Update solver.py

pypesto/hierarchical/parameter.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Literal
+from typing import Any, Literal, Optional
 
 import numpy as np
 
@@ -24,8 +24,9 @@ class InnerParameter:
     Attributes
     ----------
     coupled:
-        Whether the inner parameter is part of an observable that has both
-        an offset and scaling inner parameter.
+        If the inner parameter is part of an observable that has both
+        an offset and scaling inner parameter, this attribute points to
+        the other inner parameter. Otherwise, it is None.
     dummy_value:
         Value to be used when the optimal parameter is not yet known
         (in particular to simulate unscaled observables).
@@ -62,7 +63,7 @@ def __init__(
         See class attributes.
         """
         self.inner_parameter_id: str = inner_parameter_id
-        self.coupled = False
+        self.coupled: InnerParameter = None
         self.inner_parameter_type: str = inner_parameter_type
 
         if scale not in {LIN, LOG, LOG10}:
@@ -82,7 +83,12 @@ def __init__(
 
         self.lb: float = lb
         self.ub: float = ub
-        self.check_bounds()
+        # Scaling and offset parameters can be bounded arbitrarily
+        if inner_parameter_type not in (
+            InnerParameterType.SCALING,
+            InnerParameterType.OFFSET,
+        ):
+            self.check_bounds()
         self.ixs: Any = ixs
 
         if dummy_value is None:
@@ -114,3 +120,21 @@ def check_bounds(self):
                 f"`[{expected_lb}, {expected_ub}]`. "
                 f"All expected parameter bounds:\n{INNER_PARAMETER_BOUNDS}"
             )
+
+    def is_within_bounds(self, value):
+        """Check whether a value is within the bounds."""
+        if value < self.lb or value > self.ub:
+            return False
+        return True
+
+    def get_unsatisfied_bound(self, value) -> Optional[str]:
+        """Get the unsatisfied bound index, if any."""
+        if value < self.lb:
+            return LOWER_BOUND
+        elif value > self.ub:
+            return UPPER_BOUND
+        return None
+
+    def get_bounds(self) -> dict:
+        """Get the bounds."""
+        return {LOWER_BOUND: self.lb, UPPER_BOUND: self.ub}

pypesto/hierarchical/petab.py

@@ -39,6 +39,11 @@ def correct_parameter_df_bounds(parameter_df: pd.DataFrame) -> pd.DataFrame:
     def correct_row(row: pd.Series) -> pd.Series:
         if pd.isna(row[PARAMETER_TYPE]):
             return row
+        if row[PARAMETER_TYPE] in [
+            InnerParameterType.SCALING,
+            InnerParameterType.OFFSET,
+        ]:
+            return row
         bounds = INNER_PARAMETER_BOUNDS[row[PARAMETER_TYPE]]
         row[PETAB_LOWER_BOUND] = bounds[PYPESTO_LOWER_BOUND]
         row[PETAB_UPPER_BOUND] = bounds[PYPESTO_UPPER_BOUND]

pypesto/hierarchical/problem.py

@@ -224,34 +224,49 @@ def inner_problem_from_petab_problem(
         par.ixs = ix_matrices[par.inner_parameter_id]
 
     par_group_types = {
-        tuple(obs_pars.split(';')): {
+        tuple(obs_pars.split(';')): (
             petab_problem.parameter_df.loc[obs_par, PARAMETER_TYPE]
             for obs_par in obs_pars.split(';')
-        }
+        )
         for (obs_id, obs_pars), _ in petab_problem.measurement_df.groupby(
             [petab.OBSERVABLE_ID, petab.OBSERVABLE_PARAMETERS], dropna=True
         )
         if ';' in obs_pars  # prefilter for at least 2 observable parameters
     }
 
     coupled_pars = {
-        par
+        group
         for group, types in par_group_types.items()
         if (
             (InnerParameterType.SCALING in types)
             and (InnerParameterType.OFFSET in types)
         )
-        for par in group
     }
 
+    # Check each group is of length 2
+    for group in coupled_pars:
+        if len(group) != 2:
+            raise ValueError(
+                f"Expected exactly 2 parameters in group {group}: a scaling "
+                f"and an offset parameter."
+            )
+
+    id_to_par = {par.inner_parameter_id: par for par in inner_parameters}
+
+    # assign coupling
     for par in inner_parameters:
         if par.inner_parameter_type not in [
             InnerParameterType.SCALING,
             InnerParameterType.OFFSET,
         ]:
             continue
-        if par.inner_parameter_id in coupled_pars:
-            par.coupled = True
+        for group in coupled_pars:
+            if par.inner_parameter_id in group:
+                coupled_parameter_id = group[
+                    group.index(par.inner_parameter_id) - 1
+                ]
+                par.coupled = id_to_par[coupled_parameter_id]
+                break
 
     return AmiciInnerProblem(xs=inner_parameters, data=data, edatas=edatas)
 

pypesto/hierarchical/solver.py

@@ -13,6 +13,7 @@
     apply_offset,
     apply_scaling,
     apply_sigma,
+    compute_bounded_optimal_scaling_offset_coupled,
     compute_nllh,
     compute_optimal_offset,
     compute_optimal_offset_coupled,
@@ -91,30 +92,78 @@ def solve(
             ``problem``.
         """
         x_opt = {}
-
         data = copy.deepcopy(problem.data)
 
         # compute optimal offsets
         for x in problem.get_xs_for_type(InnerParameterType.OFFSET):
-            if x.coupled:
+            if x.coupled is not None:
                 x_opt[x.inner_parameter_id] = compute_optimal_offset_coupled(
                     data=data, sim=sim, sigma=sigma, mask=x.ixs
                 )
+
+                # calculate the optimal coupled scaling
+                coupled_scaling = x.coupled
+                x_opt[
+                    coupled_scaling.inner_parameter_id
+                ] = compute_optimal_scaling(
+                    data=data,
+                    sim=sim,
+                    sigma=sigma,
+                    mask=coupled_scaling.ixs,
+                    optimal_offset=x_opt[x.inner_parameter_id],
+                )
+
+                # check whether they both satisfy their bounds
+                if x.is_within_bounds(
+                    x_opt[x.inner_parameter_id]
+                ) and coupled_scaling.is_within_bounds(
+                    x_opt[coupled_scaling.inner_parameter_id]
+                ):
+                    continue
+                else:
+                    # if not, we need to recompute them
+                    (
+                        x_opt[coupled_scaling.inner_parameter_id],
+                        x_opt[x.inner_parameter_id],
+                    ) = compute_bounded_optimal_scaling_offset_coupled(
+                        data=data,
+                        sim=sim,
+                        sigma=sigma,
+                        s=coupled_scaling,
+                        b=x,
+                        s_opt_value=x_opt[coupled_scaling.inner_parameter_id],
+                        b_opt_value=x_opt[x.inner_parameter_id],
+                    )
+            # compute non-coupled optimal offset
             else:
                 x_opt[x.inner_parameter_id] = compute_optimal_offset(
                     data=data, sim=sim, sigma=sigma, mask=x.ixs
                 )
+                # check if the solution is within bounds
+                # if not, we set it to the unsatisfied bound
+                if not x.is_within_bounds(x_opt[x.inner_parameter_id]):
+                    x_opt[x.inner_parameter_id] = x.get_bounds()[
+                        x.get_unsatisfied_bound(x_opt[x.inner_parameter_id])
+                    ]
+
         # apply offsets
         for x in problem.get_xs_for_type(InnerParameterType.OFFSET):
             apply_offset(
                 offset_value=x_opt[x.inner_parameter_id], data=data, mask=x.ixs
             )
 
-        # compute optimal scalings
+        # compute non-coupled optimal scalings
         for x in problem.get_xs_for_type(InnerParameterType.SCALING):
-            x_opt[x.inner_parameter_id] = compute_optimal_scaling(
-                data=data, sim=sim, sigma=sigma, mask=x.ixs
-            )
+            if x.coupled is None:
+                x_opt[x.inner_parameter_id] = compute_optimal_scaling(
+                    data=data, sim=sim, sigma=sigma, mask=x.ixs
+                )
+                # check if the solution is within bounds
+                # if not, we set it to the unsatisfied bound
+                if not x.is_within_bounds(x_opt[x.inner_parameter_id]):
+                    x_opt[x.inner_parameter_id] = x.get_bounds()[
+                        x.get_unsatisfied_bound(x_opt[x.inner_parameter_id])
+                    ]
         # apply scalings
         for x in problem.get_xs_for_type(InnerParameterType.SCALING):
             apply_scaling(
@@ -186,6 +235,8 @@ def __init__(
         self.x_guesses = None
         self.dummy_lb = -1e20
         self.dummy_ub = +1e20
+        self.user_specified_lb = None
+        self.user_specified_ub = None
 
     def initialize(self):
         """(Re-)initialize the solver."""
@@ -216,21 +267,19 @@ def solve(
             Whether to scale the results to the parameter scale specified in
             ``problem``.
         """
-        pars = problem.xs.values()
-        # We currently cannot handle constraints on inner parameters correctly,
-        # and would have to assume [-inf, inf]. However, this may not be
-        # supported by all inner optimizers, so we go for some (arbitrary)
-        # large value.
-        lb = np.array(
-            [
-                0
-                if x.inner_parameter_type == InnerParameterType.SIGMA
-                else self.dummy_lb
-                for x in pars
+        pars = list(problem.xs.values())
+
+        # This has to be done only once
+        if self.user_specified_lb is None or self.user_specified_ub is None:
+            self.user_specified_lb = [
+                i for i in range(len(pars)) if pars[i].lb != -np.inf
+            ]
+            self.user_specified_ub = [
+                i for i in range(len(pars)) if pars[i].ub != np.inf
             ]
-        )
 
-        ub = np.full(shape=len(pars), fill_value=self.dummy_ub)
+        lb = [x.lb for x in pars]
+        ub = [x.ub for x in pars]
 
         x_guesses = self.sample_startpoints(problem, pars)
 
@@ -265,20 +314,33 @@ def fun(x):
         pypesto_problem = Problem(
             objective, lb=lb, ub=ub, x_names=x_names, **self.problem_kwargs
         )
-
         pypesto_problem.set_x_guesses(
             x_guesses[:, pypesto_problem.x_free_indices]
         )
 
         # perform the actual optimization
         result = minimize(pypesto_problem, **self.minimize_kwargs)
-
         best_par = result.optimize_result.list[0]['x']
 
-        if (np.isclose(best_par, lb) | np.isclose(best_par, ub)).any():
+        # Check if the index of an optimized parameter on the dummy bound
+        # is not in the list of specified bounds. If so, raise an error.
+        if any(
+            (
+                i not in self.user_specified_lb
+                for i, x in enumerate(best_par)
+                if x == self.dummy_lb
+            )
+        ) or any(
+            (
+                i not in self.user_specified_ub
+                for i, x in enumerate(best_par)
+                if x == self.dummy_ub
+            )
+        ):
             raise RuntimeError(
-                "Active bounds in inner problem optimization. This can result "
-                "in incorrect gradient computation for the outer parameters."
+                f"An optimal inner parameter is on the defualt dummy bound of numerical optimization. "
+                f"This means the optimal inner parameter is either extremely large (>={self.dummy_ub})"
+                f"or extremely small (<={self.dummy_lb}). Consider changing the inner parameter bounds."
             )
 
         x_opt = dict(zip(pypesto_problem.x_names, best_par))