add option to convert model assignments to experiments in v1 to v2 converter

Author: BSnelling

petab/v2/petab1to2.py

@@ -5,17 +5,20 @@
 import re
 import shutil
 import warnings
+from collections import defaultdict
 from contextlib import suppress
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from urllib.parse import urlparse
 from uuid import uuid4
 
 import pandas as pd
+import sympy as sp
 from pandas.io.common import get_handle, is_url
+from sbmlmath import TimeSymbol, sbml_math_to_sympy
 
 from .. import v1, v2
-from ..v1.math import sympify_petab
+from ..v1.math import petab_math_str, sympify_petab
 from ..v1.yaml import get_path_prefix, load_yaml, validate
 from ..versions import get_major_version
 from .models import MODEL_TYPE_SBML
@@ -24,7 +27,9 @@
 
 
 def petab1to2(
-    yaml_config: Path | str, output_dir: Path | str | None = None
+    yaml_config: Path | str,
+    output_dir: Path | str | None = None,
+    assignments_to_experiments: bool = False,
 ) -> v2.Problem | None:
     """Convert from PEtab 1.0 to PEtab 2.0 format.
 
@@ -46,27 +51,39 @@ def petab1to2(
     :param output_dir:
         The output directory to save the converted PEtab problem, or ``None``,
         to return a :class:`petab.v2.Problem` instance.
+    :param assignments_to_experiments:
+        Whether to replace time-dependent piecewise assignments in the model,
+        such as ``k := dose * piecewise(0, time < t_dose, 1)``, by
+        experiment periods and conditions.
 
     :raises ValueError:
         If the input is invalid or does not pass linting or if the generated
         files do not pass linting.
     """
     if output_dir is not None:
-        return petab_files_1to2(yaml_config, output_dir)
+        return petab_files_1to2(
+            yaml_config, output_dir, assignments_to_experiments
+        )
 
     with TemporaryDirectory() as tmp_dir:
-        petab_files_1to2(yaml_config, tmp_dir)
+        petab_files_1to2(yaml_config, tmp_dir, assignments_to_experiments)
         return v2.Problem.from_yaml(Path(tmp_dir, Path(yaml_config).name))
 
 
