ExperimentsToSbmlConverter: write one flat initial assignment per target

petab/v2/converters.py

@@ -199,8 +199,15 @@ def convert(self) -> Problem:
 
         self._add_preequilibration_indicator()
 
+        # Collect values for initial assignments for the different experiments.
+        #  All expressions must be combined into a single initial assignment
+        #  per target.
+        # target_id -> [(experiment_indicator, target_value), ...]
+        period0_assignments: dict[str, list[tuple[str, sp.Basic]]] = {}
         for experiment in self._new_problem.experiments:
-            self._convert_experiment(experiment)
+            self._convert_experiment(experiment, period0_assignments)
+
+        self._create_initial_assignments(period0_assignments)
 
         self._add_indicators_to_conditions()
 
@@ -209,9 +216,20 @@ def convert(self) -> Problem:
 
         return self._new_problem
 
-    def _convert_experiment(self, experiment: Experiment) -> None:
+    def _convert_experiment(
+        self,
+        experiment: Experiment,
+        period0_assignments: dict[str, list[tuple[str, sp.Basic]]],
+    ) -> None:
         """
         Convert a single experiment to SBML events or initial assignments.
+
+        :param experiment: The experiment to convert.
+        :param period0_assignments: The changes of the first period are not
+            applied directly, but are added to this dictionary
+            (``target_id -> [(experiment_indicator, target_value), ...]``).
+            They are combined across experiments into a single initial
+            assignment per target by :meth:`_create_initial_assignments`.
         """
         model = self._model
         experiment.sort_periods()
@@ -228,11 +246,6 @@ def _convert_experiment(self, experiment: Experiment) -> None:
             )
         add_sbml_parameter(model, id_=exp_ind_id, constant=False, value=0)
         kept_periods: list[ExperimentPeriod] = []
-        # Collect values for initial assignments for the different experiments.
-        #  All expressions must be combined into a single initial assignment
-        #  per target.
-        # target_id -> [(experiment_indicator, target_value), ...]
-        period0_assignments: dict[str, list[tuple[str, sp.Basic]]] = {}
 
         for i_period, period in enumerate(experiment.sorted_periods):
             if period.is_preequilibration:
@@ -281,60 +294,6 @@ def _convert_experiment(self, experiment: Experiment) -> None:
                     self._new_problem.get_changes_for_period(period),
                 )
 
-        # Create initial assignments for the first period
-        if period0_assignments:
-            free_symbols_in_assignments = set()
-            for target_id, changes in period0_assignments.items():
-                # The initial value might only be changed for a subset of
-                #  experiments. We need to keep the original initial value
-                #  for all other experiments.
-
-                # Is there an initial assignment for this target already?
-                # If not, fall back to the initial value of the target.
-                if (
-                    ia := model.getInitialAssignmentBySymbol(target_id)
-                ) is not None:
-                    default = sbml_math_to_sympy(ia.getMath())
-                else:
-                    # use the initial value of the target as default
-                    target = model.getElementBySId(target_id)
-                    default = self._initial_value_from_element(target)
-
-                # Only create the initial assignment if there is
-                #  actually something to change.
-                if expr_cond_pairs := [
-                    (target_value, sp.Symbol(exp_ind) > 0.5)
-                    for exp_ind, target_value in changes
-                    if target_value != default
-                ]:
-                    # Unlike events, we can't have different initial
-                    #  assignments for different experiments, so we need to
-                    #  combine all changes into a single piecewise
-                    #  expression.
-
-                    expr = sp.Piecewise(
-                        *expr_cond_pairs,
-                        (default, True),
-                    )
-
-                    # Create a new initial assignment if necessary, otherwise
-                    #  overwrite the existing one.
-                    if ia is None:
-                        ia = model.createInitialAssignment()
-                        ia.setSymbol(target_id)
-
-                    set_math(ia, expr)
-                    free_symbols_in_assignments |= expr.free_symbols
-
-            # the target value may depend on parameters that are only
-            #  introduced in the PEtab parameter table - those need
-            #  to be added to the model
-            for sym in free_symbols_in_assignments:
-                if model.getElementBySId(sym.name) is None:
-                    add_sbml_parameter(
-                        model, id_=sym.name, constant=True, value=0
-                    )
-
         if len(kept_periods) > 2:
             raise AssertionError("Expected at most two periods to be kept.")
 
@@ -349,6 +308,66 @@ def _convert_experiment(self, experiment: Experiment) -> None:
 
         experiment.periods = kept_periods
 
+    def _create_initial_assignments(
+        self, period0_assignments: dict[str, list[tuple[str, sp.Basic]]]
+    ) -> None:
+        """Create initial assignments for the first period of all experiments.
+
+        Unlike events, we can't have different initial assignments for
+        different experiments, so all changes of a target are combined into
+        a single piecewise expression, conditioned on the experiment
+        indicators.
+
+        :param period0_assignments: The changes of the first period of all
+            experiments
+            (``target_id -> [(experiment_indicator, target_value), ...]``).
+        """
+        model = self._model
+        free_symbols_in_assignments = set()
+        for target_id, changes in period0_assignments.items():
+            # The initial value might only be changed for a subset of
+            #  experiments. We need to keep the original initial value
+            #  for all other experiments.
+
+            # Is there an initial assignment for this target already?
+            # If not, fall back to the initial value of the target.
+            if (
+                ia := model.getInitialAssignmentBySymbol(target_id)
+            ) is not None:
+                default = sbml_math_to_sympy(ia.getMath())
+            else:
+                # use the initial value of the target as default
+                target = model.getElementBySId(target_id)
+                default = self._initial_value_from_element(target)
+
+            # Only create the initial assignment if there is
+            #  actually something to change.
+            if expr_cond_pairs := [
+                (target_value, sp.Symbol(exp_ind) > 0.5)
+                for exp_ind, target_value in changes
+                if target_value != default
+            ]:
+                expr = sp.Piecewise(
+                    *expr_cond_pairs,
+                    (default, True),
+                )
+
+                # Create a new initial assignment if necessary, otherwise
+                #  overwrite the existing one.
+                if ia is None:
+                    ia = model.createInitialAssignment()
+                    ia.setSymbol(target_id)
+
+                set_math(ia, expr)
+                free_symbols_in_assignments |= expr.free_symbols
+
+        # the target value may depend on parameters that are only
+        #  introduced in the PEtab parameter table - those need
+        #  to be added to the model
+        for sym in free_symbols_in_assignments:
+            if model.getElementBySId(sym.name) is None:
+                add_sbml_parameter(model, id_=sym.name, constant=True, value=0)
+
     @staticmethod
     def _initial_value_from_element(target: libsbml.SBase) -> sp.Basic:
         """Get the initial value of an SBML element.

tests/v2/test_converters.py

@@ -1,6 +1,8 @@
 from math import inf
 
 import pandas as pd
+import sympy as sp
+from sbmlmath import sbml_math_to_sympy
 
 from petab.v2 import Change, Condition, Experiment, ExperimentPeriod, Problem
 from petab.v2.converters import ExperimentsToSbmlConverter
@@ -79,6 +81,58 @@ def test_experiments_to_events_converter():
     ]
 
 
+def test_initial_assignments_many_experiments():
+    """The first-period changes of all experiments are combined into a single
+    flat piecewise initial assignment per target.
+
+    Nesting the piecewise once per experiment exceeded lxml's depth limit
+    (https://github.com/PEtab-dev/libpetab-python/issues/540).
+    """
+    n_experiments = 300
+    ant_model = """
+    species A = 1
+    species B
+    k = 3
+    B = 2 * k
+    """
+    problem = Problem()
+    problem.model = SbmlModel.from_antimony(ant_model)
+    for i in range(n_experiments):
+        problem.add_condition(f"c{i}", A=i + 2, B=i + 2)
+        problem.add_experiment(f"e{i}", 0, f"c{i}")
+    # an experiment that does not change the initial values
+    problem.add_condition("c_default", k=3)
+    problem.add_experiment("e_default", 0, "c_default")
+
+    converted = ExperimentsToSbmlConverter(problem).convert()
+    sbml_model = converted.model.sbml_model
+
+    indicators = {
+        sp.Symbol(ExperimentsToSbmlConverter.get_experiment_indicator(exp.id))
+        for exp in problem.experiments
+    }
+    for target_id, default in (("A", 1), ("B", 2 * sp.Symbol("k"))):
+        expr = sbml_math_to_sympy(
+            sbml_model.getInitialAssignmentBySymbol(target_id).getMath()
+        )
+        assert isinstance(expr, sp.Piecewise)
+        assert len(expr.args) == n_experiments + 1
+        assert not any(isinstance(e, sp.Piecewise) for e, _ in expr.args)
+
+        for exp_id, expected in (
+            ("e0", 2),
+            (f"e{n_experiments - 1}", n_experiments + 1),
+            ("e_default", default),
+        ):
+            active = sp.Symbol(
+                ExperimentsToSbmlConverter.get_experiment_indicator(exp_id)
+            )
+            assert (
+                expr.xreplace({ind: int(ind == active) for ind in indicators})
+                == expected
+            )
+
+
 def test_simulate_experiment_to_events():
     """
     Convert PEtab experiment to SBML events and compare BasiCO simulation