Broadcast and mapping of scale and unscale functions (#505)

* allow broadcasting in `map_scale`

* fixup

* fixup

* add docstring

Author: yannikschaelte

petab/parameters.py

@@ -352,7 +352,7 @@ def get_priors_from_df(parameter_df: pd.DataFrame,
 
 
 def scale(parameter: numbers.Number, scale_str: 'str') -> numbers.Number:
-    """Scale parameter according to scale_str
+    """Scale parameter according to `scale_str`.
 
     Arguments:
         parameter:
@@ -375,7 +375,7 @@ def scale(parameter: numbers.Number, scale_str: 'str') -> numbers.Number:
 
 
 def unscale(parameter: numbers.Number, scale_str: 'str') -> numbers.Number:
-    """Unscale parameter according to scale_str
+    """Unscale parameter according to `scale_str`.
 
     Arguments:
         parameter:
@@ -397,12 +397,49 @@ def unscale(parameter: numbers.Number, scale_str: 'str') -> numbers.Number:
     raise ValueError("Invalid parameter scaling: " + scale_str)
 
 
-def map_scale(parameters: Iterable[numbers.Number],
-              scale_strs: Iterable[str]) -> Iterable[numbers.Number]:
-    """As scale(), but for Iterables"""
+def map_scale(
+    parameters: Iterable[numbers.Number],
+    scale_strs: Union[Iterable[str], str]
+) -> Iterable[numbers.Number]:
+    """Scale the parameters, i.e. as `scale()`, but for Iterables.
+
+    Arguments:
+        parameters:
+            Parameters to be scaled.
+        scale_strs:
+            Scales to apply. Broadcast if a single string.
+
+    Returns:
+        parameters:
+            The scaled parameters.
+    """
+    if isinstance(scale_strs, str):
+        scale_strs = [scale_strs] * len(parameters)
     return map(lambda x: scale(x[0], x[1]), zip(parameters, scale_strs))
 
 
+def map_unscale(
+    parameters: Iterable[numbers.Number],
+    scale_strs: Union[Iterable[str], str]
+) -> Iterable[numbers.Number]:
+    """Unscale the parameters, i.e. as `unscale()`, but for Iterables.
+
+    Arguments:
+        parameters:
+            Parameters to be unscaled.
+        scale_strs:
+            Scales that the parameters are currently on.
+            Broadcast if a single string.
+
+    Returns:
+        parameters:
+            The unscaled parameters.
+    """
+    if isinstance(scale_strs, str):
+        scale_strs = [scale_strs] * len(parameters)
+    return map(lambda x: unscale(x[0], x[1]), zip(parameters, scale_strs))
+
+
 def normalize_parameter_df(parameter_df: pd.DataFrame) -> pd.DataFrame:
     """Add missing columns and fill in default values."""
     df = parameter_df.copy(deep=True)

tests/test_parameters.py

@@ -1,8 +1,9 @@
 """Tests for petab/parameters.py"""
 import tempfile
 import pytest
-
+import numpy as np
 import pandas as pd
+
 import petab
 from petab.C import *
 
@@ -170,3 +171,29 @@ def test_normalize_parameter_df():
 
     # check is valid petab
     petab.check_parameter_df(actual)
+
+
+def test_scale_unscale():
+    """Test the parameter scaling functions."""
+    par = 2.5
+    # scale
+    assert petab.scale(par, LIN) == par
+    assert petab.scale(par, LOG) == np.log(par)
+    assert petab.scale(par, LOG10) == np.log10(par)
+    # unscale
+    assert petab.unscale(par, LIN) == par
+    assert petab.unscale(par, LOG) == np.exp(par)
+    assert petab.unscale(par, LOG10) == 10**par
+
+    # map scale
+    assert list(petab.map_scale([par]*3, [LIN, LOG, LOG10])) == \
+        [par, np.log(par), np.log10(par)]
+    # map unscale
+    assert list(petab.map_unscale([par]*3, [LIN, LOG, LOG10])) == \
+        [par, np.exp(par), 10**par]
+
+    # map broadcast
+    assert list(petab.map_scale([par, 2*par], LOG)) == \
+        list(np.log([par, 2*par]))
+    assert list(petab.map_unscale([par, 2*par], LOG)) == \
+        list(np.exp([par, 2*par]))