-def petab_files_1to2(yaml_config: Path | str | dict, output_dir: Path | str):
+def petab_files_1to2(
+    yaml_config: Path | str | dict,
+    output_dir: Path | str,
+    assignments_to_experiments: bool = False,
+):
     """Convert PEtab files from PEtab 1.0 to PEtab 2.0.
 
 
     :param yaml_config:
         The PEtab problem as dictionary or YAML file name.
     :param output_dir:
         The output directory to save the converted PEtab problem.
+    :param assignments_to_experiments:
+        See :func:`petab1to2`.
 
     :raises ValueError:
         If the input is invalid or does not pass linting or if the generated
@@ -96,6 +113,11 @@ def petab_files_1to2(yaml_config: Path | str | dict, output_dir: Path | str):
     )
     if v1.lint_problem(petab_problem):
         raise ValueError("Provided PEtab problem does not pass linting.")
+    new_periods = (
+        _assignments_to_periods(petab_problem)
+        if assignments_to_experiments
+        else {}
+    )
 
     output_dir = Path(output_dir)
 
@@ -115,7 +137,10 @@ def petab_files_1to2(yaml_config: Path | str | dict, output_dir: Path | str):
     for file in (
         model.location for model in new_yaml_config.model_files.values()
     ):
-        _copy_file(get_src_path(file), Path(get_dest_path(file)))
+        if new_periods:
+            petab_problem.model.to_file(get_dest_path(file))
+        else:
+            _copy_file(get_src_path(file), Path(get_dest_path(file)))
 
     # Update observable table
     for observable_file in new_yaml_config.observable_files:
@@ -128,6 +153,8 @@ def petab_files_1to2(yaml_config: Path | str | dict, output_dir: Path | str):
     # Update condition table
     for condition_file in new_yaml_config.condition_files:
         condition_df = v1.get_condition_df(get_src_path(condition_file))
+        if new_periods:
+            condition_df = petab_problem.condition_df
         condition_df = v1v2_condition_df(condition_df, petab_problem.model)
         v2.write_condition_df(condition_df, get_dest_path(condition_file))
 
@@ -145,7 +172,11 @@ def create_experiment_id(sim_cond_id: str, preeq_cond_id: str) -> str:
             preeq_cond_id
             and petab_problem.condition_df.loc[preeq_cond_id].notna().any()
         )
-        if not sim_cond_exists and not preeq_cond_exists:
+        if (
+            not sim_cond_exists
+            and not preeq_cond_exists
+            and (sim_cond_id, preeq_cond_id) not in new_periods
+        ):
             # if we have only all-NaN conditions, we don't create a new
             #  experiment
             return ""
@@ -189,6 +220,16 @@ def create_experiment_id(sim_cond_id: str, preeq_cond_id: str) -> str:
                 v2.C.CONDITION_ID: sim_cond_id,
             }
         )
+        experiments.extend(
+            {v2.C.EXPERIMENT_ID: exp_id, v2.C.TIME: t, v2.C.CONDITION_ID: c}
+            for t, c in new_periods.get((sim_cond_id, preeq_cond_id), [])
+        )
+    if new_periods:
+        # drop periods of conditions whose changes were all replaced
+        has_changes = petab_problem.condition_df.notna().any(axis=1)
+        experiments = [
+            e for e in experiments if has_changes[e[v2.C.CONDITION_ID]]
+        ]
     if experiments:
         exp_table_path = output_dir / "experiments.tsv"
         if exp_table_path.exists():
@@ -339,6 +380,101 @@ def _copy_file(src: Path | str, dest: Path):
         shutil.copy(str(src), str(dest))
 
 
+def _assignments_to_periods(
+    problem: v1.Problem,
+) -> dict[tuple[str, str], list[tuple[float, str]]]:
+    """Replace time-dependent piecewise assignments by experiment periods.
+
+    Assignment rules and initial assignments containing
+    ``piecewise(before, time < t_switch, after)`` are removed from the model,
+    which keeps the value before the switch. The condition table columns they
+    depend on are replaced by new conditions that set the assignment target
+    at the start of the respective periods.
+
+    :returns: The new ``(time, condition ID)`` periods for each measured
+        ``(simulation condition ID, preequilibration condition ID)`` pair.
+    """
+    sbml_model = problem.model.sbml_model
+    condition_df = problem.condition_df
+    overrides = {
+        cid: row.dropna().to_dict() for cid, row in condition_df.iterrows()
+    }
+    pairs = list(
+        problem.get_simulation_conditions_from_measurement_df()
+        .reindex(
+            columns=[
+                v1.C.SIMULATION_CONDITION_ID,
+                v1.C.PREEQUILIBRATION_CONDITION_ID,
+            ],
+            fill_value="",
+        )
+        .itertuples(index=False)
+    )
+    parameter_ids = set(problem.parameter_df.index)
+    periods = defaultdict(list)
+    # (target ID, target value) -> new condition ID
+    new_conditions = {}
+    replaced_columns = set()
+
+    for assignment in [
+        *(r for r in sbml_model.getListOfRules() if r.isAssignment()),
+        *sbml_model.getListOfInitialAssignments(),
+    ]:
+        expr = sbml_math_to_sympy(assignment)
+        pws = [pw for pw in expr.atoms(sp.Piecewise) if pw.has(TimeSymbol)]
+        if not pws:
+            continue
+        (pw,) = pws
+        (time,) = pw.atoms(TimeSymbol)
+        (before, switch), _ = pw.args
+        t_switch = sp.solve(switch.lhs - switch.rhs, time)[0]
+
+        target_id = assignment.getId()
+        target = sbml_model.getParameter(target_id)
+        target.setValue(float(expr.subs(pw, before)))
+        target.setConstant(False)
+        assignment.removeFromParentAndDelete()
+        replaced_columns |= {str(s) for s in expr.free_symbols}
+
+        for sim_id, preeq_id in pairs:
+            t_sim = float(t_switch.subs(overrides[sim_id]))
+            values = {
+                t: expr.subs(overrides[sim_id]).subs(time, t)
+                for t in (0.0, t_sim)
+            }
+            if preeq_id:
+                values[v2.C.TIME_PREEQUILIBRATION] = expr.subs(
+                    overrides[preeq_id]
+                ).subs(time, 0)
+            elif {str(s) for s in values[t_sim].free_symbols} <= parameter_ids:
+                # the model already has the pre-switch value, and the switch
+                #  is valid as first period (only parameter table symbols)
+                del values[0.0]
+            for t, value in values.items():
+                cond_id = new_conditions.setdefault(
+                    (target_id, value), f"cond_{len(new_conditions)}"
+                )
+                periods[sim_id, preeq_id].append((t, cond_id))
+
+    problem.condition_df = pd.concat(
+        [
+            condition_df.drop(columns=replaced_columns, errors="ignore"),
+            pd.DataFrame.from_dict(
+                {
+                    cid: {
+                        target_id: float(value)
+                        if value.is_number
+                        else petab_math_str(value)
+                    }
+                    for (target_id, value), cid in new_conditions.items()
+                },
+                orient="index",
+            ),
+        ]
+    ).rename_axis(v1.C.CONDITION_ID)
+    return {pair: sorted(p) for pair, p in periods.items()}
+
+
 def v1v2_condition_df(
     condition_df: pd.DataFrame, model: v1.Model
 ) -> pd.DataFrame